"""
    SeisBandPass(in; <keyword arguments>)

Apply a bandpass filter to seismic data. Input and output data is 2D in tx domain. Filter is applied in fx domain.

# Arguments
- `in`: Input 2D data array in tx domain. Time is first dimension.

# Keyword arguments
- `dt=0.004`: time sampling interval in secs.
- `fa=0`: lower frequency cut [Hz]
- `fb=0`: lower frequency banpass [Hz]
- `fc=60`: upper frequency bandpass [Hz]
- `fd=80`: upper frequency cut [Hz]

# Example
```
julia> d = SeisLinearEvents(); SeisPlotAmplitude(d,125,0.004);
julia> d_filter = SeisBandPass(d;dt=0.004,fa=2,fb=8,fc=12,fd=20); SeisPlotAmplitude(d_filter,125,0.004)
julia> SeisPlot([d d_filter],title="Data and Filtered data")

```
"""
function SeisBandPass(d;dt=0.004,fa=0,fb=0,fc=60,fd=80)

    nt =size(d,1)
	dn = reshape(d,nt,:)
    v = size(dn)
	ntr = v[2]
    nf = iseven(nt) ? nt : nt + 1
	df = 1/nf/dt
	nw = round(Int,nf/2) + 1

	if(fd*dt*nf < nw)
		iw_max = round(Int,floor(fd*dt*nf))
	else
		iw_max = round(Int,floor(0.5/dt))
	end

	dp = pad_first_axis(dn,nf)
	m = fft(dp,1)/sqrt(size(dp,1))
	if fa > 0.
		m[1,:] *= 0.
	end
	for iw=2:iw_max
		f = df*(iw-1)
		if (f<fa)
			m[iw,:] *= 0.
		elseif (f >= fa && f < fb)
			m[iw,:] *= (f-fa)/(fb-fa)
		elseif (f >= fb && f <= fc)
			m[iw,:] *= 1.
		elseif (f > fc && f <= fd)
			m[iw,:] *= 1. - (f-fc)/(fd-fc)
		else
			m[iw,:] *= 0.
		end
	end
	m[iw_max:end,:] .= 0.

	# symmetries
	for iw=nw+1:nf
		m[iw,:] = conj(m[nf-iw+2,:])
	end
	dn = real(bfft(m,1)/sqrt(size(m,1)))
	dout = dn[1:nt,1:ntr]
	return reshape(dout,v);
end

"""
    SeisBandPass(in,out,parameters; <keyword arguments>)

# Arguments
- `in::String`: Input file - Seis format
- `out::String`: Output file - Seis format.
- `parameters` : list of the keyword arguments for the function SeisBandPass.

# Keyword arguments
- `group="gather"` : Options are all, some or gather
- `key=["imx","imy"]` : Defines type of gather
- `itrace=1` : Initial trace number
- `ntrace=10000` : Total number of traces to process at once

# Example
```
julia> param = Dict(:fa>=0,:fb=>0,:fc=>60,:fd=>80)
julia> SeisBandPass(filein,fileout, param,group="all")
```

"""
function SeisBandPass(in::String,out::String,parameters;group="gather",key=["imx","imy"],itrace=1,ntrace=10000)

    fa = get(parameters,:fa,0)
    fb = get(parameters,:fb,0)
    fc = get(parameters,:fc,60)
    fd = get(parameters,:fd,80)

    ext = SeisMain.ReadTextHeader(in);
    dt = ext.d1
    parameters = Dict(:dt=>dt,:fa=>fa,:fb=>fb,:fc=>fc,:fd=>fd )

	SeisProcessFile(in,out,[SeisBandPass],[parameters];group=group,key=key,itrace=itrace,ntrace=ntrace)
end

function pad_first_axis(a,N1)
	n1 = size(a,1)
	nx = size(a,2)
	b = zeros(N1,nx)
	for ix = 1 : nx
		b[1:n1,ix] = a[:,ix]
	end
	return b
end


##
using Pkg, Revise, BenchmarkTools
Pkg.activate(".")
##
import Images
import NPZ: npzread
using GLMakie
using ZStackAnalysis
##
using Base.Threads
nthreads()
##
file_content = npzread("./data/Fed_X63_Z2_SIM.npz")
zstack = Preprocessing.convert_to_float32(file_content["data"])
scale = file_content["scale"]
##
size(zstack)
##

##
xy_region = (1201:1500, 201:500)
actin = interpolate_zstack(zstack[1, :, xy_region...], scale)
actin_adj = Images.adjust_histogram(actin, Images.Equalization(nbins = 256, minval = 0.0, maxval = 0.7))
##
intensity, orientation = LineFilterTransform(actin_adj[1:20, 1:20, 1:20], 10, 7, 20)
@btime OrientationFilterTransform(intensity, orientation, 10, 7, 20)
##
@btime line_segment()


module Utils
using ..Types, ..scenario
using LinearAlgebra

export rotation, machNumber, temperature, density, gravity, T_MatrixQ, euler2quartenions, T_matrixE

function machNumber(v::Float64)
    ma = v / 340.0
    return ma
end
function machNumber(v::Float64, y::Float64)
    t = temperature(y)
    a = speedOfSound(t)
    ma = v / a
    return ma
end

function temperature(y::Float64;y0 = 0.0, t0 = 288.15, β = -6.5)
    t = (t0+β * (y-y0)*1e-3)#-273.15 #deg Celsius
    return t
end

function speedOfSound(t::Float64; γ = 1.4, R = 287.05287)
    a = sqrt(γ*R*t)
    return a
end
function pressure(y::Float64;p0 = 101325.0, y0 = 0.0, t0 = 288.15, β = -6.5, R = 287.05287)
    p = p0*(1+β/t0*(y-y0)*1e-3)^(-gravity(y)/(β*R))
    return p
end
function density(y::Float64; R = 287.05287)
    t = temperature(y)
    p = pressure(y)
    ρ = p/(R*t)
    return ρ
end

function gravity(y::Float64; g0 = 9.80665, r = 6356766.0)
    g = g0*(r/(r+y))^2
    return g
end

function Δrange(v0::Float64, θ::Float64, target::Target2D, trajectory::Function,proj::Projectile2D)
p= deepcopy(proj)
#println(proj.position)
#println(p.position)
vx = v0*cosd(θ)
vy = v0*sind(θ)
#p = createProjectile(0.03, 0.04,[vx,vy], [0.0,0.0])
p.velocity_m = [vx,vy]
trajectory(p,target)
#println(p.position)
Δrange = target.position[2]-p.position[2]

return Δrange
end

function bissection(target::Target2D, trajectory::Function, p::Projectile2D; θmin = 0.0, θmax = 45.0, ϵ = 0.001 )

#t2D = scenario.createTarget(0.2,[200.0,50.0])

#global θmin = 0.0
#global θmax = 45.0
#v0 = 900.0
#ϵ = 0.001
#global fmin = Δrange(inputs.v0, θmin,t2D, ExternalBallistics.PmmZeroDrag.timeOfFlight!)
#global fmax = Δrange(inputs.v0, θmax,t2D, ExternalBallistics.PmmZeroDrag.timeOfFlight!)
#global error = min(abs(fmin),abs(fmax))
fmin = Δrange(inputs.v0, θmin,target, trajectory,p)
#println(θmin," ", θmax)
fmax = Δrange(inputs.v0, θmax,target, trajectory,p)
#println(fmin," ", fmax)
error = min(abs(fmin),abs(fmax))
θ = 0.0

if sign(fmin)==sign(fmax)
  println("Bisection method fails.")
else

  while error>ϵ

    #global θmin
    #global θmax
    #global fmin
    #global fmax
    #global error
    #global θ = 1/2*(θmin+θmax)
    θ = 1/2*(θmin+θmax)
    #println(error)

    f = Δrange(inputs.v0, θ,target, trajectory,p)
    error = abs(f)
    if f*fmin < 0
            θmax = θ
        else f*fmax < 0
            θmin = θ
      end
    end

  end
  return θ

end

function rotation(v::Array{Float64,1}, T::Array{Float64,2})
    vrot = T*v
    return vrot
end

function arctangent(A::Float64, B::Float64)
    if B>0
        return atan(A/B)
    elseif B==0 && A>0
        return pi/2
    elseif B==0 && A<0
        return -pi/2
    elseif B<0
        return pi+atan(A/B)
    end
end

function quaternions(λ::String, T)
    if λ == "λ0"
        λ0 = 1/2*(1+tr(T))^(1/2)
        λ1 = 1/λ0*1/4*(T[3,2]-T[2,3])
        λ2 = 1/λ0*1/4*(T[1,3]-T[3,1])
        λ3 = 1/λ0*1/4*(T[2,1]-T[1,2])
    elseif λ =="λ1"
        λ1 = 1/2*(1+2*T[1,1]-tr(T))^(1/2)
        λ0 = 1/λ1*1/4*(T[3,2]-T[2,3])
        λ2 = 1/λ1*1/4*(T[1,2]+T[2,1])
        λ3 = 1/λ1*1/4*(T[1,3]+T[3,1])
    elseif λ == "λ2"
        λ2 = 1/2*(1+2*T[2,2]-tr(T))^(1/2)
        λ0 = 1/λ2*1/4*(T[1,3]-T[3,1])
        λ1 = 1/λ2*1/4*(T[1,2]+T[2,1])
        λ3 = 1/λ2*1/4*(T[2,3]+T[3,2])
    elseif λ == "λ3"
        λ3 = 1/2*(1+2*T[3,3]-tr(T))^(1/2)
        λ0 = 1/λ3*1/4*(T[2,1]-T[1,2])
        λ1 = 1/λ3*1/4*(T[1,3]+T[3,1])
        λ2 = 1/λ3*1/4*(T[2,3]-T[3,2])
    end
    if sign(λ0)==-1
        λ0 = -1* λ0
        λ1 = -1* λ1
        λ2 = -1* λ2
        λ3 = -1* λ3
    end

    return λ0, λ1, λ2, λ3

end

function quaternions2euler(λ0::Float64, λ1::Float64, λ2::Float64, λ3::Float64)
θ = asin(2*(λ0*λ2-λ1*λ3))
ψ_num = 2*(λ1*λ2+λ0*λ3)
ψ_denom = (λ0^2+λ1^2-λ2^2-λ3^2)
ψ = arctangent(ψ_num,ψ_denom)
Φ_num = 2*(λ2*λ3+λ0*λ1)
Φ_denom = (λ0^2-λ1^2-λ2^2+λ3^2)
Φ = arctangent(Φ_num,Φ_denom)
return θ, ψ, Φ
end
function T_matrixE(θ::Float64, ψ::Float64, Φ::Float64)
    T = Array{Float64}(undef, 3,3)
    T[1,1] = cos(θ)*cos(ψ)
    T[1,2] = -cos(Φ)*sin(ψ) + sin(Φ)*sin(θ)*cos(ψ)
    T[1,3] = sin(Φ)*sin(ψ) + cos(Φ)*sin(θ)*cos(ψ)
    T[2,1] = cos(θ)*sin(ψ)
    T[2,2] = cos(Φ)*cos(ψ) + sin(Φ)*sin(θ)*sin(ψ)
    T[2,3] = -sin(Φ)*cos(ψ) + cos(Φ)*sin(θ)*sin(ψ)
    T[3,1] = -sin(θ)
    T[3,2] = sin(Φ)*cos(θ)
    T[3,3] = cos(Φ)*cos(θ)
    return T
end

function dominant(T)
    T00 = tr(T)
    trace = Dict("λ0"=>T00,"λ1"=>T[1,1],"λ2"=>T[2,2],"λ3"=>T[3,3])
    max_trace = maximum(values(trace))
    quater = "none"
    for (k,v) in trace
        if v == max_trace
            quater = k
        end
    end
    return quater
end

function euler2quartenions(θ::Float64, ψ::Float64, Φ::Float64)

    T = T_matrixE(θ, ψ, Φ)
    dominant_quater = dominant(T)

    λ0, λ1, λ2, λ3 = quaternions(dominant_quater, T)

    return λ0, λ1, λ2, λ3

end

function T_MatrixQ(λ0::Float64, λ1::Float64, λ2::Float64, λ3::Float64)
    T = Array{Float64}(undef, 3,3)
    T[1,1] = (λ0^2+λ1^2-λ2^2-λ3^2)
    T[1,2] = 2*(λ1*λ2-λ0*λ3)
    T[1,3] = 2*(λ1*λ3+λ0*λ2)
    T[2,1] = 2*(λ1*λ2+λ0*λ3)
    T[2,2] = (λ0^2-λ1^2+λ2^2-λ3^2)
    T[2,3] = 2*(λ2*λ3-λ0*λ1)
    T[3,1] = 2*(λ1*λ3-λ0*λ2)
    T[3,2] = 2*(λ2*λ3+λ0*λ1)
    T[3,3] = (λ0^2-λ1^2-λ2^2+λ3^2)

    return T

end

end


@testset "Regular matrix algebra" begin
  G = small_group(8, 4)
  A = AlgGrp(FlintQQ, G)
  B, BtoA = Hecke.regular_matrix_algebra(A)

  @test dim(B) == dim(A)

  for i = 1:dim(A)
    for j = 1:dim(A)
      @test BtoA(B[i])*BtoA(B[j]) == BtoA(B[i]*B[j])
      @test (BtoA\A[i])*(BtoA\A[j]) == BtoA\(A[i]*A[j])
    end
  end
end


using Revise

using Test

using Pkg; Pkg.activate(".")
using HypergraphModularity
using StatsBase
using Combinatorics

# just testing that these run correctly at the moment
@testset "Ω" begin

    kmax = 10

    # general partitionize-based intensity function
    Ω = partitionIntensityFunction(p->sum(p), kmax)
    @test typeof(Ω) == IntensityFunction

    # all-or-nothing cut
    Ω = allOrNothingIntensityFunction(x->(x[1]+1)^(-x[2]), kmax)
    @test typeof(Ω) == IntensityFunction

    # number of unique elements in z: generalizes all-or-nothing
    Ω = sumOfExteriorDegreesIntensityFunction((x, α)->(x[1]+1)^(-α[x[2]]*x[2]), kmax)
    @test typeof(Ω) == IntensityFunction
end

n = 20
Z = rand(1:5, n)
ϑ = dropdims(ones(1,n) + rand(1,n), dims = 1)
μ = mean(ϑ)
kmax = 4

ω(x, α) = (x[1]+1)^(-α[x[2]]*x[2])

α = repeat([2.0], kmax)

