export i_softmax_crossentropy, i_relu, i_logsumexp

function (_::PlusEq{typeof(argmax)})(out!, x::AbstractArray)
    out! += argmax(x)
    out!, x
end

function (_::MinusEq{typeof(argmax)})(out!, x::AbstractArray)
    out! -= argmax(x)
    out!, x
end

"""
    i_softmax_crossentropy(x, p, imax, xmax, Z, out)

Softmax-Cross entropy function.
"""
@i function i_softmax_crossentropy(x, p, imax, xmax, Z, out::T) where T
    # subtract maximum
    imax += argmax(x)  # trade off space of xmax to time
    xmax += x[imax]
    # accumulate exp(x) to Z, and finally get logZ
    for i=1:length(x)
        x[i] -= xmax
        Z += Base.exp(x[i])
    end
    @routine begin
        yi ← zero(T)
        logZ ← zero(T)
        logZ += log(Z)
    end
    for i=1:length(x)
        yi += logZ
        yi -= x[i]
        out += yi * p[i]
        yi += x[i]
        yi -= logZ
    end
    ~@routine
end

"""
    i_relu(out!, x)

ReLU in machine learning.
"""
@i function i_relu(out!, x)
    @invcheckoff if (x > 0, ~)
        out! += x
    end
end

"""
    i_logsumexp(logout!, out!, xs!, inds!, x)

Compute `logout! = log(sum(exp(x)))`.

# Arguments

    * `out!`, output,
    * `logout!`, logged output,
    * `xs!`, an empty vector to cache the ascending values (same type as `x`),
    * `inds!`, an empty vector to cache the ascending indices (integer type),
    * `x`, input vector.
"""
@i function i_logsumexp(logout!, out!, xs!, inds!, x::AbstractArray{T}) where T
  	i_ascending!(xs!, inds!, x)
    @routine begin
        mx ← zero(T)
        mx += xs![end]
    end
	@invcheckoff @inbounds for i=1:length(x)
		x[i] -= mx
		out! += exp(x[i])
		x[i] += mx
	end
  	logout! += log(out!)
	logout! += mx
    ~@routine
end


##########################
### Points Conversions ###
##########################

"""
Create ECEF coordinates from a given `lla`
"""
function ECEF(lla::LLA, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84)
    ϕdeg, λdeg, h = lla.lat, lla.lon, lla.alt
    d = datum

    sinϕ, cosϕ = sind(ϕdeg), cosd(ϕdeg)
    sinλ, cosλ = sind(λdeg), cosd(λdeg)

    N = d.a / sqrt(1 - d.e² * sinϕ^2)  # Radius of curvature (meters)

    x = (N + h) * cosϕ * cosλ
    y = (N + h) * cosϕ * sinλ
    z = (N * (1 - d.e²) + h) * sinϕ

    return ECEF(x, y, z)
end
"""
Create LLA coordinates from a given `ecef`
"""
function LLA(ecef::ECEF, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84)
    x, y, z = ecef.x, ecef.y, ecef.z
    d = datum

    p = hypot(x, y)
    θ = atan(z*d.a, p*d.b)
    λ = atan(y, x)
    ϕ = atan(z + d.e′² * d.b * sin(θ)^3, p - d.e²*d.a*cos(θ)^3)

    N = d.a / sqrt(1 - d.e² * sin(ϕ)^2)  # Radius of curvature (meters)
    h = p / cos(ϕ) - N

    return LLA(rad2deg(ϕ), rad2deg(λ), h)
end

"""
Create ENU coordinates from given `ecef` coordinates
given a reference point `lla_ref` for linarization
"""
function ENU(ecef::ECEF, lla_ref::LLA, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84)
    ϕdeg, λdeg = lla_ref.lat, lla_ref.lon

    ecef_ref = ECEF(lla_ref, datum)
    ∂x = ecef.x - ecef_ref.x
    ∂y = ecef.y - ecef_ref.y
    ∂z = ecef.z - ecef_ref.z

    # Compute rotation matrix
    sinλ, cosλ = sind(λdeg), cosd(λdeg)
    sinϕ, cosϕ = sind(ϕdeg), cosd(ϕdeg)

    # R = [     -sinλ       cosλ  0.0
    #      -cosλ*sinϕ -sinλ*sinϕ cosϕ
    #       cosλ*cosϕ  sinλ*cosϕ sinϕ]
    #
    # east, north, up = R * [∂x, ∂y, ∂z]
    east  = ∂x * -sinλ      + ∂y * cosλ       + ∂z * 0.0
    north = ∂x * -cosλ*sinϕ + ∂y * -sinλ*sinϕ + ∂z * cosϕ
    up    = ∂x * cosλ*cosϕ  + ∂y * sinλ*cosϕ  + ∂z * sinϕ

    return ENU(east, north, up)
end

"""
Create ENU coordinates from given `ecef` coordinates
given a center of reference point `bounds` for linarization
"""
ENU(ecef::ECEF, bounds::Bounds{LLA}, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) = ENU(ecef, OpenStreetMapX.center(bounds), datum)

"""
Create ECEF coordinates from given `enu` coordinates
and a reference point being center of `bounds` for linearization
"""
function ECEF(enu::ENU, lla_ref::LLA, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84)
    ϕdeg, λdeg = lla_ref.lat, lla_ref.lon

    ecef_ref = ECEF(lla_ref, datum)

    # Compute rotation matrix
    sinλ, cosλ = sind(λdeg), cosd(λdeg)
    sinϕ, cosϕ = sind(ϕdeg), cosd(ϕdeg)

    # R = [-sinλ -sinϕ*cosλ  cosϕ*cosλ
    #      	cosλ -sinϕ*sinλ  cosϕ*sinλ
    #        0.0       cosϕ       sinϕ]
    #
    # x,y,z = R * [east, north, up] + [ecef_ref.x, ecef_ref.y, ecef_ref.z]

    x  = -sinλ*enu.east + -sinϕ*cosλ*enu.north + cosϕ*cosλ*enu.up + ecef_ref.x
    y = cosλ*enu.east +  -sinϕ*sinλ*enu.north + cosϕ*sinλ*enu.up + ecef_ref.y
    z    = 0.0*enu.east + cosϕ*enu.north + sinϕ*enu.up + ecef_ref.z

    return ECEF(x, y, z)
end

"""
Create ECEF coordinates from given `enu` coordinates
and a reference point being center of `bounds` for linearization
"""
ECEF(enu::ENU, bounds::Bounds{LLA}, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) = ECEF(enu, OpenStreetMapX.center(bounds), datum)

"""
Create ENU coordinates from given `lla` coordinates
and a reference point `lla_ref` for linearization
"""
ENU(lla::LLA, lla_ref::LLA, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) = ENU(ECEF(lla, datum), lla_ref, datum)

"""
Create ENU coordinates from given `lla` coordinates
and a reference point being center of `bounds` for linearization
"""
ENU(lla::LLA, bounds::Bounds{LLA}, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) = ENU(ECEF(lla, datum), bounds, datum)

"""
Create LLA coordinates from given `enu` coordinates
and a reference point `lla_ref` for linearization
"""
LLA(enu::ENU, lla_ref::LLA, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) = LLA(ECEF(enu,lla_ref))

"""
Create LLA coordinates from given `enu` coordinates
and a reference point being center of `bounds` for linearization
"""
LLA(enu::ENU, bounds::Bounds{LLA}, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) = LLA(ECEF(enu,bounds))

"""
Converts a dictionary of `LLA` `nodes` into a dictionary of `ECEF` values.
Uses a reference point `lla_ref` for linearization.
"""
function ECEF(nodes::Dict{Int,LLA}, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84)
    r = Dict{Int,ECEF}()
    sizehint!(r, ceil(Int, 1.5*length(nodes)))

    for (key, node) in nodes
        r[key] = ECEF(node, datum)
    end

    return r
end

"""
Converts a dictionary of `ECEF` `nodes` into a dictionary of `LLA` values.
"""
function LLA(nodes::Dict{Int,ECEF}, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84)
    r = Dict{Int,LLA}()
    sizehint!(r, ceil(Int, 1.5*length(nodes)))

    for (key, node) in nodes
        r[key] = LLA(node, datum)
    end

    return r
end

"""
Converts a dictionary of `LLA` and `ECEF` `nodes` into a dictionary of `ENU` values.
Uses a reference point `lla_ref` for linearization.
"""
function ENU(nodes::Dict{Int,T}, lla_ref::LLA,
            datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) where T<:Union{LLA,ECEF}
    r = Dict{Int,ENU}()
    sizehint!(r, ceil(Int, 1.5*length(nodes)))

    for (key, node) in nodes
        r[key] = ENU(node, lla_ref, datum)
    end

    return r
end

"""
Converts a dictionary of `LLA` and `ECEF` `nodes` into a dictionary of `ENU` values.
Uses the center of the given `bounds` for linearization.
"""
ENU(nodes::Dict{Int,T}, bounds::Bounds{LLA}, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) where T<:Union{LLA,ECEF} = ENU(nodes, OpenStreetMapX.center(bounds), datum)

"""
Converts a dictionary of `ENU` `nodes` into a dictionary of `ECEF` values.
Uses a reference point `lla_ref` for linearization.
"""
function ECEF(nodes::Dict{Int,ENU},lla_ref::LLA , datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84)
    r = Dict{Int,ECEF}()
    sizehint!(r, ceil(Int, 1.5*length(nodes)))

    for (key, node) in nodes
        r[key] = ECEF(node, lla_ref, datum)
    end

    return r
end

"""
Converts a dictionary of `ENU` `nodes` into a dictionary of `ECEF` values.
Uses the center of the given `bounds` for linearization.
"""
ECEF(nodes::Dict{Int,ENU}, bounds::Bounds{LLA}, datum::Ellipsoid = WGS84) = ECEF(nodes, OpenStreetMapX.center(bounds), datum)

"""
Converts a dictionary of `ENU` `nodes` into a dictionary of `LLA` values.
Uses a reference point `lla_ref` for linearization.
"""
function LLA(nodes::Dict{Int,ENU},lla_ref::LLA , datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84)
    r = Dict{Int,LLA}()
    sizehint!(r, ceil(Int, 1.5*length(nodes)))

    for (key, node) in nodes
        r[key] = LLA(node, lla_ref, datum)
    end

    return r
end

"""
Converts a dictionary of `ENU` `nodes` into a dictionary of `LLA` values.
Uses the center of the given `bounds` for linearization.
"""
LLA(nodes::Dict{Int,ENU}, bounds::Bounds{LLA}, datum::OpenStreetMapX.Ellipsoid = OpenStreetMapX.WGS84) = LLA(nodes, OpenStreetMapX.center(bounds), datum)


module GraphTheoryMips

using JuMP
using Cbc
#using Graphs
using LightGraphs

include("graphs/structural.jl")
include("graphs/flows.jl")
include("util.jl")

export

max_independent_set,
min_dominating_set,
min_vertex_cover,
distances,
bipartite_set,
bin_decisions

end


__precompile__()

module JQuantmod

export get_price, get_irx, get_sp500, get_log_return, get_log_return_sp500, get_log_irx

using RCall
using Printf

function init()
    R"""
library(quantmod)

get_irx_ratio <- function(from ,to){
    IRX = getSymbols("^IRX", from = from, to=to,
            auto.assign=FALSE)
    IRX_Adj_Close = data.frame(IRX)$IRX.Adjust
    IRX_Annual_Return = (IRX_Adj_Close/100)[2:length(IRX_Adj_Close)]
    IRX_Annual_Log_Return = log(1+IRX_Annual_Return)
    IRX_Daily_Log_Return = IRX_Annual_Log_Return/252
    IRX_Daily_Log_Return
}

get_price <- function(s, to, from){
    INSTR = getSymbols(paste("^",s,sep=""), from = from, to=to,
            auto.assign=FALSE)
    cmd = sprintf("data.frame(INSTR)$%s.Adjusted", s)
    INSTR_Adj_Close=eval(parse(text=cmd))

    INSTR_Adj_Close
}

get_price2 <- function(s, to, from){
    INSTR = getSymbols(s, from = from, to=to,
            auto.assign=FALSE)
    cmd = sprintf("data.frame(INSTR)$%s.Adjusted", s)
    INSTR_Adj_Close=eval(parse(text=cmd))

    INSTR_Adj_Close
}
"""
end

"""
    get_log_return(s::String, from::String, to::String)

    Return the daily log price for symbol `s` from `from` to `to`
"""
function get_log_return(s::String, from::String, to::String)
    init()
    p = get_price(s, from, to)
    return log.(p[2:end])-log.(p[1:end-1])
end

"""
    get_price(s::String, from::String, to::String)

    Return the daily price for symbol `s` from `from` to `to`
"""
function get_price(s::String, from::String, to::String)
    init()
    if s[1]=='^'
        s = (s[2:end])
        reval(@sprintf("r = get_price('%s',from='%s',to='%s')", s, from, to))
    else
        reval(@sprintf("r = get_price2('%s',from='%s',to='%s')", s, from, to))
    end
    @rget r
    r
end

"""
    get_log_irx(from::String, to::String)

    Return the daily risk-free log return from `from` to `to`
"""
function get_log_irx(from::String, to::String)
    init()
    reval(@sprintf("r = get_irx_ratio('%s','%s')", from, to))
    @rget r
    r
