module Density

# This implements different methods to compute density
#
# If you want to add a new method here, feel free to do so.
# Remember to also export the function name in GeoParams.jl (in addition to here)

using Parameters, LaTeXStrings, Unitful
using ..Units
using ..PhaseDiagrams
#using ..MaterialParameters
using GeoParams: AbstractMaterialParam, AbstractMaterialParamsStruct
import Base.show

abstract type AbstractDensity <: AbstractMaterialParam end

export  ComputeDensity,         # calculation routines
        ComputeDensity!,        # in place calculation
        ConstantDensity,        # constant
        PT_Density              # P & T dependent density

# Constant Density -------------------------------------------------------
"""
    ConstantDensity(ρ=2900kg/m^3)

Set a constant density:
```math
    \\rho  = cst
```
where ``\\rho`` is the density [``kg/m^3``].
"""
@with_kw_noshow mutable struct ConstantDensity <: AbstractDensity
    equation::LaTeXString   =   L"\rho = cst"
    ρ::GeoUnit              =   2900kg/m^3                # density
end

# Calculation routines
function ComputeDensity(P,T, s::ConstantDensity)
    @unpack ρ   = s
    if length(T)>1
        return Value(ρ).*ones(size(T))
    else
        return ρ*1.0
    end
end

function ComputeDensity!(rho::AbstractArray{<:AbstractFloat,N},P::AbstractArray{<:AbstractFloat,N},T::AbstractArray{<:AbstractFloat,N}, s::ConstantDensity) where N
    @unpack ρ   = s

    rho .= NumValue(ρ)

    return nothing
end

# Print info
function show(io::IO, g::ConstantDensity)
    print(io, "Constant density: ρ=$(g.ρ.val)")
end
#-------------------------------------------------------------------------

# Pressure & Temperature dependent density -------------------------------
"""
    PT_Density(ρ0=2900kg/m^3, α=3e-5/K, β=1e-9/Pa, T0=0C, P=0MPa)

Set a pressure and temperature-dependent density:
```math
    \\rho  = \\rho_0 (1.0 - \\alpha (T-T_0) + \\beta  (P-P_0) )
```
where ``\\rho_0`` is the density [``kg/m^3``] at reference temperature ``T_0`` and pressure ``P_0``,
``\\alpha`` is the temperature dependence of density and ``\\beta`` the pressure dependence.

"""
@with_kw_noshow mutable struct PT_Density <: AbstractDensity
    equation::LaTeXString   =   L"\rho = \rho_0(1.0-\alpha (T-T_0) + \beta (P-P_0)"
    ρ0::GeoUnit             =   2900kg/m^3                  # density
    α::GeoUnit              =   3e-5/K                      # T-dependence of density
    β::GeoUnit              =   1e-9/Pa                     # P-dependence of density
    T0::GeoUnit             =   0C                          # Reference temperature
    P0::GeoUnit             =   0MPa                        # Reference pressure
end

# Calculation routine
function ComputeDensity(P,T, s::PT_Density)
    @unpack ρ0,α,β,P0, T0   = s

    ρ = ρ0*(1.0 - α*(T - T0) + β*(P - P0) )

    return ρ
end

function ComputeDensity!(ρ, P,T, s::PT_Density)
    @unpack ρ0,α,β,P0, T0   = s

    ρ[:] = ρ0*(1.0 - α*(T-T0) + β*(P-P0) )

    return ρ
end

function ComputeDensity!(ρ::AbstractArray{<:AbstractFloat},P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, s::PT_Density)
    @unpack ρ0,α,β,P0, T0   = s
    ρ0 = NumValue(ρ0)
    α  = NumValue(α)
    β  = NumValue(β)
    P0 = NumValue(P0)
    T0 = NumValue(T0)

    ρ  .= ρ0*(1.0 .- α*( T .- T0) + β*(P .- P0) )

    return nothing
end

# Print info
function show(io::IO, g::PT_Density)
    print(io, "P/T-dependent density: ρ0=$(g.ρ0.val), α=$(g.α.val), β=$(g.β.val), T0=$(g.T0.val), P0=$(g.P0.val)")
end
#-------------------------------------------------------------------------

#-------------------------------------------------------------------------
# Phase diagrams
"""
    ComputeDensity(P,T, s::PhaseDiagram_LookupTable)

Interpolates density as a function of `T,P`
"""
function ComputeDensity(P,T, s::PhaseDiagram_LookupTable)
    return s.Rho.(T,P)
end

"""
    ComputeDensity!(rho::AbstractArray{<:AbstractFloat}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, s::PhaseDiagram_LookupTable)

In-place computation of density as a function of `T,P`, in case we are using a lookup table.
"""
function ComputeDensity!(rho::AbstractArray{<:AbstractFloat}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, s::PhaseDiagram_LookupTable)
    rho[:] = s.Rho.(T,P)
    return nothing
end

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

"""
    ComputeDensity!(rho::AbstractArray{<:AbstractFloat}, Phases::AbstractArray{<:Integer}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

In-place computation of density `rho` for the whole domain and all phases, in case a vector with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This assumes that the `Phase` of every point is specified as an Integer in the `Phases` array.

# Example
```julia
julia> MatParam    =   Array{MaterialParams, 1}(undef, 2);
julia> MatParam[1] =   SetMaterialParams(Name="Mantle", Phase=1,
                        CreepLaws= (PowerlawViscous(), LinearViscous(η=1e23Pa*s)),
                        Density   = PerpleX_LaMEM_Diagram("test_data/Peridotite.in"));
julia> MatParam[2] =   SetMaterialParams(Name="Crust", Phase=2,
                        CreepLaws= (PowerlawViscous(), LinearViscous(η=1e23Pa*s)),
                        Density   = ConstantDensity(ρ=2900kg/m^3));
julia> Phases = ones(Int64,400,400);
julia> Phases[:,20:end] .= 2
julia> rho     = zeros(size(Phases))
julia> T       =  ones(size(Phases))
julia> P       =  ones(size(Phases))*10
julia> ComputeDensity!(rho, Phases, P,T, MatParam)
julia> rho
400×400 Matrix{Float64}:
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46  …  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46  …  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
    ⋮                                            ⋮     ⋱                     ⋮
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46  …  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 3334.46  3334.46  3334.46  3334.46  3334.46  3334.46     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
```
The routine is made to minimize allocations:
```julia
julia> julia> @time ComputeDensity!(rho, Phases, P,T, MatParam)
0.003121 seconds (49 allocations: 3.769 MiB)
```
"""
function ComputeDensity!(rho::AbstractArray{<:AbstractFloat, N}, Phases::AbstractArray{<:Integer, N}, P::AbstractArray{<:AbstractFloat, N},T::AbstractArray{<:AbstractFloat, N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct, 1}) where N

    for i = 1:length(MatParam)

        if !isnothing(MatParam[i].Density)
            # Create views into arrays (so we don't have to allocate)
            ind = Phases .== i;
            rho_local   =   view(rho, ind )
            P_local     =   view(P  , ind )
            T_local     =   view(T  , ind )

            ComputeDensity!(rho_local, P_local, T_local, MatParam[i].Density[1] )
        end

    end

end

"""
    ComputeDensity!(rho::AbstractArray{<:AbstractFloat,N}, PhaseRatios::AbstractArray{<:AbstractFloat, M}, P::AbstractArray{<:AbstractFloat,N},T::AbstractArray{<:AbstractFloat,N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

In-place computation of density `rho` for the whole domain and all phases, in case a vector with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This assumes that the `PhaseRatio` of every point is specified as an Integer in the `PhaseRatios` array, which has one dimension more than the data arrays (and has a phase fraction between 0-1)

"""
function ComputeDensity!(rho::AbstractArray{<:AbstractFloat, N}, PhaseRatios::AbstractArray{<:AbstractFloat, M}, P::AbstractArray{<:AbstractFloat, N},T::AbstractArray{<:AbstractFloat, N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct, 1}) where {N,M}

    if M!=(N+1)
        error("The PhaseRatios array should have one dimension more than the other arrays")
    end

    rho .= 0.0;
    for i = 1:length(MatParam)

        rho_local   = zeros(size(rho))
        Fraction    = selectdim(PhaseRatios,M,i);
        if (maximum(Fraction)>0.0) & (!isnothing(MatParam[i].Density))

            ComputeDensity!(rho_local, P, T, MatParam[i].Density[1] )

            rho .= rho .+ rho_local.*Fraction
        end

    end

end

end


using PhysicalConstants.CODATA2018
using LinearAlgebra
using Combinatorics
import DimensionalData.formatdims

# Exports
export AntennaArray
export ArrayFactor
export baselines
export calibrate

function λ(f::Quantity)
    SpeedOfLightInVacuum/f |> u"m"
end

function λ(f::Real)
    SpeedOfLightInVacuum/(f*u"Hz") |> u"m"
end

function phase_location(b::AbstractArray,ϕ::Number,θ::Number,freq::Number)
    # Direction of the source
    ŝ = [sind(θ)*cosd(ϕ),sind(θ)*sind(ϕ),cosd(θ)]
    k = @. 2π/λ(freq) * ŝ
    # Phase
    cis(-k⋅b)
end

"""
   AntennaArray(locations,ϕ,θ)
Creates a new `AntennaArray` which represents the `locations` phased in the direction of `ϕ` and `θ`.
"""
function AntennaArray(locations,ϕ,θ,freq)
    AntennaArray(phase_location.(locations,ϕ,θ,freq),locations)
end

"""
    aa(ϕ,θ,freq)
Finds the relative amplitude of the array factor for the array `aa` in the azimuth ϕ,
elevation `θ` at frequency freq. Angles in degrees.
"""
function (aa::AntennaArray)(ϕ::Number,θ::Number,freq::Number)
    (phase_location.(aa.locations,ϕ,θ,freq) ⋅ aa.excitations) / sum(abs.(aa.excitations))
end

"""
    ArrayFactor(aa)
Constructs a `RadiationPattern` representative of the array factor of the array `aa`
sampled over azimuth `ϕ` and elevation `θ` at frequenc[y|ies] `freq`.
"""
function ArrayFactor(aa::AntennaArray,ϕ,θ,freq::AbstractVector)
    data = [aa(phi,theta,f) for phi ∈ ϕ, theta ∈ θ, f ∈ freq]
    RadiationPattern(data,ϕ,θ,freq)
end

function ArrayFactor(aa::AntennaArray,ϕ,θ,freq::Number)
    data = [aa(phi,theta,freq) for phi ∈ ϕ, theta ∈ θ]
    RadiationPattern(data,ϕ,θ)
end

""""
    baselines(af)
Returns all the baseline vectors for the antenna locations in `aa`.
"""
function baselines(aa::AntennaArray)
    [.-(pair...) for pair ∈ combinations(aa.locations,2)]
end

"""
    calibrate(array,phase_center)
For a given antenna array `array`, return a new `AntennaArray` whose phase center has been adjusted
to `phase_center`.
"""
function calibrate()

end


function game(decisionMatrix::DataFrame; verbose=false)::GameResult

    newDecisionMatrix = similar(decisionMatrix)

    nrow, ncol = size(decisionMatrix)

    minmat = minimum(decisionMatrix)
    modified = false

    if minmat < 0
        modified = true
        newDecisionMatrix = decisionMatrix .- minmat
    end

    dm = convert(Matrix, newDecisionMatrix)

    model = Model(Cbc.Optimizer);
    MOI.set(model, MOI.Silent(), !verbose)

    @variable(model, x[1:nrow])
    @objective(model, Min, sum(x))

    for i in 1:nrow
        @constraint(model, sum(x[1:nrow] .* dm[:, i]) >= 1)
    end

    for i in 1:nrow
        @constraint(model, x[i] >= 0)
    end

    optimize!(model);

    values = JuMP.value.(x)

    gamevalue = 1 / objective_value(model)
    p = gamevalue * values

    if modified
        gamevalue -= abs(minmat)
    end

    result = GameResult(
        p,
        gamevalue
    )

    return result
end


using OmegaGrav
using Dierckx
using Plots, LaTeXStrings
include("dodlnMh.jl")
# %% Compute the matter power spectrum using CLASS
include("compute_pk_class.jl")
# %% Compute dΩgrav/dlnMh and dΩtherm/dlnMh at z = 0
z = 0
# Define functions to return linear and non-linear power spectra
# Note: The CLASS code takes wavenumbers in units of 1/Mpc (no h) and
# return power spectra in units of Mpc^3 (no 1/h^3).
pkcb_class(kovh) = cosmo.pk_cb_lin(kovh * h0, z) * h0^3
# Spline interpolate in log(k)
lnk = log(1e-4):0.05:log(100)
pkcb = Spline1D(lnk, pkcb_class.(exp.(lnk)))
# %% Define a function to return dΩgrav/dlnMh from Halos, excluding the neutrino contribution
dΩgdlnMh(lnMh) = dogravdlnMh(x -> pkcb(log(x)), z, Ωm, lnMh, Ωcb)
# %% Define a function to return dΩtherm/dlnMh from Halos, excluding the neutrino contribution
dΩthdlnMh(lnMh) = dothermdlnMh(x -> pkcb(log(x)), z, Ωm, h0, lnMh, Ωcb)

#%% Plot results and save to figure2.pdf
fb = params["omega_b"] / (params["omega_cdm"] + params["omega_b"])
lnMh = log(1e11):0.1:log(5e15)
p = plot(
   exp.(lnMh),
   -2 * fb * dΩgdlnMh.(lnMh) / 3,
   xaxis = :log,
   yaxis = :log,
   xlab = "Halo Mass [M⊙/h]",
   ylab = L"d\Omega/d\ln M",
   lab = L"\Omega_W^{ref,halo}",
   legend = :topleft,
   legendfontsize = 12,
   labelfontsize = 15,
   lw = 5,
   c = :green,
   ls = :dot,
   ylims = [1e-10, 1e-8],
   xlims = [1e11, 5e15],
)
p = plot!(exp.(lnMh), dΩthdlnMh.(lnMh), lw = 5, c = 1, lab = L"\Omega_{th}")

# %% Compute dΩgrav/dlnMh and dΩtherm/dlnMh at z = 1
z = 1
pkcb_class(kovh) = cosmo.pk_cb_lin(kovh * h0, z) * h0^3
pkcb = Spline1D(lnk, pkcb_class.(exp.(lnk)))
dΩgdlnMh(lnMh) = dogravdlnMh(x -> pkcb(log(x)), z, Ωm, lnMh, Ωcb)
dΩthdlnMh(lnMh) = dothermdlnMh(x -> pkcb(log(x)), z, Ωm, h0, lnMh, Ωcb)
p = plot!(
   exp.(lnMh),
   -2 * fb * dΩgdlnMh.(lnMh) / 3,
   lw = 2,
   c = :green,
   ls = :dot,
   lab = "",
)
p = plot!(exp.(lnMh), dΩthdlnMh.(lnMh), lw = 2, c = 1, lab = "")
savefig("figure2.pdf")
display(p)