Ω = sumOfExteriorDegreesIntensityFunction(ω, kmax)
H = sampleSBM(Z, ϑ, Ω;α=α, kmax=kmax, kmin = 1)

@testset "HSBM Sampling" begin
    @test length(keys(H.E)) == kmax
end

# test for dyadic log likelihood
# currently broken, need to fix.

##
@testset "dyadic likelihood" begin
    ωᵢ = 1.0
    ωₒ = 0.1
    L₁ = dyadicLogLikelihood(H, Z, ωᵢ, ωₒ; weighted=false, binary=false, constants = false)
    L₂ = HypergraphModularity.dyadicLogLikelihoodDirect(H, Z, ωᵢ, ωₒ)
    @test L₁ ≈ L₂
end
##

##
##

@testset "cut" begin
    cut1 = first_term_eval(H, Z, Ω;α=α)
    Hyp, w = HypergraphModularity.hyperedge_formatting(H)
    cut2 = HypergraphModularity.first_term_v2(Hyp,w,Z,Ω;α=α)
    cut3 = HypergraphModularity.first_term_v3(H, Z, Ω;α=α)
    @test cut1 ≈ cut2
    @test cut2 ≈ cut3
    @test cut1 ≈ cut3

# CutDiff appears to be off, unclear whether related to refactor
    @testset "cutdiff" begin
        Hyp, w = hyperedge_formatting(H)
        node2edges = EdgeMap(H)
        I = rand(1:n)
        J = Z[I] + 1
        a = NaiveCutDiff(H,Z,I,J,Ω;α=α)
        b = HypergraphModularity.NaiveCutDiff2(Hyp, w, node2edges,Z,I,J,Ω;α=α)

        c = CutDiff(Hyp,w,node2edges,Z,I,J,Ω; α=α)

        @test a ≈ b
        # @test a ≈ c
        # @test b ≈ c

        # now we'll test multi-node cutdiff
        # this tests the implementation in diffs.jl.

        to_move = findall(==(3), Z)
        t = 1

        C = HypergraphModularity.evalCuts(H, Z,  Ω)

        ΔC = HypergraphModularity.cutDiff(C, to_move, t, Z, Hyp, w, node2edges, Ω)

        Z_ = copy(Z)
        Z_[to_move] .= t

        C_new = HypergraphModularity.evalCuts(H, Z_,  Ω)

        @test C_new == C + ΔC
    end
end

D = H.D

@testset "vol" begin

    @testset "sums" begin

        all_partitions = collect(Iterators.flatten([(partitions(i,j)) for i =1:kmax for j=1:i]))

        p = [3, 1, 1]

        s1 = HypergraphModularity.evalSumNaive(p, Z, D)
        s2 = HypergraphModularity.evalSumNaive2(p, Z, D)

        @test s1 == s2

        sumsNaive = HypergraphModularity.evalSumsNaive(Z,D,kmax)
        sumsPV    = HypergraphModularity.evalSumsPV(Z,D,kmax)
        sums      = HypergraphModularity.evalSums(Z,D,kmax)[3]

        @test all([sumsNaive[p] == sumsPV[p] for p in all_partitions])
        @test all([sumsNaive[p] == sums[p] for p in all_partitions])
    end

    @testset "voldiff" begin

        ℓ = maximum(Z) # number of total groups

        # let's move all the nodes in group 4 to group 5
        to_move = findall(==(4), Z)
        t = 5
        Z_ = copy(Z)

        # first we'll step through, moving each node one-by-one
        V, μ, M = evalSums(Z, D, kmax; constants=false, bigInt=false);
        for i in to_move
            ΔV, Δμ, ΔM = increments(V, μ, M, i, t, D, Z_)
            V, μ, M = HypergraphModularity.addIncrements(V, μ, M, ΔV, Δμ, ΔM)
            # carry out the change in membership
            Z_[i] = t
        end
        step_wise = M

        # next we'll compute the complete set of increments at once
        V, μ, M = evalSums(Z, D, kmax; constants=false, bigInt=false);
        ΔV, Δμ, ΔM = increments(V, μ, M, to_move, t, D, Z)
        V, μ, M = HypergraphModularity.addIncrements(V, μ, M, ΔV, Δμ, ΔM)
        batch = M
        #
        @test all([step_wise[p] == batch[p] for p in keys(M)])
    end

    @testset "second term" begin

        # aggregation
        ℓ = 0
        bigInt = true
        V, μ, M = evalSums(Z, H, ℓ, bigInt)
        M̂ = HypergraphModularity.aggregateSums(M, Ω)
        @test sum(values(M)) ≈ sum(values(M̂))

        s1 = sum(Ω.ω(Ω.aggregator(p), α)*M[p] for p in keys(M))
        s2 = sum(Ω.ω(p̂, α)*M̂[p̂] for p̂ in keys(M̂))

        @test s1 ≈ s2

        s1 = HypergraphModularity.second_term_eval(H, Z, Ω; α=α)
        s2 = HypergraphModularity.naiveSecondTerm(H, Z, Ω; α=α)

        @test s1 ≈ s2
    end
end

@testset "modularity" begin

    Q_naive = HypergraphModularity.modularityNaive(H, Z, Ω;α=α)
    Q = modularity(H, Z, Ω;α=α)
    @test Q_naive ≈ Q
end

@testset "log likelihood" begin
    Ω = estimateΩEmpirically(H, Z; aggregator = p -> [length(p) == 1, sum(p)], bigNums = true)
    L = HypergraphModularity.logLikelihoodNaive(H, Z, Ω; α=α)
    L̄ = logLikelihood(H, Z, Ω;α=α)
    @test Float64(L) ≈ Float64(sum(L̄))
end

@testset "hyperlouvain" begin
    """
    more a test that the code successfully runs than a meaningful test of correctness
    """
    Z_ = HyperLouvain(H,kmax,Ω,100,true;α=α, verbose=true)
    Z_ = SuperNodeLouvain(H,kmax,Ω,100,true;α=α, verbose=true)
end


using FileIO, Gnuplot, ColorSchemes
using GeometryBasics, LinearAlgebra, Statistics
# based on https://matplotlib.org/matplotblog/posts/custom-3d-engine/
# download this file
# https://matplotlib.org/matplotblog/posts/custom-3d-engine/bunny.obj
bunny = load("bunny.obj")
ptstemp = bunny.position
npts = length(ptstemp)
points = zeros(npts, 3)
for i in 1:npts
    points[i,:] = ptstemp[i]
end
# normalize into a box
mnpts = minimum(points, dims=1)
mxpts = maximum(points, dims=1)
# all the data needed.
verpts = (points .- (mxpts .+ mnpts)./2)./maximum(mxpts .- mnpts)
triface = decompose(TriangleFace{Int32}, bunny)
# just lines and no order
bunnyB = Vector{String}()
for i in 1:length(triface)
    polytri = verpts[triface[i],:]
    x, y = polytri[:,1], polytri[:,2] # polytri[:,3]
    push!(bunnyB, join(string.(x) .* " ".* string.(y), "\n") * "\n")
end
# perpective, views, rotations
function frustum(left, right, bottom, top, znear, zfar)
    M = zeros(Float32, (4, 4))
    M[1, 1] = 2.0 * znear / (right - left)
    M[2, 2] = 2.0 * znear / (top - bottom)
    M[3, 3] = -(zfar + znear) / (zfar - znear)
    M[1, 3] = (right + left) / (right - left)
    M[3, 2] = (top + bottom) / (top - bottom)
    M[3, 4] = -2.0 * znear * zfar / (zfar - znear)
    M[4, 3] = -1.0
    M
end
function perspective(fovy, aspect, znear, zfar)
    h = tan(fovy*π/360) * znear
    w = h * aspect
    return frustum(-w, w, -h, h, znear, zfar)
end
function translate(x, y, z)
    [1 0 0 x; 0 1 0 y; 0 0 1 z; 0 0 0 1]
end
function xrotate(θ)
    t = π * θ / 180
    c, s = cos(t), sin(t)
    [1 0 0 0; 0 c -s 0; 0 s c 0; 0 0 0 1]
end
function yrotate(θ)
    t = π * θ / 180
    c, s = cos(t), sin(t)
    [c 0 s 0; 0 1 0 0; -s 0 c 0; 0 0 0 1]
end
# Model, view, perspective (MVP)
model = xrotate(20) * yrotate(45)
view  = translate(0,0,-3.5)
proj  = perspective(25, 1, 1, 100)
MVP   = proj * view * model
verpts = hcat(verpts, ones(npts)) * MVP'
verpts ./= verpts[:,4]
# zdepth_order
zdepth = zeros(length(triface))
ptsxy = []
for i in 1:length(triface)
    polytri = verpts[triface[i],:]
    x, y = polytri[:,1], polytri[:,2]
    push!(ptsxy, [x,y])
    zdepth[i] = -mean(polytri[:,3])
end
zmin, zmax = minimum(zdepth), maximum(zdepth)
zdepth = (zdepth .- zmin)/(zmax-zmin)
zdepth_order = sortperm(zdepth)
zdepth = zdepth[zdepth_order] # used for the colormap
# polygons to gnuplot
bunnyC = Vector{String}()
for (indx,i) in enumerate(zdepth_order) # order in which polygons should be drawn
    x, y = ptsxy[i][1], ptsxy[i][2]
    push!(bunnyC, join(string.(x) .* " ".* string.(y) .* " " .*
                  string(zdepth[indx]), "\n") * "\n")
end
# let's do an egg
x = -1:0.05:1
y = -1.5:0.05:2
egg(x,y) = x^2 + y^2/(1.4 + y/5)^2
segg = [egg(x,y) for x in x, y in y]
# the bunny !
@gp "set multi layout 1,3 title 'Easter'" :-
@gp :- 2 "set style fill solid 1 noborder" title ="Bunny in color"
@gp :- "set size ratio -1" :-
@gp :- bunnyC "w filledcurves notit lc pal" palette(:spring)
@gp :- "unset colorbox; unset xtics; unset ytics; unset border"
@gp :- 3 "set style fill solid 1 noborder" title ="Bunny lines"
@gp :- "set size ratio -1" :-
@gp :- bunnyB "w l notit lc 'black' lw 0.1" :-
@gp :- "unset xtics; unset ytics; unset border"
@gp :- 1 title = "Egg" :-
@gsp :- x y segg "w l lc palette" palette(:cool) "set view map" :-
@gsp :- "set contour base;set key off" "set auto fix" :-
@gsp :- "set cntrparam levels incremental  0,0.01,1" "unset surface" :-
@gsp :- "set size ratio -1" cbrange =(0,1)
@gp :- "unset colorbox; unset xtics; unset ytics; unset border"
@gsp :- "unset xtics; unset ytics; unset border"
Gnuplot.save(term = "pngcairo size 1600,800", output = "bunnyC.png")


using Statistics
using LinearAlgebra
using Distributions
using Random

function log_bayesian_posterior(s_param, θ::Array{Float64,1}, forward::Function,
    y::Array{Float64,1},  Σ_η::Array{Float64,2},
    μ0::Array{Float64,1}, Σ0::Array{Float64,2})

    Gu = forward(s_param, θ)
    Φ = - 0.5*(y - Gu)'/Σ_η*(y - Gu) - 0.5*(θ - μ0)'/Σ0*(θ - μ0)
    return Φ

end

function d_log_bayesian_posterior(s_param, θ::Array{Float64,1}, forward::Function,
    y::Array{Float64,1},  Σ_η::Array{Float64,2},
    μ0::Array{Float64,1}, Σ0::Array{Float64,2})

    Gu, dGu = forward(s_param, θ)
    Φ = - 0.5*(y - Gu)'/Σ_η*(y - Gu) - 0.5*(θ - μ0)'/Σ0*(θ - μ0)
    dΦ = dGu'* (Σ_η\(y - Gu)) - Σ0\(θ - μ0)
    return Φ, dΦ

end

function log_likelihood(s_param, θ::Array{Float64,1}, forward::Function,
    y::Array{Float64,1},  Σ_η::Array{Float64,2})

    Gu = forward(s_param, θ)
    Φ = - 0.5*(y - Gu)'/Σ_η*(y - Gu)
    return Φ

end

"""
When the density function is Φ/Z,
The f_density function return log(Φ) instead of Φ
"""
function RWMCMC_Run(log_bayesian_posterior::Function, θ0::Array{FT,1}, τ::FT, n_ite::IT; seed::IT=11) where {FT<:AbstractFloat, IT<:Int}

    N_θ = length(θ0)
    θs = zeros(Float64, n_ite, N_θ)
    fs = zeros(Float64, n_ite)

    θs[1, :] .= θ0
    fs[1] = log_bayesian_posterior(θ0)

    Random.seed!(seed)
    for i = 2:n_ite
        θ_p = θs[i-1, :]
        θ = θ_p + τ * rand(Normal(0, 1), N_θ)

        fs[i] = log_bayesian_posterior(θ)
        α = min(1.0, exp(fs[i] - fs[i-1]))
        if α > rand(Uniform(0, 1))
            # accept
            θs[i, :] = θ
            fs[i] = fs[i]
        else
            # reject
            θs[i, :] = θ_p
            fs[i] = fs[i-1]
        end

    end

    return θs
end

"""
When the density function is Φ/Z,
The f_density function return log(Φ) instead of Φ
"""
function PCN_Run(log_likelihood::Function, θ0::Array{FT,1}, θθ0_cov::Array{FT,2}, β::FT, n_ite::IT; seed::IT=11) where {FT<:AbstractFloat, IT<:Int}

    N_θ = length(θ0)
    θs = zeros(Float64, n_ite, N_θ)
    fs = zeros(Float64, n_ite)

    θs[1, :] .= θ0
    fs[1] = log_likelihood(θ0)

    Random.seed!(seed)
    for i = 2:n_ite
        θ_p = θs[i-1, :]
        θ = sqrt(1 - β^2)*θ_p + β * rand(MvNormal(zeros(size(θ0)), θθ0_cov))

        fs[i] = log_likelihood(θ)
        α = min(1.0, exp(fs[i] - fs[i-1]))
        if α > rand(Uniform(0, 1))
            # accept
            θs[i, :] = θ
            fs[i] = fs[i]
            @info "accept i = ", i
        else
            # reject
            θs[i, :] = θ_p
            fs[i] = fs[i-1]
        end

    end

    return θs
end

