######### gradle-to-gate calculation

# the direct requirements coefficients
A = [0.4  0.2  0.1 ;  # 1/Manufacturing/US
     0.2  0.1  0.1 ;  # 2/Transport/US
     0.3  0.3  0.2 ]  # 3/Energy/US

# the satellite table
B = [2.0  0.5  1]     # CO2

# the demand vector
f = [200.0 ;
       0.0 ;
      50.0 ]

# the scaling vector
I = eye(size(A)[1])
s = (I - A) \ f

# the total LCI result
g = B * s  # 1163.27 kg CO2

# the direct contribution results
G = B * diagm(s)   # [836.735  61.2245  265.306]  kg CO2

######### additional requirements in the distribution phase

# additional requirements matrix of the distribution phase
K = [ 0.0  0.0  0.0 ;  # 1/Manufacturing/US
      0.3  0.0  0.1 ;  # 2/Transport/US
      0.0  0.0  0.0 ]  # 3/Energy/US

# demand vector of the distribution phase
fd = K * f  # [0.0  65.0  0.0]'

# the scaling vector of the distribution phase
sd = (I - A) \ fd

# the total LCI result of the distribution phase
gd = B * sd  # 160.131 kg CO2

# the direct contribution results of the distribution phase
Gd = B * diagm(sd)   # [72.0117  42.6385  45.481]  kg CO2

######### additional emissions in the distribution phase

Bd = [ 0.0  0.0  5.0]  # kg CO2

gd2 = Bd * f   # 250.0 kg CO2

Gd2 = Bd * diagm(f)  # [0.0  0.0  250.0] kg CO2


using Manopt, Manifolds, Test
@testset "Stochastic Gradient Descent" begin
    M = Sphere(2)
    # 5 point mean
    p = [0.0, 0.0, 1.0]
    s = 1.0
    pts = [
        exp(M, p, X) for
        X in [zeros(3), [s, 0.0, 0.0], [-s, 0.0, 0.0], [0.0, s, 0.0], [0.0, -s, 0.0]]
    ]
    F(y) = 1 / 2 * sum([distance(M, y, x)^2 for x in pts])
    gradF(M, y) = sum([-log(M, y, x) for x in pts])
    sgradF1(M, y) = [-log(M, y, x) for x in pts]
    function sgradF1!(M, X, y)
        for i in 1:length(pts)
            log!(M, X[i], y, pts[i])
            X[i] .*= -1
        end
        return X
    end
    sgradF2 = [((M, y) -> -log(M, y, x)) for x in pts]
    sgradF2! = [function f!(M, X, y)
        log!(M, X, y, x)
        X .*= -1
        return X
    end for x in pts]

    @testset "Constructors" begin
        p1 = StochasticGradientProblem(M, sgradF1)
        p1e = StochasticGradientProblem(M, sgradF1!; evaluation=MutatingEvaluation())
        p2 = StochasticGradientProblem(M, sgradF2)
        p2m = StochasticGradientProblem(M, sgradF2!; evaluation=MutatingEvaluation())
        @test get_gradient(p1, 1, p) == zeros(3)
        @test get_gradient(p2, 1, p) == zeros(3)
        for pr in [p1, p2, p2m]
            X = zero_vector(M, p)
            get_gradient!(pr, X, 1, p)
            @test X == get_gradient(pr, 1, p)
        end
        @test get_gradients(p1, p) == get_gradients(p2, p)
        @test get_gradients(p2m, p) == get_gradients(p2, p)
        Z = get_gradients(p2m, p)
        Z2 = similar.(Z)
        get_gradients!(p2m, Z2, p)
        @test Z == Z2
        Z3 = similar.(Z)
        get_gradients!(p1, Z3, p)
        @test Z == Z3
        Z4 = similar.(Z)
        get_gradients!(p2, Z4, p)
        @test Z == Z4
        Z5 = similar.(Z)
        get_gradients!(p1e, Z5, p)
        @test Z == Z5
        @test_throws ErrorException get_gradients(p1e, p)
        @test_throws ErrorException get_gradient!(p1e, Z4, 1, p)
        o = StochasticGradientDescentOptions(
            M, deepcopy(p), zero_vector(M, p); direction=StochasticGradient(deepcopy(p))
        )
        o.order = collect(1:5)
        o.order_type = :Linear
        initialize_solver!(p1, o)
        o.order_type = :Linear
        step_solver!(p1, o, 1)
        @test o.x == exp(M, p, get_gradient(p1, 1, p))
    end
    @testset "Momentum and Average Processor Constructors" begin
        p1 = StochasticGradientProblem(M, sgradF1)
        p2 = GradientProblem(M, F, gradF)
        m1 = MomentumGradient(M, p, StochasticGradient(zero_vector(M, p)))
        m2 = MomentumGradient(p1, p)
        @test typeof(m1.direction) == typeof(m2.direction) #both use StochasticGradient
        m3 = MomentumGradient(M, p)
        m4 = MomentumGradient(p2, p)
        @test m3.direction == m4.direction #both use Gradient
        a1 = AverageGradient(M, p, 10, StochasticGradient(zero_vector(M, p)))
        a2 = AverageGradient(p1, p, 10)
        @test typeof(a1.direction) == typeof(a2.direction) #both use StochasticGradient
        a3 = AverageGradient(M, p, 10)
        a4 = AverageGradient(p2, p, 10)
        @test a3.direction == a4.direction #both use Gradient
    end
    @testset "Comparing Stochastic Methods" begin
        x1 = stochastic_gradient_descent(M, sgradF1, p; order_type=:Linear)
        @test norm(x1) ≈ 1
        o1 = stochastic_gradient_descent(
            M, sgradF1, p; order_type=:Linear, return_options=true
        )
        x1a = get_solver_result(o1)
        @test x1 == x1a
        x2 = stochastic_gradient_descent(M, sgradF1, p; order_type=:FixedRandom)
        @test norm(x2) ≈ 1
        x3 = stochastic_gradient_descent(M, sgradF1, p; order_type=:Random)
        @test norm(x3) ≈ 1
        x4 = stochastic_gradient_descent(
            M,
            sgradF1,
            p;
            order_type=:Random,
            direction=AverageGradient(M, p, 10, StochasticGradient(zero_vector(M, p))),
        )
        @test norm(x4) ≈ 1
        x5 = stochastic_gradient_descent(
            M,
            sgradF1,
            p;
            order_type=:Random,
            direction=MomentumGradient(M, p, StochasticGradient(zero_vector(M, p))),
        )
        @test norm(x5) ≈ 1
    end
end


using Distributions, SpecialFunctions, Roots, LinearAlgebra
"""
GSmileOriginal : SMiLe rule (Faraji, et al., 2018)
    for the Gaussian task (section 2.3.1)
    - No knowledge of environment's prior (i.e. μ0 and σ0) for the calculation
        of Scc (it is there for the calculation of the χ and ν though)
    - Optimization to find γ (There is a closed formula for this task)
"""
struct GSmileOriginal <: GLearner
    m::Float64
    χ_0::Float64
    ν_0::Float64
    sigma::Float64 # std of observations
    χ_n::Array{Float64, 1}
    ν_n::Array{Float64, 1}
    mu_n::Array{Float64, 1} #parameter we are interested in
    Scc::Array{Float64, 1}
    γ_n::Array{Float64, 1}
end
function GSmileOriginal(;m = .1, mu_0 = 0., sigma_0 = 1., sigma = 0.1)
    χ_n = Array{Float64, 1}(undef, 1)
    ν_n = Array{Float64, 1}(undef, 1)
    mu_n = Array{Float64, 1}(undef, 1)
    ν_0 = calcν(sigma, sigma_0)
    χ_0 = calcχ(mu_0, ν_0, sigma)
    χ_n[1] = copy(χ_0)
    ν_n[1] = copy(ν_0)
    mu_n[1] = sigma * χ_n[1] / ν_n[1]
    GSmileOriginal(m, χ_0, ν_0, sigma, χ_n, ν_n, mu_n, [1.], [m/(1. + m)])
end
export GSmileOriginal
function update!(learnerG::GSmileOriginal, y)

    σ_n_squared = learnerG.sigma^2 / learnerG.ν_n[1]

    # The scaled likelihood is ~N(y, σ^2), assuming uniform prior (σ0 -> Inf)
    learnerG.Scc[1] = KL(learnerG.mu_n[1], σ_n_squared, y, learnerG.sigma^2)
    Bmax = KL(y, learnerG.sigma^2, learnerG.mu_n[1], σ_n_squared)
    B = learnerG.m * learnerG.Scc[1]/(1. + learnerG.m * learnerG.Scc[1]) * Bmax

    learnerG.γ_n[1] = find_γ0(y, learnerG.sigma^2, learnerG.mu_n[1], σ_n_squared, B)

    # Update using scaled likelihood: x0 = 0, and v0 = 0 (cause σ0 -> Inf)
    update_χ_ν!(learnerG, calcφ(learnerG, y), 1.)

    computemu_n!(learnerG)
end
export update!
function KL(μ1, σ1_squared, μ2, σ2_squared)
    σ_ratio = σ1_squared / σ2_squared
    max((μ1 - μ2)^2/(2. * σ2_squared) + 0.5 * (σ_ratio - 1. - log(σ_ratio)), 0.)
end
function update_χ_ν!(learnerG::GSmileOriginal, χ_0_y, ν_0_y)
    learnerG.χ_n[1] = (1. - learnerG.γ_n[1]) * learnerG.χ_n[1] + learnerG.γ_n[1] * χ_0_y
    learnerG.ν_n[1] = (1. - learnerG.γ_n[1]) * learnerG.ν_n[1] + learnerG.γ_n[1] * ν_0_y
end
function find_γ0(y, σ_squared, μ_n, σ_n_squared, B)
    # f = γ -> KL(γ .* betas .+ (1. - γ) .* alphas, alphas) - B
    ρ = σ_squared / σ_n_squared
    # w = 1. / (1. + (1. - γ)*ρ/γ)
    # μ_{n+1} = y*w + μ_n*(1. - w)
    #σ_{n+1} = 1. / (γ/σ^2 + (1. - γ)/σ_{n}^2)
    f = γ -> KL(y*get_w_smile(γ, ρ) + μ_n*(1. - get_w_smile(γ, ρ)), #μ_{n+1}
                get_σ_squared_smile(γ, σ_squared, σ_n_squared), #σ_{n+1}
                μ_n, σ_n_squared) - B
    if abs(f(0.)) < 5*eps()
        γ0 = 0.
    elseif abs(f(1.)) < 5*eps()
        γ0 = 1.
    else
        γ0 = find_zero(f, (0., 1.))
    end
    γ0
end
function get_w_smile(γ, ρ)
    1. / (1. + (1. - γ)*ρ/γ)
end
function get_σ_squared_smile(γ, σ_squared, σ_n_squared)
    1. /(γ/σ_squared + (1. - γ)/σ_n_squared)
end


function cost_no_Etot(L::Array{Float64, 1}, N::Array{Float64, 1}, Ltot::Float64, γ::Float64)
    #### Returns a Float64. It is the cost function evaluated in γ (in the case where there is no constrain on the energy) ####
    # getting the constants we will need
    n_pop = length(L)

    # Let's start the integration
    int = 0.0
    for a = 1:n_pop
        Lₐ = L[a]
        nₐ = N[a]

        x = γ * Lₐ

        int += nₐ * (Lₐ * coth(x) - 1 / γ)
    end

    return(Ltot - int)
end

function Dichotomy_step(L::Array{Float64, 1}, N::Array{Float64, 1}, Ltot::Float64, γ_min::Float64, γ_max::Float64)
    #### Returns a tuple of 2 floats. They are the new upper and lower bounds at the end of the dichotomy step ####
    # Let's get the current middle of the segment and the value of the cost function there
    γ_mid = (γ_max + γ_min) / 2
    c_mid = cost_no_Etot(L, N, Ltot, γ_mid)

    # Let's get the new bounds
    if (c_mid > 0) # We use the fact that the cost function is monotonous and decreasing
        return (γ_mid, γ_max)
    elseif (c_mid < 0)
        return (γ_min, γ_mid)
    else
        return (γ_mid, γ_mid)
    end
end

function Solution_no_Etot_dichotomy(L::Array{Float64, 1}, N::Array{Float64, 1}, Ltot::Float64)
    #### Returns a Float64. It is the optimal γ found by the dichotomy method ####
    #### This function performs the whole dichotomy                           ####

    # first let's make define the starting bounds of the dichotomy
    min_L, useless = findmin(L)
    max_L, useless = findmax(L)

    γ_min = 10.0^(-7) / max_L # these bounds come from looking at the shape of coth(x) - 1/x
    γ_max = 10.0^(2) / min_L
    γ_mid = (γ_max + γ_min) / 2

    # then let's define an acceptable relative error and maximum number of steps
    ϵ = 10^(-6)
    n_step = 100

    # then let's do the dichotomy method itself
    i = 0
    cost = Inf
    while ((cost > ϵ) && (i <= n_step))
        γ_min, γ_max = Dichotomy_step(L, N, Ltot, γ_min, γ_max) # update the bounds

        γ_mid = (γ_max + γ_min) / 2
        cost = abs(cost_no_Etot(L, N, Ltot, γ_mid)) # update the cost and test wether it is acceptable

        i += 1
    end

    return(γ_mid)