end

"""
    get_irx(from::String, to::String)

    Return the daily risk-free return from `from` to `to`
"""
function get_irx(from::String, to::String)
    init()
    r = get_log_irx(from, to)
    return exp.(r) .- 1
end

"""
    get_sp500(from::String, to::String)

    Return the daily return for S&P500 from `from` to `to`
"""
function get_sp500(from::String, to::String)
    init()
    return get_price("^GSPC", from, to)
end

"""
    get_log_return_sp500(from::String, to::String)

    Return the daily log return for S&P500 from `from` to `to`
"""
function get_log_return_sp500(from::String, to::String)
    init()
    return get_log_return("^GSPC", from, to)
end
end


struct Exp{X} <: AbstractFloat
    log::X
end

Base.iszero(::Exp) = false

Base.log(x::Exp) = x.log

function Base.show(io::IO, x::Exp)
    io = IOContext(io, :compact => true)
    print(io, "Exp(", log(x), ")")
end

Base.:*(a::Exp, b::Exp) = Exp(log(a) + log(b))

Base.:/(a::Number, b::Exp) = a * Exp(-log(b))

Base.inv(a::Exp) = Exp(-log(a))

Base.:^(a::Exp, b::Number) = Exp(b * log(a))

Base.:^(a::Exp, b::Integer) = Exp(b * log(a))

Base.:^(a::Exp, b) = Exp(b * log(a))
Base.:^(a::Exp, b::Rational) = Exp(b * log(a))

Base.promote_rule(T::Type{<:Real}, ::Type{Exp}) = T
Base.promote_rule(::Type{Exp}, T::Type{<:Real}) = T

Base.convert(::Type{<:Real}, x::Exp) = exp(log(x))

# Exp(3) / Exp(2)

# Exp(3) * Exp(2)

# 3 + Exp(2)

# 3 / Exp(4)

# Exp(2) ^ 3

# 1 / Exp(2)

# 2 + Exp(3) ^ 0


using PkgSRA

"""
Steps 1 and 3: fit a distribution of models

calc_distribution_of_models(dat,
                             numerical_grad,
                             sindy_library;
                             U=nothing,
                             selection_dist=Normal(0,10),
                             use_clustering_minimization=false,
                             chain_opt=(iterations=200,
                                        num_training_pts=100,
                                        start_ind=101),
                                        val_list = Iterators.product(1:3,1:3))

Returns both the Turing chain and the sindycModel used as the
        priors for the coefficients.
"""
function calc_distribution_of_models(dat,
                                     numerical_grad,
                                     sindy_library;
                                     U=nothing,
                                     selection_dist=Normal(0,10),
                                     use_clustering_minimization=false,
                                     chain_opt=(iterations=200,
                                                num_training_pts=100,
                                                start_ind=101),
                                                val_list = Iterators.product(1:3,1:3)) # DEBUG
   # Use AIC to find the best model
   (best_sindy_model,best_criterion,all_criteria,all_models, best_index) =
       sindyc_ensemble(dat, numerical_grad, sindy_library, val_list,
                       U=U,
                       selection_criterion=my_aicc,
                       sparsification_mode="num_terms",
                       selection_dist=selection_dist,
                       use_clustering_minimization=use_clustering_minimization)

   # Use Turing to fit a distribution of models
   turing_model = convert_sindy_to_turing_enforce_zeros(best_sindy_model;
                                   dat_noise_prior=Normal(0.0, 20.0),
                                   coef_noise_std=1.0)
   chain = generate_chain(dat, numerical_grad,
                          turing_model;
                          chain_opt...)[1]

   return (chain=chain, best_sindy_model=best_sindy_model)
end

"""
Step 2: Get the distribution of residuals

calc_distribution_of_residuals(dat, numerical_grad, chain,
                                sample_ind,
                                sindy_template)

Returns:
        residual
        sample_gradients
        sample_noise
        dat_grad
"""
function calc_distribution_of_residuals(dat, numerical_grad, chain,
                                        sample_ind,
                                        sindy_template)
    sample_gradients3d, sample_noise2d = sample_sindy_posterior_grad(chain,
                dat, sample_ind, sindy_template)

    # Calculate the residuals
    sample_gradients = transpose(drop_all_1dims(mean(sample_gradients3d, dims=1)))
        # sample_gradients_std = transpose(drop_all_1dims(std(sample_gradients3d, dims=1)))
        sample_noise = mean(sample_noise2d)
    dat_grad = numerical_grad[:, sample_ind]

    residual = dat_grad .- sample_gradients

    return (residual=residual,
            sample_gradients=sample_gradients,
            sample_noise=sample_noise,
            dat_grad=dat_grad)
end

"""
Step 3: Subsample the data using the residuals

function subsample_using_residual(residual, noise;
                                  noise_factor=1.0,
                                  min_length=1,
                                  shorten_length=1)
return accepted_ind
"""
function subsample_using_residual(residual, noise;
                                  noise_factor=1.0,
                                  min_length=1,
                                  shorten_length=1)
    partial_accepted_ind = abs.(residual) .< (noise_factor*noise)
    partial_accepted_ind = vec(prod(Int.(partial_accepted_ind), dims=1))
    _, _, all_starts, all_ends = calc_contiguous_blocks(
            partial_accepted_ind,
            minimum_length=min_length + 2*shorten_length)

    # Get the "good" indices
    accepted_ind = []
    for (s, e) in zip(all_starts, all_ends)
        accepted_ind = vcat(accepted_ind,
                (s+shorten_length):(e-shorten_length))
    end
    return accepted_ind
end

"""
Optional step 0: Randomly subsample the initial data (for the naive model)

calc_best_random_subsample(model_template, dat2, numerical_grad2, sindy_library;
                                    num_pts=400,
                                    num_subsamples=10,
                                    val_list = Iterators.product(1:3,1:3),
                                    sindyc_ensemble_params=get_sindyc_ensemble_parameters(),
                                    use_control=true,
                                    max_error=nothing)
return (best_initial_subsample=best_initial_subsample,
        best_model=all_final_models2[tmp_ind],
        all_initial_subsamples=initial_subsamples,
        all_errL2=all_errL2)
"""
function calc_best_random_subsample(model_template, dat2,
                        numerical_grad2;
                        num_pts=400,
                        num_subsamples=20,
                        val_list = Iterators.product(1:3,1:3),
                        sindyc_ensemble_params=get_sindyc_ensemble_parameters(),
                        use_control=true,
                        max_error=0)
    # Initially, randomly subsample the data
    initial_subsamples = []
    for i in 1:num_subsamples
        push!(initial_subsamples, randperm(size(dat2,2))[1:num_pts])
    end

    println("Calculating best random subsample; control=$use_control)")

    # Just do SINDY here, not Turing yet
    sz = size(initial_subsamples)
    all_final_models2 = Vector{DynamicalSystemModel}(undef,sz)
    all_errL2 = zeros(sz)
    current_error = max_error+1
    max_iter = 2
    this_iter = 1
    if use_control; f = sindyc_ensemble; else; f = sindy_ensemble end
    while current_error > max_error
        if this_iter > 1
            println("Random subsampling gave a large error ($current_error), retrying")
        end
        for (i, inds) in enumerate(initial_subsamples)
            i%10==0 && println("Testing random subsample $i/$(length(initial_subsamples))")
            (all_final_models2[i], all_errL2[i], _, _, _) =
                 f(model_template, dat2[:,inds], numerical_grad2[:, inds],
                    val_list; sindyc_ensemble_params...)
            # if sindyc_ensemble_params[:selection_criterion] == my_aicc
            #     # Calculate real L2 error
            #     this_dat = all_final_models2[i](dat2, 0)
            #     all_errL2[i] = sum(abs, numerical_grad2.-this_dat)
            # else
            #     # Should be cross validation; already L2 error
            # end
        end
        current_error = minimum(all_errL2)
        this_iter += 1
        if max_error == 0 || this_iter >= max_iter
            break
        end
    end

    # Choose the best of the above models and subsamples
    min_err, tmp_ind = findmin(all_errL2)
    best_initial_subsample = initial_subsamples[tmp_ind]

    return (best_initial_subsample=best_initial_subsample,
            best_model=all_final_models2[tmp_ind],
            all_initial_subsamples=initial_subsamples,
            all_errL2=all_errL2)
end

#####
#####
#####
export calc_distribution_of_models,
        calc_distribution_of_residuals,
        subsample_using_residual,
        calc_best_random_subsample


####################
# This example demonstrates a BandedMatrix on the GPU using CLArrays
# We construct the Matrix on the CPU first for now as this is currently
# slow when tried to do directly on the GPU.
####################

using GPUArrays, CuArrays, FillArrays, BandedMatrices
import BandedMatrices: _BandedMatrix

function cu_finitedifference(n) where T
    data = CuArray{Float32}(3, n)
    data[2,:] .= -2f0 * n^2
    data[[1,3],:] = Δt/Δx^2
    _BandedMatrix(convert(T, data), n, 1, 1)
end

function expliciteuler(L, u₀, n)
    u = copy(u₀)
    v = copy(u)
    for _ = 1:n
        mul!(v, L, u)
        u,v = v,u
    end
    u
end

####
# 32 bit
####
nₓ = 20_000
Δx = 2f0/nₓ
Δt = Δx^2/8f0
T = 1f0 / 1_000
n_t = floor(Int, T/Δt)
L = finitedifference(CLArray, nₓ, Δt)

x= range(-1, stop=1, length=nₓ)
u₀ = CLArray{Float32}(exp.(-10x.^2));

@time begin
    u = expliciteuler(L, u₀, n_t);
    GPUArrays.synchronize(u)
end # 25s

L̃ = BandedMatrix{Float32,Matrix{Float32}}(L)
ũ₀ = Vector{Float32}(u₀)
@time ũ = expliciteuler(L̃, ũ₀, n_t); # 44s

maximum(abs, Array(u)-ũ)  # 3E-3

####
# 64 bit
####
nₓ = 20_000
Δx = 2/nₓ
Δt = Δx^2/8
T = 1f0 / 1_000
n_t = floor(Int, T/Δt)
L = finitedifference(CLArray, nₓ, Δt)

x= range(-1, stop=1, length=nₓ)
u₀ = CLArray(exp.(-10x.^2));

@time begin
    u = expliciteuler(L, u₀, n_t);
    GPUArrays.synchronize(u)
end # 30s

L̃ = BandedMatrix{Float64,Matrix{Float64}}(L)
ũ₀ = Vector{Float64}(u₀)
@time ũ = expliciteuler(L̃, ũ₀, n_t); # 45s

maximum(abs, Array(u)-ũ)  # 1.72E-13


"""
    symplectic_integrate(x₀, p₀, Λ, U, δUδx, N=50, ϵ=0.1, progress=false)

Do a symplectic integration of the potential energy `U` (with gradient `δUδx`)
starting from point `x₀` with momentum `p₀` and mass matrix `Λ`. The number of
steps is `N` and the step size `ϵ`.

Returns `ΔH, xᵢ, pᵢ` corresponding to change in Hamiltonian, and final position
and momenta. If `hist` is specified a trace of requested variables throughout
each step is also returned.

"""
function symplectic_integrate(
    x₀::AbstractVector{T}, p₀, Λ, U, δUδx=x->@ondemand(Zygote.gradient)(U,x)[1];
    N=50, ϵ=T(0.1), progress=false, hist=nothing) where {T}

    xᵢ, pᵢ = x₀, p₀
    δUδxᵢ = δUδx(xᵢ)
    H(x,p) = U(x) - p⋅(Λ\p)/2

    _hist = []

    @showprogress (progress ? 1 : Inf) "Symplectic Integration: " for i=1:N
        xᵢ₊₁    = xᵢ - T(ϵ) * (Λ \ (pᵢ - T(ϵ)/2 * δUδxᵢ))
        δUδx₊₁  = δUδx(xᵢ₊₁)
        pᵢ₊₁    = pᵢ - T(ϵ)/2 * (δUδx₊₁ + δUδxᵢ)
        xᵢ, pᵢ, δUδxᵢ = xᵢ₊₁, pᵢ₊₁, δUδx₊₁

        if hist!=nothing
            histᵢ = (i=i, x=xᵢ, p=pᵢ, δUδx=δUδx₊₁, H=(((:H in hist) ? H(xᵢ,pᵢ) : nothing)))
            push!(_hist, getfield.(Ref(histᵢ),hist))
        end

    end

    ΔH = H(xᵢ,pᵢ) - H(x₀,p₀)

    if hist == nothing
        return ΔH, xᵢ, pᵢ
    else
        return ΔH, xᵢ, pᵢ, _hist
    end

end