# Stein variational gradient descent
function SVGD_Kernel(θ; h = -1)
    J, nd = size(θ)

    XY = θ*θ';
    x2= sum(θ.^2, dims=2);
    X2e = repeat(x2, 1, J);
    pairwise_dists = X2e + X2e' - 2*XY

    if h < 0
        h = median(pairwise_dists)
        h = sqrt(0.5 * h / log(J+1))
    end
    # compute the rbf kernel
    Kxy = exp.( -pairwise_dists / h^2 / 2)

    dxkxy = -Kxy * θ
    sumkxy = sum(Kxy, dims=2)
    for i = 1:nd
        dxkxy[:, i] = dxkxy[:,i] + θ[:,i].*sumkxy
    end

    dxkxy = dxkxy / (h^2)
    return Kxy, dxkxy
end

# lnprob(θ) = ∇ log p(θ)
# θs0 is a J by n_dim matrix
function SVGD_Run(θs0, lnprob, n_ite::IT = 1000, stepsize::FT = 1e-3, alpha = 0.9) where {FT<:AbstractFloat, IT<:Int}
        # Check input

    θs = copy(θs0)
    lnpgrad = similar(θs0)
    J, nd = size(θs0)
    # adagrad with momentum
    fudge_factor = 1e-6
    historical_grad = 0

    for iter = 1:n_ite
        for i = 1:J
            _, lnpgrad[i, :] = lnprob(θs[i, :])
        end
        # calculating the kernel matrix
        kxy, dxkxy = SVGD_Kernel(θs; h = -1)
        grad_theta = (kxy * lnpgrad + dxkxy) / J

        # adagrad
        if iter == 1
            historical_grad = historical_grad .+ grad_theta .^ 2
        else
            historical_grad .= alpha * historical_grad .+ (1 - alpha) * (grad_theta .^ 2)
        end

        adj_grad = grad_theta ./ (fudge_factor .+ sqrt.(historical_grad))
        θs = θs + stepsize * adj_grad
    end

    return θs
end

"""
Metropolis Adjusted Langevin Algorithm
When the density function is Φ/Z,
The f_density function return log(Φ) instead of Φ
"""
function MALA_Run(d_log_bayesian_posterior::Function, θ0::Array{FT,1}, τ::FT, n_ite::IT; seed::IT=11) where {FT<:AbstractFloat, IT<:Int}

    N_θ = length(θ0)
    θs = zeros(Float64, n_ite, N_θ)
    fs = zeros(Float64, n_ite)
    dfs = zeros(Float64, n_ite, N_θ)

    θs[1, :] .= θ0
    fs[1], dfs[1, :] = d_log_bayesian_posterior(θ0)

    Random.seed!(seed)
    for i = 2:n_ite
        θ_p = θs[i-1, :]
        θ = θ_p + τ *dfs[i-1, :] + sqrt(2τ)*rand(Normal(0,1), N_θ)

        fs[i], dfs[i, :] = d_log_bayesian_posterior(θ)
        α = min(1.0, exp( (fs[i]   - norm(θ_p - θ - τ *dfs[i, :])^2/(4τ)  ) -
                          (fs[i-1] - norm(θ - θ_p - τ *dfs[i-1, :])^2/(4τ))
                        )
                )
        # @info α , θ,  fs[i], dfs[i,:], θ_p, fs[i-1], dfs[i-1,:],
        if α > rand(Uniform(0, 1))
            # accept
            θs[i, :] = θ
            fs[i] = fs[i]
            dfs[i, :] = dfs[i, :]
        else
            # reject
            θs[i, :] = θ_p
            fs[i] = fs[i-1]
            dfs[i, :] = dfs[i-1, :]
        end

    end

    return θs
end

# θ_s: N_ens/2 by N_θ matrix, the set
# θ_c: N_ens/2 by N_θ matrix, the complementary set
function emcee_Propose(θ_s::Array{FT,2}, θ_c::Array{FT,2}; a::FT = 2.0) where {FT<:AbstractFloat, IT<:Int}
    Ns, N_θ = size(θ_s)

    zz = ((a - 1.0) * rand(Uniform(0, 1), Ns) .+ 1).^2.0 / a
    factors = (N_θ - 1.0) * log.(zz)

    rint = rand(1:Ns, Ns,)
    return θ_c[rint, :] - (θ_c[rint, :] - θ_s) .* zz, factors

end

"""
When the density function is Φ/Z,
The log_bayesian_posterior function return log(Φ) instead of Φ
θ0 : initial ensemble of size N_ens by N_θ
"""
function emcee_Run(log_bayesian_posterior::Function, θ0::Array{FT,2}, n_ite::IT; random_split::Bool = true, a::FT = 2.0, seed::IT=11) where {FT<:AbstractFloat, IT<:Int}
    Random.seed!(seed)

    N_ens, N_θ = size(θ0)
    @assert(N_ens >= 2N_θ)
    @assert(N_ens % 2 == 0)

    θs = zeros(Float64, n_ite, N_ens, N_θ)
    fs = zeros(Float64, n_ite, N_ens)

    θs[1, :, :] .= θ0
    for k = 1:N_ens
        fs[1, k] = log_bayesian_posterior(θ0[k, :])
    end

    nsplit = 2
    N_s = div(N_ens, 2)

    all_inds = Array(1:N_ens)
    inds = all_inds .% nsplit # 2 group
    log_probs = zeros(Float64, N_s)
    for i_t = 2:n_ite

        if random_split
            shuffle!(inds)
        end
        for split = 0:1
            # boolean array of the current set
            s_inds = (inds .== split)
            # boolean array for the complementary set
            c_inds = (inds .!= split)

            s, c = θs[i_t - 1, s_inds, :], θs[i_t - 1, c_inds, :]

            q, factors = emcee_Propose(s, c; a = a)
            for i = 1:N_s
                log_probs[i] = log_bayesian_posterior(q[i, :])
            end

            # Loop over the walkers and update them accordingly.
            for i = 1:N_s
                j = all_inds[s_inds][i]

                # j is the index
                # @info i, factors[i] , log_probs[i] , fs[i_t - 1, j]

                α = min(1.0, exp(factors[i] + log_probs[i] - fs[i_t - 1, j]))
                if α > rand(Uniform(0, 1))
                    # accept
                    θs[i_t, j, :] = q[i, :]
                    fs[i_t, j] = log_probs[i]
                else
                    # reject
                    θs[i_t, j, :] = θs[i_t - 1, j, :]
                    fs[i_t, j] = fs[i_t - 1, j]
                end
            end
        end
    end

    return θs
end


module SLATEC

export DRC, DRD, DRF, DRJ

const D1MACH1 = floatmin(Float64)
const D1MACH2 = floatmax(Float64)
const D1MACH3 = eps(Float64)/2
const D1MACH4 = eps(Float64)
const D1MACH5 = log10(2.)

#***BEGIN PROLOGUE  DRF
#***PURPOSE  Compute the incomplete or complete elliptic integral of the
#            1st kind.  For X, Y, and Z non-negative and at most one of
#            them zero, RF(X,Y,Z) = Integral from zero to infinity of
#                                -1/2     -1/2     -1/2
#                      (1/2)(t+X)    (t+Y)    (t+Z)    dt.
#            If X, Y or Z is zero, the integral is complete.
#***LIBRARY   SLATEC
#***CATEGORY  C14
#***TYPE      DOUBLE PRECISION (RF-S, DRF-D)
#***KEYWORDS  COMPLETE ELLIPTIC INTEGRAL, DUPLICATION THEOREM,
#             INCOMPLETE ELLIPTIC INTEGRAL, INTEGRAL OF THE FIRST KIND,
#             TAYLOR SERIES
#***AUTHOR  Carlson, B. C.
#             Ames Laboratory-DOE
#             Iowa State University
#             Ames, IA  50011
#           Notis, E. M.
#             Ames Laboratory-DOE
#             Iowa State University
#             Ames, IA  50011
#           Pexton, R. L.
#             Lawrence Livermore National Laboratory
#             Livermore, CA  94550

function DRF(X::Float64, Y::Float64, Z::Float64)

    ERRTOL = (4.0*D1MACH3)^(1.0/6.0)
    LOLIM  = 5.0 * D1MACH1
    UPLIM  = D1MACH2/5.0
    C1 = 1.0/24.0
    C2 = 3.0/44.0
    C3 = 1.0/14.0

    ans = 0.0
    if min(X,Y,Z) < 0.0
        return ans, 1
    end

    if max(X,Y,Z) > UPLIM
        return ans, 3
    end

    if min(X+Y,X+Z,Y+Z) < LOLIM
        return ans, 2
    end

    XN = X
    YN = Y
    ZN = Z
    MU = 0.
    XNDEV = 0.
    YNDEV = 0.
    ZNDEV = 0.

    while true
        MU = (XN+YN+ZN)/3.0
        XNDEV = 2.0 - (MU+XN)/MU
        YNDEV = 2.0 - (MU+YN)/MU
        ZNDEV = 2.0 - (MU+ZN)/MU
        EPSLON = max(abs(XNDEV),abs(YNDEV),abs(ZNDEV))
        if (EPSLON < ERRTOL) break end
        XNROOT = sqrt(XN)
        YNROOT = sqrt(YN)
        ZNROOT = sqrt(ZN)
        LAMDA = XNROOT*(YNROOT+ZNROOT) + YNROOT*ZNROOT
        XN = (XN+LAMDA)*0.250
        YN = (YN+LAMDA)*0.250
        ZN = (ZN+LAMDA)*0.250
    end

    E2 = XNDEV*YNDEV - ZNDEV*ZNDEV
    E3 = XNDEV*YNDEV*ZNDEV
    S  = 1.0 + (C1*E2-0.10-C2*E3)*E2 + C3*E3
    ans = S/sqrt(MU)

    return ans, 0
end

#***BEGIN PROLOGUE  DRD
#***PURPOSE  Compute the incomplete or complete elliptic integral of
#            the 2nd kind. For X and Y nonnegative, X+Y and Z positive,
#            DRD(X,Y,Z) = Integral from zero to infinity of
#                                -1/2     -1/2     -3/2
#                      (3/2)(t+X)    (t+Y)    (t+Z)    dt.
#            If X or Y is zero, the integral is complete.
#***LIBRARY   SLATEC
#***CATEGORY  C14
#***TYPE      DOUBLE PRECISION (RD-S, DRD-D)
#***KEYWORDS  COMPLETE ELLIPTIC INTEGRAL, DUPLICATION THEOREM,
#             INCOMPLETE ELLIPTIC INTEGRAL, INTEGRAL OF THE SECOND KIND,
#             TAYLOR SERIES
#***AUTHOR  Carlson, B. C.
#             Ames Laboratory-DOE
#             Iowa State University
#             Ames, IA  50011
#           Notis, E. M.
#             Ames Laboratory-DOE
#             Iowa State University
#             Ames, IA  50011
#           Pexton, R. L.
#             Lawrence Livermore National Laboratory
#             Livermore, CA  94550

function DRD(X::Float64, Y::Float64, Z::Float64)

    ERRTOL = (D1MACH3/3.0)^(1.0/6.0)
    LOLIM  = 2.0/(D1MACH2)^(2.0/3.0)
    TUPLIM = D1MACH1^(1.0E0/3.0E0)
    TUPLIM = (0.10*ERRTOL)^(1.0E0/3.0E0)/TUPLIM
    UPLIM  = TUPLIM^2.0
    C1 = 3.0/14.0
    C2 = 1.0/6.0
    C3 = 9.0/22.0
    C4 = 3.0/26.0

    ans = 0.0
    if min(X,Y) < 0.0
        return ans, 1
    end

    if max(X,Y,Z) > UPLIM
        return ans, 3
    end

    if min(X+Y,Z) < LOLIM
        return ans, 2
    end

    XN = X
    YN = Y
    ZN = Z
    SIGMA = 0.0
    POWER4 = 1.0
    MU = 0.
    XNDEV = 0.
    YNDEV = 0.
    ZNDEV = 0.

    while true
        MU = (XN+YN+3.0*ZN)*0.20
        XNDEV = (MU-XN)/MU
        YNDEV = (MU-YN)/MU
        ZNDEV = (MU-ZN)/MU
        EPSLON = max(abs(XNDEV), abs(YNDEV), abs(ZNDEV))
        if (EPSLON < ERRTOL) break end
        XNROOT = sqrt(XN)
        YNROOT = sqrt(YN)
        ZNROOT = sqrt(ZN)
        LAMDA = XNROOT*(YNROOT+ZNROOT) + YNROOT*ZNROOT
        SIGMA = SIGMA + POWER4/(ZNROOT*(ZN+LAMDA))
        POWER4 = POWER4*0.250
        XN = (XN+LAMDA)*0.250
        YN = (YN+LAMDA)*0.250
        ZN = (ZN+LAMDA)*0.250
    end

    EA = XNDEV*YNDEV
    EB = ZNDEV*ZNDEV
    EC = EA - EB
    ED = EA - 6.0*EB
    EF = ED + EC + EC
    S1 = ED*(-C1+0.250*C3*ED-1.50*C4*ZNDEV*EF)
    S2 = ZNDEV*(C2*EF+ZNDEV*(-C3*EC+ZNDEV*C4*EA))
    ans = 3.0*SIGMA + POWER4*(1.0+S1+S2)/(MU*sqrt(MU))

    return ans, 0
end

#***BEGIN PROLOGUE  DRC
#***PURPOSE  Calculate a double precision approximation to
#             DRC(X,Y) = Integral from zero to infinity of
#                              -1/2     -1
#                    (1/2)(t+X)    (t+Y)  dt,
#            where X is nonnegative and Y is positive.
#***LIBRARY   SLATEC
#***CATEGORY  C14
#***TYPE      DOUBLE PRECISION (RC-S, DRC-D)
#***KEYWORDS  DUPLICATION THEOREM, ELEMENTARY FUNCTIONS,
#             ELLIPTIC INTEGRAL, TAYLOR SERIES
#***AUTHOR  Carlson, B. C.
#             Ames Laboratory-DOE
#             Iowa State University
#             Ames, IA  50011
#           Notis, E. M.
#             Ames Laboratory-DOE
#             Iowa State University
#             Ames, IA  50011
#           Pexton, R. L.
#             Lawrence Livermore National Laboratory
#             Livermore, CA  94550

