# binary search in interval (low,high)
# assumes that if pred(x)==true, then pred(y)==true ∀y>x
# assumes pred(low)==false and pred(high)==true
function binary_search_interval(pred, low::T, high::T, precision::T) where {T<:AbstractFloat}
    while high-low > precision
        mid = (low+high)/2

        if pred(mid)
            high = mid
        else
            low = mid
        end
    end
    (low+high)/2
end
binary_search_interval(pred, low, high, precision) = binary_search_interval(pred, promote(low, high, precision)...)

# Soft thresholding operator S (See Witten et al.)
function soft_threshold(a,c)
    sign(a) .* max( abs(a)-c, 0 )
end

# computes right singular vector of X
function rightsingularvector(X)
    K = Symmetric(X'X);
    N = size(K,1);
    _,v=eig(K,N:N);
    dropdims(v; dims=2)
end

# computes
#   argmax_u   u'w
#   subject to ||u||_2^2 ≤ 1
# (When used for PCA iterations, w = Xv or w = X'u = (u'X)'.)
function argmax_unpenalized(w)
    # w/norm(w)
    w/max(1e-16,norm(w)) # avoid div by zero
end

# computes
#   argmax_u   uᵀw
#   subject to ||u||_2^2 ≤ 1
#              uᵀO = (0 0 ... 0)          (vector of zeros)
# (When used for PCA iterations, w = Xv or w = X'u = (u'X)' and U is matrix with previous components, i.e. U and V respectively.)
# See Witten et al 2009, section 3.2.
#     uₖ = Cₖ₋₁Cₖ₋₁ᵀXvₖ / ||Cₖ₋₁ᵀXvₖ||₂
#        = (I - Uₖ₋₁Uₖ₋₁ᵀ)Xvₖ / ||Cₖ₋₁ᵀXvₖ||₂
#        = (Xvₖ - Uₖ₋₁Uₖ₋₁ᵀXvₖ) / √( ||Xvₖ||₂² - ||Uₖ₋₁ᵀXvₖ||₂² )
#        = (w - UUᵀw) / √( ||w||₂² - ||Uᵀw||₂² )
#     where C is the orthogonal complement of O (not needed in computations).
function argmax_unpenalized_orthogonal(w, U)
    # Uw = U'w
    # (w-U*Uw) / sqrt(sum(w.^2)-sum(Uw.^2))

    # Simplified
    u = w - U*(U'w)
    u / norm(u)
end

# computes
#   argmax_u   u'a
#   subject to ||u||₂² ≤ 1
#              ||u||₁ ≤ c
# (When used for PCA iterations, a = Xv or a = X'u = (u'X)'.)
function argmax_l1penalized(a,c)
    # test with Δ=0 first (i.e. S(a,0)=a)
    # u = a/norm(a)
    u = a/max(1e-16, norm(a)) # avoid div by zero
    c2 = c*c # we almost only use c squared
    (length(a)<=c2 || norm(u,1)≤c) && return u

    # otherwise chose Δ s.t. ||u||₁=c, where u = S(a,Δ)/||S(a,Δ)||₂

    # we could use a heap for b to avoid sorting all of b (should save a lot for very sparse vectors!)
    heap = binary_maxheap(abs(a)) # work with sorted abs(a) since it doesn't change value of Δ and makes the search simpler

    n1 = 0.0
    n2 = 0.0 # squared
    pk=qk = 0.0
    Δ = 0.0
    b = pop!(heap)

    k = 1
    #while k<length(b)
    while k<length(a)
        # pk += b[k]*b[k]
        # qk += b[k]
        pk += b*b
        qk += b

        bNext = pop!(heap)
        # n1Next = qk - k*b[k+1]
        # n2Next = k*b[k+1]*b[k+1] + pk - 2b[k+1]*qk
        n1Next = qk - k*bNext
        n2Next = k*bNext*bNext + pk - 2bNext*qk

        n1Next*n1Next≥c2*n2Next && break # squared version, ok since everything is positive - now we know bₖ₋₁≥Δ≥bₖ

        n1,n2 = n1Next,n2Next
        b = bNext
        k+=1
    end

    A = (k*k-k*c2)
    B = 2(k*n1 - c2*n1)
    C = n1*n1-c2*n2

    p = B/A
    q = C/A

    #Δ = max(0.0, b[k]-(-p/2+sqrt(max(p*p/4-q,0.0)))) # max with zero only to avoid numerical problems
    Δ = max(0.0, b-(-p/2+sqrt(max(p*p/4-q,0.0)))) # max with zero only to avoid numerical problems

    u=soft_threshold(a,Δ);
    # u/norm(u)
    u/max(1e-16, norm(u)) # avoid div by zero
end

# numerically compute rank 1 approximation
# 2nd component are retrieved using pca_rank(X-u*s*v',...), where (u,s,v) is the result from the first, and so on.
function pcarank1l1(X,c1,c2,tol,maxIters)
    P = size(X,1)
    N = size(X,2)

    v = rightsingularvector(X) # start with optimal v for non-penalized problem.
    u = zeros(P)
    s = 0.0
    converged = false

    # for i=1:maxIters # TODO: convergence criterion
    #     u = c1==Inf ? argmax_unpenalized(X*v) : argmax_l1penalized(X*v,c1)
    #     v = c2==Inf ? argmax_unpenalized(X'u) : argmax_l1penalized(X'u,c2)
    # end
    # s = (u'X*v)[1]

    for i=1:maxIters
        prevU = u
        u = c1==Inf ? argmax_unpenalized(X*v) : argmax_l1penalized(X*v,c1)

        Xu = X'u
        prevV = v
        v = c2==Inf ? argmax_unpenalized(Xu) : argmax_l1penalized(Xu,c2)

        prevS = s
        s = (v'Xu)[1]
        #println("s: ", s)

        #println("i: ", i, ", v: ", norm(v-prevV), ", u: ", norm(u-prevU), ", s:", s-prevS)

        norm(v-prevV)<tol && norm(u-prevU)<tol && (converged=true; break)
        s-prevS < tol && (converged=true; break)

        #norm(v-prevV)<tol && norm(u-prevU)<tol && (converged=true; println(norm(v-prevV), " ", norm(u-prevU)); break)
        #s-prevS < tol && (converged=true; println(s-prevS); break)
    end

    (u,s,v,converged)
end

# numerically compute rank 1 approximation
# 2nd component is retrieved using pcarank1l1_orthogonal(X-u*s*v',c1,V), where (u,s,v) is the result from the first, and so on.
function pcarank1l1_orthogonal(X,c1,V,v0,tol,maxIters)
    P = size(X,1)
    N = size(X,2)

    v = v0

    u = zeros(P)
    s = 0.0
    converged = false

    for i=1:maxIters
        prevU = u
        u = c1==Inf ? argmax_unpenalized(X*v) : argmax_l1penalized(X*v,c1)

        Xu = X'u
        prevV = v
        v = isempty(V) ? argmax_unpenalized(Xu) : argmax_unpenalized_orthogonal(Xu,V)

        prevS = s
        s = (v'Xu)[1]

        norm(v-prevV)<tol && norm(u-prevU)<tol && (converged=true; break)
        s-prevS < tol && (converged=true; break)
    end

    (u,s,v,converged)
end

function pcarank1l1_orthogonal(X,c1,V,tol,maxIters)
    # v0 = rightsingularvector(X) # start with optimal v for non-penalized problem.

    # This enforces orthogonality to V for initial guess.
    v0 = isempty(V) ? rightsingularvector(X) :
                      rightsingularvector( X - (X*V)*V' )

    pcarank1l1_orthogonal(X,c1,V,v0,tol,maxIters)
end

# Sparce PCA with l1-constraints on U-side and orthogonality constraints on V-side.
# If a component with higher optimum than a previous component is found, we backtrack and use it as a starting guess for the previous component.
# This guarantees that S will be sorted and reduces some issues with getting stuck in local optima.
function sparsepcal1(X, k::Integer, c1; backtrack=true, tol=1e-6, maxIters=1000)
    P = size(X,1)
    N = size(X,2)
    @assert N<=P # computations are arranged to exploit N<P (or rather, N<<P)

    # set L1-constraints to Inf if they are not active (better for pcarank1l1)
    c1 = c1<√P ? c1 : Inf

    U = zeros(P,k)
    S = zeros(k)
    V = zeros(N,k)

    XOrig = backtrack ? copy(X) : similar(X,0,0) # only necessary to keep track of original matrix if we are backtracking

    v0 = rightsingularvector(X)

    i=1
    while i<=k
        U[:,i], S[i], V[:,i], converged = pcarank1l1_orthogonal(X,c1,V[:,1:i-1],v0,tol,maxIters)
        converged || @warn("sparsepca component $i did not converge")

        if S[i]<0 || any(abs(V[:,i]'V[:,1:i-1]) .> 1e-6) # In this case, we are not finding any more components. I.e. the rank(XOrig)<i.
            U[:,i], S[i], V[:,i] = .0, .0, .0
            break
        end

        if backtrack && i>1 && S[i]-S[i-1]>tol # Backtrack if we found a component with higher S[i] than the previous one.
            v0 = V[:,i] # Use this optimum as a starting point for the search on the previous matrix. The orthogonality of samples (V) guarantees that the starting point will have the same S value, i.e. we will find a better local optima.
            i = something(findfirst(S[i]-S.>tol),0) - 1 # How far do we need to backtrack?
            X = XOrig - U[:,1:i]*(diagm(S[1:i])*V[:,1:i]') # Restore matrix state s.t. only the i first components have been removed.
        else
            X = X - U[:,i]*(S[i]*V[:,i]') # Remove this component to be able to find next component.
            v0 = rightsingularvector( X - (X*V)*V' ) # Starting guess orthogonal to V.
        end
        i += 1
    end

    # change signs for stability, favouring U's that are positive on average
    sgn = (sum(U;dims=1).>=0)*2.0-1.0
    U .*= sgn
    V .*= sgn

    (U,S,V)
end


# kolmogorov_smirnov.jl
# Kolmogorov–Smirnov
#
# Copyright (C) 2014   Christoph Sawade
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

export
    ExactOneSampleKSTest,
    ApproximateOneSampleKSTest, ApproximateTwoSampleKSTest

abstract type KSTest <: HypothesisTest end
abstract type ApproximateKSTest <: KSTest end
abstract type ExactKSTest <: KSTest end

population_param_of_interest(x::KSTest) = ("Supremum of CDF differences", 0.0, x.δ) # parameter of interest: name, value under h0, point estimate
default_tail(test::KSTest) = :both

## ONE SAMPLE KS-TEST

# compute supremum of differences between target and empirical cdf before and after the jump of the empirical cdf.
function ksstats(x::AbstractVector{T}, d::UnivariateDistribution) where T<:Real
    n = length(x)
    cdfs = cdf.(Ref(d), sort(x))
    δp = maximum((1:n) / n - cdfs)
    δn = -minimum((0:n-1) / n - cdfs)
    δ = max(δn, δp)
    (n, δ, δp, δn)
end

### EXACT KOLMOGOROV SMIRNOV TEST

struct ExactOneSampleKSTest <: ExactKSTest
    n::Int      # number of observations
    δ::Float64  # supremum of CDF differences
    δp::Float64 # supremum of the positive CDF differences
    δn::Float64 # supremum of the negative CDF differences
end

"""
    ExactOneSampleKSTest(x::AbstractVector{<:Real}, d::UnivariateDistribution)

Perform a one-sample exact Kolmogorov–Smirnov test of the null hypothesis that the data in
vector `x` comes from the distribution `d` against the alternative hypothesis that the
sample is not drawn from `d`.

Implements: [`pvalue`](@ref)
"""
function ExactOneSampleKSTest(x::AbstractVector{T}, d::UnivariateDistribution) where T<:Real
    if length(x) > length(unique(x))
        @warn("This test is inaccurate with ties")
    end

    ExactOneSampleKSTest(ksstats(x, d)...)
end

testname(::ExactOneSampleKSTest) = "Exact one sample Kolmogorov-Smirnov test"

function show_params(io::IO, x::ExactOneSampleKSTest, ident="")
    println(io, ident, "number of observations:   $(x.n)")
end

function pvalue(x::ExactKSTest; tail=:both)
    if tail == :left
        pvalue(KSOneSided(x.n), x.δn; tail=:right)
    elseif tail == :right
        pvalue(KSOneSided(x.n), x.δp; tail=:right)
    elseif tail == :both
        pvalue(KSDist(x.n), x.δ; tail=:right)
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end
end

### APPROXIMATE KOLMOGOROV SMIRNOV TEST

struct ApproximateOneSampleKSTest <: ApproximateKSTest
    n::Int      # number of observations
    δ::Float64  # supremum of CDF differences
    δp::Float64 # supremum of the positive CDF differences
    δn::Float64 # suproemum of the negative CDF differences
end

"""
    ApproximateOneSampleKSTest(x::AbstractVector{<:Real}, d::UnivariateDistribution)

Perform an asymptotic one-sample Kolmogorov–Smirnov test of the null hypothesis that the
data in vector `x` comes from the distribution `d` against the alternative hypothesis
that the sample is not drawn from `d`.

Implements: [`pvalue`](@ref)
"""
function ApproximateOneSampleKSTest(x::AbstractVector{T}, d::UnivariateDistribution) where T<:Real
    if length(x) > length(unique(x))
        @warn("This test is inaccurate with ties")
    end

    ApproximateOneSampleKSTest(ksstats(x, d)...)
end

testname(::ApproximateOneSampleKSTest) = "Approximate one sample Kolmogorov-Smirnov test"

function show_params(io::IO, x::ApproximateOneSampleKSTest, ident="")
    println(io, ident, "number of observations:   $(x.n)")
    println(io, ident, "KS-statistic:             $(sqrt(x.n)*x.δ)")
end

# one-sided: http://www.encyclopediaofmath.org/index.php/Kolmogorov-Smirnov_test
function pvalue(x::ApproximateOneSampleKSTest; tail=:both)
    if tail == :left
        exp(-2*x.n*x.δn^2)
    elseif tail == :right
        exp(-2*x.n*x.δp^2)
    elseif tail == :both
        pvalue(Kolmogorov(), sqrt(x.n)*x.δ; tail=:right)
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end
end

## TWO SAMPLE KS-TEST

### APPROXIMATE KOLMOGOROV SMIRNOV TEST

struct ApproximateTwoSampleKSTest <: ApproximateKSTest
    n_x::Int    # number of observations
    n_y::Int    # number of observations
    δ::Float64  # supremum of CDF differences
    δp::Float64 # supremum of the positive CDF differences
    δn::Float64 # suproemum of the negative CDF differences
end

"""
    ApproximateTwoSampleKSTest(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Perform an asymptotic two-sample Kolmogorov–Smirnov-test of the null hypothesis that `x`
and `y` are drawn from the same distribution against the alternative hypothesis that they
come from different distributions.

Implements: [`pvalue`](@ref)

# External links

  * [Approximation of one-sided test (Encyclopedia of Mathematics)
    ](https://www.encyclopediaofmath.org/index.php/Kolmogorov-Smirnov_test)
"""
function ApproximateTwoSampleKSTest(x::AbstractVector{T}, y::AbstractVector{S}) where {T<:Real, S<:Real}
    n_x, n_y = length(x), length(y)
    if n_x+n_y > length(unique([x; y]))
        @warn("This test is inaccurate with ties")
    end

    ApproximateTwoSampleKSTest(ksstats(x, y)...)
end

testname(::ApproximateTwoSampleKSTest) = "Approximate two sample Kolmogorov-Smirnov test"

function show_params(io::IO, x::ApproximateTwoSampleKSTest, ident="")
    n = x.n_x*x.n_y/(x.n_x+x.n_y)
    println(io, ident, "number of observations:   [$(x.n_x),$(x.n_y)]")
    println(io, ident, "KS-statistic:              $(sqrt(n)*x.δ)")
end

function pvalue(x::ApproximateTwoSampleKSTest; tail=:both)
    n = x.n_x*x.n_y/(x.n_x+x.n_y)
    if tail == :left
        exp(-2*n*x.δn^2)
    elseif tail == :right
        exp(-2*n*x.δp^2)
    elseif tail == :both
        pvalue(Kolmogorov(), sqrt(n)*x.δ; tail=:right)
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end
end

# compute supremum of differences between empirical cdfs.
function ksstats(x::AbstractVector{T}, y::AbstractVector{S}) where {T<:Real, S<:Real}
    n_x, n_y = length(x), length(y)
    sort_idx = sortperm([x; y])
    pdf_diffs = [ones(n_x)/n_x; -ones(n_y)/n_y][sort_idx]
    cdf_diffs = cumsum(pdf_diffs)
    δp = maximum(cdf_diffs)
    δn = -minimum(cdf_diffs)
    δ = max(δp, δn)
    (n_x, n_y, δ, δp, δn)
end


using Zygote, Test, LinearAlgebra

@testset "basic" begin

@test gradient(x -> real(abs(x)*exp(im*angle(x))), 10+20im)[1] ≈ 1
@test gradient(x -> imag(real(x)+0.3im), 0.3)[1] ≈ 0
@test gradient(x -> imag(conj(x)+0.3im), 0.3 + 0im)[1] ≈ -1im
@test gradient(x -> imag(conj(x)+0.3im), 0.3)[1] ≈ 0  # projected to zero
@test gradient(x -> abs((imag(x)+0.3)), 0.3 + 0im)[1] ≈ 1im
@test gradient(x -> abs((imag(x)+0.3)), 0.3)[1] ≈ 0

@test gradient(a -> real((a*conj(a))), 0.3im)[1] == 0.6im
@test gradient(a -> real((a.*conj(a))), 0.3im)[1] == 0.6im
@test gradient(a -> real(([a].*conj([a])))[], 0.3im)[1] == 0.6im
@test gradient(a -> real(([a].*conj.([a])))[], 0.3im)[1] == 0.6im
@test gradient(a -> real.(([a].*conj.([a])))[], 0.3im)[1] == 0.6im

@test gradient(x -> norm((im*x) ./ (im)), 2)[1] == 1
@test gradient(x -> norm((im) ./ (im*x)), 2)[1] == -1/4
@test gradient(x -> real(det(x)), [1 2im; 3im 4])[1] ≈ [4 3im; 2im 1]
@test gradient(x -> real(logdet(x)), [1 2im; 3im 4])[1] ≈ [4 3im; 2im 1]/10
@test gradient(x -> real(logabsdet(x)[1]), [1 2im; 3im 4])[1] ≈ [4 3im; 2im 1]/10

# https://github.com/FluxML/Zygote.jl/issues/705
@test gradient(x -> imag(sum(exp, x)), [1,2,3])[1] ≈ real(im .* exp.(1:3))
@test gradient(x -> imag(sum(exp, x)), [1+0im,2,3])[1] ≈ im .* exp.(1:3)

end # @testset

fs_C_to_R = (real,
             imag,
             abs,
             abs2,
             z -> abs(z)*cos(im*angle(z)),
             z->abs(cos(exp(z))),
             z->3*real(z)^3-2*imag(z)^5
             )
@testset "C->R" begin
    for f in fs_C_to_R
        for z in (1.0+2.0im, -2.0+pi*im)
            grad_zygote = gradient(real∘f, z)[1]
            ε = 1e-8
            grad_fd = (f(z+ε)-f(z))/ε + im*(f(z+ε*im)-f(z))/ε
            @test abs(grad_zygote - grad_fd) < sqrt(ε)
        end
    end
end

fs_C_to_C_holomorphic = (cos,
                         exp,
                         log,
                         z->z^2,
                         z->(real(z)+im*imag(z))^2,
                         z->real(z)^2 - imag(z)^2 +2im*(real(z)*imag(z)),
                         z->exp(cos(log(z))),
                         z->abs(z)*exp(im*angle(z)),
                         )
@testset "C->C holomorphic" begin
    for f in fs_C_to_C_holomorphic
        for z in (1.0+2.0im, -2.0+pi*im)
            grad_zygote = gradient(real∘f, z)[1]
            ε = 1e-8
            grad_fd_r = (f(z+ε)-f(z))/ε
            grad_fd_i = (f(z+ε*im)-f(z))/(ε*im)
            @assert abs(grad_fd_r - grad_fd_i) < sqrt(ε) # check the function is indeed holomorphic
            @test abs(grad_zygote - conj(grad_fd_r)) < sqrt(ε)
        end
    end
end

fs_C_to_C_non_holomorphic = (conj,
                             z->abs(z)+0im,
                             z->im*abs(z),
                             z->abs2(z)+0im,
                             z->im*abs2(z),
                             z->z'z,
                             z->conj(z)*z^2,
                             )
@testset "C->C non-holomorphic" begin
    for f in (fs_C_to_C_holomorphic...,fs_C_to_C_holomorphic...)
        for z in (1.0+2.0im, -2.0+pi*im)
            grad_zygote = gradient(real∘f, z)[1]
            ε = 1e-8
            grad_fd = real(f(z+ε)-f(z))/ε + im*real(f(z+ε*im)-f(z))/ε
            @test abs(grad_zygote - grad_fd) < sqrt(ε)
        end
    end
end

@testset "issue 342" begin
    @test Zygote.gradient(x->real(x + 2.0*im), 3.0) == (1.0,)
    @test Zygote.gradient(x->imag(x + 2.0*im), 3.0) == (0.0,)
end

@testset "issue 402" begin
    A = [1,2,3.0]
    y, B_getindex = Zygote.pullback(x->getindex(x,2,1),Diagonal(A))
    bA = B_getindex(1)[1]
    @test bA isa Diagonal
    @test bA == [0.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0]
end

@testset "issue #917" begin
    function fun(v)
        c = v[1:3] + v[4:6]*im
        r = v[7:9]
        sum(r .* abs2.(c)) # This would be calling my actual function depending on r and c
    end
    @test Zygote.hessian(fun, collect(1:9)) ≈ [14 0 0 0 0 0 2 0 0; 0 16 0 0 0 0 0 4 0; 0 0 18 0 0 0 0 0 6; 0 0 0 14 0 0 8 0 0; 0 0 0 0 16 0 0 10 0; 0 0 0 0 0 18 0 0 12; 2 0 0 8 0 0 0 0 0; 0 4 0 0 10 0 0 0 0; 0 0 6 0 0 12 0 0 0]
end


"""
    ht_p2_p1_veclap_alt

The manufactured-solution colliding flow example from Elman et al 2014. The
Hood-Taylor formulation with quadratic triangles for the velocity and continuous
pressure on linear triangles.

This implementation is an alternative: for the velocity, a single finite element
space with multiple components (2) is used instead of multiple finite element
spaces.

The formulation is the one derived in Donea, Huerta, Introduction to the finite element
method, 1993. Page 486 ff, and Elman, et al., Finite elements and fast
iterative solvers, p. 132. In other words, it is the vector Laplacian version.
"""
module ht_p2_p1_veclap_alt

using LinearAlgebra
using StaticArrays
using MeshCore: retrieve, nrelations, nentities, ir_identity, attribute, VecAttrib
using MeshSteward: T6block, T6toT3
using MeshSteward: Mesh, attach!, baseincrel, boundary
using MeshSteward: vselect, geometry, summary, transform
using MeshSteward: vtkwrite
using Elfel.RefShapes: manifdim, manifdimv
using Elfel.FElements: FEH1_T6, FEH1_T3, refshape, Jacobian
using Elfel.FESpaces: FESpace, ndofs, setebc!, nunknowns, doftype
using Elfel.FESpaces: numberfreedofs!, numberdatadofs!, numberdofs!
using Elfel.FESpaces: scattersysvec!, makeattribute, gathersysvec!, edofcompnt
using Elfel.FESpaces: highestfreedofnum, highestdatadofnum
using Elfel.FEIterators: FEIterator, ndofsperel, elnodes, eldofs, eldofvals
using Elfel.FEIterators: jacjac, location
using Elfel.QPIterators: QPIterator, bfun, bfungrad, weight
using Elfel.Assemblers: SysmatAssemblerSparse, start!, finish!, assemble!
using Elfel.Assemblers: SysvecAssembler
using Elfel.LocalAssemblers: LocalMatrixAssembler, LocalVectorAssembler, init!
using UnicodePlots

mu = 1.0 # dynamic viscosity
A = 1.0 # half of the length of the side of the square
trueux = (x, y) -> 20 * x * y ^ 3
trueuy = (x, y) -> 5 * x ^ 4 - 5 * y ^ 4
truep = (x, y) -> 60 * x ^ 2 * y - 20 * y ^ 3

function genmesh(N)
    # Hood-Taylor pair of meshes is needed
    # This mesh will be for the velocities
    vmesh = Mesh()
    attach!(vmesh, T6block(2 * A, 2 * A, N, N), "velocity")
    ir = baseincrel(vmesh)
    transform(ir, x -> x .- A)
    # This mesh will be used for the pressures. Notice that it needs to be
    # "compatible" with the velocity mesh in the sense that they need to share
    # the nodes at the corners of the triangles.
    pmesh = Mesh()
    attach!(pmesh, T6toT3(baseincrel(vmesh, "velocity")), "pressure")
    return vmesh, pmesh
end

function assembleK(ufesp, pfesp, tndof, mu)
    function integrate!(ass, elits, qpits, mu)
        unedof, pnedof = ndofsperel.(elits)
        uedofcomp = edofcompnt(ufesp)
        kuu = LocalMatrixAssembler(unedof, unedof, 0.0)
        kup = LocalMatrixAssembler(unedof, pnedof, 0.0)
        for el in zip(elits...)
            uel, pel = el
            init!(kuu, eldofs(uel), eldofs(uel))
            init!(kup, eldofs(uel), eldofs(pel))
            for qp in zip(qpits...)
                uqp, pqp = qp
                Jac, J = jacjac(uel, uqp)
                JxW = J * weight(uqp)
                gradNu = bfungrad(uqp, Jac)
                Np = bfun(pqp)
                for j in 1:unedof, i in 1:unedof
                    if uedofcomp[i] == uedofcomp[j]
                        kuu[i, j] += (mu * JxW) * (dot(gradNu[i], gradNu[j]))
                    end
                end
                for j in 1:pnedof, i in 1:unedof
                    kup[i, j] += (-JxW * Np[j]) * gradNu[i][uedofcomp[i]]
                end
            end
            assemble!(ass, kuu)
            assemble!(ass, kup)
            assemble!(ass, transpose(kup))
        end
        return ass
    end

    elits = (FEIterator(ufesp), FEIterator(pfesp))
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(ufesp, qargs), QPIterator(pfesp, qargs))
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    integrate!(ass, elits, qpits, mu)
    return finish!(ass)
end

function solve!(U, K, F, nu)
    KT = K * U
    U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

function evaluate_pressure_error(pfesp)
    function integrate!(elit, qpit, truep)
        pnedof = ndofsperel(elit)
        E = 0.0
        for el in elit
            dofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Np = bfun(qp)
                pt = truep(location(el, qp)...)
                pa = 0.0
                for j in 1:pnedof
                    pa += (dofvals[j] * Np[j])
                end
                E += (JxW) * (pa - pt)^2
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(pfesp)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(pfesp, qargs)
    return integrate!(elit, qpit, truep)
end

function evaluate_velocity_error(ufesp)
    function integrate!(elit, qpit, trueux, trueuy)
        unedof = ndofsperel(elit)
        uedofcomp = edofcompnt(ufesp)
        E = 0.0
        for el in elit
            udofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Nu = bfun(qp)
                uxt = trueux(location(el, qp)...)
                uyt = trueuy(location(el, qp)...)
                uxa = 0.0
                uya = 0.0
                for j in 1:unedof
                    (uedofcomp[j] == 1) && (uxa += (udofvals[j] * Nu[j]))
                    (uedofcomp[j] == 2) && (uya += (udofvals[j] * Nu[j]))
                end
                E += (JxW) * ((uxa - uxt)^2 + (uya - uyt)^2)
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(ufesp)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(ufesp, qargs)
    return integrate!(elit, qpit, trueux, trueuy)
end

function run(N)
    vmesh, pmesh = genmesh(N)
    # Velocity spaces
    ufesp = FESpace(Float64, vmesh, FEH1_T6(), 2)
    locs = geometry(vmesh)
    inflate = A / N / 100
    # The entire boundary
    boxes = [[-A A -A -A], [-A -A -A A], [A A -A A], [-A A A A]]
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            setebc!(ufesp, 0, i, 1, trueux(locs[i]...))
            setebc!(ufesp, 0, i, 2, trueuy(locs[i]...))
        end
    end
    # Pressure space
    pfesp = FESpace(Float64, pmesh, FEH1_T3(), 1)
    atcenter = vselect(geometry(pmesh); nearestto = [0.0, 0.0])
    setebc!(pfesp, 0, atcenter[1], 1, 0.0)
    # Number the degrees of freedom
    numberdofs!([ufesp, pfesp])
    tndof = ndofs(ufesp) + ndofs(pfesp)
    tnunk = nunknowns(ufesp) + nunknowns(pfesp)
    # Assemble the coefficient matrix
    K = assembleK(ufesp, pfesp, tndof, mu)
    # p = spy(K, canvas = DotCanvas)
    # display(p)
    # Solve the system
    U = fill(0.0, tndof)
    gathersysvec!(U, [ufesp, pfesp])
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    scattersysvec!([ufesp, pfesp], U)
    # Postprocessing
    makeattribute(pfesp, "p", 1)
    makeattribute(ufesp, "ux", 1)
    makeattribute(ufesp, "uy", 2)
    ep = evaluate_pressure_error(pfesp)
    ev = evaluate_velocity_error(ufesp)
    vtkwrite("ht_p2_p1_veclap_alt-p", baseincrel(pmesh), [(name = "p",), ])
    vtkwrite("ht_p2_p1_veclap_alt-v", baseincrel(vmesh), [(name = "ux",), (name = "uy",)])
    return (ep, ev)
end

end

using .ht_p2_p1_veclap_alt
let
    N = 4
    for loop in 1:5
        ep, ev = ht_p2_p1_veclap_alt.run(N)
        @show N, ep, ev
        N = N * 2
    end
end


"""Note: Does not (currently) escape all of latex."""
function escape_latex(s::AbstractString)
    s #not actually escaping anything  infact
end

latex_brace(x::AbstractString) = "{" * x * "}"

to_latex(s::MathSymbolInner) = to_latex("$s") #Must define first so will be over writtern
#TODO: Special case Floats, to handle e
to_latex(s::AbstractString) = escape_latex(s)

to_latex(x::MathSymbol) = to_latex(x.sym)

function to_latex(x::MathSuperscript)
    to_latex(x.base) * "^" * latex_brace(to_latex(x.expo))
end

function to_latex(x::MathSubscript)
    to_latex(x.base) * "_" * latex_brace(to_latex(x.sub))
end

function to_latex(x::MathFrac)
    "\\frac" * latex_brace(to_latex(x.numer)) * latex_brace(to_latex(x.denom))
end

function to_latex(x::MathFunctionCall)
    name = x.fname |> string |> escape_latex
    if length(name) > 1
        name = "\\mathrm"*latex_brace(name) #Multicharacter names are typeset upright
    end
    name*to_latex(x.arg)
end

function to_latex(x::MathMatrix)
    desc = "c"^size(x.mat,1)
    rows = mapslices(y->join(y," & "), to_latex.(x.mat),2) |> vec
    latexmat = join(rows, "\\\\")

    """\\left[\\begin{array}{$desc}
    $latexmat
    \\end{array}\\right]"""
end

function to_latex(x::MathSequence)
    inner = join(to_latex.(x.fragments), " ")
    if length(x) == 1 && typeof(first(x.fragments))==MathSequence
       # skip outer-most brackets on Sequence that will be breackets internally anyway
       x.fragments |> first |> to_latex
    else
        "(" * join(to_latex.(x.fragments)," ") * ")"
    end
end


# Operators that are diagonal in a dimension can be solved in O(n) operations
# isdiagop is used to inspect an operator to see if it is diagonal

Base.isdiag{T<:Number}(B::Operator{T}) = bandinds(B)==(0,0)

# multivariate case
isdiagop(K::KroneckerOperator,k) = isdiag(K.ops[k])
isdiagop(S::WrapperOperator,k) = isdiagop(S.op,k)
isdiagop(A::Union{PlusOperator,TimesOperator},k) = all(op->isdiagop(op,k),A.ops)

# diagop gets out the op corresponding to the k-th column
# TODO: diag in x direction

diagop(A::KroneckerOperator,col)=A.ops[2]

diagop(A::PlusOperator,col) = mapreduce(op->diagop(op,col),+,A.ops)
diagop(A::TimesOperator,col) = mapreduce(op->diagop(op,col),*,A.ops)
diagop(A::SpaceOperator,col) = diagop(A.op,col)
# TODO: Space for ConstantOperator
diagop(A::ConstantOperator,col) = ConstantOperator(A.c)
diagop(A::ConstantTimesOperator,col) = A.c*diagop(A.op,col)

diagop(L,k)=error("Override diagop for "*string(typeof(L)))


include("paper/N_plots.jl")
Random.seed!(7)

model, = TrajectoryOptimization.Dynamics.dubinscar!
n, m = model.n,model.m

# Objective
x0 = [2.5;2.5;0.]
xf = [7.5;2.5;0.]
tf =  3.0
Qf = 100.0*Diagonal(I,n)
Q = (1e-3)*Diagonal(I,n)
R = (1e-2)*Diagonal(I,m)

obj = LQRObjective(Q, R, Qf, tf, x0, xf)

r = 0.5
s1 = 30; s2 = 50; s3 = 15
circles = []
for i in range(0,stop=5,length=s1)
    push!(circles,(0.,i,r))
end
for i in range(0,stop=5,length=s1)
    push!(circles,(5.,i,r))
end
for i in range(0,stop=5,length=s1)
    push!(circles,(10.,i,r))
end
for i in range(0,stop=10,length=s2)
    push!(circles,(i,0.,r))
end
for i in range(0,stop=3,length=s3)
    push!(circles,(i,5.,r))
end
for i in range(5,stop=8,length=s3)
    push!(circles,(i,5.,r))
end

n_circles = 3*s1 + s2 + 2*s3

# plt = plot(title="Escape",aspect_ratio=:equal,xlim=[-1,11],ylim=[-1,11])
# plot_obstacles(circles,:red)
# display(plt)

function cI(c,x,u)
    for i = 1:n_circles
        c[i] = circle_constraint(x,circles[i][1],circles[i][2],circles[i][3])
    end
    c
end

u_bnd = 6
obj_con_obstacles = TrajectoryOptimization.ConstrainedObjective(obj,cI=cI,u_min=-u_bnd,u_max=u_bnd)

opts = SolverOptions()
opts.resolve_feasible = false
opts.R_infeasible = 1.
opts.live_plotting = false
opts.verbose = false
N = 101

solver = Solver(model, obj_con_obstacles, integration=:rk3, N=N, opts=opts)
solver_foh = Solver(model, obj_con_obstacles, integration=:rk3_foh, N=N, opts=opts)
solver_dircol = Solver(model, obj_con_obstacles, integration=:rk3_foh, N=N)

# -Initial state and control trajectories
X_guess = [2.5 2.5 0.;4. 5. .785;5. 6.25 0.;7.5 6.25 -.261;9 5. -1.57;7.5 2.5 0.]
X0 = interp_rows(N, tf, Array(X_guess'))
U0 = rand(m,N)

@time res, stats = TrajectoryOptimization.solve(solver,X0,U0)
@time res_foh, stats_foh = TrajectoryOptimization.solve(solver_foh,X0,U0)

solver_dircol.opts.verbose = false
res_d,stats_d = solve_dircol(solver_dircol, X0, U0)

# Escape
plt = plot(title="Escape",aspect_ratio=:equal)
plot_obstacles(circles)
plot_trajectory!(to_array(res.X),width=2,color=:blue,label="rk3",aspect_ratio=:equal,xlim=[-1,11],ylim=[-1,11])
plot_trajectory!(to_array(res_foh.X),width=2,color=:purple,label="rk3_foh",aspect_ratio=:equal,xlim=[-1,11],ylim=[-1,11])
plot!(X0[1,:],X0[2,:],label="Infeasible Initialization",width=1,color=:purple,linestyle=:dash)
plot_trajectory!(res_d.X,width=2,color=:black,label="DIRCOL",linestyle=:dash)

eval_f, = gen_usrfun_ipopt(solver_dircol,:hermite_simpson)
var_zoh = DircolVars(res)
var_foh = DircolVars(res_foh)
eval_f(var_zoh.Z)
eval_f(var_foh.Z)
eval_f(res_d.Z)

@show stats["iterations"]
@show stats_foh["iterations"]

group = "dubinscar/escape"
solver_truth = Solver(model, obj_con_obstacles, N=501, integration=:rk3_foh)
dt_truth = solver_truth.dt
run_dircol_truth(solver_truth, Array(res_d.X), Array(res_d.U), group::String)

run_step_size_comparison(model, obj_con_obstacles, U0, group::String, [101,201,301]; integrations=[:rk3,:rk3_foh],dt_truth=dt_truth,opts=opts,infeasible=true,X0=X0)
plot_stat("runtime",group,legend=:topleft,title="Constrained Parallel Park")
plot_stat("iterations",group,legend=:left,title="Constrained Parallel Park")
plot_stat("error",group,yscale=:log10,legend=:right,title="Constrained Parallel Park")
plot_stat("c_max",group,yscale=:log10,legend=:right,title="Constrained Parallel Park")


# Elementwise addition of two Datasets, assuming same size
function (+)(a::MultiChannelData, b::MultiChannelData)
    if a.channellayout.nobs_per_channel != b.channellayout.nobs_per_channel
        throw(ErrorException("Data observations per channel not equal"))
    elseif a.nsamples != b.nsamples
        throw(ErrorException("Number of samples not equal"))
    else
        vals = [[c1 .+ c2 for (c1, c2) in zip(samp1, samp2)] for (samp1, samp2) in zip(a.values, b.values)]
        out = MultiChannelData(vals, a.channellayout, a.nsamples)
        return out
    end
end

function eigenvalue_ratio_test(evs :: Array{Float64, 1}; threshold=5)
	# assume evs all real, sorted from largest to smallest, strictly positive
	# Could use random matrix theory to pick thresh.
	evratios = [min(ev1/ev2, typemax(Int64)) for (ev1, ev2) in zip(evs[1:(end-1)], evs[2:end])]
	rmax = maximum(evratios)
	if rmax < threshold
		facnum=0
	else
		facnum = argmax(evratios)
	end
	return facnum
end

function LTOrthog(B, blocksize)
    Blq = lq(B[1:blocksize, 1:blocksize])
    # Extract diagonal elems, sort into descending abs val, put back in diag mat.
    Lsigns = diagm(sign.(sort(diag(Blq.L), lt=(x, y) -> abs(y) < abs(x) )))
    Q = transpose(Blq.Q) * Lsigns  # eqn 11
    return Q
end

function extract_horizontal_blocks(mat, blocksizes; get_view=false)
    rbounds = vcat([0], cumsum(blocksizes))
    if get_view
        blocks = [view(mat, start:stop, :) for (start, stop) in zip(rbounds .+ 1, rbounds[2:end])]
    else
        blocks = [mat[start:stop, :] for (start, stop) in zip(rbounds .+ 1, rbounds[2:end])]
    end
    return blocks
end

function extract_diagonal_blocks(mat, rowblocksizes, colblocksizes; get_view=false)
    rbounds = vcat([0], cumsum(rowblocksizes))
    cbounds = vcat([0], cumsum(colblocksizes))
    if get_view
        blocks =  [view( mat, rstart:rstop, cstart:cstop)
            for (rstart, rstop, cstart, cstop) in zip(rbounds .+ 1, rbounds[2:end], cbounds .+ 1, cbounds[2:end])]
    else
        blocks = [mat[rstart:rstop, cstart:cstop]
            for (rstart, rstop, cstart, cstop) in zip(rbounds .+ 1, rbounds[2:end], cbounds .+ 1, cbounds[2:end])]
    end
    return blocks
end

# function extract_blocks(mat, Ls, pjs; diag=false)
#     rbounds = vcat([0], cumsum(Ls))
#     cbounds = vcat([0], cumsum(pjs))
#     if diag
#         # Pull out diagonal blocks
#         blocks = [mat[rstart:rstop, cstart:cstop]
#                     for (rstart, rstop, cstart, cstop) in zip(rbounds .+ 1, rbounds[2:end], cbounds .+ 1, cbounds[2:end])]
#     else
#         # Split rows of matrix into pieces.
#         blocks = [mat[start:stop, :] for (start, stop) in zip(rbounds .+ 1, rbounds[2:end])]
#     end
#     return blocks
# end

function flip_signs(mat)
    diag_signs = sign.(diag(mat))
    flip= mat * diagm(diag_signs)
    return flip
end

function blockdiag(blks; sparse=false)
    if sparse
        throw("unimplemented")
    else
        ncols = [size(blk)[2] for blk in blks]
        nrows = [size(blk)[1] for blk in blks]
        blkmat = zeros(sum(nrows), sum(ncols))
        colbounds = vcat([0], cumsum(ncols))
        rowbounds = vcat([0], cumsum(nrows))
        for (i, blk) in enumerate(blks)
            cstart, cstop = colbounds[i]+ 1, colbounds[i+1]
            rstart, rstop = rowbounds[i]+ 1, rowbounds[i+1]
            blkmat[rstart:rstop, cstart:cstop] = blk
        end
    end
    return blkmat
end

function stack(data :: MultiChannelData)
    X = hcat([vcat(samp...) for samp in data.values ]...)
    @assert size(X) == (sum(data.channellayout.nobs_per_channel), data.nsamples)
    return X
end

function unstack(X :: Matrix, channellayout::MultiChannelLayout; obs_ids = [], sample_ids = [])
    nsamples = size(X)[2]
    values = Array{Array{Array{Float64, 1}, 1}, 1}(undef, nsamples)
    for samp in 1:nsamples
        v = extract_horizontal_blocks(X[:, samp], channellayout.nobs_per_channel)
        values[samp] = [dropdims(vb; dims=2) for vb in v]
    end
    data = MultiChannelData(values, channellayout, nsamples; obs_ids, sample_ids)
    return data
end

#
# function stack_and_view(data)
#     N = length(data)
#     x1, xview1 = stack_and_view_single(data[1])
#     M = length(xview1)
#     L = sum([length(xc) for xc in x1])
#     X = Matrix(zeros((L, N))) # M x N, flat matrix
#     Xobs = repeat([xview1], N) # Obs-first, N x M x L_m
#     Xchan = [[Xobs[n][m] for n in 1:N] for m in 1:M ] # Chan first, M x N x L_m
#     X[:, 1] = x1
#     for n in 2:length(data)
#         #TODO: This doesn't accomplish the desired view, as view is to data
#         # underlying x, which is copied into matrix X. Need view into X itself,
#         # which requires rewriting this.
#         x, xobs = stack_and_view_single(data[n])
#         X[:, n] = x
#         Xobs[n] = xobs
#         for m in 1:M
#             Xchan[m][n] = xobs[m]
#         end
#     end
#     return X, Xobs, Xchan
# end
#
#
# # Probably need to separaete viewing and stacking code.
# function stack_and_view_single(xs)
#     # Combine ragged 3-dim array into matrix + view with struct.
#     lengs = [length(x) for x in xs]
#     bounds = vcat([0], cumsum(lengs))
#     x = vcat(xs...)
#     xview = [view(x, start:stop) for (start, stop) in zip(bounds .+ 1, bounds[2:end])]
#     return x, xview
# end

function data_fromlong(df; should_interpolate=true, max_missingpcent = 0.3 )
	#df must have 4 columns: "sample_id", "channel", "observation_id", "value"
	samp_ids = sort(unique(df.sample_id))
	channels = unique(df.channel)
	chanouts = Array{Array{Float64, 2}, 1}(undef, length(channels))
	nobs_per_channel = zeros(length(channels))
	k = 1
	local included_obs
	for chandf in groupby(df, :channel)
		obs_ids  = sort(unique(chandf.observation_id))
		included_rows = []
		included_obs = []
		chanout = zeros(length(obs_ids), length(samp_ids))
		for (j ,obs) in enumerate(obs_ids)
			obsdf = filter(:observation_id => id -> id == obs, chandf)
			sort!(obsdf, :sample_id)
			ts = obsdf.sample_id
			vals = obsdf.value
			outvals = zeros(length(samp_ids))
			if length(ts) != length(samp_ids)
				if should_interpolate
					if length(ts) <= (1-max_missingpcent)*length(samp_ids)
						# too many missing, should not be included in model.
						print("$obs has too many missing values, removed.\n")
						# Some misingness pattern problem though. What if different
						# numbers of obs for same obs_id are missing in different channels?
						continue
					else
						int = LinearInterpolation(ts, vals, extrapolation_bc=Line())
						interpolated_values = int(samp_ids)
						# outvals is full-length vector, where the observed values
						# are present, and missing values filled in
						for (i, t) in enumerate(samp_ids)
							outvals[i] = t in ts ? vals[findfirst(ts .== t)] : interpolated_values[i]
						end
					end
				else
					throw(ErrorException("Missing observations and not interpolating!"))
				end
			else
				outvals = vals
			end
			push!(included_obs, obs)
			push!(included_rows, j)
			chanout[j, :] = outvals
		end
		nobs_per_channel[k] = length(included_obs)
		chanouts[k] = chanout[included_rows, :]
		k += 1
	end
	stacked_X = cat(chanouts...; dims=1) :: Array{Float64, 2}
	convert(Array{Float64, 2}, stacked_X)
	channellayout = MultiChannelLayout(nobs_per_channel)
	data = unstack(stacked_X, channellayout; obs_ids = included_obs, sample_ids = samp_ids)
	return data
end

function to_long(params::MCFMParams; obs_ids = Int[])
	nchannels = params.channellayout.nchannels
	h_longs = Array{DataFrame, 1}(undef, nchannels)
	g_longs = Array{DataFrame, 1}(undef, nchannels)
	for (c, Hchan) in enumerate(params.H_by_channel)
		df = DataFrame(Hchan)
		rename!(df, ["common" * string(i) for i in 1:ncol(df)])
		df.obs = length(obs_ids) == nrow(df) ? obs_ids : axes(df, 1)
		df_long = DataFrames.stack(df, Not(:obs))
		df_long[!, :channel] .= c
		rename!(df_long, :variable => "factor")
		h_longs[c] = df_long
	end
	for (c, Gchan) in enumerate(params.G_by_channel)
		df = DataFrame(Gchan)
		rename!(df, ["channel$(c)_unique" * string(i) for i in 1:ncol(df)])
		df.obs = length(obs_ids) == nrow(df) ? obs_ids : axes(df, 1)
		df_long = DataFrames.stack(df, Not(:obs))
		df_long[!, :channel] .= c
		rename!(df_long, :variable => "factor")
		g_longs[c] = df_long
	end

	Hlong = vcat(h_longs...)
	Glong = vcat(g_longs...)
	out = vcat(Hlong, Glong)
	return out
end

function to_long(factors :: MultiChannelFactors; sample_ids = Int[])
	factor_names = ["common" * string(i) for i in  1:factors.factorlayout.nchannelshared]
	for (c, nc) in enumerate(factors.factorlayout.nspecific_per_channel)
		factor_names = vcat(factor_names, ["channel$(c)_unique" * string(i) for i in 1:nc])
	end
	df = DataFrame(transpose(factors.factors))
	rename!(df, factor_names)
	df.sample_id = length(sample_ids) == nrow(df) ? sample_ids : axes(df, 1)
	df_long = DataFrames.stack(df, Not(:sample_id))
	rename!(df_long, :variable => "factor")
	return df_long
end

function to_long(data :: MultiChannelData; obs_ids = Int[], sample_ids = Int[])
	X = MultiChannelFactorAnalysis.stack(data)
	nchannels = data.channellayout.nchannels
	chanblocks = extract_horizontal_blocks(X, data.channellayout.nobs_per_channel)
	Xlongs = Array{DataFrame, 1}(undef, nchannels)
	for (c, blk) in enumerate(chanblocks)
		df = DataFrame(transpose(blk))
		df.sample_id = length(sample_ids) == nrow(df) ? sample_ids : axes(df, 1)
		df_long = DataFrames.stack(df, Not(:sample_id))
		df_long[!, :channel] .= c
		rename!(df_long, :variable => "obs_str")
		# Drop the "x" in front of obs_id (stored as Categorical by default.)
		fixed_obs = [parse(Int, levels(x)[1][2:end]) for x in df_long[:obs_str]]
		df_long[!, :obs] = fixed_obs
		if length(obs_ids) == ncol(df) -1 # added sample_id
			# index into the given ids (assuming in order).
			df_long[!, :obs] = obs_ids[df_long[!, :obs]]
		end
		select!(df_long, Not(:obs_str))
		Xlongs[c] = df_long
	end
	out = vcat(Xlongs...)
	return out
end

function parse_df(df)
    #df must have 3 columns: "sample", "channel", "observation"
    # must be sorted appropriately
    nsamples = length(unique(df.sample))
    nchannel = length(unique(df.channel))

    values = Array{Array{Array{Float64, 1}, 1}, 1}(undef, nsamples)
    n = 1
    for sampdf in groupby(df, :sample)
        sample_dat = Array{Array{Float64, 1}, 1}(undef, nchannel)
        i = 1
        for chandf in groupby(sampdf, :channel) # Potential sorting bug here
            obs = copy(chandf.observation)
            sample_dat[i] = obs
            i += 1
        end
        values[n] = sample_dat
        n +=1
    end
    channellayout = MultiChannelLayout([size(c)[1] for c in values[1]])
    data = MultiChannelData(values, channellayout, nsamples)
    return data
end

"""
	varimax(A; gamma = 1.0, minit = 20, maxit = 1000, reltol = 1e-12)
VARIMAX perform varimax (or quartimax, equamax, parsimax) rotation to the column vectors of the input matrix.
# Input Arguments
- `A::Matrix{Float64}`: input matrix, whose column vectors are to be rotated. d, m = size(A).
- `gamma::Float64`: default is 1. gamma = 0, 1, m/2, and d(m - 1)/(d + m - 2), corresponding to quartimax, varimax, equamax, and parsimax.
- `minit::Int`: default is 20. Minimum number of iterations, in case of the stopping criteria fails initially.
- `maxit::Int`: default is 1000. Maximum number of iterations.
- `reltol::Float64`: default is 1e-12. Relative tolerance for stopping criteria.
# Output Argument
- `B::Matrix{Float64}`: output matrix, whose columns are already been rotated.
Implemented by Haotian Li, Aug. 20, 2019
"""
function varimax(A; gamma = 1.0, minit = 20, maxit = 1000, reltol = 1e-12)
	# Get the sizes of input matrix
	d, m = size(A)

	# If there is only one vector, then do nothing.
	if m == 1
		return A
	end

	# Warm up step: start with a good initial orthogonal matrix T by SVD and QR
	T = Matrix{Float64}(I, m, m)
	B = A * T
	L,_,M = svd(A' * (d*B.^3 - gamma*B * Diagonal(sum(B.^2, dims = 1)[:])))
	T = L * M'
	if norm(T-Matrix{Float64}(I, m, m)) < reltol
		T,_ = qr(randn(m,m)).Q
		B = A * T
	end

	# Iteration step: get better T to maximize the objective (as described in Factor Analysis book)
	D = 0
	for k in 1:maxit
		Dold = D
		L,s,M = svd(A' * (d*B.^3 - gamma*B * Diagonal(sum(B.^2, dims = 1)[:])))
		T = L * M'
		D = sum(s)
		B = A * T
		if (abs(D - Dold)/D < reltol) && k >= minit
			break
		end
	end

	# Adjust the sign of each rotated vector such that the maximum absolute value is positive.
	for i in 1:m
		if abs(maximum(B[:,i])) < abs(minimum(B[:,i]))
			B[:,i] .= - B[:,i]
		end
	end

	return B
end


# ---
# title: 667. Beautiful Arrangement II
# id: problem667
# author: Indigo
# date: 2021-06-28
# difficulty: Medium
# categories: Array
# link: <https://leetcode.com/problems/beautiful-arrangement-ii/description/>
# hidden: true
# ---
#
# Given two integers `n` and `k`, you need to construct a list which contains
# `n` different positive integers ranging from `1` to `n` and obeys the
# following requirement:
# Suppose this list is [a1, a2, a3, ... , an], then the list [|a1 \- a2|, |a2 \-
# a3|, |a3 \- a4|, ... , |an-1 \- an|] has exactly `k` distinct integers.
#
# If there are multiple answers, print any of them.
#
# **Example 1:**
#
#
#
#     Input: n = 3, k = 1
#     Output: [1, 2, 3]
#     Explanation: The [1, 2, 3] has three different positive integers ranging from 1 to 3, and the [1, 1] has exactly 1 distinct integer: 1.
#
#
# **Example 2:**
#
#
#
#     Input: n = 3, k = 2
#     Output: [1, 3, 2]
#     Explanation: The [1, 3, 2] has three different positive integers ranging from 1 to 3, and the [2, 1] has exactly 2 distinct integers: 1 and 2.
#
#
# **Note:**
#
#   1. The `n` and `k` are in the range 1 <= k < n <= 104.
#
#
## @lc code=start
using LeetCode

function beautiful_arrangement(n, k)
    res = [1:(n-k);]
    i = res[end] + 1
    j = n
    while i < j
        push!(res, j)
        j -= 1
        push!(res, i)
        i += 1
    end
    i == j && push!(res, i)
    return res
end
## @lc code=end


using .DynamicalSystems
import .DynamicalSystems: lyapunov, lyapunovspectrum

"""
Whip up a DynamicalSystem from a Process
"""
function process2ds(P::Process)
    prob = process2problem(P)
    if getalg(P) isa FunctionMap
        d = DiscreteDynamicalSystem(prob.f, prob.u0, prob.p)
    else
        d = ContinuousDynamicalSystem(prob)
    end
end
export process2ds

"""
Calculate the largest lyapunov exponent of a process
"""
function lyapunov(P::Process)
    d = process2ds(P)
    T = gettmax(P)
    Ttr = gett0(P) - gettransient_t0(P)
    lyapunov(d, T; Ttr)
end
export lyapunov

""" Calculate the lyapunov spectrum"""
function lyapunovspectrum(P::Process, N::Integer=1000, k::Integer=length(getX0(P)))
    d = process2ds(P)
    u0 = getX0(P)
    Ttr = gett0(P) - gettransient_t0(P)
    #dt = getdt(P)
    lyapunovspectrum(d, N, k; u0, Ttr)
end
export lyapunovspectrum

"""Calculate the largest lyapunov exponent as a function of a parameter"""
function lyapunovresponse(P::Process, p, prange, N=1000, k=length(getX0(P)))
    d = process2ds(P)
    λs = Array{Float64, 2}(undef, length(prange), k)
    Threads.@threads for ip ∈ 1:length(prange)
        subP = updateparam(P, p, constantParameter, prange[ip])
        if k == 1 && !(d isa DiscreteDynamicalSystem)
            λ = lyapunov(subP)
        else
            λ = lyapunovspectrum(subP, N, k)
        end
        λs[ip, :] .= λ
    end
    return λs
end
export lyapunovresponse


abstract type GlobalFunctionApproximator end

"""
    fit!(gfa::GlobalFunctionApproximator, dataset_input::AbstractMatrix, dataset_output::AbstractVector)
Fit the global function approximator to the dataset using some optimization method and a chosen
loss function.
"""
function fit! end

"""
    compute_value(gfa::GlobalFunctionApproximator, v::AbstractVector)
Return the value of the function at some query point v, based on the global function approximator

    compute_value(gfa::GlobalFunctionApproximator, v_list::AbstractVector{V}) where V <: AbstractVector{Float64}
Return the value of the function for a list of query points, based on the global function approximator
"""
function compute_value end


export Exponent, show

"""
    Exponent(expo::SortedDict{Variable, Degree})

Define a mathematical exponent, that is a product of `Variable` and conjugated
`Variable`.

### Attributes
- `expo` : a dictionary associating `Variable` objects to a `Degree` objects,
their exponent and the exponent of their conjugates.
- `degree`: a `Degree` object indicating the global degree of the exponent :
the sum of the explicit variables degrees and of the conjugated variables
degrees.

### Exemple
```julia
julia > a, b = Variable("a", Complex), Variable("b", Complex)
julia > Exponent(SortedDict(a=>Degree(1,0), b=>Degree(1,1))) == a*b*conj(b)
```
"""
struct Exponent <: AbstractPolynomial
    expo::SortedDict{Variable, Degree}
    degree::Degree

    function Exponent(expo::SortedDict{Variable, Degree})
        degexpl, degconj = 0,0
        for (var, degree) in expo
            ((degree.explvar < 0) || (degree.conjvar < 0)) && error(LOGGER, "Exponent(): Expected non negative exponent for variable $var (got $degree)")
            (isreal(var) && degree.conjvar != 0) && error(LOGGER, "Exponent(): Expected nul conj exponent for real variable $var (got $degree)")
            (isbool(var) && degree.explvar ∉ SortedSet([0,1])) && error(LOGGER, "Exponent(): Expected boolean exponent for bool $var (got $degree)")
            if degree != Degree(0,0)
                degexpl = max(degexpl, degree.explvar)
                degconj = max(degconj, degree.conjvar)
            else
                delete!(expo, var)
            end
        end
        return new(expo, Degree(degexpl, degconj))
    end
end

#############################
## Constructors
#############################
Exponent() = Exponent(SortedDict{Variable, Degree}())
Exponent(x::Variable) = Exponent(SortedDict{Variable, Degree}(x=>Degree(1,0)))

#############################
## Equality, hash
#############################
function ==(exp1::Exponent, exp2::Exponent)
  (length(exp1.expo) == length(exp2.expo)) || return false
  return isequal(exp1.expo, exp2.expo)
  # for (varname, deg) in exp2.expo
  #   if !haskey(exp1.expo, varname) || exp1.expo[varname] != deg
  #     return false
  #   end
  # end
  # return true
end
!=(exp1::Exponent, exp2::Exponent) = !(exp1 == exp2)
hash(expo::Exponent, h::UInt) = hash(expo.degree, hash(expo.expo, h))

#############################
## Show
#############################
function Base.show(io::IO, exp::Exponent)
  if exp == Exponent()
    print(io, "1")
  else
    expo = exp.expo
    i=length(expo)
    for (var, deg) in expo
      if var.kind <: Complex && deg.conjvar>0
        print(io, "conj(", var, ")")
        if deg.conjvar > 1
          print(io, "^", deg.conjvar)
        end
        if deg.explvar > 0
          print(io, " * ")
        end
      end
      if deg.explvar == 1
        print(io, var)
      elseif deg.explvar > 1
        print(io, var, "^", deg.explvar)
      end
      if i > 1
        print(io, " * ")
      end
      i -= 1
    end
  end
end


"""
    Planar quadrotor
"""

mutable struct Quadrotor2D{T}
    n::Int
    m::Int
    d::Int

    L::T    # length
    mass::T # mass
    J::T    # inertia
    g::T    # gravity
end

function f(model::Quadrotor2D, x, u, w)
    qdd1 = -1.0 * sin(x[3]) / model.mass * (u[1] + u[2])
    qdd2 = -1.0 * model.g + cos(x[3]) / model.mass * (u[1] + u[2])
    qdd3 = model.L / model.J * (-1.0 * u[1] + u[2])

    @SVector [x[4],
              x[5],
              x[6],
              qdd1,
              qdd2,
              qdd3]
end

n, m, d = 6, 2, 0
model = Quadrotor2D(n, m, d, 1.0, 1.0, 1.0, 9.81)


@marginalrule typeof(+)(:in1_in2) (m_out::UnivariateNormalDistributionsFamily, m_in1::UnivariateNormalDistributionsFamily, m_in2::PointMass) = begin
    return (in1 = prod(ProdAnalytical(), NormalMeanVariance(mean(m_out) - mean(m_in2), var(m_out)), m_in1), in2 = m_in2)
end

@marginalrule typeof(+)(:in1_in2) (m_out::UnivariateNormalDistributionsFamily, m_in1::UnivariateNormalDistributionsFamily, m_in2::UnivariateNormalDistributionsFamily) = begin
    xi_out, W_out = weightedmean_precision(m_out)
    xi_in1, W_in1 = weightedmean_precision(m_in1)
    xi_in2, W_in2 = weightedmean_precision(m_in2)

    return MvNormalWeightedMeanPrecision([ xi_in1 + xi_out; xi_in2 + xi_out ], [ W_in1 + W_out W_out; W_out W_in2 + W_out])
end


"""
---
A^{-1}_{22} = A^{22} - A^{21}(A^{11}^{-1})A^{12}
"""
function A_22inv(Ai, n1, nc)
    n2 = size(Ai, 1) - n1
    A22 = Ai[n1+1:end, n1+1:end]
    A12 = Ai[1:n1, n1+1:end]
    A11 = Ai[1:n1, 1:n1]

    v = A11 \ vec(A12[:, 1])
    A12'v
end


using DifferentialEquations
using Plots

μ = 0.15 # Drift
σ = 1.0 # Vol of Vol
θ = 0.04 # long-term variance
κ = 1.0 # Mean reversion speed
ρ = -0.75 # spot-vol correlation
u0 = [100.0, 0.0225]
tspan = (0.0, 1.0)
dt = 1.0/365
prob = HestonProblem(μ, κ, θ, σ, ρ, u0, (0.0, 1.0))

sol = solve(prob, EM(), dt=dt)

print(sol.u[1][1])

path_matrix = hcat(sol.u)

print(typeof(path_matrix))
# print(path_matrix)


using PyCall
# %% Call the python wrapper for CLASS, `classy`, via PyCall
classy = pyimport("classy")
# Create an instance of the CLASS wrapper
cosmo = classy.Class()
# Create a dictionary of the cosmological parameters
deg_ncdm = true # 3 massive neutrinos with degenerate mass?
params = Dict(
   "output" => "mPk",
   "P_k_max_h/Mpc" => 100,
   "z_pk" => 10,
   "non linear" => "halofit",
   "A_s" => 2.097e-9,
   "n_s" => 0.9652,
   "k_pivot" => 0.05,
   "h" => 0.6737,
   "omega_b" => 0.02233,
   "omega_cdm" => 0.1198,
   "N_ncdm" => 1,
)
if deg_ncdm
   push!(params, "m_ncdm" => 0.02, "deg_ncdm" => 3, "N_ur" => 0.00641)
   mν = params["m_ncdm"] * params["deg_ncdm"]
else
   push!(params, "m_ncdm" => 0.06, "N_ur" => 2.0328)
   mν = params["m_ncdm"]
end
h0 = params["h"]
Ωcb = (params["omega_b"] + params["omega_cdm"]) / h0^2
Ων = mν / 93.14 / h0^2 #  Ωνh^2 = ∑mν / (93.14 eV)
Ωm = Ωcb + Ων
# Set the parameters and run the CLASS code
cosmo.set(params)
cosmo.compute()


## Welfare analysis. INCOMPLETE, MISSING THE 3D GRAPHICS [afaik, the best way to make it in Julia is through Generic Mapping Tools (using GMT.jl)

q=(v0_transition = rt[2].v,
q.v0_transition_t3 = rt[3].v,
q.v0_ss = rt[1].v,
q.v0_ss_final = rt[s.N].v,

#Individual-specific gains
ind_welfare = (q.v0_transition./q.v0_ss).^(1/(1-m.γ)) .- 1,
ind_welfare_t3 = (q.v0_transition_t3./q.v0_ss).^(1/(1-m.γ)) .- 1,
measure_0 = sim_0.measure,
measure_end = sim_end.measure,
ind_welfare_avg = sum(q.measure_0.*q.ind_welfare),

#Gains without transition
ind_welfare_notrans = (q.v0_ss_final./q.v0_ss).^(1/(1-m.γ)) .- 1,
ind_welfare_avg_notrans = sum(q.measure_0.*q.ind_welfare_notrans),

#Average gains
agg_welfare_avg = (sum(q.measure_0.*q.v0_transition)/sum(q.measure_0.*q.v0_ss)).^(1/(1-m.γ)) .- 1,

#Gains without transition
agg_welfare_avg_notrans = (sum(q.measure_end.*q.v0_ss_final)/sum(q.measure_0.*q.v0_ss)).^(1/(1-m.γ)) .- 1,
agg_welfare_avg_notrans_fixedmeasure = (sum(q.measure_0.*q.v0_ss_final)/sum(q.measure_0.*q.v0_ss)).^(1/(1-m.γ)) .- 1,
agg_welfare_avg_notrans_fixedvalue = (sum(q.measure_end.*q.v0_ss)/sum(q.measure_0.*q.v0_ss)).^(1/(1-m.γ)) .- 1)

### STATISTICS

# All
q=merge(q,(share_winners = sum(sim_0.measure[q.ind_welfare.>0]),
share_losers = sum(sim_0.measure[q.ind_welfare.<0]),
welfare_avg_winners = sum(sim_0.measure[q.ind_welfare.>=0].*q.ind_welfare[q.ind_welfare.>=0])/ sum(sim_0.measure[q.ind_welfare.>=0]),
welfare_avg_losers = sum(sim_0.measure[q.ind_welfare.<0].*q.ind_welfare[q.ind_welfare.<0])/ sum(sim_0.measure[q.ind_welfare.<0])))

# Exporters vs Non-exporters
q=merge(q,(welfare_avg_exp = sum(sim_0.measure[rt[1].e.==1].*q.ind_welfare[rt[1].e.==1])/sum(sim_0.measure[rt[1].e.==1]),
welfare_avg_nonexp = sum(sim_0.measure[rt[1].e.==0].*q.ind_welfare[rt[1].e.==0])/sum(sim_0.measure[rt[1].e.==0]),

share_losers_exp = sum(sim_0.measure[q.ind_welfare.<0 .& rt[1].e.==1])/ sum(sim_0.measure[rt[1].e.==1]),
share_losers_nonexp = sum(sim_0.measure[q.ind_welfare.<0 .& rt[1].e.==0])/  sum(sim_0.measure[rt[1].e.==0]),
share_winners_exp = sum(sim_0.measure[q.ind_welfare.>=0 .& rt[1].e.==1]) / sum(sim_0.measure[rt[1].e.==1]),
share_winners_nonexp = sum(sim_0.measure[q.ind_welfare.>=0 .& rt[1].e.==0]) / sum(sim_0.measure[rt[1].e.==0]),

welfare_avg_losers_exp=  sum(sim_0.measure[rt[1].e.==1].*q.ind_welfare[rt[1].e.==1].*(q.ind_welfare[rt[1].e.==1].<0))/sum(sim_0.measure[rt[1].e.==1].*(q.ind_welfare[rt[1].e.==1].<0)),
welfare_avg_winners_exp=  sum(sim_0.measure[rt[1].e.==1].*q.ind_welfare[rt[1].e.==1].*(q.ind_welfare[rt[1].e.==1].>=0))/sum(sim_0.measure[rt[1].e.==1].*(q.ind_welfare[rt[1].e.==1].>=0)),
welfare_avg_losers_nonexp=  sum(sim_0.measure[rt[1].e.==0].*q.ind_welfare[rt[1].e.==0].*(q.ind_welfare[rt[1].e.==0].<0))/sum(sim_0.measure[rt[1].e.==0].*(q.ind_welfare[rt[1].e.==0].<0)),
welfare_avg_winners_nonexp=  sum(sim_0.measure[rt[1].e.==0].*q.ind_welfare[rt[1].e.==0].*(q.ind_welfare[rt[1].e.==0].>=0))/sum(sim_0.measure[rt[1].e.==0].*(q.ind_welfare[rt[1].e.==0].>=0))))

# Rich vs Poor
q=merge(q,(measure_a_cum = cumsum(sum(sim_0.measure,dims=2),dims=1),
aP = sum(q.measure_a_cum.<0.99,dims=1) .+1,
# aP = s.a_grid_size./2 .+1; # threshold for asset

welfare_avg_poor = sum(sim_0.measure[1:q.aP-1,:].*q.ind_welfare[1:q.aP-1,:])/sum(sim_0.measure[1:q.aP-1,:]),
welfare_avg_rich = sum(sim_0.measure[q.aP:end,:].*q.ind_welfare[q.aP:end,:])/sum(sim_0.measure[q.aP:end,:]),

share_losers_poor = sum(sim_0.measure[1:q.aP-1,:].*(q.ind_welfare[1:q.aP-1,:].<0))/sum(sim_0.measure[1:q.aP-1,:]),
share_winners_poor = sum(sim_0.measure[1:q.aP-1,:].*(q.ind_welfare[1:q.aP-1,:].>=0))/sum(sim_0.measure[1:q.aP-1,:]),
welfare_avg_losers_poor =  sum(sim_0.measure[1:q.aP-1,:].*q.ind_welfare[1:q.aP-1,:].*(q.ind_welfare[1:q.aP-1,:].<0))/sum(sim_0.measure[1:q.aP-1,:].*(q.ind_welfare[1:q.aP-1,:].<0)),
welfare_avg_winners_poor =  sum(sim_0.measure[1:q.aP-1,:].*q.ind_welfare[1:q.aP-1,:].*(q.ind_welfare[1:q.aP-1,:].>=0))/sum(sim_0.measure[1:q.aP-1,:].*(q.ind_welfare[1:q.aP-1,:].>=0)),

share_losers_rich = sum(sim_0.measure[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].<0))/sum(sim_0.measure[q.aP:end,:]),
share_winners_rich = sum(sim_0.measure[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].>=0))/sum(sim_0.measure[q.aP:end,:]),
welfare_avg_losers_rich = sum(sim_0.measure[q.aP:end,:].*q.ind_welfare[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].<0))/sum(sim_0.measure[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].<0)),
welfare_avg_winners_rich = sum(sim_0.measure[q.aP:end,:].*q.ind_welfare[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].>=0))/sum(sim_0.measure[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].>=0))))

# Productive vs Non-productive
q=merge(q,(measure_z_cum = cumsum(sum(sim_0.measure,dims=1),dims=2),
zP = sum(q.measure_z_cum.<0.5,dims=1).+1,
# zP = s.z_grid_size./2 .+1; # threshold for productive

welfare_avg_NP = sum(sim_0.measure[:,1:q.zP-1].*q.ind_welfare[:,1:q.zP-1])/sum(sim_0.measure[:,1:q.zP-1]),
welfare_avg_P = sum(sim_0.measure[:,q.zP:end].*q.ind_welfare[:,q.zP:end])/sum(sim_0.measure[:,q.zP:end]),
share_losers_NP =  sum(sim_0.measure[:,1:q.zP-1].*(q.ind_welfare[:,1:q.zP-1].<0))/sum(sim_0.measure[:,1:q.zP-1]),
share_winners_NP =  sum(sim_0.measure[:,1:q.zP-1].*(q.ind_welfare[:,1:q.zP-1].>=0))/sum(sim_0.measure[:,1:q.zP-1]),
welfare_avg_winners_NP  = sum(sim_0.measure[:,1:q.zP-1].*q.ind_welfare[:,1:q.zP-1].*(q.ind_welfare[:,1:q.zP-1].>=0))/sum(sim_0.measure[:,1:q.zP-1].*(q.ind_welfare[:,1:q.zP-1].>=0)),
welfare_avg_losers_NP  = sum(sim_0.measure[:,1:q.zP-1].*q.ind_welfare[:,1:q.zP-1].*(q.ind_welfare[:,1:q.zP-1].<0))/sum(sim_0.measure[:,1:q.zP-1].*(q.ind_welfare[:,1:q.zP-1].<0)),
share_winners_P=  sum(sim_0.measure[:,q.zP:end].*(q.ind_welfare[:,q.zP:end].>=0))/sum(sim_0.measure[:,q.zP:end]),
share_losers_P =  sum(sim_0.measure[:,q.zP:end].*(q.ind_welfare[:,q.zP:end].<0))/sum(sim_0.measure[:,q.zP:end]),
welfare_avg_losers_P = sum(sim_0.measure[:,q.zP:end].*q.ind_welfare[:,q.zP:end].*(q.ind_welfare[:,q.zP:end].<0))/sum(sim_0.measure[:,q.zP:end].*(q.ind_welfare[:,q.zP:end].<0)),
welfare_avg_winners_P = sum(sim_0.measure[:,q.zP:end].*q.ind_welfare[:,q.zP:end].*(q.ind_welfare[:,q.zP:end].>=0))/sum(sim_0.measure[:,q.zP:end].*(q.ind_welfare[:,q.zP:end].>=0))))

###  Exporter vs Nonexporters and Rich vs Poor
q=merge(q,(measure_a_cum = cumsum(sum(sim_0.measure,dims=2),dims=1),
aP = sum(q.measure_a_cum .<0.99,dims=1) .+1,
# aP = s.a_grid_size/2+1, # threshold for asset
aVP = 1,

welfare_avg_expR =  sum((rt[1].e[q.aP:end,:].==1).*sim_0.measure[q.aP:end,:].*q.ind_welfare[q.aP:end,:])/sum((rt[1].e[q.aP:end,:].==1).*sim_0.measure[q.aP:end,:]),
welfare_avg_nonexpR =  sum((rt[1].e[q.aP:end,:].==0).*sim_0.measure[q.aP:end,:].*q.ind_welfare[q.aP:end,:])/sum((rt[1].e[q.aP:end,:].==0).*sim_0.measure[q.aP:end,:]),
welfare_avg_expP =  sum((rt[1].e[q.aVP+1:q.aP-1,:].==1).*sim_0.measure[q.aVP+1:q.aP-1,:].*q.ind_welfare[q.aVP+1:q.aP-1,:])/sum((rt[1].e[q.aVP+1:q.aP-1,:].==1).*sim_0.measure[q.aVP+1:q.aP-1,:]),
welfare_avg_nonexpP =  sum((rt[1].e[q.aVP+1:q.aP-1,:].==0).*sim_0.measure[q.aVP+1:q.aP-1,:].*q.ind_welfare[q.aVP+1:q.aP-1,:])/sum((rt[1].e[q.aVP+1:q.aP-1,:].==0).*sim_0.measure[q.aVP+1:q.aP-1,:]),
welfare_avg_expVP = sum((rt[1].e[1:q.aVP ,:].==1).*sim_0.measure[1:q.aVP ,:].*q.ind_welfare[1:q.aVP ,:])/sum((rt[1].e[1:q.aVP ,:].==1).*sim_0.measure[1:q.aVP ,:]),
welfare_avg_nonexpVP = sum((rt[1].e[1:q.aVP ,:].==0).*sim_0.measure[1:q.aVP ,:].*q.ind_welfare[1:q.aVP ,:])/sum((rt[1].e[1:q.aVP ,:].==0).*sim_0.measure[1:q.aVP ,:]),

share_expR = sum((rt[1].e[q.aP:end,:].==1).*sim_0.measure[q.aP:end,:]),
share_expP = sum((rt[1].e[q.aVP+1:q.aP-1,:].==1).*sim_0.measure[q.aVP+1:q.aP-1,:]),
share_nonexpR = sum((rt[1].e[q.aP:end,:].==0).*sim_0.measure[q.aP:end,:]),
share_nonexpP = sum((rt[1].e[q.aVP+1:q.aP-1,:].==0).*sim_0.measure[q.aVP+1:q.aP-1,:]),
share_expVP = sum((rt[1].e[1:q.aVP,:].==1).*sim_0.measure[1:q.aVP,:]),
share_nonexpVP = sum((rt[1].e[1:q.aVP,:].==0).*sim_0.measure[1:q.aVP,:]),

share_losers_expR = sum((q.ind_welfare[q.aP:end,:].<0) .*  (rt[1].e[q.aP:end,:].==1).*sim_0.measure[q.aP:end,:]),
share_winners_expR = sum((q.ind_welfare[q.aP:end,:].>=0) .*  (rt[1].e[q.aP:end,:].==1).*sim_0.measure[q.aP:end,:]),
share_losers_nonexpR = sum( (q.ind_welfare[q.aP:end,:].<0) .*  (rt[1].e[q.aP:end,:].==0).*sim_0.measure[q.aP:end,:]),
share_winners_nonexpR = sum( (q.ind_welfare[q.aP:end,:].>=0) .*  (rt[1].e[q.aP:end,:].==0).*sim_0.measure[q.aP:end,:]),

share_losers_expP = sum( (q.ind_welfare[q.aVP+1:q.aP-1,:].<0) .*  (rt[1].e[q.aVP+1:q.aP-1,:].==1).*sim_0.measure[q.aVP+1:q.aP-1,:]),
share_winners_expP = sum( (q.ind_welfare[q.aVP+1:q.aP-1,:].>=0) .*  (rt[1].e[q.aVP+1:q.aP-1,:].==1).*sim_0.measure[q.aVP+1:q.aP-1,:]),
share_losers_nonexpP = sum( (q.ind_welfare[q.aVP+1:q.aP-1,:].<0) .*  (rt[1].e[q.aVP+1:q.aP-1,:].==0).*sim_0.measure[q.aVP+1:q.aP-1,:]),
share_winners_nonexpP = sum((q.ind_welfare[q.aVP+1:q.aP-1,:].>=0) .*  (rt[1].e[q.aVP+1:q.aP-1,:].==0).*sim_0.measure[q.aVP+1:q.aP-1,:]),

share_losers_expVP = sum((q.ind_welfare[1:q.aVP,:].<0) .*  (rt[1].e[1:q.aVP,:].==1).*sim_0.measure[1:q.aVP,:]),
share_winners_expVP = sum((q.ind_welfare[1:q.aVP,:].>=0) .*  (rt[1].e[1:q.aVP,:].==1).*sim_0.measure[1:q.aVP,:]),
share_losers_nonexpVP = sum((q.ind_welfare[1:q.aVP,:].<0) .*  (rt[1].e[1:q.aVP,:].==0).*sim_0.measure[1:q.aVP,:]),
share_winners_nonexpVP = sum((q.ind_welfare[1:q.aVP,:].>=0) .*  (rt[1].e[1:q.aVP,:].==0).*sim_0.measure[1:q.aVP,:]),

welfare_avg_losers_expR = sum(sim_0.measure[q.aP:end,:] .* q.ind_welfare[q.aP:end,:] .* (q.ind_welfare[q.aP:end,:].<0) .*(rt[1].e[q.aP:end,:].==1))sum(sim_0.measure[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].<0) .*  (rt[1].e[q.aP:end,:].==1)),
welfare_avg_winners_expR = sum(sim_0.measure[q.aP:end,:].*q.ind_welfare[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].>=0) .* (rt[1].e[q.aP:end,:].==1))/sum(sim_0.measure[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].>=0) .* (rt[1].e[q.aP:end,:].==1)),
welfare_avg_losers_nonexpR =sum(sim_0.measure[q.aP:end,:] .* q.ind_welfare[q.aP:end,:] .* (q.ind_welfare[q.aP:end,:].<0) .*  (rt[1].e[q.aP:end,:].==0))/sum(sim_0.measure[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].<0) .*  (rt[1].e[q.aP:end,:].==0)),
welfare_avg_winners_nonexpR = sum(sim_0.measure[q.aP:end,:].*q.ind_welfare[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].>=0) .* (rt[1].e[q.aP:end,:].==0))/sum(sim_0.measure[q.aP:end,:].*(q.ind_welfare[q.aP:end,:].>=0).*(rt[1].e[q.aP:end,:].==0)),

welfare_avg_losers_expP = sum(sim_0.measure[q.aVP+1:q.aP-1,:] .*q.ind_welfare[q.aVP+1:q.aP-1,:].* (q.ind_welfare[q.aVP+1:q.aP-1,:].<0) .*  (rt[1].e[q.aVP+1:q.aP-1,:].==1))/sum(sim_0.measure[q.aVP+1:q.aP-1,:].*(q.ind_welfare[q.aVP+1:q.aP-1,:].<0) .*  (rt[1].e[q.aVP+1:q.aP-1,:].==1) ),
welfare_avg_winners_expP = sum(sim_0.measure[q.aVP+1:q.aP-1,:].*q.ind_welfare[q.aVP+1:q.aP-1,:].*(q.ind_welfare[q.aVP+1:q.aP-1,:].>=0) .* (rt[1].e[q.aVP+1:q.aP-1,:].==1))/sum(sim_0.measure[q.aVP+1:q.aP-1,:].*(q.ind_welfare[q.aVP+1:q.aP-1,:].>=0) .* (rt[1].e[q.aVP+1:q.aP-1,:].==1) ),
welfare_avg_losers_nonexpP = sum(sim_0.measure[q.aVP+1:q.aP-1,:] .* q.ind_welfare[q.aVP+1:q.aP-1,:] .* (q.ind_welfare[q.aVP+1:q.aP-1,:].<0) .*  (rt[1].e[q.aVP+1:q.aP-1,:].==0))/sum(sim_0.measure[q.aVP+1:q.aP-1,:].*(q.ind_welfare[q.aVP+1:q.aP-1,:].<0) .*  (rt[1].e[q.aVP+1:q.aP-1,:].==0) ),
welfare_avg_winners_nonexpP = sum(sim_0.measure[q.aVP+1:q.aP-1,:].*q.ind_welfare[q.aVP+1:q.aP-1,:].*(q.ind_welfare[q.aVP+1:q.aP-1,:].>=0) .* (rt[1].e[q.aVP+1:q.aP-1,:].==0))/sum(sim_0.measure[q.aVP+1:q.aP-1,:].*(q.ind_welfare[q.aVP+1:q.aP-1,:].>=0) .* (rt[1].e[q.aVP+1:q.aP-1,:].==0) ),

welfare_avg_losers_expVP = sum(sim_0.measure[1:q.aVP,:] .* q.ind_welfare[1:q.aVP,:] .* (q.ind_welfare[1:q.aVP,:].<0) .*  (rt[1].e[1:q.aVP,:].==1))/sum(sim_0.measure[1:q.aVP,:].*(q.ind_welfare[1:q.aVP,:].<0) .*  (rt[1].e[1:q.aVP,:].==1)),
q.welfare_avg_winners_expVP = sum(sim_0.measure[1:q.aVP,:].*q.ind_welfare[1:q.aVP,:].*(q.ind_welfare[1:q.aVP,:].>=0) .* (rt[1].e[1:q.aVP,:].==1))/sum(sim_0.measure[1:q.aVP,:].*(q.ind_welfare[1:q.aVP,:].>=0) .* (rt[1].e[1:q.aVP,:].==1)),
q.welfare_avg_losers_nonexpVP =sum(sim_0.measure[1:q.aVP,:] .* q.ind_welfare[1:q.aVP,:] .* (q.ind_welfare[1:q.aVP,:].<0) .*  (rt[1].e[1:q.aVP,:].==0))/sum(sim_0.measure[1:q.aVP,:].*(q.ind_welfare[1:q.aVP,:].<0) .*  (rt[1].e[1:q.aVP,:].==0)),
q.welfare_avg_winners_nonexpVP = sum(sim_0.measure[1:q.aVP,:].*q.ind_welfare[1:q.aVP,:].*(q.ind_welfare[1:q.aVP,:].>=0) .* (rt[1].e[1:q.aVP,:].==0))/sum(sim_0.measure[1:q.aVP,:].*(q.ind_welfare[1:q.aVP,:].>=0) .* (rt[1].e[1:q.aVP,:].==0))))

###  Productive vs Non-productive and Rich vs Poor
q=merge(q,(zP_low = sum(q.measure_z_cum .<0.33,dims=1) .+1,
zP_med = sum(q.measure_z_cum.<0.66,dims=1) .+1,

# shares
share_lowR = sum(sim_0.measure[q.aP:end,1:q.zP_low.-1]),
share_lowP = sum(sim_0.measure[q.aVP.+1:q.aP-1,1:q.zP_low.-1]),
share_lowVP = sum(sim_0.measure[1:q.aVP,1:q.zP_low.-1]),

share_medR = sum(sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1]),
share_medP = sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1]),
share_medVP = sum(sim_0.measure[1:q.aVP,q.zP_low:q.zP_med]),

share_highR = sum(sim_0.measure[q.aP:end,q.zP_med:end]),
share_highP = sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end]),
share_highVP = sum(sim_0.measure[1:q.aVP,q.zP_med:end])))

