using LsqFit, Distributions, DSP, DelimitedFiles

struct input_data
    t::Array{Float64,1}
	DTOF::Array{Float64,1}
	IRF::Array{Float64,1}
	DTOFerrors::Array{Float64,1}
	ρ::Float64
	nmed::Float64
	ndet::Float64
	lambda::Float64
	normpeaks::Bool
end

struct model_params
	#yerrors::Array{Float64,1}          #uncertainties in dependent variable
    model::Function                    #model to be fitted against
	initparams::Array{Float64,1}
	lb::Array{Float64,1}
	ub::Array{Float64,1}
	risefactor::Float64  			#percent of peak to fit on rising edge
	tailfactor::Float64				#perfecnt of peak to fit on failing tail
	dof::Int64
end

struct fitresult
    #Fit inputs
    xdata::Array{Float64,1}     #independent variable
    ydata::Array{Float64,1}     #dependent variable
    yerrors::Array{Float64,1}   #uncertainties in dependent variable
	model::Function             #model to be fitted against

    #Fit results
    residuals::Array{Float64,1}  	 #weighted residuals
    params::Array{Float64,1}    		#best fit parameters for model
	marginerror::Array{Float64,1}        #product of standard error and the critical value of each parameter at a certain significance level (default is 5%) from t-distribution.
    perrors::Array{Float64,1}            #uncertainties in best fit parameters
    dof::Int64                  		#number of degrees of freedom in fit
    chisq::Float64             		 #chi^2 value of fit
    probability::Float64      		  #probability of chi^2 value
end

function DT_model(t,β,ρ,nmed,ndet)
	#diffusion theory model for homogenous semi-infinite layer
	μa = β[1]
	μsp = β[2]

	if (nmed == ndet)
		Afac = 1
	elseif nmed > ndet
		costhetac = sqrt(1 - (ndet/nmed).^2)
		R0 = ((nmed/ndet-1)/(nmed/ndet +1)).^2
		Afac = (2/(1-R0) -1 + abs(costhetac.^3))./(1-abs(costhetac.^2))
	else nmed < ndet
		R0 = ((nmed/ndet-1)/(nmed/ndet +1)).^2
		Afac = 2/(1-R0) - 1
	end

	z0 = 1/(μa+μsp)
	D = 1/(3*(μa + μsp))
	zb = 2*Afac*D

	v = 29.9792345/nmed # speed of light in medium

	Rt1 = @. v*exp(-μa*v*t)
	Rt1 = @. Rt1/((4*π*D*v*t)^1.5)
	Rt1 = @. Rt1*(exp(-(z0^2 + ρ^2)/(4*D*v*t)) - exp(-((2*zb + z0)^2 + ρ^2)/(4*D*v*t)))

	Rt2 = @. 3*exp(-μa*v*t)/(2*((4*π*D*v)^1.5)*(t^2.5))
	Rt2 = @. Rt2*(z0*exp(-((z0^2 + ρ^2)/(4*D*v*t))) + (2*zb + z0)*exp(-((2*zb+z0)^2 + ρ^2)/(4*D*v*t)))

	Rt = @. abs(Rt1)+abs(Rt2)
	replace!(Rt, NaN => 0)
	m = findmax(Rt)
	Rt = Rt./m[1]
end

function conv_DT(t, β, data)

	RtDT = DT_model(data.t, β, data.ρ, data.nmed, data.ndet)

	convDT = conv(data.IRF, RtDT)
	convDT = convDT./maximum(convDT)

	tidx = findfirst(x -> x == t[1], data.t)

	if data.normpeaks == true
		peakindDTOF, peakindconvDT = findmax(data.DTOF)[2], findmax(convDT)[2]
		if peakindDTOF > peakindconvDT
			convDT = [zeros(peakindDTOF-peakindconvDT); convDT]
		elseif peakindDTOF < peakindconvDT
			convDT = convDT[peakindconvDT-peakindDTOF:end]
		end
	end

	convDT = convDT[tidx:tidx+length(t)-1]
	return log.(convDT)
end

function getfit(input_data, model_params; alpha=0.1)

	#get fitting window
	counts = input_data.DTOF
	maxvalue, maxindex = findmax(counts)

	ind1 = findfirst(x -> x > maxvalue*model_params.risefactor, counts)
	ind2 =  findlast(x -> x >maxvalue*model_params.tailfactor, counts)

	weights = 1 ./ input_data.DTOFerrors

	fit = curve_fit((t, β) -> conv_DT(t, β, input_data), input_data.t[ind1:ind2], log.(input_data.DTOF[ind1:ind2]), model_params.initparams, lower=model_params.lb, upper=model_params.ub)

	perrors = estimate_errors(fit,alpha)
	marginerror = margin_error(fit,alpha)

	chisq = sum(fit.resid.^2)
	chisqdist = Distributions.Chisq(2)
	probability = Distributions.ccdf(chisqdist,chisq)
	return fitresult(input_data.t[ind1:ind2],log.(input_data.DTOF[ind1:ind2]),input_data.DTOFerrors, DT_model, fit.resid,fit.param,marginerror,perrors,2,chisq,probability)
end

function plotfit(fit::fitresult1, data::input_data)
	scatter(fit.xdata, fit.ydata, color="black", label="DTOF", markersize=3, alpha=0.8)
	plot!(fit.xdata, conv_DT(fit.xdata, fit.params, data), color="red", lw=3, label="DT-fit", alpha=0.8)

end

function loadData(filename)
    #loads data from asc
    data_array = readdlm(filename,skipstart=10)
    t = data_array[1:end-1,1]
    Rt = data_array[1:end-1,2]
    (time = convert(Array{Float64},t), counts = convert(Array{Float64},Rt))
end


"""
Tableau for the 2-stage stochastic LobattoIIIA-IIIB-IIID method
  Tableau for the 2-stage stochastic LobattoIIIA-IIIB-IIID method
  (based on the deterministic LobattoIIIA-IIIB-IIID due to L. Jay)
  It satisfies the conditions for convergence of order 1.0 for one Wiener process,
  but it doesn't satisfy the conditions for Lagrange-d'Alembert integrators
"""
function TableauStochasticLobattoIIIABD2()

    TableauSISPRK(:StochasticLobattoIIIABD2, CoefficientsLobattoIIIA(2), CoefficientsLobattoIIIA(2),
                                         CoefficientsLobattoIIIB(2), CoefficientsLobattoIIID(2),
                                         CoefficientsLobattoIIIB(2), CoefficientsLobattoIIID(2))
end

"""
Tableau for the 2-stage modified stochastic LobattoIIIA-IIIB method
  Tableau for the 2-stage modified stochastic LobattoIIIA-IIIB method
  Satisfies the conditions for Lagrange-d'Alembert integrators
  and the conditions for convergence of order 1.0 for one Wiener process
"""
function TableauModifiedStochasticStoermerVerlet(c::Number=0.0)

    @assert c ≥ 0.0
    @assert c ≤ 1.0

    a_drift2 = [[c 0.0]
                [c 0.0]]

    b_drift2 = [c, 1-c]

    c_drift2 = [c, c]

    TableauSISPRK(:StochasticModifiedStormerVerlet,
                  CoefficientsLobattoIIIA(2), CoefficientsLobattoIIIA(2),
                  CoefficientsLobattoIIIB(2), CoefficientsRK(typeof(c), :cTableau, 1, a_drift2, b_drift2, c_drift2),
                  CoefficientsLobattoIIIB(2), CoefficientsLobattoIIIB(2))
end


"Tableau for the s-stage Gauss-Lobatto SFIRK method"
function TableauStochasticGLRK(s::Int)
    TableauSIRK(Symbol("StochasticGLRK" * string(s)), CoefficientsGLRK(s), CoefficientsGLRK(s))
end

"""
Tableau for the 2-stage stochastic symplectic DIRK method
  Tableau for the stochastic symplectic DIRK method
  Satisfies the conditions for Lagrange-d'Alembert integrators.
  Satisfies the conditions for strong convergence of order 1.0 for one Wiener process
"""
function TableauStochasticDIRK(c::Number=0.5)

    a_drift = [[c/2  0.0    ]
                [c    (1-c)/2]]

    b_drift = [c, 1-c]

    c_drift = [c/2, (1+c)/2]

    TableauSIRK(:StochasticSymplecticDIRK, 2, a_drift, b_drift, c_drift, 2, a_drift, b_drift, c_drift)
end


using Transparency

"""
    function B_ν(ν, T)

Planck's law! Radiation in LTE. Takes frequency and temperature, returns
specific intensity
"""
function B_ν(ν, T)::AbstractFloat
    return 2*h*ν^3/c_0^2 * 1/(exp(h*ν/(k_B*T)) - 1)
end

"""
    function B_ν(λ, T)

Planck's law! Radiation in LTE. Takes wavelength and temperature, returns
specific intensity
"""
function B_λ(λ::Unitful.Length, T::Unitful.Temperature)
    return 2*h*c_0^2/λ^5 * 1/(exp(h*c_0/(λ*k_B*T)) - 1)
end

"""
    function α_absorption(λ::Unitful.Length, temperature::Unitful.Temperature,
               electron_density::NumberDensity, h_ground_density::NumberDensity,
               proton_density::NumberDensity)

Extinction from photon destruction processes.
"""
function α_absorption(λ::Unitful.Length, temperature::Unitful.Temperature,
                      electron_density::NumberDensity, h_neutral_density::NumberDensity,
                      proton_density::NumberDensity)

    # α = max(0u"m^-1", Transparency.hminus_ff_stilley(λ, temperature, h_ground_density, electron_density))
    α = hminus_ff(λ, temperature, h_neutral_density, electron_density; recipe="stilley")
    α += hminus_bf(λ, temperature, h_neutral_density, electron_density; recipe="geltman")
    α += hydrogenic_ff(c_0 / λ, temperature, electron_density, proton_density, 1)
    # hydrogenic_bf ??? why return return 0 * αbf_const * species_densitys
    α += h2plus_ff(λ, temperature, h_neutral_density, proton_density)
    α += h2plus_bf(λ, temperature, h_neutral_density, proton_density)
    return α
end

"""
    α_scattering(λ::Unitful.Length,
                 electron_density::NumberDensity,
                 h_ground_density::NumberDensity)

Extinction from scattering processes.
"""
function α_scattering(λ::Unitful.Length,
                      electron_density::NumberDensity,
                      h_ground_density::NumberDensity)

   α = thomson(electron_density)
   α += rayleigh_h(λ, h_ground_density) # Gives the edge... (1217.7 Å)
   return α
end

# function α_absorption(λ::Unitful.Length, temperature::Unitful.Temperature,
                      # electron_density::NumberDensity, h_ground_density::NumberDensity,
                      # proton_density::NumberDensity)
#
    # α = max(0u"m^-1", Transparency.hminus_ff_stilley(λ, temperature, h_ground_density, electron_density))
    # α += Transparency.hminus_bf_wbr(λ, temperature, h_ground_density, electron_density)
    # α += hydrogenic_ff(c_0 / λ, temperature, electron_density, proton_density, 1)
    # α += h2plus_ff(λ, temperature, h_ground_density, proton_density)
    # α += h2plus_bf(λ, temperature, h_ground_density, proton_density)
    # return α
# end


using Intervals
using Base.Test

@test Interval(1,1) == Interval(1.0, 1.0)
@test Interval{Float64}(1,1) === Interval{Float64}(1.0, 1.0)

@test Interval(1,1)/Interval(1,1) == Interval(1,1)
@test Interval(1,1)/Interval(2,3) == Interval(1/3,0.5)

@test Interval(1,1)*Interval(1,1) == Interval(1,1)
@test Interval(1,1)*Interval(2.5,3.4) == Interval(3.4,2.5)

@test inf(Interval(2,1)) == 1
@test sup(Interval(2,1)) == 2


# [ref](https://opensourc.es/blog/random-forest)

mutable struct Node

    feature    :: Int64
    comp_value :: Float64
    data_idxs  :: Vector{Int64}
    mean       :: Float64
    left       :: Union{Nothing, Node}
    right      :: Union{Nothing, Node}

    Node() = new()

end

mutable struct Tree
    root   :: Node
    Tree() = new()
end

function create_random_forest(feature_matrix, train_ys, ntrees)

    forest             = Vector{Tree}()
    tc                 =  1
    total_nfeatures    = size(feature_matrix)[1]
    nfeatures_per_tree = 5
    ncols_per_tree     = length(train_ys)÷2

    while tc < ntrees

        feature_idx = sample(1:total_nfeatures,
                             nfeatures_per_tree;
                             replace=false)

        cols = sample(1:length(train_ys), ncols_per_tree; replace=false)

        tree = create_random_tree(feature_matrix,
                                  feature_idx,
                                  cols,
                                  train_ys,
                                  total_nfeatures)

        if tree.root.left_child != nothing
            push!(forest,tree)
            tc += 1
        else
            println("Couldn't create tree!!")
        end
        flush(stdout)
    end

    forest

end

function create_random_tree(glob_feature_matrix, ...)

    tree = Tree()
    @views feature_matrix = glob_feature_matrix[feature_idx, cols]
    root, queue = create_root_node!(...)

    while length(queue) > 0
        node = popfirst!(queue)
        left_node, right_node = compute_node!(...)
        if left_node != nothing
            queue_compute_nodes!(...)
        end
    end
    tree.root = root
    tree

end


using Manopt, Manifolds, ManifoldsBase, LinearAlgebra, Test