function DRC(X::Float64, Y::Float64)

    ERRTOL = (D1MACH3/16.0)^(1.0/6.0)
    LOLIM  = 5.0 * D1MACH1
    UPLIM  = D1MACH2 / 5.0
    C1 = 1.0/7.0
    C2 = 9.0/22.0

    ans = 0.0
    if X < 0.0 || Y <= 0.0
        return ans, 1
    end

    if max(X,Y) > UPLIM
        return ans, 3
    end

    if X+Y < LOLIM
        return ans, 2
    end

    XN = X
    YN = Y
    MU = 0.
    SN = 0.

    while true
        MU = (XN+YN+YN)/3.0
        SN = (YN+MU)/MU - 2.0
        if abs(SN) < ERRTOL break end
        LAMDA = 2.0*sqrt(XN)*sqrt(YN) + YN
        XN = (XN+LAMDA)*0.250
        YN = (YN+LAMDA)*0.250
    end

    S = SN*SN*(0.30+SN*(C1+SN*(0.3750+SN*C2)))
    ans = (1.0+S)/sqrt(MU)

    return ans,0
end

#***BEGIN PROLOGUE  DRJ
#***PURPOSE  Compute the incomplete or complete (X or Y or Z is zero)
#            elliptic integral of the 3rd kind.  For X, Y, and Z non-
#            negative, at most one of them zero, and P positive,
#             RJ(X,Y,Z,P) = Integral from zero to infinity of
#                              -1/2     -1/2     -1/2     -1
#                    (3/2)(t+X)    (t+Y)    (t+Z)    (t+P)  dt.
#***LIBRARY   SLATEC
#***CATEGORY  C14
#***TYPE      DOUBLE PRECISION (RJ-S, DRJ-D)
#***KEYWORDS  COMPLETE ELLIPTIC INTEGRAL, DUPLICATION THEOREM,
#             INCOMPLETE ELLIPTIC INTEGRAL, INTEGRAL OF THE THIRD KIND,
#             TAYLOR SERIES
#***AUTHOR  Carlson, B. C.
#             Ames Laboratory-DOE
#             Iowa State University
#             Ames, IA  50011
#           Notis, E. M.
#             Ames Laboratory-DOE
#             Iowa State University
#             Ames, IA  50011
#           Pexton, R. L.
#             Lawrence Livermore National Laboratory
#             Livermore, CA  94550

function DRJ(X::Float64, Y::Float64, Z::Float64, P::Float64)

    ERRTOL = (D1MACH3/3.0)^(1.0/6.0)
    LOLIM  = (5.0 * D1MACH1)^(1.0/3.0)
    UPLIM  = 0.30*( D1MACH2 / 5.0)^(1.0/3.0)

    C1 = 3.0/14.0
    C2 = 1.0/3.0
    C3 = 3.0/22.0
    C4 = 3.0/26.0

    ans = 0.0
    if min(X,Y,Z) < 0.0
        return ans, 1
    end

    if max(X,Y,Z,P) > UPLIM
        return ans, 3
    end

    if min(X+Y,X+Z,Y+Z,P) < LOLIM
        return ans, 2
    end

    IER = 0
    XN = X
    YN = Y
    ZN = Z
    PN = P
    SIGMA = 0.0
    POWER4 = 1.0
    MU = 0.
    XNDEV = 0.
    YNDEV = 0.
    ZNDEV = 0.
    PNDEV = 0.

    while true
        MU = (XN+YN+ZN+PN+PN)*0.20
        XNDEV = (MU-XN)/MU
        YNDEV = (MU-YN)/MU
        ZNDEV = (MU-ZN)/MU
        PNDEV = (MU-PN)/MU
        EPSLON = max(abs(XNDEV), abs(YNDEV), abs(ZNDEV), abs(PNDEV))
        if EPSLON < ERRTOL break end
        XNROOT =  sqrt(XN)
        YNROOT =  sqrt(YN)
        ZNROOT =  sqrt(ZN)
        LAMDA = XNROOT*(YNROOT+ZNROOT) + YNROOT*ZNROOT
        ALFA = PN*(XNROOT+YNROOT+ZNROOT) + XNROOT*YNROOT*ZNROOT
        ALFA = ALFA*ALFA
        BETA = PN*(PN+LAMDA)*(PN+LAMDA)
        drc,IER = DRC(ALFA,BETA)
        SIGMA = SIGMA + POWER4*drc
        POWER4 = POWER4*0.250
        XN = (XN+LAMDA)*0.250
        YN = (YN+LAMDA)*0.250
        ZN = (ZN+LAMDA)*0.250
        PN = (PN+LAMDA)*0.250
    end

    EA = XNDEV*(YNDEV+ZNDEV) + YNDEV*ZNDEV
    EB = XNDEV*YNDEV*ZNDEV
    EC = PNDEV*PNDEV
    E2 = EA - 3.0*EC
    E3 = EB + 2.0*PNDEV*(EA-EC)
    S1 = 1.0 + E2*(-C1+0.750*C3*E2-1.50*C4*E3)
    S2 = EB*(0.50*C2+PNDEV*(-C3-C3+PNDEV*C4))
    S3 = PNDEV*EA*(C2-PNDEV*C3) - C2*PNDEV*EC
    ans = 3.0*SIGMA + POWER4*(S1+S2+S3)/(MU* sqrt(MU))

    return ans, IER
end

end # module


using LinearAlgebra

mutable struct Bird <: AbstractAgent
    id::Int
    pos::NTuple{2,Float64}
    vel::NTuple{2,Float64}
    speed::Float64
    cohere_factor::Float64
    separation::Float64
    separate_factor::Float64
    match_factor::Float64
    visual_distance::Float64
end

"""
```julia
flocking(;
    n_birds = 100,
    speed = 1.0,
    cohere_factor = 0.25,
    separation = 4.0,
    separate_factor = 0.25,
    match_factor = 0.01,
    visual_distance = 5.0,
    dims = (100, 100),
)
```
Same as in [Flock model](@ref).
"""
function flocking(;
        n_birds = 100,
        speed = 1.0,
        cohere_factor = 0.25,
        separation = 4.0,
        separate_factor = 0.25,
        match_factor = 0.01,
        visual_distance = 5.0,
        dims = (100, 100),
    )

    space2d = ContinuousSpace(2; periodic = true, extend = dims)
    model = ABM(Bird, space2d, scheduler = random_activation)
    for _ in 1:n_birds
        vel = Tuple(rand(2) * 2 .- 1)
        add_agent!(
            model,
            vel,
            speed,
            cohere_factor,
            separation,
            separate_factor,
            match_factor,
            visual_distance,
        )
    end
    index!(model)
    return model, flocking_agent_step!, dummystep
end

function flocking_agent_step!(bird, model)
    ## Obtain the ids of neighbors within the bird's visual distance
    ids = nearby_ids(bird, model, bird.visual_distance)
    ## Compute velocity based on rules defined above
    bird.vel =
        (
            bird.vel .+ cohere(bird, model, ids) .+ separate(bird, model, ids) .+
            match(bird, model, ids)
        ) ./ 2
    bird.vel = bird.vel ./ norm(bird.vel)
    ## Move bird according to new velocity and speed
    move_agent!(bird, model, bird.speed)
end

distance(a1, a2) = sqrt(sum((a1.pos .- a2.pos) .^ 2))

get_heading(a1, a2) = a1.pos .- a2.pos

function cohere(bird, model, ids)
    N = max(length(ids), 1)
    birds = model.agents
    coherence = (0.0, 0.0)
    for id in ids
        coherence = coherence .+ get_heading(birds[id], bird)
    end
    return coherence ./ N .* bird.cohere_factor
end

function separate(bird, model, ids)
    seperation_vec = (0.0, 0.0)
    N = max(length(ids), 1)
    birds = model.agents
    for id in ids
        neighbor = birds[id]
        if distance(bird, neighbor) < bird.separation
            seperation_vec = seperation_vec .- get_heading(neighbor, bird)
        end
    end
    return seperation_vec ./ N .* bird.separate_factor
end

function match(bird, model, ids)
    match_vector = (0.0, 0.0)
    N = max(length(ids), 1)
    birds = model.agents
    for id in ids
        match_vector = match_vector .+ birds[id].vel
    end
    return match_vector ./ N .* bird.match_factor
end


module BetaTest

using Test
using ReactiveMP
using Distributions
using Random

import DomainSets

@testset "Generic Functional Distributions" begin

    @testset "ContinuousUnivariateLogPdf" begin

        @testset "Constructor" begin
            f  = (x) -> -x ^ 2
            d1 = ContinuousUnivariateLogPdf(f)
            d2 = ContinuousUnivariateLogPdf(DomainSets.FullSpace(), f)

            @test typeof(d1) === typeof(d2)
            @test d1 ≈ d2
        end

        @testset "pdf/logpdf" begin
            d1 = ContinuousUnivariateLogPdf(DomainSets.FullSpace(), (x) -> -x ^ 2)

            f32_points1 = range(Float32(-10.0), Float32(10.0), length = 50)
            f64_points1 = range(-10.0, 10.0, length = 50)
            bf_points1  = range(BigFloat(-10.0), BigFloat(10.0), length = 50)
            points1 = vcat(f32_points1, f64_points1, bf_points1)

            @test all(map(p -> -p^2 == d1(p), points1))
            @test all(map(p -> -p^2 == logpdf(d1, p), points1))
            @test all(map(p -> exp(-p^2) == pdf(d1, p), points1))
            @test all(map(p -> -p^2 == d1([ p ]), points1))
            @test all(map(p -> -p^2 == logpdf(d1, [ p ]), points1))
            @test all(map(p -> exp(-p^2) == pdf(d1, [ p ]), points1))

            d2 = ContinuousUnivariateLogPdf(DomainSets.HalfLine(), (x) -> -x ^ 4)

            f32_points2 = range(Float32(0.0), Float32(10.0), length = 50)
            f64_points2 = range(0.0, 10.0, length = 50)
            bf_points2  = range(BigFloat(0.0), BigFloat(10.0), length = 50)
            points2 = vcat(f32_points2, f64_points2, bf_points2)

            @test all(map(p -> -p^4 == d2(p), points2))
            @test all(map(p -> -p^4 == logpdf(d2, p), points2))
            @test all(map(p -> exp(-p^4) == pdf(d2, p), points2))
            @test all(map(p -> -p^4 == d2([ p ]), points2))
            @test all(map(p -> -p^4 == logpdf(d2, [ p ]), points2))
            @test all(map(p -> exp(-p^4) == pdf(d2, [ p ]), points2))

            @test_throws AssertionError d2(-1.0)
            @test_throws AssertionError logpdf(d2, -1.0)
            @test_throws AssertionError pdf(d2, -1.0)
            @test_throws AssertionError d2([ -1.0 ])
            @test_throws AssertionError logpdf(d2, [ -1.0 ])
            @test_throws AssertionError pdf(d2, [ -1.0 ])

            @test_throws AssertionError d2(Float32(-1.0))
            @test_throws AssertionError logpdf(d2, Float32(-1.0))
            @test_throws AssertionError pdf(d2, Float32(-1.0))
            @test_throws AssertionError d2([ Float32(-1.0) ])
            @test_throws AssertionError logpdf(d2, [ Float32(-1.0) ])
            @test_throws AssertionError pdf(d2, [ Float32(-1.0) ])

            @test_throws AssertionError d2(BigFloat(-1.0))
            @test_throws AssertionError logpdf(d2, BigFloat(-1.0))
            @test_throws AssertionError pdf(d2, BigFloat(-1.0))
            @test_throws AssertionError d2([ BigFloat(-1.0) ])
            @test_throws AssertionError logpdf(d2, [ BigFloat(-1.0) ])
            @test_throws AssertionError pdf(d2, [ BigFloat(-1.0) ])

            d3 = ContinuousUnivariateLogPdf(DomainSets.FullSpace(Float32), (x) -> -x ^ 2)

            @test all(map(p -> -p^2 == d3(p), points1))
            @test all(map(p -> -p^2 == logpdf(d3, p), points1))
            @test all(map(p -> exp(-p^2) == pdf(d3, p), points1))
            @test all(map(p -> -p^2 == d3([ p ]), points1))
            @test all(map(p -> -p^2 == logpdf(d3, [ p ]), points1))
            @test all(map(p -> exp(-p^2) == pdf(d3, [ p ]), points1))

            d4 = ContinuousUnivariateLogPdf(DomainSets.FullSpace(BigFloat), (x) -> -x ^ 2)

            @test all(map(p -> -p^2 == d4(p), points1))
            @test all(map(p -> -p^2 == logpdf(d4, p), points1))
            @test all(map(p -> exp(-p^2) == pdf(d4, p), points1))
            @test all(map(p -> -p^2 == d4([ p ]), points1))
            @test all(map(p -> -p^2 == logpdf(d4, [ p ]), points1))
            @test all(map(p -> exp(-p^2) == pdf(d4, [ p ]), points1))

            d5 = ContinuousUnivariateLogPdf(DomainSets.HalfLine{Float32}(), (x) -> -x ^ 2)

            @test all(map(p -> -p^2 == d5(p), points2))
            @test all(map(p -> -p^2 == logpdf(d5, p), points2))
            @test all(map(p -> exp(-p^2) == pdf(d5, p), points2))
            @test all(map(p -> -p^2 == d5([ p ]), points2))
            @test all(map(p -> -p^2 == logpdf(d5, [ p ]), points2))
            @test all(map(p -> exp(-p^2) == pdf(d5, [ p ]), points2))

            d6 = ContinuousUnivariateLogPdf(DomainSets.HalfLine{BigFloat}(), (x) -> -x ^ 2)

            @test all(map(p -> -p^2 == d6(p), points2))
            @test all(map(p -> -p^2 == logpdf(d6, p), points2))
            @test all(map(p -> exp(-p^2) == pdf(d6, p), points2))
            @test all(map(p -> -p^2 == d6([ p ]), points2))
            @test all(map(p -> -p^2 == logpdf(d6, [ p ]), points2))
            @test all(map(p -> exp(-p^2) == pdf(d6, [ p ]), points2))
        end

        @testset "support" begin
            d1 = ContinuousUnivariateLogPdf(DomainSets.FullSpace(), (x) -> 1.0)
            @test minimum(support(d1)) === -Inf
            @test maximum(support(d1)) === Inf

            d2 = ContinuousUnivariateLogPdf(DomainSets.HalfLine(), (x) -> 1.0)
            @test minimum(support(d2)) === 0.0
            @test maximum(support(d2)) === Inf
        end

        @testset "vague" begin
            d = vague(ContinuousUnivariateLogPdf)

            @test typeof(d) <: ContinuousUnivariateLogPdf
            @test d ≈ ContinuousUnivariateLogPdf(DomainSets.FullSpace(), (x) -> 1.0)
        end

        @testset "prod" begin
            d1 = ContinuousUnivariateLogPdf(DomainSets.FullSpace(), (x) -> 2.0 * -x ^ 2)
            d2 = ContinuousUnivariateLogPdf(DomainSets.FullSpace(), (x) -> 3.0 * -x ^ 2)

            pr1 = prod(ProdAnalytical(), d1, d2)
            pt1 = ContinuousUnivariateLogPdf(DomainSets.FullSpace(), (x) -> logpdf(d1, x) + logpdf(d2, x))

            @test isapprox(pr1, pt1, atol = 1e-12)

            d3 = ContinuousUnivariateLogPdf(DomainSets.HalfLine(), (x) -> 2.0 * -x ^ 2)
            d4 = ContinuousUnivariateLogPdf(DomainSets.HalfLine(), (x) -> 3.0 * -x ^ 2)

            pr2 = prod(ProdAnalytical(), d3, d4)
            pt2 = ContinuousUnivariateLogPdf(DomainSets.HalfLine(), (x) -> logpdf(d3, x) + logpdf(d4, x))

            @test isapprox(pr2, pt2, atol = 1e-12)

            @test !isapprox(pr1, pr2, atol = 1e-12)
        end

        @testset "convert" begin
            d = DomainSets.FullSpace()
            l = (x) -> 1.0

            c = convert(ContinuousUnivariateLogPdf, d, l)
            @test typeof(c) <: ContinuousUnivariateLogPdf
            @test isapprox(c, ContinuousUnivariateLogPdf(d, l), atol = 1e-12)

            c2 = convert(ContinuousUnivariateLogPdf, c)
            @test typeof(c2) <: ContinuousUnivariateLogPdf
            @test isapprox(c2, ContinuousUnivariateLogPdf(d, l), atol = 1e-12)

        end

    end

