using LAP_julia

viewimg(img) = colorview(Gray, img)

img_size = (250,255)

view(img)
view(imgw)

timer = TimerOutput(string(args_dict["reg_alg"]))
method_kwargs = Dict(:display => true, :timer => timer)
img, imgw, flow = gen_init(:spaghetti, img_args=[img_size], flow_args=[25])
flow_est, source_reg, timer, time_in_secs = test_registration_alg(sparse_pflap,
                                                                  img,
                                                                  imgw,
                                                                  flow,
                                                                  timer=method_kwargs[:timer],
                                                                  method_kwargs=method_kwargs)

showflow(flow_est)
showflow(flow.*(-1))

count(isnan.([real.(flow_est) imag.(flow_est)]))

LAP_julia.vec_len.(flow_est)
view(source_reg)
view(img)
view(imgw)

##
base_path_bikes = "/Users/MrTrololord/Downloads/bikes"

img = load(joinpath(base_path_bikes, "img1.png"))
for k in 2:2
    imgw = load(joinpath(base_path_bikes, "img$(k).png"))
    global imgw
end

imgw

img = Float64.(Gray.(img))
imgw = Float64.(Gray.(imgw))

timer = TimerOutput("sparse_pflap");
method_kwargs = Dict(:display => false, :timer => timer)
flow_est, source_reg, timer, time_in_secs = time_reg_alg(pflap, img, imgw, timer=method_kwargs[:timer], method_kwargs=method_kwargs)

view(imgw)
view(source_reg)
view(img)

showflow(flow_est)

##

base_path_zhang = "/Users/MrTrololord/Downloads/dataset/"
fnames = ["yard", "PGH", "testchart"]
imgs = Array{Image}(undef, size(fnames))
imgws = Array{Image}(undef, size(fnames))

for (k, name) in zip(1:length(fnames), fnames)
    imgs[k] = load(joinpath(base_path_zhang, "$(name)_t.jpg"))
    imgws[k] = load(joinpath(base_path_zhang, "$(name)_s.jpg"))
    global img, imgw
end

img = imgs[3]
imgw = imgws[3]

img = Float64.(Gray.(img))
imgw = Float64.(Gray.(imgw))

##
timer = TimerOutput("pflap");
method_kwargs = Dict(:display => false, :timer => timer)
flow_est, source_reg, timer, time_in_secs = time_reg_alg(pflap, img, imgw, timer=method_kwargs[:timer], method_kwargs=method_kwargs)

source_reg_zeros = warp_img(imgw, -real.(flow_est), -imag.(flow_est), border_strat=:zeros)
view(source_reg_zeros)

view(imgw)
view(source_reg)
viewimg(img)

showflow(imresize(flow_est, (200,200)))

##
# @info "size from $(size)"
(img, ratio), (imgw, ratio) = resize_to_diag_size(img, 400), resize_to_diag_size(imgw, 400)

img, imgw = Gray.(img), Gray.(imgw)

timer = TimerOutput("sparse_pflap");
method_kwargs = Dict(:display => false, :timer => timer)
flow_est, source_reg, timer, time_in_secs = time_reg_alg(sparse_pflap_psnr, imgw, img, timer=method_kwargs[:timer], method_kwargs=method_kwargs)

source_reg_zeros = warp_img(img, -real.(flow_est), -imag.(flow_est), border_strat=:zeros)
viewimg(source_reg_zeros)
viewimg(imgw)

flow_est

##
# load float from text

function parse_numbers(s)
    pieces = split(s, ' ', keepempty=false)
    map(pieces) do piece
        parse(Float64, piece)
    end
end

function load_H(path)
    lines = readlines(path)
    H = Array{Float64}(undef, (3,3))
    for (k, line) in zip(1:3, lines)
        H[k,:] = parse_numbers(line)
    end
    return H
end

function make_flow_from_H(H, size)
    warp_perspective(y, x) = (H[1,1]*x + H[1,2]*y + H[1,3])/(H[3,1]*x + H[3,2]*y + H[3,3]),
                             (H[2,1]*x + H[2,2]*y + H[2,3])/(H[3,1]*x + H[3,2]*y + H[3,3])
    flow = Array{Complex{Float64},2}(undef, size)
    X = ones(size[1]) * collect(range(0,1,length=size[2]))'
    Y = collect(range(0,1,length=size[1])) * ones(size[2])'
    for ind in CartesianIndices(flow)
        calc = [Y[ind], X[ind]] .- warp_perspective(Y[ind], X[ind])
        flow[ind] = calc[1] + im*calc[2]
    end
    return flow
end

H = load_H(joinpath(base_path_bikes, "H1to2p"))
out = make_flow_from_H(H, size(img))

showflow(out)

img
imgw

asdf = warp_img(img, -imag.(out), -real.(out), border_strat=:zeros)

Gray.(img) .- Gray.(asdf)

##
using CoordinateTransformations, StaticArrays, ImageTransformations, LinearAlgebra;
struct Homography{T} <: AbstractAffineMap
    m::SMatrix{3, 3, T, 9}
end

# having T in the RHS forces calling the constructor again
Homography(m::AbstractMatrix{}) = Homography(m);
function (trans::Homography{M})(x::SVector{3}) where M <: Real
    out = trans.m * x;
    out = out / out[end];
    SVector{2}(out[1:2])
end

function (trans::Homography{M})(x::SVector{2}) where M <: Real
    trans(SVector{3}([x[1], x[2], 1.0]))
end

function (trans::Homography{M})(x::CartesianIndex{2}) where M <: Real
    trans(SVector{3}([collect(x.I); 1]))
end

function (trans::Homography{M})(x::Tuple{Int, Int}) where M <: Real
    trans(CartesianIndex{2}(x))
end

function (trans::Homography{M})(x::Array{CartesianIndex{2}, 1}) where M <: Real
    CartesianIndex{2}.([tuple(y...) for y in trunc.(Int, collect.(trans.(x)))])
end

function Base.inv(trans::Homography)
    i = inv(trans.m);
    Homography(i ./ i[end])
end

Homo = Homography{Float64}(H)

img1_warp = ImageTransformations.warp(imgw, Homo)

img

##

## posun homografii

base_path_bikes = "/Users/MrTrololord/Downloads/bikes"

img = load(joinpath(base_path_bikes, "img1.png"))
k = 2
imgw = load(joinpath(base_path_bikes, "img$(k).png"))

function parse_numbers(s)
    pieces = split(s, ' ', keepempty=false)
    map(pieces) do piece
        parse(Float64, piece)
    end
end

function load_H(path)
    lines = readlines(path)
    H = Array{Float64}(undef, (3,3))
    for (k, line) in zip(1:3, lines)
        H[k,:] = parse_numbers(line)
    end
    return H
end

function make_flow_from_H(H, size)
    warp_perspective(y, x) = (H[1,1]*x + H[1,2]*y + H[1,3])/(H[3,1]*x + H[3,2]*y + H[3,3]),
                             (H[2,1]*x + H[2,2]*y + H[2,3])/(H[3,1]*x + H[3,2]*y + H[3,3])
    flow = Array{Complex{Float64},2}(undef, size)
    X = ones(size[1]) * collect(range(0,1,length=size[2]))'
    Y = collect(range(0,1,length=size[1])) * ones(size[2])'
    for ind in CartesianIndices(flow)
        calc = [Y[ind], X[ind]] .- warp_perspective(Y[ind], X[ind])
        flow[ind] = calc[1] + im*calc[2]
    end
    return flow
end

H = load_H(joinpath(base_path_bikes, "H1to2p"))
out = make_flow_from_H(H, size(img))

warped_img = warp_img(img, -imag.(out), -real.(out), border_strat=:zeros)
imgw

## posun mym alg

viewimg(img) = colorview(Gray, img)

img, imgw = Gray.(img), Gray.(imgw)

timer = TimerOutput("sparse_pflap");
method_kwargs = Dict(:display => false, :timer => timer)
flow_est, source_reg, timer, time_in_secs = time_reg_alg(sparse_pflap_psnr, imgw, img, timer=method_kwargs[:timer], method_kwargs=method_kwargs)

source_reg_zeros = warp_img(img, -real.(flow_est), -imag.(flow_est), border_strat=:zeros)
viewimg(source_reg_zeros)
viewimg(imgw)

# posun mym alg je ma lepsi vysledky?? nesmysl -> neco spatne


function bat_fit(hist, fit_function, prior; nsamples=10^5, nchains=4, params_first=true, norm_expected=true)
    likelihood = let h = hist, f = fit_function
        # Histogram counts for each bin as an array:
        observed_counts = h.weights

        # Histogram binning:
        bin_edges = h.edges[1]
        bin_edges_left = bin_edges[1:end-1]
        bin_edges_right = bin_edges[2:end]
        bin_widths = bin_edges_right - bin_edges_left
        bin_centers = midpoints(bin_edges)

        params -> begin
            if params_first
                expected_counts = bin_widths .* f(params, bin_centers)
            else
                expected_counts = bin_widths .* f(bin_centers, params)
            end
            if norm_expected
                expected_counts ./= sum(expected_counts)
                expected_counts .*= sum(observed_counts)
            end
            ll_value = sum(logpdf.(Poisson.(expected_counts), observed_counts))
            return LogDVal(ll_value)
        end
    end
    parshapes = varshape(prior)
    posterior = PosteriorDensity(likelihood, prior);
    return bat_sample(posterior, MCMCSampling(mcalg = MetropolisHastings(), nsteps = nsamples, nchains = nchains)).result;
end

#
# Wrapper for fits that might need several iterations
#
function bat_fit_wrapper(histogram, model, prior, nsamples_list; ch_str="XX", params_first=false, norm_expected=false)
    parameters, errors = [], []
    run = 1
    for ns in nsamples_list
        @info("Ch" * ch_str * " | try " * string(run) * " of " * string(length(nsamples_list)) * ": " * string(ns) * " samples")
        samples = try bat_fit(histogram, model, prior; nsamples=ns, nchains=4, params_first=params_first, norm_expected=norm_expected);
        catch
            IJulia.clear_output(true)
            run += 1
            continue
        end
        IJulia.clear_output(true)
        @info("Success: Ch" * ch_str * " | " * string(ns) * " samples")
        parameters = mode(samples)
        errors = std(samples)
        break
    end
    return [parameters, errors]
end

#
# Peak model with low E tail
#
function peak(E, par)
    n = try par[1] catch; par[1][1] end
    σ = try par[2] catch; par[2][1] end
    µ = try par[3] catch; par[3][1] end

    bkg_l = try par[4] catch; par[4][1] end
    bkg_r = try par[5] catch; par[5][1] end
    b = (bkg_r - bkg_l) / (maximum(E) - minimum(E))
    a = bkg_l - b * minimum(E)

    c = try par[6] catch; par[6][1] end
    d = try par[7] catch; par[7][1] end
    β = try par[8] catch; par[8][1] end

    d = b > 0 ? -d : d
    correction = d < 0 ? -1*minimum(d/2 .* erfc.((E .- µ) ./ (sqrt(2)*σ))) : 0.0
    result = @. n / sqrt(2*pi*σ^2) * exp(-(E - µ)^2 / (2*σ^2)) + # Gaussian
        a + b*E + # Linear bkg
        d/2 * erfc((E - µ) / (sqrt(2)*σ)) + correction +  # Low-E step
        c/(2*β) * exp((E-µ) / β + σ^2 / (2*β^2)) * erfc((E - µ) / (sqrt(2)*σ) + σ / (sqrt(2)*β)) # Low-E tail
end
function peak_twisted(par, E)
    n = try par[1] catch; par[1][1] end
    σ = try par[2] catch; par[2][1] end
    µ = try par[3] catch; par[3][1] end

    bkg_l = try par[4] catch; par[4][1] end
    bkg_r = try par[5] catch; par[5][1] end
    b = (bkg_r - bkg_l) / (maximum(E) - minimum(E))
    a = bkg_l - b * minimum(E)

    c = try par[6] catch; par[6][1] end
    d = try par[7] catch; par[7][1] end
    β = try par[8] catch; par[8][1] end

    d = b > 0 ? -d : d
    correction = d < 0 ? -1*minimum(d/2 .* erfc.((E .- µ) ./ (sqrt(2)*σ))) : 0.0
    result = @. n / sqrt(2*pi*σ^2) * exp(-(E - µ)^2 / (2*σ^2)) + # Gaussian
        a + b*E + # Linear bkg
        d/2 * erfc((E - µ) / (sqrt(2)*σ)) + correction +  # Low-E step
        c/(2*β) * exp((E-µ) / β + σ^2 / (2*β^2)) * erfc((E - µ) / (sqrt(2)*σ) + σ / (sqrt(2)*β)) # Low-E tail
end

function broad_peak(E, par)
    n = try par[1] catch; par[1][1] end
    σ = try par[2] catch; par[2][1] end
    µ = try par[3] catch; par[3][1] end

    bkg_l = try par[4] catch; par[4][1] end
    bkg_r = try par[5] catch; par[5][1] end
    b = (bkg_r - bkg_l) / (maximum(E) - minimum(E))
    a = bkg_l - b * minimum(E)

    c = try par[6] catch; par[6][1] end
    d = try par[7] catch; par[7][1] end
    β = try par[8] catch; par[8][1] end

    n2 = try par[9]  catch; par[9][1] end
    σ2 = try par[10] catch; par[10][1] end

    d = b > 0 ? -d : d
    correction = d < 0 ? -1*minimum(d/2 .* erfc.((E .- µ) ./ (sqrt(2)*σ))) : 0.0
    result = @. n / sqrt(2*pi*σ^2) * exp(-(E - µ)^2 / (2*σ^2)) + # Gaussian
        n2 / sqrt(2*pi*σ2^2) * exp(-(E - µ)^2 / (2*σ2^2)) + # Gaussian 2
        a + b*E + # Linear bkg
        d/2 * erfc((E - µ) / (sqrt(2)*σ)) .+ correction +  # Low-E step
        c/(2*β) * exp((E-µ) / β + σ^2 / (2*β^2)) * erfc((E - µ) / (sqrt(2)*σ) + σ / (sqrt(2)*β)) # Low-E tail
