#=
Solves the price function for the Lucas tree in a continuous state
setting, using piecewise linear approximation for the sequence of
candidate price functions.  The consumption endownment follows the
log linear AR(1) process

    log y' = alpha log y + sigma epsilon

where y' is a next period y and epsilon is an iid standard normal
shock. Hence

    y' = y^alpha * xi   where xi = e^(sigma * epsilon)

The distribution phi of xi is

    phi = LN(0, sigma^2) where LN means lognormal

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date : 2014-07-05

References
----------

http://quant-econ.net/jl/markov_asset.html

TODO: refactor. Python is much cleaner.
=#
using Distributions
using Grid: CoordInterpGrid, BCnan, BCnearest, InterpLinear

"""
The Lucas asset pricing model

##### Fields

-  `gam::Real` : coefficient of risk aversion in the CRRA utility function
-  `bet::Real` : Discount factor in (0, 1)
-  `alpha::Real` : Correlation coefficient in the shock process
-  `sigma::Real` : Volatility of shock process
-  `phi::Distribution` : Distribution for shock process
-  `grid::AbstractVector` : Grid of points on which to evaluate the prices. Each
point should be non-negative
-  `grid_min::Real` : Lower bound on grid
-  `grid_max::Real` : Upper bound on grid
-  `grid_size::Int` : Number of points in the grid
- `quad_nodes::Vector` : Quadrature nodes for integrating over the shock
- `quad_weights::Vector` : Quadrature weights for integrating over the shock
-  `h::Vector` : Storage array for the `h` vector in the lucas operator
"""
type LucasTree
    gam::Real
    bet::Real
    alpha::Real
    sigma::Real
    phi::Distribution
    grid::AbstractVector
    grid_min::Real
    grid_max::Real
    grid_size::Int
    quad_nodes::Vector
    quad_weights::Vector
    h::Vector

    # this needs to be an internal constructor because we need to incompletely
    # initialize the object before we can compute h.
    """
    Constructor for LucasTree

    ##### Arguments

    -  `gam::Real` : coefficient of risk aversion in the CRRA utility function
    -  `bet::Real` : Discount factor in (0, 1)
    -  `alpha::Real` : Correlation coefficient in the shock process
    -  `sigma::Real` : Volatility of shock process

    ##### Notes

    All other fields of the type are instantiated within the constructor
    """
    function LucasTree(gam::Real, bet::Real, alpha::Real, sigma::Real)
        phi = LogNormal(0.0, sigma)
        grid = make_grid(alpha, sigma)
        grid_min, grid_max, grid_size = minimum(grid), maximum(grid), length(grid)
        _int_min, _int_max = exp(-4 * sigma), exp(4 * sigma)
        n, w = qnwlege(21, _int_min, _int_max)

        # create lt object without h
        lt = new(gam, bet, alpha, sigma, phi, grid, grid_min, grid_max, grid_size,
                 n, w)

        # initialize h
        h = Array(Float64, grid_size)

        for (i, y) in enumerate(grid)
            integrand(z) = (y^alpha * z).^(1 - gam)
            h[i] = bet .* integrate(lt, integrand)
        end

        # now add h to it
        lt.h = h
        lt
    end
end

function make_grid(alpha, sigma)
    grid_size = 100
    if abs(alpha) >= 1
        grid_min, grid_max = 0.0, 10.
    else
        # Set the grid interval to contain most of the mass of the
        # stationary distribution of the consumption endowment
        ssd = sigma / sqrt(1 - alpha^2)
        grid_min, grid_max = exp(-4 * ssd), exp(4 * ssd)
    end
    return linspace_range(grid_min, grid_max, grid_size)
end

function integrate(lt::LucasTree, g::Function, qn::Array=lt.quad_nodes,
                   qw::Array=lt.quad_weights)
    int_func(x) = g(x) .* pdf(lt.phi, x)
    return do_quad(int_func, qn, qw)
end

"""
The approximate Lucas operator, which computes and returns the updated function
Tf on the grid points.

##### Arguments

- `lt::LucasTree` : An instance of the `LucasTree` type
- `f::Vector{Float64}` : A candidate function on R_+ represented as points on a
grid. It should be the same size as `lt.grid`

##### Returns

- `Tf::Vector{Float64}` : The updated function Tf

"""
function lucas_operator(lt::LucasTree, f::AbstractVector)
    grid, h, alpha, bet = lt.grid, lt.h, lt.alpha, lt.bet

    Tf = similar(f)
    Af = CoordInterpGrid(grid, f, BCnearest, InterpLinear)

    for (i, y) in enumerate(grid)
        to_integrate(z) = Af[y^alpha * z]
        Tf[i] = h[i] + bet * integrate(lt, to_integrate)
    end
    return Tf
end

"""
Compute the equilibrium price function associated with Lucas tree `lt`

##### Arguments

- `lt::LucasTree` : An instance of the `LucasTree` type
- `;kwargs...` : other arguments to be passed to `compute_fixed_point`

##### Returns

- `price::Vector{Float64}` : The price at each point in `lt.grid`

"""
function compute_lt_price(lt::LucasTree; kwargs...)
    # Simplify names, set up distribution phi
    grid, grid_size, gam = lt.grid, lt.grid_size, lt.gam

    f_init = zeros(grid)  # Initial condition
    func(x) = lucas_operator(lt, x)
    f = compute_fixed_point(func, f_init; max_iter=500, verbose=false, kwargs...)

    # p(y) = f(y) / u'(y) = f(y) * y^gamma
    price = f .* grid.^(gam)

    return price
end


using SIMD

@inline function vdot(xs, ys, ::Type{Vec{N,T}} = Vec{4,Float64}) where {N,T}
    @assert length(ys) == length(xs)
    @assert length(xs) % N == 0
    lane = VecRange{N}(0)
    vacc = zero(Vec{N, T})
    @inbounds for i in 1:N:length(xs)
        vacc = muladd(xs[lane + i], ys[lane + i], vacc)
    end
    return sum(vacc)
end


@doc raw"""
    LaplaceLikelihood(β::T=1.0)  #  Laplace likelihood with scale β

Laplace likelihood for regression:
```math
\frac{1}{2\beta} \exp(-\frac{|y-f|}{β})
```
see [wiki page](https://en.wikipedia.org/wiki/Laplace_distribution)
---
For the analytical solution, it is augmented via:
```math
p(y|f,ω) = N(y|f,ω⁻¹)
```
where ``ω \sim \text{Exp}(ω | 1/(2 β^2))``, and ``\text{Exp}`` is the [Exponential distribution](https://en.wikipedia.org/wiki/Exponential_distribution)
We use the variational distribution ``q(ω) = GIG(ω|a,b,p)``
"""
struct LaplaceLikelihood{T<:Real} <: RegressionLikelihood
    β::T
    a::T
    p::T
    function LaplaceLikelihood{T}(β::T) where {T<:Real}
        return new{T}(β, β^-2, 0.5)
    end
end

function LaplaceLikelihood(β::T=1.0) where {T<:Real}
    return LaplaceLikelihood{T}(β)
end

function implemented(
    ::LaplaceLikelihood, ::Union{<:AnalyticVI,<:QuadratureVI,<:GibbsSampling}
)
    return true
end

function (l::LaplaceLikelihood)(y::Real, f::Real)
    return pdf(Laplace(f, l.β), y)
end

function Distributions.loglikelihood(l::LaplaceLikelihood, y::Real, f::Real)
    return logpdf(Laplace(f, l.β), y)
end

function Base.show(io::IO, l::LaplaceLikelihood)
    return print(io, "Laplace likelihood (β=$(l.β))")
end

function compute_proba(
    l::LaplaceLikelihood, μ::AbstractVector{<:Real}, σ²::AbstractVector{<:Real}
)
    return μ, max.(σ², zero(eltype(σ²))) .+ 2 * l.β^2
end

## Local Updates ##
# b : Variational parameter b of GIG
# θ : Expected value of ω
function init_local_vars(::LaplaceLikelihood, batchsize::Int, T::DataType=Float64)
    return (; b=rand(T, batchsize), θ=zeros(T, batchsize))
end

function local_updates!(
    local_vars,
    l::LaplaceLikelihood,
    y::AbstractVector,
    μ::AbstractVector,
    diagΣ::AbstractVector,
)
    @. local_vars.b = diagΣ + abs2(μ - y)
    @. local_vars.θ = sqrt(l.a) / sqrt.(local_vars.b)
    return local_vars
end

function sample_local!(
    local_vars, l::LaplaceLikelihood, y::AbstractVector, f::AbstractVector
)
    @. local_vars.b = rand(GeneralizedInverseGaussian(1 / l.β^2, abs2(f - y), 0.5))
    @. local_vars.θ = inv(local_vars.b)
    return local_vars
end

@inline function ∇E_μ(::LaplaceLikelihood, ::AOptimizer, y::AbstractVector, state)
    return (state.θ .* y,)
end
@inline function ∇E_Σ(::LaplaceLikelihood, ::AOptimizer, ::AbstractVector, state)
    return (0.5 * state.θ,)
end

## ELBO ##
function expec_loglikelihood(
    ::LaplaceLikelihood,
    ::AnalyticVI,
    y::AbstractVector,
    μ::AbstractVector,
    diag_cov::AbstractVector,
    state,
)
    tot = -0.5 * length(y) * log(twoπ)
    tot += 0.5 * Zygote.@ignore(sum(log, state.θ))
    tot +=
        -0.5 * (
            dot(state.θ, diag_cov) + dot(state.θ, abs2.(μ)) - 2.0 * dot(state.θ, μ .* y) +
            dot(state.θ, abs2.(y))
        )
    return tot
end

function AugmentedKL(l::LaplaceLikelihood, state, ::Any)
    return GIGEntropy(l, state) - expecExponentialGIG(l, state)
end

GIGEntropy(l::LaplaceLikelihood, state) = GIGEntropy(l.a, state.b, l.p)

function expecExponentialGIG(l::LaplaceLikelihood, state)
    return sum(
        -log(2 * l.β^2) .-
        0.5 * (l.a .* sqrt.(state.b) + state.b .* sqrt(l.a)) ./ (l.a .* state.b * l.β^2),
    )
end

## PDF and Log PDF Gradients ##

function grad_quad(
    likelihood::LaplaceLikelihood, y::Real, μ::Real, σ²::Real, inference::AbstractInference
)
    nodes = inference.nodes * sqrt(σ²) .+ μ
    Edloglike = dot(inference.weights, ∇loglikehood.(likelihood, y, nodes))
    Ed²loglike = (1 / sqrt(twoπ * σ²)) / (likelihood.β^2)
    return -Edloglike, Ed²loglike
end

@inline function ∇loglikehood(l::LaplaceLikelihood, y::Real, f::Real)
    return sign(y - f) ./ l.β
end

@inline hessloglikelihood(::LaplaceLikelihood, ::Real, f::T) where {T<:Real} = zero(T)


module FastPow
export @fastpow

include("addition-chain-pow.jl")
include("exp_n.jl")

fastpow(x) = x
function fastpow(expr::Expr)
    if Meta.isexpr(expr, :call, 3) && expr.args[1] == :^
        if expr.args[2] === 1
            return :($exp_1($(expr.args[3])))
        elseif expr.args[2] === -1
            return :($exp_minus1($(expr.args[3])))
        elseif expr.args[2] === 2
            if expr.args[3] isa Integer && expr.args[3] < 0
                return :(exp2($(expr.args[3])))
            else
                return :($exp_2($(expr.args[3])))
            end
        elseif expr.args[2] === 10
            if expr.args[3] isa Integer && expr.args[3] < 0
                return :(exp10($(expr.args[3])))
            else
                return :($exp_10($(expr.args[3])))
            end
        elseif expr.args[3] isa Integer
            return addition_chain_pow(expr.args[2], expr.args[3])
        end
    end
    return Expr(expr.head, fastpow.(expr.args)...)