@testset "Conjugate Gradient coefficient rules" begin
    F(x) = norm(x)^2
    ∇F(x) = 2 * x
    M = Euclidean(2)
    P = GradientProblem(M, F, ∇F)
    x0 = [0.0, 1.0]
    sC = StopAfterIteration(1)
    s = ConstantStepsize(1.0)
    retr = ExponentialRetraction()
    vtm = ParallelTransport()

    ∇1 = [1.0, 1.0]
    δ1 = [0.0, 2.0]
    ∇2 = [1.0, 1.5]
    δ2 = [0.5, 2.0]
    diff = ∇2 - ∇1

    dU = SteepestDirectionUpdateRule()
    O = ConjugateGradientDescentOptions(x0, sC, s, dU, retr, vtm)
    @test O.coefficient(P, O, 1) == 0

    dU = ConjugateDescentCoefficient()
    O = ConjugateGradientDescentOptions(x0, sC, s, dU, retr, vtm)
    O.∇ = ∇1
    O.δ = δ1
    # for the first case we get zero
    @test O.coefficient(P, O, 1) == 0.0
    O.∇ = ∇2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(∇2, ∇2) / dot(-δ2, ∇1)

    dU = DaiYuanCoefficient()
    O = ConjugateGradientDescentOptions(x0, sC, s, dU, retr, vtm)
    O.∇ = ∇1
    O.δ = δ1
    # for the first case we get zero
    @test O.coefficient(P, O, 1) == 0.0
    O.∇ = ∇2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(∇2, ∇2) / dot(δ2, ∇2 - ∇1)

    dU = FletcherReevesCoefficient()
    O = ConjugateGradientDescentOptions(x0, sC, s, dU, retr, vtm)
    O.∇ = ∇1
    O.δ = δ1
    # for the first case we get zero
    @test O.coefficient(P, O, 1) == 1.0
    O.∇ = ∇2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(∇2, ∇2) / dot(∇1, ∇1)

    dU = HagerZhangCoefficient()
    O = ConjugateGradientDescentOptions(x0, sC, s, dU, retr, vtm)
    O.∇ = ∇1
    O.δ = δ1
    # for the first case we get zero
    @test O.coefficient(P, O, 1) == 0.0
    O.∇ = ∇2
    O.δ = δ2
    denom = dot(δ1, diff)
    ndiffsq = dot(diff, diff)
    @test O.coefficient(P, O, 2) ==
          dot(diff, ∇2) / denom - 2 * ndiffsq * dot(δ1, ∇2) / denom^2

    dU = HeestenesStiefelCoefficient()
    O = ConjugateGradientDescentOptions(x0, sC, s, dU, retr, vtm)
    O.∇ = ∇1
    O.δ = δ1
    @test O.coefficient(P, O, 1) == 0.0
    O.∇ = ∇2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(diff, ∇2) / dot(δ1, diff)

    dU = LiuStoreyCoefficient()
    O = ConjugateGradientDescentOptions(x0, sC, s, dU, retr, vtm)
    O.∇ = ∇1
    O.δ = δ1
    @test O.coefficient(P, O, 1) == 0.0
    O.∇ = ∇2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == -dot(∇2, diff) / dot(δ1, ∇1)

    dU = PolakRibiereCoefficient()
    O = ConjugateGradientDescentOptions(x0, sC, s, dU, retr, vtm)
    O.∇ = ∇1
    O.δ = δ1
    @test O.coefficient(P, O, 1) == 0.0
    O.∇ = ∇2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(∇2, diff) / dot(∇1, ∇1)
end
@testset "Conjugate Gradient runs – Low Rank matrix approx" begin
    A = Diagonal([2.0, 1.1, 1.0])
    M = Sphere(size(A, 1) - 1)
    F(x) = x' * A * x
    euclidean_∇F(x) = 2 * A * x
    ∇F(x) = project(M, x, euclidean_∇F(x))

    x0 = [2.0, 0.0, 2.0] / sqrt(8.0)
    x_opt = conjugate_gradient_descent(
        M,
        F,
        ∇F,
        x0;
        stepsize=ArmijoLinesearch(),
        coefficient=FletcherReevesCoefficient(),
        stopping_criterion=StopAfterIteration(15),
    )
    @test isapprox(F(x_opt), minimum(eigvals(A)); atol=2.0 * 10^-4)
    @test isapprox(x_opt, eigvecs(A)[:, size(A, 1)]; atol=3.0 * 10^-2)
    x_opt2 = conjugate_gradient_descent(
        M,
        F,
        ∇F,
        x0;
        stepsize=ArmijoLinesearch(),
        coefficient=FletcherReevesCoefficient(),
        stopping_criterion=StopAfterIteration(15),
        return_options=true,
    )
    @test get_solver_result(x_opt2) == x_opt
end


# Part 1
program = readlines("day14_input.txt")
getval(line) = last(split(line, " = "))
getmemloc(line) = parse(Int64, first(split(line, " = "))[5:end-1])
applymask(mask, val) = parse(Int64, join([m == 'X' ? v : m for (m, v) in zip(mask, bitstring(val)[end-35:end])]), base=2)

function run(prog)
    mask = repeat("0", 36)
    mem = Dict()
    for line in prog
        if startswith(line, "mask")
            mask = getval(line)
        else
            mem[getmemloc(line)] = applymask(mask, parse(Int64, getval(line)))
        end
    end
    mem
end
println("Part 1: sum = $(sum(values(run(program))))")

# Part 2
function masklocs(mask, loc)
    idx = 1
    locformat = ""
    for (m, v) in zip(mask, bitstring(loc)[end-35:end])
        if m == 'X'
            locformat *= "[$idx]"
            idx += 1
        elseif m == '1'
            locformat *= '1'
        else
            locformat *= v
        end
    end

    if idx == 1 return [parse(Int64, locformat, base=2)] end
    xcount = idx - 1

    locs = []
    for i in 0:2^xcount-1
        replaced = locformat
        for (j, x) in enumerate(collect(bitstring(i)[end-(xcount-1):end]))
            replaced = replace(replaced, "[$j]" => "$x")
        end
        push!(locs, replaced)
    end
    locs
end

function run2(prog)
    mask = repeat("0", 36)
    mem = Dict()
    for line in prog
        if startswith(line, "mask")
            mask = getval(line)
        else
            loc = getmemloc(line)
            for ml in masklocs(mask, loc)
                mem[ml] = parse(Int64, getval(line))
            end
        end
    end
    mem
end
println("Part 2: sum = $(sum(values(run2(program))))")


export Disk,𝔻

# T is Real or Complex
# TT is (Real,Real) or Complex
immutable Disk{T,TT} <: BivariateDomain{T}
    center::TT
    radius::Float64
end

Disk(c::Complex,r)=Disk{typeof(c),typeof(c)}(c,r)
Disk{T<:Real}(c::Tuple{T,T},r)=Disk{T,typeof(c)}(c,r)

Disk()=Disk((0.,0.),1.)
Disk(::AnyDomain)=Disk(NaN,(NaN,NaN))

const 𝔻=Disk()

isambiguous(d::Disk)=isnan(d.radius) && all(isnan,d.center)

#canonical is rectangle [r,0]x[-π,π]
# we assume radius and centre are zero for now
fromcanonical{T<:Real}(D::Disk{T},x,t)=D.radius*x*cos(t)+D.center[1],D.radius*x*sin(t)+D.center[2]
tocanonical{T<:Real}(D::Disk{T},x,y)=sqrt((x-D.center[1])^2+(y-D.center[2])^2)/D.radius,atan2(y-D.center[2],x-D.center[1])
checkpoints(d::Disk)=[fromcanonical(d,(.1,.2243));fromcanonical(d,(-.212423,-.3))]

# function points(d::Disk,n,m,k)
#     ptsx=0.5*(1-gaussjacobi(n,1.,0.)[1])
#     ptst=points(PeriodicInterval(),m)
#
#     Float64[fromcanonical(d,x,t)[k] for x in ptsx, t in ptst]
# end

∂(d::Disk)=Circle(Complex(d.center...),d.radius)


function sbb(borrow::UInt8, a::ArbDigit, b::ArbDigit)
    difference = SignedDoubleArbDigit(a) - SignedDoubleArbDigit(b) - SignedDoubleArbDigit(borrow)
    lowBits(difference), UInt8(difference < 0)
end

function sub2!(a::AbstractVector{ArbDigit}, b::AbstractVector{ArbDigit})
    len = min(length(a), length(b))
    a_lo, a_hi = @views a[begin:len], a[len+1:end]
    b_lo, b_hi = @views b[begin:len], b[len+1:end]

    borrow = 0x0
    for i in eachindex(a_lo)
        a_lo[i], borrow = sbb(borrow, a_lo[i], b_lo[i])
    end

    if !iszero(borrow)
        for i in eachindex(a_hi)
            a_hi[i], borrow = sbb(borrow, a_hi[i], zero(ArbDigit))
            iszero(borrow) && break
        end
    end

    (!iszero(borrow) || any(!iszero, b_hi)) && throw(ArgumentError("Cannot subtract b from a because b is larger than a."))
end

function sub!(a::ArbUInt, b::ArbUInt)
    sub2!(a.data, b.data)
    normalize!(a)
end
sub!(a::ArbUInt, b::T) where T <: Unsigned = sub!(a, ArbUInt(b))

function Base.:(-)(a::ArbUInt, b::ArbUInt)
    a_copy = deepcopy(a)
    sub2!(a_copy.data, b.data)
    a_copy
end
Base.:(-)(a::ArbUInt, b::T) where T <: Unsigned = sub!(deepcopy(a), b)
Base.:(-)(a::T, b::ArbUInt) where T <: Unsigned = sub!(deepcopy(b), a)


#!/usr/bin/env julia
#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 1, Zadanie 1.1.
#

# assert the program’s arguments
if length(ARGS) == 0
    println("error: provide the floating point precision (16, 32 or 64 bits)")
    exit(1)
end

fl_precision_string = ARGS[1]

if fl_precision_string == "16"
    fl = Float16
elseif fl_precision_string == "32"
    fl = Float32
elseif fl_precision_string == "64"
    fl = Float64
else
    println("error: provide one of [16, 32, 64]")
    exit(1)
end
# assert end

# start finding the machine epsilon at 0.5
mach_eps = fl(1/2)

# iteratively find the machine epsilon
for i in range(1, length=100)
    # if adding even smaller eps would not change the value break the loop
    if fl(1) + fl(mach_eps/2) == fl(1)
        break
    end
    # otherwise continue
    global mach_eps = fl(mach_eps/2)
end

# print the output
println("iterative method = ", mach_eps, "; eps(fl) = ", eps(fl), "; equal = ", (eps(fl) == mach_eps))


"""
    chain = mcmc(θ, reps, burnin, Prior, lnL, Proposal)

    Simple MH MCMC

    You must set the needed functions, e.g.,:
    Prior = θ -> your_prior(θ, whatever other args)
    lnL = θ -> your_log_likelihood(θ, whatever other args)
    Proposal = θ -> your_proposal(θ, whatever other args)
    (optionally) mcmcProposalDensity = (θᵗ,θ) -> your_proposal_density

    then get the chain using the above syntax, or optionally,
    (non-symmetric proposal) chain = mcmc(θ, reps, burnin, Prior, lnL, Proposal, ProposalDensity, report=true)
    (example code) mcmc(): runs a simple example. edit(mcmc,()) to see the code

"""
function mcmc()
    println("mcmc(), called with no arguments, runs a simple example")
    println("execute edit(mcmc,()) to see the code")
    # sample is from exponential, prior is lognormal, proposal is random walk lognormal
    y = rand(Exponential(3.0),30)
    # set prior, likelihood and proposal
    Prior = θ -> pdf.(Ref(LogNormal(1.0,1.0)), θ)
    lnL = θ -> sum(logpdf.(Ref(Exponential(θ)), y))
    tuning = 0.5
    Proposal = θ -> rand(LogNormal(log(θ),tuning))
    # get the chain, plot posterior, and descriptive stats
    chain = mcmc(1.0, 100000, 10000, Prior, lnL, Proposal, true) # start value, chain length, and burnin
    p = npdensity(chain[:,1]) # nonparametric plot of posterior density
    plot!(p, title="posterior density, simple MCMC example: true value = 3.0", show=true) # add a title
    dstats(chain)
    return
end

# method using threads and symmetric proposal
@views function mcmc(θ, reps::Int64, burnin::Int64, Prior::Function, lnL::Function, Proposal::Function, report::Bool, nthreads::Int64)
    perthread = reps ÷ nthreads
    chain = zeros(reps, size(θ,1)+1)
    Threads.@threads for t = 1:nthreads # collect the results from the threads
        chain[t*perthread-perthread+1:t*perthread,:] = mcmc(θ, perthread, burnin, Prior, lnL, Proposal, report)
    end
    return chain
end

# method symmetric proposal
# the main loop
@views function mcmc(θ, reps::Int64, burnin::Int64, Prior::Function, lnL::Function, Proposal::Function, report::Bool=true)
    reportevery = Int((reps+burnin)/10)
    lnLθ = lnL(θ)
    chain = zeros(reps, size(θ,1)+1)
    naccept = zeros(size(θ))
    for rep = 1:reps+burnin
        θᵗ = Proposal(θ) # new trial value
        if report
            changed = Int.(.!(θᵗ .== θ)) # find which changed
        end
        # MH accept/reject: only evaluate logL if proposal is in support of prior (avoid crashes)
        pt = Prior(θᵗ)
        accept = false
        if pt > 0.0
            lnLθᵗ = lnL(θᵗ)
            accept = rand() < exp(lnLθᵗ-lnLθ) * pt/Prior(θ)
            if accept
                θ = θᵗ
                lnLθ = lnLθᵗ
            end
        end
        if report
            naccept = naccept .+ changed .* Int.(accept)
        end
        if (mod(rep,reportevery)==0 && report)
            println("current parameters: ", round.(θ,digits=3))
            println("  acceptance rates: ", round.(naccept/reportevery,digits=3))
            naccept = naccept - naccept
        end
        if rep > burnin
            chain[rep-burnin,:] = vcat(θ, accept)
        end
    end
    return chain
end