end

function Magnetisations_no_Etot(L::Array{Float64, 1}, N::Array{Float64, 1}, γ::Float64,
                                lmax::Int64,
                                List_x::Array{Float64,1}, List_w::Array{Float64,1}, List_Y::Array{Float64, 2})
    #### Returns an array of size n_pop $ lmax_half. Each cell contains one (axisymmetric, even) magnetisation                               ####
    #### This function computes the magnetization from a solution given by the dichotomy method, but doesn't run the dichotomy method itself ####
    # getting the constants we will need
    n_pop = length(L)
    lmax_half = div(lmax, 2)
    n_quad = length(List_x)

    # Preallocating the arrays we will need
    M = zeros(n_pop * lmax_half) # Will hold the magnetisations

    # Let's start to compute the magnetisations
    for star = 1:n_pop
        Lₛ = L[star]
        nₛ = N[star]

        # first let's compute the renrmalizing factor (lower integral)
        int_down = 0.0
        for i = 1:n_quad
            weight = List_w[i]
            x = List_x[i]

            int_down += weight * exp(γ * Lₛ * x)
        end

        # then let's compute the upper integral
        for l_half = 1:lmax_half
            index = (star - 1) * lmax_half + l_half

            for i = 1:n_quad
                spherical_harmonic = List_Y[l_half, i]
                weight = List_w[i]
                x = List_x[i]

                M[index] += weight * nₛ * spherical_harmonic * exp(γ * Lₛ * x)
            end

            M[index] /= int_down
        end
    end

    return(M)
end


using RosDataProcess

ts = range(-2pi, stop=2pi, length=500)

a = TimeSeries(sin, ts)
b = differentiate(a)

#plot(a)
#plot(b)


# optimization with potential
using ColorSchemes, Revise, FourierGPE

import FourierGPE:Lgp!, nlin!, V
gr(titlefontsize=12,size=(500,300),transpose=true,colorbar=false)

# prettier plot
import FourierGPE.showpsi

bone = cgrad(ColorSchemes.bone_1.colors)
turbo = cgrad(ColorSchemes.turbo.colors)
function showpsi(x,y,ψ)
    p1 = heatmap(x,y,abs2.(ψ),aspectratio=1,c=bone)
    xlims!(x[1],x[end]);ylims!(y[1],y[end])
    xlabel!(L"x");ylabel!(L"y")
    title!(L"|\psi|^2")
    p2 = heatmap(x,y,angle.(ψ),aspectratio=1,c=turbo)
    xlims!(x[1],x[end]);ylims!(y[1],y[end])
    xlabel!(L"x");ylabel!(L"y")
    title!(L"\textrm{phase} (\psi)")
    p = plot(p1,p2,size=(600,300))
    return p
end

function Lgp!(dϕ,ϕ,sim,t)
    @unpack γ,μ,espec = sim
    nlin!(dϕ,ϕ,sim,t)
    @. dϕ = -im*(1.0 - im*γ)*(dϕ + (espec - μ)*ϕ)
    return nothing
end

function nlin!(dϕ,ϕ,sim::Sim{2},t)
    @unpack g,X,V0 = sim; x,y = X
    y = y'
    dϕ .= ϕ
    xspace!(dϕ,sim)
    @. dϕ *= V0 + V(x,y,t) + g*abs2(dϕ)
    # @. dϕ *= V(x,y,t) + g*abs2(dϕ)
    kspace!(dϕ,sim)
    return nothing
end

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

# ==== set simulation parameters
μ = 15.0

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

# Or, as a static Potential
# x,y = X
# Vs(x,y) = 0.5*(x^2 + y^2)
# ψ0(x,y,μ,g) = sqrt(μ/g)*sqrt(max(1.0-Vs(x,y)/μ,0.0)+im*0.0)
# sim = Sim(sim,V0=Vs.(x,y'))

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

# ==== evolve
@time sol = runsim(sim)

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

# ==== set simulation parameters
γ = 0.0
t = LinRange(ti,tf,Nt)
ϕi = kspace(ψg,sim)
# reltol = 1e-7
# alg = DP5()

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

# vcore = Exact(VortexDistributions.ψi,ξ)
vcore = Exact(ξ)
pv = PointVortex(rv,0.,1)
vi = ScalarVortex(vcore,pv)
psi = Torus(copy(ψg),x,y)
vortex!(psi,vi)
showpsi(x,y,psi.ψ)
ψi .= psi.ψ
ϕi = kspace(ψi,sim)

# ===== compare with Fetter JLTP 2010
ξ = 1/sqrt(μ)
Rtf = R(1)
Ωm = 3*log(Rtf/ξ/sqrt(2))/2/Rtf^2
Ωv = Ωm/(1-rv^2/Rtf^2)

# or a precession period
Tv = 2*pi/Ωv

ti = 0.; tf = Tv
t = LinRange(ti,tf,Nt)

@pack_Sim! sim

# ==== evolve
@time solv = runsim(sim)

# ==== analyse
ϕf = solv[end]
ψf = xspace(ϕf,sim)
showpsi(x,y,ψf)

anim = @animate for i in eachindex(t)
    ψ = xspace(solv[i],sim)
    showpsi(x,y,ψ)
end

saveto = joinpath(@__DIR__,"vortexprecession.gif")
gif(anim,saveto,fps=20)


using PyPlot

function plot_solution(n::Int64,sol::Array{Float64,1},desiredstate::Function,alpha::Float64)
	# Simple scatter-plot of the nodes
	# Input: solution sol=(y,p,lambda)
	#        Desired state: y_d
	#	 L2-regularization alpha

	n_block::Int64 = Int64(length(sol)/3)

	dof_map = load("data"*string(n)*".jld2", "dof_map")
	xval::Array{Float64,1} = zeros(n_block)
	yval::Array{Float64,1} = zeros(n_block)

	stateval::Array{Float64,1} = zeros(n_block)
	controlval::Array{Float64,1} = zeros(n_block)
	adjointval::Array{Float64,1} = zeros(n_block)
	lagrangeval::Array{Float64,1} = zeros(n_block)
	y_d::Array{Float64,1} = zeros(n_block)

	# Separating all variables into vectors
	for i=1:n_block
		xval[i] = dof_map[i,:][1]
		yval[i] = dof_map[i,:][2]
		y_d[i] = desiredstate(xval[i],yval[i])
		stateval[i] = sol[i]
		adjointval[i] = sol[i+n_block]
		lagrangeval[i] = sol[i+2*n_block]
		controlval[i] = (1. /alpha)*(adjointval[i] - lagrangeval[i] )
	end

	# Plot state
	figure("1")
	scatter3D(xval, yval, stateval, s=0.1, c="k", label="State")
	scatter3D(xval, yval, y_d, s=0.1, c="b", alpha=0.1, label="Desired state")
	title("State")
	xlabel("x_1")
	ylabel("x_2")
	zlabel("y")
	legend(loc="best")

	# Plot Control
	figure("2")
	scatter3D(xval, yval, controlval, s=0.1, c="k")
	title("Control")
	xlabel("x_1")
	ylabel("x_2")
	zlabel("u")
end


function WrapDist(x,y)
    # Computes the wrap-around distance on the torus [0,1).
    #
    # Inputs:
    # x: first frequency
    # y: second frequency
    #
    # Output:
    # wrap around distance between x and y

    return min(abs(x-y),abs(x-y+1))
end


using AdventOfCodeSolutions
using Test

function input(puzzle::Puzzle{2020, 5, n}) where n
    io = openInput(puzzle)
    split(read(io, String), "\n", keepempty=false)
end

seatRow(str) = map(s -> s == 'F' ? 0 : 1, collect(str[1:7]))
seatColumn(str) = map(s -> s == 'L' ? 0 : 1, collect(str[8:end]))
parseSeat(str) = parse(Int, join(str), base=2)

function decode(str)
    row = seatRow(str) |> parseSeat
    column = seatColumn(str) |> parseSeat
    return (row = row, column = column)
end

seatId(str::AbstractString) = seatId(decode(str)...)
seatId(row, column) = 8 * row + column
seatId(seat) = seatId(seat...)

solve(::Puzzle{2020, 5, 1}, input) = max(seatId.(input)...)

notVeryFrontOrBack((row, column),) = row != 0 && row != 127
hasGap((left, right),) = left == right - 2

AdjacentElements(c) = Iterators.zip(c, Iterators.drop(c, 1))

function solve(::Puzzle{2020, 5, 2}, input)
    seats = decode.(input)
    takenSeats = seatId.(filter(notVeryFrontOrBack, seats)) |> sort

    (left, right) = Iterators.filter(hasGap, AdjacentElements(takenSeats)) |> first
    return left + 1
end

@testset "Day 5" begin
    @test(seatId("BFFFBBFRRR") == 567)
    @test(seatId("FFFBBBFRRR") == 119)
    @test(seatId("BBFFBBFRLL") == 820)
    @test(solve(Puzzle(5, 1), input(Puzzle(5, 1))) == 894)
    @test(solve(Puzzle(5, 2), input(Puzzle(5, 2))) == 579)
end


# Minesweeper:

mutable struct Field
    size::Tuple{Int, Int}
    numbers::Array{Int, 2}
    possible_mines::Array{Bool, 2}
    actual_mines::Array{Bool, 2}
    visible::Array{Bool, 2}
end

function Field(x, y)
    size = (x, y)
    actual_mines = convert(Array{Bool, 2}, rand(x, y) .< 0.15)
    possible_mines = zeros(Bool, x, y)
    numbers = zeros(Int, x, y)
    visible = zeros(Bool, x, y)
    for i = 1:x
        for j = 1:y
            n = 0
            for di = -1:1
                for dj = -1:1
                    n += (0 < di+i <= x && 0 < dj+j <= y) ?
                         (actual_mines[di+i, dj+j] ? 1 : 0) : 0
                end
            end
            numbers[i, j] = n
        end
    end
    return Field(size, numbers, possible_mines, actual_mines, visible)
end

function printfield(f::Field; showall = false)
    spaces = Int(floor(log(10, f.size[2])))

    str = " "^(4+spaces)
    for i in 1:f.size[1]
        str *= string(" ", i, " ")
    end
    str *= "\n" * " "^(4+spaces) * "___"^f.size[1] * "\n"
    for j = 1:f.size[2]
        str *= " " * string(j) * " "^(floor(log(10, j)) > 0 ? 1 : spaces+1) * "|"
        for i = 1:f.size[1]
            if showall
                str *= f.actual_mines[i, j] ? " * " : "   "
            else
                if f.visible[i, j]
                    str *= " " * string(f.numbers[i, j] > 0 ? f.numbers[i, j] : " ") * " "
                elseif f.possible_mines[i, j]
                    str *= " ? "
                else
                    str *= " . "
                end
            end
        end
        str *= "\r\n"
    end
    println("Found " * string(length(f.possible_mines[f.possible_mines.==true])) *
            " of " * string(length(f.actual_mines[f.actual_mines.==true])) * " mines.\n")
    print(str)
end

function parse_input(str::String)
    input = split(chomp(str), " ")
    mode = input[1]
    println(str)
    coords =  length(input) > 1 ? (parse(Int,input[2]), parse(Int,input[3])) : (0, 0)
    return mode, coords
end

function eval_input(f::Field, str::String)
    mode, coords = parse_input(str)
    (coords[1] > f.size[1] || coords[2] > f.size[2]) && return true
    if mode == "o"
        reveal(f, coords...) || return false
    elseif mode == "m" && f.visible[coords...] == false
        f.possible_mines[coords...] = !f.possible_mines[coords...]
    elseif mode == "close"
        error("You closed the game.")
    end
    return true
end

function reveal(f::Field, x::Int, y::Int)
    (x > f.size[1] || y > f.size[2]) && return true # check for index out of bounds
    f.actual_mines[x, y] && return false # check for mines
    f.visible[x, y] = true
    if f.numbers[x, y] == 0
        for di = -1:1
            for dj = -1:1
                if (0 < di+x <= f.size[1] && 0 < dj+y <= f.size[2]) &&
                        f.actual_mines[x+di, y+dj] == false &&
                        f.visible[x+di, y+dj] == false
                    reveal(f, x+di, y+dj)
                end
            end
        end
    end
    return true
end

function play()
    print("\nWelcome to Minesweeper\n\nEnter the gridsize x  y:\n")
    s = split(readline(), " ")
    f = Field(parse.(Int,s)...)
    won = false
    while true
        printfield(f)
        print("\nWhat do you do? (\"o x y\" to reveal a field; \"m x y\" to toggle a mine; close)\n")
        eval_input(f, readline()) || break
        print("_"^80 * "\n")
        (f.actual_mines == f.possible_mines ||
              f.visible == .!f.actual_mines) && (won = true; break)
    end
    println(won ? "You won the game!" : "You lost the game:\n")
    printfield(f, showall = true)
end

play()


#=
test_tools:
- Julia version:
- Author: Berk
- Date: 2020-06-17
Tests tools.
=#

using PyCall