end

end


#!/usr/bin/julia

# Copyright (C) 2021, 2022 Matheus Fernandes Bigolin <mfrdrbigolin@disroot.org>
# SPDX-License-Identifier: MIT

"Day Five, Hydrothermal Venture."

include("Utils.jl")

using LinearAlgebra
using .Utils

const Line = Tuple{Tuple{Int64, Int64}, Tuple{Int64, Int64}}

"Calculate the endpoints of the `line`."
function calculate_endpoints(line :: Line)
    x₀, x₁ = sort([first.(line)...]) .+ 1
    y₀, y₁ = sort([last.(line)...]) .+ 1

    return ((x₀, y₀), (x₁, y₁))
end

"Calculate the slope of the `line`."
function calculate_slope(line :: Line)
    x₀, x₁ = first.(line)
    y₀, y₁ = last.(line)

    return (y₁ - y₀)/(x₁ - x₀)
end

"Superimpose the `lines` over the `plane` and update it (mutating)."
function update_lines!(lines :: Array{Line}, plane :: Array{Int64, 2})
    for l in lines
        ((x₀, y₀), (x₁, y₁)) = calculate_endpoints(l)
        slope = calculate_slope(l)

        if slope in [-Inf, 0, Inf]
            plane[y₀:y₁, x₀:x₁] += ones(Int, (y₁ - y₀ + 1, x₁ - x₀ + 1))
        else
            # Number of rotations.
            rot = Int64(div(-slope + 1, 2))
            plane[y₀:y₁, x₀:x₁] += rotr90(Diagonal(ones(Int, x₁ - x₀ + 1)), rot)
        end
    end
end

"""Calculate the number of intersections of at least two `lines` in a ℤ
cartesian plane.  Return both the axial and the total count."""
function solve(lines :: Array{Line})
    xₘₐₓ = max([max(first.(l)...) for l in lines]...)
    yₘₐₓ = max([max(last.(l)...) for l in lines]...)

    plane = zeros(Int, (yₘₐₓ + 1, xₘₐₓ + 1))

    #= In the context of this day, I used the term “axial” as referring strictly
     = to vertical and horizontal lines only.
     =#
    axials = filter(P -> P[1][1] == P[2][1] || P[1][2] == P[2][2], lines)
    update_lines!(axials, plane)
    axial_count = count(>=(2), plane)

    diagonals = filter(P -> P[1][1] != P[2][1] && P[1][2] != P[2][2], lines)
    update_lines!(diagonals, plane)
    total_count = count(>=(2), plane)

    return (axial_count, total_count)
end

const POINTS_REX = r"(\d+),(\d+) -> (\d+),(\d+)"

function main()
    usage_and_exit(length(ARGS) != 1)

    input_data = readlines(ARGS[1])
    matched_data = match.(POINTS_REX, input_data)
    parsed_data = [parse.(Int, line.captures) for line in matched_data]
    lines = [((x₀, y₀), (x₁, y₁)) for (x₀, y₀, x₁, y₁) in parsed_data]

    S₁, S₂ = solve(lines)

    @assert S₁ == 6710
    @assert S₂ == 20121

    println(S₁)
    println(S₂)
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


import Pkg
Pkg.activate("..")

using CrimsonSkyline
using Distributions: Normal, LogNormal, Poisson
using StatsBase: mean

function log_rate_model(t::Trace, ::Tuple)
    loc = sample(t, "loc", Normal())
    scale = sample(t, "scale", LogNormal())
    sample(t, "log_rate", Normal(loc, scale))
end

function count_model(t::Trace, log_rate::CompiledSF)
    log_rate_value = sample(t, "log_rate", log_rate)
    sample(t, "count", Poisson(exp(log_rate_value)))
end

function main()
    compiled_log_rate_model = compile(
        TypedSF(log_rate_model, Tuple, Float64),
        (),
        ["log_rate"]
    )
    conditioned_count_model = condition(
        count_model,
        Dict("count" => 30)
    )
    results = inference(
        conditioned_count_model, Metropolis();
        params = (compiled_log_rate_model,),
        inference_params = Dict("num_iterations" => 55000, "burn" => 5000, "thin" => 500)
    )
    @info "Posterior rate: $(exp(mean(results, "log_rate")))"

end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


"""
    fillnodata_upstream(g, toposort, data, nodata)

Fill `nodata` upstream cells with the value from the first downstream valid cell, based on
directed acyclic graph `g`, topological order `toposort`, nodata value `nodata` and `data`
containing missing values. Returns filled data `data_out`.
"""
function fillnodata_upstream(g, toposort, data, nodata)
    data_out = copy(data)
    for v in reverse(toposort)  # down- to upstream
        idx_ds = outneighbors(g, v)
        if !isempty(idx_ds)
            if data_out[v] == nodata && data_out[only(idx_ds)] != nodata
                data_out[v] = data_out[only(idx_ds)]
            end
        end
    end
    return data_out
end

"""
    stream_order(g, toposort)

Return the Strahler streamorder based on directed acyclic graph `g` and topological order
`toposort`.
"""
function stream_order(g, toposort)
    n = length(toposort)
    strord = fill(1, n)
    for v in toposort
        inds_up = inneighbors(g, v)
        if !isempty(inds_up)
            sto_up = strord[inds_up]
            if length(findall(x -> x == maximum(sto_up), sto_up)) > 1
                strord[v] = (maximum(sto_up) + 1)
            else
                strord[v] = maximum(sto_up)
            end
        end
    end
    return strord
end

"""
    subbasins(g, streamorder, toposort, min_sto)

Return subbasins with a unique id starting at 1, based on a minimum streamorder `min_sto`,
directed acyclic graph `g` and topological order `toposort`.
"""
function subbasins(g, streamorder, toposort, min_sto)
    n = length(toposort)
    subbas = fill(0, n)

    i = 1
    for v in toposort
        if streamorder[v] < min_sto
            continue
        end
        ds_nodes = outneighbors(g, v)
        if !isempty(ds_nodes)
            if streamorder[v] != streamorder[only(ds_nodes)]
                subbas[v] = i
                i += 1
            end
        else
            # also set pits (without a downstream node)
            subbas[v] = i
            i += 1
        end
    end
    return subbas
end

"""
    subbasins_order(g, outlet, max_dist)

Group subbasins (down- to upstream), starting at the `outlet` of the basin, per distance (1
to maximum distance `max_dist`) from the `outlet`, based on the directed acyclic graph `g`
of subbasins. Returns grouped subbasins `order`, ordered from `max_dist` (including
subbasins without upstream neighbor and distance < `max_dist`) to distance 0 (`outlet`).
"""
function subbasins_order(g, outlet, max_dist)
    order = Vector{Vector{Int}}(undef, max_dist + 1)
    order[1] = [outlet]
    for i = 1:max_dist
        v = Vector{Int}()
        for n in order[i]
            ups_nodes = inneighbors(g, n)
            if !isempty(ups_nodes)
                append!(v, ups_nodes)
            end
        end
        order[i+1] = v
    end

    # move subbasins without upstream neighbor (headwater) to index [max_dist+1]
    for i = 1:max_dist
        for s in order[i]
            if isempty(inneighbors(g, s))
                append!(order[max_dist+1], s)
                filter!(e -> e ≠ s, order[i])
            end
        end
    end
    return reverse(order)
end

"""
    graph_from_nodes(graph, subbas, subbas_fill)

Extract directed acyclic graph `g` representing the flow network at subbasin level from
`subbas` containing nodes with a unique id representing subbasin outlets, `subbas_fill` with
subbasin ids for the complete domain, and directed acyclic graph `graph` representing the
flow network for each subbasin cell.
"""
function graph_from_nodes(graph, subbas, subbas_fill)
    n = maximum(subbas)
    g = DiGraph(n)
    for i = 1:n
        idx = findall(x -> x == i, subbas)
        ds_idx = outneighbors(graph, only(idx))
        to_node = subbas_fill[ds_idx]
        if !isempty(to_node)
            add_edge!(g, i, only(to_node))
        end
    end
    return g
end

"""
    function kinwave_set_subdomains(config, graph, toposort, index_pit)

Setup subdomains for parallel execution (threading) of the kinematic wave calculation.
Subdomains are subbasins based on a minumum streamorder `min_sto` (see also `subbasins(g,
streamorder, toposort, min_sto)`). Subbasins are extracted for each basin outlet in Vector
`index_pit`.

# Arguments
- `config` Struct that contains parsed TOML configuration
- `graph` directed acyclic graph of the kinematic wave domain
- `toposort` topological order of `graph`
- `index_pit` Vector with basin outlets (pits)

# Output
- `subbas_order` grouped subbasin ids (`Vector{Vector{Int}}`) ordered upstream (first index)
  to downstream (last index)
- `indices_subbas` list of indices per subbasin id stored as `Vector{Vector{Int}}`
- `topo_subbas` topological order per subbasin id stored as `Vector{Vector{Int}}`
"""
function kinwave_set_subdomains(config, graph, toposort, index_pit)

    if nthreads() > 1
        min_sto = get(config.model, "min_streamorder", 4)

        # extract basins (per outlet/pit), assign unique basin id
        n_pits = length(index_pit)
        basin = fill(0, length(toposort))
        basin[index_pit] = [1:n_pits;]
        basin_fill = fillnodata_upstream(graph, toposort, basin, 0)

        # pre-allocate the Vector with indices matching the topological order of the
        # complete domain (upstream neighbors are stored at these indices)
        index_toposort = fill(0, length(toposort))
        for (i, j) in enumerate(toposort)
            index_toposort[j] = i
        end

        order_subbas = Vector{Vector{Int}}()
        indices_subbas = Vector{Vector{Int}}()
        topo_subbas = Vector{Vector{Int}}()
        index = Vector{Int}()
        total_subbas = 0
        for i = 1:n_pits
            # extract subbasins per basin, make a graph at the subbasin level, calculate the
            # maximum distance of this graph, and group and order the subbasin ids from
            # upstream to downstream
            basin = findall(x -> x == i, basin_fill)
            g, _ = induced_subgraph(graph, basin)
            toposort_b = topological_sort_by_dfs(g)
            streamorder = stream_order(g, toposort_b)
            subbas = subbasins(g, streamorder, toposort_b, min_sto)
            subbas_fill = fillnodata_upstream(g, toposort_b, subbas, 0)
            n_subbas = max(length(subbas[subbas.>0]), 1)
            if n_subbas > 1
                graph_subbas = graph_from_nodes(g, subbas, subbas_fill)
                toposort_subbas = topological_sort_by_dfs(graph_subbas)
                dist = Graphs.Experimental.Traversals.distances(
                    Graph(graph_subbas),
                    toposort_subbas[end],
                )
                max_dist = maximum([dist; 1])
                v_subbas = subbasins_order(graph_subbas, toposort_subbas[end], max_dist)
            else
                v_subbas = [[1]]
            end
            # subbasins need a unique id (in case of multiple basins/outlets in the
            # kinematic wave domain)
            for n = 1:length(v_subbas)
                v_subbas[n] .= v_subbas[n] .+ total_subbas
            end
            total_subbas += n_subbas
            append!(order_subbas, v_subbas)
            append!(index, 1:length(v_subbas))
            # in case of multiple subbasins calculate topological order per subbasin
            # (subgraph of the corresponding basin graph g), and the indices that match the
            # subbasin topological order
            if n_subbas > 1
                for s = 1:n_subbas
                    subbas_s = findall(x -> x == s, subbas_fill)
                    sg, _ = induced_subgraph(g, subbas_s)
                    toposort_sg = topological_sort_by_dfs(sg)
                    push!(topo_subbas, basin[subbas_s[toposort_sg]])
                    push!(indices_subbas, index_toposort[basin[subbas_s[toposort_sg]]])
                end
            else
                push!(topo_subbas, basin[toposort_b])
                push!(indices_subbas, index_toposort[basin[toposort_b]])
            end
        end
        # reduce the order of subbasin ids by merging groups of subbasins that have the same
        # index (multiple basins/outlets in the kinematic wave domain)
        subbas_order = Vector{Vector{Int}}(undef, maximum(index))
        for m = 1:maximum(index)
            subbas_order[m] = reduce(vcat, order_subbas[index.==m])
        end
    else
        subbas_order = [[1]]
        indices_subbas = [[1:length(toposort);]]
        topo_subbas = [toposort]
    end

    return subbas_order, indices_subbas, topo_subbas
end


abstract type DiffusionMethod end
struct ExplicitDiffusion <: DiffusionMethod end
struct CrankNickolsonDiffusion <: DiffusionMethod end
struct ADIDiffusion <: DiffusionMethod end