#=
# the main loop
function mcmc(θ, reps, burnin, Prior, lnL, Proposal::Function, ProposalDensity::Function, report=true::Bool)
    reportevery = Int((reps+burnin)/10)
    lnLθ = lnL(θ)
    chain = zeros(reps, size(θ,1)+1)
    naccept = zeros(size(θ))
    for rep = 1:reps+burnin
        θᵗ = Proposal(θ) # new trial value
        if report
            changed = Int.(.!(θᵗ .== θ)) # find which changed
        end
        # MH accept/reject: only evaluate logL if proposal is in support of prior (avoid crashes)
        pt = Prior(θᵗ)
        accept = false
        if pt > 0.0
            lnLθᵗ = lnL(θᵗ)
            accept = rand() <
            exp(lnLθᵗ-lnLθ)*
            pt/Prior(θ)*
            ProposalDensity(θ,θᵗ)/ProposalDensity(θᵗ,θ)
            if accept
                θ = θᵗ
                lnLθ = lnLθᵗ
            end
        end
        if report
            naccept = naccept .+ changed .* Int.(accept)
        end
        if (mod(rep,reportevery)==0 && report)
            println("current parameters: ", round.(θ,digits=3))
            println("  acceptance rates: ", round.(naccept/reportevery,digits=3))
            naccept = naccept - naccept
        end
        if rep > burnin
            chain[rep-burnin,:] = [θ; accept]
        end
    end
    return chain
end
=#


using Plots

f(t) = exp(im*t)
t = 0:0.01:2pi
z = f.(t)

x = real.(z)
y = imag.(z)

plot(x, y)
savefig("eulers_formula.png")


export  minmax_normalize,
        opencv_disparity_map_to_image

"Normalize all values in array based on min and max value"
function minmax_normalize(A::AbstractArray)
    min = minimum(A)
    max = maximum(A)
    (A .- min) ./ (max - min)
end

"""
    opencv_disparity_map_to_image(A::AbstractArray) -> Matrix{Gray}

OpenCV algorithms for stereo matching (create 3D image from two 2D images)
outputs grayscale images representing disparity maps. These are 16-bit integers
which represent fixed point numbers. They are encoded such that 16 means  1 and
1 means 1/16.

OpenCV images are tranposed relative to Julia images. In Julia images are column ordered,
while OpenCV is row ordered. That is Julia follows the typical Fortran convention while
OpenCV follows the C convention. This means for a Julia image if you ask for size, then
the first dimension is the  height, and the second is the width.
"""
function opencv_disparity_map_to_image(A::AbstractArray)
	grays = Gray.(minmax_normalize(A))
	img = reshape(grays, size(A, 2), size(A, 3))
	transpose(img)
end

"""
    image_to_opencv_input(img::AbstractMatrix) -> Array{UInt8, 3}

For OpenCV Image data is always stored in 3D arrays. The reason is that the color
channel represents a dimension. Thus in OpenCV a 200×300 image with RGB values
will be represented by a 3×200×300 array. If the image was grayscale however it would be
represented by a 1×200×300 array.

However in Julia this is different. The first case would be a 300×200 matrix
of `RBG` values. The type would be `Array{RGB, 2}`. Rows and columns would be transposed because rows
are specifies first for a Julia matrix.

The grayscale image would be a 300×200 matrix as well, except the element type would be `Gray`.
The type would be `Array{Gray, 2}`.

The element type of `img` could be `RGB`, `Gray` or some `Real` value. It doesn't matter
as each element will be converted to `Gray` values anyway.

# Examples
```julia
import OpenCV
import ImageView, Images, FileIO

imgL = FileIO.load("images/L1.jpg");
imgR = FileIO.load("images/R1.jpg");

imgL = image_to_opencv_input(imgL)
imgR = image_to_opencv_input(imgR)

stereo    = OpenCV.StereoBM_create(numDisparities=Int32(96), blockSize=Int32(15));
disparity = OpenCV.compute(stereo, imgL, imgR);

ImageView.imshow(opencv_disparity_map_to_image(disparity))
```
"""
function image_to_opencv_input(img::AbstractMatrix)
    h, w = size(img)
    imgᵀ = transpose(img)

    # We don't know whether we got gray scales or RGB pixes, but it doesn't matter
    # as this will convert to grayscale no matter what the input was.
    grays = Gray.(reshape(imgᵀ, 1, w, h))
    opencv_input = round.(UInt8, Float32.(grays) .* 255)
end


################################################################################
#  Copyright 2020, Tom Van Acker                                               #
################################################################################
# MultiStateSystems.jl                                                         #
# A Julia package to solve multi-state system models.                          #
# See http://github.com/timmyfaraday/MultiStateSystems.jl                      #
################################################################################

## quantile
quantile_bisect(dst::AbstractDistribution, p::Real) =
    quantile_bisect(dst, p, minimum(dst), maximum(dst), 1.0e-12)
function quantile_bisect(dst::AbstractDistribution, p::Real, lx::Real, rx::Real,
                         tol::Real)
    cl, cr = cdf(dst, lx), cdf(dst, rx)
    @assert cl <= p <= cr
    while rx - lx > tol
        m = (lx + rx) / 2
        c = cdf(dst, m)
        if p > c
            cl, lx = c, m
        else
            cr, rx = c, m
    end end
    return (lx + rx) / 2
end

## exponential
"""
    Exponential

The [*exponential distribution*](http://en.wikipedia.org/wiki/Exponential_distribution)
with scale parameter `θ` and an optional weight `ω` has a probability density
function

```math
f(x; θ, ω) = \begin{cases}
                ω/θ e^{-x/θ}    & x ≥ 0, \\
                0               & x < 0.
             \end{cases}
```
"""
# struct
struct Exponential{N<:Number, R<:Real} <: AbstractDistribution{N,R}
    θ::N            # scale
    ω::R            # weight: 0.0 < ω <= 1.0
end

# constructors
Exponential() = Exponential(1.0, 1.0)
Exponential(θ::Number) = Exponential(θ, 1.0)

# shortened constructors
𝑬() = Exponential()
𝑬(θ::Number) = Exponential(θ)
𝑬(θ::Number, ω::Real) = Exponential(θ, ω)

# functions
scale(dst::Exponential)  = dst.θ
weight(dst::Exponential) = dst.ω
params(dst::Exponential) = (dst.θ, dst.ω)

rate(dst::Exponential) = 1.0 / dst.θ

minimum(dst::Exponential) = zero(dst.θ)
maximum(dst::Exponential) = (Inf)unit(dst.θ)

xv(dst::Exponential, z::Real) = z * dst.θ
quantile(dst::Exponential, p::Real)  = -xv(dst, log1p(-p))
cquantile(dst::Exponential, p::Real) = -xv(dst, log(p))
sojourn(dst::Exponential,dφ::Number,tol::Real) =
    0.0unit(dφ):dφ:cquantile(dst,tol)

function pdf(dst::Exponential, x::Number)
    θ, ω = params(dst)
    dimension(θ)==dimension(x) || return false
    if x >= (0)unit(θ)
        y = uconvert(unit(θ),x)
        ω * (1/θ) * exp(-y/θ)
    else
        zero(R)unit(θ)
    end
end
function cdf(dst::Exponential, x::Number)
    θ, ω = params(dst)
    dimension(θ)==dimension(x) || return false
    if x >= (0)unit(θ)
        y = uconvert(unit(θ),x)
        ω * (1 - exp(-y/θ))
    else
        zero(R)
    end
end
function ccdf(dst::Exponential, x::Number)
    θ, ω = params(dst)
    dimension(θ)==dimension(x) || return false
    if x >= (0)unit(θ)
        y = uconvert(unit(θ),x)
        ω * exp(-y/θ)
    else
        ω
    end
end

## weibull
"""
    Weibull

The [*Weibull distribution*](http://en.wikipedia.org/wiki/Weibull_distribution)
with scale parameter `θ`, shape parameter `α` and optional weight `ω` has a
probability density function

```math
f(x, θ, α, ω) = \\begin{cases}
                    \\frac{αω}{θ} \\cdot \\big(\\frac{x}{θ}\\big)^{α-1} \\cdot e^{-\\big(\\frac{x}{θ}\\big)^{α}}  &\\text{if:}~x ≥ 0, \\\\
                    0                                                                                &\\text{if:}~x < 0.
                \\end{cases}
```
# Constructors
| Full              | Abbr.         | Description                                               |
| :---------------- | :------------ | :-------------------------------------------------------- |
| `Weibull(θ,α,ω)`  | `𝑾(θ,α,ω)`   | full constructor                                           |
| `Weibull(θ,α)`    | `𝑾(θ,α)`     | constructor which defaults to `Weibull(θ,α,1.0)`           |
| `Weibull(θ)`      | `𝑾(θ)`       | constructor which defaults to `Weibull(θ,1.0,1.0)`         |
| `Weibull()`       | `𝑾()`        | empty constructor which defaults to `Weibull(1.0,1.0,1.0)` |
# Examples
```julia-repl
julia> Weibull()            # default Weibull distr. with θ = 1.0, α = 1.0 and ω = 1.0
julia> 𝑾(3.0u"minute")     # Weibull distr. with θ = 3.0 min, α = 1.0 and ω = 1.0
julia> 𝑾(5.0u"yr",4.0)     # Weibull distr. with θ = 5.0 yr, α = 4.0 and ω = 1.0
julia> 𝑾(10.0,0.5,0.2)     # scaled Weibull distr. with θ = 10.0, α = 0.5 and ω = 0.2
```
"""
# struct
struct Weibull{N<:Number, R<:Real} <: AbstractDistribution{N,R}
    θ::N            # scale
    α::R            # shape
    ω::R            # weight: 0.0 < ω <= 1.0
end

# constructors
Weibull() = Weibull(1.0, 1.0, 1.0)
Weibull(θ::Number) = Weibull(θ, 1.0, 1.0)
Weibull(θ::Number, α::Real) = Weibull(θ, α, 1.0)

# shortened constructors
𝑾() = Weibull()
𝑾(θ::Number) = Weibull(θ)
𝑾(θ::Number, α::Real) = Weibull(θ, α)
𝑾(θ::Number, α::Real, ω::Real) = Weibull(θ, α, ω)

# functions
scale(dst::Weibull)  = dst.θ
shape(dst::Weibull)  = dst.α
weight(dst::Weibull) = dst.ω
params(dst::Weibull) = (dst.θ, dst.α, dst.ω)

minimum(dst::Weibull) = zero(dst.θ)
maximum(dst::Weibull) = (Inf)unit(dst.θ)

xv(dst::Weibull, z::Real) = dst.θ * z ^ (1 / dst.α)
quantile(dst::Weibull, p::Real)  = xv(dst, -log1p(-p))
cquantile(dst::Weibull, p::Real) = xv(dst, -log(p))
sojourn(dst::Weibull,dφ::Number,tol::Real) = 0.0unit(dφ):dφ:cquantile(dst,tol)

function pdf(dst::Weibull, x::Number)
    θ, α, ω = params(dst)
    dimension(θ)==dimension(x) || return false
    if x >= (0)unit(θ)
        y = (uconvert(unit(θ),x) + 1e-10unit(θ)) / θ
        ω * (α / θ) * y^(α - 1) * exp(-y^α)
    else
        zero(R)unit(θ)
    end
end
function cdf(dst::Weibull, x::Number)
    θ, α, ω = params(dst)
    dimension(θ)==dimension(x) || return false
    if x >= (0)unit(θ)
        y = uconvert(unit(θ),x) / θ
        ustrip(ω * (1 - exp(-y^α)))
    else
        zero(R)
    end
end
function ccdf(dst::Weibull, x::Number)
    θ, α, ω = params(dst)
    dimension(θ)==dimension(x) || return false
    if x >= (0)unit(θ)
        y = uconvert(unit(θ),x) / θ
        ω * exp(-y^α)
    else
        ω
    end
end

## raised cosine
# struct
"""
    Cosine

The [**raised cosine distribution**](https://en.wikipedia.org/wiki/Raised_cosine_distribution)
with mean `μ` [`Number`], maximal deviation `s` [`Number`] and optional weight
`ω` [`Real`] has a probability density function
```math
f(x, μ, σ, ω) = \\begin{cases}
                    0                                                                  &\\text{if:}~x < μ-s,        \\\\
                    \\frac{ω}{2s} \\big(1 + π \\cos\\big(\\frac{x - μ}{s}\\big)\big)   &\\text{if:}~μ-s ≤ x ≤ μ+s,  \\\\
                    0                                                                  &\\text{if:}~x > μ+s.
                \\end{cases}
```
# Constructors
| Full              | Abbr.         | Description                                               |
| :---------------- | :------------ | :-------------------------------------------------------- |
| `Cosine(μ,s,ω)`   | `𝑪(μ,s,ω)`    | full constructor                                           |
| `Cosine(μ,s)`     | `𝑪(μ,s)`      | constructor which defaults to `Cosine(μ,s,1.0)`         |
| `Cosine(μ)`       | `𝑪(μ)`        | constructor which defaults to `Cosine(μ,1.0,1.0)`         |
| `Cosine()`        | `𝑪()`         | empty constructor which defaults to `Cosine(1.0,1.0,1.0)` |
# Examples
```julia-repl
julia> Cosine()             # default Raised Cosine distr. with μ = 1.0, s = 1.0 and ω = 1.0
julia> 𝑪(3.0u"minute")      # Raised Cosine distr. with μ = 3.0 min, s = 1.0 min and ω = 1.0
julia> 𝑪(5.0u"yr",4.0u"d")  # Raised Cosine distr. with μ = 5.0 yr, s = 4.0 d and ω = 1.0
julia> 𝑪(10.0,0.5,0.2)      # scaled Raised Cosine distr. with μ = 10.0, σ = 0.5 and ω = 0.2
```
"""
struct Cosine{N<:Number, R<:Real} <: AbstractDistribution{N,R}
    μ::N            # mean
    s::N            # maximal deviation
    ω::R            # weight
end

# constructors
Cosine() = Cosine(1.0, 1.0, 1.0)
Cosine(μ::Number) = Cosine(μ, oneunit(μ), 1.0)
Cosine(μ::Number, s::Number) = Cosine(μ, uconvert(unit(μ), s), 1.0)

# shortened constructors
𝑪() = Cosine()
𝑪(μ::Number) = Cosine(μ)
𝑪(μ::Number, s::Number) = Cosine(μ, uconvert(unit(μ),s))
𝑪(μ::Number, s::Number, ω::Real) = Cosine(μ, uconvert(unit(μ), s), ω)