# %% Clean CLASS (the equivalent of the struct_free() in the `main`
# of CLASS. This step is primordial when running in a loop over different
# cosmologies, as you will saturate your memory very fast if you ommit
# it.
cosmo.struct_cleanup()
# If you want to change completely the cosmology, you should also
# clean the arguments, otherwise, if you are simply running on a loop
# of different values for the same parameters, this step is not needed
cosmo.empty()


@testset "Running PELT tests..." begin

    n = 1000;        # Number of samples
    λ = 100;         # Frequencey of changes

    ########################
    # Normal mean segments #
    ########################
    μ, σ = Normal(0.0, 10.0), 1.0
    sample, cps = @changepoint_sampler n λ Normal(μ, σ)
    seg_costs = NormalMeanSegment(sample);
    @test_nowarn PELT(seg_costs, n);
    @test_nowarn @PELT sample Normal(:?, σ)

    #######################
    # Normal var segments #
    #######################
    μ, σ = 1.0, Uniform(2.0, 15.0)
    sample, cps = @changepoint_sampler n λ Normal(μ, σ)
    seg_costs = NormalVarSegment(sample, μ)
    @test_nowarn PELT(seg_costs, n)
    @test_nowarn @PELT sample Normal(μ, :?)

    ############################
    # Exponential changepoints #
    ############################
    μ = Uniform(0.0, 10.0)
    sample, cps = @changepoint_sampler n λ Exponential(μ)
    seg_costs = ExponentialSegment(sample)
    @test_nowarn PELT(seg_costs, n)
    @test_nowarn @PELT sample Exponential(:?)

    ########################
    # Poisson changepoints #
    ########################
    μ = Uniform(0.0, 10.0)
    sample, cps = @changepoint_sampler n λ Poisson(μ)
    seg_costs = PoissonSegment(sample)
    @test_nowarn PELT(seg_costs, n)
    @test_nowarn @PELT sample Poisson(:?)

end


#@testset "BigMat.jl" begin
    A        = randn(20,5)
    Attm,err = MPT_SVD(A,[5 2 2; 5 1 1],0.0)

    # test qr
    @run Qttm,R = qr(Attm,0.0)
    Q = mpo2mat(Qttm)
    norm(qr(A).R-R)
    norm(Array(qr(A).Q)-mpo2mat(Qttm))

    Q = randTTmsubspaceiter(q,A,O,acc)
    W,S,V = svd(Attm)

    @run U,S,V    = TNrSVD(Attm,30,10,10,0.0)

#end


using FFTW, .ForwardDiff
import FFTW.fft!;

function fft!(x::AbstractArray{T1}) where {T1 <: Complex{cpx}} where {cpx <: ForwardDiff.Dual{Tg, T, N}} where {Tg, T <: Real, N}
    Xcomplex = [ForwardDiff.value(real(el)) + 1im * ForwardDiff.value(imag(el)) for el in x]
    Xder1 = reduce(hcat, [ForwardDiff.partials(real(el)) + 1im * ForwardDiff.partials(imag(el)) for el in x])'
    planned_fft = plan_fft!(similar(Xcomplex))
    Xcomplex .= planned_fft * Xcomplex
    for col in eachcol(Xder1)
        col .= planned_fft * col
    end
    npartials = size(Xder1, 2)
    epss = [ForwardDiff.Dual{Tg}(0, ntuple(i -> i == k, Val(npartials))) for k = 1:npartials]
    x .= Xcomplex .+ Xder1 * epss
    nothing
end


module GWAlpha_ML_module

using Distributed
using SharedArrays
using DelimitedFiles
using Distributions
using LinearAlgebra
using Optim
using ProgressMeter
include("significance_testing_module.jl")

"""
# _____________________________________________________________________________________________________________________________________________
# Negative log-likelihood of the Beta distribution parameters (shape1 and shape2) given the distribution of the allele frequencies across pools

`neg_log_likelihood_cdfbeta(beta::Array{Float64,1}, data_A::Array{Float64,1}, data_B::Array{Float64,1})`

# Inputs
1. beta [Array{Float64,1}]: 4-element vector of the beta distribution shape parameters, 2 for each of the data distributions
2. data_A [Array{Float64,1}; xi ∉ 0]: allele frequency distribution across pools
3. data_B [Array{Float64,1}; xi ∉ 0]: additive inverse of the allele frequency distribution across pools

# Output
1. Negative log-likelihood of the Beta distribution parameters (shape1 and shape2) given the distribution of the allele frequencies across pools

# Examples
```
beta = [1.0, 1.0, 0.5, 1.0]
f = rand(5); bins = repeat([0.2], inner=5); p = 0.5
data_A = cumsum( (f .* bins ./ p) )
data_B = cumsum( ((1 .- f) .* bins ./ (1 .- p)) )
GWAlpha_ML_module.neg_log_likelihood_cdfbeta(beta, data_A, data_B)
```
"""
function neg_log_likelihood_cdfbeta(beta::Array{Float64,1}, data_A::Array{Float64,1}, data_B::Array{Float64,1})
	-sum(
		 log.(10,
		 	 (Distributions.cdf.(Distributions.Beta(beta[1], beta[2]), data_A)) .-
		 	 (Distributions.cdf.(Distributions.Beta(beta[1], beta[2]), append!(zeros(1), data_A[1:(length(data_A)-1)])))
		     )
		 )     -sum(
		 log.(10,
		 	 (Distributions.cdf.(Distributions.Beta(beta[3], beta[4]), data_B)) .-
		 	 (Distributions.cdf.(Distributions.Beta(beta[3], beta[4]), append!(zeros(1), data_B[1:(length(data_B)-1)])))
		     )
		 )
end

"""
# _____________________________________________________________________________________________________________________________________________
# Iterating function per locus for genome-wide additive allelic effects (alpha) association using Pool-seq data

`GWAlpha_ML_iterator(COUNTS::SharedArray{Int64,3}, snp::Int64, BINS::Array{Float64,1}, MAF::Float64, MIN::Float64, MAX::Float64, SD::Float64)`

# Inputs
1. COUNTS [SharedArray{Int64,3}]: 3-dimensional array of allele counts for each locus for each of the 6 SNP alleles (A:T:C:G:N:DEL) for each pool
2. snp [Int64]: index number of the SNP
3. BINS [Array{Float64,1}]: vector of pool sizes
4. MAF [Float64]: minor allele frequency threshold to filter out alleles (MAF <= f > 1-MAF)
5. MIN [Float64]: minimum phenotype value
6. MAX [Float64]: maximum phenotype value
7. SD [Float64]: standard deviation of phenotype value

# Outputs
1. OUT_ALPHA: vector of additive allelic effects of alleles in the locus
2. OUT_allele: vector of alleles passing minor allele frequency (MAF) threshold filtering
3. OUT_snp: vector of repeated locus id for each of the allelic effects or allele
4, OUT_1: vector of ones and zeros indicating the indices of the alleles passing minor allele frequency (MAF) threshold filtering among the 6 SNP alleles per locus (A:T:C:G:N:DEL)
5, OUT_pA: vector of allele frequencies per locus

# Examples
```
NPOOLS=5; NLOCI=20; BINS=repeat([1.0/NPOOLS],inner=NPOOLS); MAF=0.001; MIN=0.00; MAX=10.00; SD=1.00;
COUNTS = SharedArrays.SharedArray{Int64,3}(NPOOLS, 6, NLOCI)
COUNTS[:,:,:] = Int.(round.(rand(NPOOLS, 6, NLOCI)*100))
for snp in 1:NLOCI
	out = GWAlpha_ML_module.GWAlpha_ML_iterator(COUNTS, snp, BINS, MAF, MIN, MAX, SD)
	println(out)
end
```
"""
function GWAlpha_ML_iterator(COUNTS::SharedArray{Int64,3}, snp::Int64, BINS::Array{Float64,1}, MAF::Float64, MIN::Float64, MAX::Float64, SD::Float64)
	OUT_ALPHA = convert(Array{Float64}, zeros(6))
	OUT_allele = zeros(6)
	OUT_snp = zeros(6)
	OUT_1 = zeros(6)
	OUT_pA = convert(Array{Float64}, zeros(6))
	# #parse allele counts from the sync file
	# COUNTS = zeros(Int64, NPOOLS, 6)
	# for i in 1:NPOOLS
	# 	COUNTS[i,:] = [parse(Int64, x) for x in split.(SYNC[snp, 4:(NPOOLS+3)], [':'])[i]]
	# end
	#convert to frequencies per pool
	counts = COUNTS[:,:,snp]
	FREQS = counts ./ ( sum(counts, dims=2) .+ 1e-10 ) #added 1e-10 to the denominator to avoid NAs in pools with no allele counts (zero depth; which should actually have been filtered out after mpileup using awk)
	allele_freqs = sum(FREQS .* BINS, dims=1)
	#iterate across alleles while filtering by MAF
	if (sum(counts) != 0.0)
		if (minimum(allele_freqs[allele_freqs .!= 0.0]) >= MAF) & (maximum(allele_freqs) < (1.0 - MAF)) #locus filtering by mean MAF
			for allele in 1:6
				if (allele_freqs[allele] > 0.0) & (maximum(FREQS[:,allele]) < 0.999999)  #allele filtering remove alleles with no counts and that the number of pools with allele frequency close to one should not occur even once!
				# if (sum(FREQS[:,allele] .== 0.0) < NPOOLS) & (sum(FREQS[:,allele] .> 0.999999) < 1) #filter-out alleles with at least 1 pool fixed for that allele because it causes a failure in the optimization
					freqA = FREQS[:, allele]
					pA = sum(freqA .* BINS)
					pB = 1 - pA
					BINA = (freqA .* BINS) ./ pA
					BINB = ( (1 .- freqA) .* BINS ) ./ (1-pA)
					percA = cumsum(BINA)
					percB = cumsum(BINB)

					### optimize (minimize) -log-likelihood of these major allele frequencies modelled as a beta distribution
					# using Nelder-Mead optimization or Box minimisation (try-catch if one or the other fails with preference to Nelder-Mead)
					lower_limits = [1e-20, 1e-20, 1e-20, 1e-20]
					upper_limits = [1.0, 1.0, 1.0, 1.0]
					initial_values = [0.1, 0.1, 0.1, 0.1]
					BETA = try
						Optim.optimize(beta->neg_log_likelihood_cdfbeta(beta, percA, percB), initial_values, NelderMead())
					catch
						try
							Optim.optimize(beta->neg_log_likelihood_cdfbeta(beta, percA, percB), lower_limits, upper_limits, initial_values)
						catch ### lower limits of 1e-20 to 1e-6 causes beta dist parameter values to shrink to zero somehow - so we're setting lower limits to 1e-5 instead
							lower_limits = [1e-5, 1e-5, 1e-5, 1e-5]
							Optim.optimize(beta->neg_log_likelihood_cdfbeta(beta, percA, percB), lower_limits, upper_limits, initial_values)
						end
					end
					MU_A = MIN + ((MAX-MIN)*BETA.minimizer[1]/(BETA.minimizer[1]+BETA.minimizer[2]))
					MU_B = MIN + ((MAX-MIN)*BETA.minimizer[3]/(BETA.minimizer[3]+BETA.minimizer[4]))

					### compute alpha
					W_PENAL = 2*sqrt(pA*pB)
					ALPHA = W_PENAL*(MU_A - MU_B) / SD
					OUT_ALPHA[allele] =  ALPHA
					OUT_allele[allele] =  allele
					OUT_snp[allele] =  snp
					OUT_1[allele] =  1
					OUT_pA[allele] =  pA
				else
					## for explicit zero effects for null (low to none) frequency alleles
					OUT_ALPHA[allele] =  0.0
					OUT_allele[allele] =  allele
					OUT_snp[allele] =  snp
					OUT_1[allele] =  0
					OUT_pA[allele] =  allele_freqs[allele]
				end
			end
		end
	end
	return([OUT_ALPHA, OUT_allele, OUT_snp, OUT_1, OUT_pA])
end