@doc doc"""
    grid_and_sample(lnP::Function; range::NamedTuple; progress=false, nsamples=1)

Interpolate the log pdf `lnP` with support on `range`, and return  the
integrated log pdf as well `nsamples` samples (drawn via inverse transform
sampling)

`lnP` should accept a NamedTuple argument and `range` should be a NamedTuple mapping
those same names to `range` objects specifying where to evaluate `lnP`, e.g.:

```julia
grid_and_sample(nt->-(nt.x^2+nt.y^2)/2, (x=range(-3,3,length=100),y=range(-3,3,length=100)))
```

The return value is `(lnP, samples, Px)` where `lnP` is an interpolated/smoothed
log PDF which can be evaluated anywhere within the original range, `Px` are
sampled points of the original PDF, and `samples` is a NamedTuple giving the
Monte-Carlo samples of each of the parameters.

(Note: only 1D sampling is currently implemented, but 2D like in the example
above is planned)
"""
function grid_and_sample(lnP::Function, range::NamedTuple{S, <:NTuple{1}}; progress=false, nsamples=1, span=0.25, rtol=1e-5) where {S}

    xs = first(range)
    xmin,xmax = first(xs),last(xs)

    # probe the pdf along the grid and interpolate
    lnPs = @showprogress (progress ? 1 : Inf) "Grid Sample: " map(x->Float64(lnP(NamedTuple{S}(x))), xs)
    lnPs .-= maximum(lnPs)
    ilnP = loess(xs,lnPs,span=span)

    # normalize the PDF. note the smoothing is done of the log PDF.
    A = @ondemand(QuadGK.quadgk)(exp∘ilnP, xmin, xmax)[1]
    lnPs .-= log(A)
    ilnP = loess(xs,lnPs,span=span)

    # draw samples via inverse transform sampling
    # (the `+ eps()`` is a workaround since Loess.predict seems to NaN sometimes when
    # evaluated right at the lower bound)
    θsamples = NamedTuple{S}(@showprogress (progress ? 1 : Inf) map(1:nsamples) do i
        r = rand()
        fzero((x->@ondemand(QuadGK.quadgk)(exp∘ilnP,xmin+sqrt(eps()),x,rtol=rtol)[1]-r),xmin+sqrt(eps()),xmax,rtol=rtol)
    end)

    if nsamples==1
        ilnP, map(first, θsamples), lnPs
    else
        ilnP, θsamples, lnPs
    end

end
grid_and_sample(lnP::Function, range, progress=false) = error("Can only currently sample from 1D distributions.")
# allow more conveniently evaluation of Loess-interpolated functions
(m::Loess.LoessModel)(x) = Loess.predict(m,x)

@doc doc"""
    sample_joint(ds::DataSet; kwargs...)

Sample the joint posterior, $\mathcal{P}(f,\phi,\theta\,|\,d)$.

Keyword arguments:

* `nsamps_per_chain` — *(required)* The number of samples per chain
* `nchains` — Run `nchains` chains in parallel *(default: 1)*
* `nchunk` — Do `nchunk` steps between parallel chain communication *(default: 1)*
* `nsavemaps` — Save maps into chain every `nsavemaps` steps *(default: 1)*
* `nburnin_always_accept` — The first `nburnin_always_accept` steps, always accept
                            HMC steps independent of integration error *(default: 0)*
* `nburnin_fixθ` — For the first `nburnin_fixθ` steps, fix θ at its starting point *(default: 0)*
* `Nϕ` — Noise to use in the HMC mass matrix. can also give `Nϕ=:qe` to use the
         EB quadratic estimate noise *(default: `:qe`)*
* `chains` — `nothing` to start a new chain; the return value from a previous call to
             `sample_joint` to resume those chains; `:resume` to resume chains
             from a file given by `filename`
* `θrange` — Range and density to grid sample parameters as a NamedTuple,
             e.g. `(Aϕ=range(0.7,1.3,length=20),)`.
* `θstart` — Starting values of parameters as a NamedTuple, e.g. `(Aϕ=1.2,)`,
             or nothing to randomly sample from θrange
* `ϕstart` — Starting ϕ, either a `Field` object, `:quasi_sample`, or `:best_fit`
* `metadata` — Does nothing, but is saved into the chain file
* `nhmc` — The number of HMC passes per ϕ Gibbs step *(default: 1)*
* `symp_kwargs` — an array of NamedTupe kwargs to pass to [`symplectic_integrate`](@ref).
                  E.g. `[(N=50,ϵ=0.1),(N=25,ϵ=0.01)]` would do 50 large steps then 25
                  smaller steps per each Gibbs pass. If specified, `nhmc` is ignored.
* `wf_kwargs` — Keyword arguments to pass to [`argmaxf_lnP`](@ref) in the Wiener Filter Gibbs step.
* `MAP_kwargs` — Keyword arguments to pass to [`MAP_joint`](@ref) when computing the starting point.
"""
function sample_joint(
    ds :: DataSet{<:FlatField{P,T,M}};
    nsamps_per_chain,
    nchains = nworkers(),
    nchunk = 1,
    nsavemaps = 1,
    nburnin_always_accept = 0,
    nburnin_fixθ = 0,
    Nϕ = :qe,
    chains = nothing,
    ϕstart = 0,
    θrange = NamedTuple(),
    θstart = nothing,
    Nϕ_fac = 2,
    pmap = (myid() in workers() ? map : pmap),
    wf_kwargs = (tol=1e-1, nsteps=500),
    nhmc = 1,
    symp_kwargs = fill((N=25, ϵ=0.01), nhmc),
    MAP_kwargs = (αmax=0.3, nsteps=40),
    metadata = nothing,
    progress = false,
    interruptable = false,
    gibbs_pass_θ::Union{Function,Nothing} = nothing,
    filename = nothing,
    storage = basetype(M)
    ) where {P,T,M}

    # save input configuration to chain
    rundat = adapt(Array, Base.@locals)

    @assert (length(θrange) in [0,1]) || (gibbs_pass_θ!=nothing) "Can only currently sample one parameter at a time, otherwise must pass custom `gibbs_pass_θ`"
    @assert progress in [false,:summary,:verbose]

    @unpack L, Cϕ = ds

    if (chains == nothing)
        if (θstart == nothing)
            θstarts = [map(range->(first(range) + rand() * (last(range) - first(range))), θrange) for i=1:nchains]
        else
            @assert θstart isa NamedTuple "θstart should be either `nothing` to randomly sample the starting value or a NamedTuple giving the starting point."
            θstarts = fill(θstart, nchains)
        end
        if (ϕstart == 0); ϕstart = zero(ds().Cϕ.diag); end
        @assert ϕstart isa Field || ϕstart in [:quasi_sample, :best_fit]
        if ϕstart isa Field
            ϕstarts = fill(ϕstart, nchains)
        else
            ϕstarts = pmap(θstarts) do θstart
                Random.seed!()
                MAP_joint(ds(;θstart...), progress=(progress==:verbose ? :summary : false), Nϕ=Nϕ, quasi_sample=(ϕstart==:quasi_sample); MAP_kwargs...)[2]
            end
        end
        chains = pmap(θstarts,ϕstarts) do θstart,ϕstart
            Any[@dictpack i=>1 f=>nothing ϕ°=>adapt(Array,ds(;θstart...).G*ϕstart) θ=>θstart seed=>deepcopy(Random.seed!())]
        end
    elseif chains == :resume
        chains = @ondemand(FileIO.load)(filename,"chains")
    elseif chains isa String
        chains = @ondemand(FileIO.load)(chains,"chains")
    end

    if (Nϕ == :qe); Nϕ = quadratic_estimate(ds()).Nϕ / Nϕ_fac; end

    swap_filename = (filename == nothing) ? nothing : joinpath(dirname(filename), ".swap.$(basename(filename))")

    dsₐ,Nϕₐ = ds,Nϕ

    # start chains
    try
        @showprogress (progress==:summary ? 1 : Inf) "Gibbs chain: " for _=1:nsamps_per_chain÷nchunk

            append!.(chains, pmap(last.(chains)) do state

                @unpack i,ϕ°,f,θ,seed = state
                f,ϕ°,ds,Nϕ = (adapt(storage, x) for x in (f,ϕ°,dsₐ,Nϕₐ))
                copy!(Random.GLOBAL_RNG, seed)
                dsθ = ds(;θ...)
                ϕ = dsθ.G\ϕ°
                pϕ°, ΔH, accept = nothing, nothing, nothing
                L = ds.L
                lnPθ = nothing
                chain = []

                for i=(i+1):(i+nchunk)

                    # ==== gibbs P(f°|ϕ°,θ) ====
                    t_f = @elapsed begin
                        Lϕ = cache(L(ϕ), ds.d)
                        f° = Lϕ * dsθ.D * argmaxf_lnP(Lϕ, dsθ; which=:sample, guess=f, progress=(progress==:verbose), wf_kwargs...)
                    end

                    # ==== gibbs P(ϕ°|f°,θ) ====
                    t_ϕ = @elapsed begin

                        Λm = (Nϕ == nothing) ? pinv(dsθ.Cϕ) : (pinv(dsθ.Cϕ) + pinv(Nϕ))

                        for kwargs in symp_kwargs

                            pϕ° = simulate(Λm)
                            (ΔH, ϕtest°) = symplectic_integrate(
                                ϕ°, pϕ°, Λm,
                                ϕ°->lnP(:mix, f°, ϕ°, θ, dsθ);
                                progress=(progress==:verbose),
                                kwargs...
                            )

                            if (i < nburnin_always_accept) || (log(rand()) < ΔH)
                                ϕ° = ϕtest°
                                accept = true
                            else
                                accept = false
                            end

                        end

                    end

                    # ==== gibbs P(θ|f°,ϕ°) ====
                    t_θ = @elapsed begin
                        if (i > nburnin_fixθ && length(θrange)>0)
                            if gibbs_pass_θ == nothing
                                lnPθ, θ = grid_and_sample(θ->lnP(:mix,f°,ϕ°,θ,ds), θrange, progress=(progress==:verbose))
                            else
                                lnPθ, θ = gibbs_pass_θ(;(Base.@locals)...)
                            end
                            dsθ = ds(;θ...)
                        end
                    end

                    # compute un-mixed maps
                    f,ϕ = unmix(f°,ϕ°,θ,ds)
                    f̃ = L(ϕ)*f

                    # save state to chain and print progress
                    timing = (f=t_f, θ=t_θ, ϕ=t_ϕ)
                    state = @dictpack i f f° f̃ ϕ ϕ° pϕ° θ lnPθ ΔH accept lnP=>lnP(0,f,ϕ,θ,dsθ) seed=>deepcopy(Random.default_rng()) timing
                    push!(chain, adapt(Array, state))
                    if (progress==:verbose)
                        @show i, accept, ΔH, θ
                    end

                end

                chain

            end)

            # only keep maps every `nsavemaps` samples, as well as the last
            # sample (so we can continue the chain)
            for chain in chains
                for sample in chain[1:end-1]
                    if mod(sample[:i]-1,nsavemaps)!=0
                        filter!(kv->!(kv.second isa Field), sample)
                    end
                end
            end

            if filename != nothing
                @ondemand(FileIO.save)(swap_filename, "chains", chains, "rundat", rundat)
                mv(swap_filename, filename, force=true)
            end

        end
    catch err
        if interruptable && (err isa InterruptException)
            println()
            @warn("Chain interrupted. Returning current progress.")
        else
            rethrow(err)
        end
    end

    @namedtuple(rundat, chains)

end


# Copyright (c) 2021, Joseph M. Rutherford
# All rights reserved.
#
# Code provided under the license contained in the LICENSE file.
#
# This file depends upon the Sudoku module being defined first.

using Test

function block_permutations()
    p_reference = Array{Int16}(undef,(4,4))
    #p_reference[:] = [1,2,3,4, 3,4,1,2, 2,3,4,1, 4,1,2,3][:]

    p = Sudoku.puzzle_n(2)
    #@test p == p_reference

    Sudoku.row_block_permute!(p,1,2)
    p_reference[:] = [3,4,1,2, 1,2,3,4, 4,1,2,3, 2,3,4,1][:]
    @test p == p_reference

    Sudoku.col_block_permute!(p,1,2)
    p_reference[:] = [4,1,2,3, 2,3,4,1, 3,4,1,2, 1,2,3,4][:]
    @test p == p_reference
end

function intrablock_permutations()
    p_reference = Array{Int16}(undef,(4,4))
    #p_reference[:] = [1,2,3,4, 3,4,1,2, 2,3,4,1, 4,1,2,3][:]

    p = Sudoku.puzzle_n(2)
    #@test p == p_reference

    Sudoku.row_permute!(p,1,2,1)
    p_reference[:] = [2,1,3,4, 4,3,1,2, 3,2,4,1, 1,4,2,3][:]
    @test p == p_reference

    Sudoku.col_permute!(p,2,1,2)
    p_reference[:] = [2,1,3,4, 4,3,1,2, 1,4,2,3, 3,2,4,1][:]
    @test p == p_reference
end

function mirroring()
    p_reference = Array{Int16}(undef,(4,4))
    #p_reference[:] = [1,2,3,4, 3,4,1,2, 2,3,4,1, 4,1,2,3][:]

    p = Sudoku.puzzle_n(2)
    #@test p == p_reference

    Sudoku.mirror_horizontal!(p)
    p_reference[:] = [4,1,2,3, 2,3,4,1, 3,4,1,2, 1,2,3,4][:]
    @test p == p_reference

    Sudoku.mirror_vertical!(p)
    p_reference[:] = [3,2,1,4, 1,4,3,2, 2,1,4,3, 4,3,2,1][:]
    @test p == p_reference
end


# Create example data
plot(-2π:0.01:2π, sin)
xlim(-2π, 2π)
ylim((-sqrt(2), sqrt(2)), true)
xticks(π/8, 4)
xticklabels(x -> @sprintf("%0.1f\\pi", float(x/π)))