# welfare
q=merge(q,(welfare_avg_lowR =  sum(sim_0.measure[q.aP:end,1:q.zP_low.-1].*q.ind_welfare[q.aP:end,1:q.zP_low.-1])/sum(sim_0.measure[q.aP:end,1:q.zP_low.-1]),
welfare_avg_lowP =  sum(sim_0.measure[q.aVP.+1:q.aP.-1,1:q.zP_low.-1].*q.ind_welfare[q.aVP.+1:q.aP.-1,1:q.zP_low.-1])/sum(sim_0.measure[q.aVP.+1:q.aP.-1,1:q.zP_low.-1]),
welfare_avg_lowVP =  sum(sim_0.measure[1:q.aVP ,1:q.zP_low.-1].*q.ind_welfare[1:q.aVP ,1:q.zP_low.-1])/sum(sim_0.measure[1:q.aVP ,1:q.zP_low.-1]),

welfare_avg_medR =  sum(sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1].*q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1])/sum(sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1]),
welfare_avg_medP =  sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].*q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1])/sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1]),
welfare_avg_medVP =  sum(sim_0.measure[1:q.aVP ,q.zP_low:q.zP_med.-1].*q.ind_welfare[1:q.aVP ,q.zP_low:q.zP_med.-1])/sum(sim_0.measure[1:q.aVP ,q.zP_low:q.zP_med.-1]),

