# This script provides an example of simulating changepoint data,
# using PELT to find these changepoints, and plotting the results
#
# Requires Gadfly plotting package

using Changepoints, Distributions
using Gadfly

n = 1000           # Sample size
λ = 100            # Frequencey of changes

# Generate a sample with changing
μ, σ = 1.0, Uniform(2.0, 15.0)
y, cps = @changepoint_sampler n λ Normal(μ, σ)
pelt_cps, cost = @PELT y Normal(μ, ?)

println(pelt_cps)
println(cps)

p = plot(y, pelt_cps)


using Makie
using Makie: textslider
using QuadGK

f(x) = exp(-x^3) # anything you want
g(x, λ) = f(x) * cos(λ * x)

function main()
    λ_slider, λ_obs = textslider(range(1, stop = 10000, step=100), "λ", start = 1)
    x = -1:0.001:1
    plotting = lift(λ_obs) do λ
        Point2f0.(x, g.(x, λ))
    end

    integral = lift(λ_obs) do λ
        res=quadgk(x->g(x,λ),-1,1)[1]
        "$res"
    end

    s1 = lines(x, f.(x))
    s2 = lines(plotting)
    Makie.text!(s2,integral,textsize=0.1,position = Point2f0(1.,1.))
    RecordEvents(
        vbox(s1,hbox(s2, λ_slider), parent=Scene()),
        "output"
    )
end

main()


####
#### Derivative operators
####

@inline ∂xᶜᵃᵃ(i, j, k, grid, u) = δxᶜᵃᵃ(i, j, k, grid, u) / Δx(i, j, k, grid)
@inline ∂xᶠᵃᵃ(i, j, k, grid, c) = δxᶠᵃᵃ(i, j, k, grid, c) / Δx(i, j, k, grid)

@inline ∂yᵃᶜᵃ(i, j, k, grid, v) = δyᵃᶜᵃ(i, j, k, grid, u) / Δy(i, j, k, grid)
@inline ∂yᵃᶠᵃ(i, j, k, grid, c) = δyᵃᶠᵃ(i, j, k, grid, c) / Δy(i, j, k, grid)

@inline ∂zᵃᵃᶜ(i, j, k, grid, w) = δzᵃᵃᶜ(i, j, k, grid, w) / ΔzF(i, j, k, grid)
@inline ∂zᵃᵃᶠ(i, j, k, grid, c) = δzᵃᵃᶠ(i, j, k, grid, c) / ΔzC(i, j, k, grid)

####
#### Derivative operators acting on functions
####

@inline ∂xᶜᵃᵃ(i, j, k, grid, f::F, args...) where F<:Function = (f(i+1, j, k, grid, args...) - f(i,   j, k, grid, args...)) / Δx(i, j, k, grid)
@inline ∂xᶠᵃᵃ(i, j, k, grid, f::F, args...) where F<:Function = (f(i,   j, k, grid, args...) - f(i-1, j, k, grid, args...)) / Δx(i, j, k, grid)

@inline ∂yᵃᶜᵃ(i, j, k, grid, f::F, args...) where F<:Function = (f(i, j+1, k, grid, args...) - f(i, j,   k, grid, args...)) / Δy(i, j, k, grid)
@inline ∂yᵃᶠᵃ(i, j, k, grid, f::F, args...) where F<:Function = (f(i, j,   k, grid, args...) - f(i, j-1, k, grid, args...)) / Δy(i, j, k, grid)

@inline ∂zᵃᵃᶜ(i, j, k, grid, f::F, args...) where F<:Function = (f(i, j, k+1, grid, args...) - f(i, j, k,   grid, args...)) / ΔzF(i, j, k, grid)
@inline ∂zᵃᵃᶠ(i, j, k, grid, f::F, args...) where F<:Function = (f(i, j, k,   grid, args...) - f(i, j, k-1, grid, args...)) / ΔzC(i, j, k, grid)

####
#### Second derivatives
####

@inline ∂²xᶜᵃᵃ(i, j, k, grid, c) = ∂xᶜᵃᵃ(i, j, k, grid, ∂xᶠᵃᵃ, c)
@inline ∂²xᶠᵃᵃ(i, j, k, grid, u) = ∂xᶠᵃᵃ(i, j, k, grid, ∂xᶜᵃᵃ, u)

@inline ∂²yᵃᶜᵃ(i, j, k, grid, c) = ∂yᵃᶜᵃ(i, j, k, grid, ∂yᵃᶠᵃ, c)
@inline ∂²yᵃᶠᵃ(i, j, k, grid, v) = ∂yᵃᶠᵃ(i, j, k, grid, ∂yᵃᶜᵃ, v)

@inline ∂²zᵃᵃᶜ(i, j, k, grid, c) = ∂zᵃᵃᶜ(i, j, k, grid, ∂zᵃᵃᶠ, c)
@inline ∂²zᵃᵃᶠ(i, j, k, grid, w) = ∂zᵃᵃᶠ(i, j, k, grid, ∂zᵃᵃᶜ, w)


@testset "GRclassic" begin
    b = 1
    M0 = 0.5
    d = GRclassic(b,M0)

    @test rate(d) == 1*log(10)
    @test location(d) == M0
    @test params(d) == (b, M0)

    @test  bvalue(d) == 1.
    @test  avalue(d, 100, 0.1) == 2.55

    @test  rand(d, 3) == [0.6547030177365674, 0.8281037053138567, 0.7333794310588824]
    @test  mean(d) == 1/rate(d) + M0

    @test pdf(d, 0.1) == 0.0
    @test round(pdf(d, 1.2), digits=2) == 0.46

    @test logpdf(d, 0.1) == -Inf
    @test round(logpdf(d,1.2), digits=2) == -0.78

    @test cdf(d, 0.1) == 0.0
    @test round(cdf(d,1.2), digits=2) == 0.8

    @test logcdf(d, 0.1) == -Inf
    @test round(logcdf(d,1.2), digits=2) == round(log(cdf(d,1.2)), digits=2)

    @test ccdf(d, 0.1) == 0.0
    @test round(ccdf(d,1.2), digits=2) == 0.2

    @test logccdf(d, 0.1) == -Inf
    @test round(logccdf(d,1.2), digits=2) == round(log(ccdf(d,1.2)), digits=2)

end


using Optim
using FANN
using Base.Test

net = ANN([10, 1], [:linear], b=0.1, errorfunc=:linear)

X = randn(10, 100)
y = sum(X, 1)
dset = DataSet(X, y)

# initial condition
w0 = weights(net)

# objective fucntion
f(w) = (weights!(net, w); mse(net, dset))
optimize(f, w0, ftol=1e-6) |> println

weights!(net, w0)
train!(net, dset, max_epochs=1000, epochs_between_reports=0, desired_error=1e-6)


using Wigxjpf
using Test

@testset "Wigxjpf.jl" begin
    Wigxjpf.wig_table_init(400, 9)
    Wigxjpf.wig_temp_init(400)

    @testset "Evaluate Wigner 3-j symbols" begin
        @test begin
            Wigxjpf.wig3jj(3, 3, 2, 3, -1, -2) ≈ -0.316227766016838
        end
        @test begin
            Wigxjpf.wig3jj(2, 4, 8, 2, -4, 2) ≈ 0.0
        end
        @test begin
            Wigxjpf.wig3jj(2, 4, 2, 2, 2, -4) ≈ 0.0
        end
        @test begin
            Wigxjpf.wig3jj(2, 4, 6, 1, -1, 0) ≈ 0.0
        end
        @test begin
            Wigxjpf.wig3jj(2, 2, 2, 2, 0, -2) ≈ -0.408248290463863
        end
        @test begin
            Wigxjpf.wig3jj(6, 6, 6, -2, -2, 4) ≈ 0.0
        end
        @test begin
            Wigxjpf.wig3jj(7, 5, 4, 7, -3, -4) ≈ -0.235702260395516
        end
        @test begin
            Wigxjpf.wig3jj(15, 15, 0, 3, -3, 0) ≈ 0.25
        end
        @test begin
            Wigxjpf.wig3jj(16, 11, 9, 4, -7, 3) ≈ -0.120622637445762
        end
        @test begin
            Wigxjpf.wig3jj(16, 14, 12, 6, 0, -6) ≈ 0.0
        end
        @test begin
            Wigxjpf.wig3jj(30, 60, 80, 4, 4, -8) ≈ -0.0190815797991916
        end
        @test begin
            Wigxjpf.wig3jj(60, 60, 60, 0, 30, -30) ≈ -0.0152857154897839
        end
        @test begin
            Wigxjpf.wig3jj(286, 200, 120, -20, 120, -100) ≈ 4.16710792865494e-10
        end
        @test begin
            Wigxjpf.wig3jj(320, 200, 120, -20, 120, -100) ≈ 3.81124616116626e-21
        end
        @test begin
            Wigxjpf.wig3jj(400, 400, 400, -20, 120, -100) ≈ 0.000749392731398951
        end
        @test begin
            Wigxjpf.wig3jj(140, 150, 160, 40, -80, 40) ≈ -0.00845794815539587
        end
    end

    @testset "Evaluate Wigner 6-j symbols" begin
        @test begin
    Wigxjpf.wig6jj(4, 4, 2, 4, 2, 2) ≈ -0.1
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 0, 4, 4, 2) ≈ 0.258198889747161
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 2, 4, 4, 2) ≈ -0.223606797749979
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 4, 4, 4, 2) ≈ 0.152752523165195
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 0, 4, 4, 4) ≈ -0.258198889747161
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 0, 4, 4, 2) ≈ 0.258198889747161
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 2, 4, 4, 4) ≈ 0.074535599249993
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 2, 4, 4, 2) ≈ -0.223606797749979
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 4, 4, 4, 4) ≈ 0.152752523165195
        end
        @test begin
            Wigxjpf.wig6jj(2, 2, 4, 4, 4, 2) ≈ 0.152752523165195
        end
        @test begin
            Wigxjpf.wig6jj(4, 4, 2, 4, 2, 4) ≈ 0.152752523165195
        end
        @test begin
            Wigxjpf.wig6jj(7, 6, 3, 2, 3, 6) ≈ 0.0
        end
        @test begin
            Wigxjpf.wig6jj(8, 8, 2, 8, 6, 2) ≈ -0.0277777777777778
        end
            @test begin
            Wigxjpf.wig6jj(12, 12, 8, 9, 7, 11) ≈ 0.0378232878232878
        end
        @test begin
            Wigxjpf.wig6jj(13, 12, 3, 6, 7, 12) ≈ 0.0
        end
        @test begin
            Wigxjpf.wig6jj(16, 16, 16, 16, 14, 14) ≈ 0.020221595293569
        end
        @test begin
            Wigxjpf.wig6jj(16, 16, 16, 16, 16, 16) ≈ -0.0126520807231535
        end
        @test begin
            Wigxjpf.wig6jj(40, 40, 40, 40, 40, 40) ≈ -0.00502940645686796
        end
        @test begin
            Wigxjpf.wig6jj(256, 240, 144, 224, 96, 160) ≈ 1.19377022891155e-17
            end
        @test begin
            Wigxjpf.wig6jj(460, 160, 300, 380, 460, 240) ≈ 2.42701001304468e-24
        end
        @test begin
            Wigxjpf.wig6jj(400, 400, 400, 400, 400, 400) ≈ 0.000155903212413242
        end
        @test begin
            Wigxjpf.wig6jj(160, 140, 120, 120, 140, 160) ≈ 0.000271672365647525
        end
    end
    @testset "Evaluate Wigner 9-j symbols" begin
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 2, 4, 6, 0, 4, 4) ≈ 0.03614031611621
        end
        @test begin
            Wigxjpf.wig9jj(1, 3, 4, 3, 1, 4, 2, 2, 0) ≈ 0.0645497224367903
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 6, 4, 8, 4, 6, 10) ≈ -0.00141609931936472
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 6, 4, 8, 4, 8, 10) ≈ 0.00793146728313409
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 6, 4, 8, 6, 4, 10) ≈ 0.00759958641837765
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 6, 4, 8, 6, 6, 10) ≈ 0.00851133416822978
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 6, 4, 8, 8, 2, 10) ≈ 0.0108588135723727
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 6, 4, 8, 8, 4, 10) ≈ 0.00335110262166189
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 6, 4, 8, 8, 6, 10) ≈ -0.0050301051199454
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 6, 4, 8, 8, 8, 10) ≈ 0.00012246297616039
        end
        @test begin
            Wigxjpf.wig9jj(2, 4, 6, 4, 6, 8, 6, 8, 10) ≈ -0.000922146636432351
        end
        @test begin
            Wigxjpf.wig9jj(2, 6, 4, 3, 5, 4, 1, 1, 0) ≈ -0.0281718084909506
        end
        @test begin
            Wigxjpf.wig9jj(6, 8, 4, 7, 7, 4, 1, 1, 2) ≈ 0.0230021853114118
        end
        @test begin
            Wigxjpf.wig9jj(2, 6, 8, 1, 7, 6, 1, 1, 2) ≈ -0.00595238095238095
        end
        @test begin
            Wigxjpf.wig9jj(6, 7, 7, 5, 6, 7, 1, 1, 2) ≈ -0.00215400852516591
        end
        @test begin
            Wigxjpf.wig9jj(17, 19, 14, 25, 16, 17, 16, 21, 19) ≈
            0.000281298301912545
        end
        @test begin
        Wigxjpf.wig9jj(30, 30, 30, 30, 6, 30, 30, 36, 20) ≈
            -7.78324615309539e-05
        end
        @test begin
        Wigxjpf.wig9jj(90, 60, 40, 90, 30, 70, 180, 90, 90) ≈
            -6.51130967998823e-07
        end
        @test begin
            Wigxjpf.wig9jj(120, 140, 260, 100, 140, 240, 120, 100, 80) ≈
            5.35253924139905e-16
        end
        @test begin
            Wigxjpf.wig9jj(200, 160, 100, 100, 200, 140, 120, 100, 200) ≈
            1.05597798065761e-07
        end
        @test begin
            Wigxjpf.wig9jj(10, 1, 9, 10, 1, 11, 18, 2, 20) ≈ 0.00257786760370727
        end
        @test begin
            Wigxjpf.wig9jj(30, 30, 60, 30, 6, 30, 30, 36, 60) ≈
            0.000168384777511567
        end
        @test begin
            Wigxjpf.wig9jj(40, 20, 60, 60, 60, 120, 20, 40, 60) ≈
            0.000268744961031981
        end
        @test begin
            Wigxjpf.wig9jj(60, 40, 20, 60, 20, 40, 120, 60, 60) ≈
            0.000268744961031981
        end
        @test begin
            Wigxjpf.wig9jj(90, 60, 40, 90, 30, 70, 180, 90, 90) ≈
            -6.51130967998823e-07
        end
        @test begin
            Wigxjpf.wig9jj(100, 102, 104, 106, 108, 110, 112, 114, 116) ≈
            4.18840599233205e-06
        end
        @test begin
            Wigxjpf.wig9jj(148, 134, 120, 164, 120, 140, 20, 40, 60) ≈
            4.94488048916531e-06
        end
        @test begin
            Wigxjpf.wig9jj(168, 180, 122, 120, 128, 80, 104, 80, 60) ≈
            -7.94917060961424e-08
        end
        @test begin
            Wigxjpf.wig9jj(188, 134, 172, 122, 146, 104, 140, 120, 160) ≈
            -1.86808574631986e-07
        end
        @test begin
            Wigxjpf.wig9jj(120, 140, 230, 100, 140, 220, 120, 100, 80) ≈
            -1.36799650122649e-06
        end
        @test begin
            Wigxjpf.wig9jj(200, 160, 100, 100, 200, 140, 120, 100, 200) ≈
            1.05597798065761e-07
        end
        @test begin
            Wigxjpf.wig9jj(40, 40, 80, 40, 40, 80, 40, 40, 80) ≈
            0.000111633558388972
        end
    end

    Wigxjpf.wig_temp_free()
    Wigxjpf.wig_table_free()