end

macro fastpow(expr)
    return esc(fastpow(expr))
end

end # module


@testset "Horn" begin
    p, q = synthetic_gradient(SynthSphere(50))
    Z = Horn()(p, q)
    @test Z[38,38] ≈ -0.0027751467456787315
    @test Z[1,1] ≈ 0.0
    @test Z[76,76] ≈ 47.05959259743248
end


function error_all_different(x::V; param = nothing, dom_size = 0
) where {T <: Number, V <: AbstractVector{T}}
    acc = Dictionary{T, Int}()
    foreach(y -> insert_or_inc(acc, y), x)
    return Float64(sum(acc .- 1))
end

concept_all_different(x) = allunique(x)

"""
    _all_different
Global constraint ensuring that all the values of a given configuration are unique.
"""
const all_different = Constraint(
    concept = concept_all_different,
    error = make_error(:all_different),
    syms = Set([:permutable]),
)


"""
    module Lattice

for defining discrete Cartesian and Polar lattice grids used in Iris.
"""
module Lattices

using ...Defaults
using Formatting
using RecipesBase
using StaticArrays

export Lattice,
Cartesian,
Polar,
get_lattice_index,
LatticeError

struct Lattice
    type::Symbol

    dx::Float64
    dy::Float64
    θ::Float64
    x0::Float64
    y0::Float64

    dr::Float64
    dθ::Float64
    r0::Float64

    sinθ::Float64 # to avoid inefficient trig calls later
    cosθ::Float64 # to avoid inefficient trig calls later
    v1::SArray{Tuple{2},Float64,1,2}
    v2::SArray{Tuple{2},Float64,1,2}
    origin::SArray{Tuple{2},Float64,1,2}

    function Lattice(type,args...)
        if type==:Cartesian
            CartesianLat(args...)
        else
            PolarLat(args...)
        end
    end

    function CartesianLat(dx::Number,dy::Number,θ::Number=0,x0::Number=0,y0::Number=0)
        sinθ, cosθ = sincos(θ)
        R = @SMatrix [ cos(θ) -sin(θ); sin(θ)  cos(θ)]
        v1 = R*SVector{2}([1,0])
        v2 = R*SVector{2}([0,1])
        origin = SVector{2}([x0,y0])
        return new(:Cartesian,dx, dy, θ, x0, y0, NaN, NaN, NaN, sinθ, cosθ, v1, v2, origin)
    end

    function PolarLat(dr::Number,dθ::Number,θ::Number=0,x0::Number=0,y0::Number=0,r0::Number=dr/2)
        return new(:Polar,NaN,NaN,θ,x0,y0,dr,dθ,r0)
    end

    # pretty printing of Lattice
    function Base.show(io::IO, lat::Lattice)
        if lat.type==:Cartesian
            if !get(io, :sub, false)
                print(io, "Cartesian Lattice", ": \n",
                "\tdx=", fmt("3.2f",lat.dx), ", dy=", fmt("3.2f",lat.dy), ", ∠", fmt("3.2f",(mod2pi(lat.θ+π/2))*180/π), "°\n",
                "\torigin: (", fmt("3.2f",lat.x0), ", ", fmt("3.2f",lat.y0), ")")
            elseif !get(io, :sub1, false)
                print(io,
                "\n\tdx=", fmt("3.2f",lat.dx), ", dy=", fmt("3.2f",lat.dy), ", ∠", fmt("3.2f",(mod2pi(lat.θ+π/2))*180/π), "°\n",
                "\torigin: (", fmt("3.2f",lat.x0), ", ", fmt("3.2f",lat.y0), ")")
            else
                print(io,
                "\n\t\tdx=", fmt("3.2f",lat.dx),", dy=", fmt("3.2f",lat.dy), ", ∠", fmt("3.2f",(mod2pi(lat.θ+π))*180/π), "°\n",
                "\t\torigin: (", fmt("3.2f",lat.x0), ", ", fmt("3.2f",lat.y0), ")")
            end
        elseif lat.type==:Polar
            if !get(io, :sub, false)
                print(io, "Polar Lattice", ": \n",
                "\tdr=", fmt("3.2f",lat.dr), ", dθ=", fmt("3.2f",lat.dθ), "=π/", fmt("3.1f",π/lat.dθ),", ∠", fmt("3.2f",(mod2pi(lat.θ))*180/π), "°\n",
                "\torigin: (", fmt("3.2f",lat.x0), ", ", fmt("3.2f",lat.y0), ")")
            elseif !get(io, :sub1, false)
                print(io,
                "\n\tdr=", fmt("3.2f",lat.dr), ", dθ=", fmt("3.2f",lat.dθ), "=π/", fmt("3.1f",π/lat.dθ),", ∠", fmt("3.2f",(mod2pi(lat.θ))*180/π), "°\n",
                "\torigin: (", fmt("3.2f",lat.x0), ", ", fmt("3.2f",lat.y0), ")")
            else
                print(io,
                "\n\t\tdr=", fmt("3.2f",lat.dr), ", dθ=", fmt("3.2f",lat.dθ), "=π/", fmt("3.1f",π/lat.dθ),", ∠", fmt("3.2f",(mod2pi(lat.θ))*180/π), "°\n",
                "\t\torigin: (", fmt("3.2f",lat.x0), ", ", fmt("3.2f",lat.y0), ")")
            end
        else
            println(io, "unknown lattice type")
        end
    end
end
"""
    Lattices.Cartesian(dx,dy,θ=0,x0=0,y0=0) -> lattice

`x0,y0` is the origin, `θ` is the rotation

----
    (::Cartesian)(i,j) -> x,y

`x,y` are coordinates of the `i,j` lattice site (`i,j` continuous)
--------
    (lat::Cartesian)(dr=lat.dr,...)
"""
Cartesian(args...) = Lattice(:Cartesian,args...)

"""
    Lattices.Polar(dr,dθ,θ=0,r0=dr/2,x0=0,y0=0) -> lattice

`r0` is the minimum radius, `x0,y0` is the origin, `θ` is the zero angle

----
    (::Polar)(i,j) -> x,y

`x,y` are coordinates of the `i,j` lattice site (`i,j` continuous)

--------
    (lat::Polar)(dr=lat.dr,...)
"""
Polar(args...) = Lattice(:Polar,args...)

"""
    (::Lattice)(i,j) -> x,y
"""
function (lat::Lattice)(i::Real,j::Real)
    if lat.type==:Cartesian
        v = lat.origin + i*lat.v1*lat.dx + j*lat.v2*lat.dy
        return v[1],v[2]
    elseif lat.type==:Polar
        r = lat.r0 + i*lat.dr
        θ = lat.θ  + j*lat.dθ
        v = r.*sincos(θ)
        return lat.x0+v[2],lat.y0+v[1]
    end
end

# functor wrappers for both Cartesian and Polar lattices
(lat::Lattice)(ijk::Tuple) = lat(ijk...)
(lat::Lattice)(k::CartesianIndex) = lat(k[1],k[2])
function (lat::Lattice)(i::Real,j::Real,v::AbstractArray)
    v[:] = lat(i,j)
    return nothing
end
function (lat::Lattice)(x::AbstractArray,y::AbstractArray)
    X = Array{Float64}(undef,size(x))
    Y = Array{Float64}(undef,size(y))
    for i ∈ eachindex(X)
        X[i], Y[i] = lat(x[i],y[i])
    end
    return X,Y
end

function (lat::Lattice)(;dx=lat.dx, dy=lat.dy, θ=lat.θ, x0=lat.x0, y0=lat.y0, dr=lat.dr, dθ=lat.dθ, r0=dr/2)
    if lat.type==:Cartesian
        return Cartesian(dx,dy,θ,x0,y0)
    else
        return Polar(dr,dθ,θ,x0,y0,r0)
    end
end

struct LatticeError <: Exception
    lat::Lattice
end
function Base.showerror(io::IO,err::LatticeError)
    print(io,"LatticeError: unrecognized lattice type ")
    print(io, err.lat.type)
    print(io, ", must be one of :Cartesian or :Polar")
end

"""
    get_lattice_index(lattice,x,y) -> nx,ny

returns indices `nx`,`ny` such that `(x,y) = (x0,y0) + nx*v1*dx + ny*v2*dx`
"""
function get_lattice_index(lat::Lattice,x::AbstractArray,y::AbstractArray)
    @assert size(x)==size(y) "x and y must have the same size"
    X = Array{Float64}(undef,size(x)...)
    Y = Array{Float64}(undef,size(y)...)
    for k ∈ eachindex(X)
        X[k],Y[k] = get_lattice_index(lat,x[k],y[k])
    end
    return X,Y
end
function get_lattice_index(lat::Lattice,x::Real,y::Real)
    i,j = lat.type==:Cartesian ? get_lattice_index_cartesian(lat,x,y) : get_lattice_index_polar(lat,x,y)
    return i,j
end
function get_lattice_index_cartesian(lat::Lattice,x::Real,y::Real)
    dx, dy, v1, v2, x0, y0 = lat.dx, lat.dy, lat.v1, lat.v2, lat.x0, lat.y0
    x += -x0
    y += -y0
    rv1 = v1[1]*x + v1[2]*y
    rv2 = v2[1]*x + v2[2]*y
    return rv1/dx, rv2/dy
end
function get_lattice_index_polar(lat::Lattice,x::Real,y::Real)
    dr, dθ, x0, y0 = lat.dr, lat.dθ, lat.x0, lat.y0
    x += -x0
    y += -y0
    r, θ = hypot(y,x), atan(y,x)+π
    return (r-lat.r0)/dr, (θ-lat.θ)/dθ
end

"""
    plot(lattice, N=[4,4])

scatter plot of lattice vectors from -`N[i]`:`+N[i]` for each dimension i.

for polar lattices, plots all angles out to radial site `N[1]`.

If `N` is a scalar, uses same span for both dimensions.
"""
@recipe f(lattice::Lattice, N::Int=4) = (lattice, [N,N])
@recipe function f(lattice::Lattice, N::Array{Int,1})
    lattice.type==:Polar ? N[2]=floor(Int,2π/lattice.dθ-eps(2π/lattice.dθ)) : nothing
    seriestype --> :scatter
    legend --> false
    aspect_ratio --> 1
    @series begin
        x, y = Float64[], Float64[]
        for n ∈ 0:N[1], m ∈ 0:N[2]
            temp = lattice(n,m)
            push!(x,temp[1]); push!(y,temp[2])
        end
        (x, y)
    end
end

end # module


export convCircKernel ,getConvCircKernel
## For the functions nImgIn, nImgOut, nFeatIn, nFeatOut, nTheta, getOp, initTheta : see AbstractConvKernel.jl
## All convKernel types are assumed to have fields nImage and sK
mutable struct convCircKernel{T} <: AbstractConvKernel{T}
    nImg :: Array{Int,1}
    sK   :: Array{Int,1}
    iK   :: Array{Int,1}
end

function nImgOut(this::convCircKernel)
    return [this.nImg[1]; this.nImg[2]; this.sK[3]]
 end

function getConvCircKernel(TYPE::Type,nImg,sK)

    n3 = [nImg[1],nImg[2],sK[3]]
    Ind = zeros(n3[1],n3[2],n3[3])
    mid = div.(n3,2)
    midK = div.(sK[1:2],2)
    Ind[mid[1]-midK[1]:mid[1]+midK[1],
        mid[2]-midK[2]:mid[2]+midK[2],:] = reshape(collect(1:prod(sK)),sK[1],sK[2],sK[3])
    Ind = circshift(Ind,1-mid)

    iK = find(Ind)
	return convCircKernel{TYPE}(nImg,sK,iK)
end