welfare_avg_highR =  sum(sim_0.measure[q.aP:end,q.zP_med:end].*q.ind_welfare[q.aP:end,q.zP_med:end])/sum(sim_0.measure[q.aP:end,q.zP_med:end]),
welfare_avg_highP =  sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end].*q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_med:end])/sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end]),
welfare_avg_highVP =  sum(sim_0.measure[1:q.aVP ,q.zP_med:end].*q.ind_welfare[1:q.aVP ,q.zP_med:end])/sum(sim_0.measure[1:q.aVP ,q.zP_med:end])))

# share of losers
q=merge(q,(share_losers_lowR = sum((q.ind_welfare[q.aP:end,1:q.zP_low.-1].<0) .* sim_0.measure[q.aP:end,1:q.zP_low.-1]),
share_losers_lowP = sum((q.ind_welfare[q.aVP.+1:q.aP.-1,1:q.zP_low.-1].<0) .* sim_0.measure[q.aVP.+1:q.aP.-1,1:q.zP_low.-1]),
share_losers_lowVP = sum((q.ind_welfare[1:q.aVP,1:q.zP_low.-1].<0) .*  sim_0.measure[1:q.aVP,1:q.zP_low-.1]),

share_losers_medR = sum((q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1].<0) .*  sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1]),
share_losers_medP = sum((q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].<0) .*sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1]),
share_losers_medVP = sum((q.ind_welfare[1:q.aVP,q.zP_low:q.zP_med.-1].<0) .*sim_0.measure[1:q.aVP,q.zP_low:q.zP_med.-1]),