function perform_step!(model::DiffusionModel, method::ExplicitDiffusion, cache::ExplicitDiffusionCache)
    u  = cache.u
    un = cache.uprev
    nx = model.grid.nx
    ny = model.grid.ny
    Δx = model.grid.Δx
    Δy = model.grid.Δy
    Δt = model.Δt
    ν  = model.parameters.ν

    copy!(un,u)
    @inbounds for j in 2:ny-1, i in 2:nx-1
        u[j,i] = un[j,i] +
            ν*Δt/(Δx*Δx)*(un[j  ,i+1]-2*un[j,i]+un[j  ,i-1]) +
            ν*Δt/(Δy*Δy)*(un[j+1,i  ]-2*un[j,i]+un[j-1,i  ])
    end

    return nothing
end
function perform_step!(model::DiffusionModel, method::CrankNickolsonDiffusion, cache::CNDiffusionCache)
    u   = cache.u
    un  = cache.uprev
    rhs = cache.rhs
    a1  = cache.a1
    a2  = cache.a2

    copy!(un,u)

    unV = @view un[2:end-1,2:end-1]
    unVC = @view unV[:]
    uV = @view u[2:end-1,2:end-1]
    uVC = @view uV[:]
    rhsV = @view rhs[2:end-1,2:end-1]
    rhsVC = @view rhsV[:]

    mul!(rhsVC, a2, unVC)
    uVC .= a1\rhsV[:]

    return nothing
end
function perform_step!(model::DiffusionModel, method::ADIDiffusion, cache::ADIDiffusionCache)
    u      = cache.u
    uhalf  = cache.uhalf
    un     = cache.uprev
    rhs    = cache.rhs
    a1x    = cache.a1x
    a1y    = cache.a1y
    a2x    = cache.a2x
    a2y    = cache.a2y
    nx, ny = model.grid.nx, model.grid.ny

    copy!(un,u)

    un_i(i) = @view un[2:end-1,i]
    u_i(i) = @view u[2:end-1,i]
    rhs_j(j) = @view rhs[j,2:end-1]
    rhs_i(i) = @view rhs[2:end-1,i]
    uhalf_j(j) = @view uhalf[j,2:end-1]

    # Step one
    for i = 2:nx-1
        mul!(rhs_i(i), a2y, un_i(i))
    end
    for j = 2:ny-1
        ldiv!(uhalf_j(j), a1x, rhs_j(j))
    end

    # Step two
    for j = 2:ny-1
        mul!(rhs_j(j), a2x, uhalf_j(j))
    end
    for i = 2:nx-1
        ldiv!(u_i(i), a1y, rhs_i(i))
    end
    return nothing
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Julia 1.2.0
#     language: julia
#     name: julia-1.2
# ---

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Semi-Lagrangian method
#
# Let us consider an abstract scalar advection equation of the form
# $$
# \frac{\partial f}{\partial t}+ a(x, t) \cdot \nabla f = 0.
# $$
# The characteristic curves associated to this equation are the solutions of the ordinary differential equations
# $$
# \frac{dX}{dt} = a(X(t), t)
# $$
# We shall denote by $X(t, x, s)$ the unique solution of this equation associated to the initial condition $X(s) = x$.
#
# The classical semi-Lagrangian method is based on a backtracking of characteristics. Two steps are needed to update the distribution function $f^{n+1}$ at $t^{n+1}$ from its value $f^n$ at time $t^n$ :
# 1. For each grid point $x_i$ compute $X(t^n; x_i, t^{n+1})$ the value of the characteristic at $t^n$ which takes the value $x_i$ at $t^{n+1}$.
# 2. As the distribution solution of first equation verifies
# $$f^{n+1}(x_i) = f^n(X(t^n; x_i, t^{n+1})),$$
# we obtain the desired value of $f^{n+1}(x_i)$ by computing $f^n(X(t^n;x_i,t^{n+1})$ by interpolation as $X(t^n; x_i, t^{n+1})$ is in general not a grid point.
#
# *[Eric Sonnendrücker - Numerical methods for the Vlasov equations](http://www-m16.ma.tum.de/foswiki/pub/M16/Allgemeines/NumMethVlasov/Num-Meth-Vlasov-Notes.pdf)*

# + {"slideshow": {"slide_type": "fragment"}}
using Plots, FFTW, LinearAlgebra, Statistics, BenchmarkTools

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Bspline
#
#

# + {"slideshow": {"slide_type": "fragment"}}
"""
    bspline(p, j, x)

Return the value at x in [0,1[ of the B-spline with integer
nodes of degree p with support starting at j.
Implemented recursively using the
[De Boor's Algorithm](https://en.wikipedia.org/wiki/De_Boor%27s_algorithm)

```math
B_{i,0}(x) := \\left\\{
\\begin{matrix}
1 & \\mathrm{if}  \\quad t_i ≤ x < t_{i+1} \\\\
0 & \\mathrm{otherwise}
\\end{matrix}
\\right.
```

```math
B_{i,p}(x) := \\frac{x - t_i}{t_{i+p} - t_i} B_{i,p-1}(x)
+ \\frac{t_{i+p+1} - x}{t_{i+p+1} - t_{i+1}} B_{i+1,p-1}(x).
```
"""
function bspline(p::Int, j::Int, x::Float64)
   if p == 0
       if j == 0
           return 1.0
       else
           return 0.0
       end
   else
       w = (x - j) / p
       w1 = (x - j - 1) / p
   end
   return (w * bspline(p - 1, j, x) + (1 - w1) * bspline(p - 1, j + 1, x))
end

# + {"slideshow": {"slide_type": "slide"}}
?bspline

# + {"slideshow": {"slide_type": "slide"}}
"""
    UniformMesh(xmin, xmax, nx)

1D uniform mesh data
"""
struct UniformMesh
   xmin  :: Float64
   xmax  :: Float64
   nx    :: Int
   dx    :: Float64
   x     :: Vector{Float64}
   function UniformMesh(xmin, xmax, nx)
      dx = (xmax - xmin) / nx
      x  = range(xmin, stop=xmax, length=nx+1)[1:end-1]
      new( xmin, xmax, nx, dx, x)
   end
end

# + {"slideshow": {"slide_type": "slide"}}
mesh = UniformMesh(-π, π, 10)
mesh.dx, mesh.nx

# + {"slideshow": {"slide_type": "slide"}}
"""
    advection!(f, p, mesh, v, dt)

function to advect the distribution function `f` with velocity `v`
during a time step `dt`. Interpolation method uses bspline periodic.
"""
function advection!(f    :: Array{Complex{Float64},2},
                    p    :: Int64,
                    mesh :: UniformMesh,
                    v    :: Vector{Float64},
                    dt   :: Float64)

   nx = mesh.nx
   nv = length(v)
   dx = mesh.dx
   modes = [2π * i / nx for i in 0:nx-1]

   # compute eigenvalues of degree p b-spline matrix
   eig_bspl  = zeros(Float64, nx)
   eig_bspl .= bspline(p, -div(p+1,2), 0.0)
   for i in 1:div(p+1,2)-1
      eig_bspl .+= bspline(p, i - (p+1)÷2, 0.0) * 2 .* cos.(i * modes)
   end
   eigalpha = zeros(Complex{Float64}, nx)

   fft!(f,1)

   for j in 1:nv
      @inbounds alpha = dt * v[j] / dx
      # compute eigenvalues of cubic splines evaluated at displaced points
      ishift = floor(-alpha)
      beta   = -ishift - alpha
      fill!(eigalpha,0.0im)
      for i in -div(p-1,2):div(p+1,2)
         eigalpha .+= (bspline(p, i-div(p+1,2), beta)
                        .* exp.((ishift+i) * 1im .* modes))
      end

      # compute interpolating spline using fft and properties of circulant matrices

      @inbounds f[:,j] .*= eigalpha ./ eig_bspl

   end

   ifft!(f,1)

end

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# # Vlasov-Poisson equation
# We consider the dimensionless Vlasov-Poisson equation for one species
# with a neutralizing background.
#
# $$
# \frac{\partial f}{\partial t}+ v\cdot \nabla_x f + E(t,x) \cdot \nabla_v f = 0, \\
# - \Delta \phi = 1 - \rho, E = - \nabla \phi \\
# \rho(t,x)  =  \int f(t,x,v)dv.
# $$
#
# - [Vlasov Equation - Wikipedia](https://en.wikipedia.org/wiki/Vlasov_equation)

# + {"slideshow": {"slide_type": "slide"}}
"""
    compute_rho(meshv, f)

Compute charge density
ρ(x,t) = ∫ f(x,v,t) dv
"""
function compute_rho(meshv::UniformMesh,
        f::Array{Complex{Float64},2})

   dv = meshv.dx
   rho = dv * sum(real(f), dims=2)
   vec(rho .- mean(rho)) # vec squeezes the 2d array returned by sum function
end

# + {"slideshow": {"slide_type": "slide"}}
"""
    compute_e(meshx, rho)
compute Ex using that -ik*Ex = rho
"""
function compute_e(meshx::UniformMesh, rho::Vector{Float64})
   nx = meshx.nx
   k =  2π / (meshx.xmax - meshx.xmin)
   modes = zeros(Float64, nx)
   modes .= k * vcat(0:div(nx,2)-1,-div(nx,2):-1)
   modes[1] = 1.0
   rhok = fft(rho)./modes
   rhok .*= -1im
   ifft!(rhok)
   real(rhok)
end

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# # Landau Damping
#
# [Landau damping - Wikipedia](https://en.wikipedia.org/wiki/Landau_damping)
# -

function landau( ϵ, kx, meshx, meshv)
    nx = meshx.nx
    nv = meshv.nx
    x  = meshx.x
    v  = meshv.x
    f  = zeros(Complex{Float64},(nx,nv))
    f .= (1.0.+ϵ*cos.(kx*x))/sqrt(2π) .* transpose(exp.(-0.5*v.^2))
    f
end

# + {"slideshow": {"slide_type": "fragment"}}
function advection_x!( f, p, meshx, v, dt)
    advection!(f, p, meshx, v, dt)
end

function advection_v!( f, fᵗ, p, meshx, meshv, ℰ, dt)
     dx = meshx.dx
     nx = meshx.nx
     ρ  = compute_rho(meshv, f)
     e  = compute_e(meshx, ρ)
     push!(ℰ, 0.5*log(sum(e.*e)*dx))
     transpose!(fᵗ, f)
     advection!(fᵗ, p, meshv, e, dt)
     transpose!(f, fᵗ)
end

# + {"slideshow": {"slide_type": "fragment"}}
function landau_damping(tf::Float64, nt::Int64)

  # Set grid
  p = 3
  nx, nv = 128, 256
  xmin, xmax = 0.0, 4π
  vmin, vmax = -6., 6.
  meshx = UniformMesh(xmin, xmax, nx)
  meshv = UniformMesh(vmin, vmax, nv)
  x, v = meshx.x, meshv.x

  # Set distribution function for Landau damping
  ϵ, kx = 0.001, 0.5
  f = landau( ϵ, kx, meshx, meshv)

  fᵗ = zeros(Complex{Float64},(nv,nx))

  # Set time domain
  dt = tf / nt

  # Run simulation
  ℰ = Float64[]

  for it in 1:nt

     advection_x!( f, p, meshx, v, 0.5dt)
     advection_v!( f, fᵗ, p, meshx, meshv, ℰ, dt)
     advection_x!( f, p, meshx, v, 0.5dt)

  end

  ℰ

end

# + {"slideshow": {"slide_type": "slide"}}
nt = 1000
tf = 100.0
t  = range(0.0, stop=tf, length=nt)
@time nrj = landau_damping(tf, nt);

# + {"slideshow": {"slide_type": "slide"}}
plot( t, nrj; label = "E")
plot!(t, -0.1533*t.-5.50; label="-0.1533t.-5.5")

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# # Callable type

# + {"slideshow": {"slide_type": "fragment"}}
"""
    Advection(f, p, mesh, v, nv, dt)

Advection type

"""
mutable struct Advection

    p        :: Int64
    mesh     :: UniformMesh
    modes    :: Vector{Float64}
    eig_bspl :: Vector{Float64}
    eigalpha :: Vector{Complex{Float64}}

    function Advection( p, mesh )
        nx        = mesh.nx
        modes     = zeros(Float64, nx)
        modes    .= [2π * i / nx for i in 0:nx-1]
        eig_bspl  = zeros(Float64, nx)
        eig_bspl  = zeros(Float64, nx)
        eig_bspl .= bspline(p, -div(p+1,2), 0.0)
        for i in 1:div(p+1,2)-1
            eig_bspl .+= bspline(p, i-(p+1)÷2, 0.0) * 2 .* cos.(i * modes)
        end
        eigalpha  = zeros(Complex{Float64}, nx)
        new( p, mesh, modes, eig_bspl, eigalpha )
    end

end

# + {"slideshow": {"slide_type": "slide"}}
function (adv :: Advection)(f    :: Array{Complex{Float64},2},
                            v    :: Vector{Float64},
                            dt   :: Float64)

   nx = adv.mesh.nx
   nv = length(v)
   dx = adv.mesh.dx

   fft!(f,1)

   @simd for j in 1:nv
      @inbounds alpha = dt * v[j] / dx
      # compute eigenvalues of cubic splines evaluated at displaced points
      ishift = floor(-alpha)
      beta   = -ishift - alpha
      fill!(adv.eigalpha,0.0im)
      for i in -div(adv.p-1,2):div(adv.p+1,2)
         adv.eigalpha .+= (bspline(adv.p, i-div(adv.p+1,2), beta)
                        .* exp.((ishift+i) * 1im .* adv.modes))
      end

      # compute interpolating spline using fft and properties of circulant matrices

      @inbounds f[:,j] .*= adv.eigalpha ./ adv.eig_bspl

   end

   ifft!(f,1)

end

# + {"slideshow": {"slide_type": "slide"}}
function landau_damping_hl(tf::Float64, nt::Int64)

  # Set grid
  p = 3
  nx, nv = 128, 256
  xmin, xmax = 0.0, 4π
  vmin, vmax = -6., 6.
  meshx = UniformMesh(xmin, xmax, nx)
  meshv = UniformMesh(vmin, vmax, nv)
  x, v = meshx.x, meshv.x
  dx = meshx.dx

  # Set distribution function for Landau damping
  ϵ, kx = 0.001, 0.5
  f = landau( ϵ, kx, meshx, meshv)
  fᵗ = zeros(Complex{Float64},(nv,nx))

  # Instantiate advection functions
  advection_x! = Advection(p, meshx)
  advection_v! = Advection(p, meshv)

  # Set time step
  dt = tf / nt

  # Run simulation
  ℰ = Float64[]

  for it in 1:nt

       advection_x!(f, v, 0.5dt)

       ρ = compute_rho(meshv, f)
       e = compute_e(meshx, ρ)

       push!(ℰ, 0.5*log(sum(e.*e)*dx))

       transpose!(fᵗ, f)
       advection_v!(fᵗ, e, dt)
       transpose!(f, fᵗ)

       advection_x!(f, v, 0.5dt)

  end

  ℰ