export Amv
function Amv(this::convCircKernel{T},theta::Array{T},Yin::Array{T}) where {T<:Number}

    nex   = div(length(Y),prod(nImgIn(this)))

    n     = nImgIn(this);

    Y     = reshape(Yin,(n[1],n[2],this.sK[3],nex))
    # compute convolution
    K = zeros(n[1],n[2],this.sK[3])

    K[this.iK] = theta[:];

    Kh = fft(K,(1,2,3))
    Yh = fft(Y,(1,2,3))

    Z  = real(ifft(Kh.*Yh,(3,2,1)))
    return Z
end

export ATmv
function ATmv(this::convCircKernel{T},theta::Array{T},Zin::Array{T}) where {T<:Number}

    nex   =  div(numel(Z),prod(nImgIn(this)))
    n     = nImgIn(this)
    Z     = reshape(Zin,(n[1],n[2],this.sK[3],nex));
    # compute convolution T
    K = zeros(n[1],n[2],this.sK[3])
    K[this.iK] = theta

    Kh = fft(K,(1,2,3))
    Zh = fft(Z,(1,2,3))

    Y  = real(ifft(conj(Kh).*Zh,(3,2,1)))
    return Y
end

function Jthetamv(this::convCircKernel{T},dtheta::Array{T},dummy::Array{T},Yin::Array{T},temp=nothing) where {T<:Number}

    nex    =  div(numel(Y),nFeatIn(this))
    Y      = reshape(Yin,:,nex)
    Z      = Amv(this,dtheta,Y)
    return Z
end

function JthetaTmv(this::convCircKernel{T},Zin::Array{T},dummy::Array{T},Yin::Array{T}) where {T<:Number}
    #  derivative of Z*(A(theta)*Y) w.r.t. theta

    nex  =  div(numel(Y),nFeatIn(this))

    n  = nImgIn(this)

    Z     = reshape(Zin,n[1],n[2],this.sK[3],nex)
    Y     = reshape(Yin,n[1],n[2],this.sK[3],nex)

    Zh = fft(Z,(1,2,3))
    Yh = fft(Y,(1,2,3))
    YZ = sum(conj(Yh).*Zh,4)
    K  = real(ifft(YZ,(3,2,1)))
    dtheta = K[this.iK];

end


"""
$(TYPEDEF)
$(TYPEDFIELDS)

"""
struct Normal <: AbstractDistDelayed
    mu::NestedFloat
    sigma::NestedFloat
end

function normal(mu::Float64, sigma::Float64)::Float64
    return first(rand(Distributions.Normal(mu, sigma), 1))
end


using KernelFunctions
using MLDataUtils
using Zygote
using Flux
using Distributions, LinearAlgebra
using Plots

Flux.@functor SqExponentialKernel
Flux.@functor ScaleTransform
Flux.@functor KernelSum
Flux.@functor Matern32Kernel

xmin = -3;
xmax = 3;
x = range(xmin, xmax; length=100)
x_test = range(xmin, xmax; length=300)
x, y = noisy_function(sinc, x; noise=0.1)
X = reshape(x, :, 1)
X_test = reshape(x_test, :, 1)
k = SqExponentialKernel(1.0)#+Matern32Kernel(2.0)
λ = [-1.0]
function f(x, k, λ)
    return kernelmatrix(k, x, X; obsdim=1) *
           inv(kernelmatrix(k, X; obsdim=1) + exp(λ[1]) * I) *
           y
end
f(X, k, 1.0)
loss(k, λ) = ŷ -> sum(y - ŷ) / length(y) + exp(λ[1]) * norm(ŷ)(f(X, k, λ))
loss(k, λ)
ps = Flux.params(k)
push!(ps, λ)
opt = Flux.Momentum(0.1)
##
for i in 1:10
    grads = Zygote.gradient(() -> loss(k, λ), ps)
    Flux.Optimise.update!(opt, ps, grads)
    p = Plots.scatter(x, y; lab="data", title="Loss = $(loss(k,λ))")
    Plots.plot!(x_test, f(X_test, k, λ); lab="Prediction", lw=3.0)
    display(p)
end


using Test
using IntervalArithmetic
using TaylorModels
# using Revise
using ImplicitDomainQuadrature

set_variables(Float64, "x", order = 2, numvars = 2)
box = IntervalBox(-1..1,2)
x0 = IntervalBox(mid(box))
tm = TaylorModelN(1,1,x0,box)
a = 5.0
b = 2.0
@test muladd(tm,a,b) == a*tm + b
@test muladd(a,tm,b) == a*tm + b
@test muladd(tm,tm,b) == tm*tm + b
@test muladd(a,b,tm) == a*b+tm
tm1 = TaylorModelN(1,1,x0,box)
tm2 = TaylorModelN(1,1,x0,box)
tm3 = TaylorModelN(1,1,x0,box)
@test muladd(tm1,a,tm2) == a*tm1 + tm2
@test muladd(tm1,tm2,tm3) == tm1*tm2 + tm3


using MAT
cd(dirname(Base.source_path()))
include("../src/utils.jl")
include("../src/1d/cpab1dIntegration.jl")
file = matopen("cpa_data_nCells_5.mat")
As = read(file, "As")
pts_src = read(file, "pts_src")
pts_src = squeeze(pts_src,2)
tw_args = read(file, "tw_args")
nCols = tw_args["nCols"]
nCells = size(As, 1)
close(file)

nSteps = 100
Nsteps = 10
tess = getTess(As, nCells, nCols)
v = getVelocityField1d(pts_src, As, nCells, nCols)

pts_fwd1 = solver1dParallelIntegration(pts_src, 0.1, tess, As, nCells, nCols, Nsteps, nSteps)
pts_fwd2 = solver1dParallelIntegration(pts_src, 0.5, tess, As, nCells, nCols, Nsteps, nSteps)
@time pts_fwd3 = solver1dParallelIntegration(pts_src, 1.0, tess, As, nCells, nCols, Nsteps, nSteps)
pts_fwd4 = solver1dParallelIntegration(pts_src, 2.0, tess, As, nCells, nCols, Nsteps, nSteps)
pts_fwd5 = solver1dParallelIntegration(pts_src, 8.0, tess, As, nCells, nCols, Nsteps, nSteps)

using PyPlot
PyPlot.figure("1D Integration",figsize=(12,6))
PyPlot.subplot(1,2,1)
PyPlot.title("Velocity Field")
PyPlot.plot(pts_src, v, color="blue")
PyPlot.xlabel("x")
PyPlot.ylabel("v")

PyPlot.subplot(1,2,2)
PyPlot.title("1D Integration for different t")
PyPlot.plot(pts_src, pts_src, label="t=0", color="black")
PyPlot.plot(pts_src, pts_fwd1, label="t=0.1", color="red")
PyPlot.plot(pts_src, pts_fwd2, label="t=0.5", color="orange")
PyPlot.plot(pts_src, pts_fwd3, label="t=1.0", color="yellow")
PyPlot.plot(pts_src, pts_fwd4, label="t=2.0", color="green")
PyPlot.plot(pts_src, pts_fwd5, label="t=8.0", color="blue")
PyPlot.xlabel("x")
PyPlot.ylabel("Transformed x")
PyPlot.legend(loc="lower right",fancybox="true")


@testset "TestWigner" begin

##
using ACE, StaticArrays, ACE.SphericalHarmonics;
using ACEbase
using ACE.SphericalHarmonics: index_y;
using ACE: evaluate
using Random, Printf, Test, ACE.Testing
using ACE, Test, Printf, LinearAlgebra, BenchmarkTools

##

SH = ACE.SphericalHarmonics.SHBasis(3)
for L = 0:3
   @info("Check correctness of Wigner-D matrix of order $L")
   println()

   φ = ACE.SphericalVector(L; T = ComplexF64)

   for ntest = 1:20
      Q, D = ACE.Wigner.rand_QD(L)
      x = randn(SVector{3, Float64})
      x = x / norm(x)
      Y1 = evaluate(SH, x)[(L^2+1):(L+1)^2]
      Dt_Y1_Q = D' * evaluate(SH, Q * x)[(L^2+1):(L+1)^2]
      print_tf(@test isapprox(Y1, Dt_Y1_Q, rtol=1e-10))
   end
   println()
end

end


export i_sum, i_mul!, i_dot, i_axpy!, i_umm!, i_norm2

"""
    i_sum(out!, x)

get the sum of `x`.
"""
@i function i_sum(out!, x::AbstractArray)
	@invcheckoff for i=1:length(x)
		@inbounds out! += x[i]
	end
end

@i function i_sum(out!, f, x::AbstractArray)
	@invcheckoff for i=1:length(x)
		@inbounds out! += f(x[i])
	end
end

"""
    i_mul!(out!, x, y)

compute `x * y` (`x` and `y` are matrices, and store results in `out!`.
"""
@i function i_mul!(out!::AbstractMatrix{T}, x::AbstractMatrix{T}, y::AbstractMatrix{T}) where T
	@safe size(x, 2) == size(y, 1) || throw(DimensionMismatch())
	@invcheckoff @inbounds for k=1:size(y,2)
	    for j=1:size(x,2)
		    for i=1:size(x,1)
				out![i,k] += x[i,j] * y[j,k]
			end
		end
	end
end

@i function i_mul!(out!::AbstractVector{T}, x::AbstractMatrix, y::AbstractVector) where T
	@safe size(x, 2) == size(y, 1) || throw(DimensionMismatch())
	@invcheckoff @inbounds for j=1:size(x,2)
		yj ← zero(T)
		yj += y[j]
		for i=1:size(x,1)
			out![i] += x[i,j] * yj
		end
		yj -= y[j]
	end
end

@i function i_dot(out!, x, y)
    @safe @assert length(x) == length(y)
    @invcheckoff @inbounds for i=1:length(x)
        out! += x[i]' * y[i]
    end
end

"""
    i_norm2(out!, x)

get the squared norm of `x`.
"""
@i function i_norm2(out!, x)
    @invcheckoff @inbounds for i=1:length(x)
        out! += abs2(x[i])
    end
end

"""
    i_axpy!(a, x, y!)

compute `y! += a * x`, where `x` and `y` are vectors.
"""
@i function i_axpy!(a, X, Y)
    @safe @assert length(X) == length(Y)
    @invcheckoff @inbounds for i=1:length(Y)
        Y[i] += a * X[i]
    end
end

"""
    i_umm!(x!, θ)

Compute unitary matrix multiplication on `x`, where the unitary matrix is parameterized by (N+1)*N/2 `θ`s.
"""
@i function i_umm!(x!::AbstractArray, θ)
    M ← size(x!, 1)
    N ← size(x!, 2)
    k ← 0
    @safe @assert length(θ) == M*(M-1)/2
    for l = 1:N
        for j=1:M
            for i=M-1:-1:j
                INC(k)
                ROT(x![i,l], x![i+1,l], θ[k])
            end
        end
    end

    k → length(θ)
end


using DescentMethods
@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

using NLPModels
using LinearOperators, LinearAlgebra
using Krylov

using Stopping
using LineSearch
using DescentMethods

function rosenbrock(x)
	n = 2; m = 2;
	f = []
	push!(f, 10 * (x[2]-x[1]^2))
	push!(f, (x[1]-1))
	return sum(f[i]^2 for i=1:m)
end
x₀ = [-1.2, 1.0]

solvers = [:CG_FR, :CG_HZ, :CG_HS, :CG_PR, :NewtonLDLT, :Newton, :Newlbfgs, :Shamanskii]

for solver in solvers
    nlp = ADNLPModel(rosenbrock, x₀)
    println("Testing $(String(solver))")
    nlpatx = NLPAtX(nlp.meta.x0)
    nlpstop = NLPStopping(nlp, Stopping.unconstrained, nlpatx)

    final_nlp_at_x, optimal = eval(solver)(nlp, nlpstop, verbose = true, linesearch = armijo_ls)
    println("optimal = $(string(optimal))")
    finalize(nlp)