end

function tail_model(x, par)
    scale = try par[1] catch; par[1][1] end
    σ     = try par[2] catch; par[2][1] end
    μ     = try par[3] catch; par[3][1] end
    C     = try par[4] catch; par[4][1] end
    l     = try par[5] catch; par[5][1] end
    e     = try par[6] catch; par[6][1] end
    t     = try par[7] catch; par[7][1] end
    d     = 0
    return @. scale * exp(-0.5 * ((x - μ)^2) / (σ^2)) / (sqrt(2 * π * σ^2)) + C * (exp(e*(x-l)) + d) / (exp((x-l)*t) +1)
end

function peak_model(x, par)
    scale = par[1]
    σ     = par[2]
    μ     = par[3]

    #step = x <= µ ? par[4] : par[5]
    if length(x) > 1
        step = []
        for i in x
            push!(step, i <= µ ? par[4] : par[5])
        end
        #mwa_step = isodd(round(length(x)/4)) ? Int.(round(length(x)/4)) : Int.(round(length(x)/4)-1)
        s = 4 * par[2] / 0.1
        mwa_step = isodd(Int.(round(s))) ? Int.(round(s)) : Int.(round(s)) + 1
        step = mwa(Array{Float64,1}(step),mwa_step)
    else
        step = x <= µ ? par[4] : par[5]
    end
    return scale .* exp.(-0.5 * ((x .- μ).^2) / (σ^2)) ./ (sqrt(2 * π * σ^2)) .+ step
end
function peak_model2(par, x)
    scale = par.scale
    σ     = par.σ
    μ     = par.µ
    cp0   = par.cp0
    #m0    = par.m0
    cp1   = par.cp1
    #m1    = par.m1
#     s     = par.s

    if length(x) > 1
        step = []
        for i in x
            push!(step, i <= µ ? cp0 : cp1)
        end
        s = 4 * σ / 0.1
#         mwa_step = isodd(round(length(x)/4)) ? Int.(round(length(x)/4)) : Int.(round(length(x)/4)-1)
        mwa_step = isodd(Int.(round(s))) ? Int.(round(s)) : Int.(round(s)) + 1
        step = mwa(Array{Float64,1}(step),mwa_step)
    else
        step = x <= µ ? cp0 : cp1
    end
    return @. scale * exp(-0.5 * ((x - μ)^2) / (σ^2)) / (sqrt(2 * π * σ^2)) + step
end

function model(x, par)
    scale = par[1]
    σ     = par[2]
    μ     = par[3]
    cp0   = par[4]
    m1    = par[5]
    return @. scale * exp(-0.5 * ((x - μ)^2) / (σ^2)) / (sqrt(2 * π * σ^2)) + m1 * x + cp0
end
function model2(par, x)
    scale = par.scale
    σ     = par.σ
    μ     = par.µ
    cp0   = par.cp0
    m1    = par.m1
    #m2    = par[9]
    return @. scale * exp(-0.5 * ((x - μ)^2) / (σ^2)) / (sqrt(2 * π * σ^2)) + m1 * x + cp0
end
function double_model(x, par)
    scale = par[1]
    σ     = par[2]
    μ     = par[3]
    scale2 = par[4]
    σ2     = par[5]
    μ2     = par[6]
    cp0   = par[7]
    m1    = par[8]
    #m2    = par[9]
    return @. scale * exp(-0.5 * ((x - μ)^2) / (σ^2)) / (sqrt(2 * π * σ^2)) + scale2 * exp(-0.5 * ((x - μ2)^2) / (σ2^2)) / (sqrt(2 * π * σ2^2)) + m1 * x + cp0# + m2 * x
end
function double_model2(par,x)
    scale = par.scale
    σ     = par.σ
    μ     = par.µ
    scale2 = par.scale2
    σ2     = par.σ2
    μ2     = par.µ2
    cp0   = par.cp0
    m1    = par.m1
    #m2    = par[9]
    return @. scale * exp(-0.5 * ((x - μ)^2) / (σ^2)) / (sqrt(2 * π * σ^2)) + scale2 * exp(-0.5 * ((x - μ2)^2) / (σ2^2)) / (sqrt(2 * π * σ2^2)) + m1 * x + cp0# + m2 * x
end
# function tail_model(x, par)
#     scale = par[1]
#     σ     = par[2]
#     μ     = par[3]
#     C     = par[4]
#     l     = par[5]
#     e     = par[6]
#     t     = par[7]
#     d     = 0
#     return @. scale * exp(-0.5 * ((x - μ)^2) / (σ^2)) / (sqrt(2 * π * σ^2)) + C * (exp(e*(x-l)) + d) / (exp((x-l)/t) +1)
# end
# function tail_model2(par, x)
#     scale = par.scale
#     σ     = par.σ
#     μ     = par.µ
#     C     = par.C
#     l     = par.l
#     e     = par.e
#     t     = par.t
#     d     = 0
#     return @. scale * exp(-0.5 * ((x - μ)^2) / (σ^2)) / (sqrt(2 * π * σ^2)) + C * (exp(e*(x-l)) + d) / (exp((x-l)/t) +1)
# end
@. linmodel(x, p) = p[1]*x + p[2]
@. hypmodel_simple(x, p) = sqrt(p[1]+p[2]/x)
@. hypmodel(x, p) = sqrt(p[1]+p[2]/(x^2))
@. poly3model(x, p) = p[1]*x^3 + p[2]*x^2 + p[3]*x + p[4]
@. atanmodel(x, p) = p[1] * atan(p[2]*x + p[3]) + p[4]
@. sqrt_fct(x,p) = sqrt( p[1] + p[2] * x + p[3] * x * x )

function fit_bat(hist, fit_function, prior, fit_ranges; nsamples=1*10^5, nchains=4, second_ff=fit_function)

    likelihood = let h = hist, f = fit_function
        # Histogram counts for each bin as an array:
        observed_counts = h.weights

        # Histogram binning:
        bin_edges = h.edges[1]
        bin_edges_left = bin_edges[1:end-1]
        bin_edges_right = bin_edges[2:end]
        bin_widths = bin_edges_right - bin_edges_left
        bin_centers = (bin_edges_right + bin_edges_left) / 2

        params -> begin
            # Log-likelihood for a single bin:
            function bin_log_likelihood(i)
                # Simple mid-point rule integration of fit function `f` over bin:
                expected_counts = bin_widths[i] * f(params, bin_centers[i])
                logpdf(Poisson(expected_counts), observed_counts[i])
            end

            # Sum log-likelihood over bins:
            idxs = eachindex(observed_counts)
            ll_value = bin_log_likelihood(idxs[1])
            for i in idxs[2:end]
                ll_value += bin_log_likelihood(i)
            end

            # Wrap `ll_value` in `LogDVal` so BAT knows it's a log density-value.
            return LogDVal(ll_value)
        end
    end
    parshapes = varshape(prior)
    posterior = PosteriorDensity(likelihood, prior);

    pretunesamples = 10000
    max_ncycles = 30
    T = Float64
    tuning = BAT.AdaptiveMetropolisTuning(
        r = 1.0,
        λ = 0.5,
        α = 0.05..0.9,
        β = 1.5,
        c = 1e-4..1e2
    )
    convergence = BAT.BrooksGelmanConvergence(
        threshold = T(5),
        corrected = false
    )
    init = BAT.MCMCInitStrategy(
        init_tries_per_chain = 8..128,
        max_nsamples_init = pretunesamples,
        max_nsteps_init = pretunesamples * 10,
        max_time_init = Inf
    )
    burnin = BAT.MCMCBurninStrategy(
        max_nsamples_per_cycle = pretunesamples,
        max_nsteps_per_cycle = pretunesamples * 10,
        max_time_per_cycle = Inf,
        max_ncycles = max_ncycles
    )
    parshapes = varshape(prior)
    posterior = PosteriorDensity(likelihood, prior);
    nsamples = 1*10^4
    nchains = 4

    samples = bat_sample(Philox4x((123, 456)), posterior, (nsamples, nchains), MetropolisHastings(),
        max_nsteps = 10 * nsamples,
        max_time = Inf,
        tuning = tuning,
        init = init,
        burnin = burnin,
        convergence = convergence,
    ).result;
    params = mode(samples)[1]
    error = std(samples)[1]
    # println("Truth: $p0")
    println("Mode: $(mode(samples))")
    println("Stddev: $(std(samples))")

    fitf = RadiationSpectra.FitFunction{Float64}( second_ff, 1, length(params));
    # set_parameter_bounds!(fitf, bounds)
    RadiationSpectra.set_fitranges!(fitf, (fit_ranges, ) )

    par = []
    err = []
    for x in eachindex(params)
        push!(par, params[x])
        push!(err, error[x])
    end
#     params = [params.scale[1], params.σ[1], params.µ[1], params.scale2[1], params.σ2[1], params.µ2[1], params.cp0[1], params.m1[1]]
#     error = [error.scale[1], error.σ[1], error.µ[1], error.scale2[1], error.σ2[1], error.µ2[1], error.cp0[1], error.m1[1]]
    fitf.fitted_parameters = par

    return par, err, fitf
end


export lsqr

"""
x,flag,err,iter,resvec = lsqr(A,b,;x=[],maxIter=20,atol=1e-10,btol=1e-10,condlim=1e6,out=1,doBidiag=false,
                              storeInterm::Bool=false)

Least Squares QR method (LSQR) for solving

   min_x || A x - b ||,

where A is assumed is either over or underdetermined.

Implementation follows and comments refer to Algorithm LSQR in:

Paige, C. C., & Saunders, M. A. (1982).
LSQR: An Algorithm for Sparse Linear Equations and Sparse Least Squares.
ACM Transactions on Mathematical Software (TOMS), 8(1), 43–71. doi:10.1145/355984.355989

Required input:

  A       - function in (flag,v,alpha,x).
            if flag=='F'
                 x <- A*v + alpha*x (overwrite x)
            elseif flag=='T'
                 x <- A'*v + alpha*x (overwriting x)
            end
  b       - right hand side vector

Optional input:

  x           - starting guess (optional)
  maxIter     - maximum number of iterations
  atol        - error tolerance for (estimated) residual norm (for compatible systems)
  btol        - error tolerance for (estimated) gradient norm (for incompatible systems)
  condlim     - for stopping based on (estimated) condition number
  out         - flag for output (-1: no output, 0: only errors, 1: final status, 2: residual norm at each iteration)
  doBidiag    - returns the bidiagonalization of A at the final iteration
  storeInterm - store and return intermediate iterates

Output:

 x       - approximate solution
 flag    - exit flag ( 0 : stopping based on atol or btol,
                      -1 : maxIter reached without converging
                      -2 : condition number became too large
                      -9 : right hand side was zero)
 his     - status at each iteration
 U,S,V   - bidiagonalization (only build if doBidiag==true)
"""
function lsqr(A::SparseMatrixCSC,b::Vector;kwargs...)
    # build function that is compatible with lsqr
    ATv = zeros(size(A,2)) # output for A'*x

    # Af(flag,x,a=0.0,v=ATv) = (flag=='F') ? A_mul_B!(1.0,A,x,a,v) : At_mul_B!(1.0,A,x,a,v)
    Af(flag,x,a=0.0,v=ATv) = (flag=='F') ? mul!(v,A,x,1.0,a) : mul!(v,transpose(A),x,1.0,a)

    return lsqr(Af,b;kwargs...)

end

function lsqr(A::Any,b::Vector;kwargs...)
    # build function that is compatible with lsqr
    Af(flag,x,a=0.0,v=0.0) = (flag=='F') ? A*x.+a*v : A'*x.+a*v

    return lsqr(Af,b;kwargs...)

end