"""
# _____________________________________________________________________________________________________________________________________________
# Genome-wide additive allelic effects (alpha) association using Pool-seq data

`GWAlpha_ML(filename_sync::String, filename_phen_py::String, MAF::Float64)`

# Inputs
1. *filename_sync* [String]: filename of the synchronized pileup file, i.e. the genotype file from pool sequencing (Pool-seq)
2. *filename_phen_py* [String]; filename of the phenotype specifications (with ".py" extension name for cross-compatibility with [GWAlpha.py](https://github.com/aflevel/GWAlpha))
3. *MAF* [Float64]: minor allele frequency threshold to filter out alleles (MAF <= f > 1-MAF)

# Output
1. Tuple of vectors of the additive allelic affects across loci with the corresponding locus identification and significance test statistics
	- CHROM
	- POS
	- ALLELE
	- FREQ
	- ALPHA
	- PVALUES
	- LOD

# Example
`OUT = GWAlpha_ML_module.GWAlpha_ML(filename_sync="test/test.sync", filename_phen_py="test/test.py", MAF=0.001)`
"""
function GWAlpha_ML(;filename_sync::String, filename_phen_py::String, MAF::Float64)
	### load the sync and phenotype files
	SYNC = DelimitedFiles.readdlm(filename_sync, '\t')
	phen = DelimitedFiles.readdlm(filename_phen_py, '\t')

	### gather phenotype specifications
	NPOOLS = length(split(phen[5], ['=', ',', '[', ']', ';'])) - 3 #less the first leading and trailing elements
	if length(split(phen[1], ['=', '\"'])) < 3
		global NAME = split(phen[1], ['=', '\''])[3]
	else
		global NAME = split(phen[1], ['=', '\"'])[3]
	end
	SD = parse.(Float64, split(phen[2], ['=',';'])[2])
	MIN = parse.(Float64, split(phen[3], ['=',';'])[2])
	MAX = parse.(Float64, split(phen[4], ['=',';'])[2])
	PERC = parse.(Float64, split(phen[5], ['=', ',', '[', ']', ';'])[3:(NPOOLS+1)])
	QUAN = parse.(Float64, split(phen[6], ['=', ',', '[', ']', ';'])[3:(NPOOLS+1)])
	BINS = append!([x for x in PERC], 1) - append!(zeros(1), PERC)

	### gather genotype (allele frequency) specifications
	NSNP = size(SYNC)[1]
	n_pools_sync = size(SYNC)[2] - 3
	if NPOOLS != n_pools_sync
		println("The number of pools with phenotype data does not match the number of pools with allele frequency data!")
		println("Please check you input files :-)")
		println("Remove leading and intervening whitespaces in the phenotype file.")
		exit()
	else
		n_pools_sync = nothing #clear out contents of this redundant n_pools variable
	end

	### creating an allele counts SharedArray from SYNC to minimize memory usage
	### input shared array
	COUNTS = SharedArrays.SharedArray{Int64,3}(NPOOLS, 6, size(SYNC)[1])
	progress_bar = ProgressMeter.Progress(NPOOLS, dt=1, desc="Converting the sync file into a SharedArray of allele counts: ",  barglyphs=BarGlyphs("[=> ]"), barlen=50, color=:yellow)
	for i in 1:NPOOLS
		COUNTS[i,:,:] = parse.(Int64, hcat(split.(SYNC[:, 4:(NPOOLS+3)], [':'])[:, i]...))
		ProgressMeter.update!(progress_bar, i)
	end
	### output shared arrays
	alpha_out = SharedArrays.SharedArray{Float64,1}(NSNP*6)
	allele_id = SharedArrays.SharedArray{Int64,1}(NSNP*6)
	locus_id = SharedArrays.SharedArray{Int64,1}(NSNP*6)
	locus_w_eff = SharedArrays.SharedArray{Int64,1}(NSNP*6)
	allele_freq = SharedArrays.SharedArray{Float64,1}(NSNP*6)

	if length(Distributed.procs()) > 1
		println("Performing GWAlpha_ML in parallel...")
	 	@time x = @sync @distributed for snp in 1:NSNP
			# println(snp)
			OUT = GWAlpha_ML_iterator(COUNTS, snp, BINS, MAF, MIN, MAX, SD)
			idx = collect(((6*snp)-5):(6*snp))
			alpha_out[idx] = OUT[1]
			allele_id[idx] = OUT[2]
			locus_id[idx] = OUT[3]
			locus_w_eff[idx] = OUT[4]
			allele_freq[idx] = OUT[5]
		end
	else
		for snp in 1:NSNP
			# println(snp)
			OUT = GWAlpha_ML_iterator(COUNTS, snp, BINS, MAF, MIN, MAX, SD)
			idx = collect(((6*snp)-5):(6*snp))
			alpha_out[idx] = OUT[1]
			allele_id[idx] = OUT[2]
			locus_id[idx] = OUT[3]
			locus_w_eff[idx] = OUT[4]
			allele_freq[idx] = OUT[5]
		end
	end
	ALPHA_OUT = alpha_out[locus_w_eff .== 1]
	ALLELE_ID_INT = allele_id[locus_w_eff .== 1]
	LOCUS_ID = locus_id[locus_w_eff .== 1]
	LOCUS_W_EFF = locus_w_eff[locus_w_eff .== 1]
	ALLELE_FREQ = allele_freq[locus_w_eff .== 1]

	### estimate heuristic p-values
	P_VALUES, LOD = significance_testing_module.estimate_pval_lod(convert(Array{Float64,1}, ALPHA_OUT))
	### output
	ALLELE_ID = repeat(["N"], inner=length(ALLELE_ID_INT))
	for i in 1:length(ALLELE_ID) #convert int allele ID into corresponding A, T, C, G, N, DEL
		if ALLELE_ID_INT[i] == 1; ALLELE_ID[i] = "A"
		elseif ALLELE_ID_INT[i] == 2; ALLELE_ID[i] = "T"
		elseif ALLELE_ID_INT[i] == 3; ALLELE_ID[i] = "C"
		elseif ALLELE_ID_INT[i] == 4; ALLELE_ID[i] = "G"
		elseif ALLELE_ID_INT[i] == 5; ALLELE_ID[i] = "N"
		elseif ALLELE_ID_INT[i] == 6; ALLELE_ID[i] = "DEL"
		end
	end
	OUT = (CHROM=convert(Array{Any,1},SYNC[LOCUS_ID,1]), POS=convert(Array{Int64,1},SYNC[LOCUS_ID,2]), ALLELE=convert(Array{Any,1},ALLELE_ID), FREQ=convert(Array{Float64,1},ALLELE_FREQ), ALPHA=convert(Array{Float64},ALPHA_OUT), PVALUES=convert(Array{Float64},P_VALUES), LOD=convert(Array{Float64},LOD))
	return(OUT)
end

end # end of GWAlpha_ML_module


using Test
using MLE
using Distributions
using Random
using Optim
using RCall

@testset "Test MLE" begin
    Random.seed!(1234)
    data = rand(Gumbel(1, 1), 1_000)

    @testset "Test Normal" begin
	f = fit_numerical_mle(Normal, data)
	@test f.converged

	f = fit_numerical_mle(Normal, data, (μ=10, σ=2))
	@test f.converged

	f = fit_numerical_mle(Normal, data, (μ=10, σ=2), bounds=(μ=[5,Inf], ))
	@test f.converged
    end

    f = fit_numerical_mle(Gumbel, data)
    @test f.converged
    @test isapprox(f.result, [0.9843, 0.9863]; atol=1e-3)

    f = fit_numerical_mle(Gumbel, data, holding=(; θ=1))
    @test f.converged

    pois_data = [1,2,1,3]

    f = fit_numerical_mle(Poisson, pois_data)
    @test f.converged
    @test isapprox(f.result, [1.75])

    R"""
    library(MASS)
    r_mle <- fitdistr($pois_data, "poisson")
    """
    @rget r_mle
    @test isapprox(f.result[1], r_mle[:estimate])
    @test isapprox(f.sd[1], r_mle[:sd])
    @test isapprox(f.loglikelihood, r_mle[:loglik])

    f = fit_numerical_mle(NegativeBinomial, pois_data)
    @test f.converged
    @test isapprox(mean(f.fit), 1.75; atol=1e-4)

    R"""
    library(MASS)
    r_mle <- fitdistr($pois_data, "negative binomial")
    """
    @rget r_mle
    @test isapprox(mean(f.fit), r_mle[:estimate][2]; atol=1e-4)

    f = fit_numerical_mle(NegativeBinomial, pois_data, holding=(;r=0.5))
    @test f.converged
    @test isapprox(mean(f.fit), 1.75; atol=1e-4)

    binom_data = [0, 1, 2, 0, 1, 2]
    f = fit_numerical_mle(Binomial, binom_data, holding=(;n=4))
    @test f.converged
    @test isapprox(f.result[1], 0.25; atol=1e-4)
    @test isapprox(f.result[1], mean(binom_data) / 4)

    f = fit_numerical_mle(Geometric, pois_data)
    @test f.converged
    @test isapprox(mean(f.fit), mean(pois_data); atol=1e-4)

    beta_data = [0.1, 0.2, 0.1, 0.2, 0.9]
    f = fit_numerical_mle(Beta, beta_data)
    @test f.converged
    @test isapprox(mean(f.fit), 0.35409; atol=1e-4)

    R"""
    library(MASS)
    r_mle <- fitdistr($beta_data, "beta", start=list(shape1=1, shape2=1))
    """
    @rget r_mle
    @test isapprox(f.result, r_mle[:estimate]; atol=1e-4)
    @test isapprox(f.vcov, r_mle[:vcov]; atol=1e-4)
    @test isapprox(f.sd, r_mle[:sd]; atol=1e-4)
    @test isapprox(f.loglikelihood, r_mle[:loglik])

    expon_data = [0.1, 0.2, 1.0, .3, .05]
    f = fit_numerical_mle(Exponential, expon_data)
    @test f.converged
    @test isapprox(mean(f.fit), mean(expon_data); atol=1e-2)

    R"""
    library(MASS)
    r_mle <- fitdistr($expon_data, "exponential")
    """
    @rget r_mle
    @test isapprox(1 ./ f.result[1], r_mle[:estimate]; atol=1e-2)
    @test isapprox(f.loglikelihood, r_mle[:loglik])

    f = fit_numerical_mle(Weibull, expon_data)
    @test f.converged

    R"""
    library(MASS)
    r_mle <- fitdistr($expon_data, "weibull")
    """
    @rget r_mle
    @test isapprox(f.result, r_mle[:estimate]; atol=1e-4)
    @test isapprox(f.sd, r_mle[:sd]; atol=1e-4)
    @test isapprox(f.loglikelihood, r_mle[:loglik])
end


function part1(input)
    state = parse_input(input, 3)
    for i = 1:6
        state = simulate_cycle(state)
    end
    return length(state)
end

function part2(input)
    state = parse_input(input, 4)
    for i = 1:6
        state = simulate_cycle(state)
    end
    return length(state)
end

function parse_input(input, dims)
    state = Set{CartesianIndex{dims}}()
    for (i, line) in enumerate(readlines(input))
        for (j, c) in enumerate(split(line, ""))
            if c == "#"
                push!(state, CartesianIndex(i, j, fill(0, dims - 2)...))
            end
        end
    end
    return state
end

function simulate_cycle(state::Set{CartesianIndex{N}}) where N
    neighbors = Dict{CartesianIndex{N}, Int}()
    neighborhood = CartesianIndices(ntuple(i -> -1:1, N))
    for position in state
        for Δ in neighborhood
            if Δ == zero(CartesianIndex{N})
                neighbors[position + Δ] = get(neighbors, position + Δ, 0) + 1
            else
                neighbors[position + Δ] = get(neighbors, position + Δ, 0) + 2
            end
        end
    end
    return Set(position
               for (position, num_neighbors) in neighbors
               if 5 <= num_neighbors <= 7)
end


module BumpOnTailSimulation

using LaTeXStrings
using LinearAlgebra
using Particles
using Particles.BumpOnTail
using Plots

# simulation parameters

const dt = 1e-1             # timestep
const T = 50                # final time
const nₜ = Int(div(T, dt))  # nb. of timesteps
const nₕ = 16               # nb. of elements
const p  = 3                # spline degree

const nₚ = Int(5e4)         # nb. of particles

const vmin = -10
const vmax = +10

# bump-on tail instability parameters: κ, ε, a, v₀, σ
# in order: spartial perturbation wave number and amplitude of s.p., fast particle fraction, speed, and temperature
const params = (κ = 0.3, ε = 0.03, a = 0.1, v₀= 4.5, σ = 0.5, χ = 1.0)

const L = 2π / params.κ     # domain length
const h = L/nₕ              # element width

const IP = VPIntegratorParameters(dt, nₜ, nₜ+1, nₕ, nₚ)

function run()

    # B-spline Poisson solver
    poisson = PoissonSolverPBSplines(p, IP.nₕ, L)

    # initial data
    P = BumpOnTail.draw_accept_reject(nₚ, params)
    # P = BumpOnTaildraw_importance_sampling(nₚ, params)

    # initial potential
    solve!(poisson, P.x, P.w)

    # plot initial particle distribution
    F = plot_particles(P.x, P.v, P.w, 0, L, vmin, vmax)
    savefig("bump-on-tail-init.png")

    F = plot_distribution(P.x, P.v, P.w, 0, L, vmin, vmax; nx=48, nv=32)
    savefig("bump-on-tail-init-f.png")

    # create integrate cache
    IC = VPIntegratorCache(IP)

    # integrate all time steps
    # for t = 1:nₜ
        integrate_vp!(P, poisson, params, IP, IC)
    # end

    # plot diagnostics
    plot(0:dt:dt*nₜ, IC.W, yaxis = :log, legend = :none)
    savefig("bump-on-tail-W.png")

    plot(0:dt:dt*nₜ, IC.K, yaxis = :log, legend = :none)
    savefig("bump-on-tail-K.png")

    plot(0:dt:dt*nₜ, IC.W .+ IC.K, yaxis = :log, legend = :none)
    savefig("bump-on-tail-E.png")

    # create animation and save to file
    anim = @animate for t in 0:IP.nₜ
        plot_particles(IC.X[:,t+1], IC.V[:,t+1], P.w, 0, L, vmin, vmax)
    end
    gif(anim, "bump-on-tail.gif", fps=10)
end

end # module

using .BumpOnTailSimulation

BumpOnTailSimulation.run()


using Test
using LinearAlgebra
using DataStructure101
const DS = DataStructure101

@testset "Graph" begin
    @testset "Di Adjacency Matrix" begin
        @testset "Basic operate" begin
            n = 10

            g = DS.Graph(n, :di_matrix)
            row_str = "\t false" ^ n
            data_str = "$row_str\n" ^ n
            @test repr(g) == "Graph(\n$data_str)"

            @test g isa DS.AbstractGraph
            @test g isa DS.DiAdjacencyMatrix
            @test DS.nv(g) == n
            @test DS.ne(g) == 0

            DS.relate!(g, 2, 5)
            @test DS.ne(g) == 1
            @test DS.neighbor(g, 2) == Int[5]

            g = DS.Graph(n, :w_di_matrix, random=false)

            row_str = "\t 0.0" ^ n
            data_str = "$row_str\n" ^ n
            @test repr(g) == "Graph(\n$data_str)"

            @test g isa DS.AbstractGraph
            @test g isa DS.WeightedDiAdjacencyMatrix
            @test DS.nv(g) == n
            @test DS.ne(g) == 0

            DS.relate!(g, 2, 5, 10.5)
            @test DS.ne(g) == 1
            @test DS.neighbor(g, 2) == Int[5]

            w = zeros(Float64, n, n)
            w[2, 5] = 10.5
            @test DS.weight(g) == w
        end

        @testset "Random Walk" begin
            n = 100
            city = DS.Graph(n, :w_di_matrix)
            p = DS.probability(city)

            test_prob = []
            for i=1:n push!(test_prob, sum(p[i, 1:end]) ≈ 1.0) end
            @test all(test_prob)

            # Random Walk Begin
            x = collect(1.0:n)
            x = DS.random_walk(city, x, 10000)
            @test p' * x ≈ x
        end

    end

    @testset "Di AdjacencyList" begin
        g = DS.Graph(10, :di_list)
        @test g isa DS.AbstractGraph
        @test g isa DS.DiAdjacencyList

        @test g[8] == 8
        @test g[11] == -1

        DS.relate!(g, 10, 9)
        DS.relate!(g, 10, 8)

        DS.relate!(g, 9, 7)
        DS.relate!(g, 9, 6)
        DS.relate!(g, 8, 5)
        DS.relate!(g, 8, 4)

        DS.relate!(g, 7, 3)
        DS.relate!(g, 7, 2)
        DS.relate!(g, 6, 1)

        @test DS.nv(g) == 10
        @test DS.ne(g) == 9
        @test DS.degree(g, 8) == 2
        @test DS.degree(g, 6) == 1

        @test DS.neighbor(g, 10) == [9, 8]

        @test repr(g) == "DiAdjacencyList(" *
            "1[], 2[], 3[], 4[], 5[], " *
            "6[1, ], 7[3, 2, ], 8[5, 4, ], 9[7, 6, ], 10[9, 8, ], )"

        DS.unrelate!(g, 6, 1)
        @test repr(g) == "DiAdjacencyList(" *
            "1[], 2[], 3[], 4[], 5[], " *
            "6[], 7[3, 2, ], 8[5, 4, ], 9[7, 6, ], 10[9, 8, ], )"
    end

    @testset "WeightedAdjacencyList" begin
        g = DS.Graph(10, :w_list)
        @test g isa DS.WeightedAdjacencyList
        @test g isa DS.AbstractGraph

        @test g[8] == 8
        @test g[11] == -1

        DS.relate!(g, 10, 9, 9.0)
        DS.relate!(g, 10, 8, 8.0)

        DS.relate!(g, 9, 7, 7.0)
        DS.relate!(g, 9, 6, 6.0)
        DS.relate!(g, 8, 5, 5.0)
        DS.relate!(g, 8, 4, 4.0)

        DS.relate!(g, 7, 3, 3.0)
        DS.relate!(g, 7, 2, 2.0)
        DS.relate!(g, 6, 1, 1.0)

        @test DS.nv(g) == 10
        @test DS.ne(g) == 9
        @test DS.degree(g, 8) == 17
        @test DS.degree(g, 6) == 7

        @test DS.neighbor(g, 10) == [9, 8]

        @test repr(g) == "WeightedAdjacencyList{Float64}(" *
            "1[6(1.0), ], 2[7(2.0), ], 3[7(3.0), ], 4[8(4.0), ], " *
            "5[8(5.0), ], 6[9(6.0), 1(1.0), ], " *
            "7[9(7.0), 3(3.0), 2(2.0), ], 8[10(8.0), 5(5.0), 4(4.0), ], " *
            "9[10(9.0), 7(7.0), 6(6.0), ], 10[9(9.0), 8(8.0), ], )"

        DS.unrelate!(g, 6, 1)
        @test DS.ne(g) == 8
        @test DS.degree(g, 1) == 0
        @test DS.degree(g, 6) == 6
        @test repr(g) == "WeightedAdjacencyList{Float64}(" *
            "1[], 2[7(2.0), ], 3[7(3.0), ], 4[8(4.0), ], " *
            "5[8(5.0), ], 6[9(6.0), ], " *
            "7[9(7.0), 3(3.0), 2(2.0), ], 8[10(8.0), 5(5.0), 4(4.0), ], " *
            "9[10(9.0), 7(7.0), 6(6.0), ], 10[9(9.0), 8(8.0), ], )"
    end