function alphac_to_expr(model, alpha, c, lse::Bool = false)
    """ Turns exponential to JuMP.NonlinearExpression. """
    n_terms, n_vars = size(alpha)
    idxs = unique([i[2] for i in findall(i->i != 0, alpha)]);
    vars = JuMP.all_variables(model)[idxs];
    if lse
        return :(x -> sum($c[i]*exp(sum($alpha[i,j]*x[j] for j in $idxs)) for i=1:$n_terms)), vars
    else
        return :(x -> sum($c[i]*prod(x[j]^$alpha[i,j] for j in $idxs) for i=1:$n_terms)), vars
    end
end

function alphac_to_varbound!(model, alpha, c, lse::Bool = false)
    """ Turns univariate exponential to JuMP variable bounds. """
    idx = unique([j for j in findall(j->j != 0, alpha)])[1]; #idx[1] is monomial index,
    var = JuMP.all_variables(model)[idx[2]];                 #idx[2] is variable index.
    val = -((sum(c)-c[idx[1]]) / c[idx[1]])^(1/alpha[idx]);
    lse && (val = log(val))
    if c[idx[1]] <= 0
        set_upper_bound(var, val);
    else
        set_lower_bound(var, val);
    end
end

function sagemark_to_GlobalModel(idx, lse::Bool = false)
    """
    Imports sagebenchmarks example from literature.solved and
    returns as a function_model.
    Arguments:
        idx:: number of benchmark
    Returns:
        gm:: a function_model
    """
    sagemarks = pyimport("sagebenchmarks.literature.solved");
    signomials, solver, run_fn = sagemarks.get_example(idx);
    f, greaters, equals = signomials;
    n_vars = size(f.alpha,2);
    model = JuMP.Model()
    @variable(model, x[1:n_vars])
    @variable(model, obj)
    @objective(model, Min, obj)
    # Assigning objective
    gm = GlobalModel(model = model, name = "sagemark" * string(idx))
    obj_fn, objvars = alphac_to_expr(gm.model, f.alpha, f.c, lse)
    add_nonlinear_constraint(gm, obj_fn, vars = objvars, dependent_var = obj)
    # Adding the rest
    for i = 1:length(greaters)
        constrexpr, constrvars = alphac_to_expr(gm.model, greaters[i].alpha, greaters[i].c, lse)
#         if sum(float(greaters[i].c .>= zeros(length(greaters[i].c)))) == 1
#         TODO: add tag for convexity
        if length(constrvars) > 1
            add_nonlinear_constraint(gm, constrexpr, vars = constrvars)
        else
            alpha = greaters[i].alpha
            c = greaters[i].c
            alphac_to_varbound!(gm.model, alpha, c, lse)
        end
    end
    for i = 1:length(equals)
        constrexpr, constrvars = alphac_to_expr(gm.model, equals[i].alpha, equals[i].c, lse)
        add_nonlinear_constraint(gm, constrexpr, vars = constrvars, equality=true)
    end
    set_optimizer(gm, SOLVER_SILENT)
    return gm
end


using DReal

import DReal:  opensmt_init, opensmt_mk_context,
               opensmt_mk_real_var, opensmt_mk_eq, opensmt_assert, opensmt_check,
               opensmt_mk_sin, opensmt_mk_cos, opensmt_get_ub, opensmt_get_lb,
               opensmt_set_precision, opensmt_mk_ite, opensmt_mk_num, opensmt_mk_gt

using Compat
# A = Var(Float64, "A", 0.0, 1.0)
# B = Var(Float64, "B", 0.0, 1.0)
# # C = Var(Float64, "C", 0.0, 1.0)
# add!(B > ifelse(A>0.4, 0.3,0.4))
# add!(A>0.4)
# model(A)
# model(B)

#
ctx = opensmt_mk_context(@compat UInt32(1))
opensmt_set_precision(ctx, 0.00001)
A = opensmt_mk_real_var( ctx, "A" , 0.0, 1.0)
B = opensmt_mk_real_var( ctx, "B" , 0.0, 1.0)

ite = opensmt_mk_ite(ctx, opensmt_mk_gt(ctx, A, opensmt_mk_num(ctx, 0.4)),
                          opensmt_mk_num(ctx, 0.3),
                          opensmt_mk_num(ctx, 0.4))
eq = opensmt_mk_gt(ctx, B, ite)
opensmt_assert( ctx, eq)
res = opensmt_check( ctx );
println("res is ", res == -1 ? "unsat" : "sat")


const OP_BIN_SCALAR_SUM = OptimizingIR.Op(+, pure=true, commutative=true, hasleftidentity=true, hasrightidentity=true, identity_element=0)
const OP_BIN_SCALAR_SUB = OptimizingIR.Op(-, pure=true, hasrightidentity=true, identity_element=0)
const OP_BIN_SCALAR_MUL = OptimizingIR.Op(*, pure=true, commutative=true, hasleftidentity=true, hasrightidentity=true, identity_element=1)
const OP_BIN_SCALAR_DIV = OptimizingIR.Op(/, pure=true, hasrightidentity=true, identity_element=1)
const OP_BIN_SCALAR_MIN = OptimizingIR.Op(min, pure=true, commutative=true)
const OP_BIN_SCALAR_MAX = OptimizingIR.Op(max, pure=true, commutative=true)
const OP_BIN_SCALAR_POW = OptimizingIR.Op(^, pure=true)

const OP_UNARY_SCALAR_SUB = OptimizingIR.Op(-, pure=true)
const OP_UNARY_SCALAR_EXP = OptimizingIR.Op(exp, pure=true)
const OP_UNARY_SCALAR_LOG = OptimizingIR.Op(log, pure=true)
const OP_UNARY_SCALAR_SQRT = OptimizingIR.Op(sqrt, pure=true)

const OP_GETINDEX = OptimizingIR.Op(Base.getindex, pure=true)

#
# Black-Scholes Model
#

const OP_BLACK_FORMULA_CALL = OptimizingIR.Op(BS.bscall, pure=true)
const OP_BLACK_FORMULA_PUT = OptimizingIR.Op(BS.bsput, pure=true)

#
# Binomial Daily Model
#

CRR_volatility_match_up_factor(σ, Δt) = exp(σ * sqrt(Δt))
const OP_CRR_VOLATILITY_MATCH_UP_FACTOR = OptimizingIR.Op(CRR_volatility_match_up_factor, pure=true)

underlying_price(s, n, i, ln_u, ln_d) = s * exp((i-1)*ln_u + (n-i)*ln_d)
const OP_BINOMIAL_UNDERLYING_PRICE_AT = OptimizingIR.Op(underlying_price, pure=true)

function risk_neutral_pricing_step(riskfree_discount_factor, risk_neutral_probability, next_step_price_up, next_step_price_down)
    df = riskfree_discount_factor
    q = risk_neutral_probability
    f_up = next_step_price_up
    f_down = next_step_price_down
    return df * ( q * f_up + (1-q) * f_down )
end
const OP_BINOMIAL_RISK_NEUTRAL_PRICING_STEP = OptimizingIR.Op(risk_neutral_pricing_step, pure=true)

function risk_neutral_probability(d, ert, u)
    @assert d <= ert <= u "No-arbitrage condition was not satisfied: d = $d, ert = $ert, u = $u."
    return ( ert - d ) / ( u - d )
end
const OP_BINOMIAL_RISK_NEUTRAL_PROBABILITY = OptimizingIR.Op(risk_neutral_probability, pure=true)


#
# Minimize total variation of a signal of R3 data.
#
# This example is part of Example 6.1 in the publication
#
# > R. Bergmann, R. Herzog, M. Silva Louzeiro, D. Tenbrinck, J. Vidal Núñez,
# > Fenchel Duality Theory and a Primal-Dual Algorithm on Riemannian Manifolds,
# > arXiv: [1908.02022](https://arxiv.org/abs/1908.02022)
#
using Manopt, Manifolds, LinearAlgebra

#
# Script Settings
experiment_name = "R3_Signal_TV_CP"
export_orig = true
export_primal = true
export_table = true
use_debug = true
#
# Automatic Script Settings
current_folder = @__DIR__
export_any = export_orig || export_primal || export_table
results_folder = joinpath(current_folder, "Signal_TV")
# Create folder if we have an export
(export_any && !isdir(results_folder)) && mkdir(results_folder)

#
# Example Settings
signal_section_size = 15
α = 5.0
σ = 0.5
τ = 0.5
θ = 1.0
γ = 0.0
max_iterations = 500
noise_level = 0.0
noise_type = :Gaussian

pixelM = Euclidean(3)
base = [1.0, 0.0, 0.0]
X = π / 4 * [0.0, 1.0, 0.0]
# Generate a signal with two sections - same signal as in S2_Signal_TV
p1 = exp(Sphere(2), base, X)
p2 = exp(Sphere(2), base, -X)
f = vcat(fill(p1, signal_section_size), fill(p2, signal_section_size))
#
# Compute exact minimizer
include("Signal_TV_commons.jl")
jump_height = distance(pixelM, f[signal_section_size], f[signal_section_size + 1])
δ = min(2 / (size(f, 1) * jump_height) * α, 1 / 2)
x_hat = shortest_geodesic(M, f, reverse(f), δ)
include("Ck.jl")

# add noise
if noise_level > 0
    f = [exp(pixelM, p, random_tangent(pixelM, p, noise_type, noise_level)) for p in f]
end
if export_orig
    orig_file = joinpath(results_folder, experiment_name * "-original.asy")
    asymptote_export_S2_data(orig_file; data=f)
    render_asymptote(orig_file)
end
#
# Initial values
m = fill(base, size(f))
n = Λ(m)
x0 = deepcopy(f)
ξ0 = ProductRepr(zero_tangent_vector(M, m), zero_tangent_vector(M, m))

storage = StoreOptionsAction((:x, :n, :ξbar))

@time o = ChambollePock(
    M,
    N,
    cost,
    x0,
    ξ0,
    m,
    n,
    prox_F,
    prox_G_dual,
    DΛ,
    adjoint_DΛ;
    primal_stepsize=σ,
    dual_stepsize=τ,
    relaxation=θ,
    acceleration=γ,
    relax=:dual,
    variant=:linearized,
    debug=if use_debug
        [
            :Iteration,
            " ",
            DebugPrimalChange(),
            " | ",
            DebugCk(storage),
            " | ",
            :Cost,
            "\n",
            100,
            :Stop,
        ]
    else
        missing
    end,
    record=if export_table
        [:Iteration, RecordPrimalChange(x0), RecordDualChange((ξ0, n)), :Cost, RecordCk()]
    else
        missing
    end,
    stopping_criterion=StopAfterIteration(max_iterations),
    return_options=true,
)
y = get_solver_result(o)
if has_record(o)
    r = get_record(o)
    Ck_values = [s[5] for s in r]
    println("The Ck Estimate lies between $(minimum(Ck_values)) and $(maximum(Ck_values))")
end
println("Distance from result to minimizer: ", distance(M, x_hat, y), "\n")

if export_primal
    orig_file = joinpath(results_folder, experiment_name * "-result.asy")
    asymptote_export_S2_data(orig_file; data=y)
    render_asymptote(orig_file)
end


