"""
```
    GammaCorrection <: AbstractHistogramAdjustmentAlgorithm
    GammaCorrection(; gamma = 1)

    adjust_histogram([T,] img, f::GammaCorrection)
    adjust_histogram!([out,] img, f::GammaCorrection)
```

Returns a gamma corrected image.

# Details

Gamma correction is a non-linear  transformation given by the relation
```math
f(x) = x^\\gamma \\quad \\text{for} \\; x \\in \\mathbb{R}, \\gamma > 0.
```
It is called a *power law* transformation because one quantity varies as a power
of another quantity.

Gamma correction has historically been used to preprocess
an image to compensate for the fact that the intensity of light generated by a
physical device is not usually a linear function of the applied signal but
instead follows a power law [1]. For example, for many Cathode Ray Tubes (CRTs) the
emitted light intensity on the display is approximately equal to the voltage
raised to the power of γ, where γ ∈ [1.8, 2.8]. Hence preprocessing a raw image with
an exponent of 1/γ  would have ensured a linear response to brightness.

Research in psychophysics has also established an [empirical  power law
](https://en.wikipedia.org/wiki/Stevens%27s_power_law)  between light intensity and perceptual
brightness. Hence, gamma correction often serves as a useful image enhancement
tool.

# Options

Various options for the parameters of the `adjust_histogram` function and the
`Gamma` type are described in more detail below.

## Choices for `img`

The function can handle a variety of input types. The returned
image depends on the input type.

For colored images, the input is converted to YIQ type and the Y channel is
gamma corrected. This is the combined with the I and Q channels and the
resulting image converted to the same type as the input.

## Choice for `gamma`

The `gamma` value must be a non-zero positive number. A `gamma` value less than
one will yield a brighter image whereas a value greater than one will produce a
darker image. If left unspecified a default value of one is assumed.

# Example

```julia
using ImageContrastAdjustment, ImageView

# Create an example image consisting of a linear ramp of intensities.
n = 32
intensities = 0.0:(1.0/n):1.0
img = repeat(intensities, inner=(20,20))'

# Brighten the dark tones.
imgadj = adjust_histogram( img, GammaCorrection(gamma = 1/2))

# Display the original and adjusted image.
imshow(img)
imshow(imgadj)
```

# References
1. W. Burger and M. J. Burge. *Digital Image Processing*. Texts in Computer Science, 2016. [doi:10.1007/978-1-4471-6684-9](https://doi.org/10.1007/978-1-4471-6684-9)
"""
Base.@kwdef struct GammaCorrection{T <: Real} <: AbstractHistogramAdjustmentAlgorithm
    gamma::T = 1.0
end

function (f::GammaCorrection)(out::GenericGrayImage, img::GenericGrayImage)
    out .= img
    correct_gamma!(out, f.gamma)
    return out
end

# TODO Expand this to more generic gray color types.
function correct_gamma!(img::AbstractArray{Gray{T}},  gamma::Real) where T <: FixedPointNumbers.Normed
    γ = Float64(gamma)
    # Create a lookup-table for the gamma transformation of the grayvalues.
    raw_type = FixedPointNumbers.rawtype(T)
    table = zeros(T, typemax(raw_type) + 1)
    for i in zero(raw_type):typemax(raw_type)
        table[i + 1] = T((i / typemax(raw_type))^γ)
    end
    # Map the pixels to their new grayvalues.
    map!(img,img) do p
        table[p.val.i + 1]
    end
end

function correct_gamma!(img::GenericGrayImage,  gamma::Real)
    γ = Float64(gamma)
    T = eltype(img)
    map!(img,img) do val
        if isnan(val)
            return val
        else
            return  T <: Integer ? round(Int, val^γ) : val^γ
        end
    end
end

function (f::GammaCorrection)(out::AbstractArray{<:Color3}, img::AbstractArray{<:Color3})
    T = eltype(out)
    yiq = convert.(YIQ, img)
    yiq_view = channelview(yiq)
    adjust_histogram!(view(yiq_view,1,:,:), f)
    out .= convert.(T, yiq)
end

(f::GammaCorrection)(out::GenericGrayImage, img::AbstractArray{<:Color3}) =
    f(out, of_eltype(Gray, img))


# Before running this, please make sure to activate and instantiate the environment
# corresponding to [this `Project.toml`](https://raw.githubusercontent.com/alan-turing-institute/DataScienceTutorials.jl/master/Project.toml) and [this `Manifest.toml`](https://raw.githubusercontent.com/alan-turing-institute/DataScienceTutorials.jl/master/Manifest.toml)
# so that you get an environment which matches the one used to generate the tutorials:
#
# ```julia
# cd("DataScienceTutorials") # cd to folder with the *.toml
# using Pkg; Pkg.activate("."); Pkg.instantiate()
# ```

using MLJ
import DataFrames
import Statistics
using PrettyPrinting
using StableRNGs

rng = StableRNG(512)
Xraw = rand(rng, 300, 3)
y = exp.(Xraw[:,1] - Xraw[:,2] - 2Xraw[:,3] + 0.1*rand(rng, 300))
X = DataFrames.DataFrame(Xraw)

train, test = partition(eachindex(y), 0.7);

@load KNNRegressor
knn_model = KNNRegressor(K=10)

knn = machine(knn_model, X, y)

fit!(knn, rows=train)
ŷ = predict(knn, X[test, :]) # or use rows=test
rms(ŷ, y[test])

evaluate!(knn, resampling=Holdout(fraction_train=0.7, rng=StableRNG(666)),
          measure=rms) |> pprint

ensemble_model = EnsembleModel(atom=knn_model, n=20);

ensemble = machine(ensemble_model, X, y)
estimates = evaluate!(ensemble, resampling=CV())
estimates |> pprint

@show estimates.measurement[1]
@show mean(estimates.per_fold[1])

params(ensemble_model) |> pprint

B_range = range(ensemble_model, :bagging_fraction,
                lower=0.5, upper=1.0)
K_range = range(ensemble_model, :(atom.K),
                lower=1, upper=20);

tm = TunedModel(model=ensemble_model,
                tuning=Grid(resolution=10), # 10x10 grid
                resampling=Holdout(fraction_train=0.8, rng=StableRNG(42)),
                ranges=[B_range, K_range])

tuned_ensemble = machine(tm, X, y)
fit!(tuned_ensemble, rows=train);

best_ensemble = fitted_params(tuned_ensemble).best_model
@show best_ensemble.atom.K
@show best_ensemble.bagging_fraction

r = report(tuned_ensemble);

using PyPlot

figure(figsize=(8,6))

res = r.plotting
vals_b = res.parameter_values[:, 1]
vals_k = res.parameter_values[:, 2]

tricontourf(vals_b, vals_k, res.measurements)
xticks(0.5:0.1:1, fontsize=12)
xlabel("Bagging fraction", fontsize=14)
yticks([1, 5, 10, 15, 20], fontsize=12)
ylabel("Number of neighbors - K", fontsize=14)

ŷ = predict(tuned_ensemble, rows=test)
rms(ŷ, y[test])

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


ENV["MPLBACKEND"]="Qt5Agg"
include("CreateMesh.jl")
using LoopVectorization
using PyPlot
pygui(true)