end


using Revise
using IMPES
using Plots

# using Plots, LinearAlgebra, IMPES
# using AnalyticalEOR

#* Define Relative Perms
kr_ww = SingleRelPerms(
        swr= 0.11,
        sor= 0.10,
        krw0= 0.2,
        kro0= 0.8,
        nw= 3.0,
        no= 2.0,
            )

kr_ow = SingleRelPerms(
        swr= 0.11,
        sor= 0.16,
        krw0= 0.4,
        kro0= 0.5,
        nw= 2.0,
        no= 3.0,
        )

#* Define reservoir
res = Reservoir(
    L = 300.,
    A = 200 * 50 * 0.3048^2,
    k = 100.0 * 9.869233e-16,
    phi = 0.2,
    )

#* Define fluids
fluid = Fluids(
            μw=1e-3,
            μo=5e-3,
)

#* Define grid
grid = Grid(
            res=res,
            Nx = 100,
            )

#* Define BC
bc = BoundaryConditions(
                        [:neumann, :dirichlet],
                        [426.5 * 0.3048 ^ 3 / 3600 / 24, 1000 * 6894.76]
                        )

tmax = 4000. * 3600 * 24 # days
dt = 1. * 3600 * 24

#* Initialize simulation and postprocessing
sim = Simulation(
                1000*ones(grid.Nx)* 6894.76,
                0.2*ones(grid.Nx),
                1*ones(grid.Nx),
                tmax,
                dt
)

post = FlowResults(sim, grid, tmax, dt)

while sim.t < tmax
    IMPES!(sim, res, kr_ww, fluid, grid, bc, post)
end

i = 1000
plot(grid.x, post.s[:,i], seriestype=:scatter)
plot(grid.x, post.p[:,i], seriestype=:scatter)


"""
    FactorisationWrapper

A small wrapper around a a factorisation `fact`, allowing to represent
multiplication by the inverse of `fact`.
"""
struct FactorisationWrapper{Fact}
    fact::Fact
end

Base.size(fact::FactorisationWrapper) = size(fact.fact)
Base.eltype(fact::FactorisationWrapper) = eltype(fact.fact)
mul!(dest, fact::FactorisationWrapper, u) = ldiv!(dest, fact.fact, u)

#TODO
#struct AdaptiveFilter{T<:Real} <: AbstractFilter{T}
#
#end

"""
    ConstantFilter

Represents the action of a modal filter on values in a nodal basis with fixed
strength.
"""
struct ConstantFilter{T<:Real,Nodal2Modal,Modal2Nodal,Tmp,FilterFunction} <: AbstractFilter{T}
    coefficients::Vector{T}
    nodal2modal::Nodal2Modal
    modal2nodal::Modal2Nodal
    tmp::Tmp
    filter::FilterFunction

    function ConstantFilter(coefficients::Vector{T}, nodal2modal::Nodal2Modal, modal2nodal::Modal2Nodal, tmp::Tmp, filter::FilterFunction) where {T<:Real,Nodal2Modal,Modal2Nodal,Tmp,FilterFunction}
        @argcheck length(coefficients) == size(tmp,1)

        new{T,Nodal2Modal,Modal2Nodal,Tmp,FilterFunction}(coefficients, nodal2modal, modal2nodal, tmp, filter)
    end
end

(filter::ConstantFilter)(u) = filter(u, filter.tmp)

function (filter::ConstantFilter)(u::AbstractVector, tmp::AbstractVector)
    @unpack coefficients, nodal2modal, modal2nodal = filter
    @boundscheck begin
        length(coefficients) == length(tmp)
    end

    mul!(tmp, nodal2modal, u)
    @inbounds tmp .*= coefficients
    mul!(u, modal2nodal, tmp)

    nothing
end

function (filter::ConstantFilter)(u::AbstractMatrix, tmp::AbstractVector)
    @boundscheck begin
        length(tmp) == size(u,1)
    end

    for j in Base.OneTo(size(u,2))
        v = view(u, :, j)
        @inbounds filter(v, tmp)
    end

    nothing
end

function (filter::ConstantFilter)(u::AbstractMatrix, tmp::AbstractMatrix)
    @unpack coefficients, nodal2modal, modal2nodal = filter
    @boundscheck begin
        size(tmp) == size(u)
    end

    mul!(tmp, nodal2modal, u)
    @inbounds tmp .*= coefficients
    mul!(u, modal2nodal, tmp)

    nothing
end

"""
    ExponentialFilter

Represents the exponential filter function `σ(η) = exp(-α*η^p)`.
"""
struct ExponentialFilter{T<:Real} <: AbstractFilterFunction
    α::T
    p::Int

    function ExponentialFilter(α::T, p::Int=2) where {T<:Real}
        α < 0 && warn("α should be nonnegative [α = $α].")
        p < 0 && warn("p should be nonnegative [p = $p].")
        new{T}(α, p)
    end
end

function ExponentialFilter(::Type{T}=Float64, p::Int=2) where {T<:Real}
    α = -log(eps(T))
    ExponentialFilter(α, p)
end

function ExponentialFilter(p::Int, ::Type{T}=Float64) where {T<:Real}
    α = -log(eps(T))
    ExponentialFilter(α, p)
end

function set_filter_coefficients!(coefficients::Vector, filter::ExponentialFilter)
    T = eltype(coefficients)
    N = length(coefficients) - 1
    @unpack α, p = filter

    @inbounds for i in Base.OneTo(length(coefficients))
        η = T(i-1) / N
        coefficients[i] = exp(-α*η^p)
    end

    nothing
end


function FaradayStep!(b, e, m, cb, nx, ny, nz)
    @turbo for k = 1:nz-1
        for j = 1:ny-1
            for i = 1:nx-1
                b[i, j, k, 1] =
                    b[i, j, k, 1] + cb[i, j, k, 1] * (e[i, j, k, 3] - e[i, j+1, k, 3] + e[i, j, k+1, 2] - e[i, j, k, 2]) - m[i, j, k, 1]
                b[i, j, k, 2] =
                    b[i, j, k, 2] + cb[i, j, k, 2] * (e[i, j, k, 1] - e[i, j, k+1, 1] + e[i+1, j, k, 3] - e[i, j, k, 3]) - m[i, j, k, 2]
                b[i, j, k, 3] =
                    b[i, j, k, 3] + cb[i, j, k, 3] * (e[i, j, k, 2] - e[i+1, j, k, 2] + e[i, j+1, k, 1] - e[i, j, k, 1]) - m[i, j, k, 3]
            end
        end
    end
    @turbo for k = 1:nz-1
        for j = 1:ny-1
            for i = nx:nx
                b[i, j, k, 1] =
                    b[i, j, k, 1] + cb[i, j, k, 1] * (e[i, j, k, 3] - e[i, j+1, k, 3] + e[i, j, k+1, 2] - e[i, j, k, 2]) - m[i, j, k, 1]
            end
        end
    end
    @turbo for k = 1:nz-1
        for j = ny:ny
            for i = 1:nx-1
                b[i, j, k, 2] =
                    b[i, j, k, 2] + cb[i, j, k, 2] * (e[i, j, k, 1] - e[i, j, k+1, 1] + e[i+1, j, k, 3] - e[i, j, k, 3]) - m[i, j, k, 2]
            end
        end
    end
    @turbo for k = nz:nz
        for j = 1:ny-1
            for i = 1:nx-1
                b[i, j, k, 3] =
                    b[i, j, k, 3] + cb[i, j, k, 3] * (e[i, j, k, 2] - e[i+1, j, k, 2] + e[i, j+1, k, 1] - e[i, j, k, 1]) - m[i, j, k, 3]
            end
        end
    end
end

            #    ex hx ey hy ez hz
function FaradayStep!(f, m, cb, nx, ny, nz)
    @turbo for k = 1:nz-1
        for j = 1:ny-1
            for i = 1:nx-1
                f[i, j, k, 2] =
                    f[i, j, k, 2] + cb[i, j, k, 1] * (f[i, j, k, 5] - f[i, j+1, k, 5] + f[i, j, k+1, 3] - f[i, j, k, 3]) - m[i, j, k, 1]
                f[i, j, k, 4] =
                    f[i, j, k, 4] + cb[i, j, k, 2] * (f[i, j, k, 1] - f[i, j, k+1, 1] + f[i+1, j, k, 5] - f[i, j, k, 5]) - m[i, j, k, 2]
                f[i, j, k, 6] =
                    f[i, j, k, 6] + cb[i, j, k, 3] * (f[i, j, k, 3] - f[i+1, j, k, 3] + f[i, j+1, k, 1] - f[i, j, k, 1]) - m[i, j, k, 3]
            end
        end
    end
    @turbo for k = 1:nz-1
        for j = 1:ny-1
            for i = nx:nx
                f[i, j, k, 4] =
                f[i, j, k, 4] + cb[i, j, k, 1] * (f[i, j, k, 3] - f[i, j+1, k, 3] + f[i, j, k+1, 2] - f[i, j, k, 2]) - m[i, j, k, 1]
            end
        end
    end
    @turbo for k = 1:nz-1
        for j = ny:ny
            for i = 1:nx-1
                f[i, j, k, 5] =
                f[i, j, k, 5] + cb[i, j, k, 2] * (f[i, j, k, 1] - f[i, j, k+1, 1] + f[i+1, j, k, 3] - f[i, j, k, 3]) - m[i, j, k, 2]
            end
        end
    end
    @turbo for k = nz:nz
        for j = 1:ny-1
            for i = 1:nx-1
                f[i, j, k, 6] =
                f[i, j, k, 6] + cb[i, j, k, 3] * (f[i, j, k, 2] - f[i+1, j, k, 2] + f[i, j+1, k, 1] - f[i, j, k, 1]) - m[i, j, k, 3]
            end
        end
    end
end


@testset "default solvers" begin
    prob = MVP.UnconstrainedProblems.examples["Powell"]
    f = objective(prob)
    g! = gradient(prob)
    h! = hessian(prob)
    @test summary(optimize(f, prob.initial_x)) == summary(NelderMead())
    @test summary(optimize(f, g!, prob.initial_x)) == summary(LBFGS())
    @test summary(optimize(f, g!, h!, prob.initial_x)) == summary(Newton())
end


# txt = readlines("day1.txt")

# txt = map(x->parse(Int64,x),txt)

# sums = Array{Int64, 1}(undef,size(txt,1)-2)
# for i in 2:(size(txt, 1)-1)
#     sums[i-1] = txt[i-1] + txt[i] + txt[i+1]
# end

# count = 0
# for i in 2:size(sums, 1)
#     if sums[i] > sums[i-1]
#         count += 1
#     end
# end

# print(count)

txt = map(x->parse(Int64,x),readlines("day1.txt"))
sums = [txt[i-1] + txt[i] + txt[i+1] for i in 2:(size(txt, 1)-1)]
sum([1 for i in 2:size(sums, 1) if sums[i] > sums[i-1]])


@testset "Cross mapping" begin

	@testset "Prediction lags" begin
	    x, y = rand(100), rand(100)
	    crossmap(x, y, η = 1)
	    crossmap(x, y, η = -1)
	    crossmap(x, y, η = 5)
	    crossmap(x, y, η = -5)
	end

	@testset "Embedding params" begin
	    x, y = rand(100), rand(100)
	    @test_throws DomainError crossmap(x, y, dim = -3)
	    @test_throws DomainError crossmap(x, y, dim = 100, τ = 2)
	    crossmap(x, y, dim = 3)
	    crossmap(x, y, dim = 10, τ = 2)
	end

	@testset "Exclusion radii" begin
	    x, y = rand(100), rand(100)
	    [crossmap(x, y, theiler_window = i) for i in rand(1:25, 10)]
	    #@test_throws DomainError crossmap(x, y, theiler_window = -1)
	end

	@testset "Replacements" begin
	    x, y = rand(100), rand(100)
	    crossmap(x, y, replace = false)
	    crossmap(x, y, replace = true)
	end

	@testset "Correspondence measures" begin
	    @testset "$i" for i in 1:100
	        x, y = rand(100), rand(100)
	        @test all(crossmap(x, y, correspondence_measure = StatsBase.rmsd) .>= 0)
	        @test all([-1 <= x <= 1 for x in crossmap(x, y, correspondence_measure = StatsBase.cor)])
	    end
	end

end


using SpecialFunctions
lgamma_(x) = logabsgamma(x)[1]

logsumexp(a,b) = (m = max(a,b); m == -Inf ? -Inf : log(exp(a-m) + exp(b-m)) + m)

function randp(p,k)
    s = 0.; for j = 1:k; s += p[j]; end
    u = rand()*s
    j = 1
    C = p[1]
    while u > C
        j += 1
        C += p[j]
    end
    @assert(j <= k)
    return j
end

function randlogp!(log_p,k)
    log_s = -Inf; for j = 1:k; log_s = logsumexp(log_s,log_p[j]); end
    p = log_p
    for j = 1:k; p[j] = exp(log_p[j]-log_s); end
    return randp(p,k)
end

function ordered_insert!(index,list,t)
    j = t
    while (j>0) && (list[j]>index)
        list[j+1] = list[j]
        j -= 1
    end
    list[j+1] = index
end

function ordered_remove!(index,list,t)
    for j = 1:t
        if list[j]>=index; list[j] = list[j+1]; end
    end
end

function ordered_next(list)
    j = 1
    while list[j]==j; j += 1; end
    return j
end