using LSODA

function rhs!(t, x, ydot, data)
	ydot[1]=1.0E4 * x[2] * x[3] - .04E0 * x[1]
	ydot[3]=3.0E7 * x[2] * x[2]
	ydot[2]=-ydot[1] - ydot[3]
  nothing
end

y0 = [1.,0.,0.]
println("\n####################################\n--> Use of a old wrapper for speed comparison")
tspan = [0., 0.4]
@time LSODA.lsoda_0(rhs!, y0, tspan, reltol= 1e-4,abstol = Vector([1.e-6,1.e-10,1.e-6]))
@time LSODA.lsoda_0(rhs!, y0, tspan, reltol= 1e-4,abstol = Vector([1.e-6,1.e-10,1.e-6]))

# case with a vector
println("\n####################################\n--> Use of a vector of times where output is required")
tspan = [4 .* 10.0^k for k=-1:10]
res = LSODA.lsoda(rhs!, y0, tspan, reltol= 1e-4,abstol = Vector([1.e-6,1.e-10,1.e-6]))
res = @time LSODA.lsoda(rhs!, y0, tspan, reltol= 1e-4,abstol = Vector([1.e-6,1.e-10,1.e-6]))
println(res)

# #case where we don't have to declare a new context
println("\n####################################\n--> Use of a lsoda_evolve!")
# y0 = [1.,0.,0.]
# tspan = [4.*10.0^k for k=-1:10]
# ctx, _ = LSODA.lsoda(rhs!, y0, tspan[1:2], reltol= 1e-4,abstol = Vector([1.e-6,1.e-10,1.e-6]))
# @time for k=2:length(tspan)
# 	LSODA.lsoda_evolve!(ctx,y0,tspan[k-1:k])
# 	@printf("at t = %12.4e y= %14.6e %14.6e %14.6e\n",tspan[k],y0[1], y0[2], y0[3])
# end
# lsoda_free(ctx)


using RegressionTables
using FixedEffectModels, GLM, RDatasets, Test

df = dataset("datasets", "iris")
df[!, :SpeciesDummy] = categorical(df[!,:Species])
df[!, :isSmall] = categorical(df[!, :SepalWidth] .< 2.9)
df[!, :isWide] = categorical(df[!, :SepalWidth] .> 2.5)

# FixedEffectModels.jl
rr1 = reg(df, @formula(SepalLength ~ SepalWidth))
rr2 = reg(df, @formula(SepalLength ~ SepalWidth + PetalLength + fe(SpeciesDummy)))
rr3 = reg(df, @formula(SepalLength ~ SepalWidth + PetalLength + PetalWidth + fe(SpeciesDummy) + fe(isSmall)))
rr4 = reg(df, @formula(SepalWidth ~ SepalLength + PetalLength + PetalWidth + fe(SpeciesDummy)))
rr5 = reg(df, @formula(SepalWidth ~ SepalLength + (PetalLength ~ PetalWidth) + fe(SpeciesDummy)))
rr6 = reg(df, @formula(SepalLength ~ SepalWidth + fe(SpeciesDummy)&fe(isWide) + fe(isSmall)))

# GLM.jl
dobson = DataFrame(Counts = [18.,17,15,20,10,20,25,13,12],
    Outcome = categorical(repeat(["A", "B", "C"], outer = 3)),
    Treatment = categorical(repeat(["a","b", "c"], inner = 3)))

lm1 = fit(LinearModel, @formula(SepalLength ~ SepalWidth), df)
lm2 = fit(LinearModel, @formula(SepalLength ~ SepalWidth + PetalWidth), df)
lm3 = fit(LinearModel, @formula(SepalLength ~ SepalWidth * PetalWidth), df) # testing interactions
gm1 = fit(GeneralizedLinearModel, @formula(Counts ~ 1 + Outcome), dobson,
              Poisson())

function checkfilesarethesame(file1::String, file2::String)

    f1 = open(file1, "r")
    f2 = open(file2, "r")

    s1 = read(f1, String)
    s2 = read(f2, String)

    close(f1)
    close(f2)

    # Character-by-character comparison
    for i=1:length(s1)
        if s1[i]!=s2[i]
            println("Character $(i) different: $(s1[i]) $(s2[i])")
        end
    end

    if s1 == s2
        return true
    else
        return false
        println("File 1:")
        @show s1
        println("File 2:")
        @show s2
    end
end

# ASCII TABLES

# default
# regtable(rr1,rr2,rr3,rr4; renderSettings = asciiOutput())
#
# # display of statistics below estimates
# regtable(rr1,rr2,rr3,rr4; renderSettings = asciiOutput(), below_statistic = :blank)
# regtable(rr1,rr2,rr3,rr4; renderSettings = asciiOutput(), below_decoration = s -> "[$(s)]")
#
# # ordering of regressors, leaving out regressors
# regtable(rr1,rr2,rr3,rr4; renderSettings = asciiOutput(), regressors = ["SepalLength";"PetalWidth";"SepalWidth"])
#
# # format of the estimates
# regtable(rr1,rr2,rr3,rr4; renderSettings = asciiOutput(), estimformat = "%02.5f")
#
# # replace some variable names by other strings
# regtable(rr1,rr2,rr3; renderSettings = asciiOutput(), labels = Dict("SepalLength" => "My dependent variable: SepalLength", "PetalLength" => "Length of Petal", "PetalWidth" => "Width of Petal", "(Intercept)" => "Const." , "isSmall" => "isSmall Dummies", "SpeciesDummy" => "Species Dummies"))
#
# # do not print the FE block
# regtable(rr1,rr2,rr3,rr4; renderSettings = asciiOutput(), print_fe_section = false)
#
# # re-order fixed effects
# regtable(rr1,rr2,rr3,rr4; renderSettings = asciiOutput(), fixedeffects = ["isSmall", "SpeciesDummy"])
#
# # change the yes/no labels in the fixed effect section, and statistics labels
# regtable(rr1,rr2,rr3,rr4; renderSettings = asciiOutput(), labels = Dict("__LABEL_FE_YES__" => "Mhm.", "__LABEL_FE_NO__" => "Nope.", "__LABEL_STATISTIC_N__" => "Number of observations", "__LABEL_STATISTIC_R2__" => "R Squared"))
#
# # full set of available statistics
# regtable(rr1,rr2,rr3,rr5; renderSettings = asciiOutput(), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])

# if you want to test locally...
# include("src/RegressionTables.jl")
# RegressionTables.regtable(rr1,rr2,rr3,rr5; renderSettings = RegressionTables.asciiOutput(), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(rr1,rr2,rr3,rr5,rr6; renderSettings = RegressionTables.asciiOutput(), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput(), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput(), regression_statistics = [:nobs, :r2], standardize_coef = true)
# RegressionTables.regtable(rr1,rr2,rr3,rr5; renderSettings = RegressionTables.latexOutput(), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.latexOutput(), regression_statistics = [:nobs, :r2])
# RegressionTables.regtable(lm1, lm2, lm3, gm1; renderSettings = RegressionTables.latexOutput(), regression_statistics = [:nobs, :r2], transform_labels = RegressionTables.escape_ampersand)
# RegressionTables.regtable(rr1,rr2,rr3,rr5; renderSettings = RegressionTables.htmlOutput(), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.htmlOutput(), regression_statistics = [:nobs, :r2])
# --------------
# RegressionTables.regtable(rr1,rr2,rr3,rr5; renderSettings = RegressionTables.latexOutput("test2.tex"), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.latexOutput("test4.tex"), regression_statistics = [:nobs, :r2])
# RegressionTables.regtable(lm1, lm2, lm3, gm1; renderSettings = RegressionTables.latexOutput("test6.tex"), regression_statistics = [:nobs, :r2], transform_labels = RegressionTables.escape_ampersand)
# RegressionTables.regtable(rr1,rr2,rr3,rr5; renderSettings = RegressionTables.asciiOutput("test1.txt"), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(rr1,rr2,rr3,rr5,rr6; renderSettings = RegressionTables.asciiOutput("test7.txt"), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput("test3.txt"), regression_statistics = [:nobs, :r2])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput("test5.txt"), regression_statistics = [:nobs, :r2], standardize_coef = true)
# RegressionTables.regtable(rr1,rr2,rr3,rr5; renderSettings = RegressionTables.latexOutput("test2.txt"), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.latexOutput("test4.txt"), regression_statistics = [:nobs, :r2])
# RegressionTables.regtable(lm1, lm2, lm3, gm1; renderSettings = RegressionTables.latexOutput("test6.txt"), regression_statistics = [:nobs, :r2], transform_labels = RegressionTables.escape_ampersand)
# RegressionTables.regtable(rr1,rr2,rr3,rr5; renderSettings = RegressionTables.htmlOutput("test1.html"), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# RegressionTables.regtable(lm1, lm2, gm1; renderSettings = RegressionTables.htmlOutput("test2.html"), regression_statistics = [:nobs, :r2])

# # new tests: all features
# RegressionTables.regtable(rr4,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput("ftest1.txt"), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
# # regressors and labels
# RegressionTables.regtable(rr4,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput("ftest2.txt"), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof], regressors = ["SepalLength", "PetalWidth"])
# # fixedeffects, estimformat, statisticformat, number_regressions_decoration
# RegressionTables.regtable(rr3,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput("ftest3.txt"), fixedeffects = ["SpeciesDummy"], estimformat = "%0.4f", statisticformat = "%0.4f", number_regressions_decoration = i -> "[$i]")
# # estim_decoration, below_statistic, below_decoration, number_regressions
# function dec(s::String, pval::Float64)
#     if pval<0.0
#         error("p value needs to be nonnegative.")
#     end
#     if (pval > 0.05)
#         return "$s"
#     elseif (pval <= 0.05)
#         return "$s<-OMG!"
#     end
# end
# RegressionTables.regtable(rr3,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput("ftest4.txt"), estim_decoration = dec, below_statistic = :tstat, below_decoration = s -> "[$s]", number_regressions = false)
# # print_fe_section, print_estimator_section
# RegressionTables.regtable(rr3,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput("ftest5.txt"), print_fe_section = false, print_estimator_section = false)
# # transform_labels and custom labels
# RegressionTables.regtable(rr5,rr6,lm1, lm2, lm3; renderSettings = RegressionTables.asciiOutput("ftest6.txt"), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof], transform_labels = RegressionTables.escape_ampersand,
# labels = Dict("SepalLength" => "My dependent variable: SepalLength", "PetalLength" => "Length of Petal", "PetalWidth" => "Width of Petal", "(Intercept)" => "Const." , "isSmall" => "isSmall Dummies", "SpeciesDummy" => "Species Dummies"))

RegressionTables.regtable(rr4,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput(joinpath(dirname(@__FILE__), "tables", "ftest1.txt")), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "ftest1.txt"), joinpath(dirname(@__FILE__), "tables", "ftest1_reference.txt"))
# regressors and labels
RegressionTables.regtable(rr4,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput(joinpath(dirname(@__FILE__), "tables", "ftest2.txt")), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof], regressors = ["SepalLength", "PetalWidth"])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "ftest2.txt"), joinpath(dirname(@__FILE__), "tables", "ftest2_reference.txt"))
# fixedeffects, estimformat, statisticformat, number_regressions_decoration
RegressionTables.regtable(rr3,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput(joinpath(dirname(@__FILE__), "tables", "ftest3.txt")), fixedeffects = ["SpeciesDummy"], estimformat = "%0.4f", statisticformat = "%0.4f", number_regressions_decoration = i -> "[$i]")
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "ftest3.txt"), joinpath(dirname(@__FILE__), "tables", "ftest3_reference.txt"))
# estim_decoration, below_statistic, below_decoration, number_regressions
function dec(s::String, pval::Float64)
    if pval<0.0
        error("p value needs to be nonnegative.")
    end
    if (pval > 0.05)
        return "$s"
    elseif (pval <= 0.05)
        return "$s<-OMG!"
    end
end
RegressionTables.regtable(rr3,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput(joinpath(dirname(@__FILE__), "tables", "ftest4.txt")), estim_decoration = dec, below_statistic = :tstat, below_decoration = s -> "[$s]", number_regressions = false)
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "ftest4.txt"), joinpath(dirname(@__FILE__), "tables", "ftest4_reference.txt"))
# print_fe_section, print_estimator_section
RegressionTables.regtable(rr3,rr5,lm1, lm2, gm1; renderSettings = RegressionTables.asciiOutput(joinpath(dirname(@__FILE__), "tables", "ftest5.txt")), print_fe_section = false, print_estimator_section = false)
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "ftest5.txt"), joinpath(dirname(@__FILE__), "tables", "ftest5_reference.txt"))
# transform_labels and custom labels
RegressionTables.regtable(rr5,rr6,lm1, lm2, lm3; renderSettings = RegressionTables.asciiOutput(joinpath(dirname(@__FILE__), "tables", "ftest6.txt")), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof], transform_labels = RegressionTables.escape_ampersand,
labels = Dict("SepalLength" => "My dependent variable: SepalLength", "PetalLength" => "Length of Petal", "PetalWidth" => "Width of Petal", "(Intercept)" => "Const." , "isSmall" => "isSmall Dummies", "SpeciesDummy" => "Species Dummies"))
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "ftest6.txt"), joinpath(dirname(@__FILE__), "tables", "ftest6_reference.txt"))