end

# + {"slideshow": {"slide_type": "slide"}}
nt = 1000
tf = 100.0
t  = range(0.0, stop=tf, length=nt)
@time nrj = landau_damping_hl(tf, nt);

# + {"slideshow": {"slide_type": "slide"}}
plot( t, nrj; label = "E")
plot!(t, -0.1533*t.-5.50; label="-0.1533t.-5.5")

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Metaprogramming

# + {"slideshow": {"slide_type": "fragment"}}
"""

    @Strang( push_t, push_v )

    Apply the second order Strang splitting

    push_t and push_v are two function calls with
    `dt` as argument.

"""
macro Strang(push_t, push_v)
    return esc(quote
        local full_dt = dt
        dt = 0.5full_dt
        $push_t
        dt = full_dt
        $push_v
        dt = 0.5full_dt
        $push_t
        dt = full_dt
    end)
end

# +
function landau_with_macro(tf::Float64, nt::Int64)

  # Set grid
  p = 3
  nx, nv = 128, 256
  xmin, xmax = 0.0, 4π
  vmin, vmax = -6., 6.
  meshx = UniformMesh(xmin, xmax, nx)
  meshv = UniformMesh(vmin, vmax, nv)
  x, v = meshx.x, meshv.x
  dx = meshx.dx

  # Set distribution function for Landau damping
  ϵ, kx = 0.001, 0.5
  f  = landau( ϵ, kx, meshx, meshv)
  fᵗ = zeros(Complex{Float64},(nv,nx))

  # Instantiate distribution functions
  advection_x! = Advection(p, meshx)
  advection_v! = Advection(p, meshv)

  push_t!( f, v, dt ) = advection_x!(f, v, dt)

  function push_v!( f, fᵗ, meshx, meshv, dt, ℰ)

      ρ = compute_rho(meshv, f)
      e = compute_e(meshx, ρ)

      push!(ℰ, 0.5*log(sum(e.*e)*dx))

      transpose!(fᵗ, f)
      advection_v!(fᵗ, e, dt)
      transpose!(f, fᵗ)

  end

  # Set time step
  dt = tf / nt

  # Run simulation
  ℰ = Float64[]

  for it in 1:nt

       @Strang( push_t!( f, v, dt ),
                push_v!( f, fᵗ, meshx, meshv, dt, ℰ))

  end

  ℰ

end
# -

nt = 1000
tf = 100.0
t  = range(0.0, stop=tf, length=nt)
@time nrj = landau_with_macro(tf, nt);
plot( t, nrj; label = "E")
plot!(t, -0.1533*t.-5.50; label="-0.1533t.-5.5")


@doc raw"""
```
t = find_threshold(Otsu(), histogram, edges)
```

Under the assumption that the histogram is bimodal the threshold is
set so that the resultant between-class variance is maximal.

# Output

Returns a real number `t` in `edges`. The `edges` parameter represents an
`AbstractRange` which specifies the intervals associated with the histogram bins.

# Details

Let ``f_i`` ``(i=1 \ldots I)`` denote the number of observations in the
``i``th bin of the histogram. Then the probability that an observation
belongs to the ``i``th bin is given by  ``p_i = \frac{f_i}{N}`` (``i = 1,
\ldots, I``), where ``N = \sum_{i=1}^{I}f_i``.

The choice of a threshold ``T`` partitions the data into two categories, ``C_0``
and ``C_1``. Let
```math
P_0(T) = \sum_{i = 1}^T p_i \quad \text{and} \quad P_1(T) = \sum_{i = T+1}^I p_i
```
denote the cumulative probabilities,
```math
\mu_0(T) = \sum_{i = 1}^T i \frac{p_i}{P_0(T)} \quad \text{and} \quad \mu_1(T) = \sum_{i = T+1}^I i \frac{p_i}{P_1(T)}
```
denote the means, and
```math
\sigma_0^2(T) = \sum_{i = 1}^T (i-\mu_0(T))^2 \frac{p_i}{P_0(T)} \quad \text{and} \quad \sigma_1^2(T) = \sum_{i = T+1}^I (i-\mu_1(T))^2 \frac{p_i}{P_1(T)}
```
denote the variances of categories ``C_0`` and ``C_1``, respectively.
Furthermore, let
```math
\mu = P_0(T)\mu_0(T) + P_1(T)\mu_1(T),
```
represent the overall mean,
```math
\sigma_b^2(T) = P_0(T)(\mu_0(T) - \mu)^2 + P_1(T)(\mu_1(T) - \mu)^2,
```
the between-category variance, and
```math
\sigma_w^2(T) = P_0(T) \sigma_0^2(T) +  P_1(T)\sigma_1^2(T)
```
the within-category variance, respectively.

Finding the discrete value ``T`` which maximises the function ``\sigma_b^2(T)``
produces the sought-after threshold value (i.e. the bin which determines the
threshold). As it turns out, that threshold value is equal to the threshold
decided by minimizing the within-category variances criterion ``\sigma_w^2(T)``.
Furthermore, that threshold is also the same as the threshold calculated by
maximizing the ratio of between-category variance to within-category variance.

# Arguments

The function arguments are described in more detail below.

##  `histogram`

An `AbstractArray` storing the frequency distribution.

##  `edges`

An `AbstractRange` specifying how the intervals for the frequency distribution
are divided.

# Example

Compute the threshold for the "cameraman" image in the `TestImages` package.

```julia
using TestImages, ImageContrastAdjustment, HistogramThresholding

img = testimage("cameraman")
edges, counts = build_histogram(img,256)
#=
  The `counts` array stores at index 0 the frequencies that were below the
  first bin edge. Since we are seeking a threshold over the interval
  partitioned by `edges` we need to discard the first bin in `counts`
  so that the dimensions of `edges` and `counts` match.
=#
t = find_threshold(Otsu(), counts[1:end], edges)
```

# Reference

1. Nobuyuki Otsu (1979). “A threshold selection method from gray-level histograms”. *IEEE Trans. Sys., Man., Cyber.* 9 (1): 62–66. [doi:10.1109/TSMC.1979.4310076](http://dx.doi.org/doi:10.1109/TSMC.1979.4310076)
"""
function find_threshold(algorithm::Otsu,  histogram::AbstractArray, edges::AbstractRange)
  N = sum(histogram)
  pdf = histogram / N
  first_bin = firstindex(pdf)
  last_bin = lastindex(pdf)
  cumulative_zeroth_moment = cumsum(pdf)
  cumulative_first_moment = cumsum(edges .* pdf)
  μ_T = cumulative_first_moment[end]
  σ²_T = sum( ((first_bin:last_bin) .- μ_T).^2  .* pdf )
  maxval = zero(eltype(first(pdf)))

  # Equation (6) for determining the probability of the first class.
  function ω(k)
    let cumulative_zeroth_moment = cumulative_zeroth_moment
      return cumulative_zeroth_moment[k]
    end
  end

  # Equation (7) for determining the mean of the first class.
  function μ(k)
   let cumulative_first_moment = cumulative_first_moment
     return cumulative_first_moment[k]
   end
  end

  # Equation (18) for determining the between-cass variance.
  function σ²(k)
    let μ_T = μ_T
      return (μ_T*ω(k) - μ(k))^2 / (ω(k) * (1-ω(k)))
    end
  end

  t = first_bin
  for k = first_bin:last_bin - 1
    val = σ²(k)
    if val  > maxval
       maxval = val
       t = k
    end
  end
  return edges[t]
end


using Test
using GeoParams
@testset "EnergyParameters.jl" begin

# This tests the MaterialParameters structure
CharUnits_GEO   =   GEO_units(viscosity=1e19, length=10km);

# Heat capacity ---------

# Constant heat capacity
cp1      =   ConstantHeatCapacity()
@test cp1.cp.val == 1050J/kg/K

Nondimensionalize!(cp1,CharUnits_GEO)
@test cp1.cp.val ≈ 1.3368075000000002e22

@test ComputeHeatCapacity(100.0, cp1) ≈ 1.3368075000000002e22  # compute

# Temperature-dependent heat capacity
# dimensional
T        =   (250:100:1250)*K;
cp2      =   T_HeatCapacity_Whittacker()
Cp       =   ComputeHeatCapacity(T,cp2)
@test sum(Cp) ≈ 11667.035717418683J/kg/K

# nondimensional
cp2_nd   =   T_HeatCapacity_Whittacker()
Nondimensionalize!(cp2_nd,CharUnits_GEO)
T_nd     =   Float64.(T/CharUnits_GEO.Temperature)
Cp_nd    =   ComputeHeatCapacity(T_nd,cp2_nd)
@test sum(Cp_nd) ≈ 1.4853886523631602e23

# Dimensionalize again and double-check the results
@test ustrip(sum(abs.(Cp_nd*CharUnits_GEO.heatcapacity - Cp))) < 1e-11

# Test with arrays
T_array     =  ustrip.(T)*ones(10)'
Cp_array    =  zeros(size(T_array))
P_array     =  zeros(size(T_array))
ComputeHeatCapacity!(Cp_array, P_array, T_array,cp1)
@test Cp_array[1] ≈ 1.3368075000000002e22

T_array     =  ustrip.(T)*ones(10)'
Cp_array    =  zeros(size(T_array))
P_array     =  zeros(size(T_array))
ComputeHeatCapacity!(Cp_array, P_array, T_array,cp2)
@test sum(Cp_array[:,1]) ≈ 11667.035717418683

T_array     =  ustrip.(T)*ones(10)'
Cp_array    =  zeros(size(T_array))
P_array     =  zeros(size(T_array))
ComputeHeatCapacity!(Cp_array, T_array,cp2)
@test sum(Cp_array[:,1]) ≈ 11667.035717418683

# Check that it works if we give a phase array
MatParam    =   Array{MaterialParams, 1}(undef, 2);
MatParam[1] =   SetMaterialParams(Name="Mantle", Phase=1,
                    HeatCapacity  = ConstantHeatCapacity());

MatParam[2] =   SetMaterialParams(Name="Crust", Phase=2,
                    HeatCapacity  = T_HeatCapacity_Whittacker());

# test computing material properties
n = 100;
Phases              = ones(Int64,n,n,n);
Phases[:,:,20:end] .= 2

Cp = zeros(size(Phases))
T  =  ones(size(Phases))*1500
P  =  zeros(size(Phases))

ComputeHeatCapacity!(Cp, Phases, P, T, MatParam)       # computation routine w/out P (not used in most heat capacity formulations)
@test sum(Cp[1,1,:]) ≈ 121399.0486067196

ComputeHeatCapacity!(Cp, Phases, T, MatParam)       # computation routine w/out P (not used in most heat capacity formulations)
@test sum(Cp[1,1,:]) ≈ 121399.0486067196

# test if we provide phase ratios
PhaseRatio  = zeros(n,n,n,3);
for i in CartesianIndices(Phases)
    iz = Phases[i]
    I = CartesianIndex(i,iz)
    PhaseRatio[I] = 1.0
end
ComputeHeatCapacity!(Cp, PhaseRatio, P, T, MatParam)
@test sum(Cp[1,1,:]) ≈ 121399.0486067196

# -----------------------

# Conductivity ----------

# Constant

# Constant conductivity
cond      =   ConstantConductivity()
@test Value(cond.k) == 3.0Watt/K/m

Nondimensionalize!(cond,CharUnits_GEO)
@test  Value(cond.k) ≈ 3.8194500000000007

@test ComputeConductivity(100.0, cond) ≈ 3.8194500000000007 # compute

# Temperature-dependent conductivity
# dimensional
T        =   (250:100:1250)*K;
cond2    =   T_Conductivity_Whittacker()
k        =   ComputeConductivity(T,cond2)
@test sum(k) ≈ 27.503366436682285Watt/K/m

# nondimensional
cond2_nd =   T_Conductivity_Whittacker()
Nondimensionalize!(cond2_nd,CharUnits_GEO)
T_nd     =   Float64.(T/CharUnits_GEO.Temperature)
k_nd     =   ComputeConductivity(T_nd,cond2_nd)
@test sum(k_nd) ≈ 35.01591097886205

# Dimensionalize again and double-check the results
@test ustrip(sum(abs.(k_nd*CharUnits_GEO.conductivity - k))) < 1e-11

# Check if we use arrays
T_array     =  ustrip.(T)*ones(100)'
k_array     =  copy(T_array)
P_array     =  copy(T_array)

ComputeConductivity!(k_array,P_array,T_array, cond)
@test k_array[1] ≈ 3.8194500000000007

ComputeConductivity!(k_array,P_array,T_array, cond2)
@test sum(k_array) ≈ 2750.3366436682285

k_TP    =   Set_TP_Conductivity["LowerCrust"]
ComputeConductivity!(k_array, P_array, T_array, k_TP)
@test sum(k_array) ≈ 2055.7129327367625

# Check that it works if we give a phase array
MatParam    =   Array{MaterialParams, 1}(undef, 3);
MatParam[1] =   SetMaterialParams(Name="Mantle", Phase=1,
                    Conductivity  = ConstantConductivity());

MatParam[2] =   SetMaterialParams(Name="Crust", Phase=2,
                    Conductivity  = T_Conductivity_Whittacker());

MatParam[3] =   SetMaterialParams(Name="MantleLithosphere", Phase=3,
                    Conductivity  = Set_TP_Conductivity["Mantle"]);

# test computing material properties
n = 100;
Phases              = ones(Int64,n,n,n);
Phases[:,:,20:end] .= 2
Phases[:,:,60:end] .= 3

PhaseRatio  = zeros(n,n,n,3);
for i in CartesianIndices(Phases)
    iz = Phases[i]
    I = CartesianIndex(i,iz)
    PhaseRatio[I] = 1.0
end

k   =  zeros(size(Phases))
T   =  ones(size(Phases))*1500
P   =  zeros(size(Phases))

ComputeConductivity!(k, Phases, P, T, MatParam)
@test sum(k) ≈ 1.9216938849389635e6

ComputeConductivity!(k, PhaseRatio, P, T, MatParam)
@test sum(k) ≈ 1.9216938849389635e6

#

######