function restricted_Gibbs!(zsa,zsb,tia,tib,tja,tjb,cia,cib,cja,cjb,ni,nj,i,j,S,ns,x,b,H,active)
# NOTE: The sufficient statistics of tia and tja must be in sync with zsa.
# Also, note that the sufficient statistics of tib and tjb are not updated in this procedure.
    if !active; tia,tja = deepcopy(tia),deepcopy(tja); end
    log_p = 0.
    for ks = 1:ns
        k = S[ks]
        if k!=i && k!=j
            if zsa[k]==cia
                ni -= 1; Theta_remove!(tia,x[k])
            else
                nj -= 1; Theta_remove!(tja,x[k])
            end
            Li = log_marginal(x[k],tia,H) - log_marginal(tia,H)
            Lj = log_marginal(x[k],tja,H) - log_marginal(tja,H)
            Pi = exp(log(ni+b)+Li - logsumexp(log(ni+b)+Li,log(nj+b)+Lj))

            if active
                if rand()<Pi
                    zsb[k] = cib
                else
                    zsb[k] = cjb
                end
            end
            if zsb[k]==cib
                ni += 1
                Theta_adjoin!(tia,x[k])
                log_p += log(Pi)
            else
                nj += 1
                Theta_adjoin!(tja,x[k])
                log_p += log(1-Pi)
            end
        end
    end
    return log_p,ni,nj
end

function split_merge!(x,z,zs,S,theta,list,N,t,H,a,b,log_v,n_split,n_merge)
    n = length(x)

    # randomly choose a pair of indices
    i = round(Int,ceil(rand()*n))
    j = round(Int,ceil(rand()*(n-1))); if j>=i; j += 1; end
    ci0,cj0 = z[i],z[j]
    ti0,tj0 = theta[ci0],theta[cj0]

    # set S[1],...,S[ns] to the indices of the points in clusters ci0 and cj0
    ns = 0
    for k = 1:n
        if z[k]==ci0 || z[k]==cj0; ns += 1; S[ns] = k; end
    end

    # find available cluster IDs for merge and split parameters
    k = 1
    while list[k]==k; k += 1; end; cm = k
    while list[k]==k+1; k += 1; end; ci = k+1
    while list[k]==k+2; k += 1; end; cj = k+2
    tm,ti,tj = theta[cm],theta[ci],theta[cj]

    # merge state
    for ks = 1:ns; Theta_adjoin!(tm,x[S[ks]]); end # get the sufficient statistics

    # randomly choose the split launch state
    zs[i] = ci; Theta_adjoin!(ti,x[i]); ni = 1
    zs[j] = cj; Theta_adjoin!(tj,x[j]); nj = 1
    for ks = 1:ns  # start with a uniformly chosen split
        k = S[ks]
        if k!=i && k!=j
            if rand()<0.5; zs[k] = ci; Theta_adjoin!(ti,x[k]); ni += 1
            else;          zs[k] = cj; Theta_adjoin!(tj,x[k]); nj += 1
            end
        end
    end
    for rep = 1:n_split  # make several moves
        log_p,ni,nj = restricted_Gibbs!(zs,zs,ti,ti,tj,tj,ci,ci,cj,cj,ni,nj,i,j,S,ns,x,b,H,true)
    end

    # make proposal
    if ci0==cj0  # propose a split
        # make one final sweep and compute it's probability density
        log_prop_ab,ni,nj = restricted_Gibbs!(zs,zs,ti,ti,tj,tj,ci,ci,cj,cj,ni,nj,i,j,S,ns,x,b,H,true)

        # probability of going from merge state to original state
        log_prop_ba = 0.0  # log(1)

        # compute acceptance probability
        log_prior_b = log_v[t+1] + lgamma_(ni+b)+lgamma_(nj+b)-2*lgamma_(a)
        log_prior_a = log_v[t] + lgamma_(ns+b)-lgamma_(a)
        log_lik_ratio = log_marginal(ti,H) + log_marginal(tj,H) - log_marginal(ti0,H)
        p_accept = min(1.0, exp(log_prop_ba-log_prop_ab + log_prior_b-log_prior_a + log_lik_ratio))
        #println("split proposal: ",p_accept)

        # accept or reject
        if rand()<p_accept # accept split
            #  z, list, N, theta, t
            for ks = 1:ns; z[S[ks]] = zs[S[ks]]; end
            ordered_remove!(ci0,list,t)
            ordered_insert!(ci,list,t-1)
            ordered_insert!(cj,list,t)
            N[ci0],N[ci],N[cj] = 0,ni,nj
            t += 1
            Theta_clear!(ti0)
        else # reject split
            Theta_clear!(ti)
            Theta_clear!(tj)
        end
        Theta_clear!(tm)

    else  # propose a merge
        # probability of going to merge state
        log_prop_ab = 0.0  # log(1)

        # compute probability density of going from split launch state to original state
        log_prop_ba,ni,nj = restricted_Gibbs!(zs,z,ti,ti0,tj,tj0,ci,ci0,cj,cj0,ni,nj,i,j,S,ns,x,b,H,false)

        # compute acceptance probability
        log_prior_b = log_v[t-1] + lgamma_(ns+b)-lgamma_(a)
        log_prior_a = log_v[t] + lgamma_(ni+b)+lgamma_(nj+b)-2*lgamma_(a)
        log_lik_ratio = log_marginal(tm,H) - log_marginal(ti0,H) - log_marginal(tj0,H)
        p_accept = min(1.0, exp(log_prop_ba-log_prop_ab + log_prior_b-log_prior_a + log_lik_ratio))
        #println("merge proposal: ",p_accept)

        # accept or reject
        if rand()<p_accept # accept merge
            #  z, list, N, theta, t
            for ks = 1:ns; z[S[ks]] = cm; end
            ordered_remove!(ci0,list,t)
            ordered_remove!(cj0,list,t-1)
            ordered_insert!(cm,list,t-2)
            N[cm],N[ci0],N[cj0] = ns,0,0
            t -= 1
            Theta_clear!(ti0)
            Theta_clear!(tj0)
        else # reject merge
            Theta_clear!(tm)
        end
        Theta_clear!(ti)
        Theta_clear!(tj)
    end
    return t
end

function sampler(options,n_total,n_keep)
    x,n = options.x,options.n
    t_max = options.t_max
    a,b,log_v = options.a,options.b,options.log_v
    use_splitmerge,n_split,n_merge = options.use_splitmerge,options.n_split,options.n_merge
    use_hyperprior = options.use_hyperprior

    model_type = options.model_type
    alpha_random,alpha = options.alpha_random,options.alpha
    sigma_alpha = 0.1 # scale for MH proposals in alpha move

    @assert(n==length(x))
    keepers = zeros(Int,n_keep)
    keepers[:] = round.(Int,range(round(Int,n_total/n_keep),stop=n_total,length=n_keep))
    keep_index = 0

    t = 1  # number of clusters
    z = ones(Int,n)  # z[i] = the cluster ID associated with data point i
    list = zeros(Int,t_max+3); list[1] = 1  # list[1:t] = the list of active cluster IDs
                                            # list is maintained in increasing order for 1:t, and is 0 after that.
                                            # Further, maximum(list[1:t]) is always <= t_max+3
    c_next = 2  # an available cluster ID to be used next
    N = zeros(Int,t_max+3); N[1] = n  # N[c] = size of cluster c

    H = construct_hyperparameters(options)
    theta = [new_theta(H)::Theta for c = 1:t_max+3]  # theta[c] = parameters for cluster c

    for i = 1:n; Theta_adjoin!(theta[1],x[i]); end

    log_p = zeros(n+1)
    zs = ones(Int,n)  # temporary variable used for split-merge assignments
    S = zeros(Int,n)  # temporary variable used for split-merge indices

    log_Nb = log.((1:n) .+ b)

    # Record-keeping variables
    t_r = zeros(Int8,n_total); @assert(t_max < 2^7)
    N_r = zeros(Int16,t_max+3,n_total); @assert(n < 2^15)
    z_r = zeros(Int8,n,n_keep); @assert(t_max < 2^7)
    theta_r = Array{Theta}(undef,0,0)

    for iteration = 1:n_total
        #  -------------- Resample H --------------
        if use_hyperprior
            update_hyperparameters!(H,theta,list,t,x,z)
        end
        if model_type=="DPM" && alpha_random
            # Metropolis-Hastings move for DP concentration parameter (using p_alpha(a) = exp(-a) = Exp(a|1))
            aprop = alpha*exp(randn()*sigma_alpha)
            top = t*log(aprop) - lgamma_(aprop+n) + lgamma_(aprop) - aprop + log(aprop)
            bot = t*log(alpha) - lgamma_(alpha+n) + lgamma_(alpha) - alpha + log(alpha)
            if rand() < min(1.0,exp(top-bot))
                alpha = aprop
            end
            #for i=1:t_max+1; log_v[i] = i*log(alpha) - lgamma_(alpha+n) + lgamma_(alpha); end
            log_v = float(1:t_max+1)*log(alpha) .- lgamma_(alpha+n) .+ lgamma_(alpha)
        end

        # -------------- Resample z's --------------
        for i = 1:n
            # remove point i from it's cluster
            c = z[i]
            N[c] -= 1
            Theta_remove!(theta[c],x[i])
            if N[c]>0
                c_prop = c_next
            else
                c_prop = c
                # remove cluster {i}, keeping list in proper order
                ordered_remove!(c,list,t)
                t -= 1
            end

            # compute probabilities for resampling
            for j = 1:t; cc = list[j]
                log_p[j] = log_Nb[N[cc]] + log_marginal(x[i],theta[cc],H) - log_marginal(theta[cc],H)
            end
            log_p[t+1] = log_v[t+1]-log_v[t] + log(a) + log_marginal(x[i],theta[c_prop],H)

            # sample a new cluster for it
            j = randlogp!(log_p,t+1)

            # add point i to it's new cluster
            if j<=t
                c = list[j]
            else
                c = c_prop
                ordered_insert!(c,list,t)
                t += 1
                c_next = ordered_next(list)
                @assert(t<=t_max, "Sampled t has exceeded t_max. Increase t_max and retry.")
            end
            # update sufficient statistics
            Theta_adjoin!(theta[c],x[i])
            z[i] = c
            N[c] += 1
        end

        # -------------- Split/merge move --------------
        if use_splitmerge
            t = split_merge!(x,z,zs,S,theta,list,N,t,H,a,b,log_v,n_split,n_merge)
            c_next = ordered_next(list)
            @assert(t<=t_max, "Sampled t has exceeded t_max. Increase t_max and retry.")
        end

        # -------------- Record results --------------
        t_r[iteration] = t
        for j = 1:t
            N_r[list[j],iteration] = N[list[j]]
        end
        if iteration==keepers[keep_index+1]
            keep_index += 1
            for i = 1:n; z_r[i,keep_index] = z[i]; end
        end
    end

    return t_r,N_r,z_r,theta_r,keepers
end


function minute(line::String)
    return parse(Int, line[16:17])
end

function foo()
    lines = sort(readlines(open("data/4.txt", "r")))
    x = Dict{Int, Array{Int, 1}}()
    guard = -1
    for line in lines
        if occursin("Guard", line)
            start = findfirst(isequal('#'), line)
            stop = findnext(isequal(' '), line, start)
            guard = parse(Int, line[start+1:stop-1])
            if !(guard in keys(x))
                x[guard] = zeros(Int, 60)
            end
        elseif occursin("falls", line)
            start = minute(line) + 1
            x[guard][start:end] .+= 1
        elseif occursin("wakes", line)
            start = minute(line) + 1
            x[guard][start:end] .-= 1
        end
    end

    sleepiest = -1
    time_asleep = -1
    for (guard, sleep) in x
        sum_sleep = sum(sleep)
        if sum_sleep > time_asleep
            sleepiest = guard
            time_asleep = sum_sleep
        end
    end

    max_minute = argmax(x[sleepiest]) - 1
    return sleepiest * max_minute
end

println(foo())


@doc raw"""
    simple_estimator(model::AshkinTeller, T::Real, Jsigma, Jtau, K)

Returns the following observables as `Dict{String, Any}`

# Observables
- `"Energy"`
    - energy density
- `"Energy^2"`
    - square of energy density
- `"|Magnetization|"`
    - absolute value of magnetization density, ``|m| = \sqrt{ (\sum_i \sigma_i )^2 + (\sum_i \tau_i)^2 } / N``
- `"|Magnetization|^2"`
    - square of magnetization density
- `"|Magnetization|^4"`
    - quadruple of magnetization density
- `"Magnetization sigma"`
    - magnetization density (sigma spin)
- `"|Magnetization sigma|"`
- `"Magnetization sigma^2"`
- `"Magnetization sigma^4"`
- `"Magnetization tau"`
    - magnetization density (tau spin)
- `"|Magnetization tau|"`
- `"Magnetization tau^2"`
- `"Magnetization tau^4"`

"""
function simple_estimator(model::AshkinTeller, T::Real, Jsigma, Jtau, K, _=nothing)
    nsites = numsites(model)
    nbonds = numbonds(model)

    Ms = mean(model.spins, dims=2)
    E = 0.0
    @inbounds for b in bonds(model)
        s1, s2 = source(b), target(b)
        bt = bondtype(b)
        E += model.spins[1,s1] * model.spins[1,s2] * Jsigma[bt]
        E += model.spins[2,s1] * model.spins[2,s2] * Jtau[bt]
        E += model.spins[1,s1] * model.spins[2,s1] * model.spins[1,s2] * model.spins[2,s2] * K[bt]
    end
    E /= nsites

    res = Measurement()
    res["|Magnetization|^2"] = Ms[1]^2 + Ms[2]^2
    res["|Magnetization|^4"] = res["|Magnetization|^2"]^2
    res["|Magnetization|"] = sqrt(res["|Magnetization|^2"])

    res["Magnetization sigma"] = Ms[1]
    res["|Magnetization sigma|"] = abs(Ms[1])
    res["Magnetization sigma^2"] = Ms[1]^2
    res["Magnetization sigma^4"] = Ms[1]^4

    res["Magnetization tau"] = Ms[2]
    res["|Magnetization tau|"] = abs(Ms[2])
    res["Magnetization tau^2"] = Ms[2]^2
    res["Magnetization tau^4"] = Ms[2]^4

    res["Energy"] = E
    res["Energy^2"] = E^2
    return res
end

default_estimator(model::AshkinTeller, update) = simple_estimator


using PyPlot
# Set up a time grid for the plot:
ntime = 10000
time = collect(linspace(0,4pi,ntime))

# exp-sine-squared:
x1 = [0.24998970200450718,0.7528029456942368,0.9604772168743339,0.6433934213976025]
x2 = [0.4508506501130407,-0.44533765457178637,1.0810460905919905,1.1810193449724977]
x3 = [0.20659251082649457,-0.20736062231716337,1.659019056376522,0.21064077694469815]
x4 = [0.07284486969656653,-0.0729380659597069,2.431485661375288,0.047122405324731394]
x5 = [0.019204025483424515,-0.01921291003735305,3.2797651575680282,0.011419793258824862]