function lsqr(A::Function,b::Vector;x=[],maxIter=20,atol=1e-10,btol=1e-10,condlim=1e6,out=0,doBidiag=false,storeInterm::Bool=false)
    m  = length(b)

    nres0  = BLAS.nrm2(m,b,1) # norm(b)
    if nres0==0; return zero(eltype(b)),-9,[0.0]; end

    # (1) initialize
    if !isempty(x)
        b = copy(b)
        b = A('F',-x,1.0,b)
    end

    nrmb  = BLAS.nrm2(m,b,1) # norm(b)
    beta  = nrmb

    u     = b/beta
    v     = A('T',u,0.0)# A'*u
    n     = length(v)
	x     = (isempty(x)) ? zeros(n) : x
	if storeInterm;   X = zeros(n,maxIter) end
	if (beta <= btol*nres0); return x,1; end
	alpha = BLAS.nrm2(n,v,1)
	if alpha==0; return x,2; end
    v    /= alpha
    w     = copy(v)

    phib  = beta
    rhob  = alpha

    # (2) perform iteration
    # maxIter = min(maxIter,n,m)
    nBk  = 0.0
    nDk  = 0.0
    nres = beta
    STOP = fill(false,3)
    if doBidiag
        U    = zeros(m,maxIter+1)
        U[:,1] = u
        V    = zeros(n,maxIter)
        bs   = zeros(maxIter)
        as   = zeros(maxIter)
    end
    his  = zeros(maxIter,4)
    if out>1
        @printf "iter\t|r_k|\t\t|A'*r_k|\t|A|\t\tcond(A)\t\tSTOP\n"
    end
    flag = -1
    iter = 1
    while iter<=maxIter
        # (3) continue the bidiagonalization
        u     = A('F',v,-alpha,u) # A*v - alpha*u
        beta  = BLAS.nrm2(m,u,1)
        BLAS.scal!(m,1/beta,u,1) # faster than: u    /= beta
        if doBidiag
            as[iter] = alpha
            bs[iter] = beta
            U[:,iter+1] = u
            V[:,iter]   = v
        end
        nBk  += alpha^2+beta^2
        v     = A('T',u,-beta,v) # A'*u - beta*v
        alpha = BLAS.nrm2(n,v,1)
        BLAS.scal!(n,1/alpha,v,1) # faster than: v    /= alpha

        # (4) construct and apply next orthogonal transformation
        rho   = sqrt(rhob^2+beta^2)
        c     = rhob/rho
        s     = beta/rho
        theta = s*alpha
        rhob  = -c*alpha
        phi   = c*phib
        phib *= s
        nDk  += dot(w,w)/(rho*rho)

        # (5) update x,w
        BLAS.axpy!(n,phi/rho,w,1,x,1) # x += (phi/rho)*w
        if storeInterm; X[:,iter] = x; end

        # the following two lines are equivalent to w  = v - (theta/rho)*w
        w = BLAS.scal!(n,-theta/rho,w,1)
        w = BLAS.axpy!(n,1.0,v,1,w,1)

        # (6) check for stopping
        nrmA = sqrt(nBk)
        nrmG = abs(phib*alpha*c)
        conA = sqrt(nBk*nDk)
        nres *= s
        his[iter,:] = [nres nrmG nrmA conA]

        STOP[1] = nres <= btol*nres0+ atol*nrmA*norm(x)
        STOP[2] = (nrmG/(nrmA*nres)) <= atol
        STOP[3] = conA > condlim

        if out>1
            @printf "%03d\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t[%d,%d,%d]\n" iter nres nrmG nrmA conA 1.0*STOP[1] 1.0*STOP[2] 1.0*STOP[3]
        end
        iter += 1

        if any(STOP)
            break;
        end
    end

    if STOP[1]; flag = 1; end
    if STOP[2] && !STOP[1]; flag = 2; end
    if STOP[3]; flag = -2; end

    if out>=0
        if flag==-1
            println(@sprintf("lsqr iterated maxIter (=%d) times without converging to desired tolerance. Estimated residual norm is %1.3e and estimated norm of gradient is %1.3e.",iter-1,his[iter-1,1],his[iter-1,2]))
        elseif flag==-2
            println(@sprintf("lsqr stopped at iteration %d because the matrix seems to be too ill-conditioned. Estimated condition number is %1.2e.",iter-1,his[iter-1,4]))
        elseif flag==1 && out>=1
            println(@sprintf("lsqr achieved desired tolerance at iteration %d. The system should be compatible. Residual norm is %1.2e.",iter-1,his[iter-1,1]))
        elseif flag==2 && out>=1
            println(@sprintf("lsqr achieved desired tolerance at iteration %d. The system should be incompatible. Gradient norm is %1.2e.",iter-1,his[iter-1,1]))
        end
    end

    x = storeInterm ? X[:,1:iter-1] : x;

    if doBidiag
        U = U[:,1:iter]
        V = V[:,1:iter-1]
		I,J,V = SparseArrays.spdiagm_internal(-1=>bs[1:iter-1],0=>as[1:iter-1])
		B = sparse(I,J,V,iter,iter-1)
        # B = spdiagm((bs[1:iter-1],as[1:iter-1]),-1:0,iter,iter-1)
        return x, flag,his[1:iter-1,:],U,B,V
    else
        return x, flag,his[1:iter-1,:]
    end
end


using LAP_julia: inds_to_points, assess_flow_quality

## testing

flow_size = size(flow)

## method
img, imgw, flow = gen_init()
timer=TimerOutput("reg alg: sparse lap");
method_kwargs = Dict(:timer => timer, :point_count => 200, :spacing => 15)
flow_est, source_reg, timer, results, (estim_at_inds, inds) = test_registration_alg(sparse_lap, img, imgw, flow, [25, (51,51)], method_kwargs, timer=timer)

showflow(flow)
showflow(flow_est)
##

estim_flow = interpolate_flow_quad(estim_at_inds, inds, size(flow))
estim_flow_rbf = interpolate_flow(estim_at_inds, inds, size(flow))

assess_flow_quality(flow, estim_flow)

assess_flow_quality(flow, estim_flow_rbf)

showflow(flow, figtitle="original")
showflow(estim_flow, figtitle="quad estimate")
showflow(estim_flow_rbf, figtitle="rbf estimate")

estim_flow = interpolate_flow_quad(estim_at_inds, inds, size(flow))
estim_flow_rbf = interpolate_flow(estim_at_inds, inds, size(flow))

assess_flow_quality(flow, estim_flow)
assess_flow_quality(flow, estim_flow_rbf)

showflow(estim_flow)
showflow(flow)
showflow(flow_truth .- estim_flow)

showflow(rbf_flow)
assess_flow_quality(flow, rbf_flow)
showflow(flow_truth .- rbf_flow)

@benchmark interpolate_flow_quad(estim_at_inds, inds, size(flow))

## benchmark both

rand_inds = rand(CartesianIndices(size(flow)), 15)
lsdkf = interpolate_flow_quad(map(x -> flow[x], rand_inds), rand_inds, size(flow))
flow
showflow(lsdkf)
showflow(flow)

# function make_flow_at_ind(ind, a, b, c, d, e, f)
#     return a +
#            b * ind[1] +
#            c * ind[2] +
#            d * ind[1]^2 +
#            e * ind[2]^2 +
#            f * ind[1]*ind[2]
# end

# function make_flow(coeffs, flow_size)
#     return reduce((coeff, first) -> coeff .* ones(flow_size) .+ first, coeffs; init=0)
# end


"""
    add_scalar_constraint(model::MOI.ModelLike,
                          func::MOI.AbstractScalarFunction,
                          set::MOI.AbstractScalarSet;
                          allow_modify_function::Bool=false)

Adds the scalar constraint obtained by moving the constant term in `func` to
the set in `model`. If `allow_modify_function` is `true` then the function
`func`, can be modified.
"""
function add_scalar_constraint end

function add_scalar_constraint(model::MOI.ModelLike, func::MOI.SingleVariable,
                               set::MOI.AbstractScalarSet;
                               allow_modify_function::Bool=false)
    return MOI.add_constraint(model, func, set)
end
function add_scalar_constraint(model::MOI.ModelLike,
                               func::Union{MOI.ScalarAffineFunction{T},
                                           MOI.ScalarQuadraticFunction{T}},
                               set::MOI.AbstractScalarSet;
                               allow_modify_function::Bool=false) where T
    set = shift_constant(set, -func.constant)
    if !allow_modify_function
        func = copy(func)
    end
    func.constant = zero(T)
    return MOI.add_constraint(model, func, set)
end


__precompile__()

module ML

export Kernel, PolynomialKernel, GaussianKernel, ExponentialKernel,
    fit!, grad, ∇, center, KRR, KernelRidgeRegressor, RFKM, RandomFourierKernelMachine,
    KPCA, KernelPCA, negative_log_likelihood, rde, mae, rmse

using Distributions: Normal
using Distances: pairwise, colwise, Euclidean, SqEuclidean, euclidean, sqeuclidean

mae(x,y) = mean(abs(x-y))

rmse(x,y) = sqrt(sumabs2(x-y)/length(x))

abstract Kernel

type PolynomialKernel <: Kernel
    α::Float64
    c::Float64
    ξ::Float64
end

PolynomialKernel(;α=1.,c=0.,ξ=2.) = PolynomialKernel(α,c,ξ)