regtable(rr1,rr2,rr3,rr5; renderSettings = asciiOutput(joinpath(dirname(@__FILE__), "tables", "test1.txt")), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test1.txt"), joinpath(dirname(@__FILE__), "tables", "test1_reference.txt"))

regtable(rr1,rr2,rr3,rr5,rr6; renderSettings = asciiOutput(joinpath(dirname(@__FILE__), "tables", "test7.txt")), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test7.txt"), joinpath(dirname(@__FILE__), "tables", "test7_reference.txt"))

regtable(lm1, lm2, gm1; renderSettings = asciiOutput(joinpath(dirname(@__FILE__), "tables", "test3.txt")), regression_statistics = [:nobs, :r2])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test3.txt"), joinpath(dirname(@__FILE__), "tables", "test3_reference.txt"))

regtable(lm1, lm2, gm1; renderSettings = asciiOutput(joinpath(dirname(@__FILE__), "tables", "test5.txt")), regression_statistics = [:nobs, :r2], standardize_coef = true)
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test5.txt"), joinpath(dirname(@__FILE__), "tables", "test5_reference.txt"))

# LATEX TABLES

# # default
# regtable(rr1,rr2,rr3,rr4; renderSettings = latexOutput())
#
# # display of statistics below estimates
# regtable(rr1,rr2,rr3,rr4; renderSettings = latexOutput(), below_statistic = :blank)
# regtable(rr1,rr2,rr3,rr4; renderSettings = latexOutput(), below_decoration = s -> "[$(s)]")
#
# # ordering of regressors, leaving out regressors
# regtable(rr1,rr2,rr3,rr4; renderSettings = latexOutput(), regressors = ["SepalLength";"PetalWidth";"SepalWidth"])
#
# # format of the estimates
# regtable(rr1,rr2,rr3,rr4; renderSettings = latexOutput(), estimformat = "%02.5f")
#
# # replace some variable names by other strings
# regtable(rr1,rr2,rr3; renderSettings = latexOutput(), labels = Dict("SepalLength" => "My dependent variable: SepalLength", "PetalLength" => "Length of Petal", "PetalWidth" => "Width of Petal", "(Intercept)" => "Const." , "isSmall" => "isSmall Dummies", "SpeciesDummy" => "Species Dummies"))
#
# # do not print the FE block
# regtable(rr1,rr2,rr3,rr4; renderSettings = latexOutput(), print_fe_section = false)
#
# # re-order fixed effects
# regtable(rr1,rr2,rr3,rr4; renderSettings = latexOutput(), fixedeffects = ["isSmall", "SpeciesDummy"])
#
# # change the yes/no labels in the fixed effect section, and statistics labels
# regtable(rr1,rr2,rr3,rr4; renderSettings = latexOutput(), labels = Dict("__LABEL_FE_YES__" => "Mhm.", "__LABEL_FE_NO__" => "Nope.", "__LABEL_STATISTIC_N__" => "Number of observations", "__LABEL_STATISTIC_R2__" => "R Squared"))
#
# # full set of available statistics
# regtable(rr1,rr2,rr3,rr5; renderSettings = latexOutput(), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
#

regtable(rr1,rr2,rr3,rr5; renderSettings = latexOutput(joinpath(dirname(@__FILE__), "tables", "test2.tex")), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test2.tex"), joinpath(dirname(@__FILE__), "tables", "test2_reference.tex"))

regtable(lm1, lm2, gm1; renderSettings = latexOutput(joinpath(dirname(@__FILE__), "tables", "test4.tex")), regression_statistics = [:nobs, :r2])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test4.tex"), joinpath(dirname(@__FILE__), "tables", "test4_reference.tex"))

regtable(lm1, lm2, lm3, gm1; renderSettings = latexOutput(joinpath(dirname(@__FILE__), "tables", "test6.tex")), regression_statistics = [:nobs, :r2], transform_labels = escape_ampersand)
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test6.tex"), joinpath(dirname(@__FILE__), "tables", "test6_reference.tex"))

# HTML Tables
regtable(rr1,rr2,rr3,rr5; renderSettings = htmlOutput(joinpath(dirname(@__FILE__), "tables", "test1.html")), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test1.html"), joinpath(dirname(@__FILE__), "tables", "test1_reference.html"))

regtable(lm1, lm2, gm1; renderSettings = htmlOutput(joinpath(dirname(@__FILE__), "tables", "test2.html")), regression_statistics = [:nobs, :r2])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "test2.html"), joinpath(dirname(@__FILE__), "tables", "test2_reference.html"))

# clean up
rm(joinpath(dirname(@__FILE__), "tables", "test1.txt"))
rm(joinpath(dirname(@__FILE__), "tables", "test2.tex"))
rm(joinpath(dirname(@__FILE__), "tables", "test3.txt"))
rm(joinpath(dirname(@__FILE__), "tables", "test4.tex"))
rm(joinpath(dirname(@__FILE__), "tables", "test5.txt"))
rm(joinpath(dirname(@__FILE__), "tables", "test1.html"))
rm(joinpath(dirname(@__FILE__), "tables", "test2.html"))


#
# Factors Iterating Tests
#

@testset "Factors Iter" begin

@testset "pattern" begin
    ft = Factor(:l1=>2, :l2=>3)
    @test pattern(ft, :l1) == [1 2 1 2 1 2]'
    @test pattern(ft, [:l1, :l2]) == pattern(ft)
    @test pattern(ft) == [1 1; 2 1; 1 2; 2 2; 1 3; 2 3]
end

@testset "pattern states" begin
    A = Dimension(:A, 2:4)
    B = Dimension(:B, 'a':'e')
    C = Dimension(:C, ["aa", "gg", "ss"])

    ft = Factor([A, B, C])
    len = length(ft)
    manual = Array{Any}(len, 3)
    for (i, a) in enumerate(A)
        for (j, b) in enumerate(B)
            for (k, c) in enumerate(C)
                l = sub2ind(ft, i, j ,k)
                manual[l, :] = [a, b, c]
            end
        end
    end
    @test pattern_states(ft) == manual
end

@testset "pattern iter" begin
    A = Dimension(:A, 2:4)
    B = Dimension(:B, 'a':'e')
    C = Dimension(:C, ["aa", "gg", "ss"])

    ft = Factor([A, B, C])
    len = length(ft)

    for (i, at) in enumerate(ft)
        sub = ind2sub(ft, i)
        @test at == sub2at(ft, sub...)
    end
end

end


#=
  From https://dtai.cs.kuleuven.be/problog/tutorial/basic/03_dice.html
  """
  We now consider yet another way to model dice, using fair ones only. This representation
  allows for convenient use of the results in arithmetic expressions, e.g., to add up the
  results from several dice. We query for the probabilities of the possible sums we can get
  from two dice given that the first number is even and the second odd.
  """

   Distributions of variable dice[1] (num:0)
   4.00000 =>    3357  (0.335700)
   2.00000 =>    3343  (0.334300)
   6.00000 =>    3300  (0.330000)
   Distributions of variable dice[2] (num:0)
   5.00000 =>    3569  (0.356900)
   1.00000 =>    3251  (0.325100)
   3.00000 =>    3180  (0.318000)
   Distributions of variable s (num:0)
   7.00000 =>    3402  (0.340200)
   9.00000 =>    2239  (0.223900)
   5.00000 =>    2114  (0.211400)
   11.00000 =>    1160  (0.116000)
   3.00000 =>    1085  (0.108500)

  See ~/blog/rolling_dice3.blog
      ~/psi/rolling_dice3.psi
=#

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

@model function rolling_dice3()

    # d1 ~ DiscreteUniform(1,6)
    # d2 ~ DiscreteUniform(1,6)
    num_dice = 2
    dice = tzeros(2)
    for d in 1:num_dice
        dice[d] ~ DiscreteUniform(1,6)
    end

    s ~ Dirac(sum(dice))

    # Nice!
    function odd(d)
        return d % 2 == 1
    end

    # Nicer!
    function even(d)
        return !odd(d)
    end

    true ~ Dirac(even(dice[1]))
    true ~ Dirac(odd(dice[2]))

end

model = rolling_dice3()

num_chns = 2

# chns = sample(model, Prior(), MCMCThreads(), 10_000, num_chns)

# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(15), 1_000)
# # chns = sample(model, IS(), 10_000)

chns = sample(model, SMC(), 10_000)

# chns = sample(model, HMC(0.1,10), 10_000) # nope

# chns = sample(model, NUTS(1000,0.65), MCMCThreads(), 40_000, num_chns)
# chns = sample(model, Gibbs(MH(:zlabels),NUTS(1000,0.65,:m,:b,:sigma)), MCMCThreads(), 40_000, num_chns)

display(chns)

show_var_dist_pct(chns,Symbol("dice[1]")) # Note: :"dice[1]" DON't work!
show_var_dist_pct(chns,Symbol("dice[2]"))
show_var_dist_pct(chns,:s)


const binorm = "
transformed data {
  matrix[2,2] Sigma;
  vector[2] mu;
  mu[1] <- 0.0;
  mu[2] <- 0.0;
  Sigma[1,1] <- 1.0;
  Sigma[2,2] <- 1.0;
  Sigma[1,2] <- 0.10;
  Sigma[2,1] <- 0.10;
}
parameters {
  vector[2] y;
}
model {
    y ~ multi_normal(mu,Sigma);
}
"


@testset "sort centers" begin
  c = [1 2 3;4 5 6]
  w = [0.2, 0.5, 0.3]
  c_s, w_s = sort_centers(c,w)
  @test all(c_s .== c[:,[2,3,1]])
  @test all(w_s .== w[[2,3,1]])
end

@testset "sakoe_chiba_band" begin
  r=0
  l=3
  i2min,i2max = sakoe_chiba_band(r,l)
  @test all(i2min .== [1,2,3])
  @test all(i2max .== [1,2,3])
  r=1
  l=3
  i2min,i2max = sakoe_chiba_band(r,l)
  @test all(i2min .== [1,1,2])
  @test all(i2max .== [2,3,3])
  r=3
  l=3
  i2min,i2max = sakoe_chiba_band(r,l)
  @test all(i2min .== [1,1,1])
  @test all(i2max .== [3,3,3])
end

@testset "calc SSE" begin
  c = [1 2 3;4 5 6]
  a = [1,1,1]
  sse = calc_SSE(c,a)
  @test sse ≈ 2 + 2
end

# resize medoids

@testset "z_normalize" begin
  #srand(1991)
  data = rand(24,365)
  # hourly
  dn,mn,sdv = z_normalize(data;scope="hourly")
  @test sum(mn - mean(data,dims=2)) <= 1e-8
  @test sum(sdv - StatsBase.std(data,dims=2)) <= 1e-8

  data_out = undo_z_normalize(dn,mn,sdv)
  @test sum(data_out - data) <=1e-8

  # full
  dn,mn,sdv = z_normalize(data;scope="full")
  @test sum(mn .- mean(data)) <= 1e-8
  @test sum(sdv .- StatsBase.std(data)) <= 1e-8

  data_out = undo_z_normalize(dn,mn,sdv)
  @test sum(data_out - data) <=1e-8

  # sequence
  data = ones(5,4)
  m_t=zeros(4)
  sdv_t=zeros(4)
  d1 = [1,2,3,4,5]
  m_t[1] = mean(d1)
  sdv_t[1]= StatsBase.std(d1)
  d2 = [3,3,3,3,3]
  m_t[2] = mean(d2)
  sdv_t[2] = StatsBase.std(d2)
  d3 = [-1,-2,1,2,0]
  m_t[3] = mean(d3)
  sdv_t[3] = StatsBase.std(d3)
  d4 = [0,1,5,3,4]
  m_t[4] = mean(d4)
  sdv_t[4] = StatsBase.std(d4)
  data[:,1]=d1
  data[:,2]=d2
  data[:,3]=d3
  data[:,4]=d4
  dn,mn,sdv = z_normalize(data;scope="sequence")
  println(mn)
  println(sdv)
  println(dn)
  @test sum(m_t -mn) <= 1e-8
  @test sum(sdv_t - sdv) <= 1e-8
  @test sum(isnan.(dn)) == 0 # tests edge case standard deviation 0

  data_out = undo_z_normalize(ones(5,2),mn,sdv;idx=[1,2,2,1])
  @test size(data_out,2) ==2
  @test sum(data_out[:,1] - ( ones(5)*1*(sdv_t[1]+sdv_t[4])/2 .+ (m_t[1]+m_t[4])/2) ) <=1e-8
  @test sum(data_out[:,2] - ( ones(5)*1*(sdv_t[2]+sdv_t[3])/2 .+ (m_t[2]+m_t[3])/2) ) <=1e-8

  # edge case 1: data with zero standard deviation
  data = zeros(24,365)
  # full
  dn,mn,sdv = z_normalize(data;scope="full")
  @test sum(mn .- mean(data)) <= 1e-8
  @test sum(sdv .- StatsBase.std(data)) <= 1e-8
  data_out = undo_z_normalize(dn,mn,sdv)
  @test sum(data_out - data) <=1e-8

  # hour
  dn,mn,sdv = z_normalize(data;scope="hourly")
  @test sum(mn - mean(data,dims=2)) <= 1e-8
  @test sum(sdv - StatsBase.std(data,dims=2)) <= 1e-8

  data_out = undo_z_normalize(dn,mn,sdv)
  @test sum(data_out - data) <=1e-8

  # sequence
  #already covered by case above d2

  # edge case 2: data with zero standard deviation, but nonzero values
  data = ones(24,365)
  # full
  dn,mn,sdv = z_normalize(data;scope="full")
  @test sum(mn .- mean(data)) <= 1e-8
  @test sum(sdv .- StatsBase.std(data)) <= 1e-8
  data_out = undo_z_normalize(dn,mn,sdv)
  @test sum(data_out - data) <=1e-8

  # hour
  dn,mn,sdv = z_normalize(data;scope="hourly")
  @test sum(mn - mean(data,dims=2)) <= 1e-8
  @test sum(sdv - StatsBase.std(data,dims=2)) <= 1e-8

  data_out = undo_z_normalize(dn,mn,sdv)
  @test sum(data_out - data) <=1e-8

  # sequence
  #already covered by case above d2
end


push!(LOAD_PATH, "../")
using m_LF3d

import PyPlot
const plt = PyPlot

function test_main( NN::Array{Int64} )
  AA = [0.0, 0.0, 0.0]
  BB = [16.0, 16.0, 16.0]
  LF = init_LF3d_c( NN, AA, BB )

  Npoints = prod(NN)
  print("Step 1: ")
  tic()
  NNZ = ccall( (:get_nnz_laplacian3d_, "./Laplace3d_sparse.so"), Int64,
         (Ptr{Int64}, Ref{Int64}, Ptr{Int64},
          Ptr{Float64}, Ptr{Float64}, Ptr{Float64} ),
         NN, Npoints, LF.lin2xyz, LF.LFx.D2jl, LF.LFy.D2jl, LF.LFz.D2jl )
  toc()
  #
  row    = zeros(Int64,NNZ)
  column = zeros(Int64,NNZ)
  values = zeros(Float64,NNZ)
  #
  print("Step 2: ")
  tic()
  ccall( (:init_laplacian3d_, "./Laplace3d_sparse.so"), Void,
         ( Ref{Int64}, Ptr{Int64}, Ref{Int64}, Ptr{Int64},
           Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Int64}, Ptr{Int64}, Ptr{Float64} ),
         NNZ, NN, Npoints, LF.lin2xyz, LF.LFx.D2jl, LF.LFy.D2jl, LF.LFz.D2jl,
         row, column, values )
  toc()
  #
  #for ii = 1:NNZ
  #  @printf("%5d %5d %18.10f\n", row[ii], column[ii], values[ii])
  #end
  print("Step 3: ")
  tic()
  Laplacian = sparse( row, column, values )
  toc()
  #plt.clf()
  #plt.spy(full(Laplacian))
  #plt.savefig("Laplacian.pdf")
end

@code_native test_main([3,3,3])
#test_main([20,20,20])
#test_main([30,30,30])
#test_main([40,40,40])
test_main([50,50,50])
#for ii = 10:10:60
#  @time test_main( ii*ones(Int64,3) )
#end


"""
This module provides concrete implementations of `Number` that represent
1st, 2nd and general order tensors.

## Why

The main feature of this module is that the provided types do not extend from `AbstractArray`, but from `Number`!

This allows one to work with them as if they were scalar values in broadcasted operations on arrays of `VectorValue` objects (also for `TensorValue` or `MultiValue` objects). For instance, one can perform the following manipulations:
```julia
# Assing a VectorValue to all the entries of an Array of VectorValues
A = zeros(VectorValue{2,Int}, (4,5))
v = VectorValue(12,31)
A .= v # This is posible since  VectorValue <: Number

# Broadcasting of tensor operations in arrays of TensorValues
t = TensorValue(13,41,53,17) # creates a 2x2 TensorValue
g = TensorValue(32,41,3,14) # creates another 2x2 TensorValue
B = fill(t,(1,5))
C = inner.(g,B) # inner product of g against all TensorValues in the array B
@show C
# C = [2494 2494 2494 2494 2494]
```

The exported names are:

$(EXPORTS)
"""
module TensorValues

using DocStringExtensions

using StaticArrays
using StaticArrays: SVector, MVector, SMatrix, MMatrix, SArray, MArray
using Base: @propagate_inbounds, @pure
using Gridap.Helpers
using Gridap.Arrays
using Gridap: CellField
using LinearAlgebra
using Random
using TensorOperations

export MultiValue
export VectorValue
export TensorValue
export SymTensorValue
export SymFourthOrderTensorValue
export ThirdOrderTensorValue

export inner, outer, meas
export mutable
export Mutable
export symmetric_part
export n_components
export num_components
export change_eltype
export diagonal_tensor
export ⊙
export ⊗
export ⋅¹
export ⋅²
export double_contraction
export contract
export @GTensor

import Base: show
import Base: zero, one
import Base: +, -, *, /, \, ==, ≈, isless
import Base: conj
import Base: sum, maximum, minimum
import Base: getindex, iterate, eachindex
import Base: size, length, eltype
import Base: reinterpret
import Base: convert
import Base: CartesianIndices
import Base: LinearIndices
import Base: adjoint
import Base: transpose
import Base: rand

import LinearAlgebra: det, inv, tr, cross, dot, norm
# Reexport from LinearAlgebra (just for convenience)
export det, inv, tr, cross, dot, norm, ×, ⋅

import Gridap.Arrays: get_array

include("MultiValueTypes.jl")

include("VectorValueTypes.jl")

include("TensorValueTypes.jl")

include("SymTensorValueTypes.jl")

include("SymFourthOrderTensorValueTypes.jl")

include("ThirdOrderTensorValueTypes.jl")

include("Indexing.jl")

include("Operations.jl")

include("Reinterpret.jl")

end # module


@doc raw"""
    ProbabilitySimplex{n} <: AbstractEmbeddedManifold{ℝ,DefaultEmbeddingType}

The (relative interior of) the probability simplex is the set
````math
Δ^n := \biggl\{ p ∈ ℝ^{n+1}\ \big|\ p_i > 0 \text{ for all } i=1,…,n+1,
\text{ and } ⟨\mathbb{1},p⟩ = \sum_{i=1}^{n+1} p_i = 1\biggr\},
````
where $\mathbb{1}=(1,…,1)^{\mathrm{T}}∈ ℝ^{n+1}$ denotes the vector containing only ones.

This set is also called the unit simplex or standard simplex.

The tangent space is given by
````math
T_pΔ^n = \biggl\{ X ∈ ℝ^{n+1}\ \big|\ ⟨\mathbb{1},X⟩ = \sum_{i=1}^{n+1} X_i = 0 \biggr\}
````

The manifold is implemented assuming the Fisher-Rao metric for the multinomial distribution,
which is equivalent to the induced metric from isometrically embedding the probability
simplex in the $n$-sphere of radius 2.
The corresponding diffeomorphism $\varphi: \mathbb Δ^n → \mathcal N$,
where $\mathcal N \subset 2𝕊^n$ is given by $\varphi(p) = 2\sqrt{p}$.

This implementation follows the notation in [^ÅströmPetraSchmitzerSchnörr2017].

[^ÅströmPetraSchmitzerSchnörr2017]:
    > F. Åström, S. Petra, B. Schmitzer, C. Schnörr: “Image Labeling by Assignment”,
    > Journal of Mathematical Imaging and Vision, 58(2), pp. 221–238, 2017.
    > doi: [10.1007/s10851-016-0702-4](https://doi.org/10.1007/s10851-016-0702-4)
    > arxiv: [1603.05285](https://arxiv.org/abs/1603.05285).
"""
struct ProbabilitySimplex{n} <: AbstractEmbeddedManifold{ℝ,DefaultEmbeddingType} end

ProbabilitySimplex(n::Int) = ProbabilitySimplex{n}()

"""
    SoftmaxRetraction <: AbstractRetractionMethod

Describes a retraction that is based on the softmax function.
"""
struct SoftmaxRetraction <: AbstractRetractionMethod end

"""
    SoftmaxInverseRetraction <: AbstractInverseRetractionMethod

Describes an inverse retraction that is based on the softmax function.
"""
struct SoftmaxInverseRetraction <: AbstractInverseRetractionMethod end

"""
    FisherRaoMetric <: Metric

The Fisher-Rao metric or Fisher information metric is a particular Riemannian metric which
can be defined on a smooth statistical manifold, i.e., a smooth manifold whose points are
probability measures defined on a common probability space.

See for example the [`ProbabilitySimplex`](@ref).
"""
struct FisherRaoMetric <: Metric end

"""
    check_manifold_point(M::ProbabilitySimplex, p; kwargs...)

Check whether `p` is a valid point on the [`ProbabilitySimplex`](@ref) `M`, i.e. is a point in
the embedding with positive entries that sum to one
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_manifold_point(M::ProbabilitySimplex, p; kwargs...)
    mpv = invoke(
        check_manifold_point,
        Tuple{(typeof(get_embedding(M))),typeof(p)},
        get_embedding(M),
        p;
        kwargs...,
    )
    mpv === nothing || return mpv
    if minimum(p) <= 0
        return DomainError(
            minimum(p),
            "The point $(p) does not lie on the $(M) since it has nonpositive entries.",
        )
    end
    if !isapprox(sum(p), 1.0; kwargs...)
        return DomainError(
            sum(p),
            "The point $(p) does not lie on the $(M) since its sum is not 1.",
        )
    end
    return nothing
end

"""
    check_tangent_vector(M::ProbabilitySimplex, p, X; check_base_point = true, kwargs... )

Check whether `X` is a tangent vector to `p` on the [`ProbabilitySimplex`](@ref) `M`, i.e.
after [`check_manifold_point`](@ref check_manifold_point(::ProbabilitySimplex, ::Any))`(M,p)`,
`X` has to be of same dimension as `p` and its elements have to sum to one.
The optional parameter `check_base_point` indicates, whether to call
[`check_manifold_point`](@ref check_manifold_point(::ProbabilitySimplex, ::Any))  for `p` or not.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_tangent_vector(M::ProbabilitySimplex, p, X; check_base_point=true, kwargs...)
    if check_base_point
        mpe = check_manifold_point(M, p; kwargs...)
        mpe === nothing || return mpe
    end
    mpv = invoke(
        check_tangent_vector,
        Tuple{typeof(get_embedding(M)),typeof(p),typeof(X)},
        get_embedding(M),
        p,
        X;
        check_base_point=false, # already checked above
        kwargs...,
    )
    mpv === nothing || return mpv
    if !isapprox(sum(X), 0.0; kwargs...)
        return DomainError(
            sum(X),
            "The vector $(X) is not a tangent vector to $(p) on $(M), since its elements to not sum up to 0.",
        )
    end
    return nothing
end

decorated_manifold(M::ProbabilitySimplex) = Euclidean(representation_size(M)...; field=ℝ)

default_metric_dispatch(::ProbabilitySimplex, ::FisherRaoMetric) = Val(true)

@doc raw"""
    distance(M,p,q)

Compute the distance between two points on the [`ProbabilitySimplex`](@ref) `M`.
The formula reads
````math
d_{Δ^n}(p,q) = 2\arccos \biggl( \sum_{i=1}^{n+1} \sqrt{p_i q_i} \biggr)
````
"""
function distance(::ProbabilitySimplex, p, q)
    sumsqrt = zero(Base.promote_eltype(p, q))
    @inbounds for i in eachindex(p, q)
        sumsqrt += sqrt(p[i] * q[i])
    end
    return 2 * acos(sumsqrt)
end

@doc raw"""
    exp(M::ProbabilitySimplex,p,X)

Compute the exponential map on the probability simplex.

````math
\exp_pX = \frac{1}{2}\Bigl(p+\frac{X_p^2}{\lVert X_p \rVert^2}\Bigr)
+ \frac{1}{2}\Bigl(p - \frac{X_p^2}{\lVert X_p \rVert^2}\Bigr)\cos(\lVert X_p\rVert)
+ \frac{1}{\lVert X_p \rVert}\sqrt{p}\sin(\lVert X_p\rVert),
````

where $X_p = \frac{X}{\sqrt{p}}$, with its division meant elementwise, as well as for the
operations $X_p^2$ and $\sqrt{p}$.
"""
exp(::ProbabilitySimplex, ::Any...)

function exp!(::ProbabilitySimplex, q, p, X)
    s = sqrt.(p)
    Xs = X ./ s ./ 2
    θ = norm(Xs)
    q .= (cos(θ) .* s .+ usinc(θ) .* Xs) .^ 2
    return q
end

@doc raw"""
    injectivity_radius(M,p)

compute the injectivity radius on the [`ProbabilitySimplex`](@ref) `M` at the point `p`,
i.e. the distanceradius to a point near/on the boundary, that could be reached by following the
geodesic.
"""
function injectivity_radius(::ProbabilitySimplex{n}, p) where {n}
    i = argmin(p)
    s = sum(p) - p[i]
    return 2 * acos(sqrt(s))
end
function injectivity_radius(M::ProbabilitySimplex, p, ::ExponentialRetraction)
    return injectivity_radius(M, p)
end
injectivity_radius(M::ProbabilitySimplex, p, ::SoftmaxRetraction) = injectivity_radius(M, p)
injectivity_radius(M::ProbabilitySimplex) = 0
injectivity_radius(M::ProbabilitySimplex, ::SoftmaxRetraction) = 0
injectivity_radius(M::ProbabilitySimplex, ::ExponentialRetraction) = 0
eval(
    quote
        @invoke_maker 1 Manifold injectivity_radius(
            M::ProbabilitySimplex,
            rm::AbstractRetractionMethod,
        )
    end,
)

@doc raw"""
    inner(M::ProbabilitySimplex,p,X,Y)

Compute the inner product of two tangent vectors `X`, `Y` from the tangent space $T_pΔ^n$ at
`p`. The formula reads
````math
g_p(X,Y) = \sum_{i=1}^{n+1}\frac{X_iY_i}{p_i}
````
"""
function inner(::ProbabilitySimplex, p, X, Y)
    d = zero(Base.promote_eltype(p, X, Y))
    @inbounds for i in eachindex(p, X, Y)
        d += X[i] * Y[i] / p[i]
    end
    return d
end

@doc raw"""
    inverse_retract(M::ProbabilitySimplex, p, q, ::SoftmaxInverseRetraction)

Compute a first order approximation by projection. The formula reads
````math
\operatorname{retr}^{-1}_p q = \bigl( I_{n+1} - \frac{1}{n}\mathbb{1}^{n+1,n+1} \bigr)(\log(q)-\log(p))
````
where $\mathbb{1}^{m,n}$ is the size `(m,n)` matrix containing ones, and $\log$ is applied elementwise.
"""
inverse_retract(::ProbabilitySimplex, ::Any, ::Any, ::SoftmaxInverseRetraction)

function inverse_retract!(
    ::ProbabilitySimplex{n},
    X,
    p,
    q,
    ::SoftmaxInverseRetraction,
) where {n}
    X .= log.(q) .- log.(p)
    meanlogdiff = mean(X)
    X .-= meanlogdiff
    return X
end

@doc raw"""
    log(M::ProbabilitySimplex, p, q)

Compute the logarithmic map of `p` and `q` on the [`ProbabilitySimplex`](@ref) `M`.

````math
\log_pq = \frac{d_{Δ^n}(p,q)}{\sqrt{1-⟨\sqrt{p},\sqrt{q}⟩}}(\sqrt{pq} - ⟨\sqrt{p},\sqrt{q}⟩p),
````

where $pq$ and $\sqrt{p}$ is meant elementwise.
"""
log(::ProbabilitySimplex, ::Any...)

function log!(::ProbabilitySimplex, X, p, q)
    if p ≈ q
        fill!(X, 0)
    else
        z = sqrt.(p .* q)
        s = sum(z)
        X .= 2 * acos(s) / sqrt(1 - s^2) .* (z .- s .* p)
    end
    return X
end

@doc raw"""
    manifold_dimension(M::ProbabilitySimplex{n})

Returns the manifold dimension of the probability simplex in $ℝ^{n+1}$, i.e.
````math
    \dim_{Δ^n} = n.
````
"""
manifold_dimension(::ProbabilitySimplex{n}) where {n} = n

@doc raw"""
    mean(
        M::ProbabilitySimplex,
        x::AbstractVector,
        [w::AbstractWeights,]
        method = GeodesicInterpolation();
        kwargs...,
    )

Compute the Riemannian [`mean`](@ref mean(M::Manifold, args...)) of `x` using
[`GeodesicInterpolation`](@ref).
"""
mean(::ProbabilitySimplex, ::Any...)

function Statistics.mean!(
    M::ProbabilitySimplex,
    p,
    x::AbstractVector,
    w::AbstractVector;
    kwargs...,
)
    return mean!(M, p, x, w, GeodesicInterpolation(); kwargs...)
end

@doc raw"""
    project(M::ProbabilitySimplex, p, Y)

project `Y` from the embedding onto the tangent space at `p` on
the [`ProbabilitySimplex`](@ref) `M`. The formula reads

````math
\operatorname{proj}_{Δ^n}(p,Y) = Y - ⟨p,Y⟩p.
````
"""
project(::ProbabilitySimplex, ::Any, ::Any)

function project!(::ProbabilitySimplex, q, p)
    if any(x -> x <= 0, p)
        throw(
            DomainError(
                p,
                "All coordinates of point from the embedding, that should be projected, must be positive, otherwise the projection is not well defined.",
            ),
        )
    end
    q .= p ./ sum(p)
    return q
end

function project!(::ProbabilitySimplex, X, p, Y)
    X .= Y .- sum(Y) .* p
    return X
end

@doc raw"""
    representation_size(::ProbabilitySimplex{n})

return the representation size of points in the $n$-dimensional probability simplex,
i.e. an array size of `(n+1,)`.
"""
representation_size(::ProbabilitySimplex{n}) where {n} = (n + 1,)

@doc raw"""
    retract(M::ProbabilitySimplex, p, X, ::SoftmaxRetraction)

Compute a first order approximation by applying the softmax function. The formula reads

````math
\operatorname{retr}_p X = \frac{p\mathrm{e}^X}{⟨p,\mathrm{e}^X⟩},
````

where multiplication, exponentiation and division are meant elementwise.
"""
retract(::ProbabilitySimplex, ::Any, ::Any, ::SoftmaxRetraction)

function retract!(::ProbabilitySimplex, q, p, X, ::SoftmaxRetraction)
    s = zero(eltype(q))
    @inbounds for i in eachindex(q, p, X)
        q[i] = p[i] * exp(X[i])
        s += q[i]
    end
    q ./= s
    return q
end

function Base.show(io::IO, ::ProbabilitySimplex{n}) where {n}
    return print(io, "ProbabilitySimplex($(n))")
end

@doc raw"""
    zero_tangent_vector(M::ProbabilitySimplex,p)

returns the zero tangent vector in the tangent space of the point `p`  from the
[`ProbabilitySimplex`](@ref) `M`, i.e. its representation by the zero vector in the embedding.
"""
zero_tangent_vector(::ProbabilitySimplex, ::Any)

zero_tangent_vector!(M::ProbabilitySimplex, v, p) = fill!(v, 0)


"""
    struct TensorLagInterp{N,Td,T}

Generic Lagrange interpolation over an `N`-dimensional tensor grid. The
implementation uses a multidimensional generalization of the barycentric
formula.

The main constructor takes an `SVector{N,Vector{Td}}` containig the `N`
one-dimensional `nodes` and an `Array{N,T}` of the function `vals` at the tensor
product grid formed by the one-dimensional `nodes`.

# Examples:
```julia
nx = 10
ny = 12
x   = [0.5+0.5cos((2k-1)*π/2nx) for k in 1:nx] # Chebyshev nodes
y   = [0.5+0.5cos((2k-1)*π/2ny) for k in 1:ny] # Chebyshev nodes
f   = (x) -> cos(x[1]*x[2])
vals = [f((x,y)) for x in x, y in y]
p   = TensorLagInterp(SVector(x,y),vals)
p((0.1,0.2)) ≈ f((0.1,0.2))
```
"""
mutable struct TensorLagInterp{N,Td,T}
    vals::Array{T,N}
    nodes1d::NTuple{N,Vector{Td}}
    weights1d::NTuple{N,Vector{Td}}
end

interpolation_nodes(p::TensorLagInterp)                 = Iterators.product(p.nodes1d...)
interpolation_nodes(p::TensorLagInterp,d::Int,i::Int)   = p.nodes1d[d][i]
function interpolation_nodes(p::TensorLagInterp{N},I::CartesianIndex{N}) where {N}
    svector(N) do d
        i = I[d]
        p.nodes1d[d][i]
    end
end

weights1d(p::TensorLagInterp) = p.weights1d
vals(p::TensorLagInterp)      = p.vals

return_type(::TensorLagInterp{_,__,T}) where {_,__,T} = T

ambient_dimension(::TensorLagInterp{N}) where {N} = N

Base.eachindex(p::TensorLagInterp)  = eachindex(vals(p))
function Base.CartesianIndices(p::TensorLagInterp)
    sz = length.(p.nodes1d)
    CartesianIndices(sz)
end

# if passed the values and nodes, construct weights automatically
function TensorLagInterp(vals::Array{T,N},nodes::NTuple{N,Vector{Td}}) where {N,Td,T}
    weights = map(barycentric_lagrange_weights,nodes)
    TensorLagInterp{N,Td,T}(vals,nodes,weights)
end
TensorLagInterp(vals::Vector,nodes::Vector) = TensorLagInterp(vals,(nodes,))
TensorLagInterp(vals::Array,nodes::NTuple)  = TensorLagInterp(vals,SVector(nodes))

# barycentric lagrange interpolation
function (p::TensorLagInterp{N,Td,T})(x::SVector) where {N,Td,T}
    return bclag_interp(x, p.vals, p.nodes1d, p.weights1d, Val{N}(), 1, length(p.vals))
end
(p::TensorLagInterp{N,Td})(x) where {N,Td} =  p(SVector{N,Td}(x))
(p::TensorLagInterp{1})(x::Number) = p((x,))

# low-level function for multidimensional barycentric lagrange interpolation
# (adapted from FastChebInterp package)
@inline function bclag_interp(x::SVector{N,Td}, vals::Array{<:Any,N},
                            nodes::NTuple{N}, weights::NTuple{N},
                            ::Val{dim}, i1, len) where {N,Td,dim}
    T   = eltype(vals)
    n   = size(vals,dim)
    @inbounds xd = x[dim]
    @inbounds W  = weights[dim]
    @inbounds X  = nodes[dim]
    num = zero(T)
    den = zero(Td)
    if dim == 1
        for i in 1:n
            ci     = vals[i1+(i-1)]
            wi     = W[i]
            x_m_xi = xd - X[i]
            if x_m_xi == 0
                return ci
            end
            num += (wi/x_m_xi)*ci
            den += (wi/x_m_xi)
        end
        return num/den
    else
        Δi = len ÷ n # column-major stride of current dimension
        # recurse down on dimension
        dim′ = Val{dim-1}()
        for i in 1:n
            ci     = bclag_interp(x, vals, nodes, weights, dim′, i1+(i-1)*Δi, Δi)
            wi     = W[i]
            x_m_xi = xd - X[i]
            if x_m_xi == 0
                return ci
            end
            num += wi/x_m_xi*ci
            den += wi/x_m_xi
        end
        return num/den
    end
end

# algorithm on https://people.maths.ox.ac.uk/trefethen/barycentric.pdf
function barycentric_lagrange_weights(x::Vector)
    n = length(x) - 1
    w = similar(x)
    w[1] = 1.0
    for j in 1:n
        for k in 0:j-1
            w[k+1] = (x[k+1] - x[j+1]) * w[k+1]
        end
        w[j+1] = prod(0:j-1) do k
            x[j+1] - x[k+1]
        end
    end
    for j in 0:n
        w[j+1] = 1/w[j+1]
    end
    return w
end

# special nodes distributions and the corresponding weights

"""
    cheb1nodes(n,a,b)

Return the `n` Chebyshev points of the first kind on the interval `[a,b]`.
"""
function cheb1nodes(n::NTuple{N},lc,uc) where {N}
    iter = cheb1nodes_iter(n,lc,uc)
    return map(x->SVector{N,Float64}(x),iter)
end
function cheb1nodes(n::NTuple{N},rec::AbstractHyperRectangle{N}) where {N}
    cheb1nodes(n,low_corner(rec),high_corner(rec))
end
function cheb1nodes(n::Integer,a::Number,b::Number)
    xcheb::Vector{Float64} = cheb1nodes(n)
    c0 = (a+b)/2
    c1 = (b-a)/2
    return @. c0 + c1*xcheb
end
function cheb1nodes(n)
    return [-cos( (2*i-1)*π / (2*n) ) for i in 1:n]
end

function cheb1nodes_iter(n::NTuple{N},lc,uc) where {N}
    nodes1d = ntuple(d->cheb1nodes(n[d],lc[d],uc[d]),N)
    return iter = Iterators.product(nodes1d...)
end
function cheb1nodes_iter(n::NTuple{N},rec::AbstractHyperRectangle{N}) where {N}
    cheb1nodes_iter(n,low_corner(rec),high_corner(rec))
end

function cheb1weights(n::NTuple{N}) where {N}
    weights1d = ntuple(d->cheb1weights(n[d]),N)
    iter = Iterators.product(weights1d...)
    return map(x->prod(x),iter)
end
function cheb1weights(i,n)
    sgn = (-1)^(i)
    val = sin((2*i-1)*π/(2*n))
    return sgn*val
end
function cheb1weights(n)
    return [(-1)^(i)*sin((2*i-1)*π/(2*n)) for i in 1:n]
end

"""
    cheb2nodes(n,a,b)

Return the `n` Chebyshev points of the second kind on the interval `[a,b]`. The
nodes are nested in the following sense: `cheb2nodes(n,a,b) ==
cheb2nodes(2n-1,a,b)[1:2:end]`.
"""
function cheb2nodes(n::NTuple{N},lc,uc) where {N}
    iter = cheb2nodes_iter(n,lc,uc)
    return map(x->SVector{N,Float64}(x),iter)
end
function cheb2nodes(n::NTuple{N},rec::AbstractHyperRectangle{N}) where {N}
    cheb2nodes(n,low_corner(rec),high_corner(rec))
end
function cheb2nodes(n::Integer,a::Number,b::Number)
    xcheb::Vector{Float64} = cheb2nodes(n)
    c0 = (a+b)/2
    c1 = (b-a)/2
    return @. c0 + c1*xcheb
end
function cheb2nodes(n)
    return [cos((i-1)*π /(n-1)) for i in 1:n]
end

function cheb2nodes_iter(n::NTuple{N},lc,uc) where {N}
    nodes1d = ntuple(d->cheb2nodes(n[d],lc[d],uc[d]),N)
    iter = Iterators.product(nodes1d...)
    return iter
end
function cheb2nodes_iter(n::NTuple{N},rec::AbstractHyperRectangle{N}) where {N}
    return cheb2nodes_iter(n,low_corner(rec),high_corner(rec))
end

function cheb2weights(n::NTuple{N}) where {N}
    weights1d = ntuple(d->cheb2weights(n[d]),N)
    iter = Iterators.product(weights1d...)
    return map(x->prod(x),iter)
end

function cheb2weights(i,n)
    sgn = (-1)^(i)
    val = i==n || i == 1 ? 0.5 : 1.0
    return sgn*val
end

function cheb2weights(n)
    return [cheb2weights(i,n) for i in 1:n]
end


using UnivariateFunctions: create_linear_interpolation, create_constant_interpolation_to_right, create_quadratic_spline
using UnivariateFunctions: create_constant_interpolation_to_left, years_from_global_base, evaluate
using Dates

tol = 10*eps()
const global_base_date = Date(2000,1,1)
StartDate = Date(2018, 7, 21)

x = StartDate .+ Dates.Day.(2 .* (1:1000 .- 1))

function ff(x::Date)
    days_between = years_from_global_base(x)
    return log(days_between) + sqrt(days_between)
end
y = ff.(x)

spline = create_linear_interpolation(x,y)
# Test if interpolating
all(abs.(evaluate.(Ref(spline), x) .- y) .< tol)

x_float = years_from_global_base.(x)
coefficient_in_first_interval = (y[2] - y[1])/(x_float[2] - x_float[1])
(coefficient_in_first_interval - spline.functions_[1].functions_[2].a_) < tol
other_coefficient = (y[11] - y[10])/(x_float[11] - x_float[10])
(other_coefficient - spline.functions_[10].functions_[2].a_) < tol

left_const = create_constant_interpolation_to_left(x,y)
right_const = create_constant_interpolation_to_right(x,y)
all(abs.(evaluate.(Ref(left_const), x_float .- tol) .- evaluate.(Ref(right_const), x_float)) .< tol)

# Making interpolation splies with date periods.
xx = Array{Month,1}([Month(1), Month(2), Month(3)])
xx2 = Array{DatePeriod,1}([Month(1), Month(2), Year(1)])
yy = [1,2,3]
abs(create_constant_interpolation_to_left(xx,yy)(Day(45)) - 2.0) < eps()
abs(evaluate(create_constant_interpolation_to_right(xx,yy), Day(15)) - 1.0) < eps()
abs(evaluate(create_linear_interpolation(xx,yy), Month(1)) - 1.0) < eps()
abs(evaluate(create_quadratic_spline(xx,yy), Month(2)) - 2.0) < eps()

abs(evaluate(create_constant_interpolation_to_left(xx2,yy), Day(45)) - 2.0) < eps()
abs(evaluate(create_constant_interpolation_to_right(xx2,yy), Day(15)) - 1.0) < eps()
abs(evaluate(create_linear_interpolation(xx2,yy), Month(1)) - 1.0) < eps()
abs(evaluate(create_quadratic_spline(xx2,yy), Month(2)) - 2.0) < eps()


type Monomial <: LsmBasisSystemPolynomType end

type MonomialFunction{I <: Integer} <: LSMBasisSystemFunction
  order::I
end

function call(m::MonomialFunction, x::Float64)
  ret = 1.0
  for i = 1:m.order
    ret *= x
  end

  return ret
end

get_type(::Monomial) = MonomialFunction{Int}

function path_basis_system!(::Monomial, order::Int, v::Vector)
  for i = 1:order + 1
    v[i] = MonomialFunction(i - 1)
  end

  return v
end


function toh05()
#
#  generalization of K.C. Toh and L. N. Trefethen
#
    k = [(-19:0)...]
    z = exp2.(k)
    p = poly(z)
    z = [z ones(20)]

    p, PolyZeros(z)
end


mutable struct State{F <: AbstractFloat, V <: AbstractVector{<:Real}}
  p::F
  beta::Bool
  gammaC::F
  gammaT::F
  etaC::V  # K
  etaT::V  # K
  lambdaC::Vector{Int}  # NC
  lambdaT::Vector{Int}  # NT
  mu::V  # K
  nu::V  # K
  omega::V  # K. alt variance.
  psi::V  # K. alt skew.
  vC::V  # NC
  vT::V  # NT
  zetaC::V  # NC
  zetaT::V  # NT
  tau::F
  sigma::V
  phi::V
end

function State(data::Data, prior::Prior)
  p = rand(prior.p)
  beta = true
  gammaC = rand(prior.gamma)
  gammaT = rand(prior.gamma)
  etaC = rand(prior.eta)
  etaT = rand(prior.eta)
  lambdaC = rand(Categorical(etaC), data.NC_finite)
  lambdaT = rand(Categorical(etaT), data.NT_finite)
  mu = rand(prior.mu, prior.K)
  nu = rand(prior.nu, prior.K)
  tau = rand(prior.tau)
  prior_omega = InverseGamma(prior.a_omega, tau)
  omega = rand(prior_omega, prior.K)
  psi = rand(prior.psi, prior.K)
  vC = [rand(Gamma(nu[k]/2, 2/nu[k])) for k in lambdaC]
  vT = [rand(Gamma(nu[k]/2, 2/nu[k])) for k in lambdaT]
  zetaC = [rand(truncated(Normal(0, sqrt(1/v)), 0, Inf)) for v in vC]
  zetaT = [rand(truncated(Normal(0, sqrt(1/v)), 0, Inf)) for v in vT]
  sigma = Util.scalefromaltskewt.(sqrt.(omega), psi)
  phi = Util.skewfromaltskewt.(sqrt.(omega), psi)

  state = State(p, beta, gammaC, gammaT, etaC, etaT, lambdaC, lambdaT,
                mu, nu, omega, psi, vC, vT, zetaC, zetaT, tau, sigma, phi)

  return state
end

ref_lambda(state::State, i::Char) = (i == 'C') ? state.lambdaC : state.lambdaT
ref_v(state::State, i::Char) = (i == 'C') ? state.vC : state.vT
ref_zeta(state::State, i::Char) = (i == 'C') ? state.zetaC : state.zetaT

ref_eta(state::State, i::Char) = (i == 'C') ? state.etaC : state.etaT
ref_gamma(state::State, i::Char) = (i == 'C') ? state.gammaC : state.gammaT


function shuffle!(r::AbstractRNG, a::AbstractVector)
    for i = length(a):-1:2
        j = rand(r, 1:i)
        a[i], a[j] = a[j], a[i]
    end
    return a
end


using JLD2, Random, LinearAlgebra, Statistics, CSV, DataFrames, FreqTables, Distributions, Kronecker, JLD, GLM, HTTP, CategoricalArrays, TexTables

# Worked With Ahmed

function q1()
# 1,A
Random.seed!(1234)
A = rand(Uniform(-5,10),10,7)
B = rand(Normal(-2,15),10,7)
C = [A[1:5,1:5] B[1:5,6:7]]
D = (A .<= 0) .*A
# 1,B
size(A)
# 1,C
unique(D)
# 1,D
E = reshape(B,70,1)
E2 = vec(B)
# 1,E
F = reshape([A B], 10, 7, 2)
# 1,F
F = permutedims(F, [3, 1, 2])
#  1,G
G = kronecker(B,C)
# 1,H
save("matrixpractice.jld", "A", A, "B", B, "C", C, "D",
 D, "E", E, "F", F, "G", G )
 # 1,I
save("firstmatrix.jld", "A", A, "B", B, "C", C, "D", D)
# 1,J
CSV.write("Cmatrix.csv", DataFrame(C,:auto))
# 1,K
CSV.write("Dmatrix.dat", DataFrame(D,:auto), delim = ' ')
# 1,L

return A, B, C, D

end

A, B, C, D = q1()

function q2(A, B, C)

#2,A
AB = [A[i,j]*B[i,j] for i in 1:10, j in 1:7]
#2,B
Cprime = [C[i,j] for i in 1:5, j in 1:7 if -5 <= C[i,j]<= 5]

C[-5 .<= C .<= 5]
#2,C
N, K, T = 15169, 6, 5

X = zeros(N,K,T);

x5 = rand(Binomial(20,0.6),N)
x6 = rand(Binomial(20,0.5),N)

    for k in 1:T
    X[:,1,k] = ones(N)

    if rand(1)[1] <= 75*(6-k)/5
            X[:,2,k] = ones(N)
    end

    X[:,3,k] = rand(Normal(15 + k - 1, 5*(k-1)),N)

    X[:,4,k] = rand(Normal(pi*(6-k)/3, 1/exp(1)),N)

    # Columns 1, 5, and 6 remain stationary over time
    X[:,5, k] = x5
    X[:,6, k] = x6
end

#2,D

β = zeros(K,T);
β[1,:] = [1:0.25:2;];
β[2,:] = [log(j) for j in 1:T];
β[3,:] = [-sqrt(j) for j in 1:T];
β[4,:] = [exp(j)-exp(j+1) for j in 1:T];
β[5,:] = [j for j in 1:T];
β[6,:] = [j/3 for j in 1:T];

β

#2,E

Y = [X[:,:,t] * β[:,t] + rand(Normal(0,0.36),N,1) for t in 1:T]

end
q2(A, B, C)

function q3()
        #3,A

       # cd("C:/Users/wat20/OneDrive/Documents/ProblemSets/PS1-julia-intro")
        df = CSV.read("nlsw88.csv", DataFrame)

        #is missing (df)
        #[print(string(":", x, ", ")) for x in names(df)]

        #3,B
        NotMarrPer = 1 - sum(df[:,"married"])/size(df)[1];
        GradPer =sum(df[:,"collgrad"])/size(df)[1];

        println("Percent not married = ", round(NotMarrPer, digits = 4) * 100,
        "\nCollege graduate percentage = ", round(GradPer, digits = 4)* 100)

        #3,C)
        #note: tabulate command is undefined
        # source of freq. tables:
        # https://stackoverflow.com/questions/34654489/rs-table-function-in-julia-for-dataframes

        prop(freqtable(df.race))

        # d)

        intrq = (describe(df,:q75)[:,2] - describe(df,:q25)[:,2])
        summarystats = [describe(df,:nmissing, :mean, :median, :std, :min, :max, :nunique) intrq]
        rename(summarystats, :x1=> :itrq)

        #3,E
        ## Cross tab between industry and occupation
        freqtable(df.industry, df.occupation)

        #3,F
        # source: https://dataframes.juliadata.org/stable/man/split_apply_combine/

        qq = df[:, [:race, :occupation, :wage]];
        gdf = groupby(qq, [:race, :occupation]);
        combine(gdf, :wage => mean)
end

q3()

#4
fmat = load("firstmatrix.jld")
#keys(fmat)
#fmat["B"]

function q4()

#4A-E
function matrixops(A, B)
    #element by element product of inputs, A and b
    #input matrices product
    #input matrices sum
            if(size(A) != size(B))
            println("inputs must have the same size")
            else
                return println(
                    "A .* B = ", A .* B,
                    "\n\nA' * B = ", A' * B,
                    "\n\nA + B = ", A + B)
            end
        end

matrixops(A, B)

#4,F
# we get the error message back matrixops(C, D)

# g)

    matrixops(
    convert(Array,df.ttl_exp ),
    convert(Array,df.wage)
    )
end

q4()


function ueberpruefeerror(Fehler,PhantomVektor,RekonstruktionVektor)

tol = 1e-12;

if norm(Fehler) - norm(PhantomVektor-RekonstruktionVektor) < tol
    println("Beachtet, dass wir Vektoren betrachten. Dafür müsst ihr die Länge der Vektoren betrachten.")
    println("Zudem habt ihr den absolten Fehler berechnet. Berechnet den Fehler relativ zum originalen Wert.")
elseif Fehler - norm(PhantomVektor-RekonstruktionVektor) < tol
    println("Ihr habt den absolten Fehler berechnet. Berechnet den Fehler relativ zum originalen Wert.")
elseif Fehler - norm(PhantomVektor-RekonstruktionVektor)/norm(PhantomVektor) < tol
    println("Fast richtig. Lasst euch das Ergebnis in Prozent ausgeben.")
elseif Fehler - norm(PhantomVektor-RekonstruktionVektor)/norm(PhantomVektor)*100 < tol
    Fehler2 = round(Fehler;digits=4);
    println("Eure Formel für den relativen Fehler ist korrekt! Der Fehler beträgt: ", Fehler2)
else
    println("Googlt noch einmal nach Fehlerarten, die ihr kennt und schaut in Workspace nach welche Variablen ihr habt, die euch hier weiterhelfen können.")
end
      # return Fehler2;

end