gamma = 1.0
model_expsin2_g1=zeros(ntime)
model_expsin2_g1 += x1[1]+x1[2].*exp(-gamma.^x1[3].*x1[4])
model_expsin2_g1 += (x2[1]+x2[2].*exp(-gamma.^x2[3].*x2[4])).*cos(2.*time)
model_expsin2_g1 += (x3[1]+x3[2].*exp(-gamma.^x3[3].*x3[4])).*cos(4.*time)
model_expsin2_g1 += (x4[1]+x4[2].*exp(-gamma.^x4[3].*x4[4])).*cos(6.*time)
model_expsin2_g1 += (x5[1]+x5[2].*exp(-gamma.^x5[3].*x5[4])).*cos(8.*time)
expsin2_g1 = exp.(-gamma.*sin(time).^2)

gamma = 3.0
model_expsin2_g3=zeros(ntime)
model_expsin2_g3 += x1[1]+x1[2].*exp(-gamma.^x1[3].*x1[4])
model_expsin2_g3 += (x2[1]+x2[2].*exp(-gamma.^x2[3].*x2[4])).*cos(2.*time)
model_expsin2_g3 += (x3[1]+x3[2].*exp(-gamma.^x3[3].*x3[4])).*cos(4.*time)
model_expsin2_g3 += (x4[1]+x4[2].*exp(-gamma.^x4[3].*x4[4])).*cos(6.*time)
model_expsin2_g3 += (x5[1]+x5[2].*exp(-gamma.^x5[3].*x5[4])).*cos(8.*time)
expsin2_g3 = exp.(-gamma.*sin(time).^2)

# Gaussian (squared-exponential):
#abest=[2.7518074963750734,1.5312499407621147,-0.34148888178746056,-0.3546256365636662,4.753504292338284,2.042571639590954,-1.3968344706386124,1.7967431129812566,-1.7641684073510928]
#abest=[2.670143,1.500258, 0.346731,-0.458705,4.125588, 1.390912, -1.211438,1.744602, 1.798436]
abest = [-0.519196,0.285478,1.383928,1.83396, 63.119512,0.0,1.479579,0.0, -32.173615,0.0,1.593691,0.0, -29.425156,0.0,1.388929,0.0]

gaussian=exp.(-time.^2/2.)
nexp = 4

model_gaussian = zeros(ntime)
for k=1:nexp
  model_gaussian += abest[1+(k-1)*4].*exp.(-abest[3+(k-1)*4].*time).*cos(abest[4+(k-1)*4].*time)
  model_gaussian += abest[2+(k-1)*4].*exp.(-abest[3+(k-1)*4].*time).*sin(abest[4+(k-1)*4].*time)
end

# Matern 5/2:
matern52=(1.0+sqrt(5.).*time+5.*time.^2./3.).*exp.(-sqrt(5.).*time)
nexp = 4
atrial = [-147.6965573,1.0097610,85.8875034,1.0621734,-4.6540017,0.8353745,67.4630744,0.9160394]
model_52 = zeros(ntime)
for k=1:nexp
  model_52 += atrial[1+(k-1)*2].*exp.(-atrial[2+(k-1)*2].*sqrt(5.)*time)
end

colors = [ "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b",
               "#e377c2", "#7f7f7f", "#bcbd22", "#17becf", ]
linewidth=3
fontsize=10

#fix,axes = subplots(2,2,figsize=[10,8])
fix,axes = subplots(1,4,figsize=[20,4])
#fix,ax = subplots()

ax = axes[1]
ax[:set_title]("(a). Simple, exact kernels",fontsize=2*fontsize)
ax[:plot](time,1.+0.*time,alpha=0.5,label=L"$1$",linewidth=linewidth,color=colors[1])
ax[:plot](time,exp.(-time/3),alpha=0.5,label=L"$e^{-t/3}$",linewidth=linewidth,color=colors[2])
ax[:plot](time,cos(time),alpha=1.0,alpha=0.5,label=L"$\cos{(\tau)}$",color=colors[3],linewidth=linewidth)
ax[:plot](time,exp.(-time./3).*cos(time),alpha=1.0,alpha=0.5,label=L"$e^{-t/3}\cos{(\tau)}$",color=colors[4],linewidth=linewidth)
ax[:legend](loc="lower center",fontsize=fontsize)
ax[:set_xlabel](L"$\tau/\tau_0$",fontsize=2*fontsize)
ax[:set_ylabel]("Dimensional ACF",fontsize=2*fontsize)
ax[:axis]([0,4pi,-1.1,1.1])

ax = axes[3]
ax[:set_title]("(b). Gaussian",fontsize=2*fontsize)
ax[:plot](time,gaussian,alpha=0.25,label=L"$G(z)$, exact",linewidth=linewidth,color=colors[1])
ax[:plot](time,model_gaussian,alpha=1.0,ls="dashed",alpha=0.5,label=L"$G(z)$, approx",color=colors[1],linewidth=linewidth)
print("ExpSq error: ",std(gaussian-model_gaussian)," ",maximum(abs(gaussian-model_gaussian)))
ax[:legend](loc="upper right",fontsize=fontsize)
ax[:set_xlabel](L"$\tau/\tau_0$",fontsize=2*fontsize)
#ax[:set_ylabel]("Dimensional ACF")
ax[:axis]([0,pi,0,1.1])

ax = axes[2]
ax[:plot](time,expsin2_g1,alpha=0.25,label=L"$E(w,1)$, exact",linewidth=linewidth,color=colors[1])
ax[:plot](time,model_expsin2_g1,alpha=1.0,label=L"$E(w,1)$, approx",ls="dashed",alpha=0.5,color=colors[1],linewidth=linewidth)
ax[:plot](time,expsin2_g3,alpha=0.25,label=L"$E(w,3)$, exact",linewidth=linewidth,color=colors[2])
ax[:plot](time,model_expsin2_g3,alpha=1.0,label=L"$E(w,3)$, approx",ls="dashed",alpha=0.5,color=colors[2],linewidth=linewidth)
ax[:set_title]("(c). Exp-sine-squared",fontsize=2*fontsize)
ax[:legend](loc="upper center",fontsize=fontsize)
ax[:set_xlabel](L"$\tau/\tau_0$",fontsize=2*fontsize)
#ax[:set_ylabel]("Dimensional ACF")
ax[:axis]([0,pi,0,1.1])

ax = axes[4]
matern32 = (1.+sqrt(3.).*time).*exp.(-sqrt(3.).*time)
b = 100.
model_32 = (1-b)*exp.(-sqrt(3.).*time)+b*exp.(-(b-1)/b.*sqrt(3.).*time)
ax[:plot](time,matern32,alpha=0.25,label=L"$K_{3/2}(\tau)$, exact",linewidth=linewidth,color=colors[1])
ax[:plot](time,model_32,alpha=1.0, label=L"$K_{3/2}(\tau)$, approx",ls="dashed",alpha=0.5,color=colors[1],linewidth=linewidth)
ax[:plot](time,matern52,alpha=0.25,label=L"$K_{5/2}(\tau)$, exact",linewidth=linewidth,color=colors[2])
ax[:plot](time,model_52,alpha=1.0, label=L"$K_{5/2}(\tau)$, approx",ls="dashed",alpha=0.5,color=colors[2],linewidth=linewidth)
ax[:set_title]("(d). Matern",fontsize=2*fontsize)
ax[:legend](loc="upper right",fontsize=fontsize)
ax[:set_xlabel](L"$\tau/\tau_0$",fontsize=2*fontsize)
#ax[:set_ylabel]("Dimensional ACF",fontsize=2*fontsize)
ax[:axis]([0,pi,0,1.1])

tight_layout()
savefig("kernel_approx.pdf", bbox_inches="tight")


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# ╔═╡ c7356460-055b-11eb-0d69-dd42369dc913
try using AddPackage catch; using Pkg; Pkg.add("AddPackage"); using AddPackage end

# ╔═╡ df4bc620-055b-11eb-39c7-e952028567da
begin
	using Base.Iterators: repeated
	@add using Colors
	@add using CUDA
	@add using Images
	@add using MLDatasets
	@add using Parameters
	@add using Flux
	using Flux.Data: DataLoader
	using Flux: @epochs, onehotbatch, onecold, logitcrossentropy, throttle
	using Statistics
end

# ╔═╡ d7a1f38e-055b-11eb-31ef-55de28cda23d
md"# Classifier for MNIST"

# ╔═╡ d8d88ce0-055d-11eb-32e8-31586d97a59c
md"## Training arguments"

# ╔═╡ 07f1cf70-055c-11eb-05ab-e77200ea8ea7
@with_kw mutable struct Args
	α::Float64 = 3e-4      # learning rate
	batchsize::Int = 1024  # batch size
	epochs::Int = 20       # number of epochs
	device::Function = gpu # set as gpu if available
	throttle::Int = 1      # throttle print every X seconds
end

# ╔═╡ f9bad300-055d-11eb-1e4d-f362cf07c6b5
md"## Data processing"

# ╔═╡ 37042240-055c-11eb-3eba-0da0d64627aa
function getdata(args)
	# load dataset
	xtrain, ytrain = MLDatasets.MNIST.traindata(Float32)
	xtest, ytest = MLDatasets.MNIST.testdata(Float32)

	# reshape data to flatten each image into a vector
	xtrain = Flux.flatten(xtrain)
	xtest = Flux.flatten(xtest)

	# one-hot encode the labels
	ytrain, ytest = onehotbatch(ytrain, 0:9), onehotbatch(ytest, 0:9)

	# split into batches
	traindata = DataLoader(xtrain, ytrain, batchsize=args.batchsize, shuffle=true)
	testdata = DataLoader(xtest, ytest, batchsize=args.batchsize)

	return traindata, testdata
end

# ╔═╡ ffa3bbb0-055d-11eb-2652-fd3b56dcf474
md"## Model"

# ╔═╡ 87cbe9b0-055c-11eb-1c34-9d9bedbe4ad5
function buildmodel(; imgsize=(28,28,1), nclasses=10)
	return Chain(
			Dense(prod(imgsize), 32, relu),
			Dense(32, nclasses))
end

# ╔═╡ 054402a2-055e-11eb-273e-1340d52bbed8
md"""
## Cost function
$$\begin{gather}
\mathcal{L}(\hat{y}, y) = -\frac{1}{n}\sum_{i=1}^n y \left(\hat{y} - \log\left(\sum e^{\hat{y}}\right)\right)\\
\mathcal{J}(\mathbf{\hat{y}}, \mathbf{y}) = \frac{1}{m}\sum \mathcal{L}(\hat{y}, y)
\end{gather}$$
"""

# ╔═╡ b391a990-055c-11eb-13b8-b3f9ad78cea7
cost(dataloader, model) = mean(logitcrossentropy(model(x), y) for (x,y) in dataloader)

# ╔═╡ ea7f0630-055e-11eb-0717-bde12a519dfc
md"## Accuracy"

# ╔═╡ db2160e0-055c-11eb-0c66-6b3cc667e18e
function accuracy(dataloader, model)
	acc = 0
	for (x,y) in dataloader
		acc += sum(onecold(cpu(model(x))) .== onecold(cpu(y))) / size(x,2)
	end
	return acc/length(dataloader)
end

# ╔═╡ ef2b0210-055e-11eb-244b-d18b38d83a71
md"## Model training"

# ╔═╡ 01bccd20-055d-11eb-07a0-73fa53c45c24
function train(; kwargs...)
	# initialize model parameters
	args = Args(; kwargs...)

	# load data
	traindata, testdata = getdata(args)

	# construct model
	m = buildmodel()
	traindata = args.device.(traindata)
	testdata = args.device.(testdata)
	m = args.device(m)
	loss(x,y) = logitcrossentropy(m(x), y)

	# callback, optimizer, and training
	callback = throttle(()->@show(cost(traindata, m)), args.throttle)
	opt = ADAM(args.α)

	@epochs args.epochs Flux.train!(loss, params(m), traindata, opt, cb=callback)

	@show accuracy(traindata, m)
	@show accuracy(testdata, m)

	return m, traindata, testdata
end

# ╔═╡ 6092c160-055d-11eb-3364-0df28086629e
m, traindata, testdata = train();

# ╔═╡ f89b15b0-055e-11eb-3a07-091d5e96e340
begin
	Markdown.parse(string("\$\$\\operatorname{Loss}_\\text{train} = ", 100round(accuracy(traindata, m), digits=5), "\\%\$\$", "\$\$\\operatorname{Loss}_\\text{test} = ", 100round(accuracy(testdata, m), digits=5), "\\%\$\$"))
end

# ╔═╡ Cell order:
# ╟─d7a1f38e-055b-11eb-31ef-55de28cda23d
# ╠═c7356460-055b-11eb-0d69-dd42369dc913
# ╠═df4bc620-055b-11eb-39c7-e952028567da
# ╟─d8d88ce0-055d-11eb-32e8-31586d97a59c
# ╠═07f1cf70-055c-11eb-05ab-e77200ea8ea7
# ╟─f9bad300-055d-11eb-1e4d-f362cf07c6b5
# ╠═37042240-055c-11eb-3eba-0da0d64627aa
# ╟─ffa3bbb0-055d-11eb-2652-fd3b56dcf474
# ╠═87cbe9b0-055c-11eb-1c34-9d9bedbe4ad5
# ╟─054402a2-055e-11eb-273e-1340d52bbed8
# ╠═b391a990-055c-11eb-13b8-b3f9ad78cea7
# ╟─ea7f0630-055e-11eb-0717-bde12a519dfc
# ╠═db2160e0-055c-11eb-0c66-6b3cc667e18e
# ╟─ef2b0210-055e-11eb-244b-d18b38d83a71
# ╠═01bccd20-055d-11eb-07a0-73fa53c45c24
# ╠═6092c160-055d-11eb-3364-0df28086629e
# ╟─f89b15b0-055e-11eb-3a07-091d5e96e340


abstract type Preconditioner end

struct IdentityPreconditioner <: Preconditioner
end

struct JacobiPreconditioner{T <: Number} <: Preconditioner
    diagonal_list::Array{Array{T, 1}, 1}
end

function apply!(precond::IdentityPreconditioner, baselines) where {T <: Number}
    # Do nothing
end

function apply!(precond::JacobiPreconditioner{T}, baselines) where {T <: Number}
    @assert length(baselines) == length(precond.diagonal_list)

    for idx in eachindex(baselines)
        baselines[idx] .*= precond.diagonal_list[idx]
    end
end

function jacobi_preconditioner(baseline_length_list, T)
    # TODO: once we specify a mask or flags, we need to take them into account here!
    diagonals = [Array{T}(undef, length(x)) for x in baseline_length_list]

    for cur_diag_idx in eachindex(diagonals)
        diagonals[cur_diag_idx] .= 1 ./ baseline_length_list[cur_diag_idx]
    end

    JacobiPreconditioner{T}(diagonals)
end


using VoronoiCells
using Base.Test

# Corners of the unit square
x = [1.0 ; 0.0 ; 0.0 ; 1.0]
y = [1.0 ; 1.0 ; 0.0 ; 0.0]