share_losers_highR = sum((q.ind_welfare[q.aP:end,q.zP_med:end].<0) .*  sim_0.measure[q.aP:end,q.zP_med:end]),
share_losers_highP = sum((q.ind_welfare[q.aVP+1:q.aP-1,q.zP_med:end].<0) .* sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end]),
share_losers_highVP = sum((q.ind_welfare[1:q.aVP,q.zP_med:end].<0) .* sim_0.measure[1:q.aVP,q.zP_med:end])))

# share of winners
q=merge(q,(share_winners_lowR = sum((q.ind_welfare[q.aP:end,1:q.zP_low.-1].>=0) .* sim_0.measure[q.aP:end,1:q.zP_low.-1]),
share_winners_lowP = sum((q.ind_welfare[q.aVP.+1:q.aP.-1,1:q.zP_low.-1].>=0) .* sim_0.measure[q.aVP.+1:q.aP.-1,1:q.zP_low.-1]),
share_winners_lowVP = sum((q.ind_welfare[1:q.aVP,1:q.zP_low.-1].>=0) .*  sim_0.measure[1:q.aVP,1:q.zP_low.-1]),

share_winners_medR = sum((q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1].>=0) .*  sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1]),
share_winners_medP = sum((q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].>=0) .*sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1]),
share_winners_medVP = sum((q.ind_welfare[1:q.aVP,q.zP_low:q.zP_med.-1].>=0) .*sim_0.measure[1:q.aVP,q.zP_low:q.zP_med.-1]),