end


struct UnitCell{Dimensions, T <: AbstractFloat}
    tag::Symbol # currently unused
    # could be used to speed up neighbor search
    basis::Vector{Point{Dimensions, T}}
    dirs::NTuple{Dimensions, Vec{Dimensions, T}}
end

function UnitCell(
        pos::Point{D, T},
        dirs::NTuple{D, Vec{D, T}}
    ) where {D, T <: AbstractFloat}
    UnitCell([pos], dirs)
end

function UnitCell(
        basis::Vector{Point{D, T}},
        dirs::NTuple{D, Vec{D, T}}
    ) where {D, T <: AbstractFloat}
    UnitCell(:None, basis, dirs)
end

"""
    get_pos(uc, a, u)

Returns R = Rₐ + ∑ᵢ uᵢaᵢ for a UnitCell uc.
"""
function get_pos(
        uc::UnitCell{D, T},
        alpha::Int64,
        uvw::NTuple{D, Int64}
    ) where {D, T}
    uc.basis[alpha] + reduce(+, map((i, v) -> i*v, uvw, uc.dirs))
    # B.pos + reduce(+, map((i, v) -> i * v, uvw, B.dirs))
end

# Returns the dimensionality of a UnitCell
dims(uc::UnitCell{D, T}) where {D, T} = D

function merge(ucs::UnitCell{D, T}...) where {D, T}
    @assert all(
        uc -> all(
            v -> v in ucs[1].dirs,
            uc.dirs
        ),
        ucs
    ) "Lattice vectors must match!"

    UnitCell(
        ucs[1].tag,
        [p for uc in ucs for p in uc.basis],
        ucs[1].dirs
    )
end

# TODO
# center! for more safety in get_neighbors

# Literally just constructors from here on out

################################################################################
### 2D Lattices
################################################################################

"""
    square()

Returns a UnitCell representing a square lattice.

Keyword arguments include
- `pos::Point{2, T} = Point{2}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `isa::Symbol = :cubic`: Lattice shape to use (:cubic or :primitive)
"""
function square(;
        pos::Point{2, T} = Point{2}(0.),
        a::T = 1.0,
        b::T = 1.0,
        isa::Symbol = :cubic
    ) where T <: AbstractFloat
    UnitCell(:cubic, pos, (a * Vec(1., 0.), b * Vec(0., 1.)))
end

"""
    triangle()

Returns a UnitCell representing a triangular lattice.

Keyword arguments include
- `pos::Point{2, T} = Point{2}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function triangle(;
        pos::Point{2, T} = Point{2}(0.),
        a::T = 1.0,
        b::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    if isa == :cubic
        return UnitCell(
            :cubic,
            [pos, pos + Point(a*sind(30.), b*cosd(30.))],
            (a * Vec(1., 0.), b * Vec(0., 1.))
        )
    else
        return UnitCell(
            :primitive,
            pos,
            (a*Vec(0., 1.), b*Vec(cosd(30.0), sind(30.0)))
        )
    end
end

"""
    honeycomb()

Returns a UnitCell representing a honeycomb lattice.

Keyword arguments include
- `pos::Point{2, T} = Point{2}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function honeycomb(;
        pos::Point{2, T} = Point{2}(0.),
        a::T = 1.0,
        b::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    if isa == :cubic
        return UnitCell(
            :cubic,
            [pos, pos + Point(0.0, b)],
            (2a * cosd(30) * Vec(1., 0.), 2b * cosd(30) * Vec(0., 1.))
        )
    else
        return UnitCell(
            :primitive,
            [pos, pos + Point(a, 0.)],
            (
                2a * cosd(30) * Vec(0., 1.),
                2b * cosd(30) * Vec(cosd(30.0), sind(30.0))
            )
        )
    end
end

################################################################################
### 3D Lattices
################################################################################

"""
    sc()

Returns a UnitCell representing a simple cubic lattice.

Keyword arguments include
- `pos::Point{3, T} = Point{3}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `c::T = 1.0`: Scaling along the third lattice vector.
- `isa::Symbol = :cubic`: Lattice shape to use (:cubic or :primitive)
"""
function sc(;
        pos::Point{3, T} = Point{3}(0.),
        a::T = 1.0,
        b::T = 1.0,
        c::T = 1.0,
        isa::Symbol = :cubic
    ) where T <: AbstractFloat
    UnitCell(
        :cubic,
        [pos],
        (
            a * Vec(1., 0., 0.),
            b * Vec(0., 1., 0.),
            c * Vec(0., 0., 1.)
        )
    )
end

"""
    bcc()

Returns a UnitCell representing a body centered cubic lattice.