@test_approx_eq density(x, y) sqrt(2)/2


"""
Defines a circular aperture with a diameter given in mm.
"""
struct ClearDiameter{T<:Real}
    clear_diameter::T
    ClearDiameter(ca) =
        ca < zero(typeof(ca)) ? error("ClearDiameter < 0") : new{typeof(ca)}(ca)
end

function is_vignetted(ray, aper::Nothing) where {T}
    return (false)
end

function is_vignetted(ray, aper::ClearDiameter{T}) where {T}
    ray_height = 2 * sqrt(ray.x^2 + ray.y^2)
    return (ray_height > aper.clear_diameter)
end


println((sum(1:100) ^ 2) - sum(map(x -> x ^ 2, 1:100)))


# Tests automatic differentiation on s2.jl:
include("../src/transit_structure.jl")
include("../src/s2.jl")
if VERSION >= v"0.7"
  using Test
else
  using Base.Test
end

using ForwardDiff
using DiffResults

# Randomizer seed
using Random
Random.seed!(42)

function s2_grad_func(r::T,b::T) where {T <: Real}
  # Computes the derivative of s_n(r,b) with respect to r, b.
  # Create a vector for use with ForwardDiff
  x=[r,b]
  # Now, define a wrapper of s2 for use with ForwardDiff:
  function diff_s2(x::Array{T,1}) where {T <: Real}
  # x should be a two-element vector with values [r,b]
  rc,bc = x
  s2c= s2(rc,bc)
  return s2c
  end

  # Set up a type to store s_n and it's Jacobian with respect to x:
  out = DiffResults.GradientResult(x)
  # Compute the Jacobian (and value):
  out = ForwardDiff.gradient!(out,diff_s2,x)
  # Place the value in the s_2 vector:
  s_2 = DiffResults.value(out)
  # And, place the Jacobian in an array:
  s2_gradient= DiffResults.gradient(out)
return s_2,s2_gradient
end

function s2_grad_num(r::T,b::T) where {T <: Real}
# Now, carry out finite-difference:
dq = big(1e-18)

# Allocate an array for s_2:
# Make BigFloat versions of r & b:
r_big = big(r); b_big = big(b)
# Compute s_n to BigFloat precision:
s_2_big= s2(r_big,b_big)
# Now, compute finite differences:
s2_grad_big= zeros(BigFloat,2)
s2_plus= s2(r_big+dq,b_big)
s2_minus= s2(r_big-dq,b_big)
s2_grad_big[1] = (s2_plus-s2_minus)*.5/dq
s2_plus=s2(r_big,b_big+dq)
s2_minus=s2(r_big,b_big-dq)
s2_grad_big[2] = (s2_plus-s2_minus)*.5/dq
return convert(Float64,s_2_big),convert(Array{Float64,1},s2_grad_big)
end

# In the following lines I test the derivatives for the few special
# cases needed for computing s_2.  Note that in the paper I report
# \Lambda, while this gets multiplied by -\pi to produce s_2.

function runtest_s2(r,b,test_name)
@testset "$test_name" begin
s_2,s2_gradient= s2_grad_func(r,b)
s2_big,s2_grad_numeric = s2_grad_num(r,b)
s2_grad_ana = zeros(2)
diff1 = s2_grad_numeric-s2_gradient
s_2,Eofk,Em1mKdm=s2!(r,b,s2_grad_ana)
diff2 = s2_grad_ana-s2_gradient
println("b : ",b," r: ",r," flux diff: ",s_2-s2_big," grad diff(num-auto): ",diff1," diff(ana-auto): ",diff2 )
#println("gradient: ",s2_gradient," expected: ",-[2,-2./3.])
@test isapprox(s2_grad_numeric[1],s2_gradient[1],atol=1e-8)
@test isapprox(s2_grad_numeric[2],s2_gradient[2],atol=1e-8)
@test isapprox(s2_grad_ana[1],s2_gradient[1],atol=1e-8)
@test isapprox(s2_grad_ana[2],s2_gradient[2],atol=1e-8)
end
return
end

# Try r=b=1/2 special case:
runtest_s2(0.5,0.5,"r=0.5;b=0.5")

# Try b=0 special case:
runtest_s2(0.3,0.0,"r=0.3;b=0.0")

# Try r+b=1 special case:
runtest_s2(0.2,0.8,"r=0.2;b=0.8")

# Try r=b <1/2 special case:
runtest_s2(0.3,0.3,"r=0.3;b=0.3")

# Try r=b > 1/2 special case:
runtest_s2(3.0,3.0,"r=3.0;b=3.0")

# Now, try a random case with b+r < 1:
b=r=2.0
while b+r > 1
  global r = 2rand(); global b= 2rand()
end
runtest_s2(r,b,"r+b < 1")

# Now, try a random case with b+r > 1:
b=r=0.
while b+r < 1
  global r = 2rand(); global b= 2rand()
end
runtest_s2(r,b,"r+b > 1")

@testset "s2 grid" begin
# Now, try out hard cases - ingress/egress of small planets:
r0 = [0.01,0.01,1.0,10.0,100.0]
nb = 200
#l_max = 20
l_max = 10
n_max = l_max^2+2*l_max

using PyPlot
fig,axes = subplots(1,length(r0))
get_cmap("plasma")
epsilon = 1e-12; delta = 1e-3
for i=1:length(r0)
  global r=r0[i]
  if r < 1.0
    bgrid = [linearspace(1e-15,epsilon,nb); linearspace(epsilon,delta,nb); linearspace(delta,r-delta,nb);
     -logarithmspace(log10(delta),log10(epsilon),nb) .+ r; linearspace(r-epsilon,r+epsilon,nb); logarithmspace(log10(epsilon),log10(delta),nb) .+ r;
     linearspace(r+delta,1-r-delta,nb); -logarithmspace(log10(delta),log10(epsilon),nb) .+ (1-r); linearspace(1-r-epsilon,1-r+epsilon,nb);
     logarithmspace(log10(epsilon),log10(delta),nb) .+ (1-r); linearspace(1-r+delta,1+r-delta,nb); - logarithmspace(log10(delta),log10(epsilon),nb) .+ (1+r);
     linearspace(1+r-epsilon,1+r-1e-15,nb)]
  else
    bgrid = [linearspace(r-1+1e-15,r-1+epsilon,nb); logarithmspace(log10(epsilon),log10(delta),nb) .+ (r-1); linearspace(r-1+delta,r-delta,nb);
     -logarithmspace(log10(delta),log10(epsilon),nb) .+ r; linearspace(r-epsilon,r+epsilon,nb); logarithmspace(log10(epsilon),log10(delta),nb) .+ r;
     linearspace(r+delta,r+1-delta,nb); -logarithmspace(log10(delta),log10(epsilon),nb) .+ (r+1); linearspace(r+1-epsilon,r+1-1e-15,nb)]
  end
  nbgrid = length(bgrid)
  s2_jac_grid = zeros(nbgrid,2)
  s2_grid = zeros(nbgrid)
  s2_grid_big = zeros(nbgrid)
  s2_jac_grid_num = zeros(nbgrid,2)
  s2_grad_ana = zeros(2)
  for j=1:nbgrid
    s_2,Eofk,Em1mKdm= s2!(r,bgrid[j],s2_grad_ana)
    s2_grid[j]=s_2
    s2_jac_grid[j,:]=s2_grad_ana
    s2_big,s2_grad_numeric = s2_grad_num(r,bgrid[j])
    s2_grid_big[j]=s2_big
    s2_jac_grid_num[j,:]= s2_grad_numeric
    @test isapprox(s_2,s2_big,atol=1e-8)
    @test isapprox(s2_grad_ana[1],s2_grad_numeric[1],atol=1e-8)
    @test isapprox(s2_grad_ana[2],s2_grad_numeric[2],atol=1e-8)
  end
# Now, make plots:
  ax = axes[i]
  ax[:semilogy](bgrid,abs.(s2_grid[:,1]-s2_grid_big[:,1]) .+1e-18,lw=1,label="s 2 ")
  ax[:semilogy](bgrid,abs.(s2_jac_grid[:,1]-s2_jac_grid_num[:,1]) .+1e-18,lw=1,label="ds2/dr")
  ax[:semilogy](bgrid,abs.(s2_jac_grid[:,2]-s2_jac_grid_num[:,2]) .+1e-18,lw=1,label="ds2/db")
  ax[:legend](loc="upper right",fontsize=6)
  ax[:set_xlabel]("b values")
  ax[:set_ylabel]("Derivative Error")
end

savefig("test_s2_gradient.pdf", bbox_inches="tight")

end


using ExactWrightFisher

@testset "test exact L2 Gamma distance functions" begin
    # @test_warn AssertionError DualOptimalFiltering.log_int_prod_2_Gammas(0.5, 0.5, 0.5, 1)
    @test DualOptimalFiltering.log_int_prod_2_Gammas(1, 1, 1, 1) == -log(2)
    @test DualOptimalFiltering.log_int_prod_2_Gammas(2, 0.5, 1, 0.5) == log(0.5^2*0.5 )
    # @test DualOptimalFiltering.log_int_prod_2_Gammas(0.1, 1.2, 1.1, 0.3) ≈ Float64(log(DualOptimalFiltering.prod_2_gammas_arb(0.1, 1.2, 1.1, 0.3) ))

    @test DualOptimalFiltering.int_prod_2_Gammas(1, 1, 1, 1) == 0.5
    @test DualOptimalFiltering.int_prod_2_Gammas(2, 0.5, 1, 0.5) ≈ 0.5^2*0.5

    @test_nowarn DualOptimalFiltering.log_L2_dist_Gamma_mixtures(log.([0.2,0.3,0.5]), [1,2,3], [1,2,4], log.([0.6,0.4]), [1,2], [4,2])
    @test !isfinite(DualOptimalFiltering.log_L2_dist_Gamma_mixtures(log.([0.2,0.3,0.5]), [1,2,3], [1,2,4], log.([0.2,0.3,0.5]), [1,2,3], [1,2,4]))

    # @test DualOptimalFiltering.log_L2_dist_Gamma_mixtures(log.([0.2,0.3,0.5]), [1,2,3], [1,2,4], log.([0.6,0.4]), [1,2], [4,2]) ≈ Float64(log(DualOptimalFiltering.L2_dist_Gamma_mixtures_arb([0.2,0.3,0.5], [1,2,3], [1,2,4], [0.6,0.4], [1,2], [4,2])))

    Random.seed!(0)
    δ = 3.
    γ = 3.
    σ = 0.5
    λ = 1.
    θ = γ/σ^2
    X = DualOptimalFiltering.generate_CIR_trajectory(range(0, stop = 2, length = 20), δ, γ, σ, λ);
    Y = map(λ -> rand(Poisson(λ),10), X);
    data = Dict(zip(range(0, stop = 2, length = 20), Y))
    Λ_of_t, wms_of_t, θ_of_t = DualOptimalFiltering.filter_CIR(δ, γ, σ, λ, data);
    Λ_of_t1, wms_of_t1, θ_of_t1 = DualOptimalFiltering.filter_CIR_keep_fixed_number(δ, γ, σ, λ, data, 1);

    time_grid = data |> keys |> collect |> sort
    logw1 = log.(wms_of_t[time_grid[20]])
    αlist_1, βlist_1 = DualOptimalFiltering.create_gamma_mixture_parameters(δ, θ_of_t[time_grid[20]], Λ_of_t[time_grid[20]][isfinite.(logw1)])

    logw2 = log.(wms_of_t1[time_grid[20]])
    αlist_2, βlist_2 = DualOptimalFiltering.create_gamma_mixture_parameters(δ, θ_of_t1[time_grid[20]], Λ_of_t1[time_grid[20]])

    @test_nowarn DualOptimalFiltering.log_L2_dist_Gamma_mixtures(logw1[isfinite.(logw1)], αlist_1, βlist_1, logw2, αlist_2, βlist_2)

    cc = DualOptimalFiltering.create_gamma_mixture_pdf(δ, θ_of_t[time_grid[20]], Λ_of_t[time_grid[20]][isfinite.(logw1)], wms_of_t[time_grid[20]][isfinite.(logw1)])

    cc1 = DualOptimalFiltering.create_gamma_mixture_pdf(δ, θ_of_t1[time_grid[20]], Λ_of_t1[time_grid[20]], wms_of_t1[time_grid[20]])

    # res = DualOptimalFiltering.L2_dist_1D(cc, cc1)
    # @test exp(DualOptimalFiltering.log_L2_dist_Gamma_mixtures(logw1[isfinite.(logw1)], αlist_1, βlist_1, logw2, αlist_2, βlist_2)) ≈ res[1] atol=res[2]
end;