module Diagnostics

    using GeometricBase

    export compute_one_form, compute_invariant, compute_invariant_error, compute_momentum_error, compute_error_drift

    """
    Takes a 1d DataSeries holding an invariant and computes the relative error `(inv(t)-inv(0))/inv(0)`.

    Returns a 1d DataSeries similar to the argument holding the time series of the relativ errors.
    """
    function compute_relative_error(invds::AbstractDataSeries{T,1}) where {T}
        errds = similar(invds)
        for i in eachindex(errds, invds)
            errds[i] = (invds[i] - invds[0]) / invds[0]
        end
        return errds
    end

    """
    Compute the one-form (symplectic potential) for the solution of a Lagrangian system.

    Arguments: `(t::TimeSeries, q::AbstractDataSeries, one_form::Function)`

    The `one_form` function needs to take three arguments `(t,q,k)` where `k` is the index
    of the one-form component.

    Returns a DataSeries similar to `q` holding the time series of the one-form.
    """
    function compute_one_form(t::TimeSeries, q::AbstractDataSeries, one_form::Function)
        ϑ = similar(q)
        try
            for i in axes(p,2)
                for k in axes(p,1)
                    ϑ[k,i] = one_form(t[i], q[:,i], k)
                end
            end
        catch ex
            if isa(ex, DomainError)
                @warn("DOMAIN ERROR: One-form diagnostics crashed.")
            else
                throw(ex)
            end
        end
        return ϑ
    end

    """
    Compute an invariant for the solution of an ODE or DAE system.

    Arguments: `(t::TimeSeries, q::AbstractDataSeries{T}, invariant::Function)`

    The `invariant` functions needs to take two arguments `(t,q)` and return the
    corresponding value of the invariant.

    Returns a 1d DataSeries holding the time series of the invariant.
    """
    function compute_invariant(t::TimeSeries, q::AbstractDataSeries{<:AbstractArray{T}}, invariant::Function) where {T}
        invds = DataSeries(T, ntime(q))
        try
            for i in eachindex(invds)
                invds[i] = invariant(t[i], q[i])
            end
        catch ex
            if isa(ex, DomainError)
                @warn("DOMAIN ERROR: Invariant diagnostics crashed.")
            else
                throw(ex)
            end
        end
        return invds
    end

    """
    Compute an invariant for the solution of a partitioned ODE or DAE system.

    Arguments: `(t::TimeSeries, q::AbstractDataSeries{T}, p::AbstractDataSeries{T}, invariant::Function)`

    The `invariant` functions needs to take three arguments `(t,q,p)` and return the
    corresponding value of the invariant.

    Returns a 1d DataSeries holding the time series of the invariant.
    """
    function compute_invariant(t::TimeSeries, q::AbstractDataSeries{<:AbstractArray{T}}, p::AbstractDataSeries{<:AbstractArray{T}}, invariant::Function) where {T}
        invds = DataSeries(T, ntime(q))
        try
            for i in eachindex(invds)
                invds[i] = invariant(t[i], q[i], p[i])
            end
        catch ex
            if isa(ex, DomainError)
                @warn("DOMAIN ERROR: Invariant diagnostics crashed.")
            else
                throw(ex)
            end
        end
        return invds
    end

    """
    Compute the relative error of an invariant for the solution of an ODE or DAE system.

    Arguments: `(t::TimeSeries, q::AbstractDataSeries{T}, invariant::Function)`

    The `invariant` functions needs to take two arguments `(t,q)` and return the
    corresponding value of the invariant.

    Returns a tuple of two 1d DataSeries holding the time series of the invariant and the relativ error, respectively.
    """
    function compute_invariant_error(t::TimeSeries, q::AbstractDataSeries{T}, invariant::Function) where {T}
        invds = compute_invariant(t, q, invariant)
        errds = compute_relative_error(invds)
        (invds, errds)
    end

    """
    Compute the relative error of an invariant for the solution of a partitioned ODE or DAE system.

    Arguments: `(t::TimeSeries, q::AbstractDataSeries{T}, p::AbstractDataSeries{T}, invariant::Function)`

    The `invariant` functions needs to take three arguments `(t,q,p)` and return the
    corresponding value of the invariant.

    Returns a tuple of two 1d DataSeries holding the time series of the invariant and the relativ error, respectively.
    """
    function compute_invariant_error(t::TimeSeries, q::AbstractDataSeries{T}, p::AbstractDataSeries{T}, invariant::Function) where {T}
        invds = compute_invariant(t, q, p, invariant)
        errds = compute_relative_error(invds)
        (invds, errds)
    end

    """
    Computes the difference of the momentum and the one-form of an implicit ODE or DAE system.

    Arguments: `(t::TimeSeries, q::AbstractDataSeries{T}, p::AbstractDataSeries{T}, one_form::Function)`

    The `one_form` function needs to take three arguments `(t,q,k)` where `k` is the index
    of the one-form component.

    Returns a DataSeries similar to `p` holding the time series of the difference between the momentum and the one-form.
    """
    function compute_momentum_error(t::TimeSeries, q::AbstractDataSeries{T}, p::AbstractDataSeries{T,2}, one_form::Function) where {T}
        e = similar(p)

        for i in axes(p,2)
            for n in axes(p,1)
                for k in eachindex(p[n,i])
                    e[n,i][k] = p[n,i][k] - one_form(t[n], q[n,i], k)
                end
            end
        end

        return e
    end

    function compute_momentum_error(t::TimeSeries, q::AbstractDataSeries{T}, p::AbstractDataSeries{T,1}, one_form::Function) where {T}
        e = similar(p)

        for n in axes(p,1)
            for k in eachindex(p[n])
                e[n][k] = p[n][k] - one_form(t[n], q[n], k)
            end
        end

        return e
    end

    """
    Computes the difference of the momentum and the one-form of an implicit ODE or DAE system.

    Arguments: `(p::AbstractDataSeries{DT}, ϑ::AbstractDataSeries{DT})`

    Returns a DataSeries similar to `p` holding the time series of the difference between `p` and `ϑ`.
    """
    function compute_momentum_error(p::AbstractDataSeries{DT}, ϑ::AbstractDataSeries{DT}) where {DT}
        @assert axes(p) == axes(ϑ)

        e = similar(p)
        parent(e) .= parent(p) .- parent(ϑ)

        return e
    end

    """
    Computes the drift in an invariant error.

    Arguments: `(t::TimeSeries, invariant_error::AbstractDataSeries{T,1}, interval_length=100)`

    The time series of the solution is split into intervals of `interval_length` time steps.
    In each interval, the maximum of the absolute value of the invariant error is determined.
    Returns a tuple of a TimeSeries that holds the centers of all intervals and a 1d DataSeries
    that holds the maxima.

    This is useful to detect drifts in invariants that are not preserved exactly but whose error
    is oscillating such as the energy error of Hamiltonian systems with symplectic integrators.
    """
    function compute_error_drift(t::TimeSeries, invariant_error::AbstractDataSeries{T,1}, interval_length=100) where {T}
        @assert ntime(t) == ntime(invariant_error)
        @assert mod(t.n, interval_length) == 0

        nint   = div(ntime(invariant_error), interval_length)
        Tdrift = TimeSeries(nint, (t[end] - t[begin]) / nint)
        Idrift = DataSeries(T, nint)

        Tdrift[0] = t[0]

        for i in 1:nint
            i1 = interval_length*(i-1)+1
            i2 = interval_length*i
            Idrift[i] = maximum(abs.(invariant_error[i1:i2]))
            Tdrift[i] = div(t[i1] + t[i2], 2)
        end

        (Tdrift, Idrift)
    end

end


include("./util.jl")

depths = map(l -> parse(Int, l), util.input_lines(1))

changes = depths[2:end] - depths[1:end-1]

num_increases = sum(map(d -> d > 0, changes))

println("Part 1: $num_increases")

windows = []

for i = 1:length(depths)
    if i + 3 <= length(depths)
        push!(windows, depths[i:i+3])
    end
end

window_sums = map(sum, windows)
window_changes = window_sums[2:end] - window_sums[1:end-1]
num_window_increases = sum(map(d -> d > 0, window_changes))

println("Part 2: $num_window_increases")


using OptiMo, Bazinga, ScSTO
using Printf
using PyPlot, PyCall

################################################################################
################################################################################
# unconstrained STO
################################################################################
################################################################################
# time interval
t0 = 0.0
tf = 12.0