Base.call(k::PolynomialKernel, x::AbstractVector, y::AbstractVector) = (k.α*vecdot(x,y) + k.c).^k.ξ
Base.call(k::PolynomialKernel, x::AbstractArray, y::AbstractArray) = (k.α*(x'*y) + k.c).^k.ξ
Base.call(k::PolynomialKernel, x::AbstractArray) = k(x,x)

type GaussianKernel <: Kernel
    σ::Float64
end

GaussianKernel(;σ=1.) = GaussianKernel(σ)

Base.call(k::GaussianKernel, x::AbstractVector, y::AbstractVector) = exp(-sqeuclidean(x,y)/(2k.σ^2))
Base.call(k::GaussianKernel, x::AbstractArray, y::AbstractVector) = exp(-colwise(SqEuclidean(),x,y)/(2k.σ^2))
Base.call(k::GaussianKernel, x::AbstractVector, y::AbstractArray) = k(y,x)
Base.call(k::GaussianKernel, x::AbstractArray, y::AbstractArray) = exp(-pairwise(SqEuclidean(),x,y)/(2k.σ^2))
Base.call(k::GaussianKernel, x::AbstractArray) = exp(-pairwise(SqEuclidean(),x)/(2k.σ^2))

function grad(k::GaussianKernel, x::AbstractVector, y::AbstractVector)
    (y-x)*(k(x,y)/k.σ^2)
end

function grad(k::GaussianKernel, x::AbstractArray, y::AbstractVector)
    K = k(x,y)
    n,m = size(x,1),length(K)
    dst = Array{eltype(x)}((n,m))
    for i in 1:m
        dst[:,i] = grad(k,x[:,i],y)
    end
    reshape(dst,(n,size(K)...))
end

function grad(k::GaussianKernel, x::AbstractVector, y::AbstractArray)
    K = k(x,y)
    n,m = size(x,1),length(K)
    dst = Array{eltype(x)}((n,m))
    for i in 1:m
        dst[:,i] = grad(k,x,y[:,i])
    end
    reshape(dst,(n,size(K)...))
end

function grad(k::GaussianKernel, x::AbstractMatrix, y::AbstractMatrix)
    n,m,p = size(x,1),size(x,2),size(y,2)
    K = k(x,y)
    dst = Array{eltype(x)}((n,m,p))
    for j in 1:p
        for i in 1:m
            dst[:,i,j] = grad(k,x[:,i],y[:,j])
        end
    end
    dst
end

# symmetric
# function grad(k::GaussianKernel, x::AbstractMatrix)
#   n,m = size(x,1),size(x,2)
#   K = k(x,x)
#   dst = Array{eltype(x)}((n,m,m))
#   for j in 1:m
#       for i in j:m
#           dst[:,i,j] = (x[:,j]-x[:,i])*(K[i,j]/k.σ^2)
#           dst[:,j,i] = -dst[:,i,j]
#       end
#   end
#   dst
# end

# function grad(k::GaussianKernel, x::AbstractArray, y::AbstractArray)
#   n = size(x,1)
#   m,p = prod(size(x)[2:end]),prod(size(y)[2:end])
#   K = k(x,y)
#   shp = size(K)
#   K = reshape(K,(m,p))
#   dst = Array{eltype(x)}((n,m,p))
#   for j in 1:p
#       for i in 1:m
#           dst[:,i,j] = (y[:,j]-x[:,i])*(K[i,j]/k.σ^2)
#       end
#   end
#   reshape(dst,(n,shp...))
# end

function rand_train_test_ind(n,k)
    ind = randperm(n)
    ind[1:k], ind[k+1:end]
end

function uniform_train_test_ind(n,d)
    train = round(Int,1:n/d:n)
    test = setdiff(1:n,train)
    train,test
end

function perturb(x,i,h)
    y = copy(x)
    y[i] += h
    y
end

function fd_grad(f,x;h=1e-4)
    f0 = f(x)
    eltype(x)[(f(perturb(x,i,h))-f0)/h for i in 1:length(x)]
end

function fd_jacobian(f,x;h=1e-4)
    f0 = f(x)
    hcat([(f(perturb(x,i,h))-f0)/h for i in 1:length(x)]...)'
end

type ExponentialKernel <: Kernel
    σ::Float64
end

ExponentialKernel(;σ=1.) = ExponentialKernel(σ)

Base.call(k::ExponentialKernel, x::AbstractVector, y::AbstractVector) = exp(-euclidean(x,y)/(2k.σ^2))
Base.call(k::ExponentialKernel, x::AbstractArray, y::AbstractVector) = exp(-colwise(Euclidean(),x,y)/(2k.σ^2))
Base.call(k::ExponentialKernel, x::AbstractVector, y::AbstractArray) = k(y,x)
Base.call(k::ExponentialKernel, x::AbstractArray, y::AbstractArray) = exp(-pairwise(Euclidean(),x,y)/(2k.σ^2))
Base.call(k::ExponentialKernel, x::AbstractArray) = exp(-pairwise(Euclidean(),x)/(2k.σ^2))

abstract Predictor

# type KernelRidgeRegressor{T,Tk<:Kernel,Ni,No} <: Predictor
#     x::Array{T,Ni}          # features
#     y::Array{T,No}          # labels
#     k::Tk                   # kernel
#     λ::Float64              # regularization strength
#     w::Array{T,No}          # weights
# end

# typealias KRR KernelRidgeRegressor

# function KernelRidgeRegressor(x  ::AbstractArray,                           # features
#                               y  ::AbstractArray;                           # labels
#                               k  ::Kernel  = GaussianKernel(),              # kernel
#                               λ  ::Float64 = 1e-4)                          # regularization strength
#     KernelRidgeRegressor{eltype(x),typeof(k),ndims(x),ndims(y)}(x,y,k,λ,zeros(y))
# end

# typealias sKRR{T,Tk<:Kernel,Ni} KRR{T,Tk,Ni,1} # single valued

# function fit!(m::sKRR)
#     K = m.k(m.x)                                        # kernel matrix
#     m.w = (K + m.λ*I)\m.y                               # solve for weights
# end

# function fit!(m::KRR)
#     K = m.k(m.x)                                        # kernel matrix
#     m.w = (K + m.λ*I)\m.y'                              # solve for weights
# end

# # fit with new features, labels
# function fit!(m::KRR,
#               x::AbstractArray,                 # features
#               y::AbstractArray)                 # labels
#     m.x,m.y = x,y
#     fit!(m)
# end

# Base.call(m::sKRR, x::AbstractVector) = dot(m.w,m.k(x,m.x))
# Base.call(m::sKRR, x::AbstractArray) = m.k(x,m.x)*m.w

# Base.call(m::KRR, x::AbstractVector) = m.w' * m.k(x,m.x)
# Base.call(m::KRR, x::AbstractArray) = m.w' * m.k(x,m.x)'

# grad(m::KRR, x::AbstractVector) = grad(m.k,x,m.x)*m.w

# function grad(m::KRR, x::AbstractMatrix)
#     n,d = size(m.x)
#     k = size(x,2)
#     reshape(reshape(grad(m.k,x,m.x),(n*k,d)) * m.w, (n,k))
# end

# ∇ = grad

mae(x,y) = mean(abs(x-y))
rmse(x,y) = sqrt(sumabs2(x-y)/length(x))

function L_matrix(n::Int, d::Int)
    d == 1 && return I
    @assert n%d == 0 "n=$n must be divisble by d=$d"
    m = div(n,d)
    L = zeros(n,m)
    for i in 1:m
        L[d*(i-1)+1:d*i,i] = 1
    end
    L
end

L_matrix(x::AbstractMatrix, y::AbstractVector) = L_matrix(size(x,2),div(size(x,2),length(y)))

type KernelRidgeRegressor{T,Tk<:Kernel,Ni,No} <: Predictor
    x::Array{T,Ni}          # features
    y::Array{T,No}          # labels
    k::Tk                   # kernel
    λ::Float64              # regularization strength
    w::Array{T,No}          # weights
end

typealias KRR KernelRidgeRegressor

function KernelRidgeRegressor(x  ::AbstractArray,                           # features
                              y  ::AbstractArray;                           # labels
                              k  ::Kernel  = GaussianKernel(),              # kernel
                              λ  ::Float64 = 1e-4)                          # regularization strength
    KernelRidgeRegressor{eltype(x),typeof(k),ndims(x),ndims(y)}(x,y,k,λ,zeros(y))
end

# 1-dim features
function KernelRidgeRegressor(x  ::Vector{Float64},                         # features
                              y  ::Vector{Float64},                         # labels
                              k  ::Kernel = GaussianKernel(),               # kernel
                              λ  ::Float64 = 1e-4)                          # regularization strength
    KernelRidgeRegressor(reshape(x,(1,length(x))),y,k=k,λ=λ)
end

typealias sKRR{T,Tk<:Kernel,Ni} KRR{T,Tk,Ni,1} # single valued

function fit!(m::sKRR)
    # m.w = (K + m.λ*I)\m.y                             # solve for weights
    K = m.k(m.x)                                        # kernel matrix
    L = L_matrix(m.x,m.y)
    KL = K*L
    Λ = I/m.λ
    KLΛ = KL*Λ
    m.w = (K + KLΛ*KL')\(KLΛ*m.y)
end

# fit with new features, labels
function fit!(m::KRR,
              x::AbstractArray,                 # features
              y::AbstractArray)                 # labels
    m.x,m.y = x,y
    fit!(m)
end

Base.call(m::sKRR, x::AbstractVector) = dot(m.w,m.k(x,m.x))
Base.call(m::sKRR, x::AbstractArray) = m.k(x,m.x)*m.w

Base.call(m::KRR, x::AbstractVector) = m.w' * m.k(x,m.x)
Base.call(m::KRR, x::AbstractArray) = m.w' * m.k(x,m.x)'

grad(m::KRR, x::AbstractVector) = grad(m.k,x,m.x)*m.w

function grad(m::KRR, x::AbstractMatrix)
    n,d = size(m.x)
    k = size(x,2)
    reshape(reshape(grad(m.k,x,m.x),(n*k,d)) * m.w, (n,k))
end

∇ = grad

# # uncentered kernel matrix
# K = kpca._get_kernel(X_train)

# # scaled kPCA eigenvectors
# A = (kpca.alphas_/sqrt(kpca.lambdas_)).T

# # centering matrix
# H = empty((Nt,Nt))
# H.fill(-1.0/Nt)
# for i in range(Nt):
#   H[i,i] += 1

# # projection matrix
# Q = dot(H, dot(dot(A.T, A), H))

# def coeff(x, kernel=None, X=None, Q=None, K=None):
#   n = K.shape[0]
#   kx = squeeze(kernel(x, X))

#   # expansion coefficients of sample over training
#   # samples in uncentered feature space
#   c = dot(Q, kx - sum(K, axis=0)/n) + 1.0/n

#   return c

# def preimage(x0, c, kernel=None, X=None, max_steps=100, tol=1E-15):
#   x = x0.copy()

#   # fixed point iteration
#   for _ in range(max_steps):
#       kx = squeeze(kernel(x, X))
#       dx = dot(c*kx, X) / dot(c, kx) - x

#       x = x + dx

#       if dot(dx, dx) < tol:
#           break

#   return x

# abstract Predictor

# type KernelRidgeRegressor{T<:Kernel} <: Predictor
#   x::Matrix{Float64}      # features
#   y::Vector{Float64}      # labels
#   k::T                    # kernel
#   λ::Float64              # regularization strength
#   β::Float64              # gradient regularization strength
#   w::Vector{Float64}      # weights
# end

# typealias KRR KernelRidgeRegressor

# function KernelRidgeRegressor(x  ::Matrix{Float64},                       # features
#                             y  ::Vector{Float64};                         # labels
#                             k  ::Kernel  = GaussianKernel(),              # kernel
#                             λ  ::Float64 = 1e-4,                          # regularization strength
#                             β  ::Float64 = 1e-4)                          # gradient regularization strength
#   KernelRidgeRegressor{typeof(k)}(x,y,k,λ,β,[])
# end

# # 1-dim features
# # function KernelRidgeRegressor(x  ::Vector{Float64},                         # features
# #                               y  ::Vector{Float64},                         # labels
# #                               # dy ::Vector{Float64};                       # gradients
# #                               k  ::Kernel = GaussianKernel(),               # kernel
# #                               λ  ::Float64 = 1e-4,                          # regularization strength
# #                               β  ::Float64 = 1e-2)                          # gradient regularization strength
# #     KernelRidgeRegressor(reshape(x,(1,length(x))),y,reshape(dy,(1,length(dy))),k=k,λ=λ,β=β)
# # end

# # function loss(m::KernelRidgeRegressor)
# #     n,d = size(m.X)
# #     K = m.k(m.X)
# #     # δK = reshape(grad(m.k,m.X,m.X),(n*d,d))
# #     function loss(w)
# #         Kw = K*w
# #         sumabs2(Kw-m.y) + m.λ*dot(w,Kw) #+ m.β*sumabs2(δK*w - vec(m.dy))
# #         # m.β*sumabs(δK*w - vec(m.dy))
# #     end
# #     loss
# # end

# function fit!(m::KernelRidgeRegressor)
#   K = m.k(m.x)                                        # kernel matrix
#   m.w = (K + m.λ*I)\m.y                               # solve for weights
# end

# # function fit!(m::KernelRidgeRegressor;
# #               method     ::Symbol  = :l_bfgs,
# #               xtol       ::Float64 = 1e-8,
# #               ftol       ::Float64 = 1e-8,
# #               grtol      ::Float64 = 1e-8,
# #               iterations ::Int     = 1000)
# #     res = optimize(loss(m),zeros(size(m.X,2)),method=method,
# #         xtol=xtol,ftol=ftol,grtol=grtol,iterations=iterations)
# #     m.w = res.minimum
# #     res
# # end

# # fit with new features, labels
# function fit!(m::KernelRidgeRegressor,
#             x::Matrix{Float64},                       # features
#             y::Vector{Float64})                       # labels
#   m.x,m.y = x,y
#   fit!(m)
# end

# # function fit!(m::KernelRidgeRegressor,
# #               x::Vector{Float64},                       # features
# #               y::Vector{Float64})                       # labels
# #     fit!(m,reshape(x,(1,length(x))),y)
# # end

# Base.call(m::KernelRidgeRegressor, x::AbstractVector) = dot(m.w,m.k(x,m.x))
# Base.call(m::KernelRidgeRegressor, x::AbstractMatrix) = m.k(x,m.x)*m.w
# # Base.call(m::KernelRidgeRegressor, x::Number) = m([x])

# grad(m::KernelRidgeRegressor, x::AbstractVector) = grad(m.k,x,m.x)*m.w

# function grad(m::KRR, x::AbstractMatrix)
#   n,d = size(m.x)
#   k = size(x,2)
#   reshape(reshape(grad(m.k,x,m.x),(n*k,d)) * m.w, (n,k))
# end

# abstract Predictor

# type KernelRidgeRegressor{T<:Kernel} <: Predictor
#   X::Matrix{Float64}      # features
#   y::Vector{Float64}      # labels
#   k::T                    # kernel
#   λ::Float64              # regularization strength
#   w::Vector{Float64}      # weights
# end

# typealias KRR KernelRidgeRegressor

# function KernelRidgeRegressor(X ::Matrix{Float64},                            # features
#                             y ::Vector{Float64};                          # labels
#                             k ::Kernel  = GaussianKernel(),               # kernel
#                             λ ::Float64 = 1e-4)                           # regularization strength
#   KernelRidgeRegressor{typeof(k)}(X,y,k,λ,[])
# end

# # 1-dim features
# function KernelRidgeRegressor(x ::Vector{Float64},                            # features
#                             y ::Vector{Float64};                          # labels
#                             k ::Kernel = GaussianKernel(),                # kernel
#                             λ ::Float64 = 1e-4)                           # regularization strength
#   KernelRidgeRegressor(reshape(x,(1,length(x))),y,k=k,λ=λ)
# end

# function loss(m::KernelRidgeRegressor)
#   K = m.k(m.X)
#   function loss(w)
#       Kw = K*w
#       sumabs2(Kw-m.y) + m.λ*dot(w,Kw)
#   end
#   loss
# end

# function fit!(m::KernelRidgeRegressor;
#             method     ::Symbol  = :l_bfgs,
#             xtol       ::Float64 = 1e-8,
#             ftol       ::Float64 = 1e-8,
#             grtol      ::Float64 = 1e-8,
#             iterations ::Int     = 1000)
#   res = optimize(loss(m),zeros(size(m.X,2)),method=method,
#       xtol=xtol,ftol=ftol,grtol=grtol,iterations=iterations)
#   m.w = res.minimum
#   res
# end

# # function fit!(m::KernelRidgeRegressor)
# #     K = m.k(m.X)                                        # kernel matrix
# #     K[diagind(K)] += m.λ                                # regularization
# #     m.w = K\m.y                                         # solve for weights
# # end

# # fit with new features, labels
# function fit!(m::KernelRidgeRegressor,
#             X::Matrix{Float64},                       # features
#             y::Vector{Float64})                       # labels
#   m.X,m.y = X,y
#   fit!(m)
# end

# function fit!(m::KernelRidgeRegressor,
#             x::Vector{Float64},                       # features
#             y::Vector{Float64})                       # labels
#   fit!(m,reshape(x,(1,length(x))),y)
# end

# Base.call(m::KernelRidgeRegressor, x::AbstractVector) = dot(m.w,m.k(x,m.X))
# Base.call(m::KernelRidgeRegressor, x::AbstractMatrix) = m.k(x,m.X)*m.w
# Base.call(m::KernelRidgeRegressor, x::Number) = m([x])

# grad(m::KernelRidgeRegressor, x::AbstractVector) = grad(m.k,x,m.X)*m.w

# function grad(m::KRR, x::AbstractMatrix)
#   n,d = size(m.X)
#   k = size(x,2)
#   reshape(reshape(grad(m.k,x,m.X),(n*k,d)) * m.w, (n,k))
# end

type RandomFourierKernelMachine <: Predictor
    X::Matrix{Float64}              # features
    y::Vector{Float64}              # labels
    φ::Matrix{Float64}              # random features
    n::Int                          # input dimensionality
    d::Int                          # number of random features
    σ::Float64                      # kernel width
    λ::Float64                      # regularization strength
    w::Vector{Float64}              # weights
end

typealias RFKM RandomFourierKernelMachine

function RandomFourierKernelMachine(X ::Matrix{Float64},            # features
                                    y ::Vector{Float64};            # labels
                                    d ::Int     = 100,              # number of random features
                                    σ ::Float64 = 1.,               # kernel width
                                    λ ::Float64 = 1e-4)             # regularization strength
  RandomFourierKernelMachine(X,y,Matrix{Float64}(),size(X,1),d,σ,λ,[])
end

function kernel_matrix(m::RFKM, x::AbstractArray)
  z = exp(m.φ*x*im)/sqrt(m.d)
  real(z'*z)
end

function kernel_matrix(m::RFKM, x::AbstractArray, y::AbstractArray)
  zx = exp(m.φ*x*im)/sqrt(m.d)
  zy = exp(m.φ*y*im)/sqrt(m.d)
  real(zx'*zy)
end

function fit!(m::RFKM)
    d,n = m.d,m.n
    dist = Normal(0,1/m.σ)                      # normal distribution
    m.φ = rand(dist, (d,n))                     # random features
    z = cos(m.φ*m.X)/sqrt(m.d)                  # projection of X onto random features
    # z0 = centered ? zeros(n) : squeeze(mean(z,2),2)
    # y0 = centered ? zeros(n) : squeeze(mean(y,2),2)
    # zc = z .- z0
    # yc = y .- y0
    m.w = (z*z' + m.λ*I)\(z*m.y)
end

function Base.call(m::RFKM, x::AbstractVector)
    z = cos(m.φ*x)/sqrt(m.d)
    # y = Φ.α'*(z .- Φ.z0) .+ Φ.y0
    dot(m.w,z)
end

function Base.call(m::RFKM, x::AbstractMatrix)
    z = cos(m.φ*x)/sqrt(m.d)
    z'*m.w
end

function grad(m::RFKM, x::AbstractVector)
  dz = m.φ*im .* exp(m.φ*x*im)/sqrt(m.d)
  real(dz'*m.w)
end

function grad(m::RFKM, x::AbstractArray)
  n,d,k = m.n,m.d,size(x,2)
  c = -sin(m.φ*x)/sqrt(d)
  dz = m.φ .* reshape(c,(d,1,k))
  reshape(m.w' * reshape(dz,(d,n*k)), (n,k))
end

# H = I + fill(-1/n,size(K))
# H*K*H

doc"""
centering kernel matrices

$\varphi = \frac{1}{n} \sum_i \varphi(x_i)$

$\tilde K_{ij} = (\varphi(x_i) - \tilde\varphi)^T (\varphi(x_j) - \tilde\varphi)$

$\tilde K_{ij} = K_{ij} - \frac{1}{n} \sum_j K_{ij} - \frac{1}{n} \sum_i K_{ji} + \frac{1}{n^2} \sum_{ij} K_{ij}$

$\tilde L_{ij} = L_{ij} - \frac{1}{n} \sum_j K_{ij} - \frac{1}{n} \sum_i L_{ji} + \frac{1}{n^2} \sum_{ij} K_{ij}$

$\tilde M_{ij} = M_{ij} - \frac{1}{n} \sum_j L_{ij} - \frac{1}{n} \sum_i L_{ji} + \frac{1}{n^2} \sum_{ij} M_{ij}$

"""
function center(K)
    n = size(K,1)
    A = sum(K,1)/n
    K .- A .- A' + sum(K)/n^2
end

# J = fill(-1/n,(m,n))
# (L + J*K)*H
function center(K,L::Matrix)
    n = size(K,1)
    L .- sum(K,1)/n .- sum(L,2)/n + sum(K)/n^2
end

function center(K,L::Vector)
    n = size(K,1)
    L .- vec(sum(K,1))/n .- sum(L)/n + sum(K)/n^2
end

# M + L*J' + J*L' + J*K*J'
function center(K,L,M)
    n = size(K,1)
    M .- sum(L,2)/n .- sum(L,2)'/n + sum(K)/n^2
end

type KernelPCA{T<:Kernel}
    k::T                        # kernel
    x::Matrix{Float64}          # features
    K::Matrix{Float64}          # kernel matrix
    u::Vector{Float64}          # eigenvalues
    v::Matrix{Float64}          # eigenvectors
    n::Int                      # number of principal components
end

typealias KPCA KernelPCA

function KernelPCA(; k = GaussianKernel(), n::Int = -1)
    KernelPCA{typeof(k)}(k,Matrix{Float64}(),Matrix{Float64}(),[],Matrix{Float64}(),n)
end

function fit!(m::KPCA, x::Matrix{Float64})
    n = m.n
    m.x = x                                                 # features
    m.K = m.k(x)                                            # compute kernel matrix
    K = center(m.K)                                         # center kernel matrix
    u,v = eig(Symmetric(K))                                 # compute principal components
    ind = sortperm(u,rev=true)                              # sort eigenvalues
    u = u[ind]
    v = v[:,ind]
    n = n == -1 ? findlast(u.>0,true) : n                   # ignore zero eigenvalues
    m.u = u[1:n]                                            # select n eigenvalues
    m.v = v[:,1:n]                                          # select n eigenvectors
    m.v ./= sqrt(m.u)'                                      # normalize eigenvectors
    return
end

function Base.call(m::KPCA, x::Vector{Float64})
    kx = center(m.K,m.k(x,m.x))         # compute centered kernel matrix
    m.v' * kx                           # project features onto principal components
end

function Base.call(m::KPCA, x::Matrix{Float64})
    K = center(m.K,m.k(x,m.x))          # compute centered kernel matrix
    (K * m.v)'                          # project features onto principal components
end

function grad(m::KPCA, x::Vector{Float64})
    n,d = size(m.x)
    δK = grad(m.k,x,m.x)
    δK = δK .- sum(δK,2)/d
    δK * m.v
end

function grad(m::KPCA, x::Matrix{Float64})
    n,d = size(m.x)
    l = size(x,2)
    δK = grad(m.k,x,m.x)
    δK = δK .- sum(δK,3)/d
    reshape(reshape(δK,(n*l,d)) * m.v, (n,l,m.n))
end

∇ = grad

function negative_log_likelihood(φ::KPCA, y::AbstractVector)
    n = length(φ.u)
    z = φ.v' * y  # projection of labels onto KPCA eigenvectors
    σ1 = d -> sum(z[1:d].^2) / d
    σ2 = d -> sum(z[d+1:n].^2) / (n-d)
    Float64[d/n * log(σ1(d)) + (n-d)/n * log(σ2(d)) for d in 1:n-1]
end

function rde(φ::KPCA, y::AbstractVector)
    nll = negative_log_likelihood(φ,y)
    d = findfirst(diff(nll).<0)
end

end # module

# abstract Predictor

# abstract Loss

# type L2RegularizedQuadraticLoss <: Loss
#   λ::Float64
# end

# L2RegularizedQuadraticLoss(;λ=1e-8) = L2RegularizedQuadraticLoss(λ)

# type QuadraticLoss <: Loss end

# type L1RegularizedQuadraticLoss <: Loss
#   λ::Float64
# end

# L1RegularizedQuadraticLoss(;λ=1e-8) = L1RegularizedQuadraticLoss(λ)

# type KernelRidgeRegressor{Tk<:Kernel,Tl<:Loss} <: Predictor
#   k::Tk                   # kernel
#   l::Tl                   # loss model
#   X::Matrix{Float64}      # features
#   y::Vector{Float64}      # labels
#   w::Vector{Float64}      # weights
# end

# function KernelRidgeRegressor(X ::Matrix{Float64},                            # features
#                             y ::Vector{Float64};                          # labels
#                             k ::Kernel = GaussianKernel(),                # kernel
#                             l ::Loss   = L2RegularizedQuadraticLoss())    # loss
#   KernelRidgeRegressor{typeof(k),typeof(l)}(k,l,X,y,[])
# end

# # 1-dim features
# function KernelRidgeRegressor(x ::Vector{Float64},                            # features
#                             y ::Vector{Float64};                          # labels
#                             k ::Kernel = GaussianKernel(),                # kernel
#                             l ::Loss   = L2RegularizedQuadraticLoss())    # loss
#   KernelRidgeRegressor(reshape(x,(1,length(x))),y,k=k,l=l)
# end

# function loss{Tk<:Kernel}(m::KernelRidgeRegressor{Tk,L1RegularizedQuadraticLoss})
#   K = m.k(m.X)                                # kernel matrix
#   L = chol(K)                                 # cholesky factorization
#   function loss(w)
#       sumabs2(K*w-m.y) + m.l.λ*sumabs(L*w)
#   end
#   loss
# end

# function loss{Tk<:Kernel}(m::KernelRidgeRegressor{Tk,L2RegularizedQuadraticLoss})
#   K = m.k(m.X)
#   function loss(w)
#       Kw = K*w
#       sumabs2(Kw-m.y) + m.l.λ*dot(w,Kw)
#   end
#   loss
# end

# function loss_gradient!{Tk<:Kernel}(m::KernelRidgeRegressor{Tk,L2RegularizedQuadraticLoss})
#   K = m.k(m.X)
#   function loss_gradient!(w, dst)
#       Kw = K*w
#       dst[:] = 2*K*(Kw-m.y) + 2*m.l.λ*Kw
#   end
#   loss_gradient!
# end

# function loss{Tk<:Kernel}(m::KernelRidgeRegressor{Tk,QuadraticLoss})
#   K = m.k(m.X)
#   loss(w) = sumabs2(K*w-m.y)
#   loss
# end

# function loss_gradient!{Tk<:Kernel}(m::KernelRidgeRegressor{Tk,QuadraticLoss})
#   K = m.k(m.X)
#   function loss_gradient!(w, dst)
#       dst[:] = 2*K*(K*w-m.y)
#   end
#   loss_gradient!
# end

# @generated function fit!(m::KernelRidgeRegressor;
#                        method     ::Symbol  = :l_bfgs,
#                        xtol       ::Float64 = 1e-8,
#                        ftol       ::Float64 = 1e-8,
#                        grtol      ::Float64 = 1e-8,
#                        iterations ::Int     = 1000)
#   if method_exists(:loss_gradient!,(m,))  # if loss gradient is implemented
#       return quote
#           res = optimize(loss(m),loss_gradient!(m),zeros(size(m.X,2)),method=method,
#               xtol=xtol,ftol=ftol,grtol=grtol,iterations=iterations)
#           m.w = res.minimum
#           res
#       end
#   else
#       return quote
#           res = optimize(loss(m),zeros(size(m.X,2)),method=method,
#               xtol=xtol,ftol=ftol,grtol=grtol,iterations=iterations)
#           m.w = res.minimum
#           res
#       end
#   end
# end

# function fit!{Tk<:Kernel}(m::KernelRidgeRegressor{Tk,L2RegularizedQuadraticLoss})
#   K = m.k(m.X)                                        # kernel matrix
#   K[diagind(K)] += m.l.λ                              # regularization
#   m.w = K\m.y                                         # solve for weights
# end

# # fit with new features, labels
# function fit!(m::KernelRidgeRegressor,
#             X::Matrix{Float64},                       # features
#             y::Vector{Float64})                       # labels
#   m.X,m.y = X,y
#   fit!(m)
# end

# function fit!(m::KernelRidgeRegressor,
#             x::Vector{Float64},                       # features
#             y::Vector{Float64})                       # labels
#   fit!(m,reshape(x,(1,length(x))),y)
# end

# predict
# Base.call(m::KernelRidgeRegressor, x::AbstractVector) = dot(m.w,m.k(x,m.X))
# Base.call(m::KernelRidgeRegressor, x::AbstractMatrix) = m.k(x,m.X)*m.w
# Base.call(m::KernelRidgeRegressor, x::Number) = m([x])


cust_z(x,μ,σ) = ismissing(x) ? 0.0 : (x - μ)/σ

function ZScores(x::AbstractArray)
    μ = mean(skipmissing(x))
    σ = std(skipmissing(x))
    cust_z.(x,μ,σ)
end

delta(vec) = vec - lag(vec)

function repeat_by!(df,f,b,cc, n)
    c_list = [Symbol(join([n, x],"_")) for x in cc]
    for c in c_list
        df[!,c] = Vector{Union{Float64,Missing}}(undef,nrow(df))
    end
    combine(groupby(df,b)) do dd
        for (col,rcol) in zip(cc,c_list)
            dd[:,rcol] = f(dd[:,col])
        end
    end
    return df
end

function elaborate(df)
    # df = get_stats()
    repeat_by!(df,delta,:name,deltas_cols,"delta")
    cc = deltas_cols
    c_list = [Symbol(join(["delta", x],"_")) for x in deltas_cols]
    cc = vcat(c_list,z_cols)
    repeat_by!(df,ZScores,:date,cc,"zscore")
    df.zscore_delta_arena_place = df.zscore_delta_arena_place .* -1
    df.zscore_total = df.zscore_delta_arena_victories .+
        # df.zscore_delta_arena_place .+
        # df.zscore_delta_arena_defences .+
        df.zscore_delta_maxed_heroes .+
        df.zscore_clan_tokens .+
        df.zscore_raid_points .+
        df.zscore_portal_stones
    cc = names(df)
    filter!(t -> occursin("zscore",string(t)),cc)
    repeat_by!(df,t->tiedrank(t;rev = true),:date,cc,"rank")
    return df
end

function update_results(clan::AbstractDict)
    full = get_stats(clan)
    elab = elaborate(full)
    return elab
end


# Solves a basic eigenvalue problem defined through different NEP types through NLEIGS

using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using NonlinearEigenproblems.RKHelper
using Test
using SparseArrays

struct CustomNLEIGSNEP <: NEP
    n::Int  # problem size; this is the only required field in a custom NEP type when used with NLEIGS
end

nep_types_test_n = 2
nep_types_test_B = Vector{Matrix{Float64}}([[1 3; 5 6], [3 4; 6 6]])
nep_types_test_C = Vector{Matrix{Float64}}([[1 0; 0 1]])

# implement a few methods used by the solver
import NonlinearEigenproblems.NEPCore.compute_Mder
import NonlinearEigenproblems.NEPCore.compute_Mlincomb
import Base.size
nep_types_test_pep = PEP([nep_types_test_B; nep_types_test_C])
compute_Mder(::CustomNLEIGSNEP, λ::Number) = compute_Mder(nep_types_test_pep, λ)
compute_Mlincomb(::CustomNLEIGSNEP, λ::Number, x::AbstractVecOrMat) = compute_Mlincomb(nep_types_test_pep, λ, x)
size(::CustomNLEIGSNEP, _) = nep_types_test_n

function nleigs_nep_types()
    n = nep_types_test_n
    B = nep_types_test_B
    C = nep_types_test_C
    f = [λ -> λ^2]
    Σ = [-10.0-2im, 10-2im, 10+2im, -10+2im]

    # define and solve the same problem in many different ways
    problems = [
        ("SPMF_NEP", SPMF_NEP([B; C], [λ -> λ^0; λ -> λ; λ -> λ^2])),
        ("PEP", PEP([B; C])),
        ("PEP + SPMF", SumNEP(PEP(B), SPMF_NEP(C, f))),
        ("PEP + LowRankFactorizedNEP", SumNEP(PEP(B), LowRankFactorizedNEP([LowRankMatrixAndFunction(sparse(C[1]), f[1])]))),
        ("Custom NEP type", CustomNLEIGSNEP(n))]

    for problem in problems
        @bench @testset "$(problem[1])" begin
            lambda, X = nleigs(problem[2], Σ, maxit=10, v=ones(n).+0im, blksize=5)
            verify_lambdas(4, problem[2], lambda, X)
        end
    end
end

@testset "NLEIGS: NEP Types" begin
    nleigs_nep_types()
end


using LowRankModels
import StatsBase: sample, Weights

# test ordistic loss

## generate data
srand(1);
m,n,k = 100,100,3;
kfit = k+1
nlevels = 5; # number of levels
d = nlevels-1 # embedding dimension
D = n*d;
# coordinates of covariates
X_real = randn(m,k)
# directions of observations
Y_real = randn(k,n)
# measurement thresholds
T_real = k*randn(d,n) # notice x^T y has variance k; so making this bigger makes the problem easier
for j=1:n # this scheme doesn't work to ensure uniform sampling
	T_real[:,j] = sort(T_real[:,j])
end

signedsums = @compat Array{Float64}(d, nlevels)
for i=1:d
    for j=1:nlevels
        signedsums[i,j] = i<j ? 1 : -1
    end
end

XY = X_real*Y_real;
XYplusT = zeros(Float64, (m,D))
A = zeros(Int, (m, n))
for i=1:m
	for j=1:n
		u = XY[i,j] + T_real[:,j]
		XYplusT[i,(j-1)*d+(1:d)] = u
		diffs = u'*signedsums
		wv = Weights(Float64[exp(-diffs[l]) for l in 1:nlevels])
		l = sample(wv)
		A[i,j] = l
	end
end
# loss is insensitive to shifts; regularizer should pick this shift
XYplusT = XYplusT .- mean(XYplusT, 2)

# and the model
losses = BvSLoss(nlevels)
rx, ry = lastentry1(QuadReg(.01)), OrdinalReg(QuadReg(.01)) #lastentry_unpenalized(QuadReg(10));
glrm = GLRM(A,losses,rx,ry,kfit, scale=false, offset=false, X=randn(kfit,m), Y=randn(kfit,D));

# fit w/o initialization
@time X,Y,ch = fit!(glrm);
XYh = X'*Y#[:,1:d:D];
@show ch.objective
println("After fitting, parameters differ from true parameters by $(vecnorm(XYplusT - XYh)/sqrt(prod(size(XYplusT)))) in RMSE")
A_imputed = impute(glrm)
println("After fitting, $(sum(A_imputed .!= A) / prod(size(A))*100)\% of imputed entries are wrong")
println("After fitting, imputed entries are off by $(sum(abs,(A_imputed - A)) / prod(size(A))*100)\% on average")
println("(Picking randomly, $((nlevels-1)/nlevels*100)\% of entries would be wrong.)\n")

# initialize
init_svd!(glrm)
XYh = glrm.X' * glrm.Y
println("After initialization with the svd, parameters differ from true parameters by $(vecnorm(XYplusT - XYh)/sqrt(prod(size(XYplusT)))) in RMSE")
A_imputed = impute(glrm)
println("After initialization with the svd, $(sum(A_imputed .!= A) / prod(size(A))*100)\% of imputed entries are wrong")
println("After initialization with the svd, imputed entries are off by $(sum(abs,(A_imputed - A)) / prod(size(A))*100)\% on average")
println("(Picking randomly, $((nlevels-1)/nlevels*100)\% of entries would be wrong.)\n")

# fit w/ initialization
@time X,Y,ch = fit!(glrm);
XYh = X'*Y#[:,1:d:D];
@show ch.objective
println("After fitting, parameters differ from true parameters by $(vecnorm(XYplusT - XYh)/sqrt(prod(size(XYplusT)))) in RMSE")
A_imputed = impute(glrm)
println("After fitting, $(sum(A_imputed .!= A) / prod(size(A))*100)\% of imputed entries are wrong")
println("After fitting, imputed entries are off by $(sum(abs,(A_imputed - A)) / prod(size(A))*100)\% on average")
println("(Picking randomly, $((nlevels-1)/nlevels*100)\% of entries would be wrong.)\n")

# test scaling
scale!(glrm.ry, 3)


### polyMap computes the deterministic part defined by the polynomial g(θ,z)=∑ᵢ₌₀ θᵢ zⁱ⁻¹
function polyMap(Θ,z)
  g = 0.
  for i in 1:1:length(Θ)
    g += Θ[i]*z^(i-1)
  end
  return g
end

### noiseMix1 simulates the noise process f₁
function noiseMix1()
  u = rand()
  z = 0.
  if u < 0.25
    z = rand(Normal(0,sqrt(0.0001)))
  elseif u < 0.5
    z = rand(Normal(0, sqrt(6*0.0001)))
  elseif u < 0.75
    z = rand(Normal(0,sqrt(11*0.0001)))
  else
    z = rand(Normal(0,sqrt(16*0.0001)))
  end
  return z
end

### noiseMix2 simulates the noise process f₂ᵣ
function noiseMix2()
  u = rand()
  z = 0.
  if u < 8/10
    z = rand(Normal(0,0.001))
  else
    z = rand(Normal(0,0.2))
  end
  return z
end

### genData simulates chaotic time series from the random map with deterministi part polyMap and
### stochastic part specified by noiseMix.
function genData(n,Θ,x₀,seed)
  Random.seed!(seed)
  x = zeros(n)
  x[1] = polyMap(Θ,x₀) + noiseMix2()
  for i in 2:1:n
    x[i] = polyMap(Θ,x[i-1]) + noiseMix2()
  end
  return x
end

### genData simulates chaotic time series from the random map with deterministi part polyMap and
### stochastic part specified by a Gaussian distribution (Parametric)
function genDataParam(n,Θ,x₀,seed)
  Random.seed!(seed)
  x = zeros(n)
  x[1] = polyMap(Θ,x₀) + rand(Normal(0.0, sqrt(0.0001)))
  for i in 2:1:n
    x[i] = polyMap(Θ,x[i-1]) + rand(Normal(0.0, sqrt(0.0001)))
  end
  return x
end

### tgeornd simulates random variates from the truncated geometric distribution
function tgeornd(p,k)
  return rand(Geometric(p)) + k
end

### range_intersection provides the intersection of the interval [a,b] with the intervals speficied by [x₁,...,xn]
function range_intersection(A::Array{Float64}, B::Array{Float64})
#=

This function is a straightforward translation in Julia from Xavier Xavier MATLAB code
Copyright (c) 2011, Xavier Xavier
All rigths reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

Purpose: Range/interval intersection

 A and B two ranges of closed intervals written
 as vectors [lowerbound1 upperbound1 lowerbound2 upperbound2]
 or as matrix [lowerbound1, lowerbound2, lowerboundn;
               upperbound1, upperbound2, upperboundn]
 A and B have to be sorted in ascending order

 out is the mathematical intersection A n B

 EXAMPLE USAGE:
   >> out=rangeIntersection([1 3 5 9],[2 9])
   	out =  [2 3 5 9]
   >> out=rangeIntersection([40 44 55 58], [42 49 50 52])
   	out =  [42 44]
=#

# Allocate, as we don't know yet the size, we assume the largest case
  out1 = zeros(length(B)+(length(A)-2))
  k = 1

  while isempty(A) .== 0 && isempty(B) .== 0
  # make sure that first is ahead second
    if A[1] .> B[1]
      temp = copy(B)
      B = copy(A)
      A = copy(temp)
    end

    if A[2] .< B[1]
      A = copy(A[3:end])
      continue
    elseif A[2] .== B[1]
      out1[k] = B[1]
      out1[k + 1] = B[1]
      k = k + 2

      A = A[3:end]
      continue
    else
      if A[2] .== B[2]
        out1[k] = B[1]
        out1[k+1] = B[2]
        k = k + 2

        A = A[3:end]
        B = B[3:end]

      elseif A[2] .< B[2]
        out1[k] = B[1]
        out1[k+1] = A[2]
        k = k + 2

        A = A[3:end]
      else
        out1[k] = B[1]
        out1[k+1] = B[2]
        k = k + 2

        B = B[3:end]
      end
     end
    end

  # Remove the tails
  out = copy(out1[1:k-1])
  return out
end

### unifmixrnd simulates from a mixture of uniforms in the intervals specified by x
function unifmixrnd(x)
  n = length(x)
  nc = 0.0
  for i in 1:2:n
    nc += x[i + 1] - x[i]
  end

  U = 0.
  u = rand()
  prob = 0.
  for i in 1:2:n
    prob += (x[i + 1] - x[i]) / nc
    if u .< prob
      U = x[i] + (x[i + 1] - x[i]) * rand()
      break
    end
  end
  return U
end

### learnc updates the concentration parameter of the DP
function learnc(ac, bc, c, ss, k)
  z = 0.
  x = rand(Beta(c + 1, ss))
  px = (ac + k - 1.0) / (bc - log(x))
  px = px / (px + ss)
  r = rand(Uniform(0.0, 1.0))
  if (r .< px)
    z = rand(Gamma(ac + k, 1.0 / (bc - log(x))))
  else
    z = rand(Gamma(ac + k - 1.0, 1.0 / (bc - log(x))))
  end
  return z
end

### synchronized samples the geometric probability from the corresponding posterior under a transformed gamma prior
function synchronized(az, bz, z, N, n)
  Ln = az - n + N - 1
  u1 = rand(Uniform(0.0, (1.0 - z)^Ln))
  u2 = rand(Uniform(0.0, exp(-bz / z)))
  if Ln .>= 0
    return (((1 - u1 ^ (1.0 / Ln)) ^ (2 * n - az) - (-bz / log(u2)) ^ (2 * n - az)) * rand() + (-bz / log(u2)) ^ (2 * n - az)) ^ (1 / (2 * n - az))
  else
    temp1 = 1 - u1^(1 / Ln)
    temp2 = -bz / log(u2)
    temp3 = max(temp1, temp2)
    return ((1. - temp3 ^ (2. * n - az)) * rand() + temp3 ^ (2. * n - az)) ^ (1. / (2. * n - az))
  end
end

function predictions(x::Array{Float64,2}, weights::Array{Float64,2})
    np, _ = size(weights)
    preds = zeros(np, size(x,2))
    @inbounds for i in 1:1:np
        preds[i,:] = g(x, weights[i,:])
    end
    stds = std(preds, dims=1)
    return preds, stds
end


function CalculateSampling(in)

	cutoff = 1e-10
	itrace = 1

	n=size(in)
	in=reshape(in,n[1],:)
        wd = zeros(Float32,size(in))
	n2=size(in,2)
	for itrace = 1 : n2
		a = sum(in[:,itrace].*in[:,itrace])
		if (a > cutoff)
			wd[:,itrace] = 1.
		end
	end
	wd=reshape(wd,n)
	return wd;
end

function CalculateSampling(in,h;cutoff=1e-10)

	itrace = 1
        n=size(in)
	in=reshape(in,n[1],:)
	wd = zeros(Float32,size(in))
	n2=size(in,2)
	for itrace = 1 : n2
		a = sqrt(sum(in[:,itrace].^2))
		if (a > cutoff)
			wd[:,itrace] = 1.
		end
	end
	wd=reshape(wd,n)
	return wd,h;
end

function CalculateSampling(in::AbstractString,wd::AbstractString;cutoff=1e-10)
	# calculate sampling operator (1's for live traces, 0's for missing traces)

	@compat parameters = Dict(:cutoff=>cutoff)
	SeisProcess(in,wd,[CalculateSampling],[parameters])

end


using jInv.InverseSolve
using Test
using jInv.Mesh
using jInv.Utils

# setup forward mesh
domain = [0;5;0;5; 0;3]
n     = [12;13;10]
Minv  = getRegularMesh(domain,n)

modFun = (expMod,boundMod,fMod)
for k=1:length(modFun)
	# checkDerivative of $(modFun[k])
	function testModFun(m,v=[])
		mc,dm = modFun[k](m)
		if !isempty(v)
			dm = dm*v
			return mc,dm
		end
		return mc
	end
	derivativeDistModel, = checkDerivative(testModFun,rand(prod(n)),out=false)
	@test derivativeDistModel
end


using Wilsonloop
using Test

function test()

    println("plaq")
    plaq = make_plaq()
    display(plaq)
    for μ=1:4
        println("μ = $μ")
        staples = make_plaq_staple(μ)
        display(staples)
    end
    show(loops_staple_prime[(4,4)])

    #return

    loop = [(1,+1)]

    println(loop)
    w = Wilsonline(loop)
    loop2 = [(1,-1)]
    w2 = Wilsonline(loop2)

    V1 = derive_U(w,1)
    V2 = derive_U(w2,1)
    show(V1)
    println("d")
    show(V2)
    println("d")

    loop = [(1,+1),(2,+1),(1,-1),(2,-1)]
    println(loop)
    w = Wilsonline(loop)
    println("P: ")
    show(w)
    println("P^+: ")
    show(w')
    println("staple")
    for μ=1:4
        println("μ = $μ")
        V1 = make_staple(w,μ)
        V2 = make_staple(w',μ)
        show(V1)
        show(V2)
    end

    println("derive w")
    for μ=1:4
        dU = derive_U(w,μ)
        for i=1:length(dU)
            show(dU[i])
        end
    end

    println("-------------------------------------------------------")
    println("C and dC/dU")
    for μ=1:4
        C = make_Cμ(w,μ)
        #=
        V1 = make_staple(w,μ)
        V2 = make_staple(w',μ)
        C = eltype(V1)[]
        for i=1:length(V1)
            push!(C,V1[i]')
        end
        for i=1:length(V2)
            push!(C,V2[i]')
        end
        =#
        println("-------------------------------------------")
        println("μ = $μ")
        for i=1:length(C)
            println("---------------------------------------")
            println("C[$i]: ")
            show(C[i])
            for ν=1:4
                println("-----------------------------")
                println("ν = $ν")
                dCdU = derive_U(C[i],ν)
                println("dC_{$μ}/dU_{$ν}: ")
                for j=1:length(dCdU)
                    show(dCdU[j])
                end
            end
        end
    end

    println("-------------------------------------------------------")
    println("C and dC/dUdag")
    for μ=1:4
        C = make_Cμ(w,μ)
        #=
        V1 = make_staple(w,μ)
        V2 = make_staple(w',μ)
        C = eltype(V1)[]
        for i=1:length(V1)
            push!(C,V1[i]')
        end
        for i=1:length(V2)
            push!(C,V2[i]')
        end
        =#
        println("-------------------------------------------")
        println("μ = $μ")
        for i=1:length(C)
            println("---------------------------------------")
            println("C[$i]: ")
            show(C[i])
            for ν=1:4
                println("-----------------------------")
                println("ν = $ν")
                dCdU = derive_Udag(C[i],ν)
                println("dC_{$μ}/dUdag_{$ν}: ")
                for j=1:length(dCdU)
                    show(dCdU[j])
                end
            end
        end
    end

    #=
    w = Wilsonline(loop)
    println("P: ")
    display(w)
    println("P^+: ")
    display(w')
    println("staple")
    make_staples(w)
    =#
end

@testset "Wilsonloop.jl" begin
    test()
    @test true
    # Write your tests here.
end


using Test
using SparseArrays

@testset "Model Generation" begin
    G = [0.0 0.0 1.0 1.0;  0.0 1.0 0.0 1.0]
    T1 = SparseArrays.sparse(Array{Int8,2}([1 1 0 0;  0 1 1 0;  0 0 1 1]))
    T2 = SparseArrays.sparse(Array{Int8,2}([1 1 1]))

    @test isnothing(CAGD.Model())

    m = CAGD.Model(G)
    @test size(m, 0, 1) == 2
    @test size(m, 0, 2) == 4
    @test size(m, 1, 1) == 0
    @test size(m, 1, 2) == 4
    @test size(m, 2, 1) == 0
    @test size(m, 2, 2) == 0
    @test_throws BoundsError size(m, 3, 1)

    m = CAGD.Model(G, [T1, T2])
    @test size(m, 0, 1) == 2
    @test size(m, 0, 2) == 4
    @test size(m, 1, 1) == 3
    @test size(m, 1, 2) == 4
    @test size(m, 2, 1) == 1
    @test size(m, 2, 2) == 3
    @test_throws BoundsError size(m, 3, 1)

    @test_throws ArgumentError CAGD.Model(G, [T1])
    @test_throws ArgumentError CAGD.Model(G, [T1, T1])
    @test_throws ArgumentError CAGD.Model(G, [T1, T1, T1])

end

@testset "Model Cloning" begin
    G = [0.0 0.0 1.0 1.0;  0.0 1.0 0.0 1.0]
    T1 = SparseArrays.sparse(Array{Int8,2}([1 1 0 0;  0 1 1 0;  0 0 1 1]))
    T2 = SparseArrays.sparse(Array{Int8,2}([1 1 1]))
    m1 = CAGD.Model(G, [T1, T2])
    m2 = deepcopy(m1)
    m3 = copy(m2)
    CAGD.addModelCell!(m2, 1, [1, 3])

    @test m1.T[1] != m2.T[1]
    @test m2.T[1] == m3.T[1]
end

@testset "Topology Manipulation" begin

    G = [
        0.0 0.0 0.0 0.0 1.0 1.0 1.0 1.0
        0.0 0.0 1.0 1.0 0.0 0.0 1.0 1.0
        0.0 1.0 0.0 1.0 0.0 1.0 0.0 1.0
    ]
    m = CAGD.Model(G)

    @testset "add Model Cells" begin
        CAGD.addModelCell!(m, 1, SparseArrays.sparse(Int8[-1; 1; 0; 0;  0; 0; 0; 0]))
        @test_throws MethodError CAGD.addModelCell!(m, 1, SparseArrays.sparse(Int8[-1 0 1 0  0 0 0 0]))
        @test_throws ArgumentError CAGD.addModelCell!(m, 1, SparseArrays.sparse(Int8[-1; 1; 0; 0;  0; 0; 0]))

        CAGD.addModelCells!(m, 1, SparseArrays.sparse(Int8[-1 0 1 0  0 0 0 0]))
        @test_throws ArgumentError CAGD.addModelCells!(m, 1, SparseArrays.sparse(Int8[-1 0 1 0  0 0 0]))
        @test size(m, 1) == (2, 8)

        @test_throws ArgumentError CAGD.addModelCells!(m, 0, SparseArrays.sparse(Int8[-1 0 1 0  0 0 0 0]))
        @test_throws ArgumentError CAGD.addModelCells!(m, 4, SparseArrays.sparse(Int8[-1 0 1 0  0 0 0 0]))
    end

    @testset "add Model Cells (EV)" begin
        CAGD.addModelCell!(m, 1, [2; 3], signed = true)
        @test m.T[1][3, :].nzval == [-1; 1]
        @test_throws ArgumentError CAGD.addModelCell!(m, 2, [1, 2, 4])
        CAGD.addModelCells!(m, 2, [[1, 2, 3]], signed = true)
        @test Array(m.T[2][1, :]) == [1; -1; 1]
        CAGD.addModelCells!(m, 1, [[2, 4], [3, 4]], signed = true)
        @test size(m, 2) == (1, 5)
        CAGD.addModelCell!(m, 2, [1; 2; 4; 5])
        @test m.T[2][2, :].nzval == [1; 1; 1; 1]
        CAGD.addModelCells!(m, 2, [[3; 4; 5]], signed = true)
        @test size(m, 3) == (0, 3)
    end

    @testset "get Model Lo Cell" begin
        @test CAGD.getModelLoCell(m, 2, 1) == [1; 2; 3]
        @test CAGD.getModelLoCell(m, 2, 2) == [1; 2; 4; 5]
        @test CAGD.getModelLoCell(m, 1, 2) == [1; 3]
        @test_throws ArgumentError CAGD.getModelLoCell(m, 0, 2)
        @test_throws ArgumentError CAGD.getModelLoCell(m, 4, 2)
        @test_throws ArgumentError CAGD.getModelLoCell(m, 3, 2)
        @test_throws ArgumentError CAGD.getModelLoCell(m, 2, -1)

        @test CAGD.getModelLoCell(m, 2, [1;2]) == [1; 2; 3; 4; 5]
        @test CAGD.getModelLoCell(m, 1, [1;2]) == [1; 2; 3]
        @test_throws ArgumentError CAGD.getModelLoCell(m, 0, [1; 2])
        @test_throws ArgumentError CAGD.getModelLoCell(m, 4, [1; 2])
        @test_throws ArgumentError CAGD.getModelLoCell(m, 2, [6; 2])
        @test_throws ArgumentError CAGD.getModelLoCell(m, 2, [2; 6])
        @test_throws ArgumentError CAGD.getModelLoCell(m, 2, [-1; 1])
    end

    @testset "get Model Cell Vertices" begin
        @test CAGD.getModelCellVertices(m, 2, 3) == [
            [0.0, 0.0, 1.0],
            [0.0, 1.0, 0.0],
            [0.0, 1.0, 1.0]
        ]

        @test CAGD.getModelCellVertices(m, 2, [1; 3]) == [
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 1.0],
            [0.0, 1.0, 0.0],
            [0.0, 1.0, 1.0]
        ]

        @test CAGD.getModelCellVertices(m, 2, [3], true)[2] == [2; 3; 4]

        @test_throws ArgumentError CAGD.getModelCellVertices(m, 0, [1; 2])
        @test_throws ArgumentError CAGD.getModelCellVertices(m, 4, [1; 2])
        @test_throws ArgumentError CAGD.getModelCellVertices(m, 2, [6; 2])
        @test_throws ArgumentError CAGD.getModelCellVertices(m, 2, [2; 6])
        @test_throws ArgumentError CAGD.getModelCellVertices(m, 2, [-1; 1])
    end

    @testset "get Model Cell Geometry" begin
        @test CAGD.getModelCellGeometry(m, 2, 3) == [
            0.0 0.0 0.0
            0.0 1.0 1.0
            1.0 0.0 1.0
        ]

        @test CAGD.getModelCellGeometry(m, 2, [1; 3]) == [
            0.0 0.0 0.0 0.0
            0.0 0.0 1.0 1.0
            0.0 1.0 0.0 1.0
        ]

        @test CAGD.getModelCellGeometry(m, 2, [3], true)[2] == [2; 3; 4]

        @test_throws ArgumentError CAGD.getModelCellGeometry(m, 0, [1; 2])
        @test_throws ArgumentError CAGD.getModelCellGeometry(m, 4, [1; 2])
        @test_throws ArgumentError CAGD.getModelCellGeometry(m, 2, [6; 2])
        @test_throws ArgumentError CAGD.getModelCellGeometry(m, 2, [2; 6])
        @test_throws ArgumentError CAGD.getModelCellGeometry(m, 2, [-1; 1])
    end

    @testset "delete Model Cells" begin
        CAGD.deleteModelCell!(m, 2, 2)
        @test size(m, 2) == (2, 5)
        CAGD.addModelCell!(m, 2, [1; 2; 4; 5])
        CAGD.deleteModelCells!(m, 1, [4; 5])
        @test size(m, 2) == (1, 3)
        @test_throws ArgumentError CAGD.deleteModelCells!(m, 1, [4; 5])
        @test_throws ArgumentError CAGD.deleteModelCells!(m, 0, [4; 5])
        @test_throws ArgumentError CAGD.deleteModelCells!(m, 4, [4; 5])
    end

end

@testset "Geometry Manipulation" begin
    G1 = [0. 1.; 0. 1.];
    G2 = [1. 0.; 0. 1.];
    m = CAGD.Model(G1);

    @testset "add model vertices" begin
        CAGD.addModelCell!(m, 1, [1, 2]);
        CAGD.addModelVertices!(m, G2)
        @test m.G == [G1 G2]
        @test size(m, 1) == (1, 4)
        @test m.T[1] == SparseArrays.sparse(Int8[1 1 0 0])
    end

    @testset "delete model vertices" begin
        CAGD.deleteModelVertex!(m, 3)
        @test m.G == [0. 1. 0.; 0. 1. 1.];
        @test size(m, 1) == (1, 3)

        CAGD.deleteModelVertices!(m, [1, 3]);
        @test m.G == [1.; 1.][:, :]
        @test size(m, 1) == (0, 1)

        @test_throws ArgumentError CAGD.deleteModelVertices!(m, [1, 3])
        @test_throws ArgumentError CAGD.deleteModelVertices!(m, [0, 1])

    end
end

@testset "Model Manipulation" begin

    @testset "Purge Model" begin
        m = CAGD.Model([0.0 0.0 1.0 0.0 2.0; 0.0 1.0 0.0 2.0 0.0], [
            SparseArrays.sparse(Int8[-1 1 0 0 0; -1 0 1 0 0; 0 -1 1 0 0; 0 -1 0 1 0]),
            SparseArrays.sparse(Int8[1 -1 1 0 ])
        ])
        CAGD.modelPurge!(m)
        @test size(m, 1) == (4, 4)

        CAGD.modelPurge!(m, depth = 1)
        @test size(m, 1) == (3, 3)
    end

    @testset "Merge Model Vertices" begin
        m = CAGD.Model([
                0.0  0.0  1.0  0.0   0.0  -1.0  0.0  -1.0  -1.0
                0.0  1.0  0.0  0.0  -1.0   0.0  1.0   0.0   1.0
            ], [
                SparseArrays.sparse(Int8[
                    -1   1  0   0   0  0   0   0  0
                    -1   0  1   0   0  0   0   0  0
                     0  -1  1   0   0  0   0   0  0
                     0   0  0  -1   1  0   0   0  0
                     0   0  0  -1   0  1   0   0  0
                     0   0  0   0  -1  1   0   0  0
                     0   0  0   0   0  0  -1   1  0
                     0   0  0   0   0  0  -1   0  1
                     0   0  0   0   0  0   0  -1  1
                ]),
                SparseArrays.sparse(Int8[
                    1  -1  1  0   0  0  0   0  0
                    0   0  0  1  -1  1  0   0  0
                    0   0  0  0   0  0  1  -1  1
                ]),
        ])

        m1 = CAGD.mergeModelVertices(m, signed_merge = false);
        m2 = CAGD.mergeModelVertices(m, signed_merge = true);

        @test size(m1, 1) == (9, 6)
        @test size(m2, 1) == (9, 6)
        @test size(m1, 2) == (3, 9)
        @test size(m2, 2) == (3, 9)

        @test m1.G == m2.G
        @test sortslices(m1.T[1], dims = 1) == sortslices(abs.(m2.T[1]), dims = 1)

        @test m1.G == [
            0.0  0.0  1.0   0.0  -1.0  -1.0
            0.0  1.0  0.0  -1.0   0.0   1.0
        ]

        @test Matrix(m2.T[1]) == [
            -1   0  0   0   1  0
            -1   0  0   1   0  0
            -1   0  1   0   0  0
            -1   1  0   0   0  0
             0  -1  0   0   0  1
             0  -1  0   0   1  0
             0  -1  1   0   0  0
             0   0  0  -1   1  0
             0   0  0   0  -1  1
        ]

        @test Matrix(m2.T[2]) == [
            1  -1  0   0  0   0   0  -1   0
            0   0  1  -1  0   0  -1   0   0
            0   0  0   0  1  -1   0   0  -1
        ]
    end
end

@testset "Models Interaction" begin
    @testset "Unite Models" begin
        m1 = CAGD.Model([0.0 0.0 1.0; 0.0 1.0 0.0], [
            SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
            SparseArrays.sparse(Int8[1 -1 1])
        ])

        m2 = CAGD.Model([0.0 0.0 -1.0; 0.0 -1.0 0.0], [
            SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
            SparseArrays.sparse(Int8[1 -1 1])
        ])

        m3 = CAGD.Model([0.0 -1.0 -1.0; 1.0 0.0 1.0], [
            SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
            SparseArrays.sparse(Int8[1 -1 1])
        ])

        m4 = CAGD.uniteModels(m1, m2)
        CAGD.uniteModels!(m1, m2)

        @test m1 == m4
        @test size(m1, 1) == (6, 6)
        @test size(m1, 2) == (2, 6)

        models = [
            CAGD.Model([0.0 0.0 1.0; 0.0 1.0 0.0], [
                SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
                SparseArrays.sparse(Int8[1 -1 1])
            ]);

            CAGD.Model([0.0 0.0 -1.0; 0.0 -1.0 0.0], [
                SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
                SparseArrays.sparse(Int8[1 -1 1])
            ]);

            CAGD.Model([0.0 -1.0 -1.0; 1.0 0.0 1.0], [
                SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
                SparseArrays.sparse(Int8[1 -1 1])
            ])
        ];

        CAGD.uniteModels!(m1, m3)

        @test m1 == CAGD.uniteMultipleModels(models)
    end

    @testset "Merge Models" begin
        models = [
            CAGD.Model([0.0 0.0 1.0; 0.0 1.0 0.0], [
                SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
                SparseArrays.sparse(Int8[1 -1 1])
            ]);

            CAGD.Model([0.0 0.0 -1.0; 0.0 -1.0 0.0], [
                SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
                SparseArrays.sparse(Int8[1 -1 1])
            ]);

            CAGD.Model([0.0 -1.0 -1.0; 1.0 0.0 1.0], [
                SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
                SparseArrays.sparse(Int8[1 -1 1])
            ])
        ]

        m = CAGD.mergeMultipleModels(models, signed_merge = true)

        @test size(m, 1) == (9, 6)
        @test size(m, 2) == (3, 9)

        @test m.G == [
            0.0  0.0  1.0   0.0  -1.0  -1.0
            0.0  1.0  0.0  -1.0   0.0   1.0
        ]

        @test Matrix(m.T[1]) == [
            -1   0  0   0   1  0
            -1   0  0   1   0  0
            -1   0  1   0   0  0
            -1   1  0   0   0  0
             0  -1  0   0   0  1
             0  -1  0   0   1  0
             0  -1  1   0   0  0
             0   0  0  -1   1  0
             0   0  0   0  -1  1
        ]

        @test Matrix(m.T[2]) == [
            1  -1  0   0  0   0   0  -1   0
            0   0  1  -1  0   0  -1   0   0
            0   0  0   0  1  -1   0   0  -1
        ]
    end
end

@testset "Model Export" begin
    modelOut = CAGD.Model([0.0 0.0 1.0; 0.0 1.0 0.0], [
        SparseArrays.sparse(Int8[-1 1 0; -1 0 1; 0 -1 1]),
        SparseArrays.sparse(Int8[1 -1 1])
    ]);

    #CAGD.jlexportModel(modelOut, "./testModelExport.jl", "modelIn");

    #include("./testModelExport.jl");

    #@test modelIn == modelOut

    #rm("./testModelExport.jl")
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using Test

abstract type HeatProblem <: AbstractProblem

end

function HeatProblem(dim::Int=1, elements=[])
    return Problem{HeatProblem}(dim, elements)
end

function get_unknown_field_name(::Type{P}) where P<:HeatProblem
    return "temperature"
end

function get_unknown_field_type(::Type{P}) where P<:HeatProblem
    return Float64
end

""" Calculate a potential Π = Wint - Wext of system. """
function get_potential_energy(problem::Problem{HeatProblem}, element::Element{Quad4}, ip::IntegrationPoint, time::Number; variation=nothing)
    k = element("temperature thermal conductivity", ip, time)
    f = element("temperature load", ip, time)
    T = element("temperature", ip, time, variation)
    c = element("temperature nonlinearity coefficient", ip, time)
    gradT = element("temperature", ip, time, Val{:grad}, variation)
    Wint = (k + c*T) * 1/2*vecdot(gradT, gradT)
    Wext = f*T
    W = Wint - Wext
    J = get_jacobian(element, ip, time)
    return W*det(J)
end

function get_potential_energy(problem::Problem{HeatProblem}, element::Element{Seg2}, ip::IntegrationPoint, time::Number; variation=nothing)
    T = element("temperature", ip, time, variation)[1]
    T_ext = element("temperature external", ip, time)[1]
    coeff = element("temperature coefficient", ip, time)[1]
    q0 = coeff*(T_ext^4 - T^4)
    Wint = 0.0
    Wext = q0*T
    W = Wint - Wext
    J = get_jacobian(element, ip, time)
    return W*norm(J)
end

function test_potential_energy_method()

    # create model -- start
    element = Quad4([1, 2, 3, 4])
    element["geometry"] = Vector[[0.0,0.0], [1.0,0.0], [1.0,1.0], [0.0,1.0]]
    element["temperature thermal conductivity"] = 6.0
    element["temperature load"] = [0.0, 0.0, 0.0, 0.0]
    element["temperature nodal load"] = [3.0, 3.0, 0.0, 0.0]
    element["temperature nonlinearity coefficient"] = 6.0
    element["temperature"] = (0.0 => zeros(Float64, 4))
    problem = HeatProblem()
    push!(problem, element)
    # create model -- end

    solve!(problem, [1, 2], 0.0)
    temp = element("temperature", [0.0, -1.0], 0.0)
    err = temp - 2/3
    @info("error: $err")
    @test isapprox(err, 0.0)
end

function test_potential_energy_method_2()

    # create model -- start
    N = Vector[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]
    element1 = Quad4([1, 2, 3, 4])
    element1["geometry"] = Vector[N[1], N[2], N[3], N[4]]
    element1["temperature thermal conductivity"] = 6.0
    element1["temperature load"] = [0.0, 0.0, 0.0, 0.0]
    element1["temperature nonlinearity coefficient"] = [0.0, 0.0, 0.0, 0.0]
    element1["temperature"] = (0.0 => zeros(Float64, 4))

    element2 = Seg2([1, 2])
    element2["geometry"] = Vector[N[1], N[2]]
    element2["temperature coefficient"] = 3.0e-8 # ~ 5.7e-8 * 0.5
    element2["temperature external"] = 100.0
    element2["temperature"] = (0.0 => zeros(Float64, 2))
    # create model -- end

    problem = HeatProblem()
    push!(problem, element1)
    push!(problem, element2)
    solve!(problem, [1, 2], 0.0)

    temp = element1("temperature", [0.0, -1.0], 0.0)
    err = temp - 0.5
    @info("error: $err")
    @test isapprox(err, 0.0, atol=1.0e-6)

    # @test isapprox(temp, 2.93509690572300E+00)  # tested using Code Aster
end


function createscenarios{T1<:Real}(mpc::MPCSystem1D{T1}, t::Int,
                                   ω::UnivariateDistribution,
                                   includemean::Bool = true)
    # Assumes each disturbance to the system is i.i.d. ω for each time step
    T = mpc.system.T
    scenarios = rand(ω, T-t, mpc.numscenarios[t+1])
    if includemean == true
        scenarios[:,1] = mean(ω)
    end
    return scenarios
end

function predictedobjective{T1<:Real}(a, x::T1, t::Int, system::DynamicSystem1D,
                                      scenario::Vector{T1})
    numsteps = length(scenario)
    @assert t+numsteps == system.T
    @assert length(a) == numsteps
    retval = 0.
    Xs = copy(x)
    for step = 1:numsteps
        @inbounds retval += system.U(t+step, Xs, a[step], scenario[step])
        @inbounds Xs = system.f(t+step, Xs, a[step], scenario[step])
    end

    retval += system.Ubar(Xs)

    return -retval # Minimizing in Optim
end

function MCobjective{T1<:Real}(a, x::T1, t::Int, system::DynamicSystem1D,
                               scenarios::Array{T1,2})
    numsteps, numscenarios = size(scenarios)
    @assert length(a) == (numsteps-1)*numscenarios+1
    retval = 0.

    for i = 1:numscenarios
        j = (numsteps-1)*(i-1)
        idxs = [1; collect(j+2:j+numsteps)]
        @inbounds retval += predictedobjective(a[idxs], x, t, system, scenarios[:,i])
    end

    return retval/numscenarios
end

#==
# On smaller tests, this seems to be slower than doing ForwardDiff directly on MCobjective?
#
function gradpredictedobjective!{T1<:Real}(grad, x::T1, t::Int,
system::DynamicSystem1D,
scenario::Vector{T1})
objective(a) = predictedobjective(a, x, t, system, scenario)
ForwardDiff.gradient!(grad, objective, a)
end

function gradMCobjective!{T1<:Real}(grad, a, x::T1, t::Int, system::DynamicSystem1D,
scenarios::Array{T1,2})
numscenarios = size(scenarios,2)
tmpgrad = zeros(grad)
for i = 1:numscenarios
grad[:] += gradpredictedobjective!(tmpgrad, a, x, t, system, scenarios[:,i])
end
scale!(grad, 1/N)
end
==#

function onlinedecision(mpc::MPCSystem1D{T1}, x::T1,
                        t::Int,
                        aguess::Vector{T1},
                        ω::UnivariateDistribution,
                        includemean::Bool = true;
                        verbose::Bool = false,
                        optimizer::To = LBFGS(),
                        maxiter::Int = 1000,
                        autodiff::Symbol = :forward) where T1<: Real where To <: Optim.AbstractOptimizer
    system = mpc.system
    w = createscenarios(mpc, t, ω, includemean)
    objective(a) = MCobjective(a, x, t, system, w)

    # aguess contains the initial guess for
    # decisions a_t,\dots,a_{T-1}
    # ainit should take into account different decisions for different scenarios
    # (i.e. a[2:numsteps] should also depend on scenario)
    # so we should have a = [at, a[t+1:T-1,1], a[t+1:T-1,2],..., a[t+1:T-1,numscenarios]]
    numsteps, numscenarios = size(w)

    ainit = zeros(T1, (numsteps-1)*numscenarios+1)
    ainit[1] = aguess[1]
    for scenario = 1:numscenarios
        j = (numsteps-1)*(scenario-1)
        ainit[j+2:j+numsteps] = aguess[2:end]
    end

    if To <: Optim.SecondOrderOptimizer
        # TODO: allow reversediff?
        df = TwiceDifferentiable(objective, ainit; autodiff = autodiff)
    elseif To <: Optim.FirstOrderOptimizer
        df = OnceDifferentiable(objective, ainit; autodiff = autodiff)
    else
        df = NonDifferentiable(objective,ainit)
    end

    # res = optimize(df, ainit,
    #                optimizer,
    #                Optim.Options(show_trace=verbose, extended_trace=verbose,
    #                              iterations = maxiter))

    res = optimize(df, fill(system.amin, length(ainit)), fill(system.amax, length(ainit)),
                   ainit,
                   Fminbox(optimizer),
                   Optim.Options(
                       show_trace=verbose, extended_trace=verbose,
                       iterations = maxiter, outer_iterations = maxiter,
                       allow_f_increases=true))

    return Optim.minimizer(res)
end


using Plots
using Printf
using Oceananigans
using Oceananigans.Units
using Oceananigans.TurbulenceClosures: ConvectiveAdjustmentVerticalDiffusivity

grid = RectilinearGrid(size=32, z=(-64, 0), topology=(Flat, Flat, Bounded))

closure = ConvectiveAdjustmentVerticalDiffusivity(convective_κz = 1,
                                                  background_κz = 1e-5,
                                                  convective_νz = 1e-3,
                                                  background_νz = 1e-4)

Qᵇ = +1e-8
Qᵘ = +1e-5

b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ))
u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))

model = HydrostaticFreeSurfaceModel(grid = grid,
                                    tracers = :b,
                                    buoyancy = BuoyancyTracer(),
                                    coriolis = FPlane(f=1e-4),
                                    boundary_conditions = (; b=b_bcs, u=u_bcs),
                                    closure = closure)

N² = 1e-5
bᵢ(x, y, z) = N² * z
set!(model, b = bᵢ)

z = znodes(model.tracers.b)

b = view(interior(model.tracers.b), 1, 1, :)
u = view(interior(model.velocities.u), 1, 1, :)
v = view(interior(model.velocities.v), 1, 1, :)

b_plot = plot(b, z, linewidth = 2, label = "t = 0", xlabel = "Buoyancy", ylabel = "z", legend=:bottomright)

u_plot = plot(u, z, linewidth = 2, label = "u at t = 0",
              linestyle = :solid, color = :black, xlabel = "velocity", ylabel = "z", legend = :bottomright)

plot!(u_plot, v, z, linewidth = 2, label = "v at t = 0",
      linestyle = :dash, color = :black)

simulation = Simulation(model, Δt = 20.0, stop_time = 48hours)

run!(simulation)

plot!(b_plot, b, z, linewidth = 2, label = @sprintf("t = %s", prettytime(model.clock.time)))

plot!(u_plot, u, z, label = @sprintf("u at t = %s", prettytime(model.clock.time)),
      linewidth = 2, linestyle = :dash, color = :blue)

plot!(u_plot, v, z, label = @sprintf("v at t = %s", prettytime(model.clock.time)),
      linewidth = 2, linestyle = :solid, color = :blue)

ub_plot = plot(b_plot, u_plot, layout=(1, 2))

display(ub_plot)