Keyword arguments include
- `pos::Point{3, T} = Point{3}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `c::T = 1.0`: Scaling along the third lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function bcc(;
        pos::Point{3, T} = Point{3}(0.),
        a::T = 1.0,
        b::T = 1.0,
        c::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    if isa == :cubic
        return UnitCell(
            :cubic,
            [pos, pos + Point{3}(0.5)],
            (
                a * Vec(1., 0., 0.),
                b * Vec(0., 1., 0.),
                c * Vec(0., 0., 1.)
            )
        )
    else
        return UnitCell(
            :primitive,
            [pos],
            (
                a * 0.5 * Vec(1., 1., -1.),
                b * 0.5 * Vec(1., -1., 1.),
                c * 0.5 * Vec(-1., 1., 1.)
            )
        )
    end
end

"""
    fcc()

Returns a UnitCell representing a face centered cubic lattice.

Keyword arguments include
- `pos::Point{3, T} = Point{3}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `c::T = 1.0`: Scaling along the third lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function fcc(;
        pos::Point{3, T} = Point{3}(0.),
        a::T = 1.0,
        b::T = 1.0,
        c::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    if isa == :cubic
        return UnitCell(
            :cubic,
            [
                pos,
                pos + Point(0.5, 0.5, 0.0),
                pos + Point(0.5, 0.0, 0.5),
                pos + Point(0.0, 0.5, 0.5)
            ],
            (
                a * Vec(1., 0., 0.),
                b * Vec(0., 1., 0.),
                c * Vec(0., 0., 1.)
            )
        )
    else
        return UnitCell(
            :primitive,
            [pos],
            (
                a * 0.5 * Vec(1., 1., 0.),
                b * 0.5 * Vec(1., 0., 1.),
                c * 0.5 * Vec(0., 1., 1.)
            )
        )
    end
end

"""
    diamond()

Returns a UnitCell representing a diamond lattice.

Keyword arguments include
- `pos::Point{3, T} = Point{3}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `c::T = 1.0`: Scaling along the third lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function diamond(;
        pos::Point{3, T} = Point{3}(0.),
        a::T = 1.0,
        b::T = 1.0,
        c::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    merge(
        fcc(pos = pos, a=a, b=b, c=c, isa=isa),
        fcc(pos = pos + Point{3}(0.25), a=a, b=b, c=c, isa=isa)
    )
end


# tests for kinetic_monte_carlo.jl
using MonteCarloX
using StatsBase
using Random

# for kldivergence
include("utils.jl")

function test_kmc_next(;verbose = false)
    pass = true

    list_rates = [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0]
    sum_rates = sum(list_rates)
    weights = ProbabilityWeights(list_rates)
    tbin = 0.5 / sum_rates
    pdf_time(t) = sum_rates * exp(-(t + tbin) * sum_rates)
    cdf_time(t) = 1.0 - exp(-(t + tbin) * sum_rates)
    pdf_event(e) = list_rates[e] / sum_rates

    if verbose
        println("... check sum(weights)[$(sum(weights))] == sum_rates [$(sum_rates)]")
    end
    pass &= abs(sum_rates - sum(weights)) < 1e-10

    samples = 10^5
    times = zeros(samples)
    events = zeros(samples)
    alg = KineticMonteCarlo()
    rng = MersenneTwister(1000)
    for i = 1:samples
        time, event = next(alg, rng, weights)
        times[i] = time
        events[i] = event
    end

    t_max = maximum(times)
    hist_time = fit(Histogram, times, 0:tbin:t_max + tbin)
    cdf_time_meas = normalize(float(hist_time), mode = :probability)
    pdf_sum = 0.0
    for i in 1:length(cdf_time_meas.weights)
        pdf_sum += cdf_time_meas.weights[i]
        cdf_time_meas.weights[i] = pdf_sum
    end
    # since measurements only performed until time_max, we have to rescale
    # pdf_time(t) to integrate to 1 in the bounds (-infty, t_max). This is
    # achieved by deviding by cdf(t_max). This then percolates to cdf(t)
    cdf_time_renorm(t) = cdf_time(t) / cdf_time(cdf_time_meas.edges[1][end])
    kld_time = kldivergence(cdf_time_meas, cdf_time_renorm)
    pass &= abs(kld_time) < 0.01
    if verbose
        println("... kldivergence cdf_time to cdf_true = $(kld_time)")
    end

    hist_event = fit(Histogram, events, 1:(length(list_rates) + 1))
    pdf_event_meas = normalize(float(hist_event))
    kld_event = kldivergence(pdf_event_meas, pdf_event)
    if verbose
        println("... kldivergence pdf_event to pdf_true = $(kld_event)")
    end
    pass &= abs(kld_event) < 0.01

    return pass
end

# TODO
function test_kmc_event_handler(;verbose = false)
    pass = true

    return pass
end


using Roots
using Test
import SpecialFunctions.erf

struct SomeInterval
    a
    b
end
Base.extrema(I::SomeInterval) = I.a < I.b ? (I.a, I.b) : (I.b, I.a)

include("./test_find_zero.jl")
include("./test_bracketing.jl")
include("./test_derivative_free.jl")
include("./test_simple.jl")
include("./test_find_zeros.jl")
include("./test_fzero.jl")
include("./test_newton.jl")
include("./test_simple.jl")

#include("./test_composable.jl")
#include("./runbenchmarks.jl")
#include("./test_derivative_free_interactive.jl")


"""
    DeathDistribution

An abstract type used to form an assumption of how deaths occur throughout a
    year. See `Balducci()`, `Uniform()`, and `Constant()` for concrete
    assumption types.
"""
abstract type DeathDistribution end

"""
    Balducci()

A `DeathDistribution` type that assumes a decreasing force of mortality
over the year.
"""
struct Balducci <: DeathDistribution end

"""
    Uniform()

A `DeathDistribution` type that assumes an increasing force of mortality
over the year.
"""
struct Uniform <: DeathDistribution end

"""
    Constant()

A `DeathDistribution` type that assumes a constant force of mortality
over the year.
"""
struct Constant <: DeathDistribution end


"""
Add a bias/intercept term to a vector on the fly without creating or copying data:

- `BiasVector(rand(10))` is roughly equivalent to `vcat(rand(10), 1.0)`
"""
immutable BiasVector{A<:AVecF} <: AVecF
  vec::A
end

Base.length(v::BiasVector) = length(v.vec) + 1
Base.size(v::BiasVector) = (length(v),)
Base.getindex(v::BiasVector, i::Integer) = i > length(v.vec) ? 1.0 : getindex(v.vec, i)
Base.setindex!(v::BiasVector, x::Real, i::Integer) = setindex!(v.vec, x, i)

"""
Adda bias/intercept term to a matrix on the fly without creating or copying data:

- `BiasMatrix(rand(10,5))` is roughly equivalent to `hcat(rand(10,5), ones(10))`
"""
immutable BiasMatrix{A<:AMatF} <: AMatF
  mat::A
end

Base.length(m::BiasMatrix) = length(m.mat) + nrows(m.mat)
Base.size(m::BiasMatrix) = (nrows(m.mat), ncols(m.mat)+1)
Base.getindex(m::BiasMatrix, i::Integer, j::Integer) = j > ncols(m.mat) ? 1.0 : getindex(m.mat, i, j)
Base.setindex!(m::BiasMatrix, x::Real, i::Integer, j::Integer) = setindex!(m.mat, x, i, j)

#-----------------------------------------------------------------------------------#
"""
Add second-order interaction terms on the fly without creating or copying data:

- `TwoWayInteractionVector(rand(p))` "adds" the `binomial(p, 2)` interaction terms
"""
immutable TwoWayInteractionVector{A<:AVecF} <: AVecF
    vec::A
end
Base.length(v::TwoWayInteractionVector) = length(v.vec) + binomial(length(v.vec), 2)
Base.size(v::TwoWayInteractionVector) = (length(v),)
function Base.getindex(v::TwoWayInteractionVector, i::Integer)
    first = 0
    second = i
    n = length(v.vec)
    if i <= n
        return v.vec[i]
    else
        while second > n
            second -= n
            first += 1
            n -= 1
        end
        return v.vec[first] * v.vec[first + second]
    end
end

"""
Add second-order interaction terms on the fly without creating or copying data:

- `TwoWayInteractionMatrix(rand(n, p))` "adds" the `binomial(p, 2)` interaction terms
to each row
"""
immutable TwoWayInteractionMatrix{A<:AMatF} <: AMatF
    mat::A
end
function Base.length(m::TwoWayInteractionMatrix)
    n, p = size(m.mat)
    n * (p + binomial(p, 2))
end
function Base.size(m::TwoWayInteractionMatrix)
    n, p = size(m. mat)
    (n, p + binomial(p, 2))
end
function Base.getindex(m::TwoWayInteractionMatrix, i::Integer, j::Integer)
    first = 0
    second = j
    n, p = size(m.mat)
    if j <= p
        return m.mat[i, j]
    else
        while second > p
            second -= p
            first += 1
            p -= 1
        end
        return m.mat[i, first] * m.mat[i, first + second]
    end
end


# Compass search as described on page 18 (402) in Kolda2003:
#  Kolda, Tamara G., Robert Michael Lewis, and Virginia Torczon. "Optimization
#  by direct search: New perspectives on some classical and modern methods.",
#  SIAM review 45.3 (2003): 385-482.
function compass_search(f, n; max_fevals = 1e6, delta_tol = 1e-20,
  step_size = 1.0, x = false, known_fmin = false)

  @assert delta_tol > 0
  @assert step_size > delta_tol

  if x == false
    x = randn(n, 1)
  end

  directions = [Matrix{Float64}(I, n,n) -Matrix{Float64}(I, n, n)]

  num_fevals = 0
  fc = fbest = Inf
  candidate = zeros(n, 1)

  while(step_size > delta_tol && num_fevals <= max_fevals)
    # Check all directions to find a better point
    found_better = false
    for direction in 1:(2*n)
      candidate = x + step_size .* directions[:, direction]

      fc = f(candidate)
      num_fevals += 1

      if fc < fbest
        found_better = true
        break
      end

    end

    if found_better
      fbest = fc
      x = candidate

      if fbest == known_fmin
        break
      end
    else
      step_size = step_size / 2
    end
  end

  return x, fbest, num_fevals
end


"""
***
# Viscosity of Liquids
***
Data are calculated by solving the equation:

``Log_{10}(\\mu_{G})= A+B/T+C*T+D*T^2``

This function can be used by two methods:
* Get the liquid viscosity by introducing the component's name and the temperature.
* Get the full list of components' names supported.
***
### Example:
`julia> LiquidViscosity("Water",298.15)`

`0.911`

`julia> LiquidViscosity()`

`String["Carbon monoxide", "1,3-Butadiene", "Cyclohexanol", "Styrene", "n-Octane", "Acetic acid", "Hydrogen", "Dichloromethane",
"1-Butanol (n-Butanol)", "1-Octene", "Methyl chloride", "Ethylene glycol", "Acrylic acid", "Ethanol", "1,1,1-Trichloroethane", "Oxygen",
"1,4-Dioxane", "n-Nonane", "Pyridine", "Hydrogen cyanide", "Ammonia", "carbon disulfide", "1-Nonane (n-Nonane)", "Hydrogen chloride",
"1-Decene", "Chlorine", "Acetic anhydride", "1,1-Dichloroethane", "Methyl isobutyl ketone", "Chloroform", "1,1,2-Trichloroethane",
 "p-Xylene", "Butyric acid", "Nitrobenzene", "1-Butene", "Ethylene oxide", "Cyclohexane", "Acetaldehyde", "1,2-Dichloroethane",
 "Methane", "n-Heptane", "n-Hexane", "Nitrogen", "Carbon tetrachloride", "Hydrogen peroxide", "n-Decane", "Ethane", "Water",
 "Diethyl ketone", "Ethyl chloride", "Naphthalene", "Fluorine", "1-Propanol (n-Propanol)", "n-Butane", "o-Xylene", "Toluene",
 "n-Pentane", "Aniline", "Methyl bromide", "m-Xylene", "Carbon dioxide", "Benzene", "Phenol", "i-Butane (iso-Butane)", "Methylamine",
 "Dimethylamine", "Acetone", "Ethyl acetate", "Ethylbenzene", "Cyclopropane", "Glycerol", "Sulfur dioxide", "Ethylene", "Diethyl ether", "Vinyl acetate"]`
"""
function LiquidViscosity(substance::String,T::Float64)
    try
        param=DataLiquidViscosity[substance]

        if T >= param[6] || T <= param[5]
            print_with_color(:red,
            @sprintf("%4.4f K is out of range %4.4f K - %4.4f K.\n",T,
            param[5],param[6]))
            return NaN
        else
            μG=10^(param[1]+param[2]/T+param[3]*T+param[4]*T^2)
        end
    catch
        print_with_color(:red,@sprintf("%s not found in database.\n",substance))
        return NaN
    end
end

function LiquidViscosity()
    println(keys(DataLiquidViscosity))
end


abstract type AbstractWald <: SequentialSamplingModel
end

"""
# Wald Constructor
- `υ`: drift rate
- `α`: decision threshold
- `θ`: a encoding-response offset
## Usage
````julia
using SequentialSamplingModels
dist = Wald(υ=3.0, α=.5, θ=.130)
rt = rand(dist, 10)
like = pdf.(dist, rt)
loglike = logpdf.(dist, rt)
````
"""
struct Wald{T1,T2,T3} <: AbstractWald
    ν::T1
    α::T2
    θ::T3
end

Wald(;ν, α, θ) = Wald(ν, α, θ)

function pdf(d::AbstractWald, t::Float64)
    return pdf(InverseGaussian(d.α/d.ν, d.α^2), t - d.θ)
end

function logpdf(d::AbstractWald, t::Float64)
    return logpdf(InverseGaussian(d.α/d.ν, d.α^2), t - d.θ)
end

function logccdf(d::Wald, t::Float64)
    return logccdf(InverseGaussian(d.α/d.ν, d.α^2), t - d.θ)
end

function cdf(d::Wald, t::Float64)
    return cdf(InverseGaussian(d.α/d.ν, d.α^2), t - d.θ)
end

rand(d::AbstractWald) = rand(InverseGaussian(d.α/d.ν, d.α^2)) + d.θ

function rand(d::AbstractWald, n::Int)
    return rand(InverseGaussian(d.α/d.ν, d.α^2), n) .+ d.θ
end

mean(d::AbstractWald) = mean(InverseGaussian(d.α/d.ν, d.α^2)) + d.θ
std(d::AbstractWald) = std(InverseGaussian(d.α/d.ν, d.α^2))

"""
# WaldMixture Constructor
- `υ`: drift rate
- `σ`: standard deviation of drift rate
- `α`: decision threshold
- `θ`: a encoding-response offset
## Usage
````julia
using SequentialSamplingModels
dist = WaldMixture(υ=3.0, σ=.2, α=.5, θ=.130)
rt = rand(dist, 10)
like = pdf.(dist, rt)
loglike = logpdf.(dist, rt)
````
## References
Steingroever, H., Wabersich, D., & Wagenmakers, E. J. (2020).
Modeling across-trial variability in the Wald drift rate parameter.
Behavior Research Methods, 1-17.
"""
struct WaldMixture{T1,T2,T3,T4} <: AbstractWald
    ν::T1
    σ::T2
    α::T3
    θ::T4
end

WaldMixture(;ν, σ, α, θ) = WaldMixture(ν, σ, α, θ)

function pdf(d::WaldMixture, t::Float64)
    @unpack ν, σ, α ,θ = d
    c1 = α/√(2*π*(t - θ)^3)
    c2 = 1/cdf(Normal(0,1), ν/σ)
    c3 = exp(-(ν*(t - θ) - α)^2/(2*(t - θ)*((t - θ)*σ^2 + 1)))
    c4 = (α*σ^2 + ν)/√(σ^2*((t - θ)*σ^2 + 1))
    return c1 * c2 * c3 * cdf(Normal(0,1), c4)
end

function logpdf(d::WaldMixture, t::Float64)
    @unpack ν, σ, α ,θ = d
    c1 = log(α) - log(√(2*π*(t - θ)^3))
    c2 = log(1) - logcdf(Normal(0,1), ν/σ)
    c3 = -(ν*(t - θ) - α)^2/(2*(t - θ)*((t - θ)*σ^2 + 1))
    c4 = (α*σ^2 + ν)/√(σ^2*((t - θ)*σ^2 + 1))
    return c1 + c2 + c3 + logcdf(Normal(0,1), c4)
end

function rand(d::WaldMixture)
    x = rand(truncated(Normal(d.ν, d.σ), 0, Inf))
    return rand(InverseGaussian(d.α/x, d.α^2)) + d.θ
end

function rand(d::WaldMixture, n::Int)
    return map(x->rand(d), 1:n)
end


using WignerSymbols
using JLD2

Nmax = 250;
A_lnm = zeros(Float32,2Nmax+1,Nmax,2Nmax+1);
B_lnm = zeros(Float32,2Nmax+1,Nmax,2Nmax+1);

#arr_to_fill = Array{Any}(nothing, l+n+1, n+1);

function f(Nmax)
    cN = 0;
    cM = 0;
    for l = 0:2Nmax
        for n = 1:Nmax
            mᵒ    = max(l-n,n+1)
            m_max = min(l+n,Nmax)
            #println("l,n: ", l, " ", n, " " , m_max-mᵒ, " " , (n+l)-mᵒ)
            cN += max(0,m_max-mᵒ)
            cM += max(0,(n+l)-mᵒ)
            #= for m = mᵒ:n+l
                println()
                fac = sqrt(2(2l+1)*(2n+1)*(2m+1));
                #println(m," ",l, " ", n)
                @inbounds A_lnm[l+1,n,m-n+1] = fac*wigner3j(m,n,l,-1,1,0);
                if l>1
                    @inbounds B_lnm[l+1,n,m-n+1] = fac*wigner3j(m,n,l,-1,-1,2);
                end
            end
            println(l," ",n) =#
        end
    end
    return cN,cM
end

global
global

@save "wignerTable.jld2" A_lnm B_lnm


jacobian(::Type{B},
         X::AbstractMatrix{F},
         u::Vector{F}) where {B<:AbUniDirBasis, F<:Real} = derv(B, X*u) .* X

jacobian(ω::B, X) where {B<:AbUniDirBasis} = derv(ω, X) .* X

function valueandjacobian(::Type{B},
                    X::Matrix{F},
                    u::Vector{F}) where {B<:AbUniDirBasis, F<:Real}
  Xu=X*u
  calc(B, Xu), derv(B, Xu).*X
end

function backpropωs(prob::Problem{L, B, CS}, ens::AbEnsemble{F, B};
    maxiter::Int=10, αbtls::F=.3, βbtls::F=.25, maxjter::Int=9, ftol::F=1e-6) where {F, L, B, CS}

    m = length(ens)
    p = sizep(prob)
    ∇L_u = Array{F}(p, m)
    currloss = loss(prob, ens.fit)
    converged = false
    # printfmt("    0: 0) L={:6.4f}", currloss)

    t = 1/sizen(prob)
    for i in 1:maxiter
        ∇L_z = derv(L(), prob.y, ens.fit)
        for j in 1:m
            ∇L_u[:, j] = ens.β[j] * jacobian(ens.ωs[j], prob.X)' * ∇L_z
        end

        normsqgrad = sum(abs2.(∇L_u))
        # printfmtln("\t\t|∇f|={:8.6f}", normsqgrad)
        jter = 0
        while jter < maxjter          #Back-tracking Line Search
            jter += 1
            newens = deepcopy(ens)
            stepalong!(newens, ∇L_u, t)
            steplnsq = ssedir(ens, newens)
            newloss = loss(prob, newens.fit)
            # printfmt("   {:2d}:{:2d}) L={:6.4f} t={:.6f}", i, jter, newloss, t)
            if newloss < currloss - αbtls*t*steplnsq
                percentgain = 1 - newloss/currloss
                # printfmtln("\t\tBTLS success %gain={:9.7f}", percentgain)
                converged = percentgain < ftol
                currloss = newloss
                ens = newens
                break
            end
            t *= βbtls
        end
        if normsqgrad < 1e-6
            # printfmtln("   {:2d}) norm too small", i)
            break
        end
        if jter == maxjter
            # printfmtln("   {:2d}) t too small", i)
            break
        end
        if converged
            # printfmtln("   {:2d}) Converged", i)
            break
        end
    end
    # println("\t\tNorms: ", round.(Int, 100*[normcs(prob, ens.ωs[j].u) for j in 1:m]))
    ens, currloss
end


using Revise
import Optim
using ADCME
mpi_init()

r = mpi_rank()
θ0 = Variable(rand(2))
θ = mpi_bcast(θ0)
a = 1.
b = 100.
x, y = θ[1], θ[2]
fxy = (a-x)^2 + b*(y-x^2)^2
local_loss = fxy * (r+1)
loss = mpi_sum(local_loss)
g = gradients(loss, θ0)

sess = Session(); init(sess)

_f_eval = 0
_g_eval = 0
function f(x0)
    global _f_eval += 1
    L = run(sess, loss, θ0=>x0)
    L
end

function g!(G, x0)
    global _g_eval += 1
    gv = run(sess, g, θ0=>x0)
    G[:] = gv
end

x0 = rand(2)
mpi_optimize(f, g!, x0)

mpi_finalize()


#---------------------------------------------------------------------#
#This module stores everything needed to build different qudrature points
#in 1D.
#Written by F.X. Giraldo on 4/10/2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey, CA 93943-5216
#---------------------------------------------------------------------#
module QuadraturePoints
    include("chebyshev_gauss.jl")
    include("equispaced_points.jl")
    include("legendre_poly.jl")
    include("legendre_gauss.jl")
    include("lobatto_gauss.jl")
    include("lagrange_basis.jl")
end


using Gadfly, Compat

plot(layer((x,y) -> x*exp(-(x-(@compat round(Int, x)))^2-y^2), -8., 8, -2., 2),
     layer((x, y) -> sqrt(hypot(x, y)), -2, 2, -1, 1),
     layer(x=[3], y=[1.5], Geom.point))


using Convex
using ECOS
# Formulate and solve a quadratic problem:
#     min ||g||^2 st. Afriat inequalities for random d

# Initialize W^c=W (consumption measurement error).
W=ones(M,T,N)
# Initialize simulated q^*=q.
qsim=zeros(M,T,N)
# Initialize simulated Afriat numbers δ=δsim.
δsim=zeros(M,T,N)
# Initialize optimal value to a high number
optimval=ones(n,ndelta+1)*10000
# Initialize matrix to verify Afriat inequalities of the simulated data.
aiverify2=zeros(n,T,T)
# Convex model, variable definition.
v=Variable(T, Positive())
c=Variable(T,K,Positive())
# Identity Matrix in Convex.
P=I+zeros(1,1)

# Loop to check all entries of deltavec per individual and save in Delta the
# value corresponding to the lowest objective value
for dt=2:ndelta+1
    for id=1:n
        Deltatemp[id]=deltavec[dt-1]
        # Convex objective value, see Appendix C.
        modvex=minimize(quadform(rho[id,1,:]'*(c[1,:]'-cve[id,1,:]),P)
            +quadform(rho[id,2,:]'*(c[2,:]'-cve[id,2,:]),P)
            +quadform(rho[id,3,:]'*(c[3,:]'-cve[id,3,:]),P)
            +quadform(rho[id,4,:]'*(c[4,:]'-cve[id,4,:]),P))
        # Define the Afriat's constraints.
        for t=1:T
            for s=1:T
                modvex.constraints+=v[t]-v[s]-Deltatemp[id]^(-(t-1))*rho[id,t,:]'*(c[t,:]'-c[s,:]')>=0
            end
        end
        # Solve the model using ECOS.
        solve!(modvex,() -> ECOS.Optimizer(verbose=false))
        # Save objective value.
        optimval[id,dt]=modvex.optval
        # Re-initialize the verification matrix.
        aiverify=zeros(n,T,T)

        # Keep the lowest objective value, associated Delta
        if (optimval[id,dt]<optimval[id,dt-1])
            Delta[id]=Deltatemp[id]
            for i=1:T
                vsim[id,i]=v.value[i]
                for j=1:K
                    cvesim[id,i,j]=c.value[i,j]
                end
            end
            # Verify Afriat's inequalities with the solution.
            for t=1:T
                for s=1:T
                    aiverify2[id,t,s]=vsim[id,t]-vsim[id,s]-Delta[id]^(-(t-1))*rho[id,t,:]'*(cvesim[id,t,:]-cvesim[id,s,:])
                end
            end
        end
    end
    #Memory cleaning.
    modvex=nothing
    GC.gc()
end

minimum(aiverify2)


function weight(PPDE, pThresh, alpha, precision)
	((PPDE >= 0) & (PPDE <= 1)) || error("PPDE has to be between 0 and 1")
	if PPDE==1
	        weight = -alpha * 2 * log(1-pThresh)
	else
	        weight = min(2 * (-log(1 - PPDE) + log(1 - pThresh)),
				-alpha * 2 * log(1 - pThresh))
	end

	return round(weight, sigdigits = precision)
end


@inline function density_integral(bound_low::Real, bound_up::Real,
                                  norm::Real, slope::Real, ρ::Real)

    # density integral (eq. 9 M01)
    nb = 4π * norm * bound_low^3 / ρ
    density = nb * ( (bound_up/bound_low)^(3.0 - slope ) - 1.0 ) / ( 3.0 - slope )

    if ( 3.0 - slope_soft ) < slope < ( 3.0 + slope_soft )
        slope_var = (slope - 3.0)/slope_soft
        density2 = nb * log10(bound_up/bound_low)
        if slope_var != 0.0
            density = density * slope_var + density2 * ( 1.0 - slope_var )
        else
            density = density2
        end
    end

    return density
end

# function calculateCRNumber(CR_N, CR_S, CR_Cut, ρ;pmin::Real=10.0, pmax::Real=1.e7,
#                            mc::Real=0.0, SelectBin::Int=-1,
#                            units::GadgetPhysical=GadgetPhysical(), verbose=true)
#     # calculates engergy per bin in cgs units.

#     CR_Num = 0.0
#     cnst_c = 2.9979e10
#     ρ *= units.m_unit / units.l_unit^3  # convert rho in cgs units

#     CR_N .*= 1.e20  # compensate for io

#     Nbins = length(CR_N)
#     bin_width = log10(pmax/pmin)/Nbins
#     bins = 10.0.^collect(log10(pmin):bin_width:log10(pmax)) .* mc
#     above_cut = findall( bins .> CR_Cut/mc )

#     if length(above_cut) > 0
#         bins[above_cut] .= CR_Cut/mc
#     end

#     if SelectBin == -1
#         CR_Num = zeros(length(CR_N))
#         for i = 1:length(CR_N)
#             CR_Num[i] = density_integral(bin[i], bin[i+1], CR_N[i], CR_S[i], ρ)
#         end
#     elseif length(SelectBin) == 1
#         CR_Num = density_integral(bin[SelectBin], bin[SelectBin+1], CR_N[SelectBin], CR_S[SelectBin], ρ)
#     elseif length(SelectBin) > 1
#         CR_Num = zeros(length(SelectBin))
#         if i ∈ SelectBin
#             CR_Num[i-SelectBin[1]+1] = density_integral(bin[i], bin[i+1], CR_N[i], CR_S[i], ρ)
#         end
#     end
#     return CR_Num
# end


"""
    function threegrouptransmission(du,u,p,t)

In-place creation of vector field for the three group transmission model.
"""
function threegrouptransmission(du,u,p,t)
    #Get state
	S₁,E₁,I₁,R₁,W₁,C₁,C_sero₁,S₂,E₂,I₂,R₂,W₂,C₂,C_sero₂,S₃,E₃,I₃,R₃,W₃,C₃,C_sero₃ = u
    #Get parameters
	R₀,α,γ,ϵ,σ,N₁,N₂,N₃,ω,schooleffect,ct_min1,ct_min2,ct_min3 = p
	#Adjust for schools being open or shut
    Reff = R₀*(t < 24 || t ≥ 319) + R₀*schooleffect*(t≥ 24 && t< 319)
	#Contact rates for each group
    ct1 = KenyaCoVSD.ct_kenya(t,ct_min1)
    ct2 = KenyaCoVSD.ct_kenya(t,ct_min2)
    ct3 = KenyaCoVSD.ct_kenya(t,ct_min3)
    #Force of infection on each group
    λ₁ = γ*Reff*(ct1*ϵ*I₁ + ct2*(1-ϵ)*I₂*(N₁/(N₁ + N₃)) + ct3*(1-ϵ)*I₃*(N₁/(N₁ + N₂)))
    λ₂ = γ*Reff*(ct1*(1-ϵ)*I₁*(N₂/(N₂ + N₃)) + ct2*ϵ*I₂ + ct3*(1-ϵ)*I₃*(N₂/(N₂ + N₁)))
    λ₃ = γ*Reff*(ct1*(1-ϵ)*I₁*(N₃/(N₂ + N₃)) + ct2*(1-ϵ)*I₂*(N₃/(N₃ + N₁)) + ct3*ϵ*I₃)
	#Dynamics
	du[1] = -ct1*(S₁/N₁)*λ₁
	du[2] = ct1*((S₁+σ*W₁)/N₁)*λ₁ - α*E₁
	du[3] = α*E₁ - γ*I₁
	du[4] = γ*I₁ - ω*R₁
	du[5] = ω*R₁ - σ*ct1*(W₁/N₁)*λ₁
	du[6] = ct1*((S₁+σ*W₁)/N₁)*λ₁
    du[7] = ct1*(S₁/N₁)*λ₁

	du[8] = -ct2*(S₂/N₂)*λ₂
	du[9] = ct2*((S₂+σ*W₂)/N₂)*λ₂ - α*E₂
	du[10] = α*E₂ - γ*I₂
	du[11] = γ*I₂ - ω*R₂
	du[12] = ω*R₂ - σ*ct2*(W₂/N₂)*λ₂
	du[13] = ct2*((S₂+σ*W₂)/N₂)*λ₂
    du[14] = ct2*(S₂/N₂)*λ₂

	du[15] = -ct3*(S₃/N₃)*λ₃
	du[16] = ct3*((S₃+σ*W₃)/N₃)*λ₃ - α*E₃
	du[17] = α*E₃ - γ*I₃
	du[18] = γ*I₃ - ω*R₃
	du[19] = ω*R₃ - σ*ct3*(W₃/N₃)*λ₃
	du[20] = ct3*((S₃+σ*W₃)/N₃)*λ₃
    du[21] = ct3*(S₃/N₃)*λ₃
	return nothing
end

u0_threegroup = [4.3e6*0.25 - 100.,100.,0.,0.,0.,0.,0.,
        4.3e6*0.5 - 100.,100.,0.,0.,0.,0.,0.,
        4.3e6*0.25 - 100.,100.,0.,0.,0.,0.,0.]
f_threegroup = ODEFunction(threegrouptransmission;syms = [:S₁,:E₁,:I₁,:R₁,:W₁,:C₁,:C_sero₁,
															:S₂,:E₂,:I₂,:R₂,:W₂,:C₂,:C_sero₂,
															:S₃,:E₃,:I₃,:R₃,:W₃,:C₃,:C_sero₃])
prob_three_group = ODEProblem(f_threegroup,u0_threegroup,(0,365))

"""
    function ll_threegroup_newvariant_infboost(θ,model::KenyaCoVSD.CoVAreaModel,seroreversionrate,cts;σ = 0.16,ω = 1/180,ct_min2 = 0.445)

Log-likelihood of the case and serology data in `model` assuming: 1) Three group dynamics, and 2) that new variant is more transmissible rather than immune evading.
"""
function ll_threegroup_newvariant_infboost(θ,model::KenyaCoVSD.CoVAreaModel,seroreversionrate;σ = 0.16,ω = 1/180)
    @unpack PCR_cases,sero_cases,baseline_sero_array,PCR_array,sero_sensitivity,sero_specificity,N,M_BB,prob,α,γ,relative_testing_rate = model
    @unpack R₀,ϵ,schooleffect,ct_min1,ct_min2,ct_min3,χ₁,χ₂,χ₃,p_test1,p_test2,p_test3,z_pch1,z_pch2,z_pch3,extra_transmissibility,influx_exposed_new_variant,p_test_boost,χ_boost,E₀ = θ

    N₁ = 0.25*N
    N₂ = 0.5*N
    N₃ = 0.25*N
    pch1,pch2,pch3 = softmax([z_pch1,z_pch2,z_pch3])

    #Set variance scalers
    clustering_factor_PCR = 0.5
    M_PCR = 30.
    T = eltype(R₀)
	# if T <: Real
	# 	T = eltype(cts)
	# end
    #Set transmission parameters with daily cts appended on end
    p = convert.(T,[R₀,α,γ,ϵ,σ,N₁,N₂,N₃,ω,schooleffect,ct_min1,ct_min2,ct_min3])
    u0 = convert.(T,[N₁-E₀,E₀,0.,0.,0.,0.,0.,
                    N₂-E₀,E₀,0.,0.,0.,0.,0.,
                    N₃-E₀,E₀,0.,0.,0.,0.,0.])

    #Create new variant introduction
    function new_variant_effect!(integrator)
        integrator.p[1] *= extra_transmissibility
        integrator.u[1] -= influx_exposed_new_variant
        integrator.u[2] += influx_exposed_new_variant
        integrator.u[1 + 7] -= influx_exposed_new_variant
        integrator.u[2 + 7] += influx_exposed_new_variant
        integrator.u[1 + 14] -= influx_exposed_new_variant
        integrator.u[2 + 14] += influx_exposed_new_variant
    end

    janendpoint = (Date(2021,1,30) - Date(2020,2,20)).value
	aprilendpoint = (Date(2021,4,30) - Date(2020,2,20)).value
    variant_cb = PresetTimeCallback([janendpoint],new_variant_effect!)

    #Sero-waning
    sero_array = vcat(baseline_sero_array[1:30],[(1-seroreversionrate)^k for k in 1:length(baseline_sero_array[31:end])])

    #Solve for daily incidence by age (not currently broken down by severity of disease)
    LL = T(0.)
    # try
        sol = solve(prob, BS3();tspan = (0,aprilendpoint),reltol = 1e-3,
					callback = variant_cb,
					u0=u0,
					p=p,
					saveat = 1,
					verbose = false,
                    isoutofdomain=(u,p,t) -> any(x -> x < 0, u))

        ι₁ = diff(sol[:C₁])
        ι₂ = diff(sol[:C₂])
        ι₃ = diff(sol[:C₃])
        ι_sero₁ = diff(sol[:C_sero₁])
        ι_sero₂ = diff(sol[:C_sero₂])
        ι_sero₃ = diff(sol[:C_sero₃])

        PCR₁ = KenyaCoVSD.simple_conv(ι₁,PCR_array)
        PCR₂ = KenyaCoVSD.simple_conv(ι₂,PCR_array)
        PCR₃ = KenyaCoVSD.simple_conv(ι₃,PCR_array)
        weighted_PCR_ptest = (p_test1.*PCR₁ .+ p_test2.*PCR₂ .+ p_test3.*PCR₃).*1e-4
        weighted_prop_PCR = pch1.*(χ₁.*PCR₁./((χ₁-1).*PCR₁ .+ N₁)) .+ pch2.*(χ₂.*PCR₂./((χ₂-1).*PCR₂ .+ N₂)) .+ pch3.*(χ₃.*PCR₃./((χ₃-1).*PCR₃ .+ N₃))
        boosted_weighted_prop_PCR = pch1.*(χ_boost.*χ₁.*PCR₁./((χ_boost.*χ₁-1).*PCR₁ .+ N₁)) .+ pch2.*(χ_boost.*χ₂.*PCR₂./((χ_boost.*χ₂-1).*PCR₂ .+ N₂)) .+ pch3.*(χ_boost.*χ₃.*PCR₃./((χ_boost.*χ₃-1).*PCR₃ .+ N₃))

    	sero = sero_sensitivity.*KenyaCoVSD.simple_conv(ι_sero₁.+ι_sero₂.+ι_sero₃,sero_array)

        #Calculate log-likelihood for PCR testing
        for t in 55:janendpoint
            #Convert from μ,α parameterisation to p,r parameterisation for using Neg. binomial model
            μ = relative_testing_rate[t]*weighted_PCR_ptest[t] + 0.001 #Covert approximately from total probability of detection in % to daily probability of detection
            σ² = μ + clustering_factor_PCR*μ^2
            p_negbin = 1 - (clustering_factor_PCR*μ^2/σ²)
            r_negbin  = 1/clustering_factor_PCR
            LL += logpdf(NegativeBinomial(r_negbin,p_negbin),PCR_cases[t,1])#likelihood contribution from PCR testing --- positive case detection

            if PCR_cases[t,2] >= 0 #Negative tests available
                p_PCR_pred = weighted_prop_PCR[t] + 0.001
                #Convert from p_PCR,inv_M_PCR parameterisation of Beta-Binomial to standard α,β parameterisation
                LL += logpdf(BetaBinomial(PCR_cases[t,2],p_PCR_pred*M_PCR,(1-p_PCR_pred)*M_PCR),PCR_cases[t,1])#likelihood contribution from PCR testing --- proportion postive
            end
        end

        #Calculate log-likelihood for PCR testing after new variant introduction
        for t in (janendpoint+1):min(size(PCR_cases,1),size(PCR₁,1))
            #Convert from μ,α parameterisation to p,r parameterisation for using Neg. binomial model
            μ = p_test_boost*relative_testing_rate[t]*weighted_PCR_ptest[t] + 0.001 #Covert approximately from total probability of detection in % to daily probability of detection
            σ² = μ + clustering_factor_PCR*μ^2
            p_negbin = 1 - (clustering_factor_PCR*μ^2/σ²)
            r_negbin  = 1/clustering_factor_PCR
            LL += logpdf(NegativeBinomial(r_negbin,p_negbin),PCR_cases[t,1])#likelihood contribution from PCR testing --- positive case detection

            if PCR_cases[t,2] >= 0 #Negative tests available
                p_PCR_pred = boosted_weighted_prop_PCR[t] + 0.001
                #Convert from p_PCR,inv_M_PCR parameterisation of Beta-Binomial to standard α,β parameterisation
                LL += logpdf(BetaBinomial(PCR_cases[t,2],p_PCR_pred*M_PCR,(1-p_PCR_pred)*M_PCR),PCR_cases[t,1])#likelihood contribution from PCR testing that day in that age group
            end
        end
        #Calculate log-likelihood contribution from serological testing
        for t in 1:min(size(sero_cases,1),size(sero,1))
            #Convert from the p_hat,M_BB parameterisation of the Betabinomial distribution to the α_BB,β_BB parameterisation
            #Overall serological testing
            if sero_cases[t,1,1] + sero_cases[t,2,1] > 0
                p_sero_pred  = sero[t]/N
                p_hat  = p_sero_pred + (1-p_sero_pred)*(1-sero_specificity)
                LL += logpdf(BetaBinomial(sero_cases[t,1,1] + sero_cases[t,2,1],M_BB*p_hat,M_BB*(1-p_hat)),sero_cases[t,1,1])#Likelihood contribution from sero testing
            end

        end
    # catch
    #     LL = T(-Inf)
    # end
    return LL::T
end

function threegroup_priors(θ)
	@unpack R₀,ϵ,schooleffect,ct_min1,ct_min2,ct_min3,χ₁,χ₂,χ₃,p_test1,p_test2,p_test3,z_pch1,z_pch2,z_pch3,extra_transmissibility,influx_exposed_new_variant,p_test_boost,χ_boost,E₀ = θ
    LP = 0.
    LP += logpdf(Gamma(3,2.5/3),R₀)
    LP += logpdf(Beta(7,3),ϵ)
    LP += logpdf(Beta(8,2),schooleffect)
    LP += logpdf(Beta(7,3),ct_min1)
    LP += logpdf(Beta(7,3),ct_min2)
    LP += logpdf(Beta(55,45),ct_min3)
    LP += logpdf(Gamma(3,4/3),χ₁)
    LP += logpdf(Gamma(3,4/3),χ₂)
    LP += logpdf(Gamma(3,1.5/3),χ₃)
    LP += logpdf(Gamma(3,1/3),p_test1)#In the region of 1% detection rate
    LP += logpdf(Gamma(3,1/3),p_test2)#In the region of 1% detection rate
    LP += logpdf(Gamma(3,5/3),p_test3)#In the region of 5% detection rate
    LP += logpdf(Normal(0,1.0),z_pch1)
	LP += logpdf(Normal(0,1.0),z_pch2)
	LP += logpdf(Normal(0,1.0),z_pch3)
    LP += logpdf(Gamma(10,1.5/10),extra_transmissibility)
    LP += logpdf(Gamma(5,100/5),influx_exposed_new_variant)
	LP += logpdf(Normal(0,0.1),log(χ_boost))
	LP += logpdf(Normal(0,0.1),log(p_test_boost))
	LP += logpdf(Gamma(3,100/3),E₀)
    return LP
end

## Create three group model
nai_three_group = deepcopy(basic_nai_model)
nai_three_group.prob = prob_three_group
nai_three_group.log_likelihood = ll_threegroup_newvariant_infboost
nai_three_group.log_priors = threegroup_priors

#Define variable transformation for three group model
trans_three_groups = as((R₀ = as(Real, 0.0, 7.0),
        ϵ = as(Real,0.5,0.999),
        schooleffect = as(Real,0.25,1.0),
        ct_min1 = as(Real, 0.1, 0.999),
        ct_min2 = as(Real, 0.1, 0.999),
        ct_min3 = as(Real, 0.1, 0.999),
        χ₁=as(Real, 0.1, 30.0),χ₂=as(Real, 0.1, 30.0),χ₃=as(Real, 0.1, 30.0),
        p_test1=as(Real, 0.0, 50.0),p_test2=as(Real, 0.0, 50.0),p_test3=as(Real, 0.0, 50.0),
        z_pch1 = asℝ,z_pch2 = asℝ,z_pch3 = asℝ,
        extra_transmissibility = as(Real, 0.5, 2.5),
        influx_exposed_new_variant = as(Real, 0.0, 10e3),
        p_test_boost = as(Real, 0.5, 3.0),
        χ_boost= as(Real, 0.5, 3.0),
        E₀=as(Real, 0.0, 10e3)))

D = Diagonal(0.1*ones(TransformVariables.dimension(trans_three_groups)))


# import Particles package
using Particles

# parameters
Δt = 0.1    # time step size
nt = 100    # number of time steps
ns = 5      # number of samples

# random initial conditions
x₀ = [ [rand() .* 2π, rand() .* 2 .- 1] for i in 1:ns ]

# vector field
function pendulum!(ẋ, x)
    ẋ[1] = x[2]
    ẋ[2] = - sin(x[1])
end

# create an Equation instance
equ = Equation(pendulum!, x₀)

# create Simulation instance
sim = Simulation(equ, ExplicitEuler, Δt, nt)

# run simulation
x = sim()

# load Plots package
# using Plots

# set plot ranges
# xlim = (-π, +π)
# ylim = (-3, +3)

# plot energy contour lines
# plt = contour(
#     LinRange(xlim..., 100),
#     LinRange(ylim..., 100),
#     (x,v) -> v^2 / 2  + 1 - cos(x),
#     xlim = xlim,
#     ylim = ylim,
#     title = "Pendulum",
#     legend = false,
#     size = (800, 600)
# )

# plot solutions
# for i in axes(x,2)
#     scatter!(plt, (x[1,i,:] .+ π) .% 2π .- π, x[2,i,:], marker=5)
# end

# save figure to file
# savefig("pendulum.png")


""" [Siebesma2003](@cite) """
Bomex_θ_liq_ice(::Type{FT}) where {FT} = z -> if z <= 520.0
        FT(298.7)
    elseif z > 520.0 && z <= 1480.0
        FT(298.7) + (z - 520) * (FT(302.4) - FT(298.7)) / (1480 - 520)
    elseif z > 1480.0 && z <= 2000
        FT(302.4) + (z - 1480) * (FT(308.2) - FT(302.4)) / (2000 - 1480)
    elseif z > 2000.0
        FT(308.2) + (z - 2000) * (FT(311.85) - FT(308.2)) / (3000 - 2000)
    else
        FT(0)
    end
""" [Siebesma2003](@cite) """
Bomex_q_tot(::Type{FT}) where {FT} = z -> if z <= 520
        (FT(17.0) + z * (FT(16.3) - 17) / 520) / 1000
    elseif z > 520.0 && z <= 1480.0
        (FT(16.3) + (z - 520) * (FT(10.7) - FT(16.3)) / (1480 - 520)) / 1000
    elseif z > 1480.0 && z <= 2000.0
        (FT(10.7) + (z - 1480) * (FT(4.2) - FT(10.7)) / (2000 - 1480)) / 1000
    elseif z > 2000.0
        (FT(4.2) + (z - 2000) * (3 - FT(4.2)) / (3000 - 2000)) / 1000
    end
""" [Siebesma2003](@cite) """
Bomex_u(::Type{FT}) where {FT} = z -> if z <= 700.0
        FT(-8.75)
    elseif z > 700.0
        FT(-8.75) + (z - 700) * (FT(-4.61) - FT(-8.75)) / (3000 - 700)
    end
""" [Siebesma2003](@cite) """
Bomex_tke(::Type{FT}) where {FT} = z -> if (z <= 2500.0)
        FT(1) - z / 3000
    else
        FT(0)
    end

# Geostrophic velocity profiles
""" :( """
Bomex_geostrophic_u(::Type{FT}) where {FT} = z -> -10 + FT(1.8e-3) * z
""" :( """
Bomex_geostrophic_v(::Type{FT}) where {FT} = z -> FT(0)

""" dTdt(Π, z) """
Bomex_dTdt(::Type{FT}) where {FT} = (Π, z) -> if z <= 1500.0
        (-2 / (3600 * 24)) * Π
    else
        (-2 / (3600 * 24) + (z - 1500) * (0 - -2 / (3600 * 24)) / (3000 - 1500)) * Π
    end

# Large-scale drying
""" :( """
Bomex_dqtdt(::Type{FT}) where {FT} = z -> if z <= 300.0
        FT(-1.2e-8)   #kg/(kg * s)
    elseif z > 300.0 && z <= 500.0
        FT(-1.2e-8) + (z - 300) * (0 - FT(-1.2e-8)) / (500 - 300) #kg/(kg * s)
    else
        FT(0)
    end

# Large scale subsidence
""" :( """
Bomex_subsidence(::Type{FT}) where {FT} = z -> if z <= 1500.0
        FT(0) + z * (FT(-0.65) / 100 - 0) / (1500 - 0)
    elseif z > 1500.0 && z <= 2100.0
        FT(-0.65) / 100 + (z - 1500) * (0 - FT(-0.65) / 100) / (2100 - 1500)
    else
        FT(0)
    end


using ConstrainedDynamics
using ConstrainedDynamics: skewplusdiag

function dynTvel(vars)
    ezg = [0; 0; -9.81]
    Δt = 0.01
    m = 1.

    vc1 = vars[1:3]
    vc2 = vars[4:6]

    return m * ((vc2 - vc1) / Δt + ezg)
end

function dynRvel(vars)
    Δt = 0.01

    J = diagm(ones(3))
    ωc1 = vars[1:3]
    ωc2 = vars[4:6]
    sq2 = sqrt(4 / Δt^2 - ωc2' * ωc2)
    sq1 = sqrt(4 / Δt^2 - ωc1' * ωc1)

    return skewplusdiag(ωc2, sq2) * (J * ωc2) - skewplusdiag(ωc1, sq1) * (J * ωc1)
end


using LinearAlgebra, DelimitedFiles, StatsBase

crabs = vec(readdlm("data_7.dat", ',', Int))

max_pos = maximum(crabs)
min_pos = minimum(crabs)
dists = Int[]
for p in min_pos:max_pos
    dist = broadcast(abs, crabs .- p)
    dist_sum = sum(dist)
    append!(dists, dist_sum)
end

println(minimum(dists))


# likelihood

logl(w, Ldx) = sum(log(Ldx*w))
dlogl(w, Ldx) = sum(Ldx./(Ldx*w), 1) |> vec

# entropy

repeatweights(w, zmult::Int) = repmat(w/zmult, zmult)
repeatweights(w, L::Matrix)  = repeatweights(w, Int(size(L,1) / length(w)))

# Monte Carlo approximation to the z-entropy
function hz(w, L, wz=repeatweights(w, L))
  rhoz = L * w           # \Int L(z|x) * pi(x) dx
  l = 0.
  for (r,w) in zip(rhoz, wz)
    r == 0 && continue
    l -= log(r)*w
  end
  l
end

# derivative of the MC approximation hz
function dhz(w, L, wz=repeatweights(w, L))
  nz, nx  = size(L)
  zmult   = Int(nz/nx)

  rhoz    = L * w
  factors = wz ./ rhoz

  d = zeros(w)
  @Threads.threads for k = 1:nx
    @simd for i = 1:nz
      @inbounds d[k] -= L[i,k] * factors[i]
    end
    @simd for m = 0:zmult-1
      @inbounds d[k] -= log(rhoz[k+nx*m]) / zmult
    end
  end
  d
end


"""
`sgf(yr,r,g)`: State Government Finance data.
```math
\\alpha_{yr,r,g}^{sgf} = \\dfrac{\\bar{sgf}_{yr,r,g}}{\\sum_{r'} \\bar{sgf}_{yr,r',g}}
```
Note: D.C. is not included in the original data set.
Assume its SGF data is equal to Maryland's.
"""
function share_sgf!(d::Dict)
    println("  Calculating sgf(yr,r,g), regional shares of State Government Finance data")
    d[:sgf] = d[:sgf] / transform_over(d[:sgf], :r)
    verify_over(d[:sgf], :r) !== true && @error("SGF shares don't sum to 1.")
    return dropnan!(d[:sgf])
end


using Combinatorics, Statistics, DataFrames, CSV

data = CSV.read("../data/fertilizer.csv")
control = data.Control
fertilizer = data.FertilizerX

subGroups = collect(combinations([control;fertilizer],10))

meanFert = mean(fertilizer)
pVal = sum([mean(i) >= meanFert for i in subGroups])/length(subGroups)
println("p-value = ", pVal)


using FFTW

export OneDPoisson

"""
    OneDPoisson

Derived type to solve the Poisson equation on 1d regular
cartesian mesh with periodic boundary conditions
- kx   : wave number in x
- nc_x : cells number in x
- dx   : x step size
- rht  : fft(rho)
"""
struct OneDPoisson

    grid::OneDGrid
    modes::Array{Float64,1}
    rht::Array{ComplexF64,1}

    function OneDPoisson(grid::OneDGrid)

        nx = grid.nx
        rht = zeros(ComplexF64, nx)
        k =  2π / grid.dimx
        modes = [1.0; k .* vcat(1:nx÷2-1,-nx÷2:-1)...]

        new(grid, modes, rht)

    end

end

export compute_e_from_rho!

"""
    compute_e_from_rho!( ex, poisson, rho )

computes electric field `e` from `rho` by solving Poisson equation.

"""
function compute_e_from_rho!(ex, poisson::OneDPoisson, rho)

   poisson.rht .= fft(rho) ./ poisson.modes
   ex .= real(ifft(-1im .* poisson.rht ))

end

export OneDPoissonPIC

"""
    OneDPoissonPIC( poisson, kernel )

- kernel : Kernel smoother taking care of charge deposition and field evaluation
- poisson : Poisson solver
- rho : Coefficients of expansion of rho (MPI global version)
- efield : Coefficients of expansion of electric field
- phi : Coefficients of expansion of potential
"""
struct OneDPoissonPIC

    ndofs::Int
    kernel::ParticleMeshCoupling2D
    poisson::OneDPoisson
    rho::Vector{Float64}
    efield::Vector{Vector{Float64}}
    phi::Vector{Float64}

    function OneDPoissonPIC(poisson::OneDPoisson, kernel::ParticleMeshCoupling1D)

        ndofs = poisson.grid.nx

        rho = zeros(ndofs)
        efield = [zeros(ndofs), zeros(ndofs)]

        new(ndofs, kernel, poisson, rho, efield)

    end

end

"""
    add_charge!(pic, position, marker_charge)

Add charge from one particle
- self : Pic Poisson solver object
- position : Position of the particle
- marker_charge : Particle weight times charge
"""
function add_charge!(pic::OneDPoissonPIC, position::Float64, marker_charge)

    GEMPIC.add_charge!(pic.rho, pic.kernel, position, marker_charge)

end

"""
    evaluate_rho!(pic, position)

Evaluate charge density at rho at one position
- self : Pic Poisson solver object
- position : Position of the particle
- func_value : Value of rho at given position
"""
function evaluate_rho!(pic :: OneDPoissonPIC, position)

    GEMPIC.evaluate!(kernel, position, pic.rho)

end

export OneDSplittingOperator

"""
Operator splitting type for 2d2v Vlasov-Poisson
- pic :: PIC poisson solver
- pg :: Particle group
"""
struct OneDSplittingOperator

    pic::OneDPoissonPIC
    pg::ParticleGroup{1,1}

end

"""
    operator_t(split, dt)
Push x
- split :: time splitting object
- dt   :: time step
"""
function operator_t!(split::OneDSplittingOperator, dt)

    for i_part = 1:split.pg.n_particles
        split.pg.array[1, i_part] += dt * split.pg.array[2, i_part]
    end

end

function operator_v!(split::OneDSplittingOperator, dt)

    qm = split.pg.q_over_m

    for i_part = 1:split.pg.n_particles

        # Evaluate efields at particle position

        xi = split.pg.array[1, i_part]

        ex = GEMPIC.evaluate(split.pic.kernel, xi, split.pic.efield)

        vx_new = split.pg.array[1, i_part]

        vx_new += dt * qm * ex

        split.pg.array[2, i_part] = vx_new

    end

end

export charge_deposition!

function charge_deposition!(split::OneDSplittingOperator)

    fill!(split.pic.rho_dofs, 0.0)
    for i_part = 1:split.pg.n_particles
        xp = split.pg.array[1, i_part]
        wp = split.pg.array[3, i_part]
        GEMPIC.add_charge!(split.pic.rho_dofs, split.pic.kernel, xp, wp)
    end

end

export solve_fields!

"""
    solve_fields!( pic )

Solve efields from rho
"""
function solve_fields!(pic :: OneDPoissonPIC)

    compute_e_from_rho!(pic.efield[1], pic.poisson, pic.rho)

end

"""
Solve Poisson's equation for the electric field
"""
solve_fields!(split::OneDSplittingOperator) = solve_fields!(split.pic)

export strang_splitting!

"""
    strang_splitting!(split, dt)

Strang splitting
- split :: time splitting object
- dt   :: time step
"""
function strang_splitting!(split::OneDSplittingOperator, dt)

    operator_t!(split, 0.5dt)
    charge_deposition!(split)
    solve_fields!(split)
    operator_v!(split, dt)
    operator_t!(split, 0.5dt)

end

export compute_field_energy

"""
    compute_field_energy

Compute the squared l2 norm of electric field
"""
compute_field_energy(pic::OneDPoissonPIC, component::Int) =
    sum(pic.efield[component] .^ 2) * pic.poisson.grid.dx


#!/usr/sbin/julia
import Primes: nextprime

struct Data
  g::BigInt
  q::BigInt
  sks::Array{Array{BigInt}}
  pks::Array{BigInt}
  secret::Array{UInt8}
end

function create_key(len::Int, data::Data)
  sk::Array{BigInt} = []
  pk::BigInt = 1

  if length(data.pks) ≥ 10
    println("Enough of key creation!")
    return
  end

  for i in 1:len
    push!(sk, rand(1:data.q))
    pk = powermod(data.g, sk[i], data.q) * pk % data.q
  end
  push!(data.sks, sk)
  push!(data.pks, pk)
  println("Key created sucessfully!")
end

function create_secret(len::Int, data::Data)
  println("Tell me your secret: ")
  readbytes!(stdin, data.secret, len)
end

function show_data(data::Data)
  println("g = ", data.g)
  println("q = ", data.q)
  println("secret = ", data.secret)
  println("pks = ", data.pks)
end

function encrypt(flag::String, data::Data)
  i::Int = 0
  c::BigInt = 1
  d::BigInt = 1
  r::BigInt = 1
  s::Int = 0

  if length(data.pks) < 4
    println("Error: You don't have enough keys to encrypt.")
    return
  end

  if length(data.secret) ≠ length(flag)
    println("Error: Invalid secret. Make sure you create a secret first!")
    return
  end

  println("Which char index you want to encrypt?")
  try i = parse(Int, readline()) catch e i = -1 end
  if !(1 ≤ i ≤ length(flag))
    println("Error: Invalid char index!")
    return
  end

  for ki in 1:2:(length(data.pks)-1)
    c = powermod(data.pks[ki], data.sks[ki + 1][i], data.q) * c % data.q
    d = powermod(data.pks[ki + 1], data.sks[ki][i], data.q) * d % data.q
  end

  r = rand(1:data.q)
  s = data.secret[i] ⊻ Int(flag[i])
  c = powermod(data.g, s + r, data.q) * c % data.q
  d = powermod(data.g, r, data.q) * d % data.q

  println("Here we go: ", (c, d))
end

function menu()::Int
  println()
  println("1- Create Key")
  println("2- Create Secret")
  println("3- Show data")
  println("4- Encrypt Flag")
  println("5- Exit")
  println()
  option::Int = 0
  try option = parse(Int, readline()) catch e return -1 end
  return option
end

function main()
  flag::String = strip(open(f->read(f, String), "flag.txt"))
  data::Data = Data(2, nextprime(rand(0: big(2)^1024)), [], [], [])
  @assert length(flag) == 40

  println("\nWelcome! What's your plan for today?")
  while true
    option::Int = menu()
    if option == 1
      create_key(length(flag), data)
    elseif option == 2
      create_secret(length(flag), data)
    elseif option == 3
      show_data(data)
    elseif option == 4
      encrypt(flag, data)
    elseif option == 5
      println("Bye!")
      return
    else
      println("Invalid option!")
    end
  end
end

main()


function sudoku(n, start, ::Val{:raw})
    N = n^2
    d = domain(1:N)

    m = model(;kind=:sudoku)

    # Add variables
    if isnothing(start)
        foreach(_ -> variable!(m, d), 1:(N^2))
    else
        foreach(((x, v),) -> variable!(m, 1 ≤ v ≤ N ? domain(v) : d), pairs(start))
    end

    e = (x; param=nothing, dom_size=N) -> error_f(
        usual_constraints[:all_different])(x; param=param, dom_size=dom_size
    )

    # Add constraints: line, columns; blocks
    foreach(i -> constraint!(m, e, (i * N + 1):((i + 1) * N)), 0:(N - 1))
    foreach(i -> constraint!(m, e, [j * N + i for j in 0:(N - 1)]), 1:N)

    for i in 0:(n - 1)
        for j in 0:(n - 1)
            vars = Vector{Int}()
            for k in 1:n
                for l in 0:(n - 1)
                    push!(vars, (j * n + l) * N + i * n + k)
                end
            end
            constraint!(m, e, vars)
        end
    end

    return m
end

function sudoku(n, start, ::Val{:MOI})
    N = n^2
    m = Optimizer()
    MOI.add_variables(m, N^2)

    # Add domain to variables
    # TODO: make it accept starting solutions
    foreach(i -> MOI.add_constraint(m, VI(i), DiscreteSet(1:N)), 1:N^2)

    # Add constraints: line, columns; blocks
    foreach(i -> MOI.add_constraint(m, VOV(map(VI, (i * N + 1):((i + 1) * N))),
        MOIAllDifferent(N)), 0:(N - 1))
    foreach(i -> MOI.add_constraint(m, VOV(map(VI, [j * N + i for j in 0:(N - 1)])),
            MOIAllDifferent(N)), 1:N)

    for i in 0:(n - 1)
        for j in 0:(n - 1)
            vars = Vector{Int}()
            for k in 1:n
                for l in 0:(n - 1)
                    push!(vars, (j * n + l) * N + i * n + k)
                end
            end
            MOI.add_constraint(m, VOV(map(VI, vars)), MOIAllDifferent(N))
        end
    end

    return m
end

function sudoku(n, start, ::Val{:JuMP})
    N = n^2
    m = JuMP.Model(CBLS.Optimizer)

    @variable(m, 1 ≤ X[1:N, 1:N] ≤ N, Int)
    if !isnothing(start)
        for i in 1:N, j in 1:N
            v_ij = start[i,j]
            if 1 ≤ v_ij ≤ N
                @constraint(m, X[i,j] == v_ij)
            end
        end
    end

    for i in 1:N
        @constraint(m, X[i,:] in AllDifferent()) # rows
        @constraint(m, X[:,i] in AllDifferent()) # columns
    end
    for i in 0:(n-1), j in 0:(n-1)
        @constraint(m, vec(X[(i*n+1):(n*(i+1)), (j*n+1):(n*(j+1))]) in AllDifferent()) # blocks
    end
    return m, X
end

"""
    sudoku(n; start= Dictionary{Int, Int}(), modeler = :JuMP)
Create a model for the sudoku problem of domain `1:n²` with optional starting values. The `modeler` argument accepts :raw, :MOI, and :JuMP (default), which refer respectively to the solver internal model, the MathOptInterface model, and the JuMP model.

```julia
# Construct a JuMP model `m` and its associated matrix `grid` for sudoku 9×9
m, grid = sudoku(3)

# Same with a starting instance
instance = [
    9  3  0  0  0  0  0  4  0
    0  0  0  0  4  2  0  9  0
    8  0  0  1  9  6  7  0  0
    0  0  0  4  7  0  0  0  0
    0  2  0  0  0  0  0  6  0
    0  0  0  0  2  3  0  0  0
    0  0  8  5  3  1  0  0  2
    0  9  0  2  8  0  0  0  0
    0  7  0  0  0  0  0  5  3
]
m, grid = sudoku(3, start = instance)

# Run the solver
optimize!(m)

# Retrieve and display the values
solution = value.(grid)
display(solution, Val(:sudoku))
```
"""
sudoku(n; start=nothing, modeler = :JuMP) = sudoku(n, start, Val(modeler))

@doc raw"""
```julia
mutable struct SudokuInstance{T <: Integer} <: AbstractMatrix{T}
```
A `struct` for SudokuInstances, which is a subtype of `AbstractMatrix`.
---
```julia
SudokuInstance(A::AbstractMatrix{T})
SudokuInstance(::Type{T}, n::Int) # fill in blank sudoku of type T
SudokuInstance(n::Int) # fill in blank sudoku of type Int
SudokuInstance(::Type{T}) # fill in "standard" 9×9 sudoku of type T
SudokuInstance() # fill in "standard" 9×9 sudoku of type Int
SudokuInstance(n::Int, P::Pair{Tuple{Int, Int}, T}...) where {T <: Integer} # construct a sudoku given pairs of coordinates and values
SudokuInstance(P::Pair{Tuple{Int, Int}, T}...) # again, default to 9×9 sudoku, constructing given pairs
```
Constructor functions for the `SudokuInstance` `struct`.
"""
mutable struct SudokuInstance{T <: Integer} <: AbstractMatrix{T}
    A::AbstractMatrix{T} where {T <: Integer}

    function SudokuInstance(A::AbstractMatrix{T}) where {T <: Integer}
        size(A, 1) == size(A, 2) || throw(error("Sodokus must be square; received matrix of size $(size(A, 1))×$(size(A, 2))."))
        isequal(sqrt(size(A, 1)), isqrt(size(A, 1))) || throw(error("SudokuInstances must be able to split into equal boxes (e.g., a 9×9 SudokuInstance has three 3×3 squares).  Size given is $(size(A, 1))×$(size(A, 2))."))
        new{T}(A)
    end
    # # fill in blank sudoku if needed
    # SudokuInstance(::Type{T}, n::Int) where {T <: Integer} = new{T}(fill(zero(T), n, n))
    # SudokuInstance(n::Int) = new{Int}(SudokuInstance(Int, n))
    # # Use "standard" 9×9 if no size provided
    # SudokuInstance(::Type{T}) where {T <: Integer} = new{T}(SudokuInstance(T, 9))
    # SudokuInstance() = new{Int}(SudokuInstance(9))
    # # Construct a sudoku given coordinates and values
    # function SudokuInstance(n::Int, P::Pair{Tuple{Int,Int},T}...) where {T <: Integer}
    #     A = zeros(T, n, n)
    #     for (i, v) in P
    #         A[i...] = v
    #     end
    #     new{T}(A)
    # end
    # # again, default to 9×9
    # SudokuInstance(P::Pair{Tuple{Int,Int},T}...) where {T <: Integer} = new{T}(SudokuInstance(9, P...))
end

"""
    SudokuInstance(X::Dictionary)

Construct a `SudokuInstance` with the values `X` of a solver as input.
"""
function SudokuInstance(X::Dictionary)
    n = isqrt(length(X))
    A = zeros(Int, n, n)
    for (k,v) in enumerate(Base.Iterators.partition(X, n))
        A[k,:] = v
    end
    return SudokuInstance(A)
end

SudokuInstance(X::Matrix{Float64}) = SudokuInstance(map(Int, X))
# # abstract array interface for SudokuInstance struct
"""
    Base.size(S::SudokuInstance)

Extends `Base.size` for `SudokuInstance`.
"""
Base.size(S::SudokuInstance) = size(S.A)

"""
    Base.getindex(S::SudokuInstance, i::Int)
    Base.getindex(S::SudokuInstance, I::Vararg{Int,N}) where {N}

Extends `Base.getindex` for `SudokuInstance`.
"""
# Base.getindex(S::SudokuInstance, i::Int) = getindex(S.A, i)
Base.getindex(S::SudokuInstance, I::Vararg{Int,N}) where {N} = getindex(S.A, I...)

"""
    Base.setindex!(S::SudokuInstance, v, i::Int)
    Base.setindex!(S::SudokuInstance, v, I::Vararg{Int,N})

Extends `Base.setindex!` for `SudokuInstance`.
"""
# Base.setindex!(S::SudokuInstance, v, i) = setindex!(S.A, v, i)
# Base.setindex!(S::SudokuInstance, v, I::Vararg) = setindex!(S.A, v, I...)

const _rules = Dict(
    :up_right_corner => '┐',
    :up_left_corner => '┌',
    :bottom_left_corner => '└',
    :bottom_right_corner => '┘',
    :up_intersection => '┬',
    :left_intersection => '├',
    :right_intersection => '┤',
    :middle_intersection => '┼',
    :bottom_intersection => '┴',
    :column => '│',
    :row => '─',
    :blank => '⋅',  # this is the character used for 0s in a SudokuInstance puzzle
)

"""
    _format_val(a)

Format an integer `a` into a string for SudokuInstance.
"""
_format_val(a) = iszero(a) ? _rules[:blank] : string(a)

"""
    _format_line_segment(r, col_pos, M)

Format line segment of a sudoku grid.
"""
function _format_line_segment(r, col_pos, M)
    sep_length = length(r)

    line = string()
    for k in axes(r, 1)
        n_spaces = 1
        Δ = maximum((ndigits(i) for i in M[:, (col_pos * sep_length) + k])) - ndigits(r[k])
        if Δ ≥ 0
            n_spaces = Δ + 1
        end
        line *= repeat(' ', n_spaces) * _format_val(r[k])
    end

    return line * ' ' * _rules[:column]
end

"""
    _format_line(r, M)

Format line of a sudoku grid.
"""
function _format_line(r, M)
    sep_length = isqrt(length(r))

    line = _rules[:column]
    for i in 1:sep_length
        abs_sep_pos = sep_length * i
        line *= _format_line_segment(r[(abs_sep_pos - sep_length + 1):abs_sep_pos], i - 1, M)
    end

    return line
end

"""
    _get_sep_line(s, pos_row, M)

Return a line separator.
"""
function _get_sep_line(s, pos_row, M)
    sep_length = isqrt(s)

    # deal with left-most edges
    sep_line = string()
    if pos_row == 1
        sep_line *= _rules[:up_left_corner]
    elseif mod(pos_row, sep_length) == 0
        if pos_row == s
            sep_line *= _rules[:bottom_left_corner]
        else
            sep_line *= _rules[:left_intersection]
        end
    end

    # rest of row seperator; TODO: make less convoluted.  ATM it works, but I think it can be simplified.
    for pos_col in 1:s
        sep_line *= repeat(_rules[:row], maximum((ndigits(i) for i in M[:, pos_col])) + 1)
        if mod(pos_col, sep_length) == 0
            sep_line *= _rules[:row]
            if pos_col == s
                if pos_row == 1
                    sep_line *= _rules[:up_right_corner]
                elseif pos_row == s
                    sep_line *= _rules[:bottom_right_corner]
                else
                    sep_line *= _rules[:right_intersection]
                end
            elseif pos_row == 1
                sep_line *= _rules[:up_intersection]
            elseif pos_row == s
                sep_line *= _rules[:bottom_intersection]
            else
                sep_line *= _rules[:middle_intersection]
            end
        end
    end

    return sep_line
end

@doc raw"""
```julia
display(io::IO, S::SudokuInstance)
display(S::SudokuInstance) # default to stdout
```
Displays an ``n\times n`` SudokuInstance.
"""
function Base.display(io, S::SudokuInstance)
    sep_length = isqrt(size(S, 1))
    max_n_digits = maximum((ndigits(i) for i in S))

    println(io, _get_sep_line(size(S, 1), 1, S))
    for (i, r) in enumerate(eachrow(S))
        println(io, _format_line(r, S))
        if iszero(mod(i, sep_length))
            println(io, _get_sep_line(size(S, 1), i, S))
        end
    end

    return nothing
end

"""
    Base.display(S::SudokuInstance)

Extends `Base.display` to `SudokuInstance`.
"""
Base.display(S::SudokuInstance) = display(stdout, S)

"""
    Base.display(X::Dictionary)

Extends `Base.display` to a sudoku configuration.
"""
Base.display(X::Dictionary) = display(SudokuInstance(X))

"""
    Base.display(X, Val(:sudoku))

Extends `Base.display` to a sudoku configuration.
"""
Base.display(X, ::Val{:sudoku}) = display(SudokuInstance(X))


@testset "Generic elliptic curve" begin

  @testset "Constructors" begin
    @test_throws ErrorException EllipticCurve([1])
    @test_throws ErrorException EllipticCurve([1, 2, 3])
    @test_throws ErrorException EllipticCurve([1, 2, 3, 4, 5, 6])
    @test_throws ErrorException EllipticCurve([0, 0])
    @test_throws ErrorException EllipticCurve([0, 0, 0, 0, 0])

    E = @inferred EllipticCurve([1, 2], check = false)
    @test typeof(E) == EllCrv{fmpq}
    @test a_invars(E) == (0, 0, 0, 1, 2)
    @test coefficients(E) == (0, 0, 0, 1, 2)

    E = @inferred EllipticCurve([1, 2, 3, 4, 5])
    @test typeof(E) == EllCrv{fmpq}
    @test a_invars(E) == (1, 2, 3, 4, 5)

    # this is Cremona: 11a2, lmfdb: 11.a1
    E11_a1 = @inferred EllipticCurve([0, -1, 1, -7820, -263580], check = false)

    # this is Cremona: 41a1, lmfdb: 41.a1
    E43_a1 = @inferred EllipticCurve([0, 1, 1, 0, 0])

    E = @inferred EllipticCurve([0, 0, 0, 1, 2])
    @test is_short_weierstrass_model(E)

    E = @inferred EllipticCurve([0, 0, 0, 1, 2], check = false)
    @test is_short_weierstrass_model(E)

    Qx, x = PolynomialRing(FlintQQ, "x")

    f1 = x^3+3*x+5
    g1 = x+2
    E = EllipticCurve(f1, g1)
    f2, g2 = hyperelliptic_polynomials(E)
    @test f1 == f2 && g1 == g2

    E = @inferred EllipticCurve(f1)
    @test E isa EllCrv{fmpq}
    E = @inferred EllipticCurve(f1, check = false)
    @test E isa EllCrv{fmpq}

    E = EllipticCurve(f1, 1)
    f2, g2 = hyperelliptic_polynomials(E)
    @test f1 == f2 && 1 == g2

    E = EllipticCurve(x^3 + 1, zero(Qx))
    @test E isa EllCrv{fmpq}

    @test_throws ArgumentError EllipticCurve(x^10-21, x^3+5)
    @test_throws ArgumentError EllipticCurve(x^3+3, x^3+5)
    @test_throws ArgumentError EllipticCurve(3*x^3 + 1)

    K, a = NumberField(x^2 - x - 1, "a")
    OK = maximal_order(K)

    E31_1_a1 = @inferred EllipticCurve([K(1), a + 1, a, a, K(0)])
    @test typeof(E31_1_a1) == EllCrv{nf_elem}

    # lmfdb: 116.1-a1
    E116_1_a1 = @inferred EllipticCurve([K(1), K(-1), a, -a, K(0)] )
    @test typeof(E31_1_a1) == EllCrv{nf_elem}

    # short example
    Eshort = @inferred EllipticCurve([4, 0])
    @test typeof(Eshort) == EllCrv{fmpq}
    @test a_invars(Eshort) == (0, 0, 0, 4, 0)
  end

  QQx, = PolynomialRing(QQ, 5)
  F = FractionField(QQx)
  a1, a2, a3, a4, a6 = F.(gens(QQx))
  E = EllipticCurve([a1, a2, a3, a4, a6])
  @test b_invars(E) == (a1^2 + 4*a2, a1*a3 + 2*a4, a3^2 + 4*a6, a1^2*a6 - a1*a3*a4 + a2*a3^2 + 4*a2*a6 - a4^2)
  b2,b4,b6,b8 = b_invars(E)
  c4, c6 = c_invars(E)
  @test (c4, c6) == c_invars(E) # to test caching
  @test (c4, c6) == (a1^4 + 8*a1^2*a2 - 24*a1*a3 + 16*a2^2 - 48*a4, -a1^6 - 12*a1^4*a2 + 36*a1^3*a3 - 48*a1^2*a2^2 + 72*a1^2*a4 + 144*a1*a2*a3 - 64*a2^3 + 288*a2*a4 - 216*a3^2 - 864*a6)
  @test c4 == b2^2 - 24*b4
  @test c6 == -b2^3 + 36*b2*b4 - 216*b6
  d = discriminant(E)
  @test d == -b2^2*b8 - 8*b4^3 - 27*b6^2 + 9*b2*b4*b6
  @test 1728*d == c4^3 - c6^2

  # base change
  E = EllipticCurve([1, 0, 0, 0, 1])
  K, = quadratic_field(-1)
  EK = @inferred base_change(K, E)
  @test a_invars(EK) == (1, 0, 0, 0, 1)
  @test base_field(EK) === K

  F = GF(2)
  f = x -> divexact(F(numerator(x)), F(denominator(x)))
  EF = base_change(f, E)
  @test a_invars(EF) == (1, 0, 0, 0, 1)
  @test base_field(EF) === F

  # The following curves will be used in later tests
  # Creation of these was tested in previous testset
  E11_a1 = EllipticCurve([0, -1, 1, -7820, -263580], check =false)

  E43_a1 = EllipticCurve([0, 1, 1, 0, 0])

  Qx, x = PolynomialRing(FlintQQ, "x")
  K, a = NumberField(x^2 - x - 1, "a")
  OK = maximal_order(K)

  E31_1_a1 = EllipticCurve([K(1), a + 1, a, a, K(0)])

  E116_1_a1 =EllipticCurve([K(1), K(-1), a, -a, K(0)] )

  Eshort = EllipticCurve([4, 0])

  @testset "Field access" begin
    @test base_field(E11_a1) == FlintQQ
    @test base_field(E43_a1) == FlintQQ
    @test base_field(E31_1_a1) == K
    @test base_field(E116_1_a1) == K
    @test base_field(Eshort) == FlintQQ
  end

  @testset "Point construction" begin
    P = @inferred E43_a1([FlintQQ(-1), FlintQQ(0)])
    @test typeof(P) == EllCrvPt{fmpq}
    @test parent(P) == E43_a1
    @test @inferred is_finite(P)
    @test @inferred !is_infinite(P)

    P = @inferred E43_a1([-1, 0], check = false)
    @test typeof(P) == EllCrvPt{fmpq}
    @test parent(P) == E43_a1
    @test @inferred is_finite(P)
    @test @inferred !is_infinite(P)

    P = @inferred E43_a1([fmpz(-1), fmpz(0)])
    @test typeof(P) == EllCrvPt{fmpq}
    @test parent(P) == E43_a1
    @test @inferred is_finite(P)
    @test @inferred !is_infinite(P)

# the error is/was from doing QQ(K(0)) - which is possible now
#    @test_throws MethodError E43_a1([gen(K), gen(K)])

    @test_throws ErrorException E43_a1([2, 2])

    P = @inferred infinity(E43_a1)
    @test @inferred !is_finite(P)
    @test @inferred is_infinite(P)

    P = @inferred E43_a1([FlintQQ(-1), FlintQQ(0)])

    P = @inferred E116_1_a1([K(0), -K(a)])
    @test typeof(P) == EllCrvPt{nf_elem}
    @test parent(P) == E116_1_a1
    @test @inferred is_finite(P)
    @test @inferred !is_infinite(P)

    P = @inferred infinity(E116_1_a1)
    @test !is_finite(P)
    @test is_infinite(P)

    @test_throws ErrorException E116_1_a1([1, 1])

    P = @inferred Eshort([2, 4], check = false)
    @test @inferred is_finite(P)
    @test typeof(P) == EllCrvPt{fmpq}
    @test parent(P) == Eshort

    E = EllipticCurve(GF(7,2),[1,2,3,4,5])
    L = @inferred points_with_x(E,0)
    @test E([0,5]) in L && E([0, 6]) in L

  end

  @testset "Equation" begin
    E = EllipticCurve( [1, 2, 3, 4, 5])
    Kxy, (x,y) = PolynomialRing(base_field(E), ["x","y"])
    @test y^2 + x*y + 3*y - x^3 - 2*x^2 - 4*x - 5 == @inferred Kxy(equation(E))
  end

  @testset "Discriminant" begin
    @test (2*a + 10)*OK == @inferred (discriminant(E116_1_a1)*OK)
    @test -43 == @inferred discriminant(E43_a1)
    @test -4096 == @inferred discriminant(Eshort)
    E43_a1 = EllipticCurve([0, 1, 1, 0, 0], check = false)
    @test discriminant(E43_a1) == -43
    Eshort = @inferred EllipticCurve([4, 0], check = false)
    @test discriminant(Eshort) == -4096
  end

  @testset "j-invariant" begin
    b = (fmpq(-215055, 58) * a - fmpq(65799, 29))
    @test  b == @inferred j_invariant(E116_1_a1)
    @test fmpq(-4096, 43) == @inferred j_invariant(E43_a1)
    @test 1728 == @inferred j_invariant(Eshort)

    E = @inferred elliptic_curve_from_j_invariant(23)
    @test j_invariant(E) == 23

    E = @inferred elliptic_curve_from_j_invariant(1728)
    @test j_invariant(E) == 1728

    E = @inferred elliptic_curve_from_j_invariant(GF(3,2)(0))
    @test j_invariant(E) == 0

    E = @inferred elliptic_curve_from_j_invariant(GF(2,4)(0))
    @test j_invariant(E) == 0

  end

  @testset "Point aritmetic" begin
    #addition
    P = @inferred E43_a1([FlintQQ(-1), FlintQQ(0)])
    O = infinity(E43_a1)

    @test P == @inferred P + O
    @test E43_a1([2, -4]) == P + P
    @test E43_a1([fmpq(-2, 9), fmpq(1, 27)]) == P + P + P

    P = @inferred E116_1_a1([K(0), -K(a)])
    @test E116_1_a1([1, -1]) == @inferred P + P
    @test E116_1_a1([K(1), -a]) == @inferred P + P + P
    @test infinity(E116_1_a1) == @inferred P + P + P + P + P
    @test P == @inferred P + infinity(E116_1_a1)

    P = @inferred Eshort([2, 4])
    O = @inferred infinity(Eshort)
    @test_throws ErrorException Eshort([-1, -1])

    @test Eshort([0, 0]) == @inferred P + P
    @test P == @inferred O + P

    P = Eshort([2, 4])
    @test Eshort([2, -4]) == @inferred -P
    P = infinity(Eshort)
    @test P == @inferred -P

    P = E43_a1([-1, 0])
    @test E43_a1([-1, -1]) == @inferred -P
    P = infinity(E43_a1)
    @test P == @inferred -P

    # inversion
    P = @inferred E116_1_a1([K(0), -K(a)])
    @test E116_1_a1([0, 0]) == @inferred -P
    P = infinity(E116_1_a1)
    @test P == @inferred -P

    # equality
    P1 = Eshort([2, 4])
    @test @inferred ==(P1, P1)
    P2 = infinity(Eshort)
    @test @inferred ==(P2, P2)
    @test @inferred !==(P2, P1)

    P1 = E43_a1([-1, 0])
    @test @inferred ==(P1, P1)
    P2 = infinity(E43_a1)
    @test @inferred ==(P2, P2)
    @test @inferred !==(P2, P1)

    P1 = E116_1_a1([K(0), -K(a)])
    @test @inferred ==(P1, P1)
    P2 = infinity(E116_1_a1)
    @test @inferred ==(P2, P2)
    @test @inferred !==(P2, P1)

    # scalar multiplication
    P1 = Eshort([2, 4])
    @test Eshort([0, 0]) == @inferred 2*P1
    @test infinity(Eshort) == @inferred 4*P1

    P1 = E43_a1([-1, 0])
    @test E43_a1([fmpq(11, 49), fmpq(-363, 343)]) == @inferred 4*P1

    P1 = E116_1_a1([K(0), -K(a)])
    @test E116_1_a1([K(0), K(0)]) == @inferred 4*P1
    @test infinity(E116_1_a1) == @inferred 5*P1

    #division

    P1 = Eshort([2, 4])
    Q = (2*P1)//2
    @test Q == P1 || Q == -P1

    P1 = E116_1_a1([K(0), -K(a)])
    @test (3*P1)//3 == P1
    @test @inferred 3*(3*P1)//(-3) == -3*P1
    @test_throws ErrorException P1//5
  end
end


# ---
# title: 1380. Lucky Numbers in a Matrix
# id: problem1380
# author: Indigo
# date: 2022-02-18
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/lucky-numbers-in-a-matrix/description/>
# hidden: true
# ---
#
# Given a `m * n` matrix of **distinct** numbers, return all lucky numbers in
# the matrix in **any** order.
#
# A lucky number is an element of the matrix such that it is the minimum element
# in its row and maximum in its column.
#
#
#
# **Example 1:**
#
#
#
#     Input: matrix = [[3,7,8],[9,11,13],[15,16,17]]
#     Output: [15]
#     Explanation: 15 is the only lucky number since it is the minimum in its row and the maximum in its column
#
#
# **Example 2:**
#
#
#
#     Input: matrix = [[1,10,4,2],[9,3,8,7],[15,16,17,12]]
#     Output: [12]
#     Explanation: 12 is the only lucky number since it is the minimum in its row and the maximum in its column.
#
#
# **Example 3:**
#
#
#
#     Input: matrix = [[7,8],[1,2]]
#     Output: [7]
#
#
#
#
# **Constraints:**
#
#   * `m == mat.length`
#   * `n == mat[i].length`
#   * `1 <= n, m <= 50`
#   * `1 <= matrix[i][j] <= 10^5`.
#   * All elements in the matrix are distinct.
#
#
## @lc code=start
using LeetCode

function lucky_numbers(mat::Matrix{Int})::Vector{Int}
    amax = argmin(mat; dims=1)
    amin = argmax(mat; dims=2)
    return mat[intersect(amax, amin)]
end
## @lc code=end


# TaylorSeries.jl
#
# Julia module for handling Taylor series of arbitrary but finite order
#
# utils_Taylor1.jl contains the constructors and methods for 1-variable expansions
#
# utils_TaylorN.jl contains the constructors and methods for N-variable expansions
#
# Last modification: 2014.06.07
#
# Luis Benet & David P. Sanders
# UNAM
#

module TaylorSeries

import Base: zero, one
import Base: convert, promote_rule, promote, eltype, length
import Base: real, imag, conj, ctranspose
import Base: rem, mod, mod2pi
import Base: sqrt, exp, log, sin, cos, tan#, square

abstract AbstractSeries{T<:Number,N} <: Number

include("utils_Taylor1.jl")

include("utils_TaylorN.jl")

## The following routines combine Taylor and TaylorN, so they must appear defining
##   Taylor and TaylorN and some of its functionalities

# infostr
infostr{T<:Number}(a::Taylor{T}) =
    string(a.order, "-order Taylor{", T, "}:\n")
infostr{T<:Number}(a::TaylorN{T}) =
    string(a.order, "-order TaylorN{", T, "} in ", a.numVars, " variables:\n")

# pretty_print
function pretty_print{T<:Number}(a::Taylor{T})
    print( infostr(a) )
    z = zero(T)
    space = " "
    a == z && (print(string( space, z)); return)
    strout = space
    ifirst = true
    for i = 0:a.order
        monom = i==0 ? "" : i==1 ? " * x0" : string(" * x0^", i)
        @inbounds c = a.coeffs[i+1]
        c == z && continue
        cadena = numbr2str(c, ifirst)
        strout = string(strout, cadena, monom, space)
        ifirst = false
    end
    println(strout)
end
function pretty_print{T<:Number}(a::TaylorN{T})
    print( infostr(a) )
    z = zero(T)
    space = " "
    a == z && (println(string( space, z)); return)
    varstring = {}
    for ivar=1:a.numVars
        push!(varstring,string(" * x",ivar))
    end
    strout = string(" ")
    ifirst = true
    iIndices = zeros(Int, a.numVars)
    for pos = 1:length(a.coeffs)
        monom = ""
        iIndices = indicesTable[end][pos]
        if pos>1
            for ivar=1:a.numVars
                powivar = iIndices[ivar]
                if powivar == 1
                    monom = string(monom, varstring[ivar])
                elseif powivar > 1
                    monom = string(monom, varstring[ivar], "^", powivar)
                end
            end
        end
        @inbounds c = a.coeffs[pos]
        c == z && continue
        cadena = numbr2str(c, ifirst)
        strout = string(strout, cadena, monom, space)
        ifirst = false
    end
    println(strout)
end
function pretty_print{T<:Number}(a::Union(Array{Taylor{T},1},Array{TaylorN{T},1}))
    for i=1:length(a)
        pretty_print(a[i])
        println("")
    end
end

# make string from a number; for complex numbers, use
function numbr2str{T<:Real}(zz::T, ifirst::Bool=false)
    plusmin = zz > 0 ? "+ " : "- "
    if ifirst
        plusmin = zz > 0 ? " " : "-"
    end
    return string(plusmin, abs(zz))
end
function numbr2str{T}(zz::Complex{T}, ifirst::Bool=false)
    zT = zero(T)
    zz == zero(Complex{T}) && return zT
    zre, zim = reim(zz)
    cadena = ""
    if zre > zT
        if ifirst
            cadena = string("   ( ", abs(zre)," ")
        else
            cadena = string(" + ( ", abs(zre)," ")
        end
        if zim > zT
            cadena = string(cadena, "+ ", abs(zim), " im )")
        elseif zim < zT
            cadena = string(cadena, "- ", abs(zim), " im )")
        else
            cadena = string(cadena, ")")
        end
    elseif zre < zT
        cadena = string(" - ( ", abs(zre), " ")
        if zim > zT
            cadena = string(cadena, "- ", abs(zim), " im )")
        elseif zim < zT
            cadena = string(cadena, "+ ", abs(zim), " im )")
        else
            cadena = string(cadena, ")")
        end
    else
        if zim > zT
            if ifirst
                cadena = string("   ( ", abs(zim), " im )")
            else
                cadena = string(" + ( ", abs(zim), " im )")
            end
        else
            cadena = string(" - ( ", abs(zim), " im )")
        end
    end
    return cadena
end

# evalTaylor(TaylorN, Array{Taylor,1})
function evalTaylor{T<:Number,S<:Number}(a::TaylorN{T}, vT::Array{Taylor{S},1})
    numVars = NUMVARS[end]
    @assert length(vT) == numVars
    R = promote_type(T,S)
    order = length(vT[1])
    sumaT = Taylor(zero(R), order)
    z = zero(sumaT)
    nCoefTot = sizeTable[end]
    iIndices = zeros(Int, numVars)
    for pos = nCoefTot:-1:1
        @inbounds iIndices[1:end] = indicesTable[end][pos]
        a.coeffs[pos] == zero(T) && continue
        val = Taylor(a.coeffs[pos], order)
        for k = 1:numVars
            @inbounds val = val*(vT[k])^iIndices[k]
        end
        sumaT += val
    end
    return sumaT
end

## Exports to Taylor and TaylorN ##
export Taylor, diffTaylor, integTaylor, evalTaylor, deriv, pretty_print
#
export TaylorN
export set_maxOrder, get_maxOrder, set_numVars, get_numVars

end


using PtFEM

ProjDir = dirname(@__FILE__)

# This example is identical to p.4.1.1.jl but re-uses the Cholesky
# decomposed global stiffness matrix in the 2nd part by calling
# p41(m, data) where the loaded_nodes values are doubled.

l = 1.0       # Total length [m]
N = 5         # Number of nodes
els = N - 1   # Number of finite elements
nod = 2       # Number of nodes per finite elements
nodof = 1     # Degrees of freedom for each node
np_types = 1  # Number of proerty types
EA = 1.0e5    # Strain stiffness
nip = 1       # Number of integration points

struct_el = :Rod
fin_el = :Line

data = Dict(
  # Rod(nxe, np_types, nip, fin_el(nod, nodof))
  :struc_el => getfield(Main, Symbol(struct_el))(els, np_types, nip,
    getfield(Main, Symbol(fin_el))(nod, nodof)),
  :properties => [EA;],
  :x_coords => 0.0:l/els:l,
  :support => [(N, [0])],
  :loaded_nodes => [
      (1, [-0.625]),
      (2, [-1.25]),
      (3, [-1.25]),
      (4, [-1.25]),
      (5, [-0.625])
    ]
);

data |> display
println()

@time fem, dis_df, fm_df = p41(data)
println()

# Update (double) the loaded_nodes values

F = -10.0
data[:loaded_nodes] = [(i, [0.0]) for i in 1:N]
for node in 1:N-1
  data[:loaded_nodes][node][2][1] += F / (2els)
  data[:loaded_nodes][node+1][2][1] += F / (2els)
end

# Pass in m so m.cfgsm can be re-used with this new load vector.

@time fem2, dis_df2, fm_df2 = p41(fem, data)
println()

display(dis_df2)
println()
display(fm_df2)


# using LazyArrays
# const AbstractOuterProductArray{T,N} = BroadcastArray{T,N,typeof(*),I} where {T,N,I}

# OuterProductArray(arrays::Vararg{<:AbstractVector{T},N}) where {T,N} =
#                  BroadcastArray(*,ntuple(k->reshape(arrays[k],ntuple(k1-> k==k1 ? length(arrays[k]) : 1,Val(N)))  ,Val(N))...)

tensorproduct(v::AbstractVector...)  =
    OuterProductArray(v...)
tensorproduct(v::GridArrays.AbstractIntervalGrid...)  =
    ProductGrid(v...)
abstract type AbstractOuterProductArray{T,N} <: AbstractArray{T,N} end

components(opa::AbstractOuterProductArray) = opa.vectors
size(A::AbstractOuterProductArray) = map(length, components(A))
axes(A::AbstractOuterProductArray) = map(x->axes(x,1), components(A))

@propagate_inbounds function Base.getindex(F::AbstractOuterProductArray, k::Integer)
    @boundscheck checkbounds(F, k)
    unsafe_array_getindex(F, (k,))
end

@propagate_inbounds function Base.getindex(F::AbstractOuterProductArray{T, N}, kj::Vararg{<:Integer, N}) where {T, N}
    @boundscheck checkbounds(F, kj...)
    unsafe_array_getindex(F, kj)
end

unsafe_array_getindex(F::AbstractOuterProductArray{T,N}, kj::NTuple{N,Int}) where {T,N} =
    prod(map(getindex, components(F), kj))

struct OuterProductArray{T,N,V} <: AbstractOuterProductArray{T,N}
    vectors   ::  V
    OuterProductArray(v::AbstractVector...) =
        new{promote_type(map(eltype,v)...),length(v),typeof(v)}(v)
end

struct OPArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end
Broadcast.BroadcastStyle(::Type{<:OuterProductArray{T,N}}) where {T,N} = OPArrayStyle{N}()
Base.copyto!(dest::OuterProductArray, bc::Broadcast.Broadcasted{<:OPArrayStyle}) =
    OuterProductArray(map(argsi->broadcast(bc.f, argsi...), zip(map(components, bc.args)...))...)
Base.copy(a::OuterProductArray) = OuterProductArray(map(copy, components(a))...)
Base.similar(a::OuterProductArray) = OuterProductArray(map(similar, components(a))...)
Base.similar(::Type{OuterProductArray{T}}, shape::Tuple{Union{Integer, Base.OneTo},Vararg{Union{Integer,Base.OneTo},N} where N}) where T =
    OuterProductArray(map(x->Base.similar(Vector{T}, x), shape)...)
Base.similar(bc::Broadcast.Broadcasted{OPArrayStyle{N}},::Type{ElType}) where {N,ElType} =
    Base.similar(OuterProductArray{ElType}, axes(bc))
function Base.copyto!(dest::OuterProductArray, src::OuterProductArray)
    for (deste,srce) in zip(components(dest),components(src))
        copy!(deste,srce)
    end
end


using Test, FinancialMonteCarlo,Statistics;
@show "Black Scholes Model"
S0=100.0;
K=100.0;
r=[0.00,0.02];
T=1.0;
d=0.01;
r=0.02;
D=90.0;

Nsim=100000;
Nstep=30;
sigma=0.2;
mc=MonteCarloConfiguration(Nsim,Nstep);
toll=0.8

rfCurve=ZeroRate(r);

FwdData=Forward(T)
EUData=EuropeanOption(T,K)
Model_base=BrownianMotion(sigma,r-d-sigma*sigma*0.5,Underlying(0.0,0.0))
S=S0.*exp.(simulate(Model_base,mc,T));

Payoff=exp(-FinancialMonteCarlo.integral(r,T))*S[:,end];
@show Price=mean(Payoff)
Model=BlackScholesProcess(sigma,Underlying(S0,d));
@show FwdPrice=pricer(Model,rfCurve,mc,FwdData);