# functions
mean(dst::Cosine)   = dst.μ
dev(dst::Cosine)    = dst.s
weight(dst::Cosine) = dst.ω
params(dst::Cosine) = (dst.μ, dst.s, dst.ω)

minimum(dst::Cosine) = dst.μ-dst.s
maximum(dst::Cosine) = dst.μ+dst.s

quantile(dst::Cosine, p::Real) = quantile_bisect(dst, p)
cquantile(dst::Cosine, p::Real) = quantile(dst, 1-p)
sojourn(dst::Cosine, dφ::Number, tol::Real) =
    dst.μ-dst.s:dφ:dst.μ+dst.s+dφ

function pdf(dst::Cosine, x::Number)
    μ, s, ω = params(dst)
    dimension(μ)==dimension(σ)==dimension(x) || return false
    if μ-s <= x <= μ+s
        y = (uconvert(unit(μ),x) - μ) / (2 * s)
        ω / (2 * s) * (1 + cospi(y))
    else
        zero(Real)/oneunit(μ)
    end
end
function cdf(dst::Cosine, x::Number)
    μ, s, ω = params(dst)
    dimension(μ)==dimension(σ)==dimension(x) || return false
    if μ-s > x
        zero(Real)
    elseif μ-s <= x <= μ+s
        y = (uconvert(unit(μ),x) - μ) / s
        ω / 2 * (1 + y + sinpi(y) / π)
    elseif x > μ+s
        ω
    end
end
function ccdf(dst::Cosine, x::Number)
    μ, s, ω = params(dst)
    dimension(μ)==dimension(σ)==dimension(x) || return false
    if μ-s > x
        ω
    elseif μ-s <= x <= μ+s
        y = (μ - uconvert(unit(μ), x)) / s
        ω / 2 * (1 + y + sinpi(y) / π)
    elseif x > μ+s
        zero(Real)
    end
end

## dirac
"""
    Dirac

The [*Dirac distribution*](https://en.wikipedia.org/wiki/Dirac_delta_function)
with offset `o` and optional weight `ω` has a probability density function

```math
f(x; o, ω) = \begin{cases}
                ω   & x = 0, \\
                0   & x ≠ 0.
             \end{cases}
```
"""
# struct
struct Dirac{N<:Number, R<:Real} <: AbstractDistribution{N,R}
    o::N            # offset
    ω::R            # weight: 0.0 < ω <= 1.0
end

# constructors
Dirac() = Dirac(0.0, 1.0)
Dirac(o::Number) = Dirac(o, 1.0)

# shortened constructors
𝑫() = Dirac()
𝑫(o::Number) = Dirac(o)
𝑫(o::Number, ω::Real) = Dirac(o, ω)

# functions
offset(dst::Dirac) = dst.o
weight(dst::Dirac) = dst.ω
params(dst::Dirac) = (dst.o, dst.ω)

minimum(dst::Dirac) = dst.o
maximum(dst::Dirac) = dst.o

quantile(dst::Dirac, p::Real)  = dst.o
cquantile(dst::Dirac, p::Real) = dst.o

function pdf(dst::Dirac, x::Number)
    o, ω = params(dst)
    dimension(o)==dimension(x) || return false
    if x == o
        ω
    else
        zero(R)
    end
end
function cdf(dst::Dirac, x::Number)
    o, ω = params(dst)
    dimension(o)==dimension(x) || return false
    if x > o
        ω
    else
        zero(R)
    end
end
function ccdf(dst::Dirac, x::Number)
    o, ω = params(dst)
    dimension(o)==dimension(x) || return false
    if x > o
        zero(R)
    else
        ω
    end
end

## uniform
"""
    Uniform

The [*continuous uniform distribution*](http://en.wikipedia.org/wiki/Uniform_distribution_(continuous))
over an interval `[a, b]` and with optional weight `ω` has a probability density
function

```math
f(x; a, b, ω) = \begin{cases}
                    ω/(b - a)   & a ≤ x ≤ b, \\
                    0           & a > x || x > b.
                \end{cases}
```
"""
# struct
struct Uniform{N<:Number, R<:Real} <: AbstractDistribution{N,R}
    a::N            # start
    b::N            # end
    ω::R            # weight: 0.0 < ω <= 1.0
end

# constructors
Uniform() = Uniform(0.0, 1.0, 1.0)
Uniform(a::Number, b::Number) = Uniform(a, b, 1.0)

# shortened constructors
𝑼() = Uniform()
𝑼(a::Number, b::Number) = Uniform(a, b)
𝑼(a::Number, b::Number, ω::Real) = Uniform(a, b, ω)

# functions
fr(dst::Uniform)     = dst.a
to(dst::Uniform)     = dst.b
weight(dst::Uniform) = dst.ω
params(dst::Uniform) = (dst.a, dst.b, dst.ω)

minimum(dst::Uniform) = dst.a
maximum(dst::Uniform) = dst.b

quantile(dst::Uniform, p::Real)  = dst.a + p * (dst.b - dst.a)
cquantile(dst::Uniform, p::Real) = dst.b + p * (dst.a - dst.b)
sojourn(dst::Uniform,dφ::Number,tol::Real) = 0.0unit(dφ):dφ:cquantile(dst,tol)

function pdf(dst::Uniform, x::Number)
    a, b, ω = params(dst)
    dimension(a)==dimension(b)==dimension(x) || return false
    if a <= x <= b
        ω / (b - a)
    else
        zero(R)unit(a)
    end
end
function cdf(dst::Uniform, x::Number)
    a, b, ω = params(dst)
    dimension(a)==dimension(b)==dimension(x) || return false
    if a > x
        zero(R)
    elseif a <= x <= b
        y = uconvert(unit(a),x)
        ω + ω * (y - b) / (b - a)
    else
        ω
    end
end
function ccdf(dst::Uniform, x::Number)
    a, b, ω = params(dst)
    dimension(a)==dimension(b)==dimension(x) || return false
    if a > x
        ω
    elseif a <= x <= b
        y = uconvert(unit(μ),x)
        ω * (b - y) / (b - a)
    else
        zero(R)
    end
end


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using Random
Random.seed!(233)

function sparse_concate(A1, A2,hcat_::Bool)
    m1,n1 = size(A1)
    m2,n2 = size(A2)
    if !isa(A1, SparseTensor); A1 = SparseTensor(A1); end
    if !isa(A2, SparseTensor); A2 = SparseTensor(A2); end
    ii1,jj1,vv1 = find(A1)
    ii2,jj2,vv2 = find(A2)
    sparse_concate_ = load_op_and_grad("./build/libSparseConcate","sparse_concate", multiple=true)
    ii1,jj1,vv1,m1_,n1_,ii2,jj2,vv2,m2_,n2_ = convert_to_tensor([ii1,jj1,vv1,m1,n1,ii2,jj2,vv2,m2,n2], [Int64,Int64,Float64,Int32,Int32,Int64,Int64,Float64,Int32,Int32])
    ii,jj,vv = sparse_concate_(ii1,jj1,vv1,m1_,n1_,ii2,jj2,vv2,m2_,n2_,constant(hcat_))
    if hcat_
        SparseTensor(ii,jj,vv, m1, n1+n2)
    else
        SparseTensor(ii,jj,vv,m1+m2,n1)
    end
end

# TODO: specify your input parameters
A1 = sprand(10,10,0.4)
A2 = sprand(10,10,0.3)
u = sparse_concate(A1, A2, true)
sess = Session(); init(sess)
@show run(sess, u)-[A1 A2]

u = sparse_concate(A1, A2, false)
sess = Session(); init(sess)
@show run(sess, u)-[A1;A2]

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(m)
    A1 = SparseTensor(ii1,jj1,m,10,10)
    A2 = SparseTensor(ii2,jj2,m,10,10)
    return sum(sparse_concate(A1, A2, false))^2
end

# TODO: change `m_` and `v_` to appropriate values
ii1 = rand(1:10,10)
jj1 = rand(1:10,10)
ii2 = rand(1:10,10)
jj2 = rand(1:10,10)
vv1 = rand(10)
vv2 = rand(10)
m_ = constant(rand(10))
v_ = rand(10)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


using DiffGPFlux
using Measurements
using Plots
using Random: randperm
using Printf

# build training and testing data,
# for X array along the 1st dimension is the feature,
# and along the 2nd dimension is the data samples,
# y array is required to be a Vector.
X = reshape(Array(range(0,2π,length=100)), 1, 100)
y = dropdims(sin.(X) + 0.01 .* randn(1, 100), dims=1)
∂y = dropdims(cos.(X) + 0.01 .* randn(1, 100), dims=1)

index = randperm(100)
train_index = index[1:20]
test_index = index[21:end]
X_train, y_train, ∂y_train = X[:, train_index], y[train_index], ∂y[train_index]
X_test, y_test, ∂y_test = X[:, test_index], y[test_index], ∂y[test_index]

γ0 = seed_duals([0.5], Float64, 1)
β = 0.1
GP = GaussProcess(γ0, β, rbf, ∇ₓrbf)
θ = [GP.γ]
indexes = [[1]]

function loss(gp, X̄, ȳ, ∂ȳ, x)
    ∂ₓf = x -> ∇ₓpredict(gp, X̄, ȳ, x)
    ∂ₓy = dropdims(∂ₓf(X̄), dims=1)
    return mean((∂ₓy-∂ȳ).^2)
end

opt = ADAM(0.01)

plt = plot(1, marker=2)
@gif for i in 1:300
    L = loss(GP, X_train, y_train, ∂y_train, X_train)
    if i%25 == 1
        @printf("Step %d, train error %.4f, test error %.4f\n",
                i, L.value,
                (loss(GP, X_test, y_test, ∂y_test, X_test)).value)
    end
    push!(plt, L.value)
    θ̄_array = [g for g in L.partials]
    θ̄ = grads(θ̄_array, indexes)
    update!(opt, θ, θ̄, 1)
    end every 10

@show(GP.γ)