share_winners_highR = sum((q.ind_welfare[q.aP:end,q.zP_med:end].>=0) .*  sim_0.measure[q.aP:end,q.zP_med:end]),
share_winners_highP = sum((q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_med:end].>=0) .* sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end]),
share_winners_highVP = sum((q.ind_welfare[1:q.aVP,q.zP_med:end].>=0) .* sim_0.measure[1:q.aVP,q.zP_med:end])))

# avg welfare of losers
q=merge(q,(welfare_avg_losers_lowR = sum(sim_0.measure[q.aP:end,1:q.zP_low-1] .* q.ind_welfare[q.aP:end,1:q.zP_low-1] .* (q.ind_welfare[q.aP:end,1:q.zP_low-1].<0))/sum(sim_0.measure[q.aP:end,1:q.zP_low-1].*(q.ind_welfare[q.aP:end,1:q.zP_low-1].<0)),
welfare_avg_losers_lowP =sum(sim_0.measure[q.aVP+1:q.aP-1,1:q.zP_low-1] .* q.ind_welfare[q.aVP+1:q.aP-1,1:q.zP_low-1] .* (q.ind_welfare[q.aVP+1:q.aP-1,1:q.zP_low-1].<0))/sum(sim_0.measure[q.aVP+1:q.aP-1,1:q.zP_low-1].*(q.ind_welfare[q.aVP+1:q.aP-1,1:q.zP_low-1].<0)),
welfare_avg_losers_lowVP = sum(sim_0.measure[1:q.aVP,1:q.zP_low-1] .* q.ind_welfare[1:q.aVP,1:q.zP_low-1] .* (q.ind_welfare[1:q.aVP,1:q.zP_low-1].<0))/sum(sim_0.measure[1:q.aVP,1:q.zP_low-1].*(q.ind_welfare[1:q.aVP,1:q.zP_low-1].<0)),

welfare_avg_losers_medR = sum(sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1] .* q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1] .* (q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1].<0))/sum(sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1].*(q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1].<0)),
welfare_avg_losers_medP = sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1] .* q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1] .* (q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].<0))/sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].*(q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].<0)),
welfare_avg_losers_medVP = sum(sim_0.measure[1:q.aVP,q.zP_low:q.zP_med.-1] .* q.ind_welfare[1:q.aVP,q.zP_low:q.zP_med.-1] .* (q.ind_welfare[1:q.aVP,q.zP_low:q.zP_med.-1].<0))/sum(sim_0.measure[1:q.aVP,q.zP_low:q.zP_med.-1].*(q.ind_welfare[1:q.aVP,q.zP_low:q.zP_med.-1].<0)),