end


using EconPDEs, Distributions

Base.@kwdef mutable struct DiTellaModel
  # Utility Function
  γ::Float64 = 5.0
  ψ::Float64 = 1.5
  ρ::Float64 = 0.05
  τ::Float64 = 0.4

  # Technology
  A::Float64 = 200.0
  σ::Float64 = 0.03

  # MoralHazard
  ϕ::Float64 = 0.2

  # Idiosyncratic
  νbar::Float64 = 0.24
  κν::Float64 = 0.22
  σνbar::Float64 = -0.13
end

function (m::DiTellaModel)(state::NamedTuple, y::NamedTuple)
  (; γ, ψ, ρ, τ, A, σ, ϕ, νbar, κν, σνbar) = m
  (; x, ν) = state
  (; pA, pAx_up, pAx_down, pAν_up, pAν_down, pAxx, pAxν, pAνν, pB, pBx_up, pBx_down, pBν_up, pBν_down, pBxx, pBxν, pBνν, p, px_up, px_down, pν_up, pν_down, pxx, pxν, pνν) = y

  # drift and volatility of state variable ν
  g = p / (2 * A)
  i = A * g^2
  μν = κν * (νbar - ν)
  σν = σνbar * sqrt(ν)
  pAν = (μν >= 0) ? pAν_up : pAν_down
  pBν = (μν >= 0) ? pBν_up : pBν_down
  pν = (μν >= 0) ? pν_up : pν_down

  pAx, pBx, px = pAx_up, pBx_up, px_up
  iter = 0
  @label start
  σX = x * (1 - x) * (1 - γ) / (γ * (ψ - 1)) * (pAν / pA - pBν / pB) * σν / (1 - x * (1 - x) * (1 - γ) / (γ * (ψ - 1)) * (pAx / pA - pBx / pB))
  σpA = pAx / pA * σX + pAν / pA * σν
  σpB = pBx / pB * σX + pBν / pB * σν
  σp = px / p * σX + pν / p * σν
  κ = (σp + σ - (1 - γ) / (γ * (ψ - 1)) * (x * σpA + (1 - x) * σpB)) / (1 / γ)
  κν = γ * ϕ * ν / x
  σA = κ / γ + (1 - γ) / (γ * (ψ - 1)) * σpA
  νA = κν / γ
  σB = κ / γ + (1 - γ) / (γ * (ψ - 1)) * σpB
  # Interest rate r
  μX = x * (1 - x) * ((σA * κ + νA * κν - 1 / pA - τ) - (σB * κ -  1 / pB + τ * x / (1 - x)) - (σA - σB) * (σ + σp))
  if (iter == 0) & (μX <= 0)
    iter += 1
    pAx, pBx, px = pAx_down, pBx_down, px_down
    @goto start
  end

  μpA = pAx / pA * μX + pAν / pA * μν + 0.5 * pAxx / pA * σX^2 + 0.5 * pAνν / pA * σν^2 + pAxν / pA * σX * σν
  μpB = pBx / pB * μX + pBν / pB * μν + 0.5 * pBxx / pB * σX^2 + 0.5 * pBνν / pB * σν^2 + pBxν / pB * σX * σν
  μp = px / p * μX + pν / p * μν + 0.5 * pxx / p * σX^2 + 0.5 * pνν / p * σν^2 + pxν / p * σX * σν
  r = (1 - i) / p + g + μp + σ * σp - κ * (σ + σp) - γ / x * (ϕ * ν)^2

  # Market Pricing
  pAt = - pA * (1 / pA  + (ψ - 1) * τ / (1 - γ) * ((pA / pB)^((1 - γ) / (1 - ψ)) - 1) - ψ * ρ + (ψ - 1) * (r + κ * σA + κν * νA) + μpA - (ψ - 1) * γ / 2 * (σA^2 + νA^2) + (2 - ψ - γ) / (2 * (ψ - 1)) * σpA^2 + (1 - γ) * σpA * σA)
  pBt = - pB * (1 / pB - ψ * ρ + (ψ - 1) * (r + κ * σB) + μpB - (ψ - 1) * γ / 2 * σB^2 + (2 - ψ - γ) / (2 * (ψ - 1)) * σpB^2 + (1 - γ) * σpB * σB)
  # algebraic constraint
  pt = - p * ((1 - i) / p - x / pA - (1 - x) / pB)
  return (; pAt, pBt, pt)
end

m = DiTellaModel()
xn = 80
νn = 10
xs = range(0.01, 0.99, length = xn)
distribution = Gamma(2 * m.κν * m.νbar / m.σνbar^2, m.σνbar^2 / (2 * m.κν))
νs = range(quantile(distribution, 0.001), quantile(distribution, 0.999), length = νn)
stategrid = OrderedDict(:x => xs, :ν => νs)
yend = OrderedDict(:pA => ones(xn, νn), :pB => ones(xn, νn), :p => ones(xn, νn))
result = pdesolve(m, stategrid, yend; is_algebraic = OrderedDict(:pA => false, :pB => false, :p => true))
@assert result.residual_norm <= 1e-5


using SBP_GPU
using CUDA
using Test
using SafeTestsets

@testset "SBP_GPU.jl" begin
    # Write your tests here.
    @test π ≈ 3.14 atol=0.01
    Nx = Ny = 10
    y_in = randn(Nx*Ny)
    y_in_GPU = CuArray(y_in)
    y_out_GPU = similar(y_in_GPU)
    y_out_GPU = D2x_GPU(y_in_GPU,y_out_GPU,Nx,Ny)
    @test Array(y_out_GPU) ≈ D2x(y_in,Nx,Ny,1/Nx)
    @test Array(Dx_GPU(y_in_GPU,y_out_GPU,Nx,Ny)) ≈ Dx(y_in,Nx,Ny,1/Nx)

    Nx = Ny = 100
    y_in = randn(Nx*Ny)
    y_in_GPU = CuArray(y_in)
    y_out_GPU = similar(y_in_GPU)
    y_out_GPU = D2x_GPU(y_in_GPU,y_out_GPU,Nx,Ny)
    @test Array(y_out_GPU) ≈ D2x(y_in,Nx,Ny,1/Nx)
    @test Array(Dx_GPU(y_in_GPU,y_out_GPU,Nx,Ny)) ≈ Dx(y_in,Nx,Ny,1/Nx)

    Nx = Ny = 1000
    y_in = randn(Nx*Ny)
    y_in_GPU = CuArray(y_in)
    y_out_GPU = similar(y_in_GPU)
    y_out_GPU = D2x_GPU(y_in_GPU,y_out_GPU,Nx,Ny)
    @test Array(y_out_GPU) ≈ D2x(y_in,Nx,Ny,1/Nx)
    @test Array(Dx_GPU(y_in_GPU,y_out_GPU,Nx,Ny)) ≈ Dx(y_in,Nx,Ny,1/Nx)
end


module TestRadialODE

export compare_to_qnm

const tolerance = 1e-3 ## tolerance we compare to

include("../src/CustomTypes.jl")
include("../src/RadialODE.jl")
include("ReadQNM.jl")
using .CustomTypes
using .ReadQNM
import .RadialODE as RO
import Test: @test

##============================================================

"""
Compare against values computed by Leo Stein's qnm code,

J.Open Source Softw. 4 (2019) 42, 1683

This is written in python so we need to bind to the library
"""
function compare_to_qnm(
      nr::myI,
      s::myI,
      n::myI,
      l::myI,
      m::myI,
      avals::Vector{myF}
   )

   println("Comparison to qnm library: s=$s, n=$n, l=$l, m=$m")

   for a in avals
      om, la = qnm(n,s,m,l,a)

      ls_c, vs_c = RO.eig_vals_vecs_c(nr, s, m, a, tomyC(om))

      #println("testing: a=$a, ω=$om\nΛ=$la, Cheb Λ=$(ls_c[1])")

      @test abs(ls_c - la)/max(1,abs(la))<tolerance
   end
end

end


struct convGauss{T,S} <: AbstractPDF{1}
    source::T
    σ::S
end

pars(d::convGauss, isfree::Bool) = pars(d.source, isfree) + pars(d.σ, isfree)
lims(d::convGauss) = lims(d.source)
updatevalueorflag(d::convGauss, s::Symbol, isfree::Bool, v=getproperty(pars(d),s)) =
    convGauss(
        ispar(d.source,s) ? updatevalueorflag(d.source,s,isfree,v) : d.source,
        ispar(d.σ,     s) ? updatevalueorflag(d.σ,     s,isfree,v) : d.σ)

function func(d::convGauss, x::NumberOrTuple; p=pars(d))
    σ = func(d.σ, x; p)
    g(z) = AlgebraPDF.standardgauss(z, σ)
    f(z) = func(d.source, z; p)
    return quadgk(y->f(x-y) * g(y), -5*σ, +5*σ)[1]
end


@muladd function perform_step!(integrator,cache::SplitEMConstantCache,f=integrator.f)
  @unpack t,dt,uprev,u,W,p = integrator
  u = dt*(integrator.f.f1(uprev,p,t) +
           integrator.f.f2(uprev,p,t)) +
           integrator.g(uprev,p,t).*W.dW + uprev
  integrator.u = u
end

@muladd function perform_step!(integrator,cache::SplitEMCache,f=integrator.f)
  @unpack rtmp1,rtmp2 = cache
  @unpack t,dt,uprev,u,W,p = integrator

  integrator.g(rtmp2,uprev,p,t)
  if is_diagonal_noise(integrator.sol.prob)
    rmul!(rtmp2,W.dW) # rtmp2 === rtmp3
    @. u = uprev + rtmp2
  else
    mul!(rtmp1,rtmp2,W.dW)
    @. u = uprev + rtmp1
  end

  integrator.f.f1(t,uprev,rtmp1)
  @. u += dt*rtmp1
  integrator.f.f2(t,uprev,rtmp1)
  @. u += dt*rtmp1
end


include("Density.jl");
include("BiasPotential.jl");

function calc_delta_potential_of_mean_force(umbrella_center, sim_umbrella_datas, sigma_matrix, kbt, input_point)
    delta_pmf = -kbt * log(calc_density(input_point,     sim_umbrella_datas, sigma_matrix))
    delta_pmf += kbt * log(calc_density(umbrella_center, sim_umbrella_datas, sigma_matrix))
    delta_pmf -= calc_bias_potential(umbrella_center, input_point, kbt)
    return delta_pmf
end

function calc_potential_of_mean_force(umbrella_centers, sample_datas, weights, sigma_matrix)
    pmf = []
    for i = 1:size(sample_datas, 1)
        sum_gaussian = 0.0
        for j = 1:size(umbrella_centers, 1)
            sum_gaussian += weights[j] * calc_gaussian(sample_datas[i], umbrella_centers[j, :], sigma_matrix)
        end
        push!(pmf, sum_gaussian)
    end

    return pmf
end


module multimaterial_beam_examples
using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule
using FinEtools.MeshExportModule