γ_values = [GP.γ[1].value]
non_dual_GP = GaussProcess(γ_values, β, rbf, ∇ₓrbf)
pred_y, pred_σ = predict(non_dual_GP, X_train, y_train, X)
pred_∂y = ∇ₓpredict(non_dual_GP, X_train, y_train, X)
plot(X', pred_y.±pred_σ, color=:blue)
plot!(X', y, seriestype=:scatter, color=:red)

plot(X', pred_∂y', color=:green)
plot!(X', ∂y, seriestype=:scatter, color=:red)


@doc raw"""
PathIntegralLinear is a path integral along a linear path

```math
\phi (\tau; q^-, q^+) = (1-\tau) q^- + \tau q^+ .
```
"""
struct PathIntegralLinear <: PathIntegral end

# function evaluate(path::PathIntegralLinear{T}, i::Int, τ::T) where {T}
#     if i == 1
#         return one(T)-τ
#     elseif i == 2
#         return τ
#     else
#         error("PathIntegralLinear has only two basis functions, not ", k, ".")
#     end
# end

# function derivative(path::PathIntegralLinear{T}, i::Int, τ::T) where {T}
#     if i == 1
#         return -one(T)
#     elseif i == 2
#         return +one(T)
#     else
#         error("PathIntegralLinear has only two basis functions, not ", k, ".")
#     end
# end

evaluate_l(path::PathIntegralLinear, τ::T) where {T} = one(T)-τ
evaluate_r(path::PathIntegralLinear, τ::T) where {T} = τ

derivative_l(path::PathIntegralLinear, τ::T) where {T} = -one(T)
derivative_r(path::PathIntegralLinear, τ::T) where {T} = +one(T)


function banbks_trans(a::Array,n::Int64,m1::Int64,m2::Int64,al::Array,indx::Vector{Int64},b::Vector)
mm=m1+m2+1
@assert(size(a)  == (mm,n))
@assert(size(al) == (m1,n))
@assert(length(indx) == n)
@assert(length(b) == n)
l=m1
for k=1:n
  i=indx[k]
  if (i != k)
    dum  = b[k]
    b[k] = b[i]
    b[i] = dum
  end
  if (l < n)
    l += 1
  end
  for i=k+1:l
    b[i] -= al[i-k,k]*b[k]
  end
end
l=1
for i in n:-1:1
  dum=b[i]
  for k=2:l
    dum -= a[k,i]*b[k+i-1]
  end
  b[i]=dum/a[1,i]
  if (l < mm)
    l += 1
  end
end
return
end


# Linear programming in matrix
# Version: 5.1
# Author: Edward J. Xu, edxu96@outlook.com
# Date: April 5th, 2019
# ----------------------------------------------------------------------------------------------------------------------
# using JuMP
# using GLPKMathProgInterface
function lp(n_x, vec_c, vec_b, mat_a)
    # Linear Programming in Matrix Form
    # vec_c: [column vector] coefficient vector in objective function
    # vec_b: [column vector] right-hand-side coefficient vector in constraint
    # mat_a:
    # Edward J. Xu, 190405
    println("-------------------------------------------------------------------------\n",
            "------------------------ 1/2. Begin Optimization ------------------------\n",
            "-------------------------------------------------------------------------\n")
    model = Model(solver = GLPKSolverLP())
    @variable(model, vec_x[1: n_x] >= 0)
    @objective(model, Min, (transpose(vec_c) * vec_x)[1])
    constraintsForDual = @constraint(model, mat_a * vec_x .>= vec_b)
    solve(model)
    vec_result_u = getdual(constraintsForDual)
    obj = getobjectivevalue(model)
    vec_result_x = getvalue(vec_x)
    println("-------------------------------------------------------------------------\n",
            "-------------------------- 2/2. Nominal Ending --------------------------\n",
            "-------------------------------------------------------------------------\n")
    return (obj, vec_result_x, vec_result_u)
end


using PyPlot

nt = 2048
t = collect(linspace(0,10,nt))
s0 = 1.0
period = 0.3
om0 = 2pi/period
q = 1000.0
alpha = s0*om0*(1 + im/sqrt(4.*q^2-1.))
beta = om0/2/q*(1+im*sqrt(4.*q^2-1.))

# Now build matrix & Cholesky:
mat = zeros(nt,nt)
for i=1:nt
  for j=1:nt
    dt = abs.(t[i]-t[j])
    mat[i,j]=exp(-real(beta)*dt)*(real(alpha)*cos(imag(beta)*dt)+imag(alpha)*sin(imag(beta)*dt))
  end
end
chmat=transpose(chol(mat))

nfft = nt
ndev = randn(nt)
sim=*(chmat,ndev)
#plot(t,sim,color="blue")

numuniquepts = ceil(Int64,(nfft+1)/2)

spec = fft(sim)
spec = spec[1:numuniquepts]
spec = abs.(spec).^2/nfft
if rem(nfft,2) == 1
  spec[2:end]=spec[2:end]*2
else
  spec[2:end-1]=spec[2:end-1]*2
end

f = (0:numuniquepts-1)/(t[2]-t[1])/nfft

loglog(f,spec)

# Now compute noisy spectrum:
noise = 50.0
simn = sim + noise*randn(nt)

specn = fft(simn)
specn = specn[1:numuniquepts]
specn = abs.(specn).^2/nfft
if rem(nfft,2) == 1
  specn[2:end]=specn[2:end]*2
else
  specn[2:end-1]=specn[2:end-1]*2
end

loglog(f,specn)
om = 2pi*f
spect = s0*om0^4./((om.^2-om0^2).^2+om0^2.*om.^2./q^2)
loglog(f,spect)


include("solenoid.jl")
include("clvs.jl")
include("lss.jl")
using Test
function test_dsolenoid()
	s = [1.,4]
	m = 10
	u_trj = rand(m,3)
	eps = 1.e-5
	du_x = (solenoid(u_trj .+ eps.*[1 0. 0.],
						s, 1)[2,:,:] -
				solenoid(u_trj .- eps.*[1 0. 0.],
						s, 1)[2,:,:])./(2*eps)
	du_y = (solenoid(u_trj .+ eps.*[0 1. 0.],
						s, 1)[2,:,:] -
				solenoid(u_trj .- eps.*[0 1. 0.],
						s, 1)[2,:,:])./(2*eps)
	du_z = (solenoid(u_trj .+ eps.*[0 0. 1.],
						s, 1)[2,:,:] -
				solenoid(u_trj .- eps.*[0 0. 1.],
						s, 1)[2,:,:])./(2*eps)
	du_fd = reshape(collect([du_x; du_y; du_z]),
						3, 3, m)
	du = dsolenoid(u_trj, s)
	@assert all(isapprox.(du_fd, du, rtol=1.e-2)) == 1
end
function test_les()
	s = [1., 4.]
	m = 1
	n = 25000
	u0 = rand(3,m)
	u_trj = lorenz63(u0, s, n)[:,:,1]
	du_trj = dlorenz63(u_trj, s)
	du_trj = permutedims(du_trj,[2,3,1])
	les, clv_trj = clvs(du_trj,3)
	println(les/0.005)
	@assert all(isapprox.(les/0.005, [0.91, 0.,
		-14.572], rtol=1.e-1)) == 1
end
function test_perturbation()
	s = [1.,4.]
	n = 10
	eps = 1.e-4
	u = rand(n,3)
	dFds_fd = (solenoid(u, s .+ eps.*[1,0], 1)[2,:,:] -
			   solenoid(u, s .- eps.*[1,0], 1)[2,:,:])/
			   (2*eps)
	dFds = perturbation(u, s)
	@test all(isapprox.(dFds_fd - dFds, 0., atol=1.e-8)) ==
	1
end
function test_lss()
	s = [3.,4.]
    m = 1
    n = 1000
    n_runup = 1000
    d = 3
    d_u = 1

	u0 = rand(m,d)
	u_init = transpose(solenoid(u0, s, n_runup
								)[end,:,:])
   	n_samples = 10
	dJds = zeros(n_samples)
	vsh = zeros(d,n+1)
	f = zeros(d,n+1)
	u = zeros(n+1,d)
	for i=1:n_samples
		println("Starting LSS, sample ", i)
		u_trj = solenoid(u_init, s, n)[:,:,1]
		du_trj = dsolenoid(u_trj, s)
		X = perturbation(u_trj,s) #ith col in T_{u_{i+1}} M
		J = objective(u_trj,s)
		dJ = dobjective(u_trj, s)
		v, dJds[i] = lss(u_trj, du_trj, X, f, J, dJ, s, d_u)
		vsh .= v
		u .= u_trj
		u_init .= reshape(u_trj[end,:],1,3)
		println(dJds[i])
	end
	# check with analytical shadowing direction
	# which is [r,0,0]
	n_spinup = 100
	u = u[n_spinup+1:end,:]
	vsh = vsh[:, n_spinup+1:end]
	n = n - n_spinup
	x, y, z = view(u,:,1), view(u,:,2),
				  view(u,:,3)
	r = sqrt.(x.*x .+ y.*y .+ z.*z)
	t_dir = [-y./r  x./r  zeros(n+1)]'
	r_dir = [x./r   y./r  zeros(n+1)]'
	prj_theta = [dot(t_dir[:,i], vsh[:,i]) for i=1:n]
	prj_z = [dot([0,0,1.], vsh[:,i]) for i=1:n]
	#@test all(isapprox.(prj_theta, 0., atol=0.01)) == 1
	#@test all(isapprox.(prj_z, 0., atol=0.01)) == 1
	@test isapprox((sum(dJds)/n_samples),1.0,rtol=0.1)
	return vsh, dJds
end


struct SpectralTransform{NF<:AbstractFloat}
    # SIZE OF SPECTRAL GRID
    trunc::Int      # Spectral truncation
    mx::Int         # Number of zonal wavenumbers
    nx::Int         # Number of total wavenumbers

    # FFT plans
    # rfft_plan::FFTW.rFFTWPlan
    # irfft_plan::AbstractFFTs.ScaledPlan

    # LEGENDRE ARRAYS
    leg_weight::Array{NF,1}          # Legendre weights
    nsh2::Array{Int,1}              # What's this?
    leg_poly::Array{Complex{NF},3}   # Legendre polynomials

    # HARMONIC AND BIHARMONIC DIFFUSION
    ∇²::Array{NF,2}          # Laplacian = l*(l+1)/(R_earth^2)
    ∇⁻²::Array{NF,2}         # inverse Laplacian
    ∇⁴::Array{NF,2}          # Laplacian squared, for biharmonic diffusion

    # Quantities required by functions grad, uvspec, and vds
    gradx::Array{NF,1}
    uvdx::Array{NF,2}
    uvdym::Array{NF,2}
    uvdyp::Array{NF,2}
    gradym::Array{NF,2}
    gradyp::Array{NF,2}
    vddym::Array{NF,2}
    vddyp::Array{NF,2}
end

struct GeoSpectral{NF<:AbstractFloat}
    geometry::Geometry{NF}
    spectral::SpectralTransform{NF}
end

function GeoSpectral(P::Parameters)
    G = Geometry(P)
    S = SpectralTransform(P,G)
    return GeoSpectral{P.NF}(G,S)
end

function SpectralTransform(P::Parameters,G::Geometry)

    @unpack nlon, nlat, nlon_half, nlat_half = G
    @unpack coslat_NH = G
    @unpack R_earth, trunc = P

    # SIZE OF SPECTRAL GRID
    mx = trunc+1
    nx = trunc+1

    # PLAN THE FFTs
    # rfft_plan = plan_rfft(rand(NF,nlon))
    # irfft_plan = plan_irfft(rand(Complex{NF},nlon_half+1),nlon)

    # LEGENDRE WEIGHTS from pole to equator (=first half or array)
    leg_weight = gausslegendre(nlat)[2][1:nlat_half]

    # Spectral packing of speedy is m',n', where m' = m, n' = m+n with m,n being the
    # conventional wavenumbers. Due to Julia's 1-based indexing subtract two, as
    # the Legendre polynomials start with
    nsh2 = zeros(Int, nx)
    for n in 1:nx
        for m in 1:mx
            if m + n - 2 <= trunc + 1
                nsh2[n] = nsh2[n] + 1
            end
        end
    end

    # Epsilon-factors for the recurrence relation of the associated Legendre polynomials.
    ε,ε⁻¹ = ε_recurrence(mx,nx)

    # Generate associated Legendre polynomials
    # get_legendre_poly computes the polynomials at a particular latitiude
    leg_poly = zeros(mx, nx, nlat_half)
    for j in 1:nlat_half
        leg_poly[:,:,j] = legendre_polynomials(j,ε,ε⁻¹,mx,nx,G)
    end

    # leg_poly = zeros(mx, nx, nlat_half)
    # for j in 1:nlat_half
    #     leg_poly[:,:,j] = AssociatedLegendrePolynomials.λlm(0:mx-1,0:nx-1,coslat_NH[j])
    # end

    # LAPLACIANS for harmonic & biharmonic diffusion
    ∇²,∇⁻²,∇⁴ = Laplacians(mx,nx,R_earth)

    gradx   = zeros(mx)          #TODO what's this?
    uvdx    = zeros(mx, nx)
    uvdym   = zeros(mx, nx)
    uvdyp   = zeros(mx, nx)
    gradym  = zeros(mx, nx)
    gradyp  = zeros(mx, nx)
    vddym   = zeros(mx, nx)
    vddyp   = zeros(mx, nx)

    for m in 1:mx
        for n in 1:nx
            m1  = m - 1
            m2  = m1 + 1
            el1 = m + n - 2
            if n == 1
                gradx[m]   = m1/R_earth
                uvdx[m,1]  = -R_earth/(m1 + 1)
                uvdym[m,1] = 0.0
                vddym[m,1] = 0.0
            else
                uvdx[m,n]   = -R_earth*m1/(el1*(el1 + 1.0))
                gradym[m,n] = (el1 - 1.0)*ε[m2,n]/R_earth
                uvdym[m,n]  = -R_earth*ε[m2,n]/el1
                vddym[m,n]  = (el1 + 1.0)*ε[m2,n]/R_earth
            end
            gradyp[m,n] = (el1 + 2.0)*ε[m2,n+1]/R_earth
            uvdyp[m,n]  = -R_earth*ε[m2,n+1]/(el1 + 1.0)
            vddyp[m,n]  = el1*ε[m2,n+1]/R_earth
        end
    end

    SpectralTransform{P.NF}(    trunc,mx,nx,
                                # rfft_plan,irfft_plan,
                                leg_weight,nsh2,leg_poly,∇²,∇⁻²,∇⁴,
                                gradx,uvdx,uvdym,uvdyp,gradym,gradyp,vddym,vddyp)
end

"""
Laplacian operator in spectral space via element-wise matrix-matrix multiplication.
"""
function ∇²(A::Array{Complex{NF},2},
            G::GeoSpectral{NF}) where {NF<:AbstractFloat}
    return -G.spectral.∇².*A
end

"""
In-place version of ∇².
"""
function ∇²!(   Out::Array{Complex{NF},2},
                In::Array{Complex{NF},2},
                G::GeoSpectral{NF}) where {NF<:AbstractFloat}

    @unpack ∇² = G.spectral

    mx,nx = size(In)
    @boundscheck (mx,nx) == size(Out) || throw(BoundsError())
    @boundscheck (mx,nx) == size(∇²) || throw(BoundsError())

    for n in 1:nx
        for m in 1:mx
            @inbounds Out[m,n] = -∇²[m,n]*In[m,n]
        end
    end
end

"""
Inverse Laplacian in spectral space via element-wise matrix-matrix multiplication.
"""
function ∇⁻²(   A::Array{Complex{NF},2},
                G::GeoSpectral{NF}) where {NF<:AbstractFloat}
    return -G.spectral.∇⁻².*A
end

"""
Transform a spectral array into grid-point space.
"""
function gridded(  input::Array{Complex{NF},2},
                   G::GeoSpectral{NF}) where {NF<:AbstractFloat}
    return fourier_inverse(legendre_inverse(input,G),G)
end

"""
Transform a gridded array into spectral space.
"""
function spectral(  input::Array{NF,2},
                    G::GeoSpectral{NF}) where {NF<:AbstractFloat}
    return legendre(fourier(input,G),G)
end

function grad!( ψ::Array{NF,2},
                psdx::Array{Complex{NF},2},
                psdy::Array{NF,2},
                G::GeoSpectral{NF}) where {NF<:AbstractFloat}

    #TODO boundscheck

    @unpack trunc, mx, nx = G.spectral
    @unpack gradx, gradyp, gradym = G.spectral

    for n in 1:nx
        psdx[:,n] = gradx.*ψ[:,n]*im
    end

    for m in 1:mx
        psdy[m,1]  =  gradyp[m,1]*ψ[m,2]
        psdy[m,nx] = -gradym[m,nx]*ψ[m,trunc+1]
    end

    for n in 2:trunc+1
        for m in 1:mx
            psdy[m,n] = -gradym[m,n]*ψ[m,n-1] + gradyp[m,n]*ψ[m,n+1]
        end
    end
end

function vds!(  ucosm::Array{NF,2},
                vcosm::Array{NF,2},
                vorm::Array{NF,2},
                divm::Array{NF,2},
                G::GeoSpectral{NF}) where {NF<:AbstractFloat}

    #TODO boundscheck

    @unpack trunc, mx, nx = G.spectral
    @unpack gradx, vddym, vddyp = G.spectral

    #TODO preallocate in a diagnosticvars struct
    zp = zeros(Complex{NF}, mx,nx)
    zc = zeros(Complex{NF}, mx,nx)

    for n in 1:nx
        zp[:,n] = gradx.*ucosm[:,n]*im
        zc[:,n] = gradx.*vcosm[:,n]*im
    end

    for m in 1:mx
        #TODO this has an implicit conversion to complex{NF}, issue?
        vorm[m,1]  = zc[m,1] - vddyp[m,1]*ucosm[m,2]
        vorm[m,nx] = vddym[m,nx]*ucosm[m,trunc+1]
        divm[m,1]  = zp[m,1] + vddyp[m,1]*vcosm[m,2]
        divm[m,nx] = -vddym[m,nx]*vcosm[m,trunc+1]
    end

    for n in 2:trunc+1
        for m in 1:mx
            #TODO same here
            vorm[m,n] =  vddym[m,n]*ucosm[m,n-1] - vddyp[m,n]*ucosm[m,n+1] + zc[m,n]
            divm[m,n] = -vddym[m,n]*vcosm[m,n-1] + vddyp[m,n]*vcosm[m,n+1] + zp[m,n]
        end
    end
end

function uvspec!(   vorm::Array{NF,2},
                    divm::Array{NF,2},
                    ucosm::Array{NF,2},
                    vcosm::Array{NF,2},
                    G::GeoSpectral{NF}) where {NF<:AbstractFloat}

    #TODO boundscheck

    @unpack trunc, mx, nx = G.spectral
    @unpack uvdx, uvdyp, uvdym = G.spectral

    #TODO preallocate elsewhere
    zp = uvdx.*vorm*im
    zc = uvdx.*divm*im

    for m in 1:mx
        ucosm[m,1]  =  zc[m,1] - uvdyp[m,1]*vorm[m,2]
        ucosm[m,nx] =  uvdym[m,nx]*vorm[m,trunc+1]
        vcosm[m,1]  =  zp[m,1] + uvdyp[m,1]*divm[m,2]
        vcosm[m,nx] = -uvdym[m,nx]*divm[m,trunc+1]
    end

    for n in 2:trunc+1
        for m in 1:mx
          vcosm[m,n] = -uvdym[m,n]*divm[m,n-1] + uvdyp[m,n]*divm[m,n+1] + zp[m,n]
          ucosm[m,n] =  uvdym[m,n]*vorm[m,n-1] - uvdyp[m,n]*vorm[m,n+1] + zc[m,n]
        end
    end
end

function vdspec!(   ug::Array{NF,2},
                    vg::Array{NF,2},
                    vorm::Array{NF,2},
                    divm::Array{NF,2},
                    kcos::Bool,
                    G::GeoSpectral{NF}) where {NF<:AbstractFloat}

    #TODO boundscheck

    @unpack nlat, nlon, cosgr, cosgr2 = G.geometry

    #TODO preallocate elsewhere
    ug1 = zeros(NF, nlon, nlat)
    vg1 = zeros(NF, nlon, nlat)

    # either cosgr or cosgr2
    cosgr = kcos ? cosgr : cosgr2

    for j in 1:nlat
        for i in 1:nlon
            ug1[i,j] = ug[i,j]*cosgr[j]
            vg1[i,j] = vg[i,j]*cosgr[j]
        end
    end

    #TODO add spectral_trans and geometry as arguments
    specu = spectral(ug1,G)
    specv = spectral(vg1,G)
    vds!(specu, specv, vorm, divm)
end

"""Truncate spectral field by seting the spectral coefficients of the lower right triangle to zero. """
function spectral_truncation!(  A::AbstractArray{Complex{NF},2},    # spectral field to be truncated
                                trunc::Int                          # truncate to total wave number `trunc`
                                ) where {NF<:AbstractFloat}         # number format NF
    m,n = size(A)
    zeero = zero(Complex{NF})   # zero element of elements in A

    @inbounds for j in 1:n
        for i in 1:m
            if i+j-2 > trunc    # if total wavenumber larger than trunc
                A[i,j] = zeero
            end
        end
    end
end

"""Spectral truncation with unpacking Geospectral struct."""
function spectral_truncation!(  A::AbstractArray{Complex{NF},2},
                                G::GeoSpectral{NF}
                                ) where NF
    spectral_truncation!(A,G.spectral.trunc)    # unpack GeoSpectral struct
end

"""Spectral truncation of a grid-point field with memory allocation."""
function spectral_truncation(   input::AbstractArray{NF,2},
                                G::GeoSpectral{NF}
                                ) where NF
    input_spectral = spectral(input,G)          # allocates memory
    spectral_truncation!(input_spectral,G)      # in-place truncation

    # allocates memory to return spectrally truncated gridded field
    return gridded(input_spectral, G)
end

"""In-place version of spectral trunction of a grid-point field."""
function spectral_truncation!(  input::AbstractArray{NF,2},
                                input_spectral::AbstractArray{Complex{NF},2},
                                G::GeoSpectral{NF}
                                ) where NF
    spectral!(input_spectral,input,G)       # in-place spectral transform from input to input_spectral
    spectral_truncation!(input_spectral,G)  # in-place truncation
    gridded!(input,input_spectral,G)        # in-place backtransform
end

function spectral_truncation!(  input::Array{NF,2},
                                G::GeoSpectral{NF}
                                ) where NF
    spectral_truncation!(input,spectral(input,G),G)
end

"""
Computes Laplacian matrix-operators (for element-wise multiplication).
Laplacian, Laplacian squared and inverse Laplacian in spectral space
correspond to a multiplication with the total wavenumber:

    ∇²_n^m = N*(N+1)/R_earth^2

with N = m+n-2 the total wavenumber -2 due to 1-based indexing. The biharmonic
operator is ∇⁴ = (∇²)², the inverse Laplacian is ∇⁻² = 1 ./ ∇².
"""
function Laplacians(mx::Integer,nx::Integer,R_earth::Real)
    ∇²   = zeros(mx, nx)
    ∇⁻²  = zeros(mx, nx)
    ∇⁴   = zeros(mx, nx)

    for n in 1:nx
        for m in 1:mx
            # total wavenumber is m+n, -2 due to Julia's 1-based indexing
            N = m+n-2
            ∇²[m,n] = N*(N+1)/R_earth^2
            ∇⁴[m,n] = ∇²[m,n]^2
        end
    end

    # inverse Laplacian, the first coefficient being zero corresponds
    # to some (?, TODO) boundary conditions
    ∇⁻²[1,1] = 0.0                  # don't divide by 0
    ∇⁻²[2:end] = 1 ./ ∇⁻²[2:end]    # all other elements in matrix

    return ∇²,∇⁻²,∇⁴
end


module TestScheduleUndirectedWiringDiagrams
using Test

using LinearAlgebra: dot, tr

using Catlab.Theories: codom
using Catlab.CategoricalAlgebra.CSets
using Catlab.WiringDiagrams
using Catlab.WiringDiagrams.ScheduleUndirectedWiringDiagrams:
  composites, composite_junction, composite_ports, parent, box_parent
using Catlab.LinearAlgebra.TensorNetworks: @tensor_network

composite_junctions(x::ACSet, c) =
  composite_junction.(Ref(x), composite_ports(x, c))

# Open linear path
##################

d = @tensor_network out[v,z] = A[v,w] * B[w,x] * C[x,y] * D[y,z]
s = schedule(d, order=reverse(boxes(d)))
@test length(composites(s)) == 3
@test parent(s) == [2,3,3]
@test box_parent(s) == [3,2,1,1]

nd = to_nested_diagram(s)
@test parent(nd) == [2,3,3]
@test box_parent(nd) == [3,2,1,1]
@test composite_junctions(nd, 1:3) == [[3,5], [2,5], [1,5]]

matrices = map(randn, [(3,4), (4,5), (5,6), (6,7)])
out = eval_schedule(nd, matrices)
@test out ≈ foldr(*, matrices)

# Closed cycle
##############

d = @tensor_network out[] = A[w,x] * B[x,y] * C[y,z] * D[z,w]
s = schedule(d)
nd = to_nested_diagram(s)
@test parent(nd) == [2,3,3]
@test box_parent(nd) == [1,1,2,3]
@test map(length, composite_ports(nd, 1:3)) == [2,2,0]
@test composite_junctions(nd, 1:2) == [[1,3], [1,4]]

matrices = map(randn, [(10,5), (5,5), (5,5), (5,10)])
out = eval_schedule(nd, matrices)
@test out[] ≈ tr(foldl(*, matrices))

# Tensor product
################

d = @tensor_network out[w,x,y,z] = A[w,x] * B[y,z]
s = schedule(d)
@test parent(s) == [1]
@test box_parent(s) == [1,1]

A, B = randn((3,4)), randn((5,6))
out = eval_schedule(s, [A, B])
@test out ≈ (reshape(A, (3,4,1,1)) .* reshape(B, (1,1,5,6)))

# Frobenius inner product
#########################

d = @tensor_network out[] = A[x,y] * B[x,y]
s = schedule(d)
@test parent(s) == [1]
@test box_parent(s) == [1,1]

A, B = randn((5,5)), randn((5,5))
out = eval_schedule(s, [A, B])
@test out[] ≈ dot(vec(A), vec(B))

end


tolerance = 10^-6

S = randomGenericSymp(4)

@test size(S) == (8, 8)
@test typeof(S[1, 1]) == BigFloat
@test typeof(S) <: Symp
@test typeof(S) <: AbstractMatrix
@test typeof(Symp(S)) <: Symp
@test typeof(Matrix(S)) <: AbstractMatrix
@test isGeneric(S) == true
@test nModes(S) == 4
@test typeof(transpose(S)) <: Symp
@test typeof(S') <: Symp
@test typeof(-S) <: Symp
@test typeof(+S) <: Symp
@test isapprox( inv(S) * S, Id(4); atol = tolerance )
@test Omega(2) == Symp([0 1 0 0; -1 0 0 0; 0 0 0 1; 0 0 -1 0])
@test dsum(Id(2), Id(2), Id(2)) == Id(6)
@test S^-1 == inv(S)


using Flux
println("NNPDE_tests")
using DiffEqFlux
println("Starting Soon!")
using ModelingToolkit
using DiffEqBase
using Test, NeuralPDE
using GalacticOptim
using Optim
using CUDA
using Quadrature

using Random
Random.seed!(100)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end
CUDA.allowscalar(false)
const gpuones = cu(ones(1))

## ODE
@parameters θ
@variables u(..)
Dθ = Differential(θ)

# 1D ODE
eq = Dθ(u(θ)) ~ θ^3 + 2*θ + (θ^2)*((1+3*(θ^2))/(1+θ+(θ^3))) - u(θ)*(θ + ((1+3*(θ^2))/(1+θ+θ^3)))

# Initial and boundary conditions
bcs = [u(0.) ~ 1.0]

# Space and time domains
domains = [θ ∈ IntervalDomain(0.0,1.0)]
# Discretization
dt = 0.1
# Neural network
inner = 12
chain = FastChain(FastDense(1,inner,Flux.σ),
                  FastDense(inner,inner,Flux.σ),
                  FastDense(inner,inner,Flux.σ),
                  FastDense(inner,1),(u,p)->gpuones .* u)

initθ = DiffEqFlux.initial_params(chain) |> gpu

strategy = NeuralPDE.GridTraining(dt)
discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             strategy;
                                             init_params = initθ
                                             )

pde_system = PDESystem(eq,bcs,domains,[θ],[u])
prob = NeuralPDE.discretize(pde_system,discretization)
symprob = NeuralPDE.symbolic_discretize(pde_system,discretization)
res = GalacticOptim.solve(prob, ADAM(1e-1); cb = cb, maxiters=1000)
phi = discretization.phi

analytic_sol_func(t) = exp(-(t^2)/2)/(1+t+t^3) + t^2
ts = [domain.domain.lower:dt/10:domain.domain.upper for domain in domains][1]
u_real  = [analytic_sol_func(t) for t in ts]
u_predict  = [first(Array(phi(t,res.minimizer))) for t in ts]

@test u_predict ≈ u_real atol = 0.2

# t_plot = collect(ts)
# plot(t_plot ,u_real)
# plot!(t_plot ,u_predict)

## 1D PDE

@parameters t x
@variables u(..)
Dt = Differential(t)
Dxx = Differential(x)^2

eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
bcs = [u(0,x) ~ cos(x),
        u(t,0) ~ exp(-t),
        u(t,1) ~ exp(-t) * cos(1)]

domains = [t ∈ IntervalDomain(0.0,1.0),
          x ∈ IntervalDomain(0.0,1.0)]

pdesys = PDESystem(eq,bcs,domains,[t,x],[u])

inner = 12
chain = FastChain(FastDense(2,inner,Flux.σ),
                  FastDense(12,12,Flux.σ),
                  FastDense(inner,1),(u,p)->gpuones .* u)

quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=HCubatureJL(),
                                                   reltol = 1e-2, abstol = 1e-2,
                                                   maxiters = 50)
initθ = initial_params(chain) |>gpu
discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             quadrature_strategy;
                                             init_params = initθ)
prob = NeuralPDE.discretize(pdesys,discretization)
symprob = NeuralPDE.symbolic_discretize(pdesys,discretization)

res = GalacticOptim.solve(prob, ADAM(0.1); cb = cb, maxiters=1000)
phi = discretization.phi

u_exact = (t,x) -> exp.(-t) * cos.(x)
ts,xs = [domain.domain.lower:0.01:domain.domain.upper for domain in domains]
u_predict = reshape([first(Array(phi([t,x],res.minimizer))) for t in ts for x in xs],(length(ts),length(xs)))
u_real = reshape([u_exact(t,x) for t in ts  for x in xs ], (length(ts),length(xs)))
diff_u = abs.(u_predict .- u_real)

@test u_predict ≈ u_real atol = 0.3

# p1 = plot(ts, xs, u_real, linetype=:contourf,title = "analytic");
# p2 = plot(ts, xs, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(ts, xs, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)


function step_sim!(θ, p_θ, dt)
    p_θ[] = p_θ[] - 9.8*sin(θ[])*dt
    θ[]   = θ[]   + p_θ[]*dt

    if (θ[] > π)
        θ[] = θ[] - 2*π
    elseif (θ[] < -π)
        θ[] = θ[] + 2*π
    end
end

function advance_sim!(θ, p_θ, dt, Δt)

  Δt_partial = zero(dt)

  while Δt_partial <= Δt
    step_sim!(θ, p_θ, dt)
    Δt_partial += dt
  end

  return Δt_partial - Δt
end

function sim!(θ, p_θ; dt=1/24/4, fps = 24, scene)

    frame_duration = 1/fps
    t_overshoot = zero(dt)
    limiter = Limiter(frame_duration)

    while scene.events.window_open[]
        t_overshoot = advance_sim!(θ, p_θ, dt, frame_duration - t_overshoot)
        limiter()
    end
end


# Laplacian

"""
    laplacian!(v,w)

Evaluate the discrete Laplacian of `w` and return it as `v`. The data `w` can be
of type dual/primary nodes or edge components or edges; `v` must be of the same type.

# Example

```jldoctest
julia> w = Nodes(Dual,(8,6));

julia> v = deepcopy(w);

julia> w[4,3] = 1.0;

julia> laplacian!(v,w)
Nodes{Dual,8,6,Float64} data
Printing in grid orientation (lower left is (1,1))
6×8 Array{Float64,2}:
 0.0  0.0  0.0   0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0   0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0   1.0  0.0  0.0  0.0  0.0
 0.0  0.0  1.0  -4.0  1.0  0.0  0.0  0.0
 0.0  0.0  0.0   1.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0   0.0  0.0  0.0  0.0  0.0
```
"""
function laplacian!(out::Nodes{Dual,NX, NY}, w::Nodes{Dual,NX, NY}) where {NX, NY}
    view(out,2:NX-1,2:NY-1) .= view(w,2:NX-1,1:NY-2) .+ view(w,1:NX-2,2:NY-1) .+
                               view(w,3:NX,2:NY-1) .+ view(w,2:NX-1,3:NY) .-
                               4.0.*view(w,2:NX-1,2:NY-1)
    #@inbounds for y in 2:NY-1, x in 2:NX-1
    #    out[x,y] = w[x,y-1] + w[x-1,y] - 4w[x,y] + w[x+1,y] + w[x,y+1]
    #end
    out
end

function laplacian!(out::Nodes{Primal,NX, NY}, w::Nodes{Primal,NX, NY}) where {NX, NY}
    view(out,2:NX-2,2:NY-2) .= view(w,2:NX-2,1:NY-3) .+ view(w,1:NX-3,2:NY-2) .+
                               view(w,3:NX-1,2:NY-2) .+ view(w,2:NX-2,3:NY-1) .-
                               4.0.*view(w,2:NX-2,2:NY-2)
    #@inbounds for y in 2:NY-2, x in 2:NX-2
    #    out[x,y] = w[x,y-1] + w[x-1,y] - 4w[x,y] + w[x+1,y] + w[x,y+1]
    #end
    out
end

function laplacian!(out::XEdges{Dual,NX, NY}, w::XEdges{Dual,NX, NY}) where {NX, NY}
    view(out,2:NX-2,2:NY-1) .= view(w,2:NX-2,1:NY-2) .+ view(w,1:NX-3,2:NY-1) .+
                               view(w,3:NX-1,2:NY-1) .+ view(w,2:NX-2,3:NY) .-
                               4.0.*view(w,2:NX-2,2:NY-1)
    out
end

function laplacian!(out::YEdges{Dual,NX, NY}, w::YEdges{Dual,NX, NY}) where {NX, NY}
    view(out,2:NX-1,2:NY-2) .= view(w,2:NX-1,1:NY-3) .+ view(w,1:NX-2,2:NY-2) .+
                               view(w,3:NX,2:NY-2)   .+ view(w,2:NX-1,3:NY-1) .-
                               4.0.*view(w,2:NX-1,2:NY-2)
    out
end

function laplacian!(out::XEdges{Primal,NX, NY}, w::XEdges{Primal,NX, NY}) where {NX, NY}
   view(out,2:NX-1,2:NY-2) .= view(w,2:NX-1,1:NY-3) .+ view(w,1:NX-2,2:NY-2) .+
                              view(w,3:NX,2:NY-2) .+ view(w,2:NX-1,3:NY-1) .-
                              4.0.*view(w,2:NX-1,2:NY-2)
    out
end

function laplacian!(out::YEdges{Primal,NX, NY}, w::YEdges{Primal,NX, NY}) where {NX, NY}
   view(out,2:NX-2,2:NY-1) .= view(w,2:NX-2,1:NY-2) .+ view(w,1:NX-3,2:NY-1) .+
                                view(w,3:NX-1,2:NY-1) .+ view(w,2:NX-2,3:NY) .-
                               4.0.*view(w,2:NX-2,2:NY-1)
    out
end

function laplacian!(out::Edges{C,NX, NY}, w::Edges{C,NX, NY}) where {C,NX, NY}
  laplacian!(out.u,w.u)
  laplacian!(out.v,w.v)
  out
end

function laplacian!(out::EdgeGradient{C,D,NX, NY}, w::EdgeGradient{C,D,NX, NY}) where {C,D,NX, NY}
  laplacian!(out.dudx,w.dudx)
  laplacian!(out.dvdx,w.dvdx)
  laplacian!(out.dudy,w.dudy)
  laplacian!(out.dvdy,w.dvdy)
  out
end

#=
function laplacian!(out::Edges{Dual,NX, NY}, w::Edges{Dual,NX, NY}) where {NX, NY}
  view(out.u,2:NX-2,2:NY-1) .= view(w.u,2:NX-2,1:NY-2) .+ view(w.u,1:NX-3,2:NY-1) .+
                          view(w.u,3:NX-1,2:NY-1) .+ view(w.u,2:NX-2,3:NY) .-
                          4.0.*view(w.u,2:NX-2,2:NY-1)
  view(out.v,2:NX-1,2:NY-2) .= view(w.v,2:NX-1,1:NY-3) .+ view(w.v,1:NX-2,2:NY-2) .+
                               view(w.v,3:NX,2:NY-2) .+ view(w.v,2:NX-1,3:NY-1) .-
                               4.0.*view(w.v,2:NX-1,2:NY-2)
  #@inbounds for y in 2:NY-1, x in 2:NX-2
  #    out.u[x,y] = w.u[x,y-1] + w.u[x-1,y] - 4w.u[x,y] + w.u[x+1,y] + w.u[x,y+1]
  #end
  #@inbounds for y in 2:NY-2, x in 2:NX-1
  #    out.v[x,y] = w.v[x,y-1] + w.v[x-1,y] - 4w.v[x,y] + w.v[x+1,y] + w.v[x,y+1]
  #end
  out
end
=#

#=
function laplacian!(out::Edges{Primal,NX, NY}, w::Edges{Primal,NX, NY}) where {NX, NY}
  view(out.u,2:NX-1,2:NY-2) .= view(w.u,2:NX-1,1:NY-3) .+ view(w.u,1:NX-2,2:NY-2) .+
                               view(w.u,3:NX,2:NY-2) .+ view(w.u,2:NX-1,3:NY-1) .-
                              4.0.*view(w.u,2:NX-1,2:NY-2)
  view(out.v,2:NX-2,2:NY-1) .= view(w.v,2:NX-2,1:NY-2) .+ view(w.v,1:NX-3,2:NY-1) .+
                                view(w.v,3:NX-1,2:NY-1) .+ view(w.v,2:NX-2,3:NY) .-
                               4.0.*view(w.v,2:NX-2,2:NY-1)
  #@inbounds for y in 2:NY-2, x in 2:NX-1
  #    out.u[x,y] = w.u[x,y-1] + w.u[x-1,y] - 4w.u[x,y] + w.u[x+1,y] + w.u[x,y+1]
  #end
  #@inbounds for y in 2:NY-1, x in 2:NX-2
  #    out.v[x,y] = w.v[x,y-1] + w.v[x-1,y] - 4w.v[x,y] + w.v[x+1,y] + w.v[x,y+1]
  #end
  out
end
=#

"""
    laplacian(w)

Evaluate the discrete Laplacian of `w`. The data `w` can be of type dual/primary
nodes or edges. The returned result is of the same type as `w`.

# Example

```jldoctest
julia> q = Edges(Primal,(8,6));

julia> q.u[2,2] = 1.0;

julia> laplacian(q)
Edges{Primal,8,6,Float64} data
u (in grid orientation)
5×8 Array{Float64,2}:
 0.0   0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0   0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0   1.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  -4.0  1.0  0.0  0.0  0.0  0.0  0.0
 0.0   0.0  0.0  0.0  0.0  0.0  0.0  0.0
v (in grid orientation)
6×7 Array{Float64,2}:
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
```
"""
function laplacian(w::Nodes{C,NX,NY}) where {C<:CellType,NX,NY}
  laplacian!(Nodes(C,w), w)
end

function laplacian(w::XEdges{C,NX,NY}) where {C<:CellType,NX,NY}
  laplacian!(XEdges(C,w), w)
end

function laplacian(w::YEdges{C,NX,NY}) where {C<:CellType,NX,NY}
  laplacian!(YEdges(C,w), w)
end

function laplacian(w::Edges{C,NX,NY}) where {C<:CellType,NX,NY}
  laplacian!(Edges(C,w), w)
end

function laplacian(w::EdgeGradient{C,D,NX,NY}) where {C<:CellType,D<:CellType,NX,NY}
  laplacian!(EdgeGradient(C,w), w)
end

"""
    laplacian_symm!(v,w)

Evaluate the symmetric 5-point discrete Laplacian of `w` and return it as `v`. The data `w` can be
of type dual nodes only for now. This symmetric Laplacian also evaluates the
partial Laplacians (using only available stencil data) on the ghost nodes.
"""
function laplacian_symm!(out::Nodes{Dual,NX, NY}, w::Nodes{Dual,NX, NY}) where {NX, NY}
    @inbounds for y in 2:NY-1, x in 2:NX-1
        out[x,y] = w[x,y-1] + w[x-1,y] - 4w[x,y] + w[x+1,y] + w[x,y+1]
    end
    @inbounds for y in 2:NY-1
        out[1,y]  = w[1,y-1]            - 4w[1,y] + w[2,y] + w[1,y+1]
        out[NX,y] = w[NX,y-1] + w[NX-1,y]- 4w[NX,y]        + w[NX,y+1]
    end
    @inbounds for x in 2:NX-1
        out[x,1]  = w[x-1,1] + w[x+1,1] - 4w[x,1] + w[x,2]
        out[x,NY] = w[x-1,NY]+ w[x+1,NY]- 4w[x,NY] + w[x,NY-1]
    end
    out[1,1] = -4w[1,1] + w[1,2] + w[2,1]
    out[NX,1] = -4w[NX,1] + w[NX-1,1] + w[NX,2]
    out[1,NY] = -4w[1,NY] + w[1,NY-1] + w[2,NY]
    out[NX,NY] = -4w[NX,NY] + w[NX,NY-1] + w[NX-1,NY]
    out
end

"""
    plan_laplacian(dims::Tuple,[with_inverse=false],[fftw_flags=FFTW.ESTIMATE],
                          [factor=1.0],[dx=1.0],[dtype=Float64])

Constructor to set up an operator for evaluating the discrete Laplacian on
dual or primal nodal data of dimension `dims`. If the optional keyword
`with_inverse` is set to `true`, then it also sets up the inverse Laplacian
(the lattice Green's function, LGF). These can then be applied, respectively, with
`*` and `\\` operations on data of the appropriate size. The optional
parameter `factor` is a scalar used to multiply the result of the operator and
divide the inverse. The optional parameter
`dx` is used in adjusting the uniform value of the LGF to match the behavior
of the continuous analog at large distances; this is set to 1.0 by default. The
type of data on which to act is floating point by default, but can also be ComplexF64.
This is specified with the optional parameter `dtype`

Instead of the first argument, one can also supply `w::Nodes` to specify the
size of the domain.

# Example

```jldoctest
julia> w = Nodes(Dual,(5,5));

julia> w[3,3] = 1.0;

julia> L = plan_laplacian(5,5;with_inverse=true)
Discrete Laplacian (and inverse) on a (nx = 5, ny = 5) grid acting on Float64 data with spacing 1.0

julia> s = L\\w
Nodes{Dual,5,5,Float64} data
Printing in grid orientation (lower left is (1,1))
5×5 Array{Float64,2}:
 0.16707    0.129276     0.106037     0.129276    0.16707
 0.129276   0.0609665   -0.00734343   0.0609665   0.129276
 0.106037  -0.00734343  -0.257343    -0.00734343  0.106037
 0.129276   0.0609665   -0.00734343   0.0609665   0.129276
 0.16707    0.129276     0.106037     0.129276    0.16707

julia> L*s ≈ w
true
```
"""
function plan_laplacian end

"""
    plan_laplacian!(dims::Tuple,[with_inverse=false],[fftw_flags=FFTW.ESTIMATE],
                          [factor=1.0][,nthreads=length(Sys.cpu_info())])

Same as [`plan_laplacian`](@ref), but operates in-place on data. The number of threads `threads` defaults to the number of
logical CPU cores on the system.
"""
function plan_laplacian! end

struct Laplacian{NX, NY, T, R,inplace}
    factor::T
    dx::Float64
    conv::Union{CircularConvolution{NX, NY, T},Nothing}
end

Base.eltype(::Laplacian{NX,NY,T}) where {NX,NY,T} = T

for (lf,inplace) in ((:plan_laplacian,false),
                     (:plan_laplacian!,true))
    @eval function $lf(dims::Tuple{Int,Int};
                   with_inverse = false, fftw_flags = FFTW.ESTIMATE, factor::Real = 1.0, dx = 1.0, dtype = Float64, nthreads = length(Sys.cpu_info()))
        NX, NY = dims
        if !with_inverse
            return Laplacian{NX, NY, dtype, false, $inplace}(convert(dtype,factor),convert(Float64,dx),nothing)
        end

        G = view(LGF_TABLE, 1:NX, 1:NY)
        Laplacian{NX, NY, dtype, true, $inplace}(convert(dtype,factor),convert(Float64,dx),CircularConvolution(G, fftw_flags,dtype=dtype,nthreads=nthreads))
    end

    @eval function $lf(nx::Int, ny::Int;
        with_inverse = false, fftw_flags = FFTW.ESTIMATE, factor = 1.0, dx = 1.0, dtype = Float64, nthreads = length(Sys.cpu_info()))
        $lf((nx, ny), with_inverse = with_inverse, fftw_flags = fftw_flags, factor = factor, dx = dx, dtype = dtype, nthreads = nthreads)
    end

    @eval function $lf(w::ScalarGridData;
        with_inverse = false, fftw_flags = FFTW.ESTIMATE, factor = 1.0, dx = 1.0, dtype = Float64, nthreads = length(Sys.cpu_info())) where {T<:CellType,NX,NY}
        $lf(size(w), with_inverse = with_inverse, fftw_flags = fftw_flags, factor = factor, dx = dx, dtype = dtype, nthreads = nthreads)
    end
end

function Base.show(io::IO, L::Laplacian{NX, NY, T, R, inplace}) where {NX, NY, T, R, inplace}
    nodedims = "(nx = $NX, ny = $NY)"
    inverse = R ? " (and inverse)" : ""
    isinplace = inplace ? " in-place" : ""
    print(io, "Discrete$isinplace Laplacian$inverse on a $nodedims grid acting on $T data with
               factor $(L.factor) and spacing $(L.dx)")
end

mul!(out::T, L::Laplacian, s::T) where T<:GridData = (laplacian!(out, s); out .*= L.factor)

*(L::Laplacian{MX,MY,T,R,false}, s::GridData) where {MX,MY,T,R} =
      L.factor*laplacian(s)

function (*)(L::Laplacian{MX,MY,T,R,true}, s::GridData) where {MX,MY,T,R}
    mul!(s,L,deepcopy(s))
end

#=
function ldiv!(out::Nodes{C,NX, NY,T},
                   L::Laplacian{MX, MY, T, true, inplace},
                   s::Nodes{C, NX, NY,T}) where {C <: CellType, NX, NY, MX, MY, T, inplace}

    mul!(out.data, L.conv, s.data)
    inv_factor = 1.0/L.factor

    # Adjust the behavior at large distance to match continuous kernel
    out.data .-= (sum(s.data)/2π)*(GAMMA+log(8)/2-log(L.dx))
    out.data .*= inv_factor
    out
end
=#
for (datatype) in (:Nodes, :XEdges, :YEdges)
  @eval function ldiv!(out::$datatype{C,NX, NY,T},
                   L::Laplacian{MX, MY, T, true, inplace},
                   s::$datatype{C, NX, NY,T}) where {C <: CellType, NX, NY, MX, MY, T, inplace}

    mul!(out.data, L.conv, s.data)
    inv_factor = 1.0/L.factor

    # Adjust the behavior at large distance to match continuous kernel
    out.data .-= (sum(s.data)/2π)*(GAMMA+log(8)/2-log(L.dx))
    out.data .*= inv_factor
    out
  end

  #@eval \(L::Laplacian{MX,MY,T,R,false},s::$datatype{C,NX,NY}) where {MX,MY,T,R,C <: CellType,NX,NY} =
  #  ldiv!($datatype(C,s), L, s)

end

function ldiv!(out::Edges{C,NX,NY},L::Laplacian,s::Edges{C,NX,NY}) where {C,NX,NY}
  ldiv!(out.u,L,s.u)
  ldiv!(out.v,L,s.v)
  out
end

function ldiv!(out::EdgeGradient{C,D,NX,NY},L::Laplacian,s::EdgeGradient{C,D,NX,NY}) where {C,D,NX,NY}
  ldiv!(out.dudx,L,s.dudx)
  ldiv!(out.dvdx,L,s.dvdx)
  ldiv!(out.dudy,L,s.dudy)
  ldiv!(out.dvdy,L,s.dvdy)
  out
end

\(L::Laplacian{MX,MY,T,R,false},s::G) where {MX,MY,T,R,G<:GridData} =
  ldiv!(G(), L, s)

\(L::Laplacian{MX,MY,T,R,true},s::GridData{NX,NY}) where {MX,MY,T,R,NX,NY} =
    ldiv!(s, L, deepcopy(s))

#=
\(L::Laplacian{MX,MY,T,R,false},s::Nodes{C,NX,NY}) where {MX,MY,T,R,C <: CellType,NX,NY} =
  ldiv!(Nodes(C,s), L, s)

\(L::Laplacian{MX,MY,T,R,true},s::Nodes{C,NX,NY}) where {MX,MY,T,R,C <: CellType,NX,NY} =
  ldiv!(s, L, deepcopy(s))
=#


### Day 3 - Problem 2
## Squares With Three Sides
## Author: Thanasis Georgiou

# Read triangle data
input = readdlm("triangles.csv")

# Start counting
count = 0

# For each triangle
for i = 0:1:size(input, 1)-1
    # Calculate vertical index of triangle
    # i + 1: 1, 2, 3, 4, 5, 6, 7, 8, 9...
    # v    : 1, 1, 1, 4, 4, 4, 7, 7, 7...
    v = i - (i % 3) + 1

    # Grab triangle (verI + 2 cells below)
    triangle = sort(
        view(input, v:v + 2, (i % 3) + 1)
    )

    # Check if it's a valid triangle
    count += triangle[1] + triangle[2] > triangle[3]
end

# Print result
println("Possible Triangles: $(count)")


"""
The goal of the module is to attach physical coordinaes to array elements, so `a[i,j]` could be interpreted as ``f(x_i, y_j)`` for some function ``f``.
"""
module SampledDomains

abstract type AbstractDomain end

import Base.:*
import Base.size

export CartesianDomain2D, make_centered_domain2D, dualRange, dualDomain

# TODO rewrite as parametric type
struct CartesianDomain2D<:AbstractDomain
    xrange::AbstractRange
    yrange::AbstractRange
end

getranges(dom::AbstractDomain) = [getfield(dom,n) for n in fieldnames(typeof(dom))]

size(dom::AbstractDomain) = tuple(length.(reverse(getranges(dom)))...) # we need reverse to follow column major rule
# size(dom::AbstractDomain, d) = length(getfield(dom,d))
size(dom::AbstractDomain, d) = size(dom)[d]

*(dom::CartesianDomain2D, s::Real) = CartesianDomain2D(dom.xrange * s, dom.yrange*s)
*( s::Real, dom::CartesianDomain2D) = CartesianDomain2D(dom.xrange * s, dom.yrange*s)

function make_centered_domain2D(xlength, ylength, pixelsize)
    xrange = ((1:xlength) .- (1 + xlength) / 2) .* pixelsize
    yrange = ((1:ylength) .- (1 + ylength) / 2) .* pixelsize
    CartesianDomain2D(xrange, yrange)
end

"""
    dualRange(xrange::AbstractRange, q::Int =1)

Construct range in Fourier-transform-dual domain with upsampling factor `q`. If sampling in `x` has step size Δx, the dual
domain is sampled in interval (-q/2Δx, q/2Δx).
"""
function dualRange(xrange::AbstractRange, q::Int =1)
    len = length(xrange) *q
    st = step(xrange) /q
    return UnitRange(-floor(Int, len / 2), ceil(Int, len / 2) - 1) / (st * len)

end

function dualDomain(dom::CartesianDomain2D, q::Int = 1)
    kxrange = dualRange(dom.xrange, q)
    kyrange = dualRange(dom.yrange, q)
    return CartesianDomain2D(kxrange,kyrange)
end

"""
    SampledDomain(vals::Array, dom::AbstractDomain)
    SampledDomain(f::Function, dom::AbstractDomain)

Contains two fileds: `vals` and `dom` representing sampled values of function `f` on `dom`ain.
"""
struct SampledDomain
    vals::Array
    dom::AbstractDomain
    SampledDomain(vals::Array, dom::AbstractDomain) = size(vals) != size(dom) ? ErrorException("Different size of the domain and values array") : new(vals, dom)
end

SampledDomain(f::Function, dom::AbstractDomain) = SampledDomain(map(x->f(x...), collect(Iterators.product(reverse(getranges(dom))...))), dom)

end


"""
Convert from a hypergraph to a dyadic graph by turning hyperedges into cliques.
Convert between hypergraph and bipartite representations.
"""

using MatrixNetworks
using SparseArrays
using LinearAlgebra

"""
`to_adjacency_matrix`
=====================

Convert a hypergraph to an adjacency matrix

Arguments
----------
    - `H::Hypergraphs`: The hypergraph that will be converted into an adjacency matrix
    - `simple::Bool (=false)`: Whether to exclude multi-edges and self-loops from the matrix (they are included by default)

Examples
----------
~~~~
to_adjacency_matrix(Hypergraphs([[1,2,3,4], [3,4,1], [1,4,5]], ["One", "Two", "Three", "Four", "Five"], 5, 3); simple=true)
~~~~
"""
function to_adjacency_matrix(H::Hypergraph; simple::Bool=false)
    A = zeros(Int64, H.n, H.n) # Create adjacency matrix of size n x n

    # Loop through each hyperedge and create edges in the new adjacency matrix
    for i = 1:H.m
        for j = 1:length(H.edges[i])-1
            for k = j+1:length(H.edges[i])
                index1 = H.edges[i][j]
                index2 = H.edges[i][k]
                A[index1,index2] += 1
                A[index2,index1] += 1
            end
        end
    end

    if simple
        A -= Diagonal(A) # Remove self-loops
        A = min.(A, 1) # Remove multi-edges
    end

    return A
end

"""
`hypergraph_to_matrixnetworks`
=====================

Convert a hypergraph to a MatrixNetwork

Arguments
----------
    - `H::Hypergraphs`: The hypergraph that will be converted into a MatrixNetwork
    - `simple::Bool (=false)`: Whether to exclude multi-edges and self-loops from the MatrixNetwork (they are included by default)

Examples
----------
~~~~
hypergraph_to_matrixnetwork(Hypergraphs([[1,2,3,4], [3,4,1], [1,4,5]], ["One", "Two", "Three", "Four", "Five"], 5, 3); simple=true)
~~~~
"""
function hypergraph_to_matrixnetwork(H::Hypergraph; simple=false)
    return MatrixNetwork(sparse(to_adjacency_matrix(H; simple=simple)))
end

"""
`hyper_to_bipartite`
====================

Converts a hypergraph into its bipartite representation. For all vertices v and
edges e in the hypergraph, an edge (v, e) exists in the bipartite if and only if
v ∈ e in the hypergraph.
"""
function hyper_to_bipartite(H::Hypergraph)
   di_edges::Vector{Vector{Int64}} = []

   for i = 1:h.m
      for j in h.edges[i]
         push!(di_edges, [i + h.n, j])
      end
   end

   return MatrixNetwork(di_edges, h.n + h.m)
end

"""
`is_bipartite`
==============

Checks whether a graph is bipartite by checking whether it is 2-colorable.
"""
function is_bipartite(A::MatrixNetwork)
   group1::Set{Int64} = Set()
   group2::Set{Int64} = Set()

   function color_neighbors(v::Int64, is_group1::Bool)
      for n in A.ci[A.rp[v]:A.rp[v+1]-1]
         if is_group1 # v is in group 2, so all neighbors should be in group 1
            if n in group2
               return false
            end
            push!(group1, n)
         else
            if n in group1
               return false
            end
            push!(group2, n)
         end
      end
      return true
   end

   is_group1 = true

   for i = 1:A.n
      is_group1 = !is_group1
      if !color_neighbors(i, is_group1)
         return false
      end
   end
   return (group1, group2)
end

"""
`bipartite_to_hyper`
====================

Returns the left and right projections of a bipartite graph into a hypergraph.
In the left projection, all vertices in group two that are adjacent to the same
vertex in group one are put together in a hyperedge. The right projection is the
same except groups one and two are swapped in their roles.
"""
function bipartite_to_hyper(A::MatrixNetwork)
   b = is_bipartite(A)
   if typeof(b) <: Bool
      error("ArgumentError: The network must be bipartite to be projected to a hypergraph.")
   end

   g1, g2 = collect.(b)
   left_edges = [map(x -> findfirst(y -> x == y, g2), A.ci[A.rp[i]:A.rp[i+1]-1]) for i in g1]
   right_edges = [map(x -> findfirst(y -> x == y, g1), A.ci[A.rp[j]:A.rp[j+1]-1]) for j in g2]

   return Hypergraphs(left_edges, A.vals[g2], length(g2), length(left_edges)),
          Hypergraphs(right_edges, A.vals[g1], length(g1), length(right_edges))
end


# ---
# title: 42. Trapping Rain Water
# id: problem42
# author: Indigo
# date: 2021-03-10
# difficulty: Hard
# categories: Array, Two Pointers, Stack
# link: <https://leetcode.com/problems/trapping-rain-water/description/>
# hidden: true
# ---
#
# Given `n` non-negative integers representing an elevation map where the width
# of each bar is `1`, compute how much water it can trap after raining.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2018/10/22/rainwatertrap.png)
#
#
#
#     Input: height = [0,1,0,2,1,0,1,3,2,1,2,1]
#     Output: 6
#     Explanation: The above elevation map (black section) is represented by array [0,1,0,2,1,0,1,3,2,1,2,1]. In this case, 6 units of rain water (blue section) are being trapped.
#
#
# **Example 2:**
#
#
#
#     Input: height = [4,2,0,3,2,5]
#     Output: 9
#
#
#
#
# **Constraints:**
#
#   * `n == height.length`
#   * `0 <= n <= 3 * 104`
#   * `0 <= height[i] <= 105`
#
#
## @lc code=start
using LeetCode

function trap_rain(height::Vector{Int})::Int
    csum = cumsum(height)
    pushfirst!(csum, 0)
    max_pos = argmax(height)
    len = length(height)
    res = 0
    r = 1
    for i in 2:max_pos
        if height[i] >= height[r]
            res += (i > r + 1) * (height[r] * (i - r) - (csum[i] - csum[r]))
            r = i
        end
    end
    l = len
    for j in (len - 1):-1:max_pos
        if height[j] >= height[l]
            res += (j < l - 1) * (height[l] * (l - j) - (csum[l + 1] - csum[j + 1]))
            l = j
        end
    end
    return res
end
## @lc code=end