welfare_avg_losers_highR = sum(sim_0.measure[q.aP:end,q.zP_med:end] .* q.ind_welfare[q.aP:end,q.zP_med:end] .* (q.ind_welfare[q.aP:end,q.zP_med:end].<0))/sum(sim_0.measure[q.aP:end,q.zP_med:end].*(q.ind_welfare[q.aP:end,q.zP_med:end].<0)),
welfare_avg_losers_highP = sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end] .* q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_med:end] .* (q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_med:end].<0))/sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end].*(q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_med:end].<0)),
welfare_avg_losers_highVP = sum(sim_0.measure[1:q.aVP,q.zP_med:end] .* q.ind_welfare[1:q.aVP,q.zP_med:end] .* (q.ind_welfare[1:q.aVP,q.zP_med:end].<0))/sum(sim_0.measure[1:q.aVP,q.zP_med:end].*(q.ind_welfare[1:q.aVP,q.zP_med:end].<0))))

# avg welfare of winners
q=merge(q,(welfare_avg_winners_lowR = sum(sim_0.measure[q.aP:end,1:q.zP_low.-1] .* q.ind_welfare[q.aP:end,1:q.zP_low.-1] .* (q.ind_welfare[q.aP:end,1:q.zP_low.-1].>=0))/sum(sim_0.measure[q.aP:end,1:q.zP_low.-1].*(q.ind_welfare[q.aP:end,1:q.zP_low.-1].>=0)),
welfare_avg_winners_lowP = sum(sim_0.measure[q.aVP.+1:q.aP.-1,1:q.zP_low.-1] .* q.ind_welfare[q.aVP.+1:q.aP.-1,1:q.zP_low.-1] .* (q.ind_welfare[q.aVP.+1:q.aP.-1,1:q.zP_low.-1].>=0))/sum(sim_0.measure[q.aVP.+1:q.aP.-1,1:q.zP_low.-1].*(q.ind_welfare[q.aVP.+1:q.aP.-1,1:q.zP_low.-1].>=0)),
welfare_avg_winners_lowVP = sum(sim_0.measure[1:q.aVP,1:q.zP_low.-1] .* q.ind_welfare[1:q.aVP,1:q.zP_low.-1] .* (q.ind_welfare[1:q.aVP,1:q.zP_low.-1].>=0))/sum(sim_0.measure[1:q.aVP,1:q.zP_low.-1].*(q.ind_welfare[1:q.aVP,1:q.zP_low.-1].>=0)),

welfare_avg_winners_medR = sum(sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1] .* q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1] .* (q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1].>=0))/sum(sim_0.measure[q.aP:end,q.zP_low:q.zP_med.-1].*(q.ind_welfare[q.aP:end,q.zP_low:q.zP_med.-1].>=0)),
welfare_avg_winners_medP = sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1] .* q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1] .* (q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].>=0))/sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].*(q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_low:q.zP_med.-1].>=0)),
welfare_avg_winners_medVP = sum(sim_0.measure[1:q.aVP,q.zP_low:q.zP_med.-1] .* q.ind_welfare[1:q.aVP,q.zP_low:q.zP_med.-1] .* (q.ind_welfare[1:q.aVP,q.zP_low:q.zP_med.-1].>=0))/sum(sim_0.measure[1:q.aVP,q.zP_low:q.zP_med.-1].*(q.ind_welfare[1:q.aVP,q.zP_low:q.zP_med.-1].>=0)),