function tplot(p, t, v)
    # Plot triangular mesh with nodes `p` and triangles `t`
    clf()
    tris = convert(Array{Int64}, hcat(t...)')
    # fig1, ax1 = plt.subplots()
    display(tripcolor(first.(p), last.(p), tris .- 1, v,
              cmap="viridis", edgecolors="none", linewidth=0))
    axis("equal")
    ylim([0, 1])
    xlim([0, 1])
    title("Low res.")
    xlabel("x")
    ylabel("y")
    colorbar()
    show()
    return
end

function qplot(x, y, v)
    clf()
    display( pcolor(x, y, v) )
    colorbar()
    xlabel("x")
    ylabel("y")
    show()
end

function Tanalytic2!( mesh, xc , yc, T, a, b, c, d, alp, bet )
    if mesh.type == "triangle" nn_el = 3.0; nn_el_int=3; end
    if mesh.type == "quad"     nn_el = 4.0; nn_el_int=4; end
    # Evaluate T analytic on barycentres
    for j=1:mesh.nn_el # as fast as above and less cryptic
        @. xc = 0.0
        @. yc = 0.0
        @. xc = xc + 1.0/mesh.nn_el * mesh.xn[mesh.e2n[j,:]]
        @. yc = yc + 1.0/mesh.nn_el * mesh.yn[mesh.e2n[j,:]]
    end
    @. T = exp(alp*sin(a*xc + c*yc) + bet*cos(b*xc + d*yc));
    # xn = mesh.xn
    # yn = mesh.yn
    # e2n = mesh.e2n
    #  for i=1:mesh.nel
    #     xc[i] = 0;
    #     yc[i] = 0;
    #     @avx for j=1:3
    #         # xc[i] += 1.0/nn_el * mesh.xn[mesh.e2n[j,i]]
    #         # yc[i] += 1.0/nn_el * mesh.yn[mesh.e2n[j,i]]
    #         xc[i] += 1.0/nn_el * xn[e2n[j,i]]
    #         yc[i] += 1.0/nn_el * yn[e2n[j,i]]
    #     end
    # end
    # for i=1:mesh.nel
    #     T[i] = exp(alp.sin(a*xc[i] + c*yc[i]) + bet*cos(b*xc[i] + d*yc[i]));
    # end
return
end

function main()

    # Create sides of mesh
    xmin, xmax = 0, 1
    ymin, ymax = 0, 1
    nx, ny     = 20, 20
    quad = false

    if quad==true
        mesh = MakeQuadMesh( nx, ny, xmin, xmax, ymin, ymax )
    elseif quad==false
        mesh = MakeTriangleMesh( nx, ny, xmin, xmax, ymin, ymax )
    end

    println(mesh.nel)

    # Generate function to be visualised on mesh
    alp = 0.1; bet = 0.3; a = 5.1; b = 4.3; c = -6.2; d = 3.4;
    xc = zeros(mesh.nel)
    yc = zeros(mesh.nel)
    T  = zeros(mesh.nel)
    ncalls = 4
    # A) Loop version with @avx
    # @printf("Looped, %d times:\n", ncalls)
    for icall=1:ncalls
        @time Tanalytic2!(mesh, xc ,yc , T, a, b, c, d, alp, bet)
    end

    if mesh.type == "triangle"  tplot(mesh.xyviz, mesh.e2nviz,                    T ) end
    if mesh.type == "quad"      qplot(mesh.xcviz, mesh.ycviz,  reshape(T, (nx,ny) ) ) end

end

main()


delete_missing(x) = [collect(skipmissing(x[:,i])) for i in 1:ncol(x)]

function bin_dissociations(data, n_bins; bin_type = "force", edges=nothing)
    F = data[1]
    t = data[2]
    # even force bins
    if edges != nothing
        bins,bin_ids = partition(F,edges)
        binned_t = [similar(t,0) for i in bins]

    elseif bin_type == "force"
        f_min = minimum(F)
        f_max = maximum(F)

        edges = range(f_min,stop=f_max,length=n_bins+1)
        bins,bin_ids = partition(F,edges)
        binned_t = [similar(t,0) for i in bins]

    elseif bin_type == "equal number"
        edges_ind = Int64.(round.(range(1,stop=length(F),length=n_bins+1)))
        edges = F[edges_ind]
        # edges = [0,9,12,15,19,30]
        bins,bin_ids = partition(F,edges)
        binned_t = [similar(t,0) for i in bins]
    end
    # h = fit(Histogram, F)

    for (i,bin) in enumerate(bin_ids)
        push!(binned_t[bin],t[i])
    end

    return bins, binned_t, bin_ids
end

function partition(data, intervals)
    ranges = intervals[1:end-1] .=> intervals[2:end]
    bins = [similar(data, 0) for _ in 1:length(ranges)]
    ids = similar(data,0)
    for x in data
        for (i, (a, b)) in pairs(ranges)
            if i != length(bins) && a <= x < b
                push!(bins[i], x)
                push!(ids,i)
                break

            elseif i == length(bins) && a <= x <= b
                push!(bins[i], x)
                push!(ids,i)
                break
        end
        end
    end
    return bins, Int64.(ids)
end

function resort(f,t)
    t_sorted = sort.(t)
    sort_id = sortperm.(t)
    f_sorted = [f[i][sort_id[i]] for i in 1:length(sort_id)]
    return f_sorted, t_sorted
end

survival_p(t) = collect((length(t) :-1:1) / length(t))

p2lnp(p) = map((x) -> log.(x), p)


# 1-dim version of the code

using Winston

#basic parameters
gamma = 1.4
cfl = 0.8
dt = 1.0e-5
dtp = dt

nzones = 100
tend = 0.2

type data1d
    x::Array{Float64,1} #cell centers
    xi::Array{Float64,1} #cell LEFT interfaces

    rho::Array{Float64,1}
    rhop::Array{Float64,1}
    rhom::Array{Float64,1}

    vel::Array{Float64,1}
    velp::Array{Float64,1}
    velm::Array{Float64,1}

    eps::Array{Float64,1}
    epsp::Array{Float64,1}
    epsm::Array{Float64,1}

    press::Array{Float64,1}
    pressp::Array{Float64,1}
    pressm::Array{Float64,1}

    q::Array{Float64,2} #conserved quantities
    qp::Array{Float64,2}
    qm::Array{Float64,2}

    n::Int64
    g::Int64 #ghost cells

    function data1d(nzones::Int64)

        x = zeros(nzones)
        xi = zeros(nzones)

        rho = zeros(nzones)
        rhop = zeros(nzones)
        rhom = zeros(nzones)

        vel = zeros(nzones)
        velp = zeros(nzones)
        velm = zeros(nzones)

        eps = zeros(nzones)
        epsp = zeros(nzones)
        epsm = zeros(nzones)

        press = zeros(nzones)
        pressp = zeros(nzones)
        pressm = zeros(nzones)

        q = zeros(nzones, 3)
        qp = zeros(nzones, 3)
        qm = zeros(nzones, 3)
        n = nzones
        g = 3

        new(x,
            xi,
            rho,
            rhop,
            rhom,
            vel,
            velp,
            velm,
            eps,
            epsp,
            epsm,
            press,
            pressp,
            pressm,
            q,
            qp,
            qm,
            n,
            g)
    end
end

#1-dim grid
function grid_setup(self::data1d, xmin, xmax)
    dx = (xmax - xmin) / (self.n -self.g*2 - 1)
    xmin = xmin - self.g*dx
    xmax = xmax + self.g*dx

    for i = 1:self.n
        self.x[i] = xmin + i*dx
    end

    for i = 1:self.n
        self.xi[i] = self.x[i] - 0.5dx
    end

    return self
end

#Shoctube initial data
function setup_tube(self::data1d)
    rchange = 0.5(self.x[self.n - self.g] - self.x[self.g + 1])

    rho1 = 1.0
    rho2 = 0.125
    press1 = 1.0
    press2 = 0.1

    for i = 1:self.n
        if self.x[i] < rchange
            self.rho[i] = rho1
            self.press[i] = press1
            self.eps[i] = press1/rho1/(gamma - 1.0)
            self.vel[i] = 0.0
        else
            self.rho[i] = rho2
            self.press[i] = press2
            self.eps[i] = press2/rho2/(gamma - 1.0)
            self.vel[i] = 0.0
        end
    end

    return self
end

#Shoctube initial data
function setup_blast(self::data1d)

    rho1 = 0.4
    press1 = 0.2
    press2 = 0.99

    self.rho[:] = rho1*ones(self.n)
    self.press[:] = press1*ones(self.n)
    self.eps[:] = press1./rho1./(gamma - 1.0)
    self.vel[:] = zeros(self.n)

    mid = int(self.n/2)
    self.press[mid-2:mid+2] = press2
    self.eps[mid-2:mid+2] = press2./rho1./(gamma - 1.0)

    return self
end

function apply_bcs(hyd::data1d)

    #arrays starting from zero
    #       |g                  |n-g #
    #[0 1 2 x x x  .....  x x x 7 8 9]

    #arrays starting from 1
    #     |g                  |n-g    #
    #[1 2 3 x x x  .....  x x x 8 9 10]
    hyd.rho[1:hyd.g] = hyd.rho[hyd.g+1]
    hyd.vel[1:hyd.g] = hyd.vel[hyd.g+1]
    hyd.eps[1:hyd.g] = hyd.eps[hyd.g+1]
    hyd.press[1:hyd.g] = hyd.press[hyd.g+1]

    hyd.rho[(hyd.n-hyd.g+1) : hyd.n] = hyd.rho[hyd.n-hyd.g]
    hyd.vel[(hyd.n-hyd.g+1) : hyd.n] = hyd.vel[hyd.n-hyd.g]
    hyd.eps[(hyd.n-hyd.g+1) : hyd.n] = hyd.eps[hyd.n-hyd.g]
    hyd.press[(hyd.n-hyd.g+1) : hyd.n] = hyd.press[hyd.n-hyd.g]

    return hyd
end

#equation of state (ideal gas)
function eos_press(rho, eps, gamma)
    press = (gamma - 1.0) .* rho .* eps
    return press
end

#soundspeed
function eos_cs2(rho, eps, gamma)
    prs = (gamma - 1.0) .* rho .* eps
    dpde = (gamma - 1.0) .* rho
    dpdrho = (gamma - 1.0) .* eps
    cs2 = dpdrho .+ dpde .* prs ./ (rho + 1.0e-30).^2.0
    return cs2
end

#1-dim
function prim2con(rho::AbstractVector,
                  vel::AbstractVector,
                  eps::AbstractVector)

    q = zeros(size(rho, 1), 3)
    q[:,1] = rho
    q[:,2] = rho .* vel
    q[:,3] = rho .* eps .+ 0.5rho .* vel.^2.0

    return q
end

#1-dim
function con2prim(q)
    rho = vec(q[:,1])
    vel = vec(q[:,2] ./ rho)
    eps = vec(q[:,3] ./ rho - 0.5vel.^2.0)
    press = eos_press(rho, eps, gamma)

    return rho, eps, press, vel
end

#time step calculation
function calc_dt(hyd::data1d, dtp)
    cs = sqrt(eos_cs2(hyd.rho, hyd.eps, gamma))
    dtnew = 1.0
    for i = (hyd.g+1):(hyd.n-hyd.g+1)
        dtnew = min(dtnew, (hyd.x[i+1] - hyd.x[i]) / max(abs(hyd.vel[i]+cs[i]), abs(hyd.vel[i]-cs[i])))
    end

    dtnew = min(cfl*dtnew, 1.05*dtp)

    return dtnew
end

function minmod(a,b)
    if a*b < 0.0
        return 0.0
    elseif abs(a) < abs(b)
        return a
    else
        return b
    end
end

signum(x,y) = y >= 0.0 ? abs(x) : -abs(x)

function tvd_mc_reconstruction(n, g, f, x, xi)
    fp = zeros(n)
    fm = zeros(n)

    for i = g:(n-g+2)
        dx_up = x[i] - x[i-1]
        dx_down = x[i+1] - x[i]
        dx_m = x[i] -xi[i]
        dx_p = xi[i+1] - x[i]
        df_up = (f[i]-f[i-1]) / dx_up
        df_down = (f[i+1]-f[i]) / dx_down

        if df_up*df_down < 0.0
            delta = 0.0
        else
            delta = signum(min(2.0abs(df_up), 2.0abs(df_down), 0.5(abs(df_up)+abs(df_down))), df_up + df_down)
        end

        fp[i] = f[i] + delta*dx_p
        fm[i] = f[i] - delta*dx_m
    end

    return fp, fm
end

function reconstruct(hyd::data1d)

    hyd.rhop, hyd.rhom = tvd_mc_reconstruction(hyd.n,
                                               hyd.g,
                                               hyd.rho,
                                               hyd.x,
                                               hyd.xi)
    hyd.epsp, hyd.epsm = tvd_mc_reconstruction(hyd.n,
                                               hyd.g,
                                               hyd.eps,
                                               hyd.x,
                                               hyd.xi)
    hyd.velp, hyd.velm = tvd_mc_reconstruction(hyd.n,
                                               hyd.g,
                                               hyd.vel,
                                               hyd.x,
                                               hyd.xi)

    hyd.pressp = eos_press(hyd.rhop, hyd.epsp, gamma)
    hyd.pressm = eos_press(hyd.rhom, hyd.epsm, gamma)

    hyd.qp = prim2con(hyd.rhop, hyd.velp, hyd.epsp)
    hyd.qm = prim2con(hyd.rhom, hyd.velm, hyd.epsm)

    return hyd
end

#Load solvers
include("rsolvers1d.jl")

function calc_rhs(hyd::data1d)
    #reconstruction and prim2con
    hyd = reconstruct(hyd)

    #compute flux difference
    fluxdiff = hlle(hyd)

    #return RHS = -fluxdiff
    return -fluxdiff
end

###############
# main program
###############

#initialize
hyd = data1d(nzones)

#set up grid
hyd = grid_setup(hyd, 0.0, 1.0)

#set up initial data
hyd = setup_tube(hyd)

#get initial timestep
dt = calc_dt(hyd, dt)

#initial prim2con
hyd.q = prim2con(hyd.rho, hyd.vel, hyd.eps)

t = 0.0
i = 1

tsave = collect(0.0:0.01:0.2)

k = 1
while t < tend

    if i % 10 == 0
        #sleep(1.0)
        #output
        println("$i $t $dt")
        p=plot(hyd.x, hyd.rho, "r-")
        p=oplot(hyd.x, hyd.vel, "b-")
        p=oplot(hyd.x, hyd.press, "g-")
        display(p)
    end

    if (t >= tsave[k])
        println("t = $t")
        wmatr = zeros(nzones, 5)

        wmatr[:,1] = hyd.x
        wmatr[:,2] = hyd.rho
        wmatr[:,3] = hyd.vel
        wmatr[:,4] = hyd.eps
        wmatr[:,5] = hyd.press

        writecsv("step_$k.csv", wmatr)

        k += 1
    end

    #calculate new timestep
    dt = calc_dt(hyd, dt)

    #save old state
    hydold = hyd
    qold = hyd.q

    #calc rhs
    k1 = calc_rhs(hyd)
    #calculate intermediate step
    hyd.q = qold + 0.5dt*k1
    #con2prim
    hyd.rho, hyd.eps, hyd.press, hyd.vel = con2prim(hyd.q)
    #boundaries
    hyd = apply_bcs(hyd)

    #calc rhs
    k2 = calc_rhs(hyd)
    #apply update
    hyd.q = qold + dt*(0.5k1 + 0.5k2)
    #con2prim
    hyd.rho, hyd.eps, hyd.press, hyd.vel = con2prim(hyd.q)
    #apply bcs
    hyd = apply_bcs(hyd)

    #update time
    t += dt
    i += 1

end


println("hello world")

function quadratic2(a::Float64, b::Float64, c::Float64)
    # unlike other languages 2a is equivalent to 2*a
    # a^2 is used instead of a**2 or pow(a,2)
    sqr_term = sqrt(b^2-4a*c)
    r1 = quadratic(a, sqr_term, b)
    r2 = quadratic(a, -sqr_term, b)
    # multiple values can be returned from a function using tuples
    # if the return keyword is omitted, the last term is returned
    r1, r2
end


# ---
# title: 334. Increasing Triplet Subsequence
# id: problem334
# author: Indigo
# date: 2021-03-12
# difficulty: Medium
# categories:
# link: <https://leetcode.com/problems/increasing-triplet-subsequence/description/>
# hidden: true
# ---
#
# Given an unsorted array return whether an increasing subsequence of length 3
# exists or not in the array.
#
# Formally the function should:
#
# > Return true if there exists _i, j, k_
# >  such that _arr[i]_ < _arr[j]_ < _arr[k]_ given 0 ≤ _i_ < _j_ < _k_ ≤ _n_ -1
# > else return false.
#
# **Note:** Your algorithm should run in O( _n_ ) time complexity and O( _1_ )
# space complexity.
#
# **Example 1:**
#
#
#
#     Input: [1,2,3,4,5]
#     Output: true
#
#
# **Example 2:**
#
#
#
#     Input: [5,4,3,2,1]
#     Output: false
#
#
#
## @lc code=start
using LeetCode

function increasing_triplet(nums::Vector{Int})
    fst_min = scd_min = nums[1]
    len = length(nums)
    idx = 2
    while idx <= len
        (nums[idx] != fst_min) && break
        idx += 1
    end
    (idx >= len) && return false
    fst_min, scd_min = minmax(fst_min, nums[idx])
    for num in @view(nums[(idx + 1):end])
        (num > scd_min) && return true
        scd_min = num
        fst_min, scd_min = minmax(fst_min, scd_min)
    end
    return false
end
## @lc code=end


# ---
# title: 670. Maximum Swap
# id: problem670
# author: Indigo
# date: 2021-06-28
# difficulty: Medium
# categories: Array, Math
# link: <https://leetcode.com/problems/maximum-swap/description/>
# hidden: true
# ---
#
# Given a non-negative integer, you could swap two digits **at most** once to
# get the maximum valued number. Return the maximum valued number you could get.
#
# **Example 1:**
#
#
#
#     Input: 2736
#     Output: 7236
#     Explanation: Swap the number 2 and the number 7.
#
#
# **Example 2:**
#
#
#
#     Input: 9973
#     Output: 9973
#     Explanation: No swap.
#
#
# **Note:**
#
#   1. The given number is in the range [0, 108]
#
#
## @lc code=start
using LeetCode

function maximum_swap670(nums::Int)
    dgs = digits(nums)
    for i in 1:length(dgs)
        cur = @view(dgs[i:end])
        max_idx = argmax(cur)
        max_idx == lastindex(cur) && continue
        dgs[end], dgs[max_idx] = dgs[max_idx], dgs[end]
        break
    end
    return [10 .^ (0:length(dgs)-1);]' * dgs
end
## @lc code=end


# From Program 10
using LinearAlgebra

function meshgrid(x, y)
    x = reshape(x, 1, length(x))
    y = reshape(y, length(y), 1)
    X = repeat(x, length(y), 1)
    Y = repeat(y, 1, length(x))
    return X, Y
end


using Test, LinearAlgebra, LazyArrays, StaticArrays, FillArrays
import LazyArrays: MulAdd, MemoryLayout, DenseColumnMajor, DiagonalLayout, SymTridiagonalLayout
import Base.Broadcast: materialize, materialize!

@testset "Mul" begin
    @testset "eltype" begin
        @inferred(eltype(Mul(zeros(Int,2,2), zeros(Float64,2)))) == Float64
        @inferred(eltype(Mul(zeros(ComplexF16,2,2),zeros(Int,2,2),zeros(Float64,2)))) == ComplexF64

        v = Mul(zeros(Int,2,2), zeros(Float64,2))
        A = Mul(zeros(Int,2,2), zeros(Float64,2,2))
        @test @inferred(axes(v)) == (@inferred(axes(v,1)),) == (Base.OneTo(2),)
        @test @inferred(size(v)) == (@inferred(size(v,1)),) == (2,)
        @test @inferred(axes(A)) == (@inferred(axes(A,1)),@inferred(axes(A,2))) == (Base.OneTo(2),Base.OneTo(2))
        @test @inferred(size(A)) == (@inferred(size(A,1)),size(A,2)) == (2,2)

        Ã = materialize(A)
        @test Ã isa Matrix{Float64}
        @test Ã == zeros(2,2)
    end

    @testset "gemv Float64" begin
        for A in (randn(5,5), view(randn(5,5),:,:), view(randn(5,5),1:5,:),
                  view(randn(5,5),1:5,1:5), view(randn(5,5),:,1:5)),
            b in (randn(5), view(randn(5),:), view(randn(5),1:5), view(randn(9),1:2:9))
            c = similar(b);

            c .= Mul(A,b)
            @test all(c .=== A*b .=== BLAS.gemv!('N', 1.0, A, b, 0.0, similar(c)))

            copyto!(c, Mul(A,b))
            @test all(c .=== A*b .=== BLAS.gemv!('N', 1.0, A, b, 0.0, similar(c)))

            b̃ = copy(b)
            copyto!(b̃, Mul(A,b̃))
            @test all(c .=== b̃)

            c .= 2.0 .* Mul(A,b)
            @test all(c .=== BLAS.gemv!('N', 2.0, A, b, 0.0, similar(c)))

            c = copy(b)
            c .= Mul(A,b) .+ c
            @test all(c .=== BLAS.gemv!('N', 1.0, A, b, 1.0, copy(b)))

            c = copy(b)
            c .= Mul(A,b) .+ 2.0 .* c
            @test all(c .=== BLAS.gemv!('N', 1.0, A, b, 2.0, copy(b)))

            c = copy(b)
            c .= 2.0 .* Mul(A,b) .+ c
            @test all(c .=== BLAS.gemv!('N', 2.0, A, b, 1.0, copy(b)))

            c = copy(b)
            c .= 3.0 .* Mul(A,b) .+ 2.0 .* c
            @test all(c .=== BLAS.gemv!('N', 3.0, A, b, 2.0, copy(b)))

            d = similar(c)
            c = copy(b)
            d .= 3.0 .* Mul(A,b) .+ 2.0 .* c
            @test all(d .=== BLAS.gemv!('N', 3.0, A, b, 2.0, copy(b)))
        end
    end

    @testset "gemv mixed array types" begin
        (A, b, c) = (randn(5,5), randn(5), 1.0:5.0)
        d = similar(b)
        d .= Mul(A,b) .+ c
        @test all(d .=== BLAS.gemv!('N', 1.0, A, b, 1.0, Vector{Float64}(c)))

        d .= Mul(A,b) .+ 2.0 .* c
        @test all(d .=== BLAS.gemv!('N', 1.0, A, b, 2.0, Vector{Float64}(c)))

        d .= 2.0 .* Mul(A,b) .+ c
        @test all(d .=== BLAS.gemv!('N', 2.0, A, b, 1.0, Vector{Float64}(c)))

        d .= 3.0 .* Mul(A,b) .+ 2.0 .* c
        @test all(d .=== BLAS.gemv!('N', 3.0, A, b, 2.0, Vector{Float64}(c)))

        @test (similar(b) .= 1 .* Mul(A,b)) ≈
            (similar(b) .= 1 .* Mul(A,b) .+ 0 .* b) ≈
            A*b
    end

    @testset "gemv Complex" begin
        for T in (ComplexF64,),
                A in (randn(T,5,5), view(randn(T,5,5),:,:)),
                b in (randn(T,5), view(randn(T,5),:))
            c = similar(b);

            c .= Mul(A,b)
            @test all(c .=== BLAS.gemv!('N', one(T), A, b, zero(T), similar(c)))

            b .= Mul(A,b)
            @test all(c .=== b)

            c .= 2one(T) .* Mul(A,b)
            @test all(c .=== BLAS.gemv!('N', 2one(T), A, b, zero(T), similar(c)))

            c = copy(b)
            c .= Mul(A,b) .+ c
            @test all(c .=== BLAS.gemv!('N', one(T), A, b, one(T), copy(b)))

            c = copy(b)
            c .= Mul(A,b) .+ 2one(T) .* c
            @test all(c .=== BLAS.gemv!('N', one(T), A, b, 2one(T), copy(b)))

            c = copy(b)
            c .= 2one(T) .* Mul(A,b) .+ c
            @test all(c .=== BLAS.gemv!('N', 2one(T), A, b, one(T), copy(b)))

            c = copy(b)
            c .= 3one(T) .* Mul(A,b) .+ 2one(T) .* c
            @test all(c .=== BLAS.gemv!('N', 3one(T), A, b, 2one(T), copy(b)))

            d = similar(c)
            c = copy(b)
            d .= 3one(T) .* Mul(A,b) .+ 2one(T) .* c
            @test all(d .=== BLAS.gemv!('N', 3one(T), A, b, 2one(T), copy(b)))
        end
    end

    @testset "gemm" begin
        for A in (randn(5,5), view(randn(5,5),:,:), view(randn(5,5),1:5,:),
                  view(randn(5,5),1:5,1:5), view(randn(5,5),:,1:5)),
            B in (randn(5,5), view(randn(5,5),:,:), view(randn(5,5),1:5,:),
                      view(randn(5,5),1:5,1:5), view(randn(5,5),:,1:5))
            C = similar(B);

            C .= Mul(A,B)
            @test all(C .=== BLAS.gemm!('N', 'N', 1.0, A, B, 0.0, similar(C)))

            B .= Mul(A,B)
            @test all(C .=== B)

            C .= 2.0 .* Mul(A,B)
            @test all(C .=== BLAS.gemm!('N', 'N', 2.0, A, B, 0.0, similar(C)))

            C = copy(B)
            C .= Mul(A,B) .+ C
            @test all(C .=== BLAS.gemm!('N', 'N', 1.0, A, B, 1.0, copy(B)))

            C = copy(B)
            C .= Mul(A,B) .+ 2.0 .* C
            @test all(C .=== BLAS.gemm!('N', 'N', 1.0, A, B, 2.0, copy(B)))

            C = copy(B)
            C .= 2.0 .* Mul(A,B) .+ C
            @test all(C .=== BLAS.gemm!('N', 'N', 2.0, A, B, 1.0, copy(B)))

            C = copy(B)
            C .= 3.0 .* Mul(A,B) .+ 2.0 .* C
            @test all(C .=== BLAS.gemm!('N', 'N', 3.0, A, B, 2.0, copy(B)))

            d = similar(C)
            C = copy(B)
            d .= 3.0 .* Mul(A,B) .+ 2.0 .* C
            @test all(d .=== BLAS.gemm!('N', 'N', 3.0, A, B, 2.0, copy(B)))
        end
    end

    @testset "gemm Complex" begin
        for T in (ComplexF64,),
                A in (randn(T,5,5), view(randn(T,5,5),:,:)),
                B in (randn(T,5,5), view(randn(T,5,5),:,:))
            C = similar(B);

            C .= Mul(A,B)
            @test all(C .=== BLAS.gemm!('N', 'N', one(T), A, B, zero(T), similar(C)))

            B .= Mul(A,B)
            @test all(C .=== B)

            C .= 2one(T) .* Mul(A,B)
            @test all(C .=== BLAS.gemm!('N', 'N', 2one(T), A, B, zero(T), similar(C)))

            C = copy(B)
            C .= Mul(A,B) .+ C
            @test all(C .=== BLAS.gemm!('N', 'N', one(T), A, B, one(T), copy(B)))

            C = copy(B)
            C .= Mul(A,B) .+ 2one(T) .* C
            @test all(C .=== BLAS.gemm!('N', 'N', one(T), A, B, 2one(T), copy(B)))

            C = copy(B)
            C .= 2one(T) .* Mul(A,B) .+ C
            @test all(C .=== BLAS.gemm!('N', 'N', 2one(T), A, B, one(T), copy(B)))

            C = copy(B)
            C .= 3one(T) .* Mul(A,B) .+ 2one(T) .* C
            @test all(C .=== BLAS.gemm!('N', 'N', 3one(T), A, B, 2one(T), copy(B)))

            d = similar(C)
            C = copy(B)
            d .= 3one(T) .* Mul(A,B) .+ 2one(T) .* C
            @test all(d .=== BLAS.gemm!('N', 'N', 3one(T), A, B, 2one(T), copy(B)))
        end
    end

    @testset "gemv mixed array types" begin
        (A, B, C) = (randn(5,5), randn(5,5), reshape(1.0:25.0,5,5))
        D = similar(B)
        D .= Mul(A,B) .+ C
        @test all(D .=== BLAS.gemm!('N', 'N', 1.0, A, B, 1.0, Matrix{Float64}(C)))

        D .= Mul(A,B) .+ 2.0 .* C
        @test all(D .=== BLAS.gemm!('N', 'N', 1.0, A, B, 2.0, Matrix{Float64}(C)))

        D .= 2.0 .* Mul(A,B) .+ C
        @test all(D .=== BLAS.gemm!('N', 'N', 2.0, A, B, 1.0, Matrix{Float64}(C)))

        D .= 3.0 .* Mul(A,B) .+ 2.0 .* C
        @test all(D .=== BLAS.gemm!('N', 'N', 3.0, A, B, 2.0, Matrix{Float64}(C)))

        @test (similar(B) .= 1 .* Mul(A,B)) ≈
            (similar(B) .= 1 .* Mul(A,B) .+ 0 .* B) ≈
            A*B
    end

    @testset "gemv adjtrans" begin
        for A in (randn(5,5), view(randn(5,5),:,:), view(randn(5,5),1:5,:),
                  view(randn(5,5),1:5,1:5), view(randn(5,5),:,1:5)),
                    b in (randn(5), view(randn(5),:), view(randn(5),1:5), view(randn(9),1:2:9)),
                    Ac in (transpose(A), A')
            c = similar(b);

            c .= Mul(Ac,b)
            @test all(c .=== BLAS.gemv!('T', 1.0, A, b, 0.0, similar(c)))

            b .= Mul(Ac,b)
            @test all(c .=== b)

            c .= 2.0 .* Mul(Ac,b)
            @test all(c .=== BLAS.gemv!('T', 2.0, A, b, 0.0, similar(c)))

            c = copy(b)
            c .= Mul(Ac,b) .+ c
            @test all(c .=== BLAS.gemv!('T', 1.0, A, b, 1.0, copy(b)))

            c = copy(b)
            c .= Mul(Ac,b) .+ 2.0 .* c
            @test all(c .=== BLAS.gemv!('T', 1.0, A, b, 2.0, copy(b)))

            c = copy(b)
            c .= 2.0 .* Mul(Ac,b) .+ c
            @test all(c .=== BLAS.gemv!('T', 2.0, A, b, 1.0, copy(b)))

            c = copy(b)
            c .= 3.0 .* Mul(Ac,b) .+ 2.0 .* c
            @test all(c .=== BLAS.gemv!('T', 3.0, A, b, 2.0, copy(b)))

            d = similar(c)
            c = copy(b)
            d .= 3.0 .* Mul(Ac,b) .+ 2.0 .* c
            @test all(d .=== BLAS.gemv!('T', 3.0, A, b, 2.0, copy(b)))
        end
    end

    @testset "gemv adjtrans mixed types" begin
        (A, b, c) = (randn(5,5), randn(5), 1.0:5.0)
        for Ac in (transpose(A), A')
            d = similar(b)
            d .= Mul(Ac,b) .+ c
            @test all(d .=== BLAS.gemv!('T', 1.0, A, b, 1.0, Vector{Float64}(c)))

            d .= Mul(Ac,b) .+ 2.0 .* c
            @test all(d .=== BLAS.gemv!('T', 1.0, A, b, 2.0, Vector{Float64}(c)))

            d .= 2.0 .* Mul(Ac,b) .+ c
            @test all(d .=== BLAS.gemv!('T', 2.0, A, b, 1.0, Vector{Float64}(c)))

            d .= 3.0 .* Mul(Ac,b) .+ 2.0 .* c
            @test all(d .=== BLAS.gemv!('T', 3.0, A, b, 2.0, Vector{Float64}(c)))
        end
    end

    @testset "gemv adjtrans Complex" begin
        for T in (ComplexF64,),
                A in (randn(T,5,5), view(randn(T,5,5),:,:)),
                b in (randn(T,5), view(randn(T,5),:)),
                (Ac,trans) in ((transpose(A),'T'), (A','C'))
            c = similar(b);

            c .= Mul(Ac,b)
            @test all(c .=== BLAS.gemv!(trans, one(T), A, b, zero(T), similar(c)))

            b .= Mul(Ac,b)
            @test all(c .=== b)

            c .= 2one(T) .* Mul(Ac,b)
            @test all(c .=== BLAS.gemv!(trans, 2one(T), A, b, zero(T), similar(c)))

            c = copy(b)
            c .= Mul(Ac,b) .+ c
            @test all(c .=== BLAS.gemv!(trans, one(T), A, b, one(T), copy(b)))

            c = copy(b)
            c .= Mul(Ac,b) .+ 2one(T) .* c
            @test all(c .=== BLAS.gemv!(trans, one(T), A, b, 2one(T), copy(b)))

            c = copy(b)
            c .= 2one(T) .* Mul(Ac,b) .+ c
            @test all(c .=== BLAS.gemv!(trans, 2one(T), A, b, one(T), copy(b)))

            c = copy(b)
            c .= 3one(T) .* Mul(Ac,b) .+ 2one(T) .* c
            @test all(c .=== BLAS.gemv!(trans, 3one(T), A, b, 2one(T), copy(b)))

            d = similar(c)
            c = copy(b)
            d .= 3one(T) .* Mul(Ac,b) .+ 2one(T) .* c
            @test all(d .=== BLAS.gemv!(trans, 3one(T), A, b, 2one(T), copy(b)))
        end
    end

    @testset "gemm adjtrans" begin
        for A in (randn(5,5), view(randn(5,5),:,:)),
            B in (randn(5,5), view(randn(5,5),1:5,:))
            for Ac in (transpose(A), A')
                C = similar(B)
                C .= Mul(Ac,B)
                @test all(C .=== BLAS.gemm!('T', 'N', 1.0, A, B, 0.0, similar(C)))

                B .= Mul(Ac,B)
                @test all(C .=== B)

                C .= 2.0 .* Mul(Ac,B)
                @test all(C .=== BLAS.gemm!('T', 'N', 2.0, A, B, 0.0, similar(C)))

                C = copy(B)
                C .= Mul(Ac,B) .+ C
                @test all(C .=== BLAS.gemm!('T', 'N', 1.0, A, B, 1.0, copy(B)))

                C = copy(B)
                C .= Mul(Ac,B) .+ 2.0 .* C
                @test all(C .=== BLAS.gemm!('T', 'N', 1.0, A, B, 2.0, copy(B)))

                C = copy(B)
                C .= 2.0 .* Mul(Ac,B) .+ C
                @test all(C .=== BLAS.gemm!('T', 'N', 2.0, A, B, 1.0, copy(B)))

                C = copy(B)
                C .= 3.0 .* Mul(Ac,B) .+ 2.0 .* C
                @test all(C .=== BLAS.gemm!('T', 'N', 3.0, A, B, 2.0, copy(B)))

                d = similar(C)
                C = copy(B)
                d .= 3.0 .* Mul(Ac,B) .+ 2.0 .* C
                @test all(d .=== BLAS.gemm!('T', 'N', 3.0, A, B, 2.0, copy(B)))
            end
            for Bc in (transpose(B), B')
                C = similar(B)
                C .= Mul(A,Bc)
                @test all(C .=== BLAS.gemm!('N', 'T', 1.0, A, B, 0.0, similar(C)))

                B .= Mul(A,Bc)
                @test all(C .=== B)

                C .= 2.0 .* Mul(A,Bc)
                @test all(C .=== BLAS.gemm!('N', 'T', 2.0, A, B, 0.0, similar(C)))

                C = copy(B)
                C .= Mul(A,Bc) .+ C
                @test all(C .=== BLAS.gemm!('N', 'T', 1.0, A, B, 1.0, copy(B)))

                C = copy(B)
                C .= Mul(A,Bc) .+ 2.0 .* C
                @test all(C .=== BLAS.gemm!('N', 'T', 1.0, A, B, 2.0, copy(B)))

                C = copy(B)
                C .= 2.0 .* Mul(A,Bc) .+ C
                @test all(C .=== BLAS.gemm!('N', 'T', 2.0, A, B, 1.0, copy(B)))

                C = copy(B)
                C .= 3.0 .* Mul(A,Bc) .+ 2.0 .* C
                @test all(C .=== BLAS.gemm!('N', 'T', 3.0, A, B, 2.0, copy(B)))

                d = similar(C)
                C = copy(B)
                d .= 3.0 .* Mul(A,Bc) .+ 2.0 .* C
                @test all(d .=== BLAS.gemm!('N', 'T', 3.0, A, B, 2.0, copy(B)))
            end
            for Ac in (transpose(A), A'), Bc in (transpose(B), B')
                C = similar(B)
                C .= Mul(Ac,Bc)
                @test all(C .=== BLAS.gemm!('T', 'T', 1.0, A, B, 0.0, similar(C)))

                B .= Mul(Ac,Bc)
                @test all(C .=== B)

                C .= 2.0 .* Mul(Ac,Bc)
                @test all(C .=== BLAS.gemm!('T', 'T', 2.0, A, B, 0.0, similar(C)))

                C = copy(B)
                C .= Mul(Ac,Bc) .+ C
                @test all(C .=== BLAS.gemm!('T', 'T', 1.0, A, B, 1.0, copy(B)))

                C = copy(B)
                C .= Mul(Ac,Bc) .+ 2.0 .* C
                @test all(C .=== BLAS.gemm!('T', 'T', 1.0, A, B, 2.0, copy(B)))

                C = copy(B)
                C .= 2.0 .* Mul(Ac,Bc) .+ C
                @test all(C .=== BLAS.gemm!('T', 'T', 2.0, A, B, 1.0, copy(B)))

                C = copy(B)
                C .= 3.0 .* Mul(Ac,Bc) .+ 2.0 .* C
                @test all(C .=== BLAS.gemm!('T', 'T', 3.0, A, B, 2.0, copy(B)))

                d = similar(C)
                C = copy(B)
                d .= 3.0 .* Mul(Ac,Bc) .+ 2.0 .* C
                @test all(d .=== BLAS.gemm!('T', 'T', 3.0, A, B, 2.0, copy(B)))
            end
        end
    end

    @testset "symv adjtrans" begin
        A = randn(100,100)
        x = randn(100)

        @test all( (similar(x) .= Mul(Symmetric(A),x)) .===
                    (similar(x) .= Mul(Hermitian(A),x)) .===
                    (similar(x) .= Mul(Symmetric(A)',x)) .===
                    (similar(x) .= Mul(Symmetric(view(A,:,:)',:L),x)) .===
                    (similar(x) .= 1.0.*Mul(Symmetric(A),x) .+ 0.0.*similar(x)) .===
                    (similar(x) .= Mul(view(Symmetric(A),:,:),x)) .===
                    BLAS.symv!('U', 1.0, A, x, 0.0, similar(x)) )

        y = copy(x)
        y .= Mul(Symmetric(A), y)
        @test all( (similar(x) .= Mul(Symmetric(A),x)) .=== y)

        @test all( (similar(x) .= Mul(Symmetric(A,:L),x)) .===
                    (similar(x) .= Mul(Hermitian(A,:L),x)) .===
                    (similar(x) .= Mul(Symmetric(A,:L)',x)) .===
                    (similar(x) .= Mul(Symmetric(view(A,:,:)',:U),x)) .===
                    (similar(x) .= 1.0.*Mul(Symmetric(A,:L),x) .+ 0.0.*similar(x)) .===
                    (similar(x) .= Mul(view(Hermitian(A,:L),:,:),x)) .===
                    BLAS.symv!('L', 1.0, A, x, 0.0, similar(x)) )
        T = ComplexF64
        A = randn(T,100,100)
        x = randn(T,100)
        @test all( (similar(x) .= Mul(Symmetric(A),x)) .===
                    (similar(x) .= Mul(transpose(Symmetric(A)),x)) .===
                    (similar(x) .= Mul(Symmetric(transpose(view(A,:,:)),:L),x)) .===
                    (similar(x) .= one(T).*Mul(Symmetric(A),x) .+ zero(T).*similar(x)) .===
                    (similar(x) .= Mul(view(Symmetric(A),:,:),x)) .===
                    BLAS.symv!('U', one(T), A, x, zero(T), similar(x)) )

        @test all( (similar(x) .= Mul(Symmetric(A,:L),x)) .===
                    (similar(x) .= Mul(transpose(Symmetric(A,:L)),x)) .===
                    (similar(x) .= Mul(Symmetric(transpose(view(A,:,:)),:U),x)) .===
                    (similar(x) .= one(T).*Mul(Symmetric(A,:L),x) .+ zero(T).*similar(x)) .===
                    (similar(x) .= Mul(view(Symmetric(A,:L),:,:),x)) .===
                    BLAS.symv!('L', one(T), A, x, zero(T), similar(x)) )

        @test all( (similar(x) .= Mul(Hermitian(A),x)) .===
                    (similar(x) .= Mul(Hermitian(A)',x)) .===
                    (similar(x) .= one(T).*Mul(Hermitian(A),x) .+ zero(T).*similar(x)) .===
                    (similar(x) .= Mul(view(Hermitian(A),:,:),x)) .===
                    BLAS.hemv!('U', one(T), A, x, zero(T), similar(x)) )

        @test all( (similar(x) .= Mul(Hermitian(A,:L),x)) .===
                    (similar(x) .= Mul(Hermitian(A,:L)',x)) .===
                    (similar(x) .= one(T).*Mul(Hermitian(A,:L),x) .+ zero(T).*similar(x)) .===
                    (similar(x) .= Mul(view(Hermitian(A,:L),:,:),x)) .===
                    BLAS.hemv!('L', one(T), A, x, zero(T), similar(x)) )

        y = copy(x)
        y .= Mul(Hermitian(A), y)
        @test all( (similar(x) .= Mul(Hermitian(A),x)) .=== y)
    end

    @testset "tri" begin
        @testset "Float * Float vector" begin
            A = randn(Float64, 100, 100)
            x = randn(Float64, 100)

            @test all((y = copy(x); y .= Mul(UpperTriangular(A),y) ) .===
                        (similar(x) .= Mul(UpperTriangular(A),x)) .===
                        BLAS.trmv!('U', 'N', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(UnitUpperTriangular(A),y) ) .===
                        (similar(x) .= Mul(UnitUpperTriangular(A),x)) .===
                        BLAS.trmv!('U', 'N', 'U', A, copy(x)))
            @test all((y = copy(x); y .= Mul(LowerTriangular(A),y) ) .===
                        (similar(x) .= Mul(LowerTriangular(A),x)) .===
                        BLAS.trmv!('L', 'N', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(UnitLowerTriangular(A),y) ) .===
                        (similar(x) .= Mul(UnitLowerTriangular(A),x)) .===
                        BLAS.trmv!('L', 'N', 'U', A, copy(x)))

            @test all((y = copy(x); y .= Mul(UpperTriangular(A)',y) ) .===
                        (similar(x) .= Mul(UpperTriangular(A)',x)) .===
                        (similar(x) .= Mul(LowerTriangular(A'),x)) .===
                        BLAS.trmv!('L', 'T', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(UnitUpperTriangular(A)',y) ) .===
                        (similar(x) .= Mul(UnitUpperTriangular(A)',x)) .===
                        (similar(x) .= Mul(UnitLowerTriangular(A'),x)) .===
                        BLAS.trmv!('L', 'T', 'U', A, copy(x)))
            @test all((y = copy(x); y .= Mul(LowerTriangular(A)',y) ) .===
                        (similar(x) .= Mul(LowerTriangular(A)',x)) .===
                        (similar(x) .= Mul(UpperTriangular(A'),x)) .===
                        BLAS.trmv!('U', 'T', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(UnitLowerTriangular(A)',y) ) .===
                        (similar(x) .= Mul(UnitLowerTriangular(A)',x)) .===
                        (similar(x) .= Mul(UnitUpperTriangular(A'),x)) .===
                        BLAS.trmv!('U', 'T', 'U', A, copy(x)))
        end

        @testset "Float * Complex vector"  begin
            T = ComplexF64
            A = randn(T, 100, 100)
            x = randn(T, 100)

            @test all((y = copy(x); y .= Mul(UpperTriangular(A),y) ) .===
                        (similar(x) .= Mul(UpperTriangular(A),x)) .===
                        BLAS.trmv!('U', 'N', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(UnitUpperTriangular(A),y) ) .===
                        (similar(x) .= Mul(UnitUpperTriangular(A),x)) .===
                        BLAS.trmv!('U', 'N', 'U', A, copy(x)))
            @test all((y = copy(x); y .= Mul(LowerTriangular(A),y) ) .===
                        (similar(x) .= Mul(LowerTriangular(A),x)) .===
                        BLAS.trmv!('L', 'N', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(UnitLowerTriangular(A),y) ) .===
                        (similar(x) .= Mul(UnitLowerTriangular(A),x)) .===
                        BLAS.trmv!('L', 'N', 'U', A, copy(x)))
            LowerTriangular(A')  == UpperTriangular(A)'

            @test all((y = copy(x); y .= Mul(transpose(UpperTriangular(A)),y) ) .===
                        (similar(x) .= Mul(transpose(UpperTriangular(A)),x)) .===
                        (similar(x) .= Mul(LowerTriangular(transpose(A)),x)) .===
                        BLAS.trmv!('L', 'T', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(transpose(UnitUpperTriangular(A)),y) ) .===
                        (similar(x) .= Mul(transpose(UnitUpperTriangular(A)),x)) .===
                        (similar(x) .= Mul(UnitLowerTriangular(transpose(A)),x)) .===
                        BLAS.trmv!('L', 'T', 'U', A, copy(x)))
            @test all((y = copy(x); y .= Mul(transpose(LowerTriangular(A)),y) ) .===
                        (similar(x) .= Mul(transpose(LowerTriangular(A)),x)) .===
                        (similar(x) .= Mul(UpperTriangular(transpose(A)),x)) .===
                        BLAS.trmv!('U', 'T', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(transpose(UnitLowerTriangular(A)),y) ) .===
                        (similar(x) .= Mul(transpose(UnitLowerTriangular(A)),x)) .===
                        (similar(x) .= Mul(UnitUpperTriangular(transpose(A)),x)) .===
                        BLAS.trmv!('U', 'T', 'U', A, copy(x)))

            @test all((y = copy(x); y .= Mul(UpperTriangular(A)',y) ) .===
                        (similar(x) .= Mul(UpperTriangular(A)',x)) .===
                        (similar(x) .= Mul(LowerTriangular(A'),x)) .===
                        BLAS.trmv!('L', 'C', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(UnitUpperTriangular(A)',y) ) .===
                        (similar(x) .= Mul(UnitUpperTriangular(A)',x)) .===
                        (similar(x) .= Mul(UnitLowerTriangular(A'),x)) .===
                        BLAS.trmv!('L', 'C', 'U', A, copy(x)))
            @test all((y = copy(x); y .= Mul(LowerTriangular(A)',y) ) .===
                        (similar(x) .= Mul(LowerTriangular(A)',x)) .===
                        (similar(x) .= Mul(UpperTriangular(A'),x)) .===
                        BLAS.trmv!('U', 'C', 'N', A, copy(x)))
            @test all((y = copy(x); y .= Mul(UnitLowerTriangular(A)',y) ) .===
                        (similar(x) .= Mul(UnitLowerTriangular(A)',x)) .===
                        (similar(x) .= Mul(UnitUpperTriangular(A'),x)) .===
                        BLAS.trmv!('U', 'C', 'U', A, copy(x)))
        end

        @testset "Float * Float Matrix" begin
            A = randn(Float64, 100, 100)
            x = randn(Float64, 100, 100)

            @test UpperTriangular(A)*x ≈ (similar(x) .= Mul(UpperTriangular(A), x))
        end
    end

    @testset "Mixed types" begin
        A = randn(5,6)
        b = rand(Int,6)
        c = Array{Float64}(undef, 5)
        c .= Mul(A,b)

        @test_throws DimensionMismatch (similar(c,3) .= Mul(A,b))
        @test_throws DimensionMismatch (c .= Mul(A,similar(b,2)))

        d = similar(c)
        mul!(d, A, b)
        @test all(c .=== d)

        copyto!(d, MulAdd(1, A, b, 0.0, d))
        @test d == copyto!(similar(d), MulAdd(1, A, b, 0.0, d)) ≈ A*b
        @test copyto!(similar(d), MulAdd(1, A, b, 1.0, d)) ≈ A*b + d

        @test all((similar(d) .= MulAdd(1, A, b, 1.0, d)) .=== copyto!(similar(d), MulAdd(1, A, b, 1.0, d)))

        B = rand(Int,6,4)
        C = Array{Float64}(undef, 5, 4)
        C .= Mul(A,B)

        @test_throws DimensionMismatch materialize!(MulAdd(1,A,B,0,similar(C,4,4)))
        @test_throws DimensionMismatch (similar(C,4,4) .= Mul(A,B))
        @test_throws DimensionMismatch (C .= Mul(A,similar(B,2,2)))

        D = similar(C)
        mul!(D, A, B)
        @test all(C .=== D)

        A = randn(Float64,20,22)
        B = randn(ComplexF64,22,24)
        C = similar(B,20,24)
        @test all((C .= Mul(A,B)  ) .=== copyto!(similar(C), MulAdd(1.0, A, B, 0.0, C)) .=== A*B)
    end

    @testset "no allocation" begin
        function blasnoalloc(c, α, A, x, β, y)
            c .= Mul(A,x)
            c .= α .* Mul(A,x)
            c .= Mul(A,x) .+ y
            c .= α .* Mul(A,x) .+ y
            c .= Mul(A,x) .+ β .* y
            c .= α .* Mul(A,x) .+ β .* y
        end

        A = randn(5,5); x = randn(5); y = randn(5); c = similar(y);
        blasnoalloc(c, 2.0, A, x, 3.0, y)
        @test @allocated(blasnoalloc(c, 2.0, A, x, 3.0, y)) == 0
        Ac = A'
        blasnoalloc(c, 2.0, Ac, x, 3.0, y)
        @test @allocated(blasnoalloc(c, 2.0, Ac, x, 3.0, y)) == 0
    end

    @testset "multi-argument mul" begin
        A = randn(5,5)
        B = materialize(Mul(A,A,A))
        @test B isa Matrix{Float64}
        @test all(B .=== A*(A*A))

        @test Mul(A,A) * A ≈ A * Mul(A,A) ≈ Mul(A) * Mul(A,A) ≈ A^3
        @test Mul(A,A) * Mul(A,A) ≈ Mul(A) * Mul(A,A,A) ≈ A^4
    end

    @testset "Diagonal and SymTridiagonal" begin
        A = randn(5,5)
        B = Diagonal(randn(5))
        @test MemoryLayout(B) == DiagonalLayout(DenseColumnMajor())
        @test materialize(Mul(A,B)) == A*B

        A = randn(5,5)
        B = SymTridiagonal(randn(5),randn(4))
        @test MemoryLayout(B) == SymTridiagonalLayout(DenseColumnMajor())
        @test materialize(Mul(A,B)) == A*B
    end

    @testset "MulArray" begin
        A = randn(5,5)
        M = MulArray(A,A)
        @test Matrix(M) ≈ A^2
    end
end


struct FourVector{T} <: AbstractVector{Int}
    fv::MVector{4,T} # px,py,pz,E
end
FourVector(x,y,z; t) = FourVector(MVector(x,y,z,t))
Particle(;E,p) = FourVector(p...; t=E)
#
function Particle(px,py,pz; E=-Inf, msq = -Inf)
    (msq == -Inf) && (E == -Inf) && error("give either msq of E")
    (E != -Inf) && return FourVector(px,py,pz,t=E)
    return FourVector(px,py,pz;t=sqrt(px^2+py^2+pz^2+msq))
end

import Base:size
size(p::FourVector{T} where T) = (4,)

import Base: getindex, setindex!
getindex(p::FourVector{T} where T, i::Int) = (i==0 ? getindex(p.fv,4) : getindex(p.fv,i))
getindex(p::FourVector{T} where T, I::Vararg) = getindex(p.fv, I...)
# setindex!(p::FourVector{T} where T, v, i::Int)	= (i==0 ? setindex!(p.fv,v,4) : setindex!(p.fv, v, i))
setindex!(p::FourVector{T} where T, v, I::Vararg{Int, 1}) = setindex!(p.fv, v, I...)

import Base:+,-,*
+(p1::FourVector, p2::FourVector) = FourVector(p1.fv+p2.fv)
*(p::FourVector, α::T where T<:Number) = FourVector(p.fv*α)
*(α::T where T<:Number, p::FourVector) = FourVector(p.fv*α)
-(p1::FourVector, p2::FourVector) = FourVector(p1.fv-p2.fv)

import Base:copy
copy(p::FourVector) = FourVector(copy(p.fv))

import Base:getproperty
function getproperty(p::FourVector, sym::Symbol)
    (sym in fieldnames(FourVector)) && return getfield(p, sym)
    (sym == :Px || sym == :px || sym == :x || sym == :X) && return getfield(p, :fv)[1]
    (sym == :Py || sym == :py || sym == :y || sym == :Y) && return getfield(p, :fv)[2]
    (sym == :Pz || sym == :pz || sym == :z || sym == :Z) && return getfield(p, :fv)[3]
    (sym == :E  || sym == :p0 || sym == :t || sym == :T) && return getfield(p, :fv)[4]
    (sym == :p  || sym == :P) && return @view getfield(p, :fv)[1:3]
    error("no property $(sym)")
end

import LinearAlgebra: dot
dot(p1::FourVector, p2::FourVector) = p1.E*p2.E - dot(p1.p,p2.p)

# properties
psq(p::FourVector) = sum(abs2,p.p)
invmasssq(p::FourVector) = p⋅p
mass(p::FourVector) = sqrt(invmasssq(p))

sphericalangles(p::FourVector) = p.p[3]/norm(p.p), atan(p.p[2],p.p[1])
boostfactor(p::FourVector) = p.E/mass(p)


using GPFlux
using LinearAlgebra
using GaussianProcesses
using Statistics
using Calculus
using Zygote
using Test, Random

function K_check(mykernel, gpkernel, X)
  myK = mykernel(X)
	gpK = cov(gpkernel, X) + Diagonal(1e-6*ones(size(X, 2)))
	@show isposdef(myK)
	@show isposdef(gpK)
	isposdef(myK) == isposdef(gpK) &&  isapprox(myK, gpK, rtol=1e-3, atol=1e-8)
end

function dKdθ_check(mykernel, gpkernel, X, v)
  N = size(X, 2)

  myps = GPFlux.params(mykernel)
  f = () -> v'*mykernel(X)*v
  mygs = gradient(f, myps)
  mygv = [mygs.grads[p][1] for p in myps]

  n_gpps = GaussianProcesses.num_params(gpkernel)
  gpps = GaussianProcesses.get_params(gpkernel)
  gpdKdθ = zeros(N, N, n_gpps)
  gpmetric = GaussianProcesses.KernelData(gpkernel, X, X)
  GaussianProcesses.grad_stack!(gpdKdθ, gpkernel, X, X, gpmetric)
  gpgv = [tr(v*v'*gpdKdθ[:, :, i]) for i in 1:n_gpps]

	# gpgv = [tr(v*v'*gpdKdθ[:, :, i]).*exp(-gpps[i]) for i in 1:n_gpps]

#  calps = [p[1] for p in myps]
#  calgv = Calculus.gradient(calps) do ps
#      ps_ = [[p] for p in ps]
#      kernel = typeof(mykernel)(ps_...)
#      v'*kernel(X)*v
#  end

	@show sort(mygv)
# 	@show sort(calgv)
	@show sort(gpgv)
#  isapprox(sort(mygv),sort(calgv),rtol=1e-3,atol=1e-8) &&
	isapprox(sort(mygv),sort(gpgv),rtol=1e-3,atol=1e-8)
end

@testset "kernel check" begin
	rng = MersenneTwister(4)
	d = 4
	n = 10
	X = rand(rng, d, n)
	v = rand(rng, n)
	kernel_pairs = [(IsoGaussKernel([4.0], [1.0]), SE(4.0, 1.0)),
									(IsoPeriodKernel([0.0], [0.0], [1.0]), Periodic(0.0, 1.0, 0.0)),
									(IsoRQKernel([0.0], [-1.0], [0.0]), RQ(0.0, 0.0, -1.0)),
									(IsoLinearKernel([0.0]), Lin(0.0)),
								  (SE_add_PeriodKernel(IsoGaussKernel([4.0],[0.0]),IsoPeriodKernel([0.0],[0.0],[1.0])), SE(4.0,0.0)+Periodic(0.0,1.0,0.0)),
									(SE_mul_PeriodKernel(IsoGaussKernel([4.0],[0.0]),IsoPeriodKernel([0.0],[0.0],[1.0])), SE(4.0,0.0)*Periodic(0.0,1.0,0.0)),
									]

	for (mykernel, gpkernel) in kernel_pairs
		@show (mykernel, gpkernel)
		@test K_check(mykernel, gpkernel, X)
		@test dKdθ_check(mykernel, gpkernel, X, v)
	end

end


# This shows how to solve the 2D incompressible Stokes equations using SNES solvers,
# using a staggered grid discretization and a Velocity-Pressure formulation
#
#   Governing equations:
#
#              dVx/dx + dVz/dz = 0                              |   Mass balance
#   -dP/dx + dTxx/dx + dTxz/dz = 0                              |   Horizontal force balance
#   -dP/dz + dTxz/dx + dTzz/dz = rho*g                          |   Vertical force balance
#
#   with:
#       Exx = dVx/dx, Ezz = dVz/dz, Exz = 0.5*(dVx/dz + dVz/dx) |   Strain rate definition
#
#      | Txx |   | 2*eta   0     0   |  | Exx |
#      | Tzz | = |   0   2*eta   0   |  | Ezz |                 |   Linear viscous (isotropic) rheology
#      | Txz |   |   0     0   2*eta |  | Exz |
#
# This example also uses the Automatic differentiation package ForwardDiff

using PETSc, MPI
using ForwardDiff, SparseArrays, SparseDiffTools

petsclib = PETSc.petsclibs[1];
PETSc.initialize(petsclib)

function FormRes!(ptr_fx_g, ptr_x_g, user_ctx)

    # Copy global to local vectors
    PETSc.update!(user_ctx.x_l, ptr_x_g,   PETSc.INSERT_VALUES)
    PETSc.update!(user_ctx.f_l, ptr_fx_g,  PETSc.INSERT_VALUES)

    # Retrieve arrays from the local vectors
    ArrayLocal_x     =   PETSc.DMStagVecGetArrayRead(user_ctx.dm,   user_ctx.x_l);      # array with all local x-data (solution array)
    ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm,       user_ctx.f_l);      # array with all local residual

    # Compute local residual
    ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx)

    # Finalize local arrays
    Base.finalize(ArrayLocal_x)
    Base.finalize(ArrayLocal_f)

    # Copy local into global residual vector
    PETSc.update!( ptr_fx_g, user_ctx.f_l, PETSc.INSERT_VALUES)

end

function  ForwardDiff_res(x, user_ctx)
    f   = zero(x)               # vector of zeros, of same type as x (local vector)

    ArrayLocal_x     =   PETSc.DMStagVecGetArray(user_ctx.dm, x);        # array with all local x-data
    ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm, f);        # array with all local residual

    ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx);

    # As the residual vector f is linked with ArrayLocal_f, we don't need to
    # pass ArrayLocal_f back to f

    return f;
end

function ComputeSparsityPatternJacobian_automatic(x_l, user_ctx)
    # This computes the sparsity pattern and coloring of the jacobian automatically
    # This will works for any equation but is slow @ high resolutions

    f_Residual  =   (x -> ForwardDiff_res(x, user_ctx));        # pass additional arguments into the routine
    J_julia     =   ForwardDiff.jacobian(f_Residual,x_l*0 .+ 1);# Compute a "blank" jacobian

    # employ sparse structure to compute jacobian - to be moved inside routines
    jac         =   sparse(J_julia);
    colors      =   matrix_colors(jac)          # the number of nonzeros per row

    return jac, colors
end

function  func(out, x, user_ctx)

    ArrayLocal_x     =   PETSc.DMStagVecGetArray(user_ctx.dm, x);          # array with all local x-data (solition array)
    ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm, out);        # array with all local residual

    ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx);

    return nothing
end

function FormJacobian!(ptr_x_g, J, P, user_ctx)
    # This requires several steps:
    #
    #   1) Extract local vector from global solution (x) vector
    #   2) Compute local jacobian from the residual routine (note that
    #       this routine requires julia vectors as input)
    #   3) Store the matrix inside the global Petsc matrix J

    # Extract the local vector
    #PETSc.DMGlobalToLocal(user_ctx.dm, cx_g,  PETSc.INSERT_VALUES,  user_ctx.x_l)
    PETSc.update!(user_ctx.x_l, ptr_x_g, PETSc.INSERT_VALUES)
    x               =   PETSc.unsafe_localarray(Float64, user_ctx.x_l.ptr;  write=false, read=true);

    # Check the sparsity pattern
    if isnothing(user_ctx.jac)
        error("You first have to define the sparsity pattern of the jacobian")
    else
        jac     =   user_ctx.jac;
        colors  =   user_ctx.colors;
    end
    out         =   similar(x);

    f_Res           =   ((out,x)->func(out, x, user_ctx));  # pass additional arguments into the routine

    if isnothing(user_ctx.jac_cache)
        # Allocate data required for jacobian computations
        user_ctx.jac_cache = ForwardColorJacCache(f_Res,x; colorvec=colors, sparsity = jac);
    end

    # Compute Jacobian using automatic differenciation
    forwarddiff_color_jacobian!(jac, f_Res, x, user_ctx.jac_cache);

    # Store global Jacobian inside Petsc Matrix
    ind             =   PETSc.LocalInGlobalIndices(user_ctx.dm);
    if PETSc.assembled(J) == false
        println("Assembling J")
        J           =   PETSc.MatSeqAIJ(sparse(jac[ind,ind]));
    else
        J           .=   sparse(jac[ind,ind]);
    end

    # Store Julia matrix and coloring
    user_ctx.jac    =   jac;
    user_ctx.colors =   colors;

   return jac, ind
end

# Define a struct that holds data we need in the local SNES routines below
mutable struct Data_Stokes2D
    # DMs and vectors
    dm
    dmCoeff
    coeff_l
    x_l
    f_l
    # physical parameters
    eta1
    eta2
    rho1
    rho2
    gz
    kappa
    # dimensions
    dx
    dz
    xlim
    zlim
    # jacobian and sparsity pattern
    jac
    jac_cache
    colors
end
user_ctx = Data_Stokes2D(nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing);  # holds data we need in the local

function ComputeLocalResidual(dm, ArrayLocal_x, ArrayLocal_f, user_ctx)
    # Compute the local residual. The vectors include ghost points

    # Compute shear stresses
    Txx,Tzz,Txz = ComputeStresses!(user_ctx, ArrayLocal_x);

    # Extracting arrays from the residual and solution vectors
    P       = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_ELEMENT, 0);
    Vx      = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_LEFT, 0);
    Vz      = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_DOWN, 0);

    f_p     = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_f,     PETSc.DMSTAG_ELEMENT, 0);
    f_x     = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_f,     PETSc.DMSTAG_LEFT, 0);
    f_z     = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_f,     PETSc.DMSTAG_DOWN, 0);

    # Extracting physical parameters
    RhoC    = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dmCoeff,user_ctx.coeff_l,     PETSc.DMSTAG_DOWN_LEFT, 1);

    # Gettins sizes
    dx = user_ctx.dx;
    dz = user_ctx.dz;

    # Getting global indices for the local vectors
    ind                 =   PETSc.DMStagGetIndices(dm);          # indices of (center/element) points, not including ghost values.
    #sx, sn  = PETSc.DMStagGetCentralNodes(user_ctx.dm); #indices of central points

    ix      =   ind.center[1];
    iz      =   ind.center[1];

    # Horizontal force balance f(x)

    # Boundary conditions (Dirichlet with no normal velocity at boundary)
    f_x[ix[1]    ,:] .= Vx[ix[1]    ,:] .- 0.0;
    f_x[ix[end]+1,:] .= Vx[ix[end]+1,:] .+ 0.0;

    # -dP/dx + dTxx/dx + dTxz/dz = 0
    f_x[ix[2]:ix[end],iz] .= .-(P[ix[2]:ix[end],iz] .- P[ix[1]:ix[end-1],iz])./dx .+
                               (Txx[2:end,:]        .- Txx[1:end-1,:])       ./dx .+
                               (Txz[2:end-1,2:end]  .- Txz[2:end-1,1:end-1]) ./dz;

    # Vertical force balance f(z)

    # Boundary conditions (Dirichlet with no normal velocity at boundary)
    f_z[:,iz[1]    ] .= Vz[:,iz[1]    ] .- 0.0;
    f_z[:,iz[end]+1] .= Vz[:,iz[end]+1] .+ 0.0;

    # -dP/dz + dTxz/dx + dTzz/dz = rho*g
    f_z[ix,iz[2]:iz[end]] .= .-(P[ix,iz[2]:iz[end]] .- P[ix,iz[1]:iz[end-1]])./dz .+
                               (Tzz[:,2:end]        .- Tzz[:,1:end-1])       ./dz .+
                               (Txz[ 2:end,2:end-1] .- Txz[1:end-1,2:end-1]) ./dx .+
                               (RhoC[ix.+1,iz[2]:iz[end]] .+ RhoC[ix,iz[2]:iz[end]  ]) .* 0.5 .* user_ctx.gz;

    # Mass balance f(p)
    # dVx/dx + dVz/dz = 0
    kappa = user_ctx.kappa; # penalty term to help solver
    f_p[ix,iz] .= (Vx[ix.+1,iz].-Vx[ix,iz])./dx .+ (Vz[ix,iz.+1].-Vz[ix,iz])./dz .+ 1/kappa*P[ix,iz];

    # Cleanup
    Base.finalize(Vx);    Base.finalize(Vz);
    Base.finalize(P);     Base.finalize(RhoC);
    Base.finalize(f_x);   Base.finalize(f_z);
    Base.finalize(Txx);   Base.finalize(Tzz);
    Base.finalize(Txz);
end

function ComputeStresses!(user_ctx, ArrayLocal_x)

    # Compute strain rates
    Exx,Ezz,Exz = ComputeStrainRates!(user_ctx, ArrayLocal_x);

    # Getting Eta at the center and corner of the cells
    EtaE    = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dmCoeff,user_ctx.coeff_l,     PETSc.DMSTAG_ELEMENT, 0);
    EtaC    = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dmCoeff,user_ctx.coeff_l,     PETSc.DMSTAG_DOWN_LEFT, 0);

    # Getting indices for center nodes (not ghost)
    ind     = PETSc.DMStagGetIndices(user_ctx.dm); #indices of central points
    #ix      =   sx[1]:sn[1];
    #iz      =   sx[2]:sn[2];
    ix  = ind.center[1];
    iz  = ind.center[1];

    # Compute shear stresses
    Txx    = 2 .* EtaE[ix,iz] .* Exx;
    Tzz    = 2 .* EtaE[ix,iz] .* Ezz;
    Txz    = 2 .* EtaC[ix[1]:ix[end]+1,iz[1]:iz[end]+1] .* Exz;

    return Txx,Tzz,Txz

    Base.finalize(Exx);   Base.finalize(Ezz);
    Base.finalize(Exz);

end

function ComputeStrainRates!(user_ctx, ArrayLocal_x)

    # Getting velocity vectors
    Vx      = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_LEFT, 0);
    Vz      = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_DOWN, 0);

    # Gettind dimensions
    dx = user_ctx.dx;
    dz = user_ctx.dz;

    # Getting indices for center nodes (not ghost)
    ind  = PETSc.DMStagGetIndices(user_ctx.dm); #indices of central points
    #ix      =   sx[1]:sn[1];
    #iz      =   sx[2]:sn[2];
    ix      =   ind.center[1];
    iz      =   ind.center[2];

    # Set ghost points (free slip)
    Vx[:,iz[1]-1]   .=  Vx[:,iz[1]];                # ghost points on left and right size (here: free slip)
    Vx[:,iz[end]+1] .=  Vx[:,iz[end]];
    Vz[ix[1]-1,:]   .=  Vz[ix[1],:];
    Vz[ix[end]+1,:] .=  Vz[ix[end],:];

    # Compute deviatoric Strain rates
    DivV  = (Vx[ix.+1,iz].-Vx[ix,iz])./dx .+ (Vz[ix,iz.+1].-Vz[ix,iz])./dz;
    Exx   = (Vx[ix.+1,iz].-Vx[ix,iz])./dx .- 1/3 .* DivV;
    Ezz   = (Vz[ix,iz.+1].-Vz[ix,iz])./dz .- 1/3 .* DivV;
    Exz   = 0.5.*((Vx[ix[1]:ix[end].+1,iz[1]:iz[end].+1].-Vx[ix[1]:ix[end].+1,iz[1].-1:iz[end]])./dz .+
                  (Vz[ix[1]:ix[end].+1,iz[1]:iz[end].+1].-Vz[ix[1].-1:ix[end],iz[1]:iz[end].+1])./dx);

    return Exx,Ezz,Exz

end

function PopulateCoefficientData!(user_ctx)

    # Create coefficient local vector and array (to store eta and rho)
    user_ctx.coeff_l    =   PETSc.createlocalvector(user_ctx.dmCoeff);
    coeff_array         =   PETSc.DMStagVecGetArray(user_ctx.dmCoeff,user_ctx.coeff_l);

    # Create Coordinates dm, local vector and array
    dm_coord  = PETSc.getcoordinateDM(user_ctx.dmCoeff);
    vec_coord = PETSc.getcoordinateslocal(user_ctx.dmCoeff);
    coord     = PETSc.DMStagVecGetArray(dm_coord, vec_coord);
   # start,n,nExtra = PETSc.DMStagGetCorners(user_ctx.dmCoeff);

    # Get the correct entries for each of our variables in local element-wise storage
    iec = PETSc.DMStagGetLocationSlot(user_ctx.dmCoeff, PETSc.DMSTAG_DOWN_LEFT, 0);  # location eta corner
    irc = PETSc.DMStagGetLocationSlot(user_ctx.dmCoeff, PETSc.DMSTAG_DOWN_LEFT, 1);  # location rho corner
    iee = PETSc.DMStagGetLocationSlot(user_ctx.dmCoeff, PETSc.DMSTAG_ELEMENT, 0);    # location eta element
    ixc = PETSc.DMStagGetLocationSlot(dm_coord, PETSc.DMSTAG_DOWN_LEFT, 0);          # location coord corner
    ixe = PETSc.DMStagGetLocationSlot(dm_coord, PETSc.DMSTAG_ELEMENT, 0);            # location coord element

    # Fill coefficient vectors with suitable value (determined by x-coordinate in function GetPhase())

    # Element nodes (eta)

    Coeff          = coeff_array[:,:,iee+1];
    X_coord        = coord[:,:,ixe+1];
    index1         = findall(x -> GetPhase(user_ctx,x,1),X_coord);
    Coeff[index1] .= user_ctx.eta1;

    index2         = findall(x -> GetPhase(user_ctx,x,2),X_coord);
    Coeff[index2] .= user_ctx.eta2;

    coeff_array[:,:,iee+1] .= Coeff;

    # Corner nodes (rho and eta)

    CoeffE         = coeff_array[:,:,iec+1];
    CoeffR         = coeff_array[:,:,irc+1];
    X_coord        = coord[:,:,ixc+1];
    index1         = findall(x -> GetPhase(user_ctx,x,1),X_coord);
    CoeffE[index1].= user_ctx.eta1;
    CoeffR[index1].= user_ctx.rho1;

    index2         = findall(x -> GetPhase(user_ctx,x,2),X_coord);
    CoeffE[index2].= user_ctx.eta2;
    CoeffR[index2].= user_ctx.rho2;

    coeff_array[:,:,iec+1] .= CoeffE;
    coeff_array[:,:,irc+1] .= CoeffR;

end

function GetPhase(ctx,x,n)
    # Divide domain in phases according to x-coordinate (similar to SolCx benchmark)
    if x < (ctx.xlim[2]-ctx.xlim[1])/2
        if n == 1 return true else return false end
    else
        if n == 1 return false else return true end
    end
end

function SetVecX!(user_ctx,x_g)
    # Set first guess for vector x to 0
    user_ctx.x_l .= 0;
    #PETSc.DMLocalToGlobal(user_ctx.dm,user_ctx.x_l, PETSc.INSERT_VALUES, x_g);
    PETSc.update!(x_g,user_ctx.x_l, PETSc.INSERT_VALUES);

end

# Main Solver
nx, nz           =   32, 32;                  # number of nodes is x and z direction
#nx, nz           =   8, 8;                  # number of nodes is x and z direction

user_ctx.xlim    =   [0,2];                   # x and z dimensions
user_ctx.zlim    =   [0,1];
xlim             =   user_ctx.xlim;
zlim             =   user_ctx.zlim;
user_ctx.dx      =   (xlim[2]-xlim[1])/nx;   # x-resolution
user_ctx.dz      =   (zlim[2]-zlim[1])/nz;   # z-resolution
user_ctx.eta1    =   1;                      # viscosity phase 1
user_ctx.eta2    =   2;                      # viscosity phase 2
user_ctx.rho1    =   3;                      # density phase 1
user_ctx.rho2    =   4;                      # density phase 2
user_ctx.gz      =   -1;                     # gravity magnitude
user_ctx.kappa   =   1e2;                    # incompressible penaly term

# Create Solution and coefficient DMs
user_ctx.dm           =   PETSc.DMStagCreate2d(petsclib,MPI.COMM_SELF,PETSc.DM_BOUNDARY_GHOSTED,PETSc.DM_BOUNDARY_GHOSTED,nx,nz,1,1,0,1,1,PETSc.DMSTAG_STENCIL_BOX,1);  # V edge, P Element
PETSc.DMStagSetUniformCoordinatesExplicit(user_ctx.dm, xlim[1], xlim[2], zlim[1], zlim[2]);            # set coordinates
user_ctx.dmCoeff      =   PETSc.DMStagCreateCompatibleDMStag(user_ctx.dm,2,0,1,0);   # rho and eta on VERTEX, eta on ELEMENT
PETSc.DMStagSetUniformCoordinatesExplicit(user_ctx.dmCoeff, xlim[1], xlim[2], zlim[1], zlim[2]);

# Populate phases
PopulateCoefficientData!(user_ctx);

# Create solution and residual vectors
x_g             =   PETSc.createglobalvector(user_ctx.dm);
f_g             =   PETSc.createglobalvector(user_ctx.dm);
user_ctx.x_l    =   PETSc.createlocalvector(user_ctx.dm);
user_ctx.f_l    =   PETSc.createlocalvector(user_ctx.dm);

# Create Petsc matrix for Jacobian
J               =   PETSc.creatematrix(user_ctx.dm);

# Set first guess values for solution vector
SetVecX!(user_ctx,x_g);

# Compute jacobian sparcity pattern
x_l               =   PETSc.unsafe_localarray(Float64, user_ctx.x_l.ptr;  write=false, read=true)

user_ctx.jac, user_ctx.colors = ComputeSparsityPatternJacobian_automatic(user_ctx.x_l, user_ctx);
ind             =   PETSc.LocalInGlobalIndices(user_ctx.dm);    # extract indices
J           =   PETSc.MatSeqAIJ(sparse(user_ctx.jac[ind,ind]));

#Setting up SNES
S = PETSc.SNES{Float64}(petsclib,MPI.COMM_SELF;
        snes_rtol=1e-12,
        snes_monitor=true,
        snes_max_it = 500,
        snes_monitor_true_residual=true,
        snes_converged_reason=true);
S.user_ctx  =       user_ctx;

PETSc.setfunction!(S, FormRes!, f_g)
PETSc.setjacobian!(S, FormJacobian!, J, J)

#Solve
PETSc.solve!(x_g, S);

# Copy solution to local vector
PETSc.update!(user_ctx.x_l, x_g,  PETSc.INSERT_VALUES);

# Extract solution
Vx  =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm, user_ctx.x_l,   PETSc.DMSTAG_LEFT,      0);
Vz  =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm, user_ctx.x_l,   PETSc.DMSTAG_DOWN,      0);
P   =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm, user_ctx.x_l,   PETSc.DMSTAG_ELEMENT,   0);
Eta =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dmCoeff, user_ctx.coeff_l,   PETSc.DMSTAG_ELEMENT,   0);

# Getting indices for center nodes (not ghost)
ind     = PETSc.DMStagGetIndices(user_ctx.dm);
ix      =   ind.center[1];
iz      =   ind.center[2];

# Getting rid of ghost points
Vx  =   Vx[ix[1]:ix[end]+1,iz];
Vz  =   Vz[ix,iz[1]:iz[end]+1];
P   =    P[ix,iz];
Eta =  Eta[ix,iz];

# Compute central average velocities
Vx_cen = (Vx[2:end,:]  + Vx[1:end-1,:])/2;
Vz_cen = (Vz[:,2:end]  + Vz[:,1:end-1])/2;

#get coordinates
dm_coord  = PETSc.getcoordinateDM(user_ctx.dmCoeff);
vec_coord = PETSc.getcoordinateslocal(user_ctx.dmCoeff);

XCoord_e = PETSc.DMStagGetGhostArrayLocationSlot(dm_coord, vec_coord,   PETSc.DMSTAG_ELEMENT  ,      0); # location coord corner
ZCoord_e = PETSc.DMStagGetGhostArrayLocationSlot(dm_coord, vec_coord,   PETSc.DMSTAG_ELEMENT  ,      1); # location coord corner
XCoord_c = PETSc.DMStagGetGhostArrayLocationSlot(dm_coord, vec_coord,   PETSc.DMSTAG_DOWN_LEFT,      0);   # location coord element
ZCoord_c = PETSc.DMStagGetGhostArrayLocationSlot(dm_coord, vec_coord,   PETSc.DMSTAG_DOWN_LEFT,      1);   # location coord element

XCoord_e = XCoord_e[ix,iz];
ZCoord_e = ZCoord_e[ix,iz];
XCoord_c = XCoord_c[ix[1]:ix[end]+1,iz[1]:iz[end]+1];
ZCoord_c = ZCoord_c[ix[1]:ix[end]+1,iz[1]:iz[end]+1];

xe_1D    = XCoord_e[:,1];
ze_1D    = ZCoord_e[1,:];
xc_1D    = XCoord_c[:,1];
zc_1D    = ZCoord_c[1,:];

# Plot
CreatePlots = false;        # set to false by default such that testing does not require Plots
if CreatePlots==true
    using Plots

    p1 = heatmap(xe_1D,ze_1D, P', xlabel="Width", ylabel="Depth", title="Pressure",aspect_ratio = 1)
    p2 = heatmap(xe_1D,zc_1D, Vz', xlabel="Width", ylabel="Depth", title="Vz",aspect_ratio = 1)
    p3 = heatmap(xc_1D,ze_1D, Vx', xlabel="Width", ylabel="Depth", title="Vx",aspect_ratio = 1)
    quiver!(XCoord_e[:],ZCoord_e[:],quiver=(Vx_cen[:]*0.02,Vz_cen[:]*0.02), color=:white,aspect_ratio = 1)
    p4 = heatmap(xe_1D,ze_1D, Eta', xlabel="Width", ylabel="Depth", title="Eta",aspect_ratio = 1)
    plot(p1, p2, p3, p4, layout = (2, 2), legend = false)

end

# check if the solution makes sense
using LinearAlgebra, Test
check =  norm(user_ctx.x_l)
@test check ≈ 31.54688642064483


mutable struct KernelMachineRegression{M, O, C}
    machine::M
    output::O
    cost::C
    result::Union{OptimizationResults, Nothing}
end

"""
    KernelMachineRegression(X::AbstractArray, Y::AbstractArray;
        kernel=additivegaussiankernel, cost=0, dims)

Return a `KernelMachineRegression` object. `X` and `Y` are training input and output.
`Y` can be either a vector or a matrix. `kernel` is required to have a method
`kernel(u, v, k=nothing)`,where `u` and `v` are matrices (datapoints are columns),
and `k` is the kernel evaluated in the previous hidden space.
See [`additivegaussiankernel`](@ref) for an example.
`cost` is used to regularize. Finally, `dims` is a tuple containing the dimensionality
of the spaces that compose the machine space.

!!! note
    `dims` here is different than in [`KernelMachine`](@ref), as it excludes
    the dimensionality of the output. For instance, setting `dims=(3, 3, 3)`
    on a `KernelMachineRegression` with scalar output (i.e., `y::AbstractVector`)
    would correspond to a machine with `dims=(3, 3, 3, 1)`.

Use `fit!(kr::KernelMachineRegression)` to fit the `KernelMachineRegression` object, and
`predict(kr::KernelMachineRegression, input)` to get the model prediction on a dataset.
If no `input` is given, it will default to the training data.
"""
function KernelMachineRegression(X::AbstractArray, Y::AbstractArray;
    kernel=additivegaussiankernel, cost=0, dims)

    machine = KernelMachine(kernel, permutedims(X); dims=(dims..., size(Y, 2)))
    output = permutedims(Y)
    return KernelMachineRegression(machine, output, cost, nothing)
end

Base.show(io::IO, d::KernelMachineRegression) = print(io, "KernelMachineRegression {...}")

loss(k::KernelMachineRegression) = loss(k.machine, k.output, k.cost)

function loss(km::KernelMachine, output, cost)
    pred, sq_norm = km(nothing)
    mean(abs2, pred - output) + cost * sq_norm
end

function updatemachine!(k::KernelMachineRegression, w)
    km = k.machine
    aug, cs = km.augmenter, km.cs
    n_aug, n_cs = length(aug), length(cs)
    copyto!(aug, 1, w, 1, n_aug)
    copyto!(cs, 1, w, n_aug + 1, n_cs)
    return
end

function compute_fg(k::KernelMachineRegression)
    function fg!(_, G, w)
        updatemachine!(k, w)
        if isnothing(G)
            l = loss(k)
        else
            l, back = pullback(loss, k.machine, k.output, k.cost)
            gs, = back(one(l))
            copyto!(G, gs.augmenter)
            copyto!(G, length(gs.augmenter) + 1, gs.cs)
        end
        return l
    end
end

const default_method = ConjugateGradient()

function fit!(k::KernelMachineRegression,
    method::AbstractOptimizer=default_method,
    options::Options=Options(; default_options(method)...))

    fg! = compute_fg(k)
    aug, cs = k.machine.augmenter, k.machine.cs
    init = vcat(vec(aug), vec(cs))
    res = optimize(only_fg!(fg!), init, method, options)
    updatemachine!(k, minimizer(res))
    k.result = res
    return k
end

function predict(k::KernelMachineRegression, input=nothing)
    km = k.machine
    if isnothing(input)
        res, _ = km(nothing)
    else
        nsamples = size(km.data, 2)
        res_full, _ = km(permutedims(input))
        res = res_full[:, nsamples+1:end]
    end
    return permutedims(res)
end


@par function cross!(outx::AbstractArray{T},outy::AbstractArray{T},outz::AbstractArray{T},
                     ux::AbstractArray{T},uy::AbstractArray{T},uz::AbstractArray{T},
                     vx::AbstractArray{T},vy::AbstractArray{T},vz::AbstractArray{T},
                     s::@par(AbstractSimulation)) where {T<:Complex}
    @mthreads for ind in ZYRANGE
        jj,kk = Tuple(ind)
        k = ZRANGE[kk]
        j = YRANGE[jj]
            @inbounds @msimd for i in XRANGE
                outx[i,j,k] = uy[i,j,k]*vz[i,j,k] - uz[i,j,k]*vy[i,j,k]
                outy[i,j,k] = uz[i,j,k]*vx[i,j,k] - ux[i,j,k]*vz[i,j,k]
                outz[i,j,k] = ux[i,j,k]*vy[i,j,k] - uy[i,j,k]*vx[i,j,k]
            end
    end
end

@par function crossk!(outx,outy,outz,vx,vy,vz,s::@par(AbstractSimulation))
    @mthreads for ind in ZYRANGE
        jj,kk = Tuple(ind)
        k = ZRANGE[kk]
        j = YRANGE[jj]
            @inbounds @msimd for i in XRANGE
                outx[i,j,k] = im*(KY[j]*vz[i,j,k] - KZ[k]*vy[i,j,k])
                outy[i,j,k] = im*(KZ[k]*vx[i,j,k] - KX[i]*vz[i,j,k])
                outz[i,j,k] = im*(KX[i]*vy[i,j,k] - KY[j]*vx[i,j,k])
            end
    end
end

function ccross!(out::VectorField,u::VectorField,v::VectorField,s::AbstractSimulation)
    cross!(out.cx,out.cy,out.cz,u.cx,u.cy,u.cz,v.cx,v.cy,v.cz,s)
    return out
end

function curl!(out::VectorField,u::VectorField,s::AbstractSimulation)
    crossk!(out.cx,out.cy,out.cz,u.cx,u.cy,u.cz,s)
    return out
end

@par function grad!(out::VectorField,f::AbstractArray{<:Complex,3},s::@par(AbstractSimulation))
    outx = out.c.x
    outy = out.c.y
    outz = out.c.z
    @mthreads for ind in ZYRANGE
        jj,kk = Tuple(ind)
        k = ZRANGE[kk]
        j = YRANGE[jj]
            @inbounds @msimd for i in XRANGE
                outx[i,j,k] = f[i,j,k]*im*KX[i]
                outy[i,j,k] = f[i,j,k]*im*KY[j]
                outz[i,j,k] = f[i,j,k]*im*KZ[k]
            end
    end
    #dx!(out.cx,f,s)
    #dy!(out.cy,f,s)
    #dz!(out.cz,f,s)
    #dealias!(out,s)
end

@par function dx!(out::AbstractArray{<:Complex,3},f::AbstractArray{<:Complex,3},s::@par(AbstractSimulation))
    @mthreads for ind in ZYRANGE
        jj,kk = Tuple(ind)
        k = ZRANGE[kk]
        j = YRANGE[jj]
            @inbounds @msimd for i in XRANGE
                out[i,j,k] = f[i,j,k]*im*KX[i]
            end
    end
end

@par function dy!(out::AbstractArray{<:Complex,3},f::AbstractArray{<:Complex,3},s::@par(AbstractSimulation))
    @mthreads for ind in ZYRANGE
        jj,kk = Tuple(ind)
        k = ZRANGE[kk]
        j = YRANGE[jj]
            @inbounds @msimd for i in XRANGE
                out[i,j,k] = f[i,j,k]*im*KY[j]
            end
    end
end

@par function dz!(out::AbstractArray{<:Complex,3},f::AbstractArray{<:Complex,3},s::@par(AbstractSimulation))
    @mthreads for ind in ZYRANGE
        jj,kk = Tuple(ind)
        k = ZRANGE[kk]
        j = YRANGE[jj]
            @inbounds @msimd for i in XRANGE
                out[i,j,k] = f[i,j,k]*im*KZ[k]
            end
    end
end

@par function div!(out::AbstractArray{Complex{Float64},3},ux,uy,uz,w,mdρdz::Real, s::@par(AbstractSimulation))
    @mthreads for ind in ZYRANGE
        jj,kk = Tuple(ind)
        k = ZRANGE[kk]
        j = YRANGE[jj]
            @inbounds @msimd for i in XRANGE
               # out[i,j,k] = mim*(kx[i]*ux[i,j,k] + ky[j]*uy[i,j,k] + kz[k]*uz[i,j,k]) + mdœÅdz*w[i,j,k]
                out[i,j,k] = muladd(im,muladd(KX[i], ux[i,j,k], muladd(KY[j], uy[i,j,k], KZ[k]*uz[i,j,k])), mdρdz*w[i,j,k])
            end
    end
end

div!(out::AbstractArray{<:Complex,3},u::VectorField,s) = div!(out,u.cx,u.cy,u.cz,s)

@par function div!(out::AbstractArray{Complex{Float64},3},ux,uy,uz, s::@par(AbstractSimulation))
    @mthreads for ind in ZYRANGE
        jj,kk = Tuple(ind)
        k = ZRANGE[kk]
        j = YRANGE[jj]
            @inbounds @msimd for i in XRANGE
                #out[i,j,k] = -im*(kx[i]*ux[i,j,k] + ky[j]*uy[i,j,k] + kz[k]*uz[i,j,k])
                out[i,j,k] = im*muladd(KX[i], ux[i,j,k], muladd(KY[j], uy[i,j,k], KZ[k]*uz[i,j,k]))
            end
    end
end


using Distributed
using Printf

mutable struct Node{T<:AbstractFloat}
    split::T
    low::T
    up::T
    logit::Union{Nothing, T}
    left::Union{Node, Nothing}
    right::Union{Node, Nothing}
    function Node(
        split::T, low::T, up::T, logit::Union{T, Nothing} = nothing
    ) where T<:AbstractFloat
        new{T}(split, low, up, logit, nothing, nothing)
    end
end

##
isleaf(v::Node) = isnothing(v.left) && isnothing(v.right)
descend(v::Node{T}, x::T) where T<:AbstractFloat = (x ≤ v.split) ? (v.left, 0) : (v.right, 1)

# tostr(x::Union{Node, Float64, Nothing}; add_children=false) = begin
#     if isnothing(x)
#         return "(*)"
#     elseif isa(x, Float64)
#         return @sprintf("%.2f", x)
#     else
#         s = "Node($(tostr(x.split)), $(tostr(x.logit)))\n"
#         if add_children
#             s *= "|--$(tostr(x.left))\n"
#             s *= "|--$(tostr(x.right))"
#         end
#         return s
#     end
# end
# Base.show(io::IO, v::Node) = show(io, "Node(id=$(objectid(v)))")
# Base.print(io::IO, v::Node) = print(io, tostr(v, add_children=true))

function getval(v::Node{T}, x::T) where {T<:AbstractFloat}
    @assert !isnothing(v.logit) "can only evaluate nodes with values, use bfs_fill_tree"
    (x ≤ v.split) ? v.logit : -v. logit
end

function uniform_binary_splits(xmin::T, xmax::T, depth::Int) where T <: AbstractFloat
    lows = T[]
    mids = T[]
    ups = T[]
    levs = Int[]
    splitvals = collect(LinRange(xmin, xmax, 2^depth + 1))
    for d in 1:depth
        width = 2 ^ (depth - d)
        num_intervals = 2 ^ (d - 1)
        for j in 1:num_intervals
            start = 2 * (j - 1) * width + 1
            push!(lows, splitvals[start])
            push!(mids, splitvals[start + width])
            push!(ups, splitvals[start + 2width])
            push!(levs, d - 1)
        end
    end
    (lows=lows, mids=mids, ups=ups, levs=levs, splitvals=splitvals)
end

function eval_logprob(root::Node{T}, x::T) where T<:AbstractFloat
    @assert !isleaf(root) "must start from non-leaf"
    logprob = zero(T)
    v = root
    while !isnothing(v)
        β = getval(v, x)
        if β ≥ zero(T)
            logprob += -log(one(T) + exp(-β))
        else
            logprob += β - log(one(T)+ exp(β))
        end
        v, _ = descend(v, x)
    end
    return logprob
end

function eval_logdens(root::Node{T}, x::T) where T<:AbstractFloat
    @assert !isleaf(root) "must start from non-leaf"
    logprob = zero(T)
    v = root
    prev = v
    descdir = 0
    while !isnothing(v)
        β = getval(v, x)
        if β ≥ zero(T)
            logprob += -log(one(T) + exp(-β))
        else
            logprob += β - log(one(T)+ exp(β))
        end
        prev = v
        v, descdir = descend(v, x)
    end
    δ = (descdir == 0) ? prev.split - prev.low : prev.up - prev.split
    return logprob - log(δ)
end

##
function print_bfs(v::Node)
    visited = Queue{Node}()
    enqueue!(visited, v)
    while !isempty(visited)
        curr = dequeue!(visited)
        println(curr)
        !isnothing(curr.left) && enqueue!(visited, curr.left)
        !isnothing(curr.right) && enqueue!(visited, curr.right)
    end
end

##

function make_tree_from_bfs(
    lows::Vector{T},
    mid::Vector{T},
    ups::Vector{T},
    logits::Vector{S} = fill(nothing, length(mid))
) where {T<:AbstractFloat, S<:Union{Nothing, AbstractFloat}}
    num_nodes = length(mid)
    next_parent = Queue{Node}()
    root = Node(mid[1], lows[1], ups[1], logits[1])
    curr = root  # current parent
    for i in 2:num_nodes
        v = Node(mid[i], lows[i], ups[i], logits[i])
        while !(ups[i] ≈ curr.split || lows[i] ≈ curr.split)
            @assert !isempty(next_parent)  "disconnected tree!"
            curr = dequeue!(next_parent)
        end
        if ups[i] ≈ curr.split
            curr.left = v
        elseif lows[i] ≈ curr.split
            curr.right = v
        end
        enqueue!(next_parent, v)
    end
    return root
end

function make_tree(
    lows::Vector{T},
    mid::Vector{T},
    ups::Vector{T},
    logits::Vector{S} = fill(nothing, length(mid))
) where {T<:AbstractFloat, S<:Union{Nothing, AbstractFloat}}
    num_nodes = length(mid)
    nodes = [Node(mid[i], lows[i], ups[i], logits[i]) for i in 1:num_nodes]
    is_root = Set{Int}(1:num_nodes)
    for i in 1:num_nodes
        # serch for parent upwards
        child = nodes[i]
        child_lower, child_upper = lows[i], ups[i]
        for j in 1:num_nodes
            parent = nodes[j]
            parent_mid = mid[j]
            parent_lower, parent_upper = lows[j], ups[j]
            if child_lower ≈ parent_mid && child_upper ≈ parent_upper
                parent.right = child
                pop!(is_root, i)
                break
            elseif child_upper ≈ parent_mid && child_lower ≈ parent_lower
                parent.left = child
                pop!(is_root, i)
                break
            end
        end
    end
    if isempty(is_root)
        error("circular graph")
    elseif length(is_root) > 1
        error("multiple roots")
    end
    nodes[pop!(is_root)]
end


using RegionTrees
using StaticArrays
using PyPlot
using Statistics
using LinearAlgebra
using SparseArrays
using VectorizedRoutines
using ProgressBars
using MAT

struct ray3d{T}
    origin::T
    dir::T
    inv_direction::T
    sign::Union{Array{Int64,1},SArray{Tuple{3},Int64,1,3}}
end

function ray3d(p1::T, p2::T) where {T}
    dir::T = p2 - p1
    invdir::T = 1.0 ./ dir
    sign = (invdir .< 0) .* 1
    return ray3d(p1, dir, invdir, sign)
end

#An Efficient and Robust Ray–Box Intersection Algorithm by
#Amy Williams Steve Barrus R.Keith Morley Peter Shirley.
@inline @inbounds function intersects(
    r::ray3d,
    b::HyperRectangle{3,Float64},
)::Float64
    Q = [b.origin, b.origin + b.widths]
    tmin = (Q[r.sign[1]+1][1] - r.origin[1]) * r.inv_direction[1]
    tmax = (Q[1-r.sign[1]+1][1] - r.origin[1]) * r.inv_direction[1]
    tymin = (Q[r.sign[2]+1][2] - r.origin[2]) * r.inv_direction[2]
    tymax = (Q[1-r.sign[2]+1][2] - r.origin[2]) * r.inv_direction[2]
    if ((tmin > tymax) || (tymin > tmax))
        return -1.0
    end
    if (tymin > tmin)
        tmin = tymin
    end
    if (tymax < tmax)
        tmax = tymax
    end
    tzmin = (Q[r.sign[3]+1][3] - r.origin[3]) * r.inv_direction[3]
    tzmax = (Q[1-r.sign[3]+1][3] - r.origin[3]) * r.inv_direction[3]

    if ((tmin > tzmax) || (tzmin > tmax))
        return -1.0
    end
    if (tzmin > tmin)
        tmin = tzmin
    end
    if (tzmax < tmax)
        tmax = tzmax
    end
    if (~isfinite(tmax) || ~isfinite(tmin))
        error("An ambiguous ray trajectory detected: ray touches the voxel at its boundary. ")
    end
    return norm((r.origin + r.dir * tmin) - (r.origin + r.dir * tmax))

end

@inline function intersects(a::HyperRectangle{N,Float64}, b::HyperRectangle{N,Float64}) where {N}
    maxA = a.origin + a.widths
    minA = a.origin
    maxB = b.origin + b.widths
    minB = b.origin
    return (all(maxA .>= minB) && all(maxB .> minA))

end

function inittree3d(n)
    empty::Array{Int64,1} = []
    oct = Cell(SVector(-1.0, -1, -1), SVector(2.0, 2, 2), empty)
    for _ = 1:n
        for leaf in allleaves(oct)
            split!(leaf)
        end
    end
    for node in allcells(oct)
        node.data = copy(node.data)
    end
    return oct
end

function printdata(tree::Cell)
    for i = 1:length(tree.children)
        if (tree[i].children !== nothing)
            for j = 1:length(tree[i].children)
                println(i, " ", j, " ", tree[i][j].data)
            end
        else
            println(i, " ", tree[i].data)
        end
    end
end

function inserttotree!(
    box::HyperRectangle{N,Float64},
    tree::Cell,
    index::Int64,
) where {N}
    if (isleaf(tree))
        if (intersects(box, tree.boundary))
            push!(tree.data, index)
        end
    elseif (intersects(box, tree.boundary))
        for i = 1:length(tree.children)
            inserttotree!(box, tree[i], index)
        end
    end
    return nothing
end

function possiblevoxels(r::ray3d, tree::Cell)::Vector{Int64}
    if (isleaf(tree) && intersects(r, tree.boundary) > 0.0)
        return tree.data

    elseif (~(tree.parent === nothing) && intersects(r, tree.boundary) > 0.0)
        N = length(tree.children)
        v = Vector{Int64}()
        for i = 1:N
            append!(v, possiblevoxels(r, tree[i]))
        end
        return v

    elseif ((tree.parent === nothing) && intersects(r, tree.boundary) > 0.0)
        N = length(tree.children)
        v = Vector{Int64}()
        for i = 1:N
            append!(v, possiblevoxels(r, tree[i]))
        end
        return collect(Set(v))

    else
        return []
    end

end

function voxelray(
    r::ray3d,
    vv::Vector{HyperRectangle{3,Float64}},
    checklist::Vector{Int64},
)
    N = length(checklist)
    indices = Vector{Int64}()
    tlens = Vector{Float64}()
    for i = 1:N
        q = intersects(r, vv[checklist[i]])
        if (q>0.0)
            push!(tlens, q)
            push!(indices, checklist[i])
        end
    end
    return (indices, tlens)
end

@inline function rodrigues(k::T,v::T,theta::Float64)::T where {T}
    kn = k/norm(k)
    return v*cos(theta) + cross(kn,v)*sin(theta) + kn*dot(kn,v)*(1.0-cos(theta))
end

function drawrays(N, theta, fii)
    @assert length(theta) == length(fii)
    Na = length(theta)
    if (~isa(theta, Union{SArray,Array}))
        theta = [theta]
        fii = [fii]
    end
    r = 0.5 * sqrt(12)
    v1 = [1.0, 0, 0]
    v2 = [0, 1.0, 0]
    v3 = [0, 0, 1.0]
    span = range(-r, stop = r, length = N)
    using3D()
    fig = figure()
    ax = fig.add_subplot(111, projection = "3d")
    for a = 1:Na
        x1 = r * sin(theta[a]) * cos(fii[a])
        y1 = r * sin(theta[a]) * sin(fii[a])
        z1 = r * cos(theta[a])
        d = [x1, y1, z1]
        fullangle = acos(dot(d, v1) / (norm(d) * norm(v1)))
        k = cross(v1, [x1, y1, z1])
        mainaxle = rodrigues(k, v1, fullangle)
        axle2 = rodrigues(k, v2, fullangle)
        axle3 = rodrigues(k, v3, fullangle)
        center = @SVector [mainaxle[1] * r, mainaxle[2] * r, mainaxle[3] * r]
        for i = 1:N
            for j = 1:N
                p1 = center + axle2 * span[j] + axle3 * span[i]
                p2 = p1 - 2 * mainaxle * r
                ax.plot3D([p1[1], p2[1]], [p1[2], p2[2]], [p1[3], p2[3]])

            end
        end
    end

end

function constructmatrix(
    tree::Cell,
    vv::Vector{HyperRectangle{3,Float64}},
    N::Int64,
    theta,
    fii,
)
    @assert length(theta) == length(fii)
    v1 = [1.0, 0, 0]
    v2 = [0, 1.0, 0]
    v3 = [0, 0, 1.0]
    r = 0.5 * sqrt(12)

    Na = length(theta)
    if (~isa(theta, Union{SArray,Array}))
        theta = [theta]
        fii = [fii]
    end

    span = range(-r, stop = r, length = N)
    Ntotal = Na * N^2
    Nth = Threads.nthreads()
    rows = Vector{Vector{Int64}}(undef, Nth)
    cols = Vector{Vector{Int64}}(undef, Nth)
    vals = Vector{Vector{Float64}}(undef, Nth)
    for p = 1:Nth
        rows[p] = []
        cols[p] = []
        vals[p] = []
    end

    t = time()
    if (Na > 1)
        pb = ProgressBar(1:Na)
    else
        pb = 1:1
    end
    for a in pb
        x1 = r * sin(theta[a]) * cos(fii[a])
        y1 = r * sin(theta[a]) * sin(fii[a])
        z1 = r * cos(theta[a])
        d = [x1, y1, z1]
        fullangle = acos(dot(d, v1) / (norm(d) * norm(v1)))
        k = cross(v1, [x1, y1, z1])
        mainaxle = rodrigues(k, v1, fullangle)
        axle2 = rodrigues(k, v2, fullangle)
        axle3 = rodrigues(k, v3, fullangle)
        center = @SVector [mainaxle[1] * r, mainaxle[2] * r, mainaxle[3] * r]
        Threads.@threads for i = 1:N
            for j = 1:N
                rayindex = (a - 1) * N^2 + (i - 1) * N + j
                p1 = center + axle2 * span[j] + axle3 * span[i]
                p2 = p1 - 2 * mainaxle * r
                or = ray3d(p1, p2)
                checklist = possiblevoxels(or, tree)
                indices, tlens = voxelray(or, vv, checklist)
                # append!(rows,rayindex*ones(Int64,length(tlens)))
                # append!(cols,indices)
                # append!(vals,tlens)

                append!(rows[Threads.threadid()],rayindex * ones(Int64, length(tlens)))
                append!(cols[Threads.threadid()], indices)
                append!(vals[Threads.threadid()], tlens)

            end
        end
    end
    rows = vcat(rows...)
    cols = vcat(cols...)
    vals = vcat(vals...)
    M = sparse(rows, cols, vals, Ntotal, length(vv))
    println("Matrix constructed in ", time() - t, " seconds.")
    return M
end

function setupvoxels(sizex::Int64,sizey::Int64,sizez::Int64,octsize::Int64)
    oct = inittree3d(octsize)
    l = Array([[1.0, 0.0, 0.0] [0,0,1] [1, 1, 0] [0, 0, 0] [0, 1, 1] [1, 0, 1] [1, 1, 1] [0, 1, 0]]')
    pm = mean(l, dims = 1)
    size = [sizex, sizey, sizez]
    cellsize = 2.0 ./ size
    width = @SVector[cellsize[1], cellsize[2], cellsize[3]]
    for i = 1:3
        l[:, i] =cellsize[i] .* (l[:, i] .- pm[i]) .- cellsize[i] / 2.0 * size[i] .+ 0.5 * cellsize[i]
    end
    voxelvector = Vector{HyperRectangle{3,Float64}}(undef, sizex * sizey * sizez)
    t = time()
    for k = 1:sizez
        for j = 1:sizey
            for i = 1:sizex
                ll = l .+ reshape(cellsize .* [i - 1, j - 1, k - 1], 1, 3)
                origin = @SVector[
                    minimum(ll[:, 1]),
                    minimum(ll[:, 2]),
                    minimum(ll[:, 3]),
                ]
                voxel = HyperRectangle(origin, width)
                index = j + (i - 1) * sizey + (k - 1) * sizey * sizex
                inserttotree!(voxel, oct, index)
                voxelvector[index] = voxel
            end
        end
    end
    println("Voxels initialized in ", time() - t, " seconds.")
    return (oct, voxelvector)
end

function sub2ind(sizes,i::Int64,j::Int64,k::Int64)
    @assert i > 0 && i<= sizes[1]
    @assert j > 0 && j<= sizes[2]
    @assert k > 0 && k<= sizes[3]
    return (k-1)sizes[1]*sizes[2] + sizes[1]*(j-1) + i
end

function test()
    Nx = 100; Ny = 100; Nz = 100; Os = 6; Ns = 100
    (oct,voxelvector)=setupvoxels(Nx,Ny,Nz,Os)
    theta = [0,pi/2]; fii = [0, pi/2]
    M=constructmatrix(oct,voxelvector,Ns,theta,fii)

    l = Array([[1.0, 0.0, 0.0] [0,0,1] [1, 1, 0] [0, 0, 0] [0, 1, 1] [1, 0, 1] [1, 1, 1] [0, 1, 0]]')
    pm = mean(l, dims = 1)
    size = [Nx,Ny,Nz]
    cellsize = 2.0 ./ size
    qx=Vector(range(-1 + cellsize[1]/2.0,stop=1- cellsize[1]/2.0,length=Nx))
    qy=Vector(range(-1+ cellsize[2]/2.0,stop=1- cellsize[2]/2.0,length=Ny))
    qz=Vector(range(-1+ cellsize[3]/2.0,stop=1- cellsize[3]/2.0,length=Nz))

    mgrid = Matlab.meshgrid(qx,qy,qz)
    x = mgrid[1]; y = mgrid[2]; z = mgrid[3]

    object = @. 5.0.*(x^2/1^2 + (y)^2/0.5^2 + z^2/1^2 < 0.25)
    fig = plt.figure()
    imshow(object[:,:,Nz÷2+1])
    objf = object[:]
    sino = M*objf
    sino1 = sino[1:Ns^2]
    sinor1 = reshape(sino1,Ns,Ns)
    fig = plt.figure()
    imshow(sinor1)

    fig = plt.figure()
    sino2 = sino[Ns^2+1:end]
    sinor2 = reshape(sino2,Ns,Ns)
    imshow(sinor2)
    return  nothing

end

test()


using SparseArrays
using LinearAlgebra
using jInv.Mesh
using Multigrid
using Multigrid.DomainDecomposition
using Helmholtz

domain = [0.0, 1.0, 0.0, 1.0];
n = [64,64];
Mr = getRegularMesh(domain,n)

pad = 8;
m = ones(tuple(n.+1...));

maxOmega = getMaximalFrequency(m,Mr);
ABLamp = maxOmega;

w = 0.9*maxOmega;
println("omega*h:");
println(w*Mr.h*sqrt(maximum(m)));
pad = pad*ones(Int64,Mr.dim);

println("************************************************* Hybrid Keczmarz for Acoustic Helmholtz 2D ******************************************************");

Sommerfeld = true;
NeumannAtFirstDim = true;
H,gamma = GetHelmholtzOperator(Mr,m,w,w*ones(size(m))*0.01,NeumannAtFirstDim,pad,ABLamp,Sommerfeld);
shift = 0.5;
SHT = sparse((H .+ GetHelmholtzShiftOP(m, real(w),shift))');

numCores = 2;
innerKatzIter = 3;
inner = 3;

q = getAcousticPointSource(Mr,ComplexF64)[1];
x = zeros(ComplexF64,size(q));
N = size(SHT,2);

HKparam = getHybridKaczmarz(ComplexF64,Int64, SHT,Mr,[4,4], getNodalIndicesOfCell,0.95,numCores,innerKatzIter);
prec = getHybridKaczmarzPrecond(HKparam,SHT);
x[:] .= FGMRES_relaxation(SHT,q[:],x[:],inner,prec,1e-3,true,true,numCores)[1]
println("Relative norm is: ",norm(SHT'*x[:] - q[:])./norm(q[:]))

println("In float precision:")

x[:].=0.0;
x = convert(Array{ComplexF32},x);
q = convert(Array{ComplexF32},q);
SHT = convert(SparseMatrixCSC{ComplexF32,Int64},SHT);

HKparam = getHybridKaczmarz(ComplexF32,Int64, SHT,Mr,[4,4], getNodalIndicesOfCell,0.95,numCores,innerKatzIter);
prec = getHybridKaczmarzPrecond(HKparam,SHT);
x[:] .= FGMRES_relaxation(SHT,q[:],x[:],inner,prec,1e-3,true,true,numCores)[1]
println("Relative norm is: ",norm(SHT'*x[:] - q[:])./norm(q[:]))


module TweedieDistributions

using Distributions
using Random

import Distributions: mean, var, succprob, failprob, insupport
import Base: rand, minimum, maximum

export Tweedie,
    CompoundPoissonGamma,
    rand,
    mean,
    var,
    succprob,
    failprob

abstract type AbstractTweedie{T} <: Distribution{Univariate,Continuous} end

"""
    Tweedie{T}
       μ::T
       p::T
       ϕ::T

Tweedie distribution. Typically used with 1 < ξ < 2 to give compound
gamma-Poisson. Has positive mass at zero, and is continuous for positive
Reals.
"""
struct Tweedie{T} <: AbstractTweedie{T}
    μ::T
    p::T
    ϕ::T
end

"""
    CompoundPoissonGamma{T} <: AbstractTweedie
        μ::T
        p::T
        ϕ::T

A Tweedie distribution with ``1 < p < 2`` is a compound Poisson-Gamma
distribution, making it more straightforward to draw random samples from.
"""
struct CompoundPoissonGamma{T} <: AbstractTweedie{T}
    μ::T
    p::T
    ϕ::T

    function CompoundPoissonGamma(μ::T, p::T, ϕ::T) where T
        μ ≥ 0 || throw("Mean must be nonnegative")
        ϕ > 0 || throw("Variance parameter must be positive")
        1 < p < 2 ||
            throw("CPG requires 1 < p < 2, otherwise use Tweedie.")
        new{T}(μ, p, ϕ)
    end
end

function CompoundPoissonGamma(Tw::Tweedie)
    CompoundPoissonGamma(Tw.μ, Tw.p, Tw.ϕ)
end

function CompoundPoissonGamma(μ, p, ϕ)
    μ_prom, p_prom, ϕ_prom = promote(μ, p, ϕ)
    CompoundPoissonGamma(μ_prom, p_prom, ϕ_prom)
end

@inline function _cpg_poissonmean(CPG::CompoundPoissonGamma)
    CPG.μ^(2 - CPG.p) / ((2 - CPG.p) * CPG.ϕ)
end

@inline function _cpg_gammashape(CPG::CompoundPoissonGamma)
    gam_shape = (2 - CPG.p) / (CPG.p - 1)
end

@inline function _cpg_gammarate(CPG::CompoundPoissonGamma)
    CPG.ϕ * (CPG.p - 1) * CPG.μ ^ (CPG.p - 1)
end

sampler(cpg::CompoundPoissonGamma) = CompoundPoissonGammaSampler(cpg)
rand(rng::AbstractRNG, cpg::CompoundPoissonGamma) = rand(rng, sampler(cpg))

# function rand(rng::Random.AbstractRNG, CPG::CompoundPoissonGamma{T}) where T
#     pois_mean = _cpg_poissonmean(CPG)
#     N = rand(rng, Poisson(pois_mean))
#     N == 0 && return zero(T)
#     gam_shape = N * _cpg_gammashape(CPG)
#     gam_rate = _cpg_gammarate(CPG)
# end

mean(Tw::AbstractTweedie) = Tw.μ
var(Tw::AbstractTweedie) = Tw.ϕ * Tw.μ ^ Tw.p

# rand(CPG::CompoundPoissonGamma) = rand(Random.GLOBAL_RNG, CPG)

function succprob(CPG::CompoundPoissonGamma{T}) where T
    N = _cpg_poissonmean(CPG)
    ccdf(Poisson(N), zero(T))
end

function failprob(CPG::CompoundPoissonGamma{T}) where T
    N = _cpg_poissonmean(CPG)
    cdf(Poisson(N), zero(T))
end

minimum(CPG::CompoundPoissonGamma{T}) where T = zero(T)
maximum(CPG::CompoundPoissonGamma) = Inf
# Need to restrict `x` to ::Real or dispatch is ambiguous
insupport(CPG::CompoundPoissonGamma, x::Real) = 0 ≤ x < Inf

struct CompoundPoissonGammaSampler{T} <: Sampleable{Univariate, Continuous}
    λ::T
    α::T
    θ::T
end

function CompoundPoissonGammaSampler(CPG::CompoundPoissonGamma{T}) where T
    λ = _cpg_poissonmean(CPG)
    α = _cpg_gammashape(CPG)
    θ = _cpg_gammarate(CPG)
    CompoundPoissonGammaSampler{T}(λ, α, θ)
end

function rand(rng::Random.AbstractRNG, cpgs::CompoundPoissonGammaSampler{T}) where T
    N = rand(rng, Poisson(cpgs.λ))
    N == 0 && return zero(T)
    rand(rng, Gamma(N * cpgs.α, cpgs.θ))
end
rand(cpgs::CompoundPoissonGammaSampler) = rand(Random.GLOBAL_RNG, cpgs)

end # module