# control sequence [nu x N]
uvec = Matrix([repeat([0.0; 1.0], 5, 1); 0.0]')
nu = size(uvec, 1)
N = size(uvec, 2)

# cost matrix
C = Matrix([1.0 0.0 -1.0 0.0; 0.0 1.0 0.0 -1.0])
Q = Matrix(C' * C)

# initial state
x0 = [0.5; 0.7; 1; 1] # nx x 1
nx = length(x0)

# system dynamics
# nx, nu  ->  nx
# nx, nu  -> nx x nx
function dynam(x::Vector{Float64}, u::Vector{Float64})
    n = length(x)
    f = zeros(n)
    f[1] = x[1] - x[1] * x[2] - 0.4 * x[1] * u[1]
    f[2] = -x[2] + x[1] * x[2] - 0.2 * x[2] * u[1]
    return f
end
function d_dynam(x::Vector{Float64}, u::Vector{Float64})
    n = length(x)
    df = zeros(n, n)
    df[1, 1] = 1.0 - x[2] - 0.4 * u[1]
    df[1, 2] = -x[1]
    df[2, 1] = x[2]
    df[2, 2] = -1.0 + x[1] - 0.2 * u[1]
    return df
end

# switching cost
swc = 0.1

# number of points on the fixed grid
ngrid_grid = [100; 200; 1000]
ng = length(ngrid_grid)

# solver
solver = Bazinga.ZEROFPR( max_iter = 50, verbose = true )

# allocation
nt = 2000;
swdelta = Array{Float64}(undef, N, ng);
swtau = Array{Float64}(undef, N - 1, ng);
xsim = Array{Float64}(undef, nx, nt, ng);
usim = Array{Float64}(undef, nu, nt, ng);

for k = 1:ng

    ngrid = ngrid_grid[k]

    prob = scstoproblem(
        x0,
        dynam,
        d_dynam,
        uvec,
        ngrid = ngrid,
        t0 = t0,
        tf = tf,
        Q = Q,
        swc = swc,
    )

    out = solver(prob)
    print(out)

    swdelta[:, k] .= out.x
    swtau[:, k], _ = gettau(prob, swdelta[:, k])

    xsim[:, :, k], _, _, t = simulate(prob, swtau[:, k], length = nt)
    usim[:, :, k], _ = simulateinput(prob, swtau[:, k], t)

    @printf "\n\n"

end

t = collect(range(t0, stop = tf, length = nt))

figure()
subplot(3, 1, 1)
for k = 1:ng
    plot(t, xsim[1, :, k])
end
ylim(0, 1.75)
xlim(0, 12)
yticks([0; 1])
ylabel(L"x_1")
subplot(3, 1, 2)
for k = 1:ng
    plot(t, xsim[2, :, k])
end
ylabel(L"x_2")
ylim(0, 1.75)
xlim(0, 12)
yticks([0; 1])
subplot(3, 1, 3)
for k = 1:ng
    plot(t, usim[1, :, k], label = "n = $(ngrid_grid[k])")
end
legend()
ylim(-0.2, 1.2)
yticks([0; 1])
xlim(0, 12)
ylabel(L"u")
xlabel(L"$\mathrm{Time}\; [s]$")
gcf()
#savefig(foldername * "ScSTO.jl/test/data/fishing_ngrid.pdf")

################################################################################
# make switching times vectors comparable and compare results

using LinearAlgebra

function cleandelta(delta::Vector{Float64})
    # !!! for binary controls only !!!
    # uvec = 0, 1, 0, 1, 0, 1, ...
    deltanew = copy( delta )
    n = length( delta )
    for k in 1:n-2
        if deltanew[k] > 0.0 && deltanew[k+1] == 0.0
            deltanew[k+2] += deltanew[k]
            deltanew[k] = 0.0
        end
    end
    return deltanew
end
function cleandelta(delta::Matrix{Float64})
    deltanew = copy( delta )
    m = size( delta,2 )
    for k in 1:m
        deltanew[:,k] .= cleandelta( delta[:,k] )
    end
    return deltanew
end

swdelta_clean = cleandelta( swdelta )
global swdelta_diff = Array{Float64}(undef, ng-1);
for k in 1:ng-1
    swdelta_diff[k] = norm( swdelta_clean[:,k] - swdelta_clean[:,ng] , Inf )
end


#=
Popis
  Funkcia vypocita priemernu hodnotu z najkratsich ciest v grafe.

Parametre
  g : datova struktura grafu :: GenericGraph
  dist_key : nazov atributu, v ktorom je ulozene ohodnotenie hrany :: String
  components : parameter pre prepinanie graf/komponent :: Bool

Navratova hodnota
  avg_distance : priemerna hodnota najkratsej cesty pre cely graf :: Float64
  avg_distance : slovnik s mapovanim (komponent_grafu => avg_distance) :: Dict{Array{ExVertex,1},Float64}

Poznamky
  Priemernu hodnotu najkratsej cesty je mozne nastavenim parametra components vypocitat bud pre cely graf
  alebo pre jednotlive komponenty.
=#

function average_shortest_distance(g::GenericGraph; dist_key::AbstractString="weight", components::Bool=false)

  epi = AttributeEdgePropertyInspector{Float64}(dist_key)

  if components

    comp = connected_components(g)
    avg_distance = [v => 0.0 for v = comp]
    paths = Float64[]

    for vector in comp

      for v in vector

        paths = [filter(x -> x != Inf, dijkstra_shortest_paths(g, epi, [v]).dists) for v = vector ]

      end

      avg_distance[vector] = sum([sum(x) for x = paths]) / (length(vector)*(length(vector) - 1))
      empty!(paths)

    end

    return avg_distance

  elseif !(components)

    avg_distance = 0.0
    paths = [filter(x -> x != Inf, dijkstra_shortest_paths(g, epi, [v]).dists) for v = vertices(g) ]
    avg_distance = sum([sum(x) for x = paths]) / (length(vertices(g))*(length(vertices(g)) - 1))

    return avg_distance

  end

end


## get nonstat stuff first
close("all")
using Statistics
M = transD_GP.assembleTat1(opt, :fstar, temperaturenum=1)
mns = reshape(mean(M), length(img.y), length(img.x))
stdns = reshape(std(M), length(img.y), length(img.x))
imshow(mns, extent=[img.x[1],img.x[end],img.y[end],img.y[1]])
colorbar()

##
close("all")
rownum, colnum = 195, 85
transD_GP.plot_image_posterior(opt, optlog10λ, img, burninfrac=0.5, rownum=rownum, colnum=colnum, nbins=100)
figure(1)
ax = gcf().axes
ax[7].plot(img.x, img.f[rownum,:], "--y", alpha=0.9)
savefig("post_rows_ns.png", dpi=300)
figure(2)
ax = gcf().axes
ax[7].plot(img.f[:,colnum], img.y,"--y", alpha=0.9)
savefig("post_col_ns.png", dpi=300)
# sampling stats
## plot stats
close("all")
transD_GP.getchi2forall(opt, fsize=8, alpha=0.5)
gcf().text(0.02, 0.9, "a.", fontsize=14, color="red")
linidx = .!isnan.(img.d)
r = vec(img.d[linidx] - img.f[linidx])
χ² = length(r)
figure(1)
ax = gcf().axes;
ax[2].set_ylim(χ²/2 - 100, χ²/2 + 300)
ax[2].plot(xlim(), [χ²/2 , χ²/2], "--", color="gray")
figure(2)
ax = gcf().axes
ax[2].set_ylim(χ²/2 - 100, χ²/2 + 300)
ax[2].plot(xlim(), [χ²/2 , χ²/2], "--", color="gray")
savefig("img_conv_ns_1.png", dpi=300)
##
close("all")
transD_GP.getchi2forall(optlog10λ, fsize=8, alpha=0.5)
gcf().text(0.02, 0.9, "b.", fontsize=14, color="red")
figure(1)
ax = gcf().axes
ax[2].set_ylim(χ²/2 - 100, χ²/2 + 300)
ax[2].plot(xlim(), [χ²/2 , χ²/2], "--", color="gray")
figure(2)
ax = gcf().axes
ax[2].set_ylim(χ²/2 - 100, χ²/2 + 300)
ax[2].plot(xlim(), [χ²/2 , χ²/2], "--", color="gray")
savefig("img_conv_ns_2.png", dpi=300)


function evalmaximum(f, intervals::Vector{Tuple{arb, arb}};
                     evaltype = :ball,
                     n = 4,
                     absmax = false)
    evals = Vector{Tuple{arb, arb}}(undef, length(intervals))
    if evaltype == :ball
        maybeabs = ifelse(absmax, abs, identity)
        Threads.@threads for i in 1:length(intervals)
            evals[i] = getinterval(maybeabs(f(setinterval(intervals[i]...))))
        end
    elseif evaltype == :taylor
        Threads.@threads for i in 1:length(intervals)
            evals[i] = getinterval(maximumtaylor(f, intervals[i], n, absmax = absmax))
        end
    end

    evals
end

function enclosemaximum(f,
                        a::arb,
                        b::arb;
                        # A priori computed lower bound of maximum
                        lower_bound::arb = neginf(a),
                        absmax = false,
                        evaltype = :ball,
                        n = 4,
                        # FIXME: This should be epsilon in the precision of the input parameters
                        atol = sqrt(eps()),
                        rtol = sqrt(eps()),
                        maxevals = 10^3,
                        store_trace = false,
                        show_trace = false,
                        callback = nothing,
                        extended_trace = false)
    @assert isfinite(a) && isfinite(b) && a < b

    intervals = [(a, b)]
    maxenclosure = setinterval(neginf(a), posinf(a))
    E = radius(maxenclosure)
    numevals = 0
    iteration = 0
    imprecise_input_last_iteration = false
    imprecise_input = false

    trace = MaximumTrace()
    if show_trace
        print(trace)
    end

    while (!(E < atol)
           && !(E/abs(maxenclosure) <= rtol)
           && numevals <= maxevals
           && !imprecise_input)
        iteration += 1

        # Lower and upper bound maximum on each interval
        evals = evalmaximum(f, intervals, evaltype = evaltype, n = n, absmax = absmax)
        numevals += length(intervals)

        # Find global lower and upper bound for maximum
        maxlower = neginf(a)
        maxupper = neginf(a)
        for (lower, upper) in evals
            if isfinite(lower) && !(lower < maxlower)
                maxlower = max(maxlower, lower)
            end

            if !(upper < maxupper)
                maxupper = max(maxupper, upper)
            end
        end

        if maxupper < lower_bound
            @error "given lower bound is larger than maximum"
            return parent(a)(NaN)
        end

        maxenclosure = setinterval(maxlower, maxupper)

        maxlower = max(maxlower, lower_bound)

        # Split intervals where the maximum could be located
        nextintervals = Vector{eltype(intervals)}()
        for i in 1:length(intervals)
            if !(evals[i][2] < maxlower)
                midpoint = 0.5*(intervals[i][1] + intervals[i][2])
                push!(nextintervals, (intervals[i][1], midpoint))
                push!(nextintervals, (midpoint, intervals[i][2]))
            end
        end

        # Update trace
        if !extended_trace
            dt = Dict()
        else
            dt = Dict("intervals" => intervals, "evaluations" => evals)
        end
        update!(trace,
                MaximumState(iteration, length(intervals), maxenclosure, dt),
                store_trace,
                show_trace,
                callback)

        intervals = nextintervals

        # If using Taylor series and no substantial improvement is
        # made from previous iteration this usually hints at working
        # with to low precision
        if evaltype == :taylor && isfinite(maxenclosure) && radius(maxenclosure) > 0.8*E
            if imprecise_input_last_iteration
                imprecise_input = true
            end
            imprecise_input_last_iteration = true
        else
            imprecise_input_last_iteration = false
        end

        E = radius(maxenclosure)
    end

    if imprecise_input && !(E < atol || E/abs(maxenclosure) <= rtol || !(numevals <= maxevals))
        @warn "Maximum likely needs to be computed with higher precision than $(precision(parent(a)))"
    end

    if store_trace
        return maxenclosure, trace
    end

    maxenclosure
end


@testset "Filter tests" begin
  using Test
  using CLIMA
  import GaussQuadrature
  using MPI
  using LinearAlgebra

  let
    # Values computed with:
    #   https://github.com/tcew/nodal-dg/blob/master/Codes1.1/Codes1D/Filter1D.m
    W = [0x3fe98f3cd0d725e8   0x3fddfd863c6c9a44   0xbfe111110d0fd334   0x3fddbe357bce0b5c   0xbfc970267f929618
         0x3fb608a150f6f927   0x3fe99528b1a1cd8d   0x3fcd41d41f8bae45   0xbfc987d5fabab8d5   0x3fb5da1cd858af87
         0xbfb333332eb1cd92   0x3fc666666826f178   0x3fe999999798faaa   0x3fc666666826f176   0xbfb333332eb1cd94
         0x3fb5da1cd858af84   0xbfc987d5fabab8d4   0x3fcd41d41f8bae46   0x3fe99528b1a1cd8e   0x3fb608a150f6f924
         0xbfc970267f929618   0x3fddbe357bce0b5c   0xbfe111110d0fd333   0x3fddfd863c6c9a44   0x3fe98f3cd0d725e8]
    W = reinterpret.(Float64, W)

    N = size(W, 1) - 1

    topology = CLIMA.Mesh.Topologies.BrickTopology(MPI.COMM_SELF, -1.0:2.0:1.0);

    grid = CLIMA.Mesh.Grids.DiscontinuousSpectralElementGrid(topology;
                                                        polynomialorder = N,
                                                        FloatType = Float64,
                                                        DeviceArray = Array)

    filter = CLIMA.Mesh.Filters.ExponentialFilter(grid, 0, 32)
    @test filter.filter ≈ W
  end

  let
    # Values computed with:
    #   https://github.com/tcew/nodal-dg/blob/master/Codes1.1/Codes1D/Filter1D.m
    W = [0x3fd822e5f54ecb62   0x3fedd204a0f08ef8   0xbfc7d3aa58fd6968   0xbfbf74682ac4d276
         0x3fc7db36e726d8c1   0x3fe59d16feee478b   0x3fc6745bfbb91e20   0xbfa30fbb7a645448
         0xbfa30fbb7a645455   0x3fc6745bfbb91e26   0x3fe59d16feee478a   0x3fc7db36e726d8c4
         0xbfbf74682ac4d280   0xbfc7d3aa58fd6962   0x3fedd204a0f08ef7   0x3fd822e5f54ecb62]
    W = reinterpret.(Float64, W)

    N = size(W, 1) - 1

    topology = CLIMA.Mesh.Topologies.BrickTopology(MPI.COMM_SELF, -1.0:2.0:1.0);
    grid = CLIMA.Mesh.Grids.DiscontinuousSpectralElementGrid(topology;
                                                        polynomialorder = N,
                                                        FloatType = Float64,
                                                        DeviceArray = Array)

    filter = CLIMA.Mesh.Filters.ExponentialFilter(grid, 1, 4)
    @test filter.filter ≈ W
  end

  let
    T = Float64
    N = 5
    Nc = 4

    topology = CLIMA.Mesh.Topologies.BrickTopology(MPI.COMM_SELF, -1.0:2.0:1.0);
    grid = CLIMA.Mesh.Grids.DiscontinuousSpectralElementGrid(topology;
                                                        polynomialorder = N,
                                                        FloatType = T,
                                                        DeviceArray = Array)

    ξ = CLIMA.Mesh.Grids.referencepoints(grid)
    a, b = GaussQuadrature.legendre_coefs(T, N)
    V = GaussQuadrature.orthonormal_poly(ξ, a, b)

    Σ = ones(T, N+1)
    Σ[(Nc:N).+1] .= 0

    W = V*Diagonal(Σ)/V

    filter = CLIMA.Mesh.Filters.CutoffFilter(grid, Nc)
    @test filter.filter ≈ W
  end
end


function epsilon_dominates_clear(u, v, epsilon)
  veps = v .+ epsilon
  all(u .<= veps) && any(u .< veps)
end

function epsilon_dominates_fast(u, v, epsilon)
  pareto_dominates_fast(u, v .+ epsilon)
end

function pareto_dominates_clear(u, v)
  all(u .<= v) && any(u .< v)
end

function pareto_dominates_fast(u, v)
  one_smaller = false
  for i in 1:length(u)
    if u[i] > v[i]
      return false
    elseif u[i] < v[i]
      one_smaller = true
    end
  end
  return one_smaller
end

# Returns two booleans where the first indicates e-box dominance and the 2nd indicates
# epsilon progress
function epsilon_box_dominates_and_epsilon_progress(u, v, epsilon)
  uindex = floor(u ./ epsilon)
  vindex = floor(v ./ epsilon)
  inner_epsilon_box_dominates_and_epsilon_progress(
    u, uindex, (epsilon .* uindex),
    v, vindex, (epsilon .* vindex))
end

epsilon_box_dominates(u, v, epsilon) = epsilon_box_dominates_and_epsilon_progress(u, v, epsilon)[1]
epsilon_box_progress(u, v, epsilon) = epsilon_box_dominates_and_epsilon_progress(u, v, epsilon)[2]

# Inner boolean expression when the intermediate results have already been calculated.
# This is to speed up processing when there is a whole archive of solutions to compare
# a new solution to.
function inner_epsilon_box_dominates_and_epsilon_progress(
  u, uindex, uindextimesepsilon,
  v, vindex, vindextimesepsilon)

  if pareto_dominates_fast(uindex, vindex)
    return (true, true)
  elseif all(uindex .== vindex)
    smaller_distance = norm(u .- uindextimesepsilon) < norm(v .- vindextimesepsilon)
    return (smaller_distance, false)
  else
    return (false, false)
  end

end


###############################
# kmedoids, Voronoi iteration #
###############################

function kmedoids_voronoi(D::Matrix{Float}, k::Int;
                    verbose::Bool = true,
                    epsi::Float = Float(1e-10),
                    avg_iters::Int = Int(1),
                    max_iters::Int = Int(1000),
                    calc_emp::Bool = true,
                    medoid_init::Vector{Int} = Int[],
                    init::Symbol = :per)

  n = size(D,1)

  # initialize optimization
  opt = Float[]
  queue = Array{Float, 1}()
  diff_ratio = Float(1)
  avg_objective = Float(0)
  medoid_idx = zeros(Int, k)

  max_val = maximum(D)
  if verbose
    println("initialization...")
  end
  if init == :rand
    init_idx = 1:n
  elseif init == :cen
    ~, init_idx = findmin(max(D,max_val*eye(n)),2)
    ~, init_idx = ind2sub(size(D), init_idx[:])
  elseif init == :per
    ~, init_idx = findmax(D,2)
    ~, init_idx = ind2sub(size(D), init_idx[:])
  else
    error("no such initialization")
  end
  medoid_idx = (length(medoid_init) == k) ? medoid_init : sample(init_idx, k, replace = false)
  ~,clust_idx = findmin(D[:, medoid_idx],2)

  if verbose
    @printf("%-20s %-20s %-20s\n", "iteration", "objective", "difference ratio");
  end
  local clust_idx
  local objective
  local min_clust_idx, min_medoid_idx, t
  min_obj = 1.
  mu = mean(D)
  for t in 0:max_iters

    if t > 0
      # find new medoids
      for c in 1:k
        cidx = find(clust_idx .== c)
        medoid_idx[c] = (length(cidx) > 0) ? cidx[indmin(mean(D[cidx, cidx],1))] : 0
      end
      # reassign empty medoids
      empt_ind = medoid_idx .== 0
      empt_num = sum(empt_ind)
      if empt_num > 0
        medoid_idx[empt_ind] = sample(init_idx, empt_num, replace = (length(init_idx) <= empt_num))
      end
    end

    # nearest cluster
    min_dist,clust_idx = findmin(D[:, medoid_idx],2)
    ~, clust_idx = ind2sub(size(D), clust_idx[:])
    # objective
    objective = mean(fixnan.(min_dist./mu))

    if objective < min_obj
      min_obj = objective
      min_clust_idx = clust_idx
      min_medoid_idx = medoid_idx
    end

    if t > 0
      # finalization
      Base.push!(opt, objective)

      if t > avg_iters
        diff_ratio = abs(objective - avg_objective)/fixdiv(avg_objective);
      end

      Base.push!(queue, objective/avg_iters)
      avg_objective = avg_objective + objective/avg_iters
      if t > avg_iters
        avg_objective = avg_objective - shift!(queue)
      end
    end

    if verbose
      if t > avg_iters
        @printf("%-20d %-20.6f %-20.6f\n", t, objective, diff_ratio)
      else
        @printf("%-20d %-20.6f %-20s\n", t, objective, "-")
      end
    end

    if (diff_ratio < epsi)
      break
    end

  end

  return min_clust_idx, min_medoid_idx, min_obj, t
end

#################
# kmedoids, PAM #
#################

function kmedoids_pam(D::Matrix{Float}, k::Int;
                    verbose::Bool = true,
                    avg_iters::Int = Int(10),
                    max_iters::Int = Int(1000),
                    min_obj::Float = eps(Float),
                    calc_emp::Bool = true,
                    medoid_init::Vector = Int[],
                    init::Symbol = :per)

  n = size(D,1)

  # initialize optimization
  opt = Float[]
  queue = Array{Float, 1}()
  diff_ratio = Float(1)
  avg_objective = Float(0)
  medoid_idx = zeros(Int, k)

  max_val = maximum(D)
  if length(medoid_init) != k
    if verbose
      println("initialization...")
    end
    if init == :rand
      init_idx = 1:n
    elseif init == :cen
      ~, init_idx = findmin(max(D,max_val*eye(n)),2)
      ~, init_idx = ind2sub(size(D), init_idx[:])
    elseif init == :per
      ~, init_idx = findmax(D,2)
      ~, init_idx = ind2sub(size(D), init_idx[:])
    else
      error("no such initialization")
    end
    medoid_idx = sample(init_idx, k, replace = false)
  else
    medoid_idx = deepcopy(medoid_init)
  end
  mu = mean(D)
  local clust_idx
  local objective
  new_min_dist = zeros(n)

  if verbose
    @printf("%-20s %-20s\n", "iteration", "objective");
  end

  total_samp = 0
  best_samp = 0
  switched = true
  t = 0
  while (switched) && (t<=max_iters)
    t += 1

    ###############
    # build phase #
    ###############

    # nearest mediod
    min_dist,clust_idx = findmin(D[:, medoid_idx],2)
    ~, clust_idx = ind2sub(size(D), clust_idx[:])
    # objective
    objective = fixnan(mean(min_dist)./mu)
    total_samp += 1
    best_samp = total_samp

    if verbose
      if t > avg_iters
        @printf("%-20d %-20.6f\n", t, objective)
      else
        @printf("%-20d %-20s\n", t, objective)
      end
    end

    if objective < min_obj
      break
    end

    ##############
    # swap phase #
    ##############

    # find new medoids
    switched = false
    for i in randperm(k)
      new_min_dist[:] = minimum(D[:, medoid_idx[1:k .!= i]],2)
      for j in setdiff(randperm(n), medoid_idx)
        new_objective = fixnan(mean(min.(new_min_dist, D[:, j]))./mu)
        total_samp += 1
        if new_objective < objective
          switched = true
          medoid_idx[i] = j
        end
        if switched
          break
        end
      end
      if switched
        break
      end
    end
  end

  return clust_idx, medoid_idx, objective, best_samp, total_samp
end

nothing


using CompScienceMeshes
using Test

vrts = [
    point(0.0, 0.0, 0.0),
    point(1.0, 0.0, 0.0),
    point(0.0, 1.0, 0.0),
]

idcs = [
    index(1,2,3)
]

axis = 0.5π * [1,0,0]
G1 = Mesh(vrts, idcs)
G2 = rotate(G1, axis)
G3 = rotate(G1, 2*axis)

G = weld(G1, G2);
G = weld(G1, G2, G3)

@test numcells(G) == numcells(G1) + numcells(G2) + numcells(G3)
@test numvertices(G) == 5

edges = skeleton(G,1)
pairs = cellpairs(G, edges, dropjunctionpair=true)

@test size(pairs) == (2,8)

mesh = meshrectangle(1.0, 1.0, 0.5)
edges = skeleton(mesh,1)
cps = cellpairs(mesh, edges)
@test size(cps) == (2,16)

# width, height = 1.0, 0.5; h = 0.1
# mesh = meshrectangle(width, height, h)
# edges = skeleton(mesh,1)
# cps = cellpairs(mesh, edges, dropjunctionpair=true)


using Statistics: mean
using Oceananigans.Grids: Center, Face
using Oceananigans.Fields: Field, location
using Oceananigans.OutputReaders: FieldTimeSeries, FieldDataset

"""
    coarse_grain(Φ, n, ::Center)

Average or coarse grain a `Center`-centered field `Φ` down to size `n`. `Φ` is required to have evenly spaced points and `n` needs to evenly divide `length(Φ)`.
"""
function coarse_grain(Φ, n, ::Center)
    N = length(Φ)
    Δ = Int(N / n)
    Φ̅ = similar(Φ, n)
    for i in 1:n
        Φ̅[i] = mean(Φ[Δ*(i-1)+1:Δ*i])
    end
    return Φ̅
end

"""
    coarse_grain(Φ, n, ::Face)

Average or coarse grain a `Face`-centered field `Φ` down to size `n`. `Φ` is required to have evenly spaced points. The values at the left and right endpoints of `Φ` will be preserved in the output.
"""
function coarse_grain(Φ, n, ::Face)
    N = length(Φ)
    Φ̅ = similar(Φ, n)
    Δ = (N-2) / (n-2)

    Φ̅[1] = Φ[1]
    Φ̅[n] = Φ[N]

    if isinteger(Δ)
        Φ̅[2:n-1] .= coarse_grain(Φ[2:N-1], n-2, Center())
    else
        for i in 2:n-1
            i1 = 2 + (i-2)*Δ
            i2 = 2 + (i-1)*Δ

            # Like modf but with ::Int integer part.
            f1, i1 = rem(i1, 1), trunc(Int, i1)
            f2, i2 = rem(i2, 1), trunc(Int, i2)

            left_contrib = (1 - f1) * Φ[i1]
            right_contrib = f2 * Φ[i2]
            middle_contrib = sum(Φ[i1+1:i2-1])

            Φ̅[i] = (left_contrib + middle_contrib + right_contrib) / Δ
        end
    end

    return Φ̅
end

@deprecate coarse_grain(Φ, n, loc::Type{Center}) coarse_grain(Φ, n, loc::Center)
@deprecate coarse_grain(Φ, n, loc::Type{Face}) coarse_grain(Φ, n, loc::Face)

"""
    coarse_grain(field::Field, new_grid; dims=3)

Coarse grain a `field` onto a `new_grid` along `dims`. Returns a new `Field`.
"""
function coarse_grain(field::Field{X, Y, Center}, new_grid; dims=3) where {X, Y}

    # TODO: Generalize to x and y.
    @assert dims == 3
    @assert new_grid.Nx == 1
    @assert new_grid.Ny == 1

    # TODO: Generalize `coarse_grain` to non-integer ratios.
    r = field.grid.Nz / new_grid.Nz
    @assert isinteger(r)
    r = Int(r)

    coarse_field = Field(location(field)..., field.architecture, new_grid, field.boundary_conditions)

    coarse_data = zeros(size(coarse_field))
    field_interior = interior(field)[1, 1, :]
    coarse_data[1, 1, :] .= OceanParameterizations.coarse_grain(field_interior, new_grid.Nz, Center())

    set!(coarse_field, coarse_data)

    return coarse_field
end

function coarse_grain(field::Field{X, Y, Face}, new_grid; dims=3) where {X, Y}

    # TODO: Generalize to x and y.
    @assert dims == 3
    @assert new_grid.Nx == 1
    @assert new_grid.Ny == 1

    r = field.grid.Nz / new_grid.Nz

    coarse_field = Field(location(field)..., field.architecture, new_grid, field.boundary_conditions)

    coarse_data = zeros(size(location(coarse_field), new_grid))
    field_interior = interior(field)[1, 1, :]
    coarse_data[1, 1, :] .= OceanParameterizations.coarse_grain(field_interior, new_grid.Nz+1, Face())

    set!(coarse_field, coarse_data)

    return coarse_field
end

function coarse_grain(fts::FieldTimeSeries, new_grid)
    fts_new = FieldTimeSeries(new_grid, location(fts), fts.times; ArrayType=Array{Float32})

    Nt = size(fts, 4)
    for n in 1:Nt
        fts_new.data[:, :, :, n] .= coarse_grain(fts[n], new_grid).data
    end

    return fts_new
end

function coarse_grain(fds::FieldDataset, new_grid)
    coarse_fields = Dict{String, FieldTimeSeries}(
        name => coarse_grain(fts, new_grid)
        for (name, fts) in fds.fields
    )
    return FieldDataset(coarse_fields, fds.metadata, fds.filepath)
end


function test_l2_regularizer(backend::Backend, T, eps)
  println("-- Testing L2 regularizer on $(typeof(backend)){$T}...")

  coef = rand()
  param = rand(T, 2,3,4,5) - convert(T, 0.5)
  param_blob = make_blob(backend, param)
  regu = L2Regu(coef)

  loss = forward(backend, regu, 1.0, param_blob)
  expected_loss = coef * vecnorm(param)^2
  @test -eps < loss - expected_loss < eps

  grad_blob = make_zero_blob(backend, T, size(param))
  backward(backend, regu, 1.0, param_blob, grad_blob)
  grad = zeros(T, size(param))
  copy!(grad, grad_blob)

  @test all(-eps .< grad - 2coef*param .< eps)
end

function test_l2_regularizer(backend::Backend)
  test_l2_regularizer(backend, Float32, 1e-5)
  test_l2_regularizer(backend, Float64, 1e-10)
end

if test_cpu
  test_l2_regularizer(backend_cpu)
end
if test_gpu
  test_l2_regularizer(backend_gpu)
end


@testset "Pressure and temp wrapper" begin

segments = 100
MDs = range(0, stop = 5000, length = segments) |> collect
incs = repeat([0], inner = segments)
TVDs = range(0, stop = 5000, length = segments) |> collect

well = Wellbore(MDs, incs, TVDs, 2.441)

model = WellModel(wellbore = well, roughness = 0.0006,
                    temperature_method = "Shiu", geothermal_gradient = 1.0, BHT = 200,
                    pressurecorrelation = HagedornAndBrown, WHP = 350 - pressure_atmospheric, dp_est = 25,
                    q_o = 100, q_w = 500, GLR = 1200, APIoil = 35, sg_water = 1.1, sg_gas = 0.8)

pressures = pressure_and_temp!(model)
temps = model.temperatureprofile

@test length(pressures) == length(temps) == segments
@test pressures[1] == 350 - pressure_atmospheric
@test pressures[end] ≈ (1068 - pressure_atmospheric) atol = 1
@test temps[end] == 200
@test temps[1] ≈ 181 atol = 1

end #testset for pressure & temp wrapper

@testset "Gas lift wrappers" begin

segments = 100
MDs = range(0, stop = 5000, length = segments) |> collect
incs = repeat([0], inner = segments)
TVDs = range(0, stop = 5000, length = segments) |> collect

well = Wellbore(MDs, incs, TVDs, 2.441)

testpath = joinpath(dirname(dirname(pathof(PressureDrop))), "test/testdata")
valvepath = joinpath(testpath, "valvedata_wrappers_1.csv")
valves = read_valves(path = valvepath, delim = ',', skiplines = 1) #implicit read_valves test

model = WellModel(wellbore = well, roughness = 0.0, valves = valves,
                    temperature_method = "Shiu", geothermal_gradient = 1.0, BHT = 200,
                    pressurecorrelation = HagedornAndBrown, WHP = 350 - pressure_atmospheric, dp_est = 25,
                    q_o = 0, q_w = 500, GLR = 4500, APIoil = 35, sg_water = 1.0, sg_gas = 0.8, CHP = 1000, naturalGLR = 0)

tubing_pressures, casing_pressures, valvedata = gaslift_model!(model, find_injectionpoint = true, dp_min = 100) #also an implied test for 100% water cut

Δmds = [MDs[i] - MDs[i-1] for i in 81:length(MDs)]
ΔPs = [tubing_pressures[i] - tubing_pressures[i-1] for i in 81:length(MDs)]
gradients = ΔPs ./ Δmds
mean(x) = sum(x) / length(x)

expected_gradient = 0.433 / GouldWaterVolumeFactor(mean(tubing_pressures[81:end]), mean(model.temperatureprofile[81:end]))
@test mean(gradients) ≈ expected_gradient atol = 0.005

valve_table(valvedata)

#%% implicit plot test
if test_plots
    plot_gaslift(model.wellbore, tubing_pressures, casing_pressures, model.temperatureprofile, valvedata, nothing) |> x->draw(SVG("plot-gl-core.svg", 5inch, 4inch), x)
end

end #testset for gas lift wrappers


@testset "Identity inflation" begin
    A = IdentityInflation()

    Nx = 10
    Ne = 20
    X = zeros(Nx, Ne)
    X .= randn(Nx, Ne)

    Xinflation = deepcopy(X)
    A(X)

    @test norm(X - Xinflation)<1e-14
end

@testset "Additive inflation I" begin
    Ny = 10
    Nx= 10
    Ne = 20
    A = AdditiveInflation(Nx)

    @test A.Nx==Nx
    @test mean(A) == zeros(Nx)
    @test cov(A) == Diagonal(ones(Nx))

    X = zeros(Ny + Nx, Ne)
    X .= randn(Ny + Nx, Ne)

    Xadd = deepcopy(X)
    A(Xadd, Ny+1, Ny+Nx)
    @test size(Xadd) == (Ny + Nx, Ne)
    @test Xadd != X
    @test norm(Xadd[1:Ny,:] -  X[1:Ny,:]) < 10*eps()
    @test norm(Xadd[Ny+1:Ny+Nx,:] -  X[Ny+1:Ny+Nx,:]) > 10*eps()

    Σ = zeros(Nx,Nx)

    for i=1:Nx
        for j=1:Nx
            Σ[i,j] = 1.0/(i + j)
        end
    end

    A = AdditiveInflation(Nx, zeros(Nx), Σ)

    @test A.Nx==10
    @test mean(A) == zeros(Nx)
    @test cov(A) == Σ
end

@testset "Additive inflation II" begin
    Nx= 10
    Ne = 20
    A = AdditiveInflation(Nx)

    @test A.Nx==10
    @test mean(A) == zeros(Nx)
    @test cov(A) == Diagonal(ones(Nx))

    X = zeros(Nx, Ne)
    X .= randn(Nx, Ne)

    Xadd = deepcopy(X)
    A(Xadd)
    @test size(Xadd) == (Nx, Ne)
    @test Xadd != X

    Σ = zeros(Nx,Nx)

    for i=1:Nx
        for j=1:Nx
            Σ[i,j] = 1.0/(i + j)
        end
    end

    A = AdditiveInflation(Nx, zeros(Nx), Σ)

    @test A.Nx==10
    @test mean(A) == zeros(Nx)
    @test cov(A) == Σ
end

@testset "Multiplicative inflation I" begin
    Ny = 10
    Nx= 10
    Ne = 20
    A = MultiplicativeInflation(1.2)

    @test A.β==1.2

    X = zeros(Ny + Nx, Ne)
    X .= randn(Ny + Nx, Ne)

    Xmul = deepcopy(X)
    A(Xmul, Ny+1, Ny+Nx)

    for i=1:Ne
        x̄ = deepcopy(mean(X, dims = 2))[Ny+1:Ny+Nx,1]
        @test norm(Xmul[1:Ny,i] - X[1:Ny,i])<1e-10
        @test norm(Xmul[Ny+1:Ny+Nx,i] - (x̄ + A.β*(X[Ny+1:Ny+Nx,i] - x̄)))<1e-10
    end
end

@testset "Multiplicative inflation II" begin
    Nx= 10
    Ne = 20
    A = MultiplicativeInflation(1.2)

    @test A.β==1.2

    X = zeros(Nx, Ne)
    X .= randn(Nx, Ne)

    Xmul = deepcopy(X)
    A(Xmul)

    for i=1:Ne
        x̄ = deepcopy(mean(X, dims = 2))[:,1]
        @test norm(Xmul[:,i] - (x̄ + A.β*(X[:,i] - x̄)))<1e-10
    end
end

@testset "MultiAdditive Inflation" begin

    T = MultiAddInflation(10)

    @test size(T)== 10
    @test T.β == 1.0

    @test mean(T) == zeros(10)
    @test cov(T) ==  Diagonal(ones(10))

    T = MultiAddInflation(10, 1.01, zeros(10), 1.01)

    @test T.β == 1.01
    @test mean(T) == zeros(10)
    @test cov(T) == Diagonal(1.01^2*ones(10))
end


include("custom_graph.jl")

"""
    colorGraph(g)

Find a coloring of the graph g such that no two vertices connected
by an edge have the same color.
"""
function colorGraph(G::CGraph)
    colornumber = 0
    V = num_vertices(G)
    colors = zeros(Int64,V)

    while (V > 0)
        x = max_degree_vertex(G)
        colornumber = colornumber + 1
        colors[x] = colornumber
        nn = non_neighbors(G,x)
        while (length(nn) > 0)
            maxcn = -1
            ydegree = -1
            for z in nn
                cn = length_common_neighbor(G,z,x)
                if (cn > maxcn) || (cn == maxcn && vertex_degree(G,z) < ydegree)
                    y = z
                    ydegree = vertex_degree(G,y)
                    maxcn = cn
                end
            end
            if (maxcn == 0)
                y = max_degree_vertex(G,nn)
            end
            colors[y] = colornumber
            contract!(G,y,x)
            nn = non_neighbors(G,x)
        end
        rem_vertex!(G,x)
        V = num_vertices(G)
    end
    return colors
end

"""
    num_vertices(G)

Find the total number of vertices present in graph G.
"""
function num_vertices(G::CGraph)
    return length(vertices(G))
end

"""
    max_degree_vertex(G, nn)

Find the vertex in the group nn of vertices belonging to the
graph G which has the highest degree.
"""
function max_degree_vertex(G::CGraph,nn::Array{Int64,1})
    max_degree = -1
    max_degree_vertex = -1
    for v in nn
        v_degree = length(neighbors(G,v))
        if v_degree > max_degree
            max_degree = v_degree
            max_degree_vertex = v
        end
    end
    return max_degree_vertex
end

"""
    max_degree_vertex(G)

Find the vertex in graph with highest degree.
"""
function max_degree_vertex(G::CGraph)
    max_degree = -1
    max_degree_vertex = -1
    for v in vertices(G)
        v_degree = length(neighbors(G,v))
        if v_degree > max_degree
            max_degree = v_degree
            max_degree_vertex = v
        end
    end
    return max_degree_vertex
end

"""
    non_neighbors(G,x)

Find the set of vertices belonging to the graph G which do
not share an edge with the vertex x.
"""
function non_neighbors(G::CGraph, x::Int64)
    num_non_neighbors = num_vertices(G) - vertex_degree(G,x) - 1
    nn = zeros(Int64,num_non_neighbors)
    index = 1

    for i in vertices(G)
        if ((indexin([i], neighbors(G,x)))[1] == nothing) && (i != x)
            nn[index] = i
            index += 1
        end
    end
    return nn
end

"""
    length_common_neighbor(G,z,x)

Find the number of vertices that share an edge with both the
vertices z and x belonging to the graph G.
"""
function length_common_neighbor(G::CGraph,z::Int64, x::Int64)
    z_neighbors = neighbors(G,z)
    x_neighbors = neighbors(G,x)
    common_vertices = indexin(z_neighbors, x_neighbors)
    num_common_vertices = 0
    for i in common_vertices
        if i != nothing
            num_common_vertices += 1
        end
    end
    return num_common_vertices
end

"""
    vertex_degree(G,z)

Find the degree of the vertex z which belongs to the graph G.
"""
function vertex_degree(G::CGraph,z::Int64)
    return length(neighbors(G,z))
end

"""
    contract!(G,y,x)

Contract the vertex y to x, both of which belong to graph G, that is
delete vertex y and join x with the neighbors of y if they are not
already connected with an edge.
"""
function contract!(G::CGraph,y::Int64, x::Int64)
    for v in neighbors(G,y)
        if has_edge(G,v,x) == false
            add_edge!(G,v,x)
        end
    end
    rem_vertex!(G,y)

end


"""
    Implements all needed parts for constructing DAGMM model
        (https://sites.cs.ucsb.edu/~bzong/doc/iclr18-dagmm.pdf) for detecting anomalies.
    Code is inspired by multiple reimplementations
    - https://github.com/danieltan07/dagmm
    - https://github.com/tnakae/DAGMM/
    - https://github.com/Newcomer520/tf-dagmm
    - https://github.com/mperezcarrasco/PyTorch-DAGMM
"""

using Flux: mse, softmax, unsqueeze, stack
using LinearAlgebra
using StatsBase
using Statistics
using Distributions
using DistributionsAD
using MLDataPattern: RandomBatches

struct DAGMM{E,D,S}
    encoder::E
    decoder::D
    estimator::S
end

Flux.@functor DAGMM

function DAGMM(;idim::Int=1, activation="tanh", hdim=60, zdim=1, ncomp=2,
        nlayers::Int=3, dropout=0.5, init_seed=nothing, kwargs...)
    # if seed is given, set it
    (init_seed !== nothing) ? Random.seed!(init_seed) : nothing

    encoder = build_mlp(idim, hdim, zdim, nlayers, activation=activation, lastlayer="linear")
    decoder = build_mlp(zdim, hdim, idim, nlayers, activation=activation, lastlayer="linear")
    estimator = build_mlp(zdim+2, hdim, ncomp, nlayers, activation=activation, lastlayer="linear")

    if dropout > 0.0
        estimator = Chain(estimator[1:end-1]..., Dropout(dropout), estimator[end])
    end

    # reset seed
    (init_seed !== nothing) ? Random.seed!() : nothing

    DAGMM(encoder, decoder, estimator)
end

function Base.show(io::IO, model::DAGMM)
    # to avoid the show explosion
    print(io, "DAGMM(...)")
end

norm2(x; dims=1) = sqrt.(sum(abs2, x; dims=dims))

function cosine_similarity(x₁, x₂; dims=1, ε=1e-8)
    T = eltype(x₁)
    nx₁ = norm2(x₁; dims=dims)
    nx₂ = norm2(x₂; dims=dims)
    sum(x₁ .* x₂; dims=dims) ./ max.(nx₁ .* nx₂, T(ε))
end

function compute_reconstruction(x, x̂)
    relrec = norm2(x .- x̂; dims=1) ./ norm2(x; dims=1)
    cosim = cosine_similarity(x, x̂; dims=1)
    relrec, cosim
end

function (model::DAGMM)(x)
    z_c = model.encoder(x)
    x̂ = model.decoder(z_c)
    rec_1, rec_2 = compute_reconstruction(x, x̂)
    z = cat([z_c, rec_1, rec_2]..., dims=1)
    gamma = softmax(model.estimator(z))
    z_c, x̂, z, gamma
end

"""
Computing the parameters phi, mu and gamma for sample energy function

K: number of Gaussian mixture components
N: Number of samples
D: Latent dimension

z = DxN
gamma = KxN

this will be a little bit different in Julia
"""
function compute_params(z, gamma)
    #phi = D
    gamma_sum = sum(gamma; dims=2)
    phi = gamma_sum ./ size(gamma, 2)

    #mu = D x K x 1
    mu = sum(unsqueeze(gamma,1) .* unsqueeze(z,2); dims=3) ./ gamma_sum'

    # z_mu = D x K x N
    z_mu = unsqueeze(z,2) .- mu

    # z_mu = D x D x K x N
    z_mu_z_mu_t = unsqueeze(z_mu, 2) .* unsqueeze(z_mu, 1)

    #cov = K x D x D
    cov = sum(unsqueeze(unsqueeze(gamma, 1), 1) .* z_mu_z_mu_t; dims=4)
    cov = dropdims(cov; dims=4) ./ unsqueeze(gamma_sum', 1)

    phi, mu, cov
end

function compute_energy(z, phi, mu, cov; ε=1e-7)
    T = eltype(z)
    z_mu = unsqueeze(z,2) .- mu
    D, K, _ = size(mu)

    cov = cov .+ Diagonal(ones(Float32, D) .* T(ε))

    cov_inverse = stack([inv(cov[:,:,k]) for k in 1:K], 3)
    det_cov = unsqueeze(unsqueeze([det(cov[:,:,k] .* 2 * T.(π)) for k in 1:K],1),1)
    cov_diag = sum([sum(1 ./ diag(cov[:,:,k])) for k in 1:K])

    E_z = -T(0.5) .* sum(sum(unsqueeze(z_mu, 1) .* cov_inverse, dims=2) .* unsqueeze(z_mu, 2), dims=1)
    E_z = exp.(E_z)
    E_z = -log.(sum(unsqueeze(phi', 1).* E_z ./ sqrt.(det_cov), dims=3) .+ T(ε))

    E_z, cov_diag
end

function loss(model::DAGMM, x, λ₁, λ₂)
    _, x̂, z, gamma = model(x)
    reconst_loss = mse(x, x̂)
    phi, mu, cov = compute_params(z, gamma)
    sample_energy, cov_diag = compute_energy(z, phi, mu, cov)
    reconst_loss + λ₁ * mean(sample_energy) + λ₂ * cov_diag
end

loss(model::DAGMM, x, λ₁, λ₂, batchsize::Int) =
    mean(map(y -> loss(model, y, λ₁, λ₂), Flux.Data.DataLoader(x, batchsize=batchsize)))

# during testing we use params fitted from training
# testmode is not really necessary because dropout is not part if the autoencoder
function StatsBase.predict(model::DAGMM, x, phi, mu, cov)
    testmode!(model, true)
    _, _, z, _ = model(x)
    testmode!(model, false)
    compute_energy(z, phi, mu, cov)[1][:]
end

function StatsBase.fit!(model::DAGMM, data::Tuple; max_train_time=82800,
                        batchsize=64, patience=200, check_interval::Int=1,
                        wreg=1f-6, lambda_e=0.1f0, lambda_d=0.005f0, lr=1f-4, kwargs...)
    # add regularization through weight decay in optimizer
    opt = (wreg > 0) ? ADAMW(lr, (0.9, 0.999), wreg) : Flux.ADAM(lr)

    trn_model = deepcopy(model)
    ps = Flux.params(trn_model);

    X = data[1][1]
    # filter only normal data from validation set
    X_val = data[2][1][:, data[2][2] .== 0.0f0]
    val_N = size(X_val, 2)

    history = MVHistory()
    _patience = patience

    best_val_loss = Inf
    i = 1
    start_time = time()
    frmt(v) = round(v, digits=4)
    for batch in RandomBatches(X, batchsize)
        batch_loss = 0f0

        grad_time = @elapsed begin
            gs = gradient(() -> begin
                batch_loss = loss(trn_model, batch, lambda_e, lambda_d)
            end, ps)
            Flux.update!(opt, ps, gs)
        end

        push!(history, :batch_loss, i, batch_loss)
        push!(history, :grad_time, i, grad_time)

        if (i%check_interval == 0)
            testmode!(trn_model, true)
            val_loss_time = @elapsed val_loss = loss(trn_model, X_val, lambda_e, lambda_d)
            # val_loss_time = @elapsed val_loss = val_N > 5000 ? loss(trn_model, X_val, lambda_e, lambda_d, 512) : loss(trn_model, X_val, lambda_e, lambda_d)
            testmode!(trn_model, false)

            @info "$i - loss: $(frmt(batch_loss)) (batch) | $(frmt(val_loss)) (validation) || $(frmt(grad_time)) (t_grad) | $(frmt(val_loss_time)) (t_val)"

            if isnan(val_loss) || isinf(val_loss) || isnan(batch_loss) || isinf(batch_loss)
                @info "Stopped training after $(i) iterations, $((time() - start_time)/3600) hours due to invalid values."
                error("Encountered invalid values in loss function.")
            end

            push!(history, :validation_loss, i, val_loss)
            push!(history, :val_loss_time, i, val_loss_time)

            if val_loss < best_val_loss
                best_val_loss = val_loss
                _patience = patience

                model = deepcopy(trn_model)
            else # else stop if the model has not improved for `patience` iterations
                _patience -= 1
                if _patience == 0
                    @info "Stopped training after $(i) iterations, $((time() - start_time)/3600) hours."
                    break
                end
            end
        end

        # time limit for training
        if time() - start_time > max_train_time
            @info "Stopped training after $(i) iterations, $((time() - start_time)/3600) hours due to time constraints."
            model = deepcopy(trn_model)
            break
        end

        i += 1
    end

    (history=history, niter=i, model=model, npars=sum(map(p->length(p), ps)))
end


#=
# Reduction

The `reduce` function takes in a binary operator `⊕` and a ordered collection, applying the operator to that collection effectively reducing it to one final value.

For example, the operator ⊕ can be `minimum` and the collection can be an array of Integers.
⊕ can also be `addition` or `xor` and the collection may represent any type of object as long as the operator makes sense in the context.
=#

using CUDA

a = rand(10)

#-

reduce(min, a)

#-

reduce(*, a)

#-

@doc reduce(*, a)

#=
Writing `reduce` for a CPU is quite straightforward with a single `for-loop`. We will focus on writing a reduction for a linear array in this tutorial. We will iteratively develop a performant version using everything we have learnt in the previous tutorials.

---

# Reduction 1 : Divide and Conquer

The first step is to write something that works on the GPU. If we were given a dual-core machine and expected to parallelize this we would split the input array into two halves and feed each half into a different CPU. Similarly the best hint is to use the [divide-and-conquer](https://en.wikipedia.org/wiki/Divide-and-conquer_algorithm) approach.
By envisioning the reduction process as a binary tree we get:

![reduction-1](../assets/reduction1.png)

It would be a good exercise to try to write model the above process in pseudocode.

One such approach is:

The only issue with this approach on the GPU is that after each step we need to synchronize which won't be possible with arrays which span over a single thread block (1024 threads is the maximum threads in a block). Hence, we will have to use a recursive approach.

Assume for now we have $1024$ threads per block and process one element per thread. If we have $2048$ threads then we will run our algorithm with two thread blocks, storing the results in an intermediate array. After our first kernel is done we will perform a reduction on the intermediate array. And if we have an array whose length is greater than $1024^2$ we will have another level of recursion. If there are $1024*1024 + 1$ elements then the $1^{st}$ level of reduction will return an intermediate array of size $1025$ which will take another level of recursion to process.
=#

function reduction1(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, threadsPerBlock))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda shmem=shmem threads=threadsPerBlock blocks=blocks reduction1_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction1(op, sums)[1]
    end

    CUDA.@allowscalar return sums[1]
end

#-

function reduction1_kernel(op, a, sums)
    shmem = @cuDynamicSharedMem(eltype(a), (blockDim().x, ))
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    index = threadIdx().x
    len = blockDim().x

    ## Adjust length for the last block
    if blockIdx().x == gridDim().x
        len = mod1(length(a), blockDim().x)
    end

    if tid <= length(a)
        shmem[index] = a[tid]
    else
        return
    end
    sync_threads()

    steps = floor(Int, CUDA.log2(convert(Float32, len)))
    for i = 0:steps
        if mod(index - 1, 2^(i + 1)) == 0 && (index + 2^i) <= len
            shmem[index] = op(shmem[index], shmem[index + 2^i])
        end
        sync_threads()
    end

    if index == 1
        sums[blockIdx().x] = shmem[1]
    end
    return
end

#-

a = CUDA.ones(1025);
reduction1(+, a)

#-

a = CUDA.rand(100_000);
println(reduction1(+, a),"  ",reduce(+, a))

#=
The two results above not being exactly equal is expected since IEEE floats are neither associative nor commutative. Since associativity is tough to achieve on a parallel algorithm we can expect some deviation.

---

# Reduction 2 : Strided Index

One problem with the last reduction was divergent branching, for example thread three is active for exactly one computation (`a[3] op a[4]`) and is never used again. Because GPU's operate warpwise we want to use all the resources of a warp instead of a small subset. When threads in a warp do different things it has diverged and it's efficiency drops. In this case with each iteration half the number of threads go inactive.

A simple fix is to change the way threads map to the elements by using a strided index

![reduction-2](../assets/reduction2.png)
=#

function reduction2_kernel(op, a, sums)
    shmem = @cuDynamicSharedMem(eltype(a), (blockDim().x, ))
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    index = threadIdx().x
    len = blockDim().x

    ## Adjust length for the last block
    if blockIdx().x == gridDim().x
        len = mod1(length(a), blockDim().x)
    end

    if tid <= length(a)
        shmem[index] = a[tid]
    else
        return
    end
    sync_threads()

    steps = floor(Int, CUDA.log2(convert(Float32, len)))
    for i = 0:steps
        stride = 2^i
        sindex = 2*stride*(index - 1) + 1
        if sindex + stride <= len
            shmem[sindex] = op(shmem[sindex], shmem[sindex + stride])
        end
        sync_threads()
    end

    if index == 1
        sums[blockIdx().x] = shmem[1]
    end
    return
end

#-

function reduction2(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, threadsPerBlock))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda shmem=shmem threads=threadsPerBlock blocks=blocks reduction2_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction2(op, sums)
    end

    CUDA.@allowscalar return sums[1]
end

#-

a = CUDA.ones(100_000);
reduction2(+, a)

#-

@time CUDA.@sync reduction1(+, a);
@time CUDA.@sync reduction2(+, a);

#=
# Reduction 3 : Sequential access

In both the above implementations our memory-access pattern is *strided* which is difficult to coalesce. We discussed *coalesced* memory access in the **Shared Memory** tutorial.

**TL;DR** When consecutive threads access consecutive locations in memory, the GPU combines several transactions into a fewer transactions which is called coalesced memory access. When memory accesses are not consecutive which happens when the locations are non-sequantial, sparse or misaligned the GPU hardware is unable to reduce the number of transactions. Since transactions are serviced sequentially there is a significant performance penalty for non-coalesced access.

To make use of sequantial access instead of `stride` iterating from 1 to `length ÷ 2` we can do it the other way around (`length ÷ 2`:1)

**NOTE**: The algorithm below assumes that the `blockDim` is a power of two. Transforming it to become friendly with non-power of two can be done as an exercise.
=#

function reduction3_kernel(op, a, sums)
    shmem = @cuDynamicSharedMem(eltype(a), (blockDim().x, ))
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    index = threadIdx().x
    len = blockDim().x

    ## Adjust length for the last block
    if blockIdx().x == gridDim().x
        len = mod1(length(a), blockDim().x)
    end

    if tid <= length(a)
        @inbounds shmem[index] = a[tid]
    else
        return
    end
    sync_threads()

    stride = len ÷ 2
    while stride > 0
        if index <= stride && index + stride <= len
            shmem[index] = op(shmem[index], shmem[index + stride])
        end
        stride = stride ÷ 2
        sync_threads()
    end

    if index == 1
        @inbounds sums[blockIdx().x] = shmem[1]
    end

    return
end

#-

function reduction3(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, threadsPerBlock))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda shmem=shmem threads=threadsPerBlock blocks=blocks reduction3_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction3(op, sums)[1]
    end

    CUDA.@allowscalar return sums[1]
    return sums
end

#-

a = CUDA.ones(1024);
reduction3(+, a)

#-

a = CUDA.ones(Int, 1024 * 1024);
reduction3(+, a)

#-

CUDA.@time reduction2(+, a);
CUDA.@time reduction3(+, a);

#=
# Reduction 4 : Warp Shuffle

A powerful feature in modern GPUs is the ability to communicate within warps with the help of special instructions. Currently we transfer data with the help of shared memory which obviously requires `sync_threads()` and access to shared memory. Warp shuffle functions allow transferring memory within a warp without the use of shared memory also being much faster and not requiring any explicit barrier. The only drawback is that only the following primitive types are supported: `Int32, UInt32, Int64, UInt64, Float32, Float64` and any arbitrary source to destination lane mapping is not permitted.

There are four shuffle methods.
- `shfl_sync`
- `shfl_up_sync`
- `shfl_down_sync`
- `shfl_xor_sync`
=#

@doc shfl_sync

#=
`shfl_sync` acts as a broadcast transferring a lane's value to all other lane.
=#

function broadcast_gpu(lane)
    id = threadIdx().x
    val = id
    mask = typemax(UInt32) # 0xffffffff
    newval = shfl_sync(mask, val, lane)
    @cuprint("id: ", id, "\t value: ", val, "\t new value: ", newval, "\n")
    return
end

@cuda threads=32 blocks=1 broadcast_gpu(19)

#=
`shfl_up_sync` and `shfl_down_sync` copy the value from lane = current_lane ± delta. If lane is out of bounds from the warp then
=#

@doc shfl_down_sync

#-

function shfldown_gpu(delta)
    id = threadIdx().x
    val = id
    mask = typemax(UInt32) # 0xffffffff
    newval = shfl_down_sync(mask, val, delta)
    @cuprint("id: ", id, "\t old value: ", val, "\t new value: ", newval, "\n")
    return
end

@cuda threads=32 blocks=1 shfldown_gpu(2)
synchronize()

#=
We can use `shfl_down_sync` to reduce a warp much faster than shared memory.
=#

@inline function reducewarp(op, val, mask = typemax(UInt32))
    val = op(val, shfl_down_sync(mask, val, 1))
    val = op(val, shfl_down_sync(mask, val, 2))
    val = op(val, shfl_down_sync(mask, val, 4))
    val = op(val, shfl_down_sync(mask, val, 8))
    val = op(val, shfl_down_sync(mask, val, 16))
    return val
end

function reduction4_kernel(op, a, sums)
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    lane_id = tid % 32
    warp_id = cld(tid, 32)

    ## exit
    tid > length(a) && return

    ## set essential
    tid <= length(a) && (val = a[tid])

    val = reducewarp(op, val)
    lane_id == 1 && (sums[warp_id] = val)
    return
end

function reduction4(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, 32))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda threads=threadsPerBlock blocks=blocks reduction4_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction4(op, sums)
    end

    CUDA.@allowscalar return sums[1]
    return sums
end

#-

a = CUDA.ones(Int, 320_000)
reduction4(+, a)

#=
There is one big problem with our implementation, the input length is expected to be a multiple of 32. This problem can be solved either by defining a neutral element for `op` (`zero(eltype(a))` for `+`, `one(eltype(a))` for `*`) however this won't work with `xor`. Another is to force the last warp's computation via shared memory like earlier examples.
We correct this by having only 1 thread work for the last warp; this is a simple solution that is inefficient but when the number of warps is large the performance hit shouldn't be too high.
=#

function reduction5_kernel(op, a, sums)
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    lane_id = tid % 32
    warp_id = cld(tid, 32)

    ## exit non essential
    tid > cld(length(a), 32)*32 && return

    ## set essential
    tid <= length(a) && (val = a[tid])

    ## Manage last warp
    if warp_id*32 > length(a)
        if lane_id == 1
            for i=(tid + 1):length(a)
                val = op(val, a[i])
            end
        end
    else
        val = reducewarp(op, val)
    end

    lane_id == 1 && (sums[warp_id] = val)
    return
end

function reduction5(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, 32))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda threads=threadsPerBlock blocks=blocks reduction5_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction5(op, sums)
    end

    CUDA.@allowscalar return sums[1]
    return sums
end

#-

a = CUDA.ones(Int, 5_000_000)
reduction5(+, a)

#-

@time CUDA.@sync reduction3(+, a);
@time CUDA.@sync reduction5(+, a);