function multimaterial_beam_algo()
    println("""
    Multi-material beam. Rubber-like and metal-like halves,
    clamped, with shear traction at free end.
    """)
    E1 = 0.29e3;
    nu1 = 0.49;
    E2 = 0.4e4;
    nu2 = 0.3;
    W = 4.1;
    L = 12.;
    t = 6.5;
    nl = 2; nt = 1; nw = 1; ref = 9;
    p =   200.0/W/t;
    #  Loading in the Z direction
    loadv = [0;0;p]; dir = 3;
    tolerance  = t/1000;

    fens,fes  = H20block(L,W,t, nl*ref,nw*ref,nt*ref)

    # Clamped end of the beam
    l1  = selectnode(fens; box = [0 0 -Inf Inf -Inf Inf], inflate  =  tolerance)
    e1 = FDataDict("node_list"=>l1, "component"=>1, "displacement"=>0.0)
    e2 = FDataDict("node_list"=>l1, "component"=>2, "displacement"=>0.0)
    e3 = FDataDict("node_list"=>l1, "component"=>3, "displacement"=>0.0)

    # Traction on the opposite edge
    boundaryfes  =   meshboundary(fes);
    Toplist   = selectelem(fens,boundaryfes, box =  [L L -Inf Inf -Inf Inf], inflate =   tolerance);
    el1femm  =   FEMMBase(IntegDomain(subset(boundaryfes,Toplist), GaussRule(2, 2)))
    flux1 = FDataDict("femm"=>el1femm, "traction_vector"=>loadv)

    r1list   = selectelem(fens,fes, box =  [0 L/2. -Inf Inf -Inf Inf], inflate =   tolerance);
    r2list   = selectelem(fens,fes, box =  [L/2. L -Inf Inf -Inf Inf], inflate =   tolerance);

    # Model reduction type
    MR = DeforModelRed3D

    # Make region 1
    region1 = FDataDict("femm"=>FEMMDeforLinear(MR,
    IntegDomain(subset(fes,r1list), GaussRule(3,2)),
    MatDeforElastIso(MR, 0.0, E1, nu1, 0.0)))

    # Make region 2
    region2 = FDataDict("femm"=>FEMMDeforLinear(MR,
    IntegDomain(subset(fes,r2list), GaussRule(3,2)),
    MatDeforElastIso(MR, 0.0, E2, nu2, 0.0)))

    # Make model data
    modeldata =  FDataDict(
    "fens"=> fens, "regions"=>  [region1, region2],
    "essential_bcs"=>[e1, e2, e3], "traction_bcs"=>  [flux1])

    # Call the solver
    modeldata = AlgoDeforLinearModule.linearstatics(modeldata)
    geom = modeldata["geom"]
    u = modeldata["u"]

    # Write out mesh with displacements
    modeldata["postprocessing"] = FDataDict("file"=>"multimaterial_beam")
    modeldata = AlgoDeforLinearModule.exportdeformation(modeldata)

    # Write out mesh with stresses
    modeldata["postprocessing"] = FDataDict("file"=>"multimaterial_beam_xy",
    "quantity"=> :Cauchy, "component"=> :xy)
    modeldata = AlgoDeforLinearModule.exportstress(modeldata)

    # Write out mesh with stresses
    modeldata["postprocessing"] = FDataDict("file"=>"multimaterial_beam_xz",
    "quantity"=> :Cauchy, "component"=> :xz)
    modeldata = AlgoDeforLinearModule.exportstress(modeldata)

    # Write out mesh with von Mises stresses
    modeldata["postprocessing"] = FDataDict("file"=>"multimaterial_beam_vm",
    "quantity"=> :vm)
    modeldata = AlgoDeforLinearModule.exportstress(modeldata)

    # Write out mesh with von Mises stresses, elementwise
    modeldata["postprocessing"] = FDataDict("file"=>"multimaterial_beam_vm-ew",
    "quantity"=> :vm)
    modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)

end # multimaterial_beam_algo

function allrun()
    println("#####################################################")
    println("# multimaterial_beam_algo ")
    multimaterial_beam_algo()
    return true
end # function allrun

end # module multimaterial_beam_examples


# # Подготовка данных для тестирования алгоритма выбора начального приближения

# Здесь мы готовим данные, для тестирования алгоритмов.
# Подразумевается, что скрипт будет исполнятся очень редко. Все свои результаты
# он запишет в соответствующие файлы в текущей папке.

using NonLinearReactionAdvectionDiffusionWithFrontData;
using DelimitedFiles;
using Plots; pyplot();

# Выберем параметры таким образом, чтобы решение точно было хорошим.
# Здесь мы генерируем экспериментальные данные, поэтому решим прямую задачу
# всего раз для каждого случая.

x_tp    = 0.05;         # Стартовое местоположение фронта
T_end   = 0.58         # Регулируем конечное местоположение фронта
Nx      = 15000;         # Число интервалов по ``X``
Mt      = 30000;        # Число интервалов по ``T``
α       = 0.005;        # Параметр регуляризации в нахождении начального приближения

# Генерируем данные для набора малых параметров

plts = Any[];           # Вектор для хранения графиков Plots

for (eps, x_tp) in [(0.03, 0.04), (0.01, 0.03), (0.001,0.02)]
    @info "Генерируем данные для ε = $(eps), x_tp = $(x_tp)"
    ε       = eps;          # Крутизна фронта
#+
    a, b, t₀, T, N, M, ε, Xₙ, Tₘ, qₙ, ulₘ, urₘ, u₀ = dparams(x_tp = x_tp,
                                                            ε = ε,
                                                            Nx = Nx,
                                                            Mt = Mt,
                                                            T_end = T_end,
                                                            qfunc = x -> sin(3pi*x)
                                                            );
    u, XX, TP = solve(u₀, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ);
    ϕl, ϕr, ϕ, f1_data, f2_data = generate_obs_data(u, Xₙ, N, Tₘ, M, qₙ, ulₘ, urₘ);
#+
    # Убеждаемся в корректности прямого решения на глаз и пишем данные в txt
    directP = draft(u, Xₙ, N, Tₘ, M, title = "Эскиз прямого решения, \\varepsilon = $(ε)")
    savefig("direct_eps_$(ε)_sin.png");
    writedlm("tm_f1_m_eps_$(ε)_sin.txt",[Tₘ f1_data])
    writedlm("xn_qn_eps_$(ε).txt",[Xₙ qₙ])
#+
    # Находим начальные приближения
    #  q_guess = initial_guess(f1_data, Xₙ, N, Tₘ, M, ulₘ, urₘ, α);
    #  p = plot(Xₙ, qₙ, label="Истинное",  title = "Начальное приближение, \\varepsilon = $(ε)");
    #  p = plot!(Xₙ, q_guess, label="Найденное")
    #  savefig("guess_eps_$(ε)_sin.png");
    #  push!(plts, p);
#+
    # Два гаусиана разной высоты, которые частично перекрываются
    gauss_init(x) = 2 * exp(-((x-5.0/13.0)^2)/(1.5/13.0^2)) +
        1.2 * exp((-(x-8.0/13.0)^2)/(3.0/13.0^2));
    qₙ = [gauss_init(x) for x in Xₙ]
    plot(Xₙ, qₙ)
#+
    u, XX, TP = solve(u₀, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ);
    ϕl, ϕr, ϕ, f1_data, f2_data = generate_obs_data(u, Xₙ, N, Tₘ, M, qₙ, ulₘ, urₘ);
#+
    # Убеждаемся в корректности прямого решения на глаз и пишем данные в text
    directP = draft(u, Xₙ, N, Tₘ, M, title = "Эскиз прямого решения, \\varepsilon = $(ε)")
    savefig("direct_eps_$(ε)_gauss.png");
    writedlm("tm_f1_m_eps_$(ε)_gauss.txt",[Tₘ f1_data])
    writedlm("xn_qn_eps_$(ε)_gauss.txt",[Xₙ qₙ])
#+
    # Находим начальные приближения
    #  q_guess = initial_guess(f1_data, Xₙ, N, Tₘ, M, ulₘ, urₘ, α);
    #  p = plot(Xₙ, qₙ, label="Истинное",  title = "Начальное приближение, \\varepsilon = $(ε)");
    #  p = plot!(Xₙ, q_guess, label="Найденное")
    #  savefig("guess_eps_$(ε)_gauss.png");
    #  push!(plts, p);

end

# График
plts[1]

# График
plts[2]

# График
plts[3]

# График
plts[4]

# График
plts[5]

# График
plts[6]


# ---
# title: 654. Maximum Binary Tree
# id: problem654
# author: Indigo
# date: 2021-06-28
# difficulty: Medium
# categories: Tree
# link: <https://leetcode.com/problems/maximum-binary-tree/description/>
# hidden: true
# ---
#
# Given an integer array with no duplicates. A maximum tree building on this
# array is defined as follow:
#
#   1. The root is the maximum number in the array.
#   2. The left subtree is the maximum tree constructed from left part subarray divided by the maximum number.
#   3. The right subtree is the maximum tree constructed from right part subarray divided by the maximum number.
#
# Construct the maximum tree by the given array and output the root node of this
# tree.
#
# **Example 1:**
#
#
#
#     Input: [3,2,1,6,0,5]
#     Output: return the tree root node representing the following tree:
#
#           6
#         /   \
#        3     5
#         \    /
#          2  0
#            \
#             1
#
#
# **Note:**
#
#   1. The size of the given array will be in the range [1,1000].
#
#
## @lc code=start
using LeetCode

function construct_maximum_binary_tree(nums::AbstractVector{Int})
    isempty(nums) && return nothing
    max_idx = argmax(nums)
    return TreeNode(nums[max_idx],
                    construct_maximum_binary_tree(@view(nums[1:(max_idx - 1)])),
                    construct_maximum_binary_tree(@view(nums[(max_idx + 1):end])))
end
## @lc code=end


function SphericalVertexEigenfunction(
    domain::SphericalTriangle,
    vertex::Int;
    stride::Int = 1,
)
    vertex >= 1 && vertex <= 3 ||
        throw(ErrorException("vertex must be between 1 and 3 not $vertex"))
    SphericalVertexEigenfunction(domain, vertex, stride, arb[])
end

function Base.show(io::IO, u::SphericalVertexEigenfunction)
    println(io, "Spherical Vertex eigenfunction from vertex $(u.vertex)")
    if !haskey(io, :compact) || !io[:compact]
        println(io, "domain: $(u.domain)")
        print(io, "number of set coefficients: $(length(u.coefficients))")
    end
end

function active_boundaries(domain::SphericalTriangle, u::SphericalVertexEigenfunction)
    if domain === u.domain
        return u.vertex:u.vertex
    else
        return 1:3
    end
end

"""
    mu(eigenfunction::SphericalVertexEigenfunction,
       k::Integer = 1)
> Return k*μ0 as an arb ball, the parameter used for the Legendre
  function.
"""
function mu(u::SphericalVertexEigenfunction{fmpq}, k::Integer = 1)
    u.domain.parent(-k * inv(u.domain.angles[u.vertex]))
end

function mu(u::SphericalVertexEigenfunction{arb}, k::Integer = 1)
    u.domain.parent(-k * inv(u.domain.angles[u.vertex])) * u.domain.parent(π)
end

function coordinate_transformation(
    u::SphericalVertexEigenfunction,
    xyz::AbstractVector{T},
) where {T<:Union{arb,arb_series}}
    # TODO: The performance could most likely be improved
    if u.vertex == 1
        return xyz
    elseif u.vertex == 2
        # Rotate by α along the y-axis so that the second vertex ends
        # up at the north pole, then rotate by β around the z-axis so
        # that the boundary ends up parallel to the y-axis.
        α = -acos(vertex(u.domain, 2)[3])
        β = angle(u.domain, 2) - u.domain.parent(π)
        L = LinearMap(RotZY(β, α))
        return L(xyz)
    elseif u.vertex == 3
        # Rotate by α along the z-axis so that the third vertex
        # ends up parallel to the y-axis, then rotate by β around
        # the y-axis so that it ends up at the north pole. Finally
        # rotate by γ so along the z-axis so that the boundary
        # ends up parallel to the y-axis.
        α = -angle(u.domain, 1)
        β = -acos(vertex(u.domain, 3)[3])
        γ = u.domain.parent(π)
        L = LinearMap(RotZYZ(γ, β, α))
        return L(xyz)
    else
        throw(ErrorException("vertex must be between 1 and 3 not $vertex"))
    end
end

function coordinate_transformation(
    u::SphericalVertexEigenfunction,
    θ::T,
    ϕ::T,
) where {T<:Union{arb,arb_series}}
    if u.vertex == 1
        return θ, ϕ
    else
        return spherical(coordinate_transformation(u, cartesian(θ, ϕ)))
    end