@testset "test exact L2 distance functions for Dirichlet mixtures" begin

    @test DualOptimalFiltering.lmvbeta([1,1]) == 0.0
    @test DualOptimalFiltering.log_int_prod_2_Dir([1,1], [1,1]) == 0.0
    @test DualOptimalFiltering.log_int_prod_2_Dir([1,1,1], [1,1,1]) ≈ 0.6931471805599453

    @test_nowarn DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.6,0.4]), [[1,2,1],[4,2,3]])
    @test_nowarn DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]])

    Random.seed!(0);
    α_vec = [1.2, 1.4, 1.3]
    # α = α_vec
    K = length(α_vec)
    Pop_size_WF3 = 10
    Nparts = 100
    time_grid_WF3 = range(0, stop = 1, length = 5)
    wfchain = Wright_Fisher_K_dim_exact_trajectory([0.2, 0.4, 0.4], time_grid_WF3, α_vec)
    wfobs_WF3 = [rand(Multinomial(Pop_size_WF3, wfchain[:,k])) for k in 1:size(wfchain,2)] |> l -> hcat(l...)
    data_WF3 = Dict(zip(time_grid_WF3 , [wfobs_WF3[:,t]' |> collect for t in 1:size(wfobs_WF3,2)]))

    # log_ν_dict_arb, log_Cmmi_dict_arb, log_binomial_coeff_dict_arb = DualOptimalFiltering.precompute_log_terms_arb(data_WF3, sum(α_vec); digits_after_comma_for_time_precision = 4)
    #
    #
    # Λ_of_t_arb, wms_of_t_arb = DualOptimalFiltering.filter_WF_precomputed(α_vec, data_WF3, log_ν_dict_arb, log_Cmmi_dict_arb, log_binomial_coeff_dict_arb)

    Λ_of_t, wms_of_t = DualOptimalFiltering.filter_WF_adaptive_precomputation_keep_fixed_number(α_vec, data_WF3, 30)

    times = keys(Λ_of_t) |> collect |> sort

    # res = DualOptimalFiltering.compute_L2_distance_between_two_Dirichlet_mixtures(α_vec, Λ_of_t_arb[times[5]], wms_of_t_arb[times[5]], Λ_of_t[times[5]], wms_of_t[times[5]])
    @test_nowarn DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.(wms_of_t[times[5]]), DualOptimalFiltering.create_dirichlet_mixture(α_vec, Λ_of_t[times[5]]), log.(wms_of_t[times[4]]),  DualOptimalFiltering.create_dirichlet_mixture(α_vec, Λ_of_t[times[4]]))

    # @test_nowarn DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(α_vec, Λ_of_t[times[2]], wms_of_t[times[2]], Λ_of_t[times[3]], wms_of_t[times[3]])

    # @test exp(DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.(wms_of_t_arb[times[3]]), DualOptimalFiltering.create_dirichlet_mixture(α_vec, Λ_of_t_arb[times[3]]), log.(wms_of_t[times[3]]),  DualOptimalFiltering.create_dirichlet_mixture(α_vec, Λ_of_t[times[3]]))) ≈ res[1] atol=res[2]

end;


struct AliasTable{S} <: Sampleable{Univariate,Discrete}
    accept::Vector{Float64}
    alias::Vector{Int}
    isampler::S
end
ncategories(s::AliasTable) = length(s.accept)

function AliasTable(probs::AbstractVector{T}) where T<:Real
    n = length(probs)
    n > 0 || throw(ArgumentError("The input probability vector is empty."))
    accp = Vector{Float64}(undef, n)
    alias = Vector{Int}(undef, n)
    StatsBase.make_alias_table!(probs, 1.0, accp, alias)
    AliasTable(accp, alias, Random.RangeGenerator(1:n))
end

function rand(rng::AbstractRNG, s::AliasTable)
    i = rand(rng, s.isampler) % Int
    u = rand(rng)
    @inbounds r = u < s.accept[i] ? i : s.alias[i]
    r
end

show(io::IO, s::AliasTable) = @printf(io, "AliasTable with %d entries", ncategories(s))


function check_manifold(M)
    # check random points, vectors
    for i in 1:10
        x = randomMPoint(M)
        @test is_manifold_point(M, x)

        ξ = randomTVector(M, x)
        @test is_tangent_vector(M, x, ξ, atol = 1e-10)
        @test isapprox(norm(M, x, ξ), 1)

        # check projection
        q = rand(representation_size(M)...)

        η = zero_tangent_vector(M, x)
        project!(M, η, x, q)
        @test is_tangent_vector(M, x, η; atol = 1e-10)

        # check retraction
        y = retract(M, x, ξ)
        @test is_manifold_point(M, y)
    end

    return nothing
end


using Latexify
import MacroTools
using LaTeXStrings

const transition_idx_script = Ref(:^)

"""
    transition_superscript(::Bool)

Specify whether the indices in a [`Transition`](@ref) operator should be
printed as superscript. Default is `true`. If set to `false`, the indices
corresponding to the levels are printed as subscript.
"""
function transition_superscript(x::Bool)
    if x
        transition_idx_script[] = :^
    else
        transition_idx_script[] = :_
    end
    return x
end

function _postwalk_func(x)
    if x==:𝟙
        return "\\mathbb{1}"
    elseif x==:im
        return :i
    elseif MacroTools.@capture(x, dagger(arg_))
        s = "$(arg)^\\dagger"
        return s
    elseif MacroTools.@capture(x, Transition(arg_,i_,j_))
        s = "{$(arg)}$(transition_idx_script[]){{$(i)$(j)}}"
        return s
    else
        return x
    end
end

function _postwalk_average(x)
    if MacroTools.@capture(x, AVG(arg_))
        arg = MacroTools.postwalk(_postwalk_func, arg)
        # TODO: clean up; tricky because of nested string conversion of eg average(dagger(a))
        s = string(arg)
        s = replace(s, "\"" => "")
        s = replace(s, "\\\\" => "\\")
        s = replace(s, "*" => "")
        s = "\\langle " * s * "\\rangle "
        return s
    end
    return x
end

@latexrecipe function f(de::AbstractMeanfieldEquations)
    # Options
    env --> :align
    cdot --> false

    lhs = getfield.(de.equations, :lhs)
    rhs = getfield.(de.equations, :rhs)
    lhs, rhs = _latexify(lhs, rhs, de.iv)
    return lhs, rhs
end

function _latexify(lhs_, rhs_, t)

    # Convert eqs to Exprs
    rhs = _to_expression.(rhs_)
    rhs = [MacroTools.postwalk(_postwalk_func, eq) for eq=rhs]
    rhs = [MacroTools.postwalk(_postwalk_average, eq) for eq=rhs]
    lhs = _to_expression.(lhs_)
    lhs = [MacroTools.postwalk(_postwalk_func, eq) for eq=lhs]
    lhs = [MacroTools.postwalk(_postwalk_average, eq) for eq=lhs]

    tsym = Symbol(t)
    dt = Symbol(:d, tsym)
    ddt = :(d/$(dt))
    lhs = [:($ddt * ($l)) for l=lhs]

    return lhs, rhs
end
_latexify(lhs,rhs) = _latexify([lhs],[rhs])

@latexrecipe function f(op::QNumber)
    # Options
    cdot --> false

    ex = _to_expression(op)
    ex = MacroTools.postwalk(_postwalk_func, ex)
    ex = MacroTools.postwalk(_postwalk_average, ex)
    ex = isa(ex,String) ? latexstring(ex) : ex
    return ex
end

@latexrecipe function f(s::SymbolicUtils.Symbolic{<:CNumber})
    # Options
    cdot --> false

    ex = _to_expression(s)
    ex = MacroTools.postwalk(_postwalk_func, ex)
    ex = MacroTools.postwalk(_postwalk_average, ex)
    ex = isa(ex,String) ? latexstring(ex) : ex
    return ex
end

@latexrecipe function f(c::CorrelationFunction)
    avg = average(c.op1*c.op2)
    return latexify(avg)
end

@latexrecipe function f(S::Spectrum)
    l = latexstring(L"\mathcal{F}(", latexify(S.corr), L")(\omega)")
    return l
end

_to_expression(x::Number) = x
function _to_expression(x::Complex) # For brackets when using latexify
    iszero(x) && return x
    if iszero(real(x))
        return :( $(imag(x))*im )
    elseif iszero(imag(x))
        return real(x)
    else
        return :( $(real(x)) + $(imag(x))*im )
    end
end
_to_expression(op::QSym) = op.name
_to_expression(op::Create) = :(dagger($(op.name)))
_to_expression(op::Transition) = :(Transition($(op.name),$(op.i),$(op.j)) )
function _to_expression(t::QMul)
    args = if SymbolicUtils._isone(t.arg_c)
        t.args_nc
    else
        SymbolicUtils.arguments(t)
    end
    return :( *($(_to_expression.(args)...)) )
end
_to_expression(t::QAdd) = :( +($(_to_expression.(t.arguments)...)) )

_to_expression(p::Parameter) = p.name
function _to_expression(s::SymbolicUtils.Symbolic)
    if SymbolicUtils.istree(s)
        f = SymbolicUtils.operation(s)
        fsym = if f === sym_average
            :AVG
        else
            Symbol(f)
        end
        args = map(_to_expression, SymbolicUtils.arguments(s))
        return :( $(fsym)($(args...)) )
    else
        return nameof(s)
    end
end


# println("Testing vec2state! and state2vec!...")
# srand(1)
# P = 20
# K = 10
# Sigma_X = spdiagm(rand(P) .+ 0.2)
# A = sparse(1:P, Int64[ceil(i/2) for i in 1:P], ones(P))
# Sigma_L = FactorAnalysis.randcor(K, 0.2)
# d = CFADistribution(Sigma_X, A, Sigma_L)
# N = 1000
# X = rand(d, N)
# S = X*X' / N
# Base.cov2cor!(S, sqrt(diag(S)))
# function ll_test(d, S, N)
#     Theta = inv(d.Sigma_X + d.A*d.Sigma_L*d.A')
#     logdet(Theta) - trace(S*Theta)
# end
# @test abs(loglikelihood(d, S, N) - ll_test(d, S, N)) < 1e-8

println("Testing CFADistributionSigma: vec2state! and state2vec!...")
srand(10)
P = 20
K = 10
Sigma_X = spdiagm(rand(P) .+ 0.2)
A = sparse(1:P, Int64[ceil(i/2) for i in 1:P], ones(P))
Sigma_L = FactorAnalysis.randcor(K, 0.2)

d = CFADistributionSigma(Sigma_X, A, Sigma_L)
x = zeros(P + length(A.nzval) + div((K+1)*K,2))
FactorAnalysis.state2vec!(x, d)
x_copy = copy(x)
FactorAnalysis.vec2state!(d, x)
FactorAnalysis.state2vec!(x, d)
@test maximum(abs(x .- x_copy)) <= 1e-8

g = FactorAnalysis.CFADistributionSigmaGradient(Sigma_X, A, Sigma_L)
x = zeros(P + length(A.nzval) + div((K+1)*K,2))
FactorAnalysis.state2vec!(x, g)
x_copy = copy(x)
FactorAnalysis.vec2state!(g, x)
FactorAnalysis.state2vec!(x, g)
@test maximum(abs(x .- x_copy)) <= 1e-8

println("Testing CFADistributionSigma: loglikelihood...")
N = 100000
X = rand(d, N)
S = X*X' / N
Base.cov2cor!(S, sqrt(diag(S)))
@test loglikelihood(d, S, N) < 0.0
@test loglikelihood(d, S, N) > loglikelihood(CFADistributionSigma(Sigma_X, A, Sigma_L .+ eye(K)*10), S, N)
function ll_test(d, S, N)
    Theta = inv(d.Sigma_X + d.A*d.Sigma_L*d.A')
    logdet(Theta) - trace(S*Theta)
end
@test abs(loglikelihood(d, S, N) - ll_test(d, S, N)) < 1e-8

println("Testing CFADistributionSigma: rand...")
d = CFADistributionSigma(Sigma_X, A, Sigma_L)
@test size(rand(d, N)) == (P,N)

println("Testing CFADistributionSigma: gradient (compare to finite difference)...")
g = dloglikelihood(d, S, N)
eps = 1e-8
tol = 1e-4
for i in 1:P
    tmp = Sigma_X[i,i]
    Sigma_X[i,i] = tmp + eps
    l1 = loglikelihood(CFADistributionSigma(Sigma_X, A, Sigma_L), S, N)
    Sigma_X[i,i] = tmp - eps
    l2 = loglikelihood(CFADistributionSigma(Sigma_X, A, Sigma_L), S, N)
    Sigma_X[i,i] = tmp
    @test abs(g.Sigma_X[i,i] - (l1-l2)/(2*eps)) < tol
    #@test abs(dSigma_X[i,i] - (l1-l2)/(2*eps)) < tol
end
for i in 1:K, j in i+1:K
    tmp = Sigma_L[j,i]
    Sigma_L[j,i] = Sigma_L[i,j] = tmp + eps
    l1 = loglikelihood(CFADistributionSigma(Sigma_X, A, Sigma_L), S, N)
    Sigma_L[j,i] = Sigma_L[i,j] = tmp - eps
    l2 = loglikelihood(CFADistributionSigma(Sigma_X, A, Sigma_L), S, N)
    Sigma_L[j,i] = Sigma_L[i,j] = tmp
    @test abs(g.Sigma_L[j,i] - (l1-l2)/(2*eps)) < tol
end
rows = rowvals(A)
vals = nonzeros(A)
for col = 1:K
    for j in nzrange(A, col)
        tmp = vals[j]
        vals[j] = tmp + eps
        l1 = loglikelihood(CFADistributionSigma(Sigma_X, A, Sigma_L), S, N)
        vals[j] = tmp - eps
        l2 = loglikelihood(CFADistributionSigma(Sigma_X, A, Sigma_L), S, N)
        vals[j] = tmp
        @test abs(g.A[rows[j],col] - (l1-l2)/(2*eps)) < tol
        #@test abs(dA[rows[j],col] - (l1-l2)/(2*eps)) < tol
    end
end

println("Testing CFADistributionSigma: fit_mle...")
truthLL = loglikelihood(d, S, N)
dopt = fit_mle(CFADistributionSigma, ones(K,K), spones(A), S, N, show_trace=false, iterations=1000)
@test loglikelihood(dopt, S, N) > truthLL

# println("Testing CFADistributionSigma: fit_map...")
# dopt2 = fit_map(Normal(0, 1.0), CFADistributionSigma, spones(A), S, N, show_trace=true, iterations=1000)
# @test loglikelihood(dopt2, S, N) < loglikelihood(dopt, S, N)
# @test loglikelihood(dopt2, S, N) > truthLL

println("Testing CFADistributionSigma: fit_mle with a fully connected latent variable...")
srand(10)
P = 40
K = 20
Sigma_X = spdiagm(rand(P) .+ 0.2)
A = sparse(1:P, Int64[ceil(i/2) for i in 1:P], ones(P))
Sigma_L = FactorAnalysis.randcor(K, 0.2)
d = CFADistributionSigma(Sigma_X, A, Sigma_L)
N = 100000
X = rand(d, N)

# add a confounding latent factor
confounder = randn(P)
for i in 1:N
    X[:,i] .+= randn()*confounder
end

S = X*X' / N
Base.cov2cor!(S, sqrt(diag(S)))
truthLL = loglikelihood(d, S, N)

# create an independent component to handle the confounder
maskSigma_L = ones(K+1,K+1)
for i in 1:K
    maskSigma_L[K+1,i] = maskSigma_L[i,K+1] = 0
end
A2 = sparse([1:P; 1:P], [Int64[ceil(i/2) for i in 1:P]; ones(Int64, P)*(K+1)], ones(2*P))

dopt = fit_mle(CFADistributionSigma, maskSigma_L, A2, S, N, show_trace=false, iterations=15000)
Th = inv(dopt.Sigma_X + dopt.A*dopt.Sigma_L*dopt.A')
# display(dopt.A)
# tmp = deepcopy(dopt.Sigma_L)
# dopt.Sigma_L = dopt.Sigma_L * (2*eye(K+1))
# #Base.cov2cor!(dopt.Sigma_L, 2*ones(K+1))
# FactorAnalysis.normalize_Sigma_L!(dopt::CFADistributionSigma)
# display(dopt.A)
# display(tmp)
# println()
# display(dopt.Sigma_L)
# @test maximum(abs(tmp .- dopt.Sigma_L)) < 1e-8
#
# display(dopt.Sigma_X + dopt.A*dopt.Sigma_L*dopt.A')
# println()
# println("GGG2")
# display(S)
# println()
#
# println(logdet(Th) - trace(S*Th))
# println(logdet(inv(S)) - trace(S*inv(S)))
# # display(Sigma_X)
# display(dopt.Sigma_X)
# display(Sigma_L)
# println()
# display(dopt.Sigma_L)
# display(dopt.A)
@test loglikelihood(dopt, S, N) > truthLL
@test FactorAnalysis.area_under_pr(inv(Sigma_L), inv(dopt.Sigma_L[1:K,1:K])) > 0.9

# println(truthLL)
#println(v)
# println()
# v = FactorAnalysis.area_under_pr(abs(FactorAnalysis.upper(inv(Sigma_L))) .> 0.01, abs(FactorAnalysis.upper(inv(dopt.Sigma_L + 0.01*eye(K)))))
# println(v)

#display(Sigma_L)