welfare_avg_winners_highR = sum(sim_0.measure[q.aP:end,q.zP_med:end] .* q.ind_welfare[q.aP:end,q.zP_med:end] .* (q.ind_welfare[q.aP:end,q.zP_med:end].>=0))/sum(sim_0.measure[q.aP:end,q.zP_med:end].*(q.ind_welfare[q.aP:end,q.zP_med:end].>=0)),
welfare_avg_winners_highP = sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end] .* q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_med:end] .* (q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_med:end].>=0))/sum(sim_0.measure[q.aVP.+1:q.aP.-1,q.zP_med:end].*(q.ind_welfare[q.aVP.+1:q.aP.-1,q.zP_med:end].>=0)),
welfare_avg_winners_highVP = sum(sim_0.measure[1:q.aVP,q.zP_med:end] .* q.ind_welfare[1:q.aVP,q.zP_med:end] .* (q.ind_welfare[1:q.aVP,q.zP_med:end].>=0))/sum(sim_0.measure[1:q.aVP,q.zP_med:end].*(q.ind_welfare[1:q.aVP,q.zP_med:end].>=0))))

q=merge(q,(table_exp=zeros(6,6),table_prod=zeros(6,6)))

### TABLES EXPORT STATUS & NET WORTH
q.table_exp[1,:] = [q.share_nonexpVP q.share_nonexpP q.share_nonexpR q.share_expVP q.share_expP q.share_expR]
q.table_exp[2,:] = [q.welfare_avg_nonexpVP q.welfare_avg_nonexpP q.welfare_avg_nonexpR q.welfare_avg_expVP q.welfare_avg_expP q.welfare_avg_expR]
q.table_exp[3,:] = [q.share_winners_nonexpVP q.share_winners_nonexpP q.share_winners_nonexpR q.share_winners_expVP q.share_winners_expP q.share_winners_expR]
q.table_exp[4,:] = [q.welfare_avg_winners_nonexpVP q.welfare_avg_winners_nonexpP q.welfare_avg_winners_nonexpR q.welfare_avg_winners_expVP q.welfare_avg_winners_expP q.welfare_avg_winners_expR]
q.table_exp[5,:] = [q.share_losers_nonexpVP q.share_losers_nonexpP q.share_losers_nonexpR q.share_losers_expVP q.share_losers_expP q.share_losers_expR]
q.table_exp[6,:] = [q.welfare_avg_losers_nonexpVP q.welfare_avg_losers_nonexpP q.welfare_avg_losers_nonexpR q.welfare_avg_losers_expVP q.welfare_avg_losers_expP q.welfare_avg_losers_expR]

### TABLES PRODUCTIVITY & NET WORTH

q.table_prod[1,:] = [q.share_lowVP q.share_lowP q.share_lowR q.share_medVP q.share_medP q.share_medR q.share_highVP q.share_highP q.share_highR]
q.table_prod[2,:] = [q.welfare_avg_lowVP q.welfare_avg_lowP q.welfare_avg_lowR q.welfare_avg_medVP q.welfare_avg_medP q.welfare_avg_medR q.welfare_avg_highVP q.welfare_avg_highP q.welfare_avg_highR]
q.table_prod[3,:] = [q.share_winners_lowVP q.share_winners_lowP q.share_winners_lowR q.share_winners_medVP q.share_winners_medP q.share_winners_medR q.share_winners_highVP q.share_winners_highP q.share_winners_highR]
q.table_prod[4,:] = [q.welfare_avg_winners_lowVP q.welfare_avg_winners_lowP q.welfare_avg_winners_lowR q.welfare_avg_winners_medVP q.welfare_avg_winners_medP q.welfare_avg_winners_medR q.welfare_avg_winners_highVP q.welfare_avg_winners_highP q.welfare_avg_winners_highR]
q.table_prod[5,:] = [q.share_losers_lowVP q.share_losers_lowP q.share_losers_lowR q.share_losers_medVP q.share_losers_medP q.share_losers_medR q.share_losers_highVP q.share_losers_highP q.share_losers_highR]
q.table_prod[6,:] = [q.welfare_avg_losers_lowVP q.welfare_avg_losers_lowP q.welfare_avg_losers_lowR q.welfare_avg_losers_medVP q.welfare_avg_losers_medP q.welfare_avg_losers_medR q.welfare_avg_losers_highVP q.welfare_avg_losers_highP q.welfare_avg_losers_highR]

### Bins
sav = 0
zbin = 5
abin = 5

q=merge(q,(measure_bins_5=zeros(abin,zbin),measure_bins_5_end=zeros(abin,zbin),avg_welfare_bins_5=zeros(abin,zbin),winners_bins_5=zeros(abin,zbin),losers_bins_5=zeros(abin,zbin)))

for j = 1:zbin
    for i = 1:abin
        pL = (j-1)/zbin*s.z_grid_size+1
        pH = j/zbin*s.z_grid_size
        aL = (i-1)/abin*s.a_grid_size+1
        aH = i/abin*s.a_grid_size

        q.measure_bins_5[i,j] = sum(sim_0.measure[aL:aH,pL:pH]);
        q.measure_bins_5_end[i,j] = sum(sim_end.measure[aL:aH,pL:pH]),
        q.avg_welfare_bins_5[i,j] = sum(sim_0.measure[aL:aH,pL:pH] .* q.ind_welfare[aL:aH,pL:pH])./sum(sim_0.measure[aL:aH,pL:pH]),
        q.winners_bins_5[i,j] = sum(sim_0.measure[aL:aH,pL:pH] .* (q.ind_welfare[aL:aH,pL:pH].>=0))/sum(sim_0.measure[aL:aH,pL:pH]),
        q.losers_bins_5[i,j] = sum(sim_0.measure[aL:aH,pL:pH] .* (q.ind_welfare[aL:aH,pL:pH].<0))/sum(sim_0.measure[aL:aH,pL:pH])
    end
end

#FROM HERE UNTIL "3D bar (10 bins)" PENDING

figure(1)
b1 = bar3(q.measure_bins_5); colorbar
for k = 1:length(b1)
    zdata = b1(k).ZData;
    b1(k).CData = zdata;
    b1(k).FaceColor = 'interp';
end
ylabel('assets')
xlabel('productivity')
# view(32, 22)
if sav == 1
        saveas(gcf,'Figures/measure_FF_tau_c_5','epsc')
        saveas(gcf,'Figures/measure_FF_tau_c_5','pdf')
end

figure(2)
b2 = bar3(q.avg_welfare_bins_5); colorbar
for k = 1:length(b2)
    zdata = b2(k).ZData;
    b2(k).CData = zdata;
    b2(k).FaceColor = 'interp';
end
xlabel('productivity')
ylabel('assets')
# view(32, 22)
if sav == 1
        saveas(gcf,'Figures/welfare_FF_tau_c_5','epsc')
        saveas(gcf,'Figures/welfare_FF_tau_c_5','pdf')
end

figure(3)
b3 = bar3(q.losers_bins_5); colorbar
for k = 1:length(b3)
    zdata = b3(k).ZData;
    b3(k).CData = zdata;
    b3(k).FaceColor = 'interp';
end
xlabel('productivity')
ylabel('assets')
set(gca,'Zlim',[0,1])
caxis([0 1])
# view(32, 22)
if sav == 1
        saveas(gcf,'Figures/welfare_FF_tau_c_5','epsc')
        saveas(gcf,'Figures/welfare_FF_tau_c_5','pdf')
end

figure(4)
imagesc(q.winners_bins_5); colorbar
xlabel('productivity')
ylabel('assets')
# view(32, 22)
if sav == 1
        saveas(gcf,'Figures/winners_FF_tau_c_5','epsc')
        saveas(gcf,'Figures/winners_FF_tau_c_5','pdf')
end

figure(5)
imagesc(q.ind_welfare); colorbar
if sav == 1
        saveas(gcf,'Figures/welfare_image','epsc')
        saveas(gcf,'Figures/welfare_image','pdf')
end

### 3D bar (10 bins)
zbin = 10;
abin = 10;

q=merge(q,(measure_bins_10=zeros(abin,zbin),avg_welfare_bins_10=zeros(abin,zbin),winners_bins_10=zeros(abin,zbin),losers_bins_10=zeros(abin,zbin)))

for j = 1:zbin
    for i = 1:abin
        pL = (j-1)/zbin*s.z_grid_size+1;
        pH = j/zbin*s.z_grid_size;
        aL = (i-1)/abin*s.a_grid_size+1;
        aH = i/abin*s.a_grid_size;

        q.measure_bins_10[i,j] = sum(sim_0.measure[aL:aH,pL:pH])
        q.avg_welfare_bins_10[i,j] = sum(sim_0.measure[aL:aH,pL:pH] .* q.ind_welfare[aL:aH,pL:pH])./sum(sim_0.measure[aL:aH,pL:pH])
        q.winners_bins_10[i,j] = sum(sim_0.measure[aL:aH,pL:pH] .* (q.ind_welfare[aL:aH,pL:pH].>=0))/ sum(sim_0.measure[aL:aH,pL:pH])
        q.losers_bins_10[i,j] = sum(sim_0.measure[aL:aH,pL:pH] .* (q.ind_welfare[aL:aH,pL:pH].<0))/sum(sim_0.measure[aL:aH,pL:pH])
    end
end
if s.graphs_matlab ==1
file=matopen("data_welfaregraphs.mat","w")
write(file,"measure_bins_10",q.measure_bins_10)
close(file)
else

#FROM HERE UNTIL THE END PENDING
figure(5)
b1 = bar3(q.measure_bins_10); colorbar
for k = 1:length(b1)
    zdata = b1(k).ZData;
    b1(k).CData = zdata;
    b1(k).FaceColor = 'interp';
end
ylabel('assets')
xlabel('productivity')
# view(32, 22)
if sav == 1
        saveas(gcf,'Figures/measure_FF_tau_c_10','epsc')
        saveas(gcf,'Figures/measure_FF_tau_c_10','pdf')
end

figure(6)
b2 = bar3(q.avg_welfare_bins_10); colorbar
for k = 1:length(b2)
    zdata = b2(k).ZData;
    b2(k).CData = zdata;
    b2(k).FaceColor = 'interp';
end
xlabel('productivity')
ylabel('assets')
# view(32, 22)
if sav == 1
        saveas(gcf,'Figures/welfare_FF_tau_c_10','epsc')
        saveas(gcf,'Figures/welfare_FF_tau_c_10','pdf')
end

figure(7)
imagesc(q.winners_bins_10); colorbar
# xlabel('productivity')
# ylabel('assets')
# view(32, 22)
if sav == 1
        saveas(gcf,'Figures/winners_FF_tau_c_10','epsc')
        saveas(gcf,'Figures/winners_FF_tau_c_10','pdf')
end

end


let
    while true
        @parameters t σ ρ β
        @variables x(t) y(t) z(t)
        D = Differential(t)

        eqs = [D(D(x)) ~ σ*(y-x) + x^0.000000000000135,
               D(y) ~ x*(ρ-z)-y,
               D(z) ~ x*y - β*z]

        @named sys = ODESystem(eqs)
        sys = structural_simplify(sys)

        u0 = [D(x) => 2.0,
              x => 1.0,
              y => 0.0,
              z => 0.0]

        p  = [σ => 28.0,
              ρ => 10.0,
              β => 8/3]

        tspan = (0.0,100.0)
        prob = ODEProblem(sys,u0,tspan,p,jac=true)

        break
    end
end


function ftau!(champu :: Array{ComplexF64,2},
               champv :: Array{ComplexF64,2},
               m      :: MicMac,
               t      :: Float64,
               fft_u  :: Vector{ComplexF64},
               fft_v  :: Vector{ComplexF64})

    m.v  .= exp.(1im * t * m.A1)
    m.u .= -1im * m.llambda * m.A2 .* conj.(m.v)

    m.v .*= fft_u
    ifft!(m.v,1)
    m.ut .= m.v .* m.matr

    m.v  .= exp.(1im * t * m.A1)
    m.v .*= fft_v
    ifft!(m.v,1)
    m.vt .= m.v .* m.matr

    m.ut .+= conj.(m.vt)
    m.ut .*= 0.5
    m.vt = abs.(m.ut) .^ (2 * m.sigma) .* m.ut

    m.ut .= m.conjmatr .* m.vt
    fft!(m.ut,1)
    m.ut .*= m.u

    m.vt .= m.conjmatr .* conj.(m.vt)
    fft!(m.vt,1)
    m.vt .*= m.u

    transpose!(champu, m.ut)
    transpose!(champv, m.vt)

end

function ftau!(m     :: MicMac,
               t     :: Float64,
               fft_u :: Array{ComplexF64,2},
               fft_v :: Array{ComplexF64,2})

    m.v .= exp.(1im * t * m.A1)

    transpose!(m.ut, fft_u)
    transpose!(m.vt, fft_v)

    m.ut .*= m.v
    m.vt .*= m.v

    ifft!(m.ut,1)
    ifft!(m.vt,1)

    m.ut .*= m.matr
    m.vt .*= m.matr

    m.ut .+= conj.(m.vt)
    m.ut .*= 0.5
    m.vt .= abs.(m.ut) .^ (2 * m.sigma) .* m.ut

    m.u .= -1im * m.llambda .* m.A2 .* conj.(m.v)

    m.ut .= m.conjmatr .* m.vt
    m.vt .= m.conjmatr .* conj.(m.vt)

    fft!(m.ut,1)
    fft!(m.vt,1)

    m.ut .*= m.u
    m.vt .*= m.u

    transpose!(fft_u, m.ut)
    transpose!(fft_v, m.vt)

end


using JuMP, Cbc, NPZ

function Solve_func(year, tol)

	println("Solving weights for $year ...\n\n")

	array = npzread(string(year, "_input.npz"))

	A1 = array["A1"]
	A2 = array["A2"]
	b = array["b"]

	model = Model(Cbc.Optimizer)
	set_optimizer_attribute(model, "logLevel", 1)
	N = size(A1)[2]

	@variable(model, r[1:N] >= 0)
	@variable(model, s[1:N] >= 0)

	@objective(model, Min, sum(r[i] + s[i] for i in 1:N))

	# bound on top by tolerance
	@constraint(model, [i in 1:N], r[i] + s[i] <= tol)

	# Ax = b
	@constraint(model, [i in 1:length(b)], sum(A1[i,j] * r[j] + A2[i,j] * s[j]
		                          for j in 1:N) == b[i])

	optimize!(model)
	termination_status(model)

	r_vec = value.(r)
	s_vec = value.(s)

	npzwrite(string(year, "_output.npz"), Dict("r" => r_vec, "s" => s_vec))

	println("\n")

end

year_list = [x for x in 2014:2030]
tol = 0.70

# Run solver function for all years and tolerances (in order)
for i in year_list
	Solve_func(i, tol)
end


# Test disjoint set

using DataStructures
using Base.Test

s = DisjointSets{Int}(1:10)

@test length(s) == 10
@test num_groups(s) == 10

r = [find_root(s, i) for i in 1 : 10]
@test isa(r, Vector{Int})
@test isequal(r, [1:10])

for i = 1 : 5
    x = 2 * i - 1
    y = 2 * i
    union!(s, x, y)
end

@test length(s) == 10
@test num_groups(s) == 5

r0 = [1, 1, 3, 3, 5, 5, 7, 7, 9, 9]
r = [find_root(s, i) for i in 1 : 10]
@test isa(r, Vector{Int})
@test isequal(r, r0)