end

function (u::SphericalVertexEigenfunction)(
    xyz::AbstractVector{T},
    λ::arb,
    k::Integer;
    boundary = nothing,
    notransform::Bool = false,
) where {T<:Union{arb,arb_series}}
    if !isnothing(boundary) && boundary != u.vertex
        if T == arb
            return u.domain.parent(0)
        else
            return 0 * xyz[1]
        end
    end
    k = 1 + (k - 1) * u.stride

    ν::arb = -0.5 + sqrt(0.25 + λ)
    μ::arb = mu(u, k)

    if !notransform
        xyz = coordinate_transformation(u, xyz)
    end

    ϕ = atan(xyz[2], xyz[1])
    legendre_p_safe(ν, μ, xyz[3]) * sin(μ * ϕ)
end

function (u::SphericalVertexEigenfunction)(
    θ::T,
    ϕ::T,
    λ::arb,
    k::Integer;
    boundary = nothing,
    notransform::Bool = false,
) where {T<:Union{arb,arb_series}}
    if !isnothing(boundary) && boundary != u.vertex
        if T == arb
            return u.domain.parent(0)
        else
            return 0 * θ
        end
    end
    k = 1 + (k - 1) * u.stride

    ν::arb = -0.5 + sqrt(0.25 + λ)
    μ::arb = mu(u, k)

    if !notransform
        θ, ϕ = coordinate_transformation(u, θ, ϕ)
    end

    legendre_p_safe(ν, μ, cos(θ)) * sin(μ * ϕ)
end

function norm(domain::SphericalTriangle, u::SphericalVertexEigenfunction, λ::arb)
    ϕ_integral = -domain.parent(π) / (2 * mu(u, 1))

    CC = ComplexField(domain.parent.prec)
    # The integrals goes from zero to the lower bound for θ. However
    # the function has a branch cut at zero and Arb has problem
    # handling this. We therefore integrate a small distance away from
    # zero.
    a = CC(1e-1)

    # Create a new domain which has the vertex we are expanding from
    # on the north pole.
    domain = SphericalTriangle(
        domain.angles[[u.vertex, mod1(u.vertex + 1, 3), mod1(u.vertex + 2, 3)]],
        domain.parent,
    )
    b = CC(theta_bound(domain))

    θ_integral = domain.parent(0)
    for k = 1:min(4, length(u.coefficients))
        ν::arb = -0.5 + sqrt(0.25 + λ)
        μ::arb = mu(u, 1 + (k - 1) * u.stride)
        c2 = u.coefficients[k]^2

        f = θ -> begin
            θ = real(θ)
            if !isnonpositive(cos(θ) - 1)
                return CC(NaN)
            end

            CC(c2 * sin(θ)legendre_p_safe(ν, μ, cos(θ))^2)
        end
        y = Nemo.integrate(CC, f, a, b)
        θ_integral += real(y)
    end

    sqrt(ϕ_integral * θ_integral)
end


export newton, is_success, NewtonResult, NewtonCache

"""
    NewtonResult

Result returned by [`newton`](@ref).

## Fields

* `return_code::Symbol`: Can be `:success`, `:rejected` or `:max_iters`.
* `x::Vector{ComplexF64}`: The last value obtained.
* `accuracy::Float64`: Estimate of the distance of `x` to a true zero.
* `iters::Int` Number of iterations performed.
* `contraction_ratio::Float64`: The value `|xᵢ - xᵢ₋₁| / |xᵢ₋₁ - xᵢ₋₂| `.
"""
struct NewtonResult
    return_code::Symbol
    x::Vector{ComplexF64}
    accuracy::Float64
    residual::Float64
    iters::Int
    contraction_ratio::Float64
end

Base.show(io::IO, ::MIME"application/prs.juno.inline", r::NewtonResult) = r
Base.show(io::IO, result::NewtonResult) = print_fieldnames(io, result)

"""
    is_success(R::NewtonResult)

Returns `true` if the [`newton`](@ref) was successfull.
"""
is_success(R::NewtonResult) = R.return_code == :success

"""
    NewtonCache(F::AbstractSystem)

Pre-allocates the necessary memory for [`newton`](@ref).
"""
struct NewtonCache{M}
    x::Vector{ComplexF64}
    Δx::Vector{ComplexF64}
    x_ext::Vector{ComplexDF64}
    J::MatrixWorkspace{M}
    r::Vector{ComplexF64}
end
function NewtonCache(F::AbstractSystem)
    m, n = size(F)
    x = zeros(ComplexF64, n)
    Δx = copy(x)
    x_ext = zeros(ComplexDF64, n)
    J = MatrixWorkspace(m, n)
    r = zeros(ComplexF64, m)
    NewtonCache(x, Δx, x_ext, J, r)
end

"""
    newton(
        F::AbstractSystem,
        x₀::AbstractVector,
        norm::AbstractNorm = InfNorm(),
        cache::NewtonCache = NewtonCache(F, x₀);
        options...
    )

An implemenetation of a local Newton's method with various options to specify
convergence criteria.
Returns a [`NewtonResult`](@ref).
The computations are always performed in complex arithmetic with
double precision, i.e., using `Complex{Float64}`.
The optional `cache` argument pre-allocates the necessary memory.
This is useful if the method is called repeatedly.

### Options

* `abs_tol::Float64 = 1e-8`: The method is declared converged if `norm(xᵢ₊₁ - xᵢ) < abs_tol`.
* `rel_tol::Float64 = abs_tol`: The method is declared converged if
  `norm(xᵢ₊₁ - xᵢ) < abs_tol * norm(x₀)`.
* `max_iters::Int = 20`: The maximal number of iterations.
* `extended_precision::Bool = false`: An optional use of extended precision for the evaluation
  of `F(x)`. This can increase the achievable accuracy.
* `contraction_factor::Float64 = 1.0`: The Newton updates have to satisfy
  ``|xᵢ₊₁ - xᵢ| < a^{2^(i-1)}|x₁ - x₀|`` for ``i ≥ 1`` where ``a`` is `contraction_factor`.
* `min_contraction_iters::Int = typemax(Int)`:  The minimal number of iterations the
  `contraction_factor` has to be satisfied. If after `min_contraction_iters` many
  iterations the contraction factor is not satisfied the step is accepted anyway.
* `max_abs_norm_first_update::Float64 = Inf`: The initial guess `x₀` is rejected if
  `norm(x₁ - x₀) >  max_abs_norm_first_update`
* `max_rel_norm_first_update::Float64 = max_abs_norm_first_update`: The initial guess `x₀`
  is rejected if `norm(x₁ - x₀) >  max_rel_norm_first_update * norm(x₀)`
"""
newton(f::System, args...; kwargs...) = newton(ModelKitSystem(f), args...; kwargs...)
function newton(
    F::AbstractSystem,
    x₀::AbstractVector,
    norm::AbstractNorm = InfNorm(),
    cache::NewtonCache = NewtonCache(F);
    extended_precision::Bool = false,
    abs_tol::Float64 = 1e-8,
    rel_tol::Float64 = abs_tol,
    max_iters::Int = 20,
    contraction_factor::Float64 = 1.0,
    min_contraction_iters::Int = typemax(Int),
    max_abs_norm_first_update::Float64 = Inf,
    max_rel_norm_first_update::Float64 = max_abs_norm_first_update,
)
    @unpack x, Δx, x_ext, J, r = cache
    x .= x₀
    a = contraction_factor
    norm_x = norm(x)
    norm_Δxᵢ = norm_Δxᵢ₋₁ = NaN
    res = NaN
    for i = 1:max_iters
        evaluate_and_jacobian!(r, matrix(J), F, x)
        if extended_precision
            x_ext .= x
            evaluate!(r, F, x_ext)
        end

        updated!(J)
        LA.ldiv!(Δx, J, r)
        norm_Δxᵢ = norm(Δx)

        x .= x .- Δx

        if i > 1
            if norm_Δxᵢ < a * norm_Δxᵢ₋₁
                norm_Δxᵢ₋₁ = norm_Δxᵢ
                a *= a
            elseif i > min_contraction_iters ||
                   norm_Δxᵢ < rel_tol * norm_x ||
                   norm_Δxᵢ < abs_tol
                return NewtonResult(
                    :success,
                    x,
                    norm_Δxᵢ,
                    norm(r),
                    i,
                    norm_Δxᵢ / norm_Δxᵢ₋₁,
                )
            else
                return NewtonResult(
                    :rejected,
                    x,
                    norm_Δxᵢ,
                    norm(r),
                    i,
                    norm_Δxᵢ / norm_Δxᵢ₋₁,
                )
            end
        elseif i == 1 && (
            norm_Δxᵢ > max_rel_norm_first_update * norm_x ||
            norm_Δxᵢ > max_abs_norm_first_update
        )
            return NewtonResult(:rejected, x, norm(r), norm_Δxᵢ, i, NaN)
        else
            norm_Δxᵢ₋₁ = norm_Δxᵢ
        end
    end
    return NewtonResult(:max_iters, x, norm_Δxᵢ, norm(r), max_iters, norm_Δxᵢ / norm_Δxᵢ₋₁)
end


# Rigid-body motion utilities

export RegularCartesianSpatialSampling
export spatial_sampling, coord, upscale, downscale

## Cartesian domain type

struct RegularCartesianSpatialSampling{T}<:AbstractCartesianSpatialSampling{T}
    o::NTuple{3,T} # integer origin wrt to middle point
    n::NTuple{3,Integer}
    h::NTuple{3,T}
end

spatial_sampling(T::DataType, n::NTuple{3,Integer}; h::NTuple{3,<:Real}=(1,1,1), o::NTuple{3,<:Real}=(0,0,0)) = RegularCartesianSpatialSampling{T}(T.(o), n, T.(h))

## Utils

Base.size(X::RegularCartesianSpatialSampling) = X.n

function coord(X::RegularCartesianSpatialSampling{T}; mesh::Bool=false) where {T<:Real}
    idx_orig = X.o.+idx_orig_default(X.n)
    x = reshape((collect(1:X.n[1]).-idx_orig[1])*X.h[1], :,1,1)
    y = reshape((collect(1:X.n[2]).-idx_orig[2])*X.h[2], 1,:,1)
    z = reshape((collect(1:X.n[3]).-idx_orig[3])*X.h[3], 1,1,:)
    if mesh
        x = repeat(reshape(vec(x),:,1,1); outer=(1,X.n[2],X.n[3]))
        y = repeat(reshape(vec(y),1,:,1); outer=(X.n[1],1,X.n[3]))
        z = repeat(reshape(vec(z),1,1,:); outer=(X.n[1],X.n[2],1))
    end
    return x, y, z
end

idx_orig_default(n::Integer) = div(n,2)+1
idx_orig_default(n::NTuple{3,Integer}) = idx_orig_default.(n)

upscale(X::RegularCartesianSpatialSampling{T}; fact::Integer=1) where {T<:Real} = spatial_sampling(T, Integer.(X.n.*2.0^fact); h=X.h.*T(2)^-fact, o=X.o.*T(2)^fact)

function downscale(X::RegularCartesianSpatialSampling{T}; fact::Integer=1) where {T<:Real}
    (fact == 0) && (return X)
    if mod.(X.n, 2^fact) == (0,0,0)
        return spatial_sampling(T, Integer.(X.n.*2.0^-fact); h=X.h.*T(2)^fact, o=X.o.*T(2)^-fact)
    else
        error("Dimensions not multiple of scaling factors")
    end
end


#=
  From
  Yi Wu, Lei Li, Stuart Russell, Rastislav Bodik
  "Swift: Compiled Inference for Probabilistic Programming Languages"
  Page 3

  Summary Statistics
  parameters      mean       std   naive_se       mcse         ess       rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64?    Float64?   Float64?      Float64?

     cluster    3.5614    1.4721     0.0147     0.0347   1385.1836     1.0003      111.8346
       mu[1]   -0.1752    7.9247     0.0792     0.1902   1447.7438     1.0002      116.8855
       mu[2]    0.1621    8.1310     0.0813     0.2017   1361.4057     1.0007      109.9149
        z[1]    2.2715    1.2967     0.0130     0.0352   1192.9498     1.0004       96.3144
        z[2]    2.2670    1.2947     0.0129     0.0341   1242.9694     1.0001      100.3528
        z[3]    2.2901    1.2907     0.0129     0.0340   1366.0522     1.0006      110.2900
   x_post[1]    0.0905    1.4039     0.0140     0.0223   4274.1990     1.0006      345.0831
   x_post[2]    0.0648    7.9544     0.0795     0.2009   1411.9730     1.0021      113.9975
   x_post[3]   -0.2165    8.3757     0.0838     0.2278   1344.0788     1.0010      108.5160
       mu[3]    0.1385    8.4502     0.0998     0.2856     missing    missing       missing
       mu[4]    0.1468    8.7843     0.1334     0.3392     missing    missing       missing
       mu[5]   -0.9117    8.9978     0.1853     0.5347     missing    missing       missing
       mu[6]   -0.1977    9.2914     0.2807     0.9323     missing    missing       missing
       mu[7]   -0.3155    8.8300     0.4176     0.6646     missing    missing       missing
       mu[8]    4.1820   12.8344     1.1005    missing     missing    missing       missing
       mu[9]    5.5301    9.7624     1.4240    missing     missing    missing       missing
      mu[10]    2.6899    2.7582     0.6501    missing     missing    missing       missing
   ...

  Distributions of variable cluster (num:0)
  3.00000 =>    2838  (0.283800)
  2.00000 =>    2825  (0.282500)
  4.00000 =>    1979  (0.197900)
  5.00000 =>    1262  (0.126200)
  6.00000 =>     649  (0.064900)
  7.00000 =>     311  (0.031100)
  8.00000 =>      89  (0.008900)
  9.00000 =>      29  (0.002900)
  10.00000 =>      18  (0.001800)

   cf ~/webppl/infinite_gaussian_mixture.wppl

=#

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

#=
BLOG model:

distinct Data D[20]; // 20 data points
#Cluster ~ Poisson(4); // Number of clusters

random Real mu(Cluster c) ~ Gaussian(0.0,10.0); // Cluster mean
random Cluster z(Data d) ~ UniformChoice({c for Cluster c});

random Real x(Data d) ~ Gaussian(mu(z(d)), 1.0); // data

=#
@model function infinite_gaussian_mixture(x=[0.1],n=3)
    # n = length(x)

    cluster ~ Truncated(Poisson(3),1,Inf) # Number of clusters
    mu ~ filldist(Normal(0,10),cluster) # Cluster mean

    # Number of clusters per data points
    z = tzeros(n)
    for i in 1:n
        z[i] ~ DiscreteUniform(1,cluster)
    end

    # Observe the data.
    # z[i] returns a float, hence the round
    for i in 1:length(x)
        zi = round(Int,z[i])
        x[i] ~ Normal(mu[zi], 1.0)
    end

    # Here we want n posterior data points
    x_post = tzeros(n)
    for i in 1:n
        zi = round(Int,z[i])
        x_post[i] ~ Normal(mu[zi], 1.0)
    end

    # x1 ~ Dirac(x_post[1])
    # x2 ~ Dirac(x_post[2])
    # x3 ~ Dirac(x_post[3])

end

# x=[0.1,0.1,0.1]
x = [0.1]
n = 3 # we want 3 predictions
model = infinite_gaussian_mixture(x,n)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 100_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,:cluster)


# ---
# title: 477. Total Hamming Distance
# id: problem477
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Bit Manipulation
# link: <https://leetcode.com/problems/total-hamming-distance/description/>
# hidden: true
# ---
#
# The [Hamming distance](https://en.wikipedia.org/wiki/Hamming_distance) between
# two integers is the number of positions at which the corresponding bits are
# different.
#
# Now your job is to find the total Hamming distance between all pairs of the
# given numbers.
#
# **Example:**
#
#
#
#     Input: 4, 14, 2
#
#     Output: 6
#
#     Explanation: In binary representation, the 4 is 0100, 14 is 1110, and 2 is 0010 (just
#     showing the four bits relevant in this case). So the answer will be:
#     HammingDistance(4, 14) + HammingDistance(4, 2) + HammingDistance(14, 2) = 2 + 2 + 2 = 6.
#
#
# **Note:**
#
#   1. Elements of the given array are in the range of `0 ` to `10^9`
#   2. Length of the array will not exceed `10^4`.
#
#
## @lc code=start
using LeetCode

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


using Test

#--- parameters
para = ( ϵ  = 0.1, μ = 0.1)  # physical parameters
paraX= ( N  = 4, L  = π)   # mesh with 6 collocation points: x=[-π, -π/2, 0.0, π/2]
paraT= ( T  = 1e-2, dt = 1e-3) # timegrid with ten instants: t=[0.0:1.0:10.0]/10^3
param = merge(para,paraX)  # used to construct models
parap = merge(paraX,paraT) # used to construct problems

#--- initial data
init     = Init(x->cos.(x),x-> sin.(x) )

#--- model
model=WWn(param)

#--- reference problem
pb0 = Problem( model, init, parap )
solve!(pb0,verbose=false)

#--- tests on Runge-Kutta 4 solvers
@testset "RK4 solvers" begin
    # different ways to build RK4 solver
    solvers = TimeSolver[]
    push!(solvers, RK4(model.mapto(init)) )
    push!(solvers, RK4(paraX) )
    push!(solvers, RK4(paraX,2) )
    push!(solvers, RK4_naive() )

    # check all RK4 solvers generate the same data
    for solver in solvers
        pb = Problem( model, init, parap;
                    solver = solver );
        solve!(pb,verbose=false)
        @test pb.data.U == pb0.data.U
    end
end

#--- tests on explicit Euler solvers
@testset "explicit Euler solver" begin
    # build explicit Euler solver
    pb1 = Problem( model, init, parap;
                solver = Euler(model) )
    solve!(pb1,verbose=false)
    # check explicit Euler is of order ≈ dt*T
    order = paraT.dt*paraT.T
    @test isapprox(pb0.data.U[end] , pb1.data.U[end] , rtol = order)
    @test !isapprox(pb0.data.U[end] , pb1.data.U[end] , rtol = order/4)

    # different ways to build explicit Euler solver
    solvers = TimeSolver[]
    push!(solvers, Euler(model.mapto(init)) )
    push!(solvers, Euler(paraX) )
    push!(solvers, Euler(paraX,2) )
    push!(solvers, Euler_naive() )

    # check all Euler solvers generate the same data
    for solver in solvers
        pb = Problem( model, init, parap;
                    solver = solver );
        solve!(pb,verbose=false)
        @test pb.data.U == pb1.data.U
    end
end

#--- tests on symplectic Euler solvers
@testset "symplectic Euler solver" begin
    # build symplectic Euler solver
    pb1 = Problem( model, init, parap;
                solver = EulerSymp(model) )
    solve!(pb1,verbose=false)
    # check symplectic Euler is of order ≈ dt*T
    order = paraT.dt*paraT.T
    @test isapprox(pb0.data.U[end] , pb1.data.U[end] , rtol = order)
    @test !isapprox(pb0.data.U[end] , pb1.data.U[end] , rtol = order/4)

    # different ways to build symplectic Euler solver
    solvers = TimeSolver[]
    push!(solvers, EulerSymp(model.mapto(init)) )
    push!(solvers, EulerSymp(paraX) )

    # check all symplectic Euler solvers generate the same data
    for solver in solvers
        pb = Problem( model, init, parap;
                    solver = solver );
        solve!(pb,verbose=false)
        @test pb.data.U == pb1.data.U
    end

    # change number of iterations in the implicit step
    N=5
    pb2 = Problem( model, init, parap;
                solver = EulerSymp(model, Niter=N) )
    solve!(pb2,verbose=false)

    # check the difference is of order ≈ (dt)^(N+1)*T
    order = (paraT.dt)^(N+1)*paraT.T
    @test isapprox(pb2.data.U[end] , pb1.data.U[end], rtol = order )

    # change the equations solved by implicit method
    N=5
    pb3 = Problem( model, init, parap;
                solver = EulerSymp(model, implicit=2) )
    solve!(pb3,verbose=false)

    # check the difference is of order ≈ dt*T
    order = paraT.dt*paraT.T
    @test isapprox(pb3.data.U[end] , pb1.data.U[end], rtol = order )
    @test !isapprox(pb0.data.U[end] , pb1.data.U[end] , rtol = order/4)
end


#=

  From BLOG example/poisson-ball.blog
  """
  Model file for balls in an urn, allowing observation errors.
  This version uses a Poisson prior for the number of balls.
  """

  Summary Statistics
     parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
         Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

   ...
       numBalls    5.8985    1.8454     0.0185    0.0695    677.0798    1.0002      125.6178
   ...
        numTrue    5.0069    1.5379     0.0154    0.0547    816.5438    1.0006      151.4924
       obsPost1    1.3249    0.4684     0.0047    0.0088   2952.0720    1.0000      547.6942
       obsPost2    1.6475    0.4778     0.0048    0.0096   2963.4386    1.0006      549.8031

  Distributions of variable numBalls (num:0)
  5.00000 =>    1996  (0.199600)
  8.00000 =>    1678  (0.167800)
  7.00000 =>    1588  (0.158800)
  6.00000 =>    1548  (0.154800)
  4.00000 =>    1371  (0.137100)
  3.00000 =>     827  (0.082700)
  9.00000 =>     752  (0.075200)
  2.00000 =>     204  (0.020400)
  1.00000 =>      36  (0.003600)

  Distributions of variable numTrue (num:0)
  5.00000 =>    2486  (0.248600)
  4.00000 =>    2180  (0.218000)
  6.00000 =>    2040  (0.204000)
  7.00000 =>    1216  (0.121600)
  3.00000 =>    1205  (0.120500)
  8.00000 =>     368  (0.036800)
  2.00000 =>     261  (0.026100)
  1.00000 =>     130  (0.013000)
  9.00000 =>      94  (0.009400)
  10.00000 =>      11  (0.001100)
  0.00000 =>       9  (0.000900)

  Distributions of variable obsPost1
  blue       =>    6751  (0.675100)
  green      =>    3249  (0.324900)

  Distributions of variable obsPost2
  green      =>    6475  (0.647500)
  blue       =>    3525  (0.352500)

  Cf ~/webppl/poisson_Ball.wppl

=#
using Turing, StatsPlots
include("jl_utils.jl")

@model function poisson_ball(obs=[1,2,1,2,1,2,1,2,1,2])
    n = length(obs)
    Blue = 1
    Green = 2;
    colors = [Blue,Green]

    # True color for each ball
    TrueColor ~ filldist(Categorical([0.5,0.5]),n) # colors

    # Number of balls
    numBalls ~ truncated(Poisson(6),0,n-1)

    # Which balls is drawn
    BallDrawn ~ filldist(DiscreteUniform(1,numBalls),n)

    # We observe Blue,Green,Blue,Green,Blue,Green,Blue,Green,Blue,Green
    for d in 1:n
        trueColor = TrueColor[BallDrawn[d]]
        if trueColor == Blue
            obs[d] ~ Categorical([0.8,0.2])
        else
            obs[d] ~ Categorical([0.2,0.8])
        end
    end

    # Posterior, first 2 balls
    obs_post = tzeros(2)
    for d in 1:2
        trueColor = TrueColor[BallDrawn[d]]
        if trueColor == Blue
            obs_post[d] ~ Categorical([0.8,0.2])
        else
            obs_post[d] ~ Categorical([0.2,0.8])
        end
    end

    numTrue ~ Dirac(sum([obs[d] == TrueColor[d] ? 1 : 0 for d in 1:n]))

    obsPost1 ~ Dirac(obs_post[1])
    obsPost2 ~ Dirac(obs_post[2])

end

model = poisson_ball()
# 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,:numBalls)
show_var_dist_pct(chns,:numTrue)
show_var_dist_pct(chns,:obsPost1,["blue","green"])
show_var_dist_pct(chns,:obsPost2,["blue","green"])


using MultivariateFunctions
using DataStructures: OrderedDict
tol = 1e-15
fun1 = PE_Unit(1.0,1.0,1)
fun2 = PE_Unit(1.0,1.0,1)
fun3 = PE_Unit(1.0,2.0,3)
fun4 = PE_Unit(1.0,3.0,2)
fun6 = PE_Unit(0.0,1.0,1)
fun7 = PE_Unit(0.0,2.0,4)
fun8 = PE_Unit(0.0,2.0,2)

mfun1 = PE_Function(1.8, Dict(Symbol.(["x", "y", "z"]) .=> [fun1, fun2, fun3]))
abs(mfun1.multiplier_ - 1.2*1.5) < 1e-10
mfun2 = PE_Function(4.5, Dict(Symbol.(["x", "y"]) .=> [fun3, fun4]))
mfun3 = PE_Function(3.0, Dict(Symbol.(["z", "y"]) .=> [fun4, fun1]))
mfun4 = PE_Function(6.0, Dict(Symbol.(["x", "y"]) .=> [fun6, fun7]))
mfun5 = PE_Function(6.0, Dict(Symbol.(["x", "z"]) .=> [fun6, fun8]))

msum = mfun5 + mfun4

func_array = Array{Union{Sum_Of_Functions,Missing},3}(undef,3,2,2)
func_array[1,1,1] = Sum_Of_Functions(mfun1)
func_array[1,1,2] = Sum_Of_Functions(mfun2)
func_array[1,2,1] = mfun2 + mfun1
func_array[1,2,2] = mfun2 + mfun4
func_array[2,1,1] = Sum_Of_Functions(2*mfun1)
func_array[2,1,2] = mfun1 + mfun5
func_array[2,2,1] = mfun1 + mfun5
func_array[2,2,2] = mfun2 + mfun1 + mfun5
func_array[3,:,:] .= Missing()

thresholds_ = OrderedDict{Symbol,Array{Float64,1}}([:x, :y, :z] .=> [[0.0, 0.5, 1.0], [-5.0, -1.0], [0.0, 5.0]])

pw_func = Piecewise_Function(func_array, thresholds_)

# Testing zeroing.
coordinates = Dict{Symbol,Float64}([:w, :x, :y, :z] .=> [1.4, 0.4, 2.0, 3.0])
hypercube =  Dict{Symbol,Tuple{Float64,Float64}}([:x, :z] .=> [(0.5,Inf), (1.0,6.0)])
zero_outside_hypercube = Piecewise_Function(pw_func, hypercube, true)
abs(evaluate(pw_func, coordinates) - evaluate(zero_outside_hypercube, coordinates)) > 5.0
coordinates_in_cube = Dict{Symbol,Float64}([:w, :x, :y, :z] .=> [1.4, 0.7, 2.0, 3.0])
abs(evaluate(pw_func, coordinates_in_cube) - evaluate(zero_outside_hypercube, coordinates_in_cube)) < tol

# creating a pw with a pw
functions_ = Array{Any,2}(undef,3,2)
functions_[1,1] = Sum_Of_Functions(mfun1)
functions_[1,2] = Sum_Of_Functions(3.0*mfun2)
functions_[2,1] = pw_func
functions_[2,2] = mfun2 + mfun4
functions_[3,1] = pw_func + 1.0
functions_[3,2] = mfun2 + mfun5
thresholds_ = OrderedDict{Symbol,Array{Float64,1}}([:x, :y] .=> [[0.0, 0.6, 1.0], [-5.0, -1.0]])
pw_func2 = Piecewise_Function(functions_, thresholds_)
coords_21 = Dict{Symbol,Float64}([:w, :x, :y, :z] .=> [1.4, 0.7, -3.0, 3.0])
coords_22 = Dict{Symbol,Float64}([:w, :x, :y, :z] .=> [1.4, 0.7, 0.0, 3.0])
abs(evaluate(pw_func, coords_21) - evaluate(pw_func2, coords_21)) < tol
abs(evaluate(pw_func, coords_22) - evaluate(pw_func2, coords_22)) > 1.0

###############################################################################
coordinates = Dict(Symbol.(["w","x", "y", "z"]) .=> [1.4, 0.5, 2.0, 3.0])
function test_result(func, eval_to, len = 1)
    val_test = abs(evaluate(func, coordinates) - eval_to) < 1e-05
    if (!val_test)
        println("Failed Val Test")
    end
    len_test = length(func.functions_) == len
    if (!len_test)
        println(string("Failed length Test, Length is ", length(func.functions_)))
    end
    return all([val_test,len_test])
end
ismissing(evaluate(pw_func, Dict(Symbol.(["w","x", "y", "z"]) .=> [-Inf, -Inf, -Inf, -Inf])))
test_result(pw_func, evaluate( mfun1 + mfun5, coordinates), 12)
deriv = Dict{Symbol,Int}(Symbol.(["x", "z"]) .=> [2,1])
abs(evaluate(derivative(pw_func, deriv), coordinates) - evaluate(derivative(mfun1 + mfun5, deriv), coordinates)) < tol

# Integration of piecewise function
integration_limits = Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.75),(-4.0,-1.0),(0.0,3.0)])
integral_of_each_piece = integral(Sum_Of_Functions(mfun1) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.5),(-4.0,-1.0),(0.0,3.0)]))
integral_of_each_piece = integral_of_each_piece + integral(Sum_Of_Functions(2*mfun1) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.5,0.75),(-4.0,-1.0),(0.0,3.0)]))
integral_of_piecewise = integral(pw_func, integration_limits)
abs(integral_of_each_piece - integral_of_piecewise)  < tol
# Integration of piecewise function
integration_limits = Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.75),(-4.0,-1.0),(0.0,8.0)])
integral_of_each_piece = integral(Sum_Of_Functions(mfun1) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.5),(-4.0,-1.0),(0.0,5.0)]))
integral_of_each_piece = integral_of_each_piece + integral(Sum_Of_Functions(2*mfun1) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.5,0.75),(-4.0,-1.0),(0.0,5.0)]))
integral_of_each_piece = integral_of_each_piece + integral(Sum_Of_Functions(mfun2) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.5),(-4.0,-1.0),(5.0,8.0)]))
integral_of_each_piece = integral_of_each_piece + integral(mfun1 + mfun5 ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.5,0.75),(-4.0,-1.0),(5.0,8.0)]))
integral_of_piecewise = integral(pw_func, integration_limits)
abs(integral_of_each_piece - integral_of_piecewise)  < tol

# test algebra between Piecewise_Functions and PE_Functions.
test_result(pw_func + mfun4, evaluate( mfun1 + mfun5, coordinates) + evaluate(mfun4, coordinates), 12)
test_result(pw_func - mfun4, evaluate( mfun1 + mfun5, coordinates) - evaluate(mfun4, coordinates), 12)
test_result(pw_func * mfun4, evaluate( mfun1 + mfun5, coordinates) * evaluate(mfun4, coordinates), 12)
test_result(mfun4 + pw_func, evaluate( mfun1 + mfun5, coordinates) + evaluate(mfun4, coordinates), 12)
test_result(mfun4 - pw_func,  evaluate(mfun4, coordinates) - evaluate( mfun1 + mfun5, coordinates), 12)
test_result(mfun4 * pw_func, evaluate( mfun1 + mfun5, coordinates) * evaluate(mfun4, coordinates), 12)
# Multivariate Sum of Functions
test_result(pw_func + msum, evaluate( mfun1 + mfun5, coordinates) + evaluate(msum, coordinates), 12)
test_result(pw_func - msum, evaluate( mfun1 + mfun5, coordinates) - evaluate(msum, coordinates), 12)
test_result(pw_func * msum, evaluate( mfun1 + mfun5, coordinates) * evaluate(msum, coordinates), 12)
test_result(msum + pw_func, evaluate( mfun1 + mfun5, coordinates) + evaluate(msum, coordinates), 12)
test_result(msum - pw_func,  evaluate(msum, coordinates) - evaluate( mfun1 + mfun5, coordinates), 12)
test_result(msum * pw_func, evaluate( mfun1 + mfun5, coordinates) * evaluate(msum, coordinates), 12)

# Algebra between different Multivariate_Piecewise_Functions
func_array = Array{Union{Missing,Sum_Of_Functions},3}(undef,3,2,2)
func_array[1,1,1] = Sum_Of_Functions(mfun5)
func_array[1,1,2] = Sum_Of_Functions(mfun3)
func_array[1,2,1] = mfun3 + mfun1
func_array[1,2,2] = Sum_Of_Functions(2*mfun1)
func_array[2,1,1] = mfun2 + mfun4
func_array[2,1,2] = mfun1 + mfun5
func_array[2,2,1] = mfun5 + mfun5
func_array[2,2,2] = mfun4 + mfun1 + mfun5
func_array[3,:,:] .= mfun4 + mfun5 + mfun1

pw_func2 = Piecewise_Function(func_array, OrderedDict([:x, :y, :w] .=> [[0.2, 0.5, 1.0], [-4.0, -1.0], [0.0, 5.0]]))
test_result(pw_func + pw_func2, evaluate( (mfun1 + mfun5) + (mfun5 + mfun5), coordinates), 3*4*2*2)
test_result(pw_func - pw_func2, evaluate( (mfun1 + mfun5) - (mfun5 + mfun5), coordinates), 3*4*2*2)
test_result(pw_func * pw_func2, evaluate( (mfun1 + mfun5) * (mfun5 + mfun5), coordinates), 3*4*2*2)

## Testing of Sum_Of_Piecewise_Functions
spwf = Sum_Of_Piecewise_Functions([pw_func, pw_func2])
abs(evaluate(spwf, coordinates) - evaluate(pw_func, coordinates) - evaluate(pw_func2, coordinates) ) < tol
underlying_dimensions(spwf) == union(underlying_dimensions(pw_func), underlying_dimensions(pw_func2))
converted = convert(Piecewise_Function, spwf)
added = (pw_func + pw_func2)
converted.functions_[2,2,2,1].functions_[1] ≂ added.functions_[2,2,2,1].functions_[1]

test_result( (spwf + 1) - 1, evaluate( spwf, coordinates), 2)
test_result( 1 + (1 - spwf ), 2 - evaluate( spwf, coordinates), 2)
test_result( (spwf + 1.0) - 1.0, evaluate( spwf, coordinates), 2)
test_result( 1.0 + (1.0 - spwf ), 2 - evaluate( spwf, coordinates), 2)
test_result( (3 *spwf * 2) / 4, evaluate( spwf, coordinates) * 1.5, 2)
test_result( (3.0 *spwf * 2.0) / 4.0, evaluate( spwf, coordinates) * 1.5, 2)
test_result( (spwf^2), evaluate( spwf, coordinates)^2, 48)
deriv_spec = Dict{Symbol,Int}(:x => 1)
abs(evaluate(derivative(spwf, deriv_spec), coordinates) -  evaluate(derivative(pw_func,deriv_spec), coordinates) - evaluate(derivative(pw_func2, deriv_spec), coordinates)) < 1e-09
integration_limits = Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :w, :z] .=> [(0.25,0.75),(-4.0,-1.0),(0.0,8.0), (0.5, 2.4)])
evaluate(integral(spwf, integration_limits) -  integral(pw_func,integration_limits) - integral(pw_func2, integration_limits), coordinates)     < 1e-09