# TP-dependent conductivity for different predefines cases
T       =   (250:100:1250)*K;
P       = 1e6Pa*ones(size(T))
List    = ["LowerCrust"   "Mantle"        "OceanicCrust"  "UpperCrust"]
Sol_kT  = [20.55712932736763 28.700405819019323 20.55712932736763 19.940302462417037]*Watt/K/m
for i=1:length(List)
    k_TP    =   Set_TP_Conductivity[List[i]]
    k       =   ComputeConductivity(P,T,k_TP)
    @test sum(k) ≈ Sol_kT[i]

    k_TP_nd  =  deepcopy(k_TP)
    Nondimensionalize!(k_TP_nd,CharUnits_GEO)
    T_nd     =   Float64.(T/CharUnits_GEO.Temperature)
    P_nd     =   Float64.(P/CharUnits_GEO.stress)
    k_nd     =   ComputeConductivity(P_nd,T_nd,k_TP_nd)

    @test ustrip(sum(abs.(k_nd*CharUnits_GEO.conductivity - k))) < 1e-11

end

T = [200 300; 400 500]*K
k1        =   ComputeConductivity(T,cond2)

# -----------------------

# Latent heat -----------
a = ConstantLatentHeat()
Q_L = ComputeLatentHeat(a)
@test Q_L == 400kJ/kg

Nondimensionalize!(a,CharUnits_GEO)
Q_L = ComputeLatentHeat(a)
@test Q_L ≈ 4e21
# -----------------------

# Radioactive heat ------
a = ConstantRadioactiveHeat()
H_r = ComputeRadioactiveHeat(a)
@test H_r ≈ 1.0e-6Watt/m^3

Nondimensionalize!(a,CharUnits_GEO)
H_r = ComputeRadioactiveHeat(a)
@test H_r == 0.1
# -----------------------

# Shear heating -------
Χ       = ConstantShearheating(1.0)

# Define parameters as vectors
τ       = [1 2 3 4]*1e6*Pa
ε       = [1 0.1 0.1 1]*1/s
ε_el    = [0.01 0.01 0.01 0.01]*1/s

τ_2D       = [1 2; 3 4]*1e6*Pa
ε_2D       = [1 0.1; 0.1 1]*1/s
ε_el_2D    = [0.01 0.01; 0.01 0.01]*1/s

# With elasticity
H_s1 = ComputeShearheating(τ,   ε,    ε_el,     Χ)
H_s2 = ComputeShearheating(τ_2D,ε_2D, ε_el_2D,  Χ)
@test H_s1 ≈ 5.4e6Pa/s
@test H_s2 ≈ 5.4e6Pa/s

# No elasticity
H_s3 = ComputeShearheating(τ,   ε,     Χ)
H_s4 = ComputeShearheating(τ_2D,ε_2D,  Χ)
@test H_s3 ≈ 5.5e6Pa/s
@test H_s4 ≈ 5.5e6Pa/s

# Now in non-dimensional units
τ       = [1 2 3 4]
ε       = [1 0.1 0.1 1]
ε_el    = [0.01 0.01 0.01 0.01]

τ_2D       = [1 2; 3 4]
ε_2D       = [1 0.1; 0.1 1]
ε_el_2D    = [0.01 0.01; 0.01 0.01]
Nondimensionalize!(Χ,CharUnits_GEO)

H_s1 = ComputeShearheating(τ,   ε,    ε_el,     Χ)
H_s2 = ComputeShearheating(τ_2D,ε_2D, ε_el_2D,  Χ)
H_s3 = ComputeShearheating(τ,   ε,     Χ)
H_s4 = ComputeShearheating(τ_2D,ε_2D,  Χ)
@test H_s1 ≈ 5.4
@test H_s2 ≈ 5.4
@test H_s3 ≈ 5.5
@test H_s4 ≈ 5.5
# -----------------------

end


"""
    exceeds_chi_squares_threshold(sol_dict::Dict, data::Dict; prob_false::Float64=0.05, suppress_display::Bool=false)

Standard bad data detection method that consists of performing a Chi squares analysis on the objective value, i.e.,
the sum of the residuals. The outcome of the analysis depends on the degrees of freedom of the problem, which are calculated calling
the `get_degrees_of_freedom` function (see below).

This function returns:
- a Boolean. If `true`, there are probably bad data points, if `false` there probably are not.
- the sum of the weighted squares of the residuals (which might be the optimization objective or not, depending on the settings) and the Χ² threshold value.

Arguments:
- `sol_dict`: solution dictionary, i.e., the default output dictionary of state estimation calculations,
- `data`: data input of the state estimation calculations, used to calculate the degrees of freedom,
- `prob_false`: probability of errors allowed in the Chi squares test, defaults to 0.05
- `suppress_display`: if `false`, the function also displays the result of the analysis, otherwise this is suppressed.

"""
function exceeds_chi_squares_threshold(sol_dict::Dict, data::Dict; prob_false::Float64=0.05, suppress_display::Bool=false)
    dof = get_degrees_of_freedom(data)
    chi2 = _DST.Chisq(dof)
    norm_residual = []
    for (m, meas) in data["meas"]
        h_x = sol_dict["solution"][string(meas["cmp"])][string(meas["cmp_id"])][string(meas["var"])]
        z = _DST.mean.(meas["dst"])
        r = abs.(h_x-z)
        σ = _DST.std.(meas["dst"])
        sol_dict["solution"]["meas"][m]["norm_res"] = [r[i]^2/σ[i]^2 for i in 1:length(meas["dst"])]
        for i in 1:length(meas["dst"]) push!(norm_residual, r[i]^2/σ[i]^2) end
    end
    sol_dict["J"] = sum(norm_residual)
    if !suppress_display
        sol_dict["J"] >= _STT.quantile(chi2, 1-prob_false) ? display("Chi-square test indicates presence of bad data.") : display("No bad data detected by Chi-square test.")
    end
    return sol_dict["J"] >= _STT.quantile(chi2, 1-prob_false), sol_dict["J"], _STT.quantile(chi2, 1-prob_false)
end
"""
    get_degrees_of_freedom(data::Dict)
Calculates the degrees of freedom of the state estimation problem.
These equal `m-n` where `m` is the number of measurements and `n` that of the variables.
In general, the system is fully described by the voltage variables of all its buses.
Each bus has two variables for each terminal: voltage angle and magnitude in polar form, real and imag part in rectangular form.
Thus, a three-phase bus has 6 variables, a single-phase bus has 2.
An exception to this are reference buses, i.e., those buses that have a fixed voltage angle that serves as a reference. It is assumed
that there is only one reference bus in the network, and therefore its 3 (if three-phase) or 1 (if single-phase) angle variables are removed
from the count of `n`.
Furthermore, since zero-injection buses are taken care of by equality constraints, these do not count as variables, and thus are not
included in the calculation of `n`, either. Zero-injection buses are defined as buses that have no loads or generators connected to it.
`m` is simply the sum of the measurements in `data["meas"]` where again, if a measurement refers to three phases `m+=3` else `m+=1`.
"""
function get_degrees_of_freedom(data::Dict)

    ref_bus = [bus for (_,bus) in data["bus"] if bus["bus_type"] == 3]

    @assert length(ref_bus) == 1 "There is more than one reference bus, double-check model"

    load_buses = ["$(load["load_bus"])" for (_, load) in data["load"]] # buses with demand (incl. negative demand, i.e., generation passed as negative load)
    gen_slack_buses = ["$(gen["gen_bus"])" for (_, gen) in data["gen"]] # buses with generators, including the slackbus
    non_zero_inj_buses = unique(vcat(load_buses, gen_slack_buses))

    @assert !isempty(non_zero_inj_buses) "This network has no active connected component, no point doing state estimation"

    n = sum([length(bus["terminals"]) for (b, bus) in data["bus"] if b ∈ non_zero_inj_buses ])*2-length(ref_bus[1]["terminals"]) # system variables: two variables per bus (e.g., angle and magnitude) per number of phases on the bus minus the angle variable(s) of a ref bus, which are fixed
    m = sum([length(meas["dst"]) for (_, meas) in data["meas"]])
    @assert m-n > 0 "system underdetermined or just barely determined, cannot perform bad data detection with this method"

    return m-n

end


require("DataFrames")
using DataFrames

df = read_table("demo/senate112.csv")
senator_names = df[:, 425]

dm = DataArray(df[:, 1:424])

u, d, v = svd(dm, 2)

ideal_points = u * diagm(d)

df = DataFrame()
df["Senator"] = senator_names
df["X"] = DataArray(ideal_points[:, 1])
df["Y"] = DataArray(ideal_points[:, 2])

write_table("demo/ideal_points.tsv", df)


@testset verbose=true "H" begin

    @testset "Test H(0) ($T)" for T in [BigFloat, Float64, Float32]
        # We have H_{n}^{m′,m}(0) = (-1)^m′ δ_{m′,m}
        let β = zero(T)
            expiβ = cis(β)
            for ℓₘₐₓ in 0:6  # Expect overflows for higher ℓ with Float32
                for m′ₘₐₓ in 0:ℓₘₐₓ
                    Hw = fill(T(NaN), WignerHsize(ℓₘₐₓ, m′ₘₐₓ))
                    H!(Hw, expiβ, ℓₘₐₓ, m′ₘₐₓ, H_recursion_coefficients(ℓₘₐₓ, T))
                    for n in 0:ℓₘₐₓ
                        for m′ in -min(n, m′ₘₐₓ):min(n, m′ₘₐₓ)
                            for m in abs(m′):n
                                H_rec = Hw[WignerHindex(n, m′, m, m′ₘₐₓ)]
                                if m′ == m
                                    @test H_rec ≈ T(-1)^m′ atol=eps(T)
                                else
                                    @test H_rec ≈ zero(T) atol=eps(T)
                                end
                            end
                        end
                    end
                end
            end
        end
    end

    @testset "Test H(π) ($T)" for T in [BigFloat, Float64, Float32]
        # We have H_{n}^{m′,m}(0) = (-1)^m′ δ_{m′,m}
        let β = T(π)
            expiβ = cis(β)
            for ℓₘₐₓ in 0:6  # Expect overflows for higher ℓ with Float32
                for m′ₘₐₓ in 0:ℓₘₐₓ
                    Hw = fill(T(NaN), WignerHsize(ℓₘₐₓ, m′ₘₐₓ))
                    H!(Hw, expiβ, ℓₘₐₓ, m′ₘₐₓ, H_recursion_coefficients(ℓₘₐₓ, T))
                    for n in 0:ℓₘₐₓ
                        for m′ in -min(n, m′ₘₐₓ):min(n, m′ₘₐₓ)
                            for m in abs(m′):n
                                H_rec = Hw[WignerHindex(n, m′, m, m′ₘₐₓ)]
                                if m′ == -m
                                    @test H_rec ≈ T(-1)^(n+m′) atol=eps(T(π))
                                else
                                    @test H_rec ≈ zero(T) atol=1.5eps(T(π))
                                end
                            end
                        end
                    end
                end
            end
        end
    end

    @testset "Compare H to explicit d ($T)" for T in [BigFloat, Float64, Float32]
        # This compares the H obtained via recurrence with the explicit Wigner d
        # d_{\ell}^{n,m} = \epsilon_n \epsilon_{-m} H_{\ell}^{n,m},
        for β in βrange(T)
            expiβ = cis(β)
            for ℓₘₐₓ in 0:2  # 2 is the max explicitly coded ℓ
                for m′ₘₐₓ in 0:ℓₘₐₓ
                    Hw = fill(T(NaN), WignerHsize(ℓₘₐₓ, m′ₘₐₓ))
                    H!(Hw, expiβ, ℓₘₐₓ, m′ₘₐₓ, H_recursion_coefficients(ℓₘₐₓ, T))
                    for n in 0:ℓₘₐₓ
                        for m′ in -min(n, m′ₘₐₓ):min(n, m′ₘₐₓ)
                            for m in -n:n
                                d_expl = ExplicitWignerMatrices.d_explicit(n, m′, m, expiβ)
                                d_rec = epsilon(m′) * epsilon(-m) * Hw[WignerHindex(n, m′, m, m′ₘₐₓ)]
                                @test d_rec ≈ d_expl atol=30eps(T) rtol=30eps(T)
                            end
                        end
                    end
                end

                # @show ℓₘₐₓ m′ₘₐₓ
                # for n in 0:ℓₘₐₓ
                #     @show n
                #     mat = fill(zero(T), 2n+1, 2n+1)
                #     for m′ in -min(n, m′ₘₐₓ):min(n, m′ₘₐₓ)
                #         for m in abs(m′):n
                #             mat[m′+n+1, m+n+1] =
                #                 Hw[SphericalFunctions.WignerHindex(n, m′, m, m′ₘₐₓ)].val
                #         end
                #     end
                #     println("H:")
                #     display(mat)
                #     println()
                #     mat = fill(zero(T), 2n+1, 2n+1)
                #     for m′ in -min(n, m′ₘₐₓ):min(n, m′ₘₐₓ)
                #         for m in abs(m′):n
                #             mat[m′+n+1, m+n+1] =
                #                 𝔇[SphericalFunctions.WignerDindex(n, m′, m, m′ₘₐₓ)].val
                #         end
                #     end
                #     println("𝔇 :")
                #     display(mat)
                #     println()
                #     println()
                # end
                # println()
            end
        end
    end

    @testset "Compare H to formulaic d ($T)" for T in [BigFloat, Float64, Float32]
        # This compares the H obtained via recurrence with the formulaic Wigner d
        # d_{\ell}^{n,m} = \epsilon_n \epsilon_{-m} H_{\ell}^{n,m},
        tol = ifelse(T === BigFloat, 100, 1) * 30eps(T)
        for β in βrange(T)
            expiβ = cis(β)
            for ℓₘₐₓ in 0:6  # Expect overflows for higher ℓ with Float32
                for m′ₘₐₓ in 0:ℓₘₐₓ
                    Hw = fill(T(NaN), WignerHsize(ℓₘₐₓ, m′ₘₐₓ))
                    H!(Hw, expiβ, ℓₘₐₓ, m′ₘₐₓ, H_recursion_coefficients(ℓₘₐₓ, T))
                    for n in 0:ℓₘₐₓ
                        for m′ in -min(n, m′ₘₐₓ):min(n, m′ₘₐₓ)
                            for m in -n:n
                                d_form = ExplicitWignerMatrices.d_formula(n, m′, m, expiβ)
                                d_rec = epsilon(m′) * epsilon(-m) * Hw[WignerHindex(n, m′, m, m′ₘₐₓ)]
                                @test d_rec ≈ d_form atol=tol rtol=tol
                            end
                        end
                    end
                end
            end
        end
    end

end