union!(s, 1, 4)
union!(s, 3, 5)
union!(s, 7, 9)

@test length(s) == 10
@test num_groups(s) == 2

r0 = [1, 1, 1, 1, 1, 1, 7, 7, 7, 7]
r = [find_root(s, i) for i in 1 : 10]
@test isa(r, Vector{Int})
@test isequal(r, r0)

push!(s, 17)

@test length(s) == 11
@test num_groups(s) == 3

r0 = [1, 1, 1, 1, 1, 1, 7, 7, 7, 7, 11]
r = [find_root(s, i) for i in [1 : 10, 17] ]
@test isa(r, Vector{Int})
@test isequal(r, r0)


using Test
using LinearAlgebra, SparseArrays, Random

using LuxurySparse, StaticArrays
import LuxurySparse: staticize
using StaticArrays: SVector, SMatrix

Random.seed!(2)

@testset "staticize" begin
    # permmatrix
    m = pmrand(ComplexF64, 4)
    sm = m |> staticize
    @test sm isa SPermMatrix{4, ComplexF64}
    @test sm.perm isa SVector
    @test sm.vals isa SVector
    @test sm.perm == m.perm
    @test sm.vals == m.vals
    dm = sm |> dynamicize
    @test dm isa PermMatrix{ComplexF64}
    @test dm.perm isa Vector
    @test dm.vals isa Vector
    @test dm.perm == m.perm
    @test dm.vals == m.vals

    # csc
    m = sprand(ComplexF64, 4,4, 0.5)
    sm = m |> staticize
    @test sm.colptr isa SVector
    @test sm.rowval isa SVector
    @test sm.nzval isa SVector
    @test sm.nzval == m.nzval
    @test sm.rowval == m.rowval
    @test sm.colptr == m.colptr

    dm = sm |> dynamicize
    @test dm.colptr isa Vector
    @test dm.rowval isa Vector
    @test dm.nzval isa Vector
    @test dm.nzval == m.nzval
    @test dm.rowval == m.rowval
    @test dm.colptr == m.colptr

    # diagonal
    m = Diagonal(randn(ComplexF64, 4))
    sm = m |> staticize
    @test sm.diag isa SVector
    @test sm.diag == m.diag
    dm = sm |> dynamicize
    @test dm.diag isa Vector
    @test dm.diag == m.diag

    # dense vector
    m = randn(ComplexF64, 4)
    sm = m |> staticize
    @test sm isa SVector
    @test sm == m
    dm = sm |> dynamicize
    @test dm isa Vector
    @test dm == m

    # dense matrix
    m = randn(ComplexF64, 4, 4)
    sm = m |> staticize
    @test sm isa SMatrix
    @test sm == m
    dm = sm |> dynamicize
    @test dm isa Matrix
    @test dm == m

    m = @SMatrix rand(2, 2)
    @test staticize(m) === m
    m = rand(2, 2)
    @test dynamicize(m) === m
end

@testset "SSparseMatrixCSC" begin
    m = sprand(ComplexF64, 4,4, 0.5)
    sm = m |> staticize
    @test nnz(sm) == nnz(m)
    @test nonzeros(sm) == nonzeros(m)
    @test isdense(sm) == false
    @test issparse(sm)
    @test dropzeros!(sm) == sm
    @test sm == m
end


# abstract types
abstract type PointProcess end;

"""
    DeterminantalPointProcess(L::AbstractMatrix; parallelize=false)
    DeterminantalPointProcess(L::Eigen; parallelize=false)

Given a symmetric matrix `L`, creates a `DeterminantalPointProcess` (DPP). One can also pass the eigen object directly

---

    DeterminantalPointProcess(kernel::Kernel, X::AbstractVector; parallelize=false)
    DeterminantalPointProcess(kernel::Kernel, X::AbstractMatrix; obsdim=1; parallelize=false)

Similar to the basic constructor, will first build the kernel matrix with `kernel` on observations `X`.
If your input is a `AbstractMatrix` you can pass the `obsdim` argument to indicate if rows or columns represent the samples (see docs of [KernelFunctions](https://juliagaussianprocesses.github.io/KernelFunctions.jl/stable/))
"""
struct DeterminantalPointProcess{T,parallelize,TL<:AbstractMatrix{T},Tfact} <: PointProcess
    L::TL
    Lfact::Tfact
    size::Int
end

function DeterminantalPointProcess(L::AbstractMatrix; parallelize=false)
    issymmetric(L) || error("Given matrix is not symmetric")
    Lfact = eigen(L)
    return DeterminantalPointProcess{eltype(L),parallelize,typeof(L),typeof(Lfact)}(L, Lfact, length(Lfact.values))
end

function DeterminantalPointProcess(Lfact::Eigen; parallelize=false)
    L = Symmetric((Lfact.vectors .* Lfact.values') * Lfact.vectors')
    return DeterminantalPointProcess{eltype(L),parallelize,typeof(L),typeof(Lfact)}(L, Lfact, length(Lfact.values))
end

"""
    kDeterminantalPointProcess(k::Int, dpp::DeterminantalPointProcess)

Create a k-DPP where the size of the subsets is fixed to k.
You can also create a k-DPP by calling `dpp(k)`
"""
struct kDeterminantalPointProcess{T,parallelize,Tdpp<:DeterminantalPointProcess{T,parallelize}} <: PointProcess
    k::Int
    dpp::Tdpp
end

(dpp::DeterminantalPointProcess{T,p})(k::Int) where {T,p} = kDPP{T,p,typeof(dpp)}(k, dpp)

# aliases
const DPP = DeterminantalPointProcess
const kDPP = kDeterminantalPointProcess
# const KDPP = KroneckerDeterminantalPointProcess
const MCMCState = Tuple{BitArray{1},Matrix{Float64}}


using Printf

export Grid, State

struct State
    x::Int # x-position
    y::Int # y-position
    function State(x::Int, y::Int)::State
        return new(x, y);
    end
end

@inline Base.:(==)(s1::State, s2::State)::Bool = (s1.x == s2.x) && (s1.y == s2.y);

struct Grid
    dims::Tuple{Int,Int}     # size of the grid
    null_state::State        # terminal state off the grid
    p_transition::Real       # probability of transitioning correctly
    rewards::Matrix{Float64} # reward matrix
    function Grid(dims::Tuple{Int,Int} = (10, 10);
                  null_state::State = State(-1, -1),
                  p_transition::Real = 0.70)::Grid
        # initialize empty rewards matrix
        rewards = zeros(dims)
        return new(dims, null_state, p_transition, rewards);
    end
end

function Base.getindex(g::Grid, i::Int, j::Int)::Float64
    if (1 ≤ i ≤ g.dims[1]) && (1 ≤ j ≤ g.dims[2])
        return g.rewards[i,j];
    end
    throw(BoundsError(g));
end

function Base.setindex!(g::Grid, v::Real, i::Int, j::Int)
    value = Float64(v)
    if (1 ≤ i ≤ g.dims[1]) && (1 ≤ j ≤ g.dims[2])
        g.rewards[i,j] = value;
    else
        throw(BoundsError(g));
    end
end

function Base.show(io::IO, g::Grid)::Nothing
    println(io, "\033[2J") # clear repl
    println(io, "\033[$(displaysize(stdout)[1])A")
    rows, cols = g.dims
    border = "\t" * " "^7 * "-"^(cols * 11)
    println(io, border)
    for i in rows:-1:1
        print("\t")
        for j in 1:cols
            if j == 1
                @printf "%-5i  | %5.1f   |" i g.rewards[j,i]
            else
                @printf "  %5.1f   |" g.rewards[j,i]
            end
        end
        println(io, "\n" * border)
    end
    print(io, "\t" * " "^4)
    for k in 1:cols
        @printf "    %5.0f  " k
    end
end

@inline statespace(g::Grid)::Vector{State} = [[State(x,y) for x=1:g.dims[1], y=1:g.dims[2]]..., g.null_state];
@inline stateindex(g::Grid, s::State)::Int = findall(x->x == s, statespace(g))[1];
@inline inbounds(g::Grid, s::State)::Bool = (1 ≤ s.x ≤ g.dims[1]) && (1 ≤ s.y ≤ g.dims[2]);


@testset "PLS2 Prediction Tests (in sample)" begin

	@testset "Single Column Prediction Test" begin

		X        = MLJ.table([1; 2; 3.0][:,:])
		Y        = MLJ.table([1 1; 2 2; 3.0 3.0])

        pls_pipe       = MLJ.@pipeline prediction_type=:deterministic MLJ.Standardizer() PLSRegressor.PLS(n_factors=1)  target = MLJ.Standardizer()

		pls_machine    = MLJ.machine(pls_pipe, X, Y)

		train = range(1,stop=length(X))
		MLJ.fit!(pls_machine, rows=train,force=true)
		pred = MLJ.predict(pls_machine, rows=train);
		pred = MLJ.matrix(pred)

		@test isequal(round.(pred),[1 1; 2 2; 3 3.0])

	end

	@testset "Linear Prediction Tests " begin

		X        = MLJ.table([1 2; 2 4; 4 6.0])
		Y        = MLJ.table([4 2;6 4;8 6.0])

        pls_pipe       = MLJ.@pipeline prediction_type=:deterministic MLJ.Standardizer() PLSRegressor.PLS(n_factors=2)  target = MLJ.Standardizer()

		pls_machine    = MLJ.machine(pls_pipe, X, Y)

        train = range(1,stop=length(X))
        MLJ.fit!(pls_machine, rows=train,force=true)
        pred = MLJ.predict(pls_machine, rows=train);
		pred = MLJ.matrix(pred)

		@test isequal(round.(pred),[4 2;6 4;8 6.0])

		X           = MLJ.table([1 -2; 2 -4; 4 -6.0])
		Y           = MLJ.table([-4 -2;-6 -4;-8 -6.0])

        pls_pipe       = MLJ.@pipeline prediction_type=:deterministic MLJ.Standardizer() PLSRegressor.PLS(n_factors=2)  target = MLJ.Standardizer()

		pls_machine    = MLJ.machine(pls_pipe, X, Y)

        train = range(1,stop=length(X))
        MLJ.fit!(pls_machine, rows=train,force=true)
        pred = MLJ.predict(pls_machine, rows=train);
		pred = MLJ.matrix(pred)

		@test isequal(round.(pred),[-4 -2;-6 -4;-8 -6.0])

	end

end

@testset "PLS2 Pediction Tests (out of sample)" begin

	@testset "Linear Prediction Tests (Ax + b) | A>0" begin

		X        = MLJ.table([1 2;2 4;3 6;6 12;7 14.0;4 8;5 10.0])
		Y        = MLJ.table([2 2;4 4;6 6;12 12;14 14.0;8 8;10 10.0])

        pls_pipe       = MLJ.@pipeline prediction_type=:deterministic MLJ.Standardizer() PLSRegressor.PLS(n_factors=2)  target = MLJ.Standardizer()

		pls_machine    = MLJ.machine(pls_pipe, X, Y)

		train = range(1,stop=5)
		test  = range(6,stop=7)

        MLJ.fit!(pls_machine, rows=train,force=true)
        pred = MLJ.predict(pls_machine, rows=test);
		pred = MLJ.matrix(pred)

		@test isequal(round.(pred),[8 8;10 10.0])

		X        = MLJ.table([1 2;2 4;3 6;6 12;7 14.0; 4 8;5 10.0])
		Y        = MLJ.table([2 4;4 6;6 8;12 14;14 16.0; 8 10;10 12.0])

    	pls_pipe       = MLJ.@pipeline prediction_type=:deterministic MLJ.Standardizer() PLSRegressor.PLS(n_factors=2)  target = MLJ.Standardizer()

		pls_machine    = MLJ.machine(pls_pipe, X, Y)

		train = range(1,stop=5)
		test  = range(6,stop=7)

        MLJ.fit!(pls_machine, rows=train,force=true)
        pred = MLJ.predict(pls_machine, rows=test);
		pred = MLJ.matrix(pred)

		@test isequal(round.(pred),[8 10;10 12.0])

	end

	@testset "Linear Prediction Tests (Ax + b) | A<0" begin

		X        = MLJ.table([1 -2;2 -4;3 -6;6 -12;7 -14.0; 4 -8;5 -10.0])
		Y        = MLJ.table([2 -2;4 -4;6 -6;12 -12;14 -14.0; 8 -8;10 -10.0])

        pls_pipe       = MLJ.@pipeline prediction_type=:deterministic MLJ.Standardizer() PLSRegressor.PLS(n_factors=2)  target = MLJ.Standardizer()

		pls_machine    = MLJ.machine(pls_pipe, X, Y)

		train = range(1,stop=5)
		test  = range(6,stop=7)

        MLJ.fit!(pls_machine, rows=train,force=true)
        pred = MLJ.predict(pls_machine, rows=test);
		pred = MLJ.matrix(pred)

	    @test isequal(round.(pred),[8 -8;10 -10.0])

		X        = MLJ.table([1 -2;2 -4;3 -6;6 -12;7 -14.0; 4 -8;5 -10.0])
		Y        = MLJ.table([2 -4;4 -6;6 -8;12 -14;14 -16.0; 8 -10;10 -12.0])

        pls_pipe       = MLJ.@pipeline prediction_type=:deterministic MLJ.Standardizer() PLSRegressor.PLS(n_factors=2)  target = MLJ.Standardizer()

		pls_machine    = MLJ.machine(pls_pipe, X, Y)

		train = range(1,stop=5)
		test  = range(6,stop=7)

        MLJ.fit!(pls_machine, rows=train,force=true)
        pred = MLJ.predict(pls_machine, rows=test);
		pred = MLJ.matrix(pred)

		@test isequal(round.(pred),[8 -10;10 -12.0])

	end

end;


# coding: utf-8

# In[1]:

using Random
using Plots

# In[2]:

#bogo, what the heck does it even mean
#i d call it bullocks sort
#it is not really sorting
#it keeps shuffling the list until the list is sorted
#kinda feel like monte carlo simulation
#the time complexity is (n+1)!
#for a large list, we may never wait for the day it gets sorted

# In[3]:

function bogo_sort(arr)

    #deep copy is crucial
    arr_copy=deepcopy(arr)
    sort!(arr_copy)
    counter=0

    while arr_copy!=arr

        shuffle!(arr)
        counter+=1

    end

    return counter

end

# In[4]:

test_arr=rand(0:100,5)

#dont forget to deep copy
test_stats=[bogo_sort(deepcopy(test_arr)) for i in 1:100000]

# In[5]:

#visualize the ridiculous result of bogo sort by histogram
#it can go as far as 1400 times of shuffle to get it sorted from my test
#this is merely a five-element list

# In[6]:

#just out of curiosity
#could this be a power law distribution
histogram(test_stats,bins=10000,
    normed=:true,fill=(:blue, true),
    la=0,legend=:none,
    xlabel="iterations to reach sorted state",
    ylabel="frequency in 100k simulations")


using RecipesBase

# Plot a 2D IntervalBox:
@recipe function f(xx::IntervalBox{2}) #; customcolor = :green, alpha=0.5)

    (x, y) = xx

    alpha --> 0.5
    seriestype := :shape

    x = [x.lo, x.hi, x.hi, x.lo]
    y = [y.lo, y.lo, y.hi, y.hi]

    x, y
end

# Plot a vector of 2D IntervalBoxes:
@recipe function f(v::Vector{T}) where T<:IntervalBox{2}

    seriestype := :shape

    xs = Float64[]
    ys = Float64[]

    # build up coords:  # (alternative: use @series)
    for xx in v
        (x, y) = xx

        # use NaNs to separate
        append!(xs, [x.lo, x.hi, x.hi, x.lo, NaN])
        append!(ys, [y.lo, y.lo, y.hi, y.hi, NaN])

    end

    alpha --> 0.5

    #x = xs
    #y = ys

    #x, y

    xs, ys

end