# constant used for computation
EPSSCALE = 10

# Computes the Wasserstein distance between a univariate discrete distribution
# and a univariate distribution with an implemented function 'cdf'.
#
# Args:
#   Q - univariate SteinDiscrete distribution; support of Q must be sorted
#   targetcdf - CDF of the (other) univariate distribution P
#
# Returns:
#   Tuple of Wasserstein distance and upper bound on numerical integration
#   error.
function wasserstein1d(Q::SteinDiscrete,
                       targetcdf::Function;
                       lower=-Inf,
                       upper=Inf)
    X = Q.support;
    if size(X,2) != 1
        error("function only defined by univariate distributions")
    end
    if !issorted(X)
        error("Support of Q must be sorted")
    end
    q = Q.weights;
    n = length(q);

    # Compute \int_t |Q(t) - P(t)| dt where Q(t) is the cdf of Q and P(t)
    # is the cdf of P
    # First compute contribution of the interval (lower, x_1)
    # We need to be watchful for machine precision issues
    (integral, max_error) = quadgk(t->thresholdpos(targetcdf(t), EPSSCALE*eps(0.0)), lower, X[1])
    # Cumulative Q weight
    weight = 0;
    for i in 1:(n-1)
        # Compute integral contributed by each (x_i, x_{i+1}] interval
        weight = weight + q[i];
        (estimate, error) = quadgk(t->thresholdpos(abs(targetcdf(t) - weight), EPSSCALE*eps(weight)), X[i], X[i+1]);
        integral = integral + estimate;
        max_error = max_error + error;
    end
    # Add contribution of (x_n, upper)
    (estimate, error) = quadgk(t->thresholdpos(1-targetcdf(t), EPSSCALE*eps(-1.0)), X[n], upper)
    integral = integral + estimate;
    max_error = max_error + error;
    (integral, max_error)
end

# Computes the Wasserstein distance between a univariate discrete distribution
# (represented by points and weights) and a univariate distribution with an
# implemented function 'cdf'.
#
# Args:
#   points - n x 1 array of support points
#   weights - n x 1 array of real-valued weights associated with support points
#   target - univariate SteinDistribution implementing the function 'cdf'
#
# Returns:
#   Tuple of Wasserstein distance and upper bound on numerical integration
#   error.
function wasserstein1d(points::Array{Float64};
                       weights=fill(1/size(points,1), size(points,1)),
                       targetcdf=nothing)
    isa(targetcdf, Function) ||
        error("targetcdf must be the cdf of the continuous distribution.")
    wasserstein1d(SteinDiscrete(points, weights), targetcdf)
end

# This is a helper function that used to avoid weird
# machine precision issues when integrating the cdf.
function thresholdpos(x::Float64, threshold::Float64)
    (x <= threshold) ? 0.0 : x
end

# This function computes the earth movers distance between two
# discrete distributions.
#
# It solves a large LP in order to compute the EMD.
function wassersteindiscrete{T<:Number}(;
    xpoints::AbstractArray{T,2}=[],
    xweights::Array{Float64,1}=fill(1.0/size(xpoints,1), size(xpoints,1)),
    ypoints::AbstractArray{T,2}=[],
    yweights::Array{Float64,1}=fill(1.0/size(ypoints,1), size(ypoints,1)),
    solver=nothing)
    # make sure samples are same dimensionality
    @assert size(xpoints, 2) == size(ypoints, 2)
    # make sure solver is defined
    isa(solver, AbstractMathProgSolver) ||
        error("Must specify solver of type AbstractMathProgSolver")
    # setup discrete distributions
    xdist = SteinDiscrete(xpoints, xweights)
    nx = size(xdist.support, 1)
    ydist = SteinDiscrete(ypoints, yweights)
    ny = size(ydist.support, 1)
    # prepare the manhattan extension graph
    complete_support = [xdist.support; ydist.support]
    graph_support, edgeset = makel1spanner(complete_support)
    # compute the supply for the network problem
    ntotal = size(graph_support,1)
    qx = zeros(ntotal)
    qx[1:nx] = xdist.weights
    qy = zeros(ntotal)
    qy[(nx+1):(nx+ny)] = ydist.weights
    # now solve the min-cost max-flow problem
    (emd, flow, status) = mincostflow(graph_support,
                                      edgeset,
                                      qx,
                                      qy,
                                      solver)
    return (emd, size(graph_support,1), length(edgeset), status)
end

# This function computes the earth movers distance between two discrete
# univariate distributions.
# This is simply int |F1(z) - F2(z)| dz.
function wasserstein1d{T<:Number}(;
    xpoints::AbstractArray{T,1}=[],
    xweights::Array{Float64,1}=fill(1.0/length(xpoints), length(xpoints)),
    ypoints::AbstractArray{T,1}=[],
    yweights::Array{Float64,1}=fill(1.0/size(ypoints,1), size(ypoints,1)))
    # setup discrete distributions
    xdist = SteinDiscrete(xpoints, xweights)
    ydist = SteinDiscrete(ypoints, yweights)
    # now grab the goodies
    xsupp = xdist.support
    qx = xdist.weights
    ysupp = ydist.support
    qy = ydist.weights
    # initialize the indices we walk up the vectors
    ix = 0; iy = 0
    wasserstein = 0.0
    Fx = 0.0
    Fy = 0.0
    currval = -Inf
    if xsupp[1] <= ysupp[1]
        ix = 1
        Fx = qx[1]
        currval = xsupp[1]
    end
    if xsupp[1] >= ysupp[1]
        iy = 1
        Fy = qy[1]
        currval = ysupp[1]
    end
    while ix != length(xsupp) || iy != length(ysupp)
        # handle edge case when we've cycled through the distribution
        if ix == length(xsupp)
            iy += 1
            wasserstein += (ysupp[iy] - currval) * (1.0 - Fy)
            Fy += qy[iy]
            currval = ysupp[iy]
            continue
        end
        if iy == length(ysupp)
            ix += 1
            wasserstein += (xsupp[ix] - currval) * (1.0 - Fx)
            Fx += qx[ix]
            currval = xsupp[ix]
            continue
        end
        # now do the look ahead
        nextval = min(xsupp[ix+1], ysupp[iy+1])
        wasserstein += (nextval - currval) * abs(Fx - Fy)
        if xsupp[ix+1] <= ysupp[iy+1]
            ix += 1
            Fx += qx[ix]
            currval = xsupp[ix]
        else
            iy +=1
            Fy += qy[iy]
            currval = ysupp[iy]
        end
    end
    (wasserstein, ix + iy, ix + iy - 1, :Optimal)
end

# This function approximates the wasserstein metric by using two facts:
# 1. We can compute the EMD for two different discrete distributions.
# 2. The empirical wasserstein metric is root{n}-consistent for the
# true wasserstein metric.
function approxwasserstein{T<:Number}(;
    points::AbstractArray{T,2}=[],
    weights::Array{Float64,1}=fill(1.0/size(points,1), size(points,1)),
    targetsamplegen=nothing,
    replicates::Int=30,
    alpha::Float64=0.95,
    solver=nothing)
    # make sure targetsamplegen is given
    isa(targetsamplegen, Function) ||
        error("Must supply targetsamplegen as a function to generate random samples.")
    # compute a few estimates of wasserstein metric
    wassersteinestimates = Array{Float64}(replicates)
    for ii in 1:replicates
        sample = targetsamplegen()
        (emd, numnodes, numedges, status) =
            wassersteindiscrete(xpoints=points,
                                xweights=weights,
                                ypoints=sample,
                                solver=solver)
        wassersteinestimates[ii] = emd
    end
    # compute CI
    noise = Distributions.TDist(replicates - 1)
    wassmean = mean(wassersteinestimates)
    wassstd = std(wassersteinestimates) / sqrt(replicates)
    zalpha = abs(Distributions.quantile(noise, (1.0 - alpha) / 2))
    return (wassmean - zalpha * wassstd, wassmean + zalpha * wassstd)
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

module TestGivens

using Test, LinearAlgebra, Random
using LinearAlgebra: rmul!, lmul!, Givens

# Test givens rotations
@testset for elty in (Float32, Float64, ComplexF32, ComplexF64)
    if elty <: Real
        raw_A = convert(Matrix{elty}, randn(10,10))
    else
        raw_A = convert(Matrix{elty}, complex.(randn(10,10),randn(10,10)))
    end
    @testset for A in (raw_A, view(raw_A, 1:10, 1:10))
        Ac = copy(A)
        R = LinearAlgebra.Rotation(LinearAlgebra.Givens{elty}[])
        for j = 1:8
            for i = j+2:10
                G, _ = givens(A, j+1, i, j)
                lmul!(G, A)
                rmul!(A, adjoint(G))
                lmul!(G, R)

                @test lmul!(G,Matrix{elty}(I, 10, 10)) == [G[i,j] for i=1:10,j=1:10]

                @testset "transposes" begin
                    @test G'*G*Matrix(elty(1)I, 10, 10)   ≈ Matrix(I, 10, 10)
                    @test (G*Matrix(elty(1)I, 10, 10))*G' ≈ Matrix(I, 10, 10)
                    @test copy(R')*(R*Matrix(elty(1)I, 10, 10)) ≈ Matrix(I, 10, 10)
                    @test_throws ErrorException transpose(G)
                    @test_throws ErrorException transpose(R)
                end
            end
        end
        @test_throws ArgumentError givens(A, 3, 3, 2)
        @test_throws ArgumentError givens(one(elty),zero(elty),2,2)
        G, _ = givens(one(elty),zero(elty),11,12)
        @test_throws DimensionMismatch lmul!(G, A)
        @test_throws DimensionMismatch rmul!(A, adjoint(G))
        @test abs.(A) ≈ abs.(hessenberg(Ac).H)
        @test opnorm(R*Matrix{elty}(I, 10, 10)) ≈ one(elty)

        I10 = Matrix{elty}(I, 10, 10)
        G, _ = givens(one(elty),zero(elty),9,10)
        @test (G*I10)' * (G*I10) ≈ I10
        K, _ = givens(zero(elty),one(elty),9,10)
        @test (K*I10)' * (K*I10) ≈ I10

        @testset "Givens * vectors" begin
            if isa(A, Array)
                x = A[:, 1]
            else
                x = view(A, 1:10, 1)
            end
            G, r = givens(x[2], x[4], 2, 4)
            @test (G*x)[2] ≈ r
            @test abs((G*x)[4]) < eps(real(elty))
            @inferred givens(x[2], x[4], 2, 4)

            G, r = givens(x, 2, 4)
            @test (G*x)[2] ≈ r
            @test abs((G*x)[4]) < eps(real(elty))
            @inferred givens(x, 2, 4)

            G, r = givens(x, 4, 2)
            @test (G*x)[4] ≈ r
            @test abs((G*x)[2]) < eps(real(elty))
        end
    end
end

# 36430
# dimensional correctness:
const BASE_TEST_PATH = joinpath(Sys.BINDIR, "..", "share", "julia", "test")
isdefined(Main, :Furlongs) || @eval Main include(joinpath($(BASE_TEST_PATH), "testhelpers", "Furlongs.jl"))
using .Main.Furlongs

@testset "testing dimensions with Furlongs" begin
    @test_throws MethodError givens(Furlong(1.0), Furlong(2.0), 1, 2)
end

const TNumber = Union{Float64,ComplexF64}
struct MockUnitful{T<:TNumber} <: Number
    data::T
    MockUnitful(data::T) where T<:TNumber = new{T}(data)
end
import Base: *, /, one, oneunit
*(a::MockUnitful{T}, b::T) where T<:TNumber = MockUnitful(a.data * b)
*(a::T, b::MockUnitful{T}) where T<:TNumber = MockUnitful(a * b.data)
*(a::MockUnitful{T}, b::MockUnitful{T}) where T<:TNumber = MockUnitful(a.data * b.data)
/(a::MockUnitful{T}, b::MockUnitful{T}) where T<:TNumber = a.data / b.data
one(::Type{<:MockUnitful{T}}) where T = one(T)
oneunit(::Type{<:MockUnitful{T}}) where T = MockUnitful(one(T))

@testset "unitful givens rotation unitful $T " for T in (Float64, ComplexF64)
    g, r = givens(MockUnitful(T(3)), MockUnitful(T(4)), 1, 2)
    @test g.c ≈ 3/5
    @test g.s ≈ 4/5
    @test r.data ≈ 5.0
end

end # module TestGivens


struct Unicycle{T<:Real}
    ΔT::T
end

function Base.getproperty(system::Unicycle, sym::Symbol)
    if sym === :n_states
        4
    elseif sym === :n_controls
        2
    else
        getfield(system, sym)
    end
end

function DynamicsModelInterface.next_x(system::Unicycle, x_t, u_t)
    ΔT = system.ΔT
    @assert only(size(x_t)) == 4
    @assert only(size(u_t)) == 2
    px, py, v, θ = x_t
    Δv, Δθ = u_t
    [px + ΔT * v * cos(θ), py + ΔT * v * sin(θ), v + Δv, θ + Δθ]
end

# These constraints encode the dynamics of a unicycle with state layout x_t = [px, pyL, v, θ] and
# inputs u_t = [Δv, Δθ].
function DynamicsModelInterface.add_dynamics_constraints!(system::Unicycle, opt_model, x, u)
    ΔT = system.ΔT
    T = size(x, 2)

    # auxiliary variables for nonlinearities
    cosθ = @variable(opt_model, [1:T])
    @NLconstraint(opt_model, [t = 1:T], cosθ[t] == cos(x[4, t]))
    sinθ = @variable(opt_model, [1:T])
    @NLconstraint(opt_model, [t = 1:T], sinθ[t] == sin(x[4, t]))

    @constraint(
        opt_model,
        [t = 1:(T - 1)],
        x[:, t + 1] .== [
            x[1, t] + ΔT * x[3, t] * cosθ[t],
            x[2, t] + ΔT * x[3, t] * sinθ[t],
            x[3, t] + u[1, t],
            x[4, t] + u[2, t],
        ]
    )
end

function DynamicsModelInterface.add_dynamics_jacobians!(system::Unicycle, opt_model, x, u)
    ΔT = system.ΔT
    n_states, T = size(x)
    n_controls = size(u, 1)

    # auxiliary variables for nonlinearities
    cosθ = @variable(opt_model, [1:T])
    @NLconstraint(opt_model, [t = 1:T], cosθ[t] == cos(x[4, t]))
    sinθ = @variable(opt_model, [1:T])
    @NLconstraint(opt_model, [t = 1:T], sinθ[t] == sin(x[4, t]))

    # jacobians of the dynamics in x
    dfdx = @variable(opt_model, [1:n_states, 1:n_states, 1:T])
    @constraint(
        opt_model,
        [t = 1:T],
        dfdx[:, :, t] .== [
            1 0 ΔT*cosθ[t] -ΔT*x[3, t]*sinθ[t]
            0 1 ΔT*sinθ[t] +ΔT*x[3, t]*cosθ[t]
            0 0 1 0
            0 0 0 1
        ]
    )

    # jacobians of the dynamics in u
    dfdu = [
        0 0
        0 0
        1 0
        0 1
    ] .* reshape(ones(T), 1, 1, :)

    (; dx = dfdx, du = dfdu)
end

#========================================== Visualization ==========================================#

function TrajectoryVisualization.trajectory_data(::Unicycle, x, player = 1)
    [
        (; px = xi[1], py = xi[2], v = xi[3], θ = xi[4], player, t) for
        (t, xi) in enumerate(eachcol(x))
    ]
end


mutable struct MonteCarloDiffEqModel{absdiffeqmodel <: DiffEqBase.AbstractEnsembleProblem} <: ScalarMonteCarloProcess{Float64}
    model::absdiffeqmodel
    final_trasform::Function
    underlying::AbstractUnderlying
    function MonteCarloDiffEqModel(model::absdiffeqmodel, final_trasform::Function, underlying::AbstractUnderlying) where {absdiffeqmodel <: DiffEqBase.AbstractEnsembleProblem}
        return new{absdiffeqmodel}(model, final_trasform, underlying)
    end
    function MonteCarloDiffEqModel(model::absdiffeqmodel, underlying::AbstractUnderlying) where {absdiffeqmodel <: DiffEqBase.AbstractEnsembleProblem}
        func(x) = identity(x)
        return new{absdiffeqmodel}(model, func, underlying)
    end
end

export MonteCarloDiffeEqModel;

function simulate!(X, mcProcess::MonteCarloDiffEqModel, rfCurve::ZeroRate, mcBaseData::MonteCarloConfiguration{type1, type2, type3, type4, type5}, T::numb) where {numb <: Number, type1 <: Number, type2 <: Number, type3 <: AbstractMonteCarloMethod, type4 <: BaseMode, type5 <: Random.AbstractRNG}
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    Dt = T / Nstep
    @assert T > 0.0
    diffeqmodel = mcProcess.model
    sol = solve(diffeqmodel, SOSRI(), EnsembleThreads(), trajectories = Nsim, dt = Dt, adaptive = false)
    X .= [path.u[j] for path in sol.u, j = 1:(Nstep+1)]
    f(x) = mcProcess.final_trasform(x)
    broadcast!(f, X, X)
    nothing
end


"""
General observers
"""

abstract type AbstractObserver{T} end

"""
    LinearObserver(H::Matrix[, R::Matrix])
    LinearObserver(H::AbstractFloat[, R::AbstractFloat])
    LinearObserver(H::Vector[, R::Matrix])
    LinearObserver(H::RowVector[, R::Union{Matrix, AbstractFloat}])

Linear observer with the form

\$y_k = H x_k + v_k \\quad \\text{where} \\quad  v_k \\sim N(0, R)\$

A linear observer can be constructed with both matrix and scalar inputs. The
observer will not model measurement noise if the covariance R is not provided.
"""
struct LinearObserver{T<:AbstractFloat} <: AbstractObserver{T}
    H::Matrix{T}
    R::Covariance{T}

    function LinearObserver(H::Matrix{T}, R::Covariance{T}) where T
        if (size(H,1), size(H,1)) != size(R)
            throw(DimensionMismatch("Incompatible size of system matrices."))
        end
        new{T}(H, R)
    end
end
LinearObserver{T}(H::Matrix{T}) = LinearObserver(H,zeros(T,size(H,1),size(H,1)))
LinearObserver{T<:AbstractFloat}(H::T) =
    LinearObserver(reshape([H],1,1), zeros(T,1,1))
LinearObserver{T<:AbstractFloat}(H::T, R::T) =
    LinearObserver(reshape([H],1,1), reshape([R],1,1))
LinearObserver{T<:AbstractFloat}(H::Matrix{T}, R::T) =
    LinearObserver(H, reshape([R],1,1))

LinearObserver(H::Vector) = LinearObserver(reshape(H,length(H),1))
LinearObserver{T<:AbstractFloat}(H::Vector{T}, R::Covariance{T}) =
    LinearObserver(reshape(H,length(H),1), R)

LinearObserver(H::RowVector) = LinearObserver(Array(H))
LinearObserver{T<:AbstractFloat}(H::RowVector{T}, R::T) =
    LinearObserver(Array(H), reshape([R],1,1))
LinearObserver{T<:AbstractFloat}(H::RowVector{T}, R::Covariance{T}) =
    LinearObserver(Array(H), R)

"""
    NonlinearObserver(H::Function, dH_dx::Function, R::Matrix)
    NonlinearObserver(H::Function, dH_dx::Function, R::AbstractFloat)

Nonlinear observer with the form

\$y_k = H(k, x_k) + v \\quad  \\text{where} \\quad v \\sim N(0, R)\$

NonlinearObserver constructors require both the measurement function,
`H(t::Number, x::Vector)`, and its Jacobian, `dH_dx(t::Number, x::Vector)`. Both
`H()` and `dH_dx()` must return a vector.
"""
struct NonlinearObserver{T<:AbstractFloat} <: AbstractObserver{T}
    H::Function
    dH_dx::Function
    R::Covariance{T}

    function NonlinearObserver(H::Function, dH_dx::Function,
                               R::Covariance{T}) where T
        new{T}(H, dH_dx, R)
    end
end
NonlinearObserver(H::Function, dH_dx::Function, R::AbstractFloat) =
    NonlinearObserver(H, dH_dx, reshape([R], 1, 1))
function NonlinearObserver(lin_obs::LinearObserver)
    H(t,x) = lin_obs.H*x
    dH_dx(t,x) = lin_obs.H
    return NonlinearObserver(H, dH_dx, lin_obs.R)
end

"""
    assert_compatibility(obs::LinearObserver{T}, state::AbstractState{T})

Require linear observer dimensions to be compatible with input state.

    assert_compatibility(state::AbstractState{T}, obs::LinearObserver{T})

Require linear observer dimensions to be left-compatible with input state.
"""
function assert_compatibility{T}(obs::LinearObserver{T},state::AbstractState{T})
    if size(obs.H, 2) != length(state.x)
        throw(DimensionMismatch("Observer incompatible with input state."))
    end
end
function assert_compatibility{T}(state::AbstractState{T},obs::AbstractObserver{T})
    if size(obs.R, 1) != length(state.x)
        throw(DimensionMismatch("Observer incompatible with input state."))
    end
end

"""
    observable(A::Matrix, H::Matrix)
    observable(sys::LinearSystem, obs::LinearObserver)

Evaluate the observability of a linear model. Returns a boolean True/False
result.
"""
function observable(A, H)
    n = size(A, 2)
    p = size(H, 1)
    observability = zeros(n*p, n)

    observability[1:p, :] .= H
    for idx = 1:n-1
        observability[p*idx+1:p*(idx+1), :] .= H * A
    end

    rank(observability) == n ? true : false
end
observable(sys::LinearSystem, obs::LinearObserver) = observable(sys.A, obs.H)

"""
    predict(state::AbstractState, obs::AbstractObserver)

Predict a state through an arbitrary observer. The `predict` method will advance
both the state vector and, if necessary, its covariance through the input
measurement model. For NonlinearObservers, the covariance prediction is only a
linear approximation.
"""
function predict(::AbstractObserver) end
function predict{T}(state::DiscreteState{T}, obs::LinearObserver{T})
    return DiscreteState(obs.H*state.x, state.t)
end
function predict{T}(state::ContinuousState{T}, obs::LinearObserver{T})
    return ContinuousState(obs.H*state.x, state.t)
end
function predict{T}(state::UncertainDiscreteState{T}, obs::LinearObserver{T})
    return UncertainDiscreteState(
        obs.H*state.x, obs.H*state.P*obs.H'+obs.R, state.t)
end
function predict{T}(state::UncertainContinuousState{T}, obs::LinearObserver{T})
    return UncertainContinuousState(
        obs.H*state.x, obs.H*state.P*obs.H'+obs.R, state.t)
end
function predict{T}(state::DiscreteState{T}, obs::NonlinearObserver{T})
    return DiscreteState(obs.H(state.t, state.x), state.t)
end
function predict{T}(state::ContinuousState{T}, obs::NonlinearObserver{T})
    return ContinuousState(obs.H(state.t, state.x), state.t)
end
function predict{T}(state::UncertainDiscreteState{T}, obs::NonlinearObserver{T})
    jac = obs.dH_dx(state.t, state.x)
    return UncertainDiscreteState(
        obs.H(state.t, state.x), jac*state.P*jac'+obs.R, state.t)
end
function predict{T}(state::UncertainContinuousState{T},
                    obs::NonlinearObserver{T})
    jac = obs.dH_dx(state.t, state.x)
    return UncertainContinuousState(
        obs.H(state.t, state.x), jac*state.P*jac'+obs.R, state.t)
end

"""
   simulate(state::AbstractState, obs::AbstractObserver)

Simulate a state observation with initial state error and process noise. This
method returns an absolute state by sampling the initial state (if uncertain),
and propagating that state through the input system in the presence of
measurement noise.
"""
function simulate(::AbstractObserver) end
function simulate(state::AbstractAbsoluteState, obs::AbstractObserver)
    return simulate(make_uncertain(state), obs)
end
function simulate(state::AbstractUncertainState, obs::LinearObserver)
    return sample(predict(state, obs))
end
function simulate(state::AbstractUncertainState, obs::NonlinearObserver)
    return sample(predict(make_uncertain(sample(state)), obs))
end


import UMAP, Plots

function umap(X)
    X2 = convert(Array{Float64,2}, X)
    Y = UMAP.umap(X2, 3)  # ,2;n_neighbors=5);
    f2 = Plots.plot(Y[1, :], Y[2, :], seriestype = :scatter)

    Y = Y';
    # theplot = scatter(Y[:,1], Y[:,2], marker=(2,2,:auto,stroke(0)))  # , color=Int.(allabels[1:size(Y,1)]))
    #=
    f = scatter3d(
        Y[:, 1],
        Y[:, 2],
        Y[:, 3],
        marker = (2, 2, :auto, stroke(0)),
        color = Int.(sum(X, dims = 1)),
    )
    =#
    return Y
end


##################################################
## evaluate binary operations of symbolic objects
## XXX -- this may prove too narryw with the use of ::Sym

+(x::SymbolicObject, y::SymbolicObject) = pycall(sympy.Add, Sym, x, y)
*(x::SymbolicObject, y::SymbolicObject) = pycall(sympy.Mul, Sym, x, y)
-(x::SymbolicObject, y::SymbolicObject) = x + (-y)
-(x::SymbolicObject)                    =  (-1) * x
/(x::SymbolicObject, y::SymbolicObject) = x * inv(y)
^(x::SymbolicObject, y::SymbolicObject) = pycall(sympy.Pow, Sym, x, y)::Sym
^(x::SymbolicObject, y::Rational) = x^convert(Sym,y)
#^(x::SymbolicObject, y::Integer) = x^convert(Sym,y) # no Union{Integer, Rational}, as that has ambiguity
//(x::SymbolicObject, y::Int) = x / Sym(y)
//(x::SymbolicObject, y::Rational) = x / Sym(y)
//(x::SymbolicObject, y::SymbolicObject) = x / y

\(x::SymbolicObject, y::SymbolicObject) = (y'/x')' # ?

#inv(x::Sym) = x\one(x)
#inv(x::Sym) = x^(-1)
inv(x::Sym) = pycall(sympy.Pow, Sym, x, -1)


kolmogorov_smirnov(cdf_1_evaluated, pm_1, cdf_2_evaluated, pm_2, grid) = maximum(abs.(cdf_1_evaluated-cdf_2_evaluated))
function L1(f1_evaluated, pm_1, f2_evaluated, pm_2, grid)
    d = f1_evaluated-f2_evaluated
    d = sum(abs.(d[1:end-1,:] + d[2:end,:])./2.0.*diff(grid))
    return log10(d)
end
function L2(f1_evaluated, pm_1, f2_evaluated, pm_2, grid)
    d = f1_evaluated-f2_evaluated
    d = sum(((d[1:end-1,:] + d[2:end,:]).^2)./2.0.*diff(grid))
    return log10(d)
end
function L1_cell_probs(f1_evaluated, pm_1, f2_evaluated, pm_2, grid)
    d = sum(abs.(f1_evaluated-f2_evaluated))
    d += sum(abs.(pm_1[:]-pm_2[:]))
    return log10(d)
end


"""
Canonical polyadic N-mode tensor decomposition (CANDECOMP/PARAFAC).
"""
struct CANDECOMP{T<:Number, N} <: TensorDecomposition{T, N}
  factors::NTuple{N, Matrix{T}} # rank X i-th tensor mode factor matrices
  lambdas::Vector{T}
  props::Dict{Symbol, Any}      # extra properties

  CANDECOMP{T, N}(factors::NTuple{N, Matrix{T}}, lambdas::Vector{T}) where {T<:Number, N} =
    new(factors, lambdas, Dict{Symbol, Any}())
end

CANDECOMP(factors::NTuple{N, Matrix{T}}, lambdas::Vector{T}) where {T<:Number, N} =
    CANDECOMP{T, N}(factors, lambdas)

"""
Returns CANDECOMP rank.
"""
LinearAlgebra.rank(decomp::CANDECOMP) = length(decomp.lambdas)

"""
Re-composes the tensor from CANDECOMP decomposition.
"""
@generated function compose!(dest::Array{T,N}, factors::NTuple{N, Matrix{T}}, lambdas::Vector{T}) where {T <: Number, N}
  quote
    csize = @ntuple($N, k -> size(factors[k], 2))
    @nall($N, k -> csize[k] == length(lambdas)) ||
        throw(DimensionMismatch("length of lambdas ($(length(lambdas))) doesn't match the factors columns ($csize)"))
    rsize = @ntuple($N, k -> size(factors[k], 1))
    size(dest) == rsize ||
        throw(DimensionMismatch("dest dimensions $(size(dest)) do not match the factors rows ($rsize)"))
    @nloops $N i dest begin
        elm = zero(T)
        # FIXME traversal of factors is not very efficient, would be better to have them transposed
        @inbounds for j in eachindex(lambdas)
            elm += lambdas[j] * (*(@ntuple($N, k -> factors[k][i_k, j])...))
        end
        @inbounds @nref($N, dest, i) = elm
    end
    return dest
  end
end

compose(factors::NTuple{N, Matrix{T}}, lambdas::Vector{T}) where {T <: Number, N} =
    compose!(Array{T, N}(undef, ntuple(i -> size(factors[i], 1), N)), factors, lambdas)

compose(decomp::CANDECOMP) = compose(decomp.factors, decomp.lambdas)

compose!(dest::Array{T,N}, decomp::CANDECOMP{T,N}) where {T <: Number, N} =
    compose!(dest, decomp.factors, decomp.lambdas)

"""
Calculates canonical polyadic tensor decomposition (CANDECOMP/PARAFAC).

Returns:
  `CANDECOMP` object
"""
function candecomp(tnsr::AbstractArray{T,N},
                   r::Integer,
                   initial_guess::NTuple{N, Matrix{T}};
                   method::Symbol=:ALS,
                   tol::Float64=1e-5,
                   maxiter::Integer=100,
                   compute_error::Bool=false,
                   verbose::Bool=true) where {T <: Number, N}

    _check_tensor(tnsr, r)
    verbose && @info("initializing factor matrices...")
    all([(size(tnsr, i), r) == size(initial_guess[i]) for i in 1:N]) || throw(ArgumentError("dimension of initial guess does not match input tensor."))
    factors = collect(Matrix{T}, initial_guess)
    verbose && @info("applying candecomp $method method...")
    res = _candecomp(Val{method}, tnsr, r, factors, tol, maxiter, verbose)
    compute_error && _set_rel_residue(res, tnsr)
    return res
end

"Stub for non-implemented CANDECOMP algorithms."
_candecomp(
  method::Val{S},
  tnsr::AbstractArray{T,N},
  r::Integer,
  factors::Vector{Matrix{T}},
  tol::Float64,
  maxiter::Integer,
  verbose::Bool) where {T,N,S<:Symbol} =
    throw(ArgumentError("Unknown CANDECOMP method: $method"))

"""
Computes CANDECOMP by ALS (Alternating Least Squares) method.
"""
function _candecomp(
    method::Type{Val{:ALS}},
    tnsr::AbstractArray{T,N},
    r::Integer,
    factors::Vector{Matrix{T}},
    tol::Float64,
    maxiter::Integer,
    verbose::Bool) where {T <: Number, N}

    gram = [F'F for F in factors]
    tnsr_norm = norm(tnsr)
    tnsr_flat = _col_unfold(tnsr, N)
    tnsr_size = size(tnsr)
    niters = 0
    converged = false
    resid = tnsr_norm
    lbds = Matrix{T}(undef, 1, r)
    V = Matrix{T}(undef, length(tnsr) ÷ minimum(tnsr_size), r)
    pb = Progress(maxiter, "ALS iterations ")
    while !converged && niters < maxiter
        update!(pb, niters)
        nVi = 0
        @inbounds for i in 1:N
            idx = [N:-1:i + 1; i - 1:-1:1]
            nVi = length(tnsr) ÷ tnsr_size[i]
            Vi = view(V, 1:nVi, :)
            khatrirao!(Vi, factors[idx])
            factors[i] .= _row_unfold(tnsr, i) * Vi / reduce((x, y) -> x .* y, gram[idx])
            sum!(abs, lbds, factors[i])
            factors[i] ./= lbds
            mul!(gram[i], factors[i]', factors[i])
        end
        resid_old = resid
        resid = norm(view(V, 1:nVi, :) * permutedims((factors[N] .* lbds)) .- tnsr_flat)
        converged = abs(resid - resid_old) < tol * resid_old
        niters += 1
    end
    finish!(pb)

    verbose && _iter_status(converged, niters, maxiter)

    return CANDECOMP(tuple(factors...), dropdims(lbds, dims=1))
end

"""
Computes CANDECOMP by SGSD (Simultaneous Generalized Schur Decomposition) method.
"""
function _candecomp(
    method::Type{Val{:SGSD}},
    tnsr::AbstractArray{T, N},
    r::Int,
    factors::Vector{<:StridedMatrix},
    tol::Float64,
    maxiter::Integer,
    verbose::Bool) where {T <: Number, N}

    ndims(tnsr) == 3 || throw(ArgumentError("This algorithm only applies to 3-mode tensors."))
    length(factors) == 3 || throw(ArgumentError("3 factor matrices expected."))
    (n1, n2, n3) = size(tnsr)
    IB1 = min(n1 - 1, r)
    IB2 = (n2 == r) ? 2 : 1
    Q = qr(factors[1]).Q * Matrix(I, n1, n1)
    Z = qr(reverse(factors[2], dims=2)).Q * Matrix(I, n2, n2)
    q = Matrix{T}(undef, n1, n1)
    z = Matrix{T}(undef, n2, n2)

    @tensor R[4,5,3] := tnsr[1,2,3] * Q[1,4] * Z[2,5]
    R2 = similar(R)

    res = norm(tnsr)
    converged = false
    niters = 0
    pb = Progress(maxiter, "SGSD iterations ")
    @inbounds while !converged && niters < maxiter
        update!(pb, niters)
        # reset to diagonal
        @inbounds for i in 1:n1, j in 1:n1
            q[i, j] = ifelse(i == j, 1.0, 0.0)
        end
        @inbounds for i in 1:n2, j in 1:n2
            z[i, j] = ifelse(i == j, 1.0, 0.0)
        end

        for i in 1:IB1
            q[:, i:n1] *= svd(view(q, :, i:n1)' * view(R, :, n2 - r + i, :)).U
        end

        @tensor R2[4,2,3] = R[1,2,3] * q[1,4]
        for i in r:-1:IB2
            z[:, 1:n2 - r + i] *= reverse(svd(view(R2, i, :, :)' * view(z, :, 1:n2 - r + i)).V, dims=2)
        end
        Q *= q
        Z *= z

        @tensor R[4,5,3] = tnsr[1,2,3] * Q[1,4] * Z[2,5]

        res_old = res
        res = norm(tril(dropdims(sum(abs2, R, dims=3), dims=3), n2 - r - 1))
        converged = abs(res - res_old) < tol * res_old
        niters += 1
    end
    finish!(pb)
    verbose && _iter_status(converged, niters, maxiter)

    R = R[1:r, n2-r+1:n2, :]
    M = Array{T,3}(undef, r, r, 2)
    @inbounds for k in 1:2, j in 1:r, i in 1:r
        M[i, j, k] = i == j ? 1.0 : 0.0
    end
    @inbounds for i in (r - 1):-1:1, j = (i + 1):r
        if i + 1 < j
            d = (i + 1):(j - 1)
            M[i, j, :] = hcat(view(R, j, j, :), view(R, i, i, :)) \
                (view(R, i, j, :) .- mapslices(R3 -> sum(view(M, i, d, 1) .* diag(R3)[d] .* view(M, d, j, 2)), R, dims=(1, 2)))
        end
    end

    factors[1] = view(Q, :, 1:r) * view(M, :, :, 1)
    factors[2] = view(Z, :, n2 - r + 1:n2) * view(M, :, :, 2)'
    factors[3] = _row_unfold(tnsr, 3) * khatrirao(factors[2], factors[1]) /
        ((factors[2]'factors[2]) .* (factors[1]'factors[1]))

    # normalize
    lbds = ones(1, r)
    for i in 1:3
        lbd = mapslices(norm, factors[i], dims=1)
        factors[i] ./= lbd
        lbds .*= lbd
    end

    return CANDECOMP(tuple(factors...), dropdims(lbds, dims=1))
end

"""
Codes of implemented CANDECOMP methods.
"""
const CANDECOMP_methods = Symbol[t.parameters[1] for t in filter(t -> isa(t, DataType),
                                                                 [Base.unwrap_unionall(m.sig).parameters[2].parameters[1] for m in methods(_candecomp)])] |> Set{Symbol}


export CLT, learn_chow_liu_tree, parent_vector
using Graphs: SimpleGraph, SimpleDiGraph, complete_graph, kruskal_mst,
    bfs_tree, center, connected_components, induced_subgraph, src, dst
using SimpleWeightedGraphs: SimpleWeightedGraph
using MetaGraphs: MetaDiGraph, set_prop!, props, nv, ne, edges, add_edge!, vertices

#####################
# Learn a Chow-Liu tree from (weighted) data
#####################

"""
Chow-Liu Tree
"""
const CLT = MetaDiGraph

"""
learn a Chow-Liu tree from training set `train_x`, with Laplace smoothing factor `α`, specifying the tree root by `clt_root`
return a `CLT`
"""
function learn_chow_liu_tree(train_x; α = 1.0, clt_root="graph_center",
        weight=ones(Float64, num_examples(train_x)))::CLT
    features_num = num_features(train_x)

    # calculate mutual information
    (dis_cache, MI) = mutual_information(train_x, weight; α = α)

    # maximum spanning tree/ forest
    g = SimpleWeightedGraph(complete_graph(features_num))
    mst_edges = kruskal_mst(g,- MI)
    tree = SimpleGraph(features_num)
    map(mst_edges) do edge
        add_edge!(tree, src(edge), dst(edge))
    end

    # Build rooted tree / forest
    if clt_root == "graph_center"
        clt = SimpleDiGraph(features_num)
        for c in filter(c -> (length(c) > 1), connected_components(tree))
            sg, vmap = induced_subgraph(tree, c)
            sub_root = vmap[center(sg)[1]]
            clt = union(clt, bfs_tree(tree, sub_root))
        end
    elseif clt_root == "rand"
        roots = [rand(c) for c in connected_components(tree)]
        clt = SimpleDiGraph(features_num)
        for root in roots clt = union(clt, bfs_tree(tree, root)) end
    else
        error("Cannot learn CLT with root $clt_root")
    end

    clt = MetaDiGraph(clt)
    parent = parent_vector(clt)
    for (c, p) in enumerate(parent)
        set_prop!(clt, c, :parent, p)
    end

    for v in vertices(clt)
        p = parent[v]
        cpt_matrix = get_cpt(p, v, dis_cache)
        set_prop!(clt, v, :cpt, cpt_matrix)
    end

    return clt
end

"""
Calculate CPT of `child` conditioned on `parent` from `dis_cache`
"""
function get_cpt(parent, child, dis_cache)
    if parent == 0
        p = dis_cache.marginal[child, :]
        return Dict(0=>p[1], 1=>p[2])
    else
        p = dis_cache.pairwise[child, parent, :] ./ [dis_cache.marginal[parent, :]; dis_cache.marginal[parent, :]]
        @. p[isnan(p)] = 0; @. p[p==Inf] = 0; @. p[p == -Inf] = 0
        return Dict((0,0)=>p[1], (1,0)=>p[3], (0,1)=>p[2], (1,1)=>p[4]) #p(child|parent)
    end
end

"Get parent vector of a tree"
function parent_vector(tree::CLT)::Vector{Int64}
    v = zeros(Int64, nv(tree)) # parent of roots is 0
    foreach(e->v[dst(e)] = src(e), edges(tree))
    return v
end

import LogicCircuits: print_tree
"Print edges and vertices of a ChowLiu tree"
function print_tree(clt::CLT)
    for e in edges(clt) print(e); print(" ");end
    if clt isa MetaDiGraph
        for v in vertices(clt) print(v); print(" "); println(props(clt, v)) end
    end
end


Lar = LinearAlgebraicRepresentation

"""
	frag_face_channel(in_chan, out_chan, V, EV, FE, sp_idx)

Parallel fragmentation of faces in `FE` against faces in `sp_idx`.
"""
function frag_face_channel(in_chan, out_chan, V, EV, FE, sp_idx)
    run_loop = true
    while run_loop
        sigma = take!(in_chan)
        if sigma != -1
            put!(out_chan, frag_face(V, EV, FE, sp_idx, sigma))
        else
            run_loop = false
        end
    end
end

"""
	frag_face(V, EV, FE, sp_idx, sigma)

`sigma` face fragmentation against faces in `sp_idx[sigma]`
"""
function frag_face(V, EV, FE, sp_idx, sigma)

    vs_num = size(V, 1)

	# 2D transformation of sigma face
    sigmavs = (abs.(FE[sigma:sigma,:])*abs.(EV))[1,:].nzind
    sV = V[sigmavs, :]
    sEV = EV[FE[sigma, :].nzind, sigmavs]
    M = Lar.Arrangement.submanifold_mapping(sV)
    tV = ([V ones(vs_num)]*M)[:, 1:3]  # folle convertire *tutti* i vertici
    sV = tV[sigmavs, :]
    # sigma face intersection with faces in sp_idx[sigma]
    for i in sp_idx[sigma]
        tmpV, tmpEV = Lar.Arrangement.face_int(tV, EV, FE[i, :])

        sV, sEV = Lar.skel_merge(sV, sEV, tmpV, tmpEV)
    end

    # computation of 2D arrangement of sigma face
    sV = sV[:, 1:2]
    nV, nEV, nFE = planar_arrangement(sV, sEV, sparsevec(ones(Int8, length(sigmavs))))
    if nV == nothing ## not possible !! ... (each original face maps to its decomposition)
        return [], spzeros(Int8, 0,0), spzeros(Int8, 0,0)
    end
    nvsize = size(nV, 1)
    nV = [nV zeros(nvsize) ones(nvsize)]*inv(M)[:, 1:3] ## ????
    return nV, nEV, nFE
end

function merge_vertices(V::Lar.Points, EV::Lar.ChainOp, FE::Lar.ChainOp, err=1e-4)
    vertsnum = size(V, 1)
    edgenum = size(EV, 1)
    facenum = size(FE, 1)
    newverts = zeros(Int, vertsnum)
    # KDTree constructor needs an explicit array of Float64
    V = Array{Float64,2}(V)
    W = convert(Lar.Points, LinearAlgebra.transpose(V))
    kdtree = KDTree(W)
	# remove vertices congruent to a single representative
    todelete = []
    i = 1
    for vi in 1:vertsnum
        if !(vi in todelete)
            nearvs = Lar.inrange(kdtree, V[vi, :], err)
            newverts[nearvs] .= i
            nearvs = setdiff(nearvs, vi)
            todelete = union(todelete, nearvs)
            i = i + 1
        end
    end
    nV = V[setdiff(collect(1:vertsnum), todelete), :]

    # translate edges to take congruence into account
    edges = Array{Tuple{Int, Int}, 1}(undef, edgenum)
    oedges = Array{Tuple{Int, Int}, 1}(undef, edgenum)
    for ei in 1:edgenum
        v1, v2 = EV[ei, :].nzind
        edges[ei] = Tuple{Int, Int}(sort([newverts[v1], newverts[v2]]))
        oedges[ei] = Tuple{Int, Int}(sort([v1, v2]))
    end
    nedges = union(edges)
    # remove edges of zero length
    nedges = filter(t->t[1]!=t[2], nedges)
    nedgenum = length(nedges)
    nEV = spzeros(Int8, nedgenum, size(nV, 1))

    etuple2idx = Dict{Tuple{Int, Int}, Int}()
    for ei in 1:nedgenum
    	begin
        	nEV[ei, collect(nedges[ei])] .= 1
        	nEV
        end
        etuple2idx[nedges[ei]] = ei
    end
    for e in 1:nedgenum
    	v1,v2 = findnz(nEV[e,:])[1]
    	nEV[e,v1] = -1; nEV[e,v2] = 1
    end

    # compute new faces to take congruence into account
    faces = [[
        map(x->newverts[x], FE[fi, ei] > 0 ? oedges[ei] : reverse(oedges[ei]))
        for ei in FE[fi, :].nzind
    ] for fi in 1:facenum]

    visited = []
    function filter_fn(face)

        verts = []
        map(e->verts = union(verts, collect(e)), face)
        verts = Set(verts)

        if !(verts in visited)
            push!(visited, verts)
            return true
        end
        return false
    end

    nfaces = filter(filter_fn, faces)

    nfacenum = length(nfaces)
    nFE = spzeros(Int8, nfacenum, size(nEV, 1))

    for fi in 1:nfacenum
        for edge in nfaces[fi]
            ei = etuple2idx[Tuple{Int, Int}(sort(collect(edge)))]
            nFE[fi, ei] = sign(edge[2] - edge[1])
        end
    end

    return Lar.Points(nV), nEV, nFE
end

function spatial_arrangement_1(
		V::Lar.Points,
		copEV::Lar.ChainOp,
		copFE::Lar.ChainOp, multiproc::Bool=false)

	# spaceindex computation
	FV = Lar.compute_FV( copEV, copFE )
	model = (convert(Lar.Points,V'), FV)
	sp_idx = Lar.spaceindex(model)

	# initializations
    fs_num = size(copFE, 1)
    rV = Array{Float64,2}(undef,0,3)
    rEV = SparseArrays.spzeros(Int8,0,0)
    rFE = SparseArrays.spzeros(Int8,0,0)

	# multiprocessing of face fragmentation
    if (multiproc == true)
        in_chan = Distributed.RemoteChannel(()->Channel{Int64}(0))
        out_chan = Distributed.RemoteChannel(()->Channel{Tuple}(0))
        @async begin
            for sigma in 1:fs_num
                put!(in_chan, sigma)
            end
            for p in Distributed.workers()
                put!(in_chan, -1)
            end
        end
        for p in Distributed.workers()
            @async Base.remote_do(
                frag_face_channel, p, in_chan, out_chan, V, EV, FE, sp_idx)
        end
        for sigma in 1:fs_num
            rV, rEV, rFE = Lar.skel_merge(rV, rEV, rFE, take!(out_chan)...)
        end
    else
	# sequential (iterative) processing of face fragmentation
        for sigma in 1:fs_num
            #print(sigma, "/", fs_num, "\r")
            nV, nEV, nFE = Lar.Arrangement.frag_face(V, copEV, copFE, sp_idx, sigma)
			#nV, nEV, nFE = Lar.fragface(V, copEV, copFE, sp_idx, sigma)
			nV = convert(Lar.Points, nV)
            a,b,c = Lar.skel_merge( rV,rEV,rFE,  nV,nEV,nFE )
            rV=a;  rEV=b;  rFE=c
        end
    end
	# merging of close vertices, edges and faces (3D congruence)
	rV, rEV, rFE = merge_vertices(rV, rEV, rFE)
    return rV, rEV, rFE
end

function spatial_arrangement_2(
		rV::Lar.Points,
		rcopEV::Lar.ChainOp,
		rcopFE::Lar.ChainOp, multiproc::Bool=false)

	#rcopCF = Lar.build_copFC(rV, rcopEV, rcopFE)
	rcopCF = minimal_3cycles(rV, rcopEV, rcopFE)
    return rV, rcopEV, rcopFE, rcopCF
end

"""
    spatial_arrangement(V::Points, copEV::ChainOp, copFE::ChainOp; [multiproc::Bool])

Compute the arrangement on the given cellular complex 2-skeleton in 3D.

A cellular complex is arranged when the intersection of every possible pair of cell
of the complex is empty and the union of all the cells is the whole Euclidean space.
The function returns the full arranged complex as a list of vertices V and a chain of borders EV, FE, CF.

## Additional arguments:
- `multiproc::Bool`: Runs the computation in parallel mode. Defaults to `false`.
"""
function spatial_arrangement(
		V::Lar.Points, # by rows
		copEV::Lar.ChainOp,
		copFE::Lar.ChainOp, multiproc::Bool=false)

	# face subdivision
	rV, rcopEV, rcopFE = Lar.Arrangement.spatial_arrangement_1( V, copEV, copFE, multiproc ) # copFE global
	bicon_comps = Lar.Arrangement.biconnected_components(rcopEV)
	#W,bicon_comps = Lar.biconnectedComponent((W,EV))
	#@error "comps# = $(length(bicon_comps))"
	# 3-complex and containment graph
	rV, rEV, rFE, rCF = Lar.Arrangement.spatial_arrangement_2(rV, rcopEV, rcopFE)
end


@inline function policy(m::KrusellSmith,
                        net_return::Float64, wage::Float64,
                        l_in::Vector{Float64}, g::Vector{Float64}, KF::Array{Float64,2};
                        damp::Float64 = 0.8, tol::Float64 = 1e-5)
    @assert 0 <= damp <= 1

    nw::Int64                = get_setting(m, :nw)
    ns::Int64                = get_setting(m, :ns)
    wgrid::DSGE.Grid         = get_grid(m, :wgrid)
    sgrid::DSGE.Grid         = get_grid(m, :sgrid)

    zhi::Float64             = get_setting(m, :zhi)
    zlo::Float64             = get_setting(m, :zlo)
    smoother::Float64        = get_setting(m, :smoother)
    γ::Float64               = m[:γ]
    β::Float64               = m[:β]

    dist::Float64            = 1.0         # distance between l_in and l_out on decision rules
    l_out::Array{Float64,1}  = copy(l_in)  # initialize l_out

    ap = Array{Float64,1}(undef, nw) # savings
    q  = Array{Float64,2}(undef, nw, nw * ns) # auxilary variable to compute LHS of euler

    count = 1

    while dist > tol && count < 5000

        # compute c(w) given guess for l_in = β*R*E[u'(c_{t+1})]
        c::Array{Float64,1} = min.(l_in.^(-1 / γ), wgrid.points)

        # compute ap(w) given guess for l_in
        ap = wgrid.points - c

        # The term contained inside the (⋅) next to g in equation (9)
        q = euler_lhs(q, ap, wgrid.points, sgrid.points, zhi, zlo, smoother, wage, net_return, nw, ns)

        # Equation (9)
        l_out = β * net_return * q * kron(sgrid.weights .* (g / wage),
                                              wgrid.weights .* (c.^-γ))
        dist = maximum(abs.(l_out - l_in))
        l_in = damp * l_out + (1.0 - damp) * l_in;

        count += 1
    end
    if count == 5000
        @warn "Euler iteration did not converge"
    end

    # Equation (10)
    # You can re-use most of the functional form of the Euler equation to
    # back out the KF for μ(w)
    KF = kolmogorov_forward(q, KF, g, sgrid.weights, wage, nw, ns)

    return ap, l_out, KF
end

function steadystate!(m::KrusellSmith, tol::Float64 = 1e-5, maxit::Int64 = 100, Kdamp::Float64 = 0.1)
    # Load settings and grids
    nw::Int64  = get_setting(m, :nw)

    wgrid::DSGE.Grid = get_grid(m, :wgrid)
    sgrid::DSGE.Grid = get_grid(m, :sgrid)

    zhi::Float64      = get_setting(m, :zhi)
    zlo::Float64      = get_setting(m, :zlo)
    smoother::Float64 = get_setting(m, :smoother)

    α::Float64 = m[:α]
    δ::Float64 = m[:δ]
    β::Float64 = m[:β]

    g = trunc_lognpdf(sgrid, m[:μ_s].value, m[:σ_s].value)      # Density g (of skill, s) evaluated on sgrid

    m[:Lstar] = L::Float64 = quadrature_sum(g, sgrid)           # Aggregate Labor endowment

    K::Float64 = L*(α /((1.0 / β) - 1.0 + δ))^(1.0 / (1.0 - α)) # Aggregate Capital Stock

    l_in = fill(0.3, nw)        # This corresponds to l^i_t = βE_t[R_{t+1} u'(c^i_{t+1})]
                                #                           = βE_t[R_{t+1} min(l^i_{t+1}, (w^i_{t+1}^(-1/γ))]

    ap   = Array{Float64,1}(undef, nw)
    KF   = Array{Float64,2}(undef, nw, nw) # The Kolmogorov Forward Equation for μ
    μ    = Array{Float64,1}(undef, nw)     # The cross-sectional density of cash on hand, w
    newK = 0.0

    diseqm = 1000.
    count = 1

    # Equation (11)
    while abs(diseqm) > tol && count < maxit # clearing markets

        wage = MPL(1.0, L, K, α)
        net_return = MPK(1.0, L, K, α, δ)

        ap, l_in, KF = policy(m, net_return, wage, l_in, g, KF)

        # sspolicy returns the consumption decision rule, a prime
        # decision rule, l_in = updated β R u'(c_{t+1}),
        # KF is the Kolmogorov forward operator and is the map which moves you
        # from cash on hand distribution today to cash on hand dist tomorrow

        # find eigenvalue closest to 1
        D,V = eigen(wgrid.weights[1] * KF)

        if norm(D[1] - 1.0) > 2e-1 # that's the tolerance we are allowing
            @warn "your eigenvalue is too far from 1, something is wrong"
        end
        μ = real(@view(V[:,1]))  # Pick the eigenvector associated with the largest
                                 # eigenvalue and moving it back to values

        μ = μ / (wgrid.weights' * μ)        # Scale of eigenvectors not determinate:
                                            # rescale to integrate to exactly 1
        newK  = wgrid.weights' * (μ .* ap)  # compute excess supply of savings, which is a fn of w
        diseqm = newK - K

        K += Kdamp * diseqm
        count += 1
    end

    m[:lstar] = l_in
    m[:cstar] = min.(m[:lstar].value.^(-1.0/m[:γ]), wgrid.points)
    m[:μstar] = μ
    m[:Kstar] = newK
    m[:KFstar] = KF

    nothing
end

@inline function mollifier(z::Float64,zhi::Float64,
                   zlo::Float64, smoother::Float64; In::Float64=0.443993816237631)
    # mollifier function
    if z < zhi && z > zlo
        temp = (-1.0 + 2.0 * (z - zlo) / (zhi - zlo)) / smoother
        out = (2.0/(zhi - zlo)) * exp(-1.0 / (1.0 - temp*temp)) / (In * smoother);
    else
        out = 0.0
    end
    return out
end

# Derivative of the mollifier
@inline function dmollifier(x::Float64, zhi::Float64, zlo::Float64, smoother::Float64;
                            In::Float64=0.443993816237631)
    temp = (-1.0 + 2.0*(x-zlo)/(zhi-zlo))/smoother
    dy  = -(2*temp./((1 - temp.^2).^2)).*(2/(smoother*(zhi-zlo))).*mollifier(x, zhi, zlo, smoother)
end

@inline function trunc_lognpdf(grid::Grid, mu::Float64, sig::Float64)
    x = grid.points
    xlo = x[1]
    xhi = x[end]

    # truncated log normal pdf
    logn = LogNormal(mu, sig)
    return (x .- xlo .>= -eps()) .* (x .- xhi .<= eps()) .* pdf.(logn,x) / (cdf(logn,xhi) .- cdf(logn,xlo))
end

# ω_t
@inline function MPL(Z::Float64, L::Float64, K::Float64, α::Float64)
    (1.0 - α) * Z * (K^α) * (L^(-α))  # Marginal product of labor
end

# r_t + (1 - δ)
@inline function MPK(Z::Float64, L::Float64, K::Float64, α::Float64, δ::Float64)
    α * Z * (K^(α - 1.0)) * (L^(1.0 - α)) + (1.0 - δ)   # Marginal product of capital + 1-δ
end

####### Is this function ever used? #########
#=
# Re-computing
@inline function kolmogorov_forward(g::Vector{Float64}, sweights::Vector{Float64},
                            ap::Vector{Float64}, wgrid::Vector{Float64}, sgrid::Vector{Float64},
                            zhi::Float64, zlo::Float64, smoother::Float64,
                            wage::Float64, R::Float64, nw::Int64, ns::Int64)

    KF = Array{Float64, 2}(nw, nw)   # transition matrix mapping todays dist of w to w'
    for iw=1:nw
        for iwp=1:nw
            sumns = 0.0
            cursum = (wgrid[iwp] - R*ap[iw]) / wage
            for isp=1:ns
                sumns += mollifier(cursum - sgrid[isp], zhi, zlo, smoother) * (g[isp]) * sweights[isp]
            end
            @inbounds KF[iwp,iw] = sumns / wage
        end
    end
    return KF
end
=#

# Indexing into q since its already been computed
@inline function kolmogorov_forward(q::Matrix{Float64}, KF::Array{Float64,2}, g::Vector{Float64},
                                    sweights::Vector{Float64}, wage::Float64, nw::Int64, ns::Int64)
    g_sweights_prod::Array{Float64,1} = g / wage .* sweights
    nw_array = nw * (0:ns-1)
    for iw = 1:nw
        for iwp = 1:nw
            @inbounds KF[iwp, iw] = sum(q[iw, nw_array .+ iwp] .* g_sweights_prod)
        end
    end
    return KF
end

@inline function euler_lhs(q::Matrix{Float64}, ap::Vector{Float64}, wgrid::Vector{Float64},
                   sgrid::Vector{Float64}, zhi::Float64, zlo::Float64, smoother::Float64,
                   wage::Float64, net_return::Float64, nw::Int64, ns::Int64)
    for iwp = 1:nw
        for iw = 1:nw
            wgrid_term = (wgrid[iwp] - net_return * ap[iw]) / wage
            for iss = 1:ns
                @inbounds q[iw, nw * (iss - 1) + iwp] = mollifier(wgrid_term - sgrid[iss], zhi, zlo, smoother)
            end
        end
    end
    return q
end


include("auto_encoder_task.jl")

# ::: NETWORK
networks = Dict(
    "500_300_150" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
        ],[
        DenseLayer(LATENT_SIZE, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
   "600_400_200_100" => ([
        DenseLayer(INPUT_SIZE, 600),
        ActivationLayer(600, relu, Δrelu),
        DenseLayer(600, 400),
        ActivationLayer(400, relu, Δrelu),
        DenseLayer(400, 200),
        ActivationLayer(200, relu, Δrelu),
        DenseLayer(200, 100),
        ActivationLayer(100, relu, Δrelu),
        DenseLayer(100, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
        ],[
        DenseLayer(LATENT_SIZE, 100),
        ActivationLayer(100, relu, Δrelu),
        DenseLayer(100, 200),
        ActivationLayer(200, relu, Δrelu),
        DenseLayer(200, 400),
        ActivationLayer(400, relu, Δrelu),
        DenseLayer(400, 600),
        ActivationLayer(600, relu, Δrelu),
        DenseLayer(600, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
    "500_500_250_250_125_125" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 250),
        ActivationLayer(250, relu, Δrelu),
        DenseLayer(250, 250),
        ActivationLayer(250, relu, Δrelu),
        DenseLayer(250, 125),
        ActivationLayer(125, relu, Δrelu),
        DenseLayer(125, 125),
        ActivationLayer(125, relu, Δrelu),
        DenseLayer(125, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
       ],[
        DenseLayer(LATENT_SIZE, 125),
        ActivationLayer(125, relu, Δrelu),
        DenseLayer(125, 125),
        ActivationLayer(125, relu, Δrelu),
        DenseLayer(125, 250),
        ActivationLayer(250, relu, Δrelu),
        DenseLayer(250, 250),
        ActivationLayer(250, relu, Δrelu),
        DenseLayer(250, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
    "500" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
       ],[
        DenseLayer(LATENT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]))

# ::: TRAINING
runtime_plotting = false
plot_intervall = 20
begin
    tm1 = [NaN for _ in 1:TRAIN_EPISODES];
    tm2 = [NaN for _ in 1:TRAIN_EPISODES];
    tm3 = [NaN for _ in 1:TRAIN_EPISODES];
    tm4 = [NaN for _ in 1:TRAIN_EPISODES];

    println("---"); @time for e in 1:TRAIN_EPISODES
        t = time()

        z1 = forward!(networks["500_300_150"][1], train_data_x[e])
        z2 = forward!(networks["600_400_200_100"][1], train_data_x[e])
        z3 = forward!(networks["500_500_250_250_125_125"][1], train_data_x[e])
        z4 = forward!(networks["500"][1], train_data_x[e])

        y1 = forward!(networks["500_300_150"][2], z1)
        y2 = forward!(networks["600_400_200_100"][2], z2)
        y3 = forward!(networks["500_500_250_250_125_125"][2], z3)
        y4 = forward!(networks["500"][2], z4)

        dloss1 = -Δmse(train_data_x[e], y1)
        dloss2 = -Δmse(train_data_x[e], y2)
        dloss3 = -Δmse(train_data_x[e], y3)
        dloss4 = -Δmse(train_data_x[e], y4)

        dz1 = backward!(reverse(networks["500_300_150"][2]), dloss1, opt)
        _ = backward!(reverse(networks["500_300_150"][1]), dz1, opt)
        dz2 = backward!(reverse(networks["600_400_200_100"][2]), dloss2, opt)
        _ = backward!(reverse(networks["600_400_200_100"][1]), dz2, opt)
        dz3 = backward!(reverse(networks["500_500_250_250_125_125"][2]), dloss3, opt)
        _ = backward!(reverse(networks["500_500_250_250_125_125"][1]), dz3, opt)
        dz4 = backward!(reverse(networks["500"][2]), dloss4, opt)
        _ = backward!(reverse(networks["500"][1]), dz4, opt)

        tm1[e] = mse(train_data_x[e], y1)
        tm2[e] = mse(train_data_x[e], y2)
        tm3[e] = mse(train_data_x[e], y3)
        tm4[e] = mse(train_data_x[e], y4)

        # if runtime_plotting && e % plot_intervall == 0
        #     p = plot(FC_train_metric)
        #     display(p)
        # end

        # println("episode = $e : loss = $loss : dt = $(time() - t)")
    end
end

begin
    gr = 30
    plot([sum(tm1[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], ylabel="MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)", xlabel="training iterations", label="FC [500_300_150]")
    plot!([sum(tm2[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC [600_400_200_100]")
    plot!([sum(tm3[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC [500_500_250_250_125_125]")
    plot!([sum(tm4[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC [500]")
end

generate_images(10, "/Users/BBM2H16AHO/Documents/programming/DNN-Framework/auto_encoder/images/relu_sigmoid_test")


using Plots
using Random
using DifferentialEquations
using LinearAlgebra
using DynamicalSystems
using Parameters
using BasicReservoirComputing

include("utilities.jl")

function test_rc(system::DynamicalSystem, params::NamedTuple; Δt = 0.01)
    D=length(system.u0)
    utrain, utest_arr = make_data(system; train_time = 100, test_time = 25, n_test = 100)

    @unpack N, SR, ρA, α, σ, β, σb, nspin = params

    # Build and Train RC
    reservoir = rc(D, N, ρA, Δt, SR, α, σ, σb, β)
    train_RC(reservoir, utrain, spinup=nspin)

    # Forecast RC
    utrue = utest_arr[1][:, nspin:end]
    nsteps = size(utrue)[2]
    upred = forecast_RC(reservoir, nsteps, uspin = utest_arr[1][:, 1:nspin])
    t = collect(0:nsteps-1)*Δt

    p1 = plot_prediction(t, upred, utrue)

    # Test on many points
    vpt = test_RC(reservoir, utest_arr; ϵ=0.3, spinup=nspin)
    p2 = histogram(vpt, legend=false,
                   overwrite_figure=false)
    xlabel!(p2, "Valid Prediction Time")
    pfinal = plot(p1, p2, layout = (1, 2), show=true)
    display(pfinal)
    return reservoir
end

params = (N=200, σ=0.084, α=0.6, SR=0.8, ρA=0.02, β=8.5e-8, σb=1.6, nspin=200)
system = Systems.lorenz()
reservoir = test_rc(system, params);


type Setting
    img
    row::Int
    col::Int
end

function mod_fun(row_map, arr, col, scores)
    j = rand(1:size(arr)[1])        # We randomly choose a column.
    while arr[j, 1] == -1 || j == col
        j = rand(1:size(arr)[1])
    end
    elts = row_map[j]               # We get the elements on j-th column.
    nb_elt_on_col = size(elts)[1]   # We get the number of elements.
    start_index = arr[j, 3]         # We get the nearest index from the start point.
    # Randomly selecting a white pixel according to the normal distribution.
    rnd = randn()
    index = int(rnd * (nb_elt_on_col / 2) + start_index)
    index = max(index, 1)
    index = min(index, nb_elt_on_col)

    i = elts[index]

    return (i,j,index)
end

function cost_fun(orig, new, arr, pos, start)
    # Compute the local cost of the modification
    # Orig : the original image
    # new : the modified image
    # v : the candidate value for position `pos`
    # pos : position of the modification
    # start : hint about the start of the line given to the program

    rows, cols = size(new)
    row, col = start
    r, c = pos

    alpha = 1
    beta = 20

    cost = 0
    # We score the distance between the start point and the pixel.
    if c > 1 && arr[c - 1, 1] == -1
        cost += alpha * abs(r - row)
    end
    # We score the distance between the previous marked pixel and the current
    # one, weighted with the square of root of the distance with the start
    # pixel. We want the relation between neighbors to be important, but we
    # also want to be close to the line to match.
    if c > 1 && arr[c - 1, 1] != -1
        cost += beta * abs(r - arr[c - 1, 1]) * (abs(r - row) + 1) ^ (1/2)
    end

    cost
end

function extract(file_path, row, col, out_path)

    println("extracting line from: $(file_path)")
    println("starting point: (row = $(row), col = $(col))")

    # Read input image to sparse matrix
    img = convert(Array, float64(imread(file_path)))
    img[row, col] = 255                       # In order to be sure there is at least one pixel on this column.
    orig = sparse(img)

    # Get number of rows and columns of the image.
    rows, cols = size(orig)

    # Candidate
    I, J, V = findnz(orig)
    arr = zeros(cols, 3)                      # arr contains the tuples (row, index, start_index) marked, one by column.
    row_map = map((x) -> I[J .== x], 1:cols)  # row_map contains the indexes of the white pixels.

    # Initial random candidate
    for j = 1:cols
        elts = row_map[j]                     # We get elements on j-th column.
        nb_elt_on_col = size(elts)[1]         # We get the number of elements.

        start_index = 0
        # If there isn't any white pixel on the column, we don't mark it.
        if (nb_elt_on_col == 0)
            i = -1
            index = -1
        else
            # We initialize with the nearest (row speaking) pixel from the
            # start pixel.
            start_index = max(size(row_map[j][row_map[j] .<= row], 1), 1)
            # Then we mark a random one, according to the normal distribution.
            rnd = randn()
            index = int(rnd * (nb_elt_on_col / 2) + start_index)
            index = max(index, 1)
            index = min(index, nb_elt_on_col)

            i = elts[index]
        end
        arr[j,:] = [i, index, start_index]    # We mark this pixel.
    end
    # Set start point.
    arr[col, 1] = row

    obj = sparse(I, J, V)

    ###########################
    # Run simulated annealing #
    ###########################

    # Simulated annealing settings
    t = 2500.0
    t_step = 0.999999
    t_stop = 5
    max_epoc = 10000000000
    epoc = 0

    println("Compute score matrix")

    score = zeros(Float64, cols)
    for i = 1:cols
        score[i] = cost_fun(orig, obj, arr, (arr[i], i), (row, col))
    end

    tot_score = sum(score)
    println("Total score: $(tot_score)")

    println("Run Simulated Annealing")

    while epoc < max_epoc && t > t_stop

        # Choose neighbor
        r,c,index = mod_fun(row_map, arr, col, score)

        # Compute local score
        score_mod = cost_fun(orig, obj, arr, (r, c), (row, col))

        # Compute energy delta
        mod_e = score_mod - score[c]

        # Keep the candidate or not
        if mod_e < 0 || rand() <= exp(-mod_e / t)
            arr[c,:] = [r, index, arr[c,3]]
            tot_score = tot_score - score[c] + score_mod
            score[c] = score_mod
        end

        # Decrease temperature
        if epoc % 100 == 0
            t *= t_step
        end
        if epoc % 1000000 == 0
            println("Iteration number: $(epoc); cost: $(tot_score), temp: $(t)")
        end

        epoc += 1
    end

    println("temperature: $t")
    println("energy: $tot_score")
    println("epocs: $epoc")

    ###################
    # Display results #
    ###################

    img *= 0.0005
    for i = 1:cols
        if (arr[i] != -1)
            r = int(arr[i])
            img[r, i] = 255
        end
    end

    # Display result
    img2 = convert(Image, img)
    display(img2)
    #imwrite(img2, out_path)
end


#######################################################################################################################################################################################################
#
# Changes made to this function
# General
#     2021-Oct-22: add function to compute leaf level PAR and APAR
#     2022-Jan-13: use LeafBiophysics directly in the function rather than Leaf
#     2022-Feb-02: fix documentation
#     2022-Feb-02: unpack CONSTANTS only
#
#######################################################################################################################################################################################################
"""

    leaf_PAR(bio::LeafBiophysics{FT}, wls::WaveLengthSet{FT}, rad::HyperspectralRadiation{FT}; APAR_car::Bool = true) where {FT<:AbstractFloat}

Return leaf level PAR and APAR, given
- `bio` `ClimaCache.LeafBiophysics` type struct that contains leaf biophysical parameters
- `wls` `ClimaCache.WaveLengthSet` type struct that contains wave length bins
- `rad` `ClimaCache.HyperspectralRadiation` type struct that contains incoming radiation information
- `APAR_car` If true (default), account carotenoid absorption as APAR; otherwise, APAR is only by chlorophyll

---
# Examples
```julia
wls = WaveLengthSet{Float64}();
bio = LeafBiophysics{Float64}(wls);
rad = HyperspectralRadiation{Float64}();
par,apar = leaf_PAR(bio, wls, rad);
par,apar = leaf_PAR(bio, wls, rad; APAR_car=false);
```
"""
function leaf_PAR(bio::LeafBiophysics{FT}, wls::WaveLengthSet{FT}, rad::HyperspectralRadiation{FT}; APAR_car::Bool = true) where {FT<:AbstractFloat}
    @unpack IΛ_PAR, ΔΛ_PAR, Λ_PAR = wls;

    # APAR absorption feature
    _α = (APAR_car ? view(bio.α_cabcar, IΛ_PAR) : view(bio.α_cab, IΛ_PAR));

    # PAR energy from direct  and diffuse light
    _e_par_dir  = view(rad.e_direct , IΛ_PAR) .* view(bio.α_SW, IΛ_PAR);
    _e_par_diff = view(rad.e_diffuse, IΛ_PAR) .* view(bio.α_SW, IΛ_PAR);
    _par_dir    = photon.(Λ_PAR, _e_par_dir);
    _par_diff   = photon.(Λ_PAR, _e_par_diff);

    # absorbed PAR energy from direct and diffuse light
    _apar_dir  = _α .* _par_dir;
    _apar_diff = _α .* _par_diff;

    # total PAR and APAR in μmol photons m⁻² s⁻¹
    _∑par_dir   = numerical∫(_par_dir, ΔΛ_PAR);
    _∑par_diff  = numerical∫(_par_diff, ΔΛ_PAR);
    _∑apar_dir  = numerical∫(_apar_dir, ΔΛ_PAR);
    _∑apar_diff = numerical∫(_apar_diff, ΔΛ_PAR);

    return 1000 * (_∑par_dir + _∑par_diff), 1000 * (_∑apar_dir + _∑apar_diff)
end


using Distributed

@everywhere include("parallel-binomial.jl")
using PyPlot

n = 100
p = 0.5
reps = 1000
rounds = 1000

println("Starting computation")

@time c = parallel_binomial_counts(n, p, reps, rounds)

a = minimum(keys(c))
b = maximum(keys(c))
s = sum(c)

x = collect(a:b)
y = [c[t] / s for t in x]

bar(x, y)
grid()
title("Binomial distribution B($n, $p)")
savefig("binomial.pdf")


"""
    BlackKarasinski()

[Black-Karasinski model](https://en.wikipedia.org/wiki/Black%E2%80%93Karasinski_model).
"""
struct BlackKarasinski <: Model end

"""
    evaluate(IRD, BlackKarasinski(), n = 12)

Evaluate interest rate derivative using `BlackKarasinski` model.

# Arguments
- `IRD::InterestRateDerivative`: interest rate derivative
- `n`: number of paths to simulate
"""
function evaluate(IRD::InterestRateDerivative, m::BlackKarasinski, n::Int64 = 12)
    Δt = IRD.t / n
    rates = [IRD.r]
    for i = 0:n
        dr = IRD.k * (IRD.θ - log(rates[end])) * Δt + IRD.σ * rates[end] * randn()
        append!(rates, rates[end] + dr)
    end
    return rates
end


"""
    Res0dWithBN <: Block
Res0dWithBN is a residual whose neuron processing 0-d data. which has 2 BatchNorm0d blocks

# Structure
       Rⁱ → Rᵐ                  Rᵐ → Rⁱ
        ┌───┐   ┌────┐   ┌───┐   ┌───┐   ┌────┐   ┌───┐   ┌───┐
    ─┬─►│ W ├──►│ BN ├──►│f()├──►│ W ├──►│ BN ├──►│ + ├──►│f()├──►
     │  └───┘   └────┘   └───┘   └───┘   └────┘   └─┬─┘   └───┘
     └─────────────────────────►────────────────────┘
"""
mutable struct Res0dWithBN <: Block
    blocks::Vector
    function Res0dWithBN(i::Int, m::Int; type::Type=Array{Float32})
        a1 = Affine(i, m, type=type)
        b1 = BatchNorm0d(m)
        a2 = Affine(m, i, type=type)
        b2 = BatchNorm0d(i)
        new([a1, b1, a2, b2])
    end
    function Res0dWithBN()
        new(Vector(undef,4))
    end
end

@extend(Res0dWithBN, blocks)

function clone(this::Res0dWithBN; type::Type=Array{Float32})
    cloned = Res0dWithBN()
    cloned[1] = clone(this[1], type=type)
    cloned[2] = clone(this[2], type=type)
    cloned[3] = clone(this[3], type=type)
    cloned[4] = clone(this[4], type=type)
    return cloned
end

function forward(m::Res0dWithBN, x0)
    x1 = forward(m[1], x0) # Affine
    x2 = forward(m[2], x1) # batchnorm0d
    x3 =          relu(x2) # activation
    x4 = forward(m[3], x3) # Affine
    x5 = forward(m[4], x4) # batchnorm0d
    return relu(x5 + x0)   # short connectin + activation
end

function predict(m::Res0dWithBN, x0)
    x1 = predict(m[1], x0) # Affine
    x2 = predict(m[2], x1) # batchnorm0d
    x3 =          relu(x2) # activation
    x4 = predict(m[3], x3) # Affine
    x5 = predict(m[4], x4) # batchnorm0d
    return relu(x5 + x0)   # short connectin + activation
end


using DelimitedFiles

Random.seed!(1234)

δ = 1.0
φ = 0.95
ση = 1.0
σε = 1.0

θ0 = [δ;atanh(φ);ση;σε]

β = [0.0;1.0]

N = 1000

@testset "Kalman filter regression tests" begin
    arh = ARH(β)

    X,Y = Filtering.simulate(arh,θ0,N,fill([1.0],N))

    kf = kalman_filter(arh,θ0,Y,fill([1.0],length(Y)))

    @testset "Filter" begin

        xf0 = readdlm("regressions/ar/kf_xf.csv")[:,1]
        Pf0 = readdlm("regressions/ar/kf_Pf.csv")[:,1]
        xp0 = readdlm("regressions/ar/kf_xp.csv")[:,1]
        Pp0 = readdlm("regressions/ar/kf_Pp.csv")[:,1]
        Σ0  = readdlm("regressions/ar/kf_Σ.csv")[:,1]
        ε0  = readdlm("regressions/ar/kf_ε.csv")[:,1]

        @test all(vcat(kf.xf...) .== xf0)
        @test all(vcat(kf.Pf...) .== Pf0)
        @test all(vcat(kf.xp...) .== xp0)
        @test all(vcat(kf.Pp...) .== Pp0)
        @test all(vcat(kf.Σ...)  .== Σ0)
        @test all(vcat(kf.ε...)  .== ε0)
    end

    ks = kalman_smoother(arh,θ0,kf)

    @testset "Smoother" begin

        x0  = readdlm("regressions/ar/ks_x.csv")[:,1]
        P0  = readdlm("regressions/ar/ks_P.csv")[:,1]

        @test all(vcat(ks.x...)  .== x0)
        @test all(vcat(ks.P...)  .== P0)
    end

end


using Test
using Quon: QuonConst, Zero, Pi, HalfPi, NegHalfPi, InfZero, InfPi
using Quon
using Yao

to_Rz(p::Complex) = Rz(p/im)
function to_Rz(p::QuonConst)
    p === Pi && return Rz(π)
    p === Zero && return Rz(0.0)
    p === HalfPi && return Rz(π/2)
    p === NegHalfPi && return Rz(-π/2)
    p === InfZero && return matblock(projector(0))
    p === InfPi && return matblock(projector(1))
end
to_Rx(p::Complex) = Rx(p/im)
function to_Rx(p::QuonConst)
    p === Pi && return Rx(π)
    p === Zero && return Rx(0.0)
    p === HalfPi && return Rx(π/2)
    p === NegHalfPi && return Rx(-π/2)
    p === InfZero && return matblock((1/sqrt(2) + 0.0im)*[1 1; 1 1])
    p === InfPi && return matblock((1/sqrt(2) + 0.0im)*[1 -1; -1 1])
end

function test_quon_const(a, b, c)
    circ1 = chain(1, to_Rz(a), to_Rx(b), to_Rz(c))
    p, q, r = yang_baxter_param(a, b, c)
    circ2 = chain(1, to_Rx(p), to_Rz(q), to_Rx(r))
    return operator_fidelity(matblock(mat(circ1)), matblock(mat(circ2)))/
        sqrt(operator_fidelity(matblock(mat(circ1)), matblock(mat(circ1))) * operator_fidelity(matblock(mat(circ2)), matblock(mat(circ2))))
end

AllConsts = (Zero, Pi, HalfPi, NegHalfPi, InfZero, InfPi)
FiniteConsts = (Zero, Pi, HalfPi, NegHalfPi)

test_quon_const(Pi, Pi, Pi)

test_123 = [test_quon_const(p1, p2, p3) for p1 in AllConsts, p2 in AllConsts, p3 in AllConsts]
test_12 = [test_quon_const(p1, p2, Complex(p3)) for p1 in AllConsts, p2 in AllConsts, p3 in FiniteConsts]
test_23 = [test_quon_const(Complex(p1), p2, p3) for p1 in FiniteConsts, p2 in AllConsts, p3 in AllConsts]
test_13 = [test_quon_const(p1, Complex(p2), p3) for p1 in AllConsts, p2 in FiniteConsts, p3 in AllConsts]
test_1 = [test_quon_const(p1, Complex(p2), Complex(p3)) for p1 in AllConsts, p2 in FiniteConsts, p3 in FiniteConsts]
test_2 = [test_quon_const(Complex(p1), p2, Complex(p3)) for p1 in FiniteConsts, p2 in AllConsts, p3 in FiniteConsts]
test_3 = [test_quon_const(Complex(p1), Complex(p2), p3) for p1 in FiniteConsts, p2 in FiniteConsts, p3 in AllConsts]

@test sum(test_123 .≈ 1) == length(test_123) - 4
@test sum(test_12 .≈ 1) == length(test_12)
@test sum(test_13 .≈ 1) == length(test_13) - 4
@test sum(test_23 .≈ 1) == length(test_23)
@test sum(test_1 .≈ 1) == length(test_1)
@test sum(test_2 .≈ 1) == length(test_2)
@test sum(test_3 .≈ 1) == length(test_3)


"""
    dpclustering(y::Array{Real, 1}, N::Array{Real, 1}; <keyword arguments>)
Perform dirichlet clustering on the variant allele frequency distribution of cancer sequencing data and find the number of clusters that the data supports, y is a vector of the number of reads reporting each mutant, N is the total depth at each locus.

...
## Arguments
- `iterations = 1000`: number of iterations of the gibbs samples
- `C = 30`: Max number of clusters to consider
- `burninstart = round(Int64, iterations/2)`: Burn in of the gibbs samples
- `bw = 0.01`: Bandwidth of density estimation
- `maxxaxis = 0.7`:
- `cutoffweight = 0.05`: Minimum weight to be called a cluster
- `verbose = true`: Show progress of gibbs sampling with `ProgressMeter` package
- `A = 0.01`: Hyperparameter for α, see Nik-Zainal et al
- `B = 0.01`: Hyperparameter for α, see Nik-Zainal et al
...
"""
function dpclustering(y, N;
    iterations = 20000,
    C = 30, #max number of clusters
    burninstart = round(Int64, 5000),
    bw = 0.01, # bandwidth of density estimation
    maxxaxis = 0.7,
    cutoffweight = 0.05, #minimum weight to be called a cluster
    verbose = true,
    A = 0.01, # Hyperparameters for alpha
    B = 0.01 # Hyperparameters for alpha
    )

    sum(y .== 0) == 0 || error("Some mutations have VAF = 0.0, make sure these mutations are removed before clustering")

    nummuts = length(y)

    # Set up arrays and matrices for recording samples
    y = map(Float64, y) # make sure read counts and depth are floats
    N = map(Float64, N)
    VAF = y ./ N
    π = zeros(Float64, iterations, C)
    PrS = zeros(Float64, nummuts, C)
    α = zeros(Float64, iterations)
    mutburdens = zeros(Float64, iterations, C, nummuts)
    V = ones(Float64, iterations, C)
    clusterassignment = zeros(Int64, iterations, nummuts)

    #random initial clusters between 0 and 1
    π[1, :] = rand(Uniform(0.0, maximum(VAF)), C)
    for c in 1:C
        mutburdens[1, c, :] .= π[1, c]
    end

    α[1] = 1.0
    V[1, 1:(C - 1)] .= 0.5

    if verbose == true
      p = Progress(iterations, 1, "Gibbs sampling progress: ", 30)
    end

    for m in 2:iterations
        for k in 1:nummuts
            #Binomial log-likelihood
            PrS[k, 1] = log.(V[m .- 1, 1]) .+ (y[k] .* log.(mutburdens[m-1, 1, k])) .+
            (N[k] .- y[k]) .* log.(1 .- mutburdens[m - 1, 1, k])
            allocate!(PrS, V, mutburdens, y, N, k, C, m)
            takemax!(PrS, k, C)
            exp!(PrS, k, C)
            normalize!(PrS, k, C)
        end

        multinomsample!(clusterassignment, PrS, nummuts, m, C)

        # Update stick-breaking weights
        updatestick!(V, clusterassignment, α, C, m)

        V[m, [V[m, 1:(C-1)] .== 1.0; false]] .= 0.9999

        countsPerCopyNum = N

        mutburdens[m, :, :] = mutburdens[m - 1, :, :]
        @fastmath @inbounds for c in unique(clusterassignment[m, :])
          idx = findall((in)(c), clusterassignment[m, :])
          αp = sum(y[idx])
          βp = 1 ./ sum(countsPerCopyNum[idx])
          π[m, c] = minimum([rand(Gamma(αp, βp)), 0.999])
          mutburdens[m, c, :] .= π[m, c]
        end
        x1 = C .+ A .- 1

        x2 = sum(log.(1 .- V[m, 1:(C.-1)]))
        x3 = B .- x2

        α[m] = rand(Gamma(C .+ A .- 1, 1/(B .- sum(log.(1 .- V[m, 1:(C.-1)])))))
        #α[m] = rand(Gamma(x1, x3))

        if verbose == true
          next!(p)
        end

    end

    dp = DPout(clusterassignment, V, π, α)

    DF, wts =
    getdensity(dp, iterations; burninstart = burninstart, bw = bw, maxxaxis = maxxaxis)
    wtsout, clonefreq, allwts, allfreq =
    summariseoutput(dp, wts, iterations; burninstart = burninstart, cutoffweight = cutoffweight)

    sortind = sortperm(clonefreq)
    return DPresults(DF, wts, length(wtsout), wtsout[sortind],
    clonefreq[sortind], allwts, allfreq, dp, TargetData(y, N, VAF));
end

function multinomsample!(clusterassignment, PrS, nummuts, m, C)
  @fastmath @inbounds for k in 1:nummuts
    clusterassignment[m, k] = sum(rand(Multinomial(1, PrS[k, :])) .* collect(1:C))
  end
end

function takemax!(PrS, k, C)
  maxPrS = maximum(PrS[k, :])
  @fastmath @inbounds for i in 1:C
    PrS[k, i] = PrS[k, i] - maxPrS
  end
end

function exp!(PrS, k, C)
  @fastmath @inbounds for i in 1:C
    PrS[k, i] = exp(PrS[k, i])
  end
end

function normalize!(PrS, k, C)
  sumPrS = sum(PrS[k, :])
  @fastmath @inbounds for i in 1:C
    PrS[k, i] = PrS[k, i] / sumPrS
  end
end

function updatestick!(V, clusterassignment, α, C, m)
  @fastmath @inbounds for h in 1:(C-1)
    V[m, h] = rand(Beta(1+sum(clusterassignment[m, :] .== h), α[m - 1] + sum(clusterassignment[m, :] .> h)))
  end
end

function allocate!(PrS, V, mutburdens, obsy, obsN, k, C, m)
    @fastmath @inbounds for j in 2:C
        PrS[k, j] = log(V[m-1, j]) +
        sum(log.(1 .- view(V, m-1, 1:(j-1)))) +
        obsy[k] * log(mutburdens[m-1, j, k]) +
        (obsN[k] - obsy[k]) * log(1-mutburdens[m-1, j, k])
    end
end

function getdensity(dp, iterations; burninstart = 500, bw = 1.0, maxxaxis = 0.5)

    wts = zeros( size(dp.V)[1], size(dp.V)[2]);
    wts[:, 1] = dp.V[:, 1]
    wts[:, 2] = dp.V[:, 2] .* (1 .- dp.V[:, 1])

    for i in 3:size(wts)[2]
        wts[:, i] = dp.V[:, i] .* prod((1 .- dp.V[:, (1:i .- 1)]), dims = 2)
    end

    postints = zeros(512, iterations -  burninstart + 1)

    xx = kde(dp.π[burninstart - 1, :],
    weights = wts[burninstart, :]./(sum(wts[burninstart, :])),
                npoints = 512, boundary = (0, maxxaxis), bandwidth = bw).x

    for i in burninstart:iterations
        postints[:, i - burninstart + 1] = kde(dp.π[i - 1, :], weights = wts[i, :]./(sum(wts[i, :])),
        npoints = 512, boundary = (0, maxxaxis), bandwidth = bw).density
    end

    meanv = mean(postints, dims = 2)[:]
    lq = mapslices(x -> quantile(x, 0.025), postints, dims = 2)[:]
    uq = mapslices(x -> quantile(x, 0.975), postints, dims = 2)[:]

    DF = DataFrame(mean = meanv, lq = lq, uq = uq, x = collect(xx))

    return DF, wts
end

function summariseoutput(dp, wts, iterations; burninstart = 1000, cutoffweight = 0.05)

    postwts = wts[burninstart:iterations, :]
    meanwts = mean(postwts, dims = 1)
    clonewts = meanwts[meanwts.>cutoffweight]

    clonefrequency = mean(dp.π[burninstart:iterations, :], dims = 1)
    largeclonefrequency = clonefrequency[meanwts.>cutoffweight]

    clonefrequency = clonefrequency[:]
    meanwts = meanwts[:]

    sortind = sortperm(clonefrequency)

    return clonewts, largeclonefrequency, meanwts[sortind], clonefrequency[sortind]
end


using Revise, FourierGPE
using Test

## Initialize simulation
L = (12.0,12.0)
N = (256,256)
sim = Sim(L,N)
@unpack_Sim sim

# set simulation parameters
μ = 12.0

# Time dependent potential function (here trivial t dep)
import FourierGPE.V
V(x,y,t) = 0.5*(x^2 + y^2)
ψ0(x,y,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y,0.0)/μ,0.0)+im*0.0)

# make initial state
x,y = X
ψi = ψ0.(x,y',μ,g)
ϕi = kspace(ψi,sim)
@pack_Sim! sim
##

## evolve
@time sol = runsim(sim)

# ground state

ϕg = sol[end]
ψg = xspace(ϕg,sim)
showpsi(x,y,ψg)
##

##
using VortexDistributions
R(w) = sqrt(2*μ/w^2)
R(1)
rv = 2.
healinglength(x,y,μ,g) = 1/sqrt(g*abs2(ψ0(x,y,μ,g)))
ξ0 = healinglength.(0.,0.,μ,g)
ξ = healinglength(rv,0.,μ,g)

pv = PointVortex(rv,0.,1)
vi = ScalarVortex(ξ,pv)

psi = Torus(copy(ψg),x,y)
vortex!(psi,vi)
showpsi(x,y,psi.ψ)

ψi .= psi.ψ
ϕi = kspace(ψi,sim) |> fftshift
kx,ky = K .|> fftshift
x,y = X

## Transform onto polar coordinates in k space
# grids

# make transform to polar grid

function cart2pol(ψi,x,y,kx,ky)
    kmax = maximum(kx)
    nx = length(x)
    nk = nx
    Nk = nk/2 |> Int
    Nθ = 2*nk
    k = LinRange(0,kmax,Nk)
    θ = LinRange(0,2*pi,Nθ+1)[1:Nθ]
    xi = reshape(x,1,1,nx)
    ki = reshape(k,1,1,Nk)
    θi = reshape(θ,1,1,Nθ)
    dθ = diff(θ)[1]
    dk = diff(k)[1]
    dx = diff(x)[1]
    dy = diff(y)[1]

    T = zeros(eltype(ψi),Nk,Nθ,nx)
    S = zeros(eltype(ψi),nx,Nk,Nθ)
    @. T = exp(-im*xi*k*cos(θ'))*dx/sqrt(2*pi)
    @. S = exp(-im*y*k'sin(θi))*dy/sqrt(2*pi)

    ϕp = zeros(eltype(ψi),Nk,Nθ)
    for i in eachindex(x), j in eachindex(y)
        t = @view T[:,:,i]
        s = @view S[j,:,:]
        @. ϕp += t*ψi[i,j]*s
    end
    return ϕp
end

##
@time ϕp = cart2pol(ψi,x,y,kx,ky)
sum(abs2.(ϕp).*k)*dθ*dk
sum(abs2.(ψi))*dx*dy

##
heatmap(log.(eps.() .+abs2.(ϕp)))
Ek = sum(abs2.(ϕp).*k.^3*dθ,dims = 2)
plot(k,log.(Ek))
plot!(k,-k .+5)
vline!([2*pi/ξ0])


@userplot vfpoil

@recipe function f(h::vfpoil;ϵs=0.0006, Tss=80, ∇Ts=1, GORs=false, dis=2.99, Qns=10, Dns=50, tols=0.05, sgs=0.6)

        Depth,Thp,Bsw,OilPVT,GasPVT,WaterPVT, Qmax=h.args

DepthRange, Qrange, Pwf, l, P=OilOutflowSen(Depth,Thp,Bsw,OilPVT,GasPVT,WaterPVT, Qmax;ϵ=ϵs, Ts=Tss, ∇T=∇Ts, GOR=GORs, di=dis, Qn=Qns, Dn=Dns, tol=tols, sg=sgs)

    xlabel --> "Rate [bbl/d]"
    ylabel --> "Pwf [psi]"
    xformatter --> :plain
    label --> l
    @series begin
        seriestype := :path
        Qrange, Pwf'
    end
    end

@userplot vfpgas

@recipe function f(h::vfpgas; dis=2.99, Tss=80, ∇Ts=1, ϵs=0.0006, Qns=10, Dns=50, tols=0.05, sgs=0.6, Incs=0)

    Depth,Thp,GasPVT, Qmax=h.args

DepthRange, Qrange, Pwf, leg, P = GasOutflow(Depth,Thp,GasPVT, Qmax; di=dis, Ts=Tss, ∇T=∇Ts, ϵ=ϵs, Qn=Qns, Dn=Dns, tol=tols, sg=sgs, Inc=Incs)

xlabel --> "Rate [Mscfd]"
ylabel --> "Pwf [psi]"
xformatter --> :plain
label --> leg
@series begin
    seriestype := :path
    Qrange, Pwf'
end
end


using Gridap
using Gridap.Inference
using Gridap.Fields
using Test

u(x) = x

p1 = Point(1,2)
p2 = Point(2,1)
p3 = Point(4,3)
p4 = Point(6,1)
x = [p1,p2,p3,p4]

f = FunctionField(u)
test_field(f,x,x)

Fill(f,4)


using LinearAlgebra
using Pathfinder: WoodburyPDMat
using PDMats
using Test

include("test_utils.jl")

function decompose(W)
    n, m = size(W.B)
    k = min(n, m)
    UA = W.UA
    Q = W.Q
    UC = W.UC
    X = zeros(eltype(UC), n, n)
    X[diagind(X)] .= true
    X[1:k, 1:k] .= UC'
    T = UA' * Q * X
    return T
end

function test_decompositions(W)
    UA = cholesky(W.A).U
    Q, R = qr(UA' \ W.B)
    C = I + R * W.D * R'
    UC = cholesky(Symmetric(C)).U
    @test W.UA ≈ UA
    @test W.Q ≈ Q
    @test W.UC ≈ UC
end

@testset "WoodburyPDMat" begin
    @testset "A $Atype, D $Dtype" for T in (Float64, Float32),
        Atype in (:dense, :diag),
        Dtype in (:dense, :diag),
        n in (5, 10),
        k in (5, 10)

        A = Atype === :diag ? rand_pd_diag_mat(T, n) : rand_pd_mat(T, n)
        B = randn(T, n, k)
        D = Dtype === :diag ? rand_pd_diag_mat(T, k) : rand_pd_mat(T, k)
        W = @inferred WoodburyPDMat{T} WoodburyPDMat(A, B, D)
        Wmat = A + B * D * B'

        @testset "basic" begin
            @test eltype(W) === T
            @test eltype(Wmat) === T
            @test size(W) == (n, n)
            @test Matrix(W) ≈ Wmat
            @test W[3, 5] ≈ Wmat[3, 5]
            @test W ≈ Wmat
            test_decompositions(W)
            M = decompose(W)
            @test M * M' ≈ Wmat
            @test WoodburyPDMat(A, B, big.(D)) isa WoodburyPDMat{BigFloat}
            @test Matrix(WoodburyPDMat(A, B, big.(D))) ≈ Wmat
            Wbig = convert(AbstractMatrix{BigFloat}, W)
            @test Wbig isa WoodburyPDMat{BigFloat}
            @test Wbig ≈ Wmat
            test_decompositions(W)
            @test convert(AbstractMatrix{T}, W) === W
        end

        @testset "inv" begin
            invW = @inferred WoodburyPDMat{T} inv(W)
            @test eltype(invW) === T
            invWmat = inv(Matrix(W))
            @test invW isa WoodburyPDMat
            @test invW ≈ invWmat
            test_decompositions(invW)
        end

        @testset "determinant" begin
            @test @inferred(det(W)) ≈ det(Wmat)
            @test @inferred(logdet(W)) ≈ logdet(Wmat)
            @test @inferred(logabsdet(W))[1] ≈ logabsdet(Wmat)[1]
            @test logabsdet(W)[2] == logabsdet(Wmat)[2]
        end

        @testset "diag" begin
            @test @inferred(diag(W)) ≈ diag(Wmat)
        end

        @testset "adjoint/transpose" begin
            @test W' === W
            @test transpose(W) === W
        end

        @testset "lmul!" begin
            x = randn(T, n)
            y = Wmat * x
            @test lmul!(W, x) === x
            @test x ≈ y

            X = randn(T, n, 5)
            Y = Wmat * X
            @test lmul!(W, X) === X
            @test X ≈ Y
        end

        @testset "mul!" begin
            x = randn(T, n)
            y = similar(x)
            @test mul!(y, W, x) === y
            @test y ≈ Wmat * x

            X = randn(T, n, 5)
            Y = similar(X)
            @test mul!(Y, W, X) === Y
            @test Y ≈ Wmat * X
        end

        @testset "*" begin
            @inferred Union{WoodburyPDMat{Float64},Matrix{Float64}} W * 5.0
            @test W * 5.0 isa WoodburyPDMat
            @test W * 5.0 ≈ Wmat * 5
            test_decompositions(W * 5.0)
            test_decompositions(W * 3)
            @test W * -2 isa Matrix
            @test W * -2 ≈ Wmat * -2

            x = randn(T, n)
            @test W * x ≈ Wmat * x

            X = randn(T, n)
            @test W * X ≈ Wmat * X
        end

        @testset "PDMats.dim" begin
            @test PDMats.dim(W) == n
        end

        @testset "unwhiten" begin
            M = decompose(W)

            x = randn(T, n)
            @test @inferred(unwhiten(W, x)) ≈ M * x

            X = randn(T, n, 100)
            @test @inferred(unwhiten(W, X)) ≈ M * X
        end

        @testset "invunwhiten!" begin
            M = decompose(W)

            x = randn(T, n)
            @test @inferred(Pathfinder.invunwhiten!(similar(x), W, x)) ≈ M' \ x

            X = randn(T, n, 100)
            @test @inferred(Pathfinder.invunwhiten!(similar(X), W, X)) ≈ M' \ X
        end

        @testset "PDMats.quad" begin
            x = randn(T, n)
            @test @inferred(quad(W, x)) ≈ dot(x, Wmat, x)

            u = randn(T, n)
            @test quad(W, Pathfinder.invunwhiten!(similar(u), W, u)) ≈ dot(u, u)

            X = randn(T, n, 100)
            @test @inferred(quad(W, X)) ≈ quad(PDMats.PDMat(Symmetric(Wmat)), X)

            U = randn(T, n, 100)
            @test quad(W, Pathfinder.invunwhiten!(similar(U), W, U)) ≈
                vec(sum(abs2, U; dims=1))
        end

        @testset "PDMats.invquad" begin
            x = randn(T, n)
            @test @inferred(invquad(W, x)) ≈ dot(x, inv(Wmat), x)

            u = randn(T, n)
            @test invquad(W, unwhiten(W, u)) ≈ dot(u, u)

            X = randn(T, n, 100)
            @test @inferred(invquad(W, X)) ≈ invquad(PDMats.PDMat(Symmetric(Wmat)), X)

            U = randn(T, n, 100)
            @test invquad(W, unwhiten(W, U)) ≈ vec(sum(abs2, U; dims=1))
        end
    end
end


__precompile__()

module SpatialFields

import ForwardDiff
import MultiPoly
import Base: convert
import DataStructures: OrderedDict
import StaticArrays: SVector, MVector

export InterpolatingSurface,
    HermiteRadialField,
    PolynomialScalarField,
    PolynomialVectorField

derivative(f::Function) = x -> ForwardDiff.derivative(f, x)
gradient(f::Function) = x -> ForwardDiff.gradient(f, x)

abstract AbstractScalarField{N} <: Function
abstract AbstractVectorField{N} <: Function

type ScalarField{N, F <: Function} <: AbstractScalarField{N}
    func::F
end
ScalarField{F <: Function}(N::Integer, f::F) = ScalarField{N, F}(f)
(field::ScalarField{N}){N}(x::AbstractVector) = field.func(x)

type VectorField{N, F <: Function} <: AbstractVectorField{N}
    func::F
end
VectorField{F <: Function}(N::Integer, f::F) = VectorField{N, F}(f)
(field::VectorField{N}){N}(x::AbstractVector) = field.func(x)

gradient{N}(field::AbstractScalarField{N}) = VectorField(N, x -> ForwardDiff.gradient(field, x))

include("radial_functions.jl")
include("hrbf.jl")
include("interpolating.jl")
include("polynomial.jl")

end # module


using BrainFlow

# enable logs
BrainFlow.enable_dev_logger(BrainFlow.BOARD_CONTROLLER)
BrainFlow.enable_dev_logger(BrainFlow.DATA_HANDLER)

params = BrainFlowInputParams()
board_shim = BrainFlow.BoardShim(BrainFlow.SYNTHETIC_BOARD, params)
sampling_rate = BrainFlow.get_sampling_rate(BrainFlow.SYNTHETIC_BOARD)

BrainFlow.prepare_session(board_shim)
BrainFlow.start_stream(board_shim)
sleep(10)
BrainFlow.stop_stream(board_shim)
data = BrainFlow.get_current_board_data(BrainFlow.get_nearest_power_of_two(sampling_rate), board_shim)
BrainFlow.release_session(board_shim)

eeg_channels = BrainFlow.get_eeg_channels(BrainFlow.SYNTHETIC_BOARD)
data_first_channel = data[eeg_channels[1], :]

# returns tuple of wavelet coeffs and lengths
wavelet_data = BrainFlow.perform_wavelet_transform(data_first_channel, BrainFlow.DB3, 2, BrainFlow.SYMMETRIC)
restored_wavelet_data = BrainFlow.perform_inverse_wavelet_transform(wavelet_data, length(data_first_channel),
																	BrainFlow.DB3, 2, BrainFlow.SYMMETRIC)

fft_data = BrainFlow.perform_fft(data_first_channel, BrainFlow.NO_WINDOW)
restored_fft_data = BrainFlow.perform_ifft(fft_data)

println("Original Data")
println(data_first_channel)
println("Restored from Wavelet Data")
println(restored_wavelet_data)
println("Restored from FFT Data")
println(restored_fft_data)


function getInteractionHamIPicture(
    Ham_S::Operator,
    Ham_int::Operator,
    t::AbstractFloat,
)::Operator
    U_op = evolutionOperator(Ham_S, t)
    return U_op' * Ham_int * U_op
end

function getInteractionHamIPictureA(
    Ham_S::Operator,
    Ham_int::Operator,
    t::AbstractFloat,
)::Array
    U_op = evolutionOperator(Ham_S, t)
    return (U_op' * Ham_int * U_op).data
end

function getInteractionHamIPictureA(
    Ham_int::Array,
    H_lambda::Array,
    H_S::Array,
    H_Sinv::Array,
    t::AbstractFloat,
)
    U_op = evolutionOperatorA(H_lambda, H_S, H_Sinv, t)
    U_opd = adjoint(U_op)
    return U_opd * Ham_int * U_op
end


# homogeneous velocity and density model
nz = 101; nx = 201;
vel = 3000 * ones(nz, nx);  # m/s
vel[51:end,:] .= 3500;  # m/s
rho = 2000 * ones(nz, nx);  # kg/m^3

# vertical and horizontal grid size
dz = 10; dx = 10;

# time step size and maximum modelling length
dt = 0.001; tmax = 1.0;  # use second as unit

# receiver location
irx = collect(1 :2 :nx);
irz = 2 * ones(length(irx));

# ==============================================================================
#                   dot-product test for the adjoint wavefield
# ==============================================================================
test_type = [(true , Float32, 1.0e-6 ),
             (true , Float64, 1.0e-12),
             (false, Float32, 1.0e-6 ),
             (false, Float64, 1.0e-12)
            ]
@testset "multi-step forward modelling and its adjoint operator" begin

    for (free_surface, data_format, tol) in test_type

        # generate finite difference stencil
        params = TdParams(rho, vel, free_surface, dz, dx, dt, tmax;
                 data_format=data_format, fd_flag="taylor", order=3, npml=20, apml=900.);

        # allocate snapshots for forward and adjoint operator
        spt1_f = Snapshot(params);
        spt2_f = Snapshot(params);
        spt1_b = Snapshot(params);
        spt2_b = Snapshot(params);

        # initialize spt1_f and spt2_b with random numbers
        N = params.Nz * params.Nx
        spt1_f.vz .= randn(params.data_format, N); spt2_b.vz .= randn(params.data_format, N);
        spt1_f.vx .= randn(params.data_format, N); spt2_b.vx .= randn(params.data_format, N);
        spt1_f.pz .= randn(params.data_format, N); spt2_b.pz .= randn(params.data_format, N);
        spt1_f.px .= randn(params.data_format, N); spt2_b.px .= randn(params.data_format, N);

        # temporary variables
        tmp    = zeros(params.data_format, params.Nz * params.Nx);
        tmp_z1 = zeros(params.data_format, params.Nz);
        tmp_z2 = zeros(params.data_format, params.Nz);
        tmp_x1 = zeros(params.data_format, params.Nx);
        tmp_x2 = zeros(params.data_format, params.Nx);

        # nt-step forward
        nt = 78
        for it = 1 : nt
            one_step_forward!(spt2_f, spt1_f, params, tmp_z1, tmp_z2, tmp_x1, tmp_x2);
            copy_snapshot!(spt1_f, spt2_f);
        end

        # nt-step adjoint
        for it = 1 : nt
            one_step_adjoint!(spt1_b, spt2_b, params, tmp, tmp_z1, tmp_z2, tmp_x1, tmp_x2);
            copy_snapshot!(spt2_b, spt1_b);
        end

        # inner product
        tmp1 = (dot(spt1_f.vz, spt1_b.vz) + dot(spt1_f.vx, spt1_b.vx)
              + dot(spt1_f.pz, spt1_b.pz) + dot(spt1_f.px, spt1_b.px))

        tmp2 = (dot(spt2_f.vz, spt2_b.vz) + dot(spt2_f.vx, spt2_b.vx)
              + dot(spt2_f.pz, spt2_b.pz) + dot(spt2_f.px, spt2_b.px))

        @test (tmp1-tmp2) / tmp1 < tol
    end
end

# ==============================================================================
#           adjoint Pz = Px in computational area
# ==============================================================================
@testset "test pz == px for the adjoint wavefield" begin

    for (free_surface, data_format, tol) in test_type

        # generate finite difference stencil
        params = TdParams(rho, vel, free_surface, dz, dx, dt, tmax;
                 data_format=data_format, fd_flag="taylor", order=3, npml=20, apml=900.);

        # generate band-limited random recordings
        w   = ricker(12.0, params.dt)
        hl  = floor(Int64, length(w)/2)
        rec = Recordings(irz, irx, params);
        idx_o= [1]

        # generate data trace by trace
        for i = 1 : rec.nr
            tmp = conv(randn(params.nt)*1000, w)
            copyto!(rec.p, idx_o[1], tmp, hl+1, params.nt)
            idx_o[1] = idx_o[1] + params.nt
        end

        # save adjoint pz and px as an 3D cube
        path_spt = joinpath(homedir(), "Desktop/snapshot.rsf")
        multi_step_adjoint(rec, params; path_spt=path_spt)

        # # read the adjoint snapshot cube
        (hdr, d) = read_RSdata(path_spt);

        zl = params.ntop+1; zu = zl + params.nz - 1;
        xl = params.npml+1; xu = xl + params.nx - 1;

        @test norm(d[zl:zu,xl:xu,3,:] - d[zl:zu,xl:xu,4,:]) / norm(d[zl:zu,xl:xu,3,:]) < tol
    end
end


using PyPlot
mutable struct Single_Stack
    L::Float64
    N::Int64
    t::Float64

    ρ::Array{Float64,1}
    w::Array{Float64,1}
    ρθ::Array{Float64,1}

    _grav::Float64

    Π::Function

end

# compute tendency
function spatial_residual(ss::Single_Stack, ρ, w, ρθ; t = 0.0)

    N, L = ss.N, ss.L
    Δz = L/N
    _grav = ss._grav
    Π = ss.Π

    ρ_res, w_res, ρθ_res = copy(ρ), copy(w), copy(ρθ)

    for i = 1:N
        # ρ₊ = (i == N ? ρ[i] : (ρ[i] + ρ[i+1])/2.0)
        # ρ₋ = (i == 1 ? ρ[i] : (ρ[i] + ρ[i-1])/2.0)

        ρ₊ = (i == N ? 0.0 : (ρ[i] + ρ[i+1])/2.0)
        ρ₋ = (i == 1 ? 0.0 : (ρ[i] + ρ[i-1])/2.0)

        ρ_res[i] = -(w[i+1]ρ₊ - w[i]ρ₋)/Δz

        # ρθ₊ = (i == N ? ρθ[i] : (ρθ[i] + ρθ[i+1])/2.0)
        # ρθ₋ = (i == 1 ? ρθ[i] : (ρθ[i] + ρθ[i-1])/2.0)

        ρθ₊ = (i == N ? 0.0 : (ρθ[i] + ρθ[i+1])/2.0)
        ρθ₋ = (i == 1 ? 0.0 : (ρθ[i] + ρθ[i-1])/2.0)

        ρθ_res[i] = -(w[i+1]ρθ₊ - w[i]ρθ₋)/Δz
    end

    w_res[1] = w_res[N+1] = 0
    for i = 2:N
        w_res[i] = -1/2*(ρθ[i-1]/ρ[i-1] + ρθ[i]/ρ[i]) *(Π(ρθ[i]) - Π(ρθ[i-1]))/Δz - _grav
    end

    return ρ_res,w_res, ρθ_res
end

# update with RK4
function temporal_update!(ss::Single_Stack, Δt::Float64)
    ρ, w, ρθ = ss.ρ, ss.w, ss.ρθ

    ρ_res1, w_res1, ρθ_res1 = spatial_residual(ss::Single_Stack, ρ, w, ρθ; t = ss.t)
    ρ_res2, w_res2, ρθ_res2 = spatial_residual(ss::Single_Stack, ρ + Δt/2*ρ_res1, w + Δt/2*w_res1, ρθ + Δt/2*ρθ_res1; t = ss.t+Δt/2)
    ρ_res3, w_res3, ρθ_res3 = spatial_residual(ss::Single_Stack, ρ + Δt/2*ρ_res2, w + Δt/2*w_res2, ρθ + Δt/2*ρθ_res2; t = ss.t+Δt/2)
    ρ_res4, w_res4, ρθ_res4 = spatial_residual(ss::Single_Stack, ρ + Δt*ρ_res3, w + Δt*w_res3, ρθ + Δt*ρθ_res3; t = ss.t+Δt)

    ρ  .+=  Δt/6*(ρ_res1  + 2ρ_res2  + 2ρ_res3  + ρ_res4)
    w  .+=  Δt/6*(w_res1  + 2w_res2  + 2w_res3  + w_res4)
    ρθ .+=  Δt/6*(ρθ_res1 + 2ρθ_res2 + 2ρθ_res3 + ρθ_res4)

    ss.t += Δt
    @info "ρ ∈ ",   minimum(ρ), maximum(ρ)
    @info "w ∈ ",   minimum(w), maximum(w)
    @info "ρθ ∈ ",  minimum(ρθ), maximum(ρθ)

end

function compute_wave_speed(ρ, w, ρθ, _R_m, _C_p, γ, _MSLP)
    p =  ( ρθ * _R_m / _MSLP^(_R_m/_C_p) ).^(γ)
    c = sqrt.(γ*p./ρ)

    return maximum(w) + maximum(c)
end

# Decaying temperature profile
function DecayingTemperatureProfile(z::Float64; T_virt_surf, T_min_ref, _R_d, _grav, _MSLP)
    # Scale height for surface temperature
    H_sfc = _R_d * T_virt_surf / _grav
    H_t = H_sfc

    z′ = z / H_t
    tanh_z′ = tanh(z′)

    ΔTv =  T_virt_surf -  T_min_ref
    Tv =  T_virt_surf - ΔTv * tanh_z′

    ΔTv′ = ΔTv / T_virt_surf
    p = -H_t * (z′ + ΔTv′ * (log(1 - ΔTv′ * tanh_z′) - log(1 + tanh_z′) + z′))
    p /= H_sfc * (1 - ΔTv′^2)
    p = _MSLP * exp(p)
    ρ = p/(_R_d*Tv)
    return (Tv, p, ρ)
end

# update discrete_hydrostatic_balance!
function discrete_hydrostatic_balance!(ρ, w, ρθ, Δz::Float64, _grav::Float64, Π::Function)
    for i = 1:length(ρ)-1
        ρ[i+1] = ρθ[i+1]/(-2Δz*_grav/(Π(ρθ[i+1]) - Π(ρθ[i])) - ρθ[i]/ρ[i])
    end
end

###################################### Start
ndays = 10

# whether use discrete hydrostatic balance correction
DHB = true

N = 30
L = 30e3
Δz = L/N

_MSLP = 1e5
_grav = 9.8
_R_m = 287.058
γ = 1.4
_C_p = _R_m * γ / (γ - 1)
_C_v = _R_m / (γ - 1)
_R_d = _R_m

Π = (ρθ) -> _C_p*(_R_d*ρθ/_MSLP)^(_R_m/_C_v)

###################################### Initialization
ρ = zeros(Float64, N)
w = zeros(Float64, N+1)
ρθ = zeros(Float64, N)

# set the bottom velocity
# w[1] = 1.0

for i = 1:N
    z = (i-0.5)*L/N
    Tvᵢ, pᵢ, ρᵢ = DecayingTemperatureProfile(z; T_virt_surf = 280.0, T_min_ref = 230.0, _R_d = _R_d, _grav = _grav, _MSLP = _MSLP)
    ρ[i]  = ρᵢ
    ρθ[i] = ρᵢ*Tvᵢ*(_MSLP/pᵢ)^(_R_m/_C_p)
end

if DHB
    discrete_hydrostatic_balance!(ρ, w, ρθ, Δz, _grav, Π)
end

ss = Single_Stack(L, N, 0.0, copy(ρ), copy(w), copy(ρθ), _grav, Π)
max_wave_speed = compute_wave_speed(ρ, w, ρθ, _R_m, _C_p, γ, _MSLP)

@info "maximum Δt (cfl = 1) = ",  Δz / max_wave_speed

###################################### Advance
Δt = 3.0
Nt = Int64(round(86400  * ndays / Δt))
t_end
for i = 1:Nt
    temporal_update!(ss, Δt)
end
T = Nt * Δt

zz = Array(LinRange(Δz/2, L - Δz/2, N))
zz_h = Array(LinRange(0, L, N+1))

fig, (ax1, ax2, ax3) = PyPlot.subplots(ncols = 3, nrows=1, sharex=false, sharey=false, figsize=(12,6))
ax1.plot(ρ, zz,  "-o", fillstyle = "none", label = "T=0")
ax1.plot(ss.ρ, zz, "-", fillstyle = "none", label = "T=$(ndays) days")
ax1.legend()
ax1.set_ylabel("ρ")

ax2.plot(w, zz_h, "-o", fillstyle = "none", label = "T=0")
ax2.plot(ss.w, zz_h,  "-", fillstyle = "none", label = "T=$(ndays) days")
ax2.legend()
ax2.set_ylabel("w")

ax3.plot(ρθ, zz,  "-o", fillstyle = "none", label = "T=0")
ax3.plot(ss.ρθ, zz, "-", fillstyle = "none", label = "T=$(ndays) days")
ax3.legend()
ax3.set_ylabel("ρθ")

fig.tight_layout()
fig.savefig("HB-Theta-$(ndays)-HB-$(DHB).png")


using .RationalFunctionFields: RationalFunctionField, compute_groebner!,
                    contains_using_groebner, contains_randomized

logger = Logging.SimpleLogger(stderr, Logging.Debug)
Logging.global_logger(logger)

@testset "Basics tests" begin

end

@testset "Contains tests" begin
    FF = Nemo.QQ
    R, (x1, x2) = Nemo.PolynomialRing(FF, ["x1", "x2"])
    # set = [ (a^2 + b^2) // 1, a*b // 1 ]
    # set = [ a*b // (a + b), (a + b) // 1, (a^3 - b) // 1 ]
    set = [
          ( x1^2 + x2^2 ) // R(1),
          ( x1^3 + x2^3 ) // R(1),
          ( x1^4 + x2^4 ) // R(1)
    ]

    FF = RationalFunctionField(set)
    compute_groebner!(FF)

#    @test contains_using_groebner( FF, x1*x2 // R(1) )
 #   @test contains_using_groebner( FF, (x1 + x2) // R(1) )

    # note that may fail occasionaly
    @test contains_randomized( FF, (x1*x2) // R(1) )
    @test contains_randomized( FF, (x1 + x2 + x1*x2) // R(1) )

end