module ThreeOmegaMethod

using QuadGK

export ThreeOmega, LayerTOM, HeaterGeometry, Source

export three_omega

abstract type LayerTOMInformation end
struct LayerTOM{T1} <: LayerTOMInformation where {T1<:Float64}
    ky::T1
    kxy::T1
    d::T1
    ρC::T1
end

abstract type HeaterInformation end
struct HeaterGeometry{T1} <: HeaterInformation where {T1<:Float64}
    b::T1
    l::T1
    ρh::Array{T1}
end

abstract type HeaterSource end
struct Source{T1} <: HeaterSource where {T1<:Float64}
    p::T1
    f::Array{T1}
end

abstract type OutputT end
struct ThreeOmega{T1} <: OutputT where {T1<:Float64}
    f::Array{T1}
    ΔThreal::Array{T1}
    ΔThimag::Array{T1}
    integerror::Array{T1}
end

function three_omega(
    Layers::Array{T1,N1},
    heater::T2,
    source::T3,
    thresistances::Array{T4,N4};
    int_limit::T5=1.0e6,
) where {T1<:LayerTOM, N1, T2<:HeaterGeometry, T3<:Source, T4<:Number, N4, T5<:Number}
    numLayers::Int64 = size(Layers, 1)
    # Check input
    numLayers > 1 || throw(DimensionMismatch("There should be at least one layer below the heater."))
    length(thresistances) == numLayers-1 || throw(DimensionMismatch("The number of interface-resistances should be lower than the number of layers: length(thresistances) == size(layers,2)-1."))
    # Power per unit length
    p_l::Float64 = source.p / heater.l
    # Reverse order of input for bottom-up calculation
    Layers = reverse(Layers, dims=1)
    thresistances = reverse(vec(thresistances), dims=1)
    # Lambda parameters
    λ = [lambda(Layers[i].ky, Layers[i].kxy, Layers[i].ky / Layers[i].ρC) for i = 1 : numLayers]
    # Construct B parameter for the system of LayerTOMs
    B = _concatenate_B(λ, Layers, thresistances, numLayers)
    # Construct integration term. The half term comes from the heater at the top of the stack
    F::Function = (η,ω) -> (1.0./λ[end](η,ω).*(B(η,ω)[1] + B(η,ω)[2])./(0.5*B(η,ω)[2] - 0.5*B(η,ω)[1]).*(sin.(heater.b*η)./(heater.b*η)).^2)
    # return integration
    sol = _integrate_temperature(
            source.f,
            int_limit,
            F,
            heater.ρh,
            p_l/π/2,
            p_l/2/heater.b,
            im*4*π.*source.f,
    )
    return sol
end # ThreeOmegaMethod

"""Lambda parameter."""
lambda(ky::T1, kxy::T1, α::T1) where {T1<:Float64} = (η,ω) -> (ky*sqrt.(kxy*η^2 + im*2*ω/α))

"""Interface matrix."""
function Λ(λ1::T0, λ2::T0, r::T1) where {T0<:Function, T1<:Float64}
    x11 = (η,ω) -> λ2(η,ω) + λ1(η,ω) + λ2(η,ω).*λ1(η,ω).*r
    x12 = (η,ω) -> λ2(η,ω) - λ1(η,ω) - λ2(η,ω).*λ1(η,ω).*r
    x21 = (η,ω) -> λ2(η,ω) - λ1(η,ω) + λ2(η,ω).*λ1(η,ω).*r
    x22 = (η,ω) -> λ2(η,ω) + λ1(η,ω) - λ2(η,ω).*λ1(η,ω).*r
    X = (η,ω) -> (0.5./λ2(η,ω).*[x11(η,ω) x12(η,ω);
                                 x21(η,ω) x22(η,ω)])
    return X
end

"""Transfer matrix."""
function U(λ::T0, ky::T1, d::T1) where {T0<:Function, T1<:Float64}
    X = (η,ω) -> ([exp.(-λ(η,ω)/ky*d) 0.;
                   0. exp.(λ(η,ω)/ky*d)])
    return X
end

"""Integration of the temperature term."""
function _integrate_temperature(
    f::Array{T0,N0},
    int_limit::T0,
    F::T1,
    ρh::Array{T0,N2},
    plπ::T0,
    plb::T0,
    h_param::Array{T2,N0},
) where {T0<:Float64, N0, T1<:Function, T2<:ComplexF64, N2}
    ΔTh = Array{ComplexF64,1}(undef, length(f))
    int_error = Array{Float64,1}(undef, length(f))
    @inbounds for i in eachindex(f)
        temp0::Tuple{ComplexF64,Float64} = quadgk((η) -> F(η, 2*π*f[i]), 0, int_limit, rtol=sqrt.(eps()))
        temp1::ComplexF64 = temp0[1]
        int_error[i] = temp0[2]
        ΔTh[i] = ( plπ*temp1 + ρh[1]*plb ) ./ ( 1 + ρh[2]*h_param[i] * (ρh[1] + temp1/plb)  )
    end
    sol = ThreeOmega(f, real(ΔTh), imag(ΔTh), int_error)
    return sol
end # function _integrate_temperature()

"""

    Concatenate B terms depending on the number of LayerTOMs. The boundary condition for
    the substrate is semi-infinite.

"""
function _concatenate_B(
    λ::Array{T0,N0},
    Layers::Array{T1,N1},
    thresistances::Array{T2,N2},
    numLayers::T3,
) where {T0<:Function, N0, T1<:LayerTOM, N1, T2<:Number, N2, T3<:Int64}
    B = Λ(λ[end-1], λ[end], thresistances[end])
    @inbounds for i = (numLayers-1):-1:2
        tempU::Function = U(λ[i], Layers[i].ky, Layers[i].d)
        tempΛ::Function = Λ(λ[i-1], λ[i], thresistances[i])
        B = let B = B
            (η,ω) -> B(η,ω) * tempU(η,ω) * tempΛ(η,ω)
        end # B = let B = B
    end # i = numLayers-1 : -1 : 2
    # Add semi-infinite boundary condition
    B = let B = B
        (η,ω) -> B(η,ω) * [0.; 1.]
    end # B = let B = B
    return B
end # function _concatenate_B()

end # module


using SparseArrays
using ViewerGL
GL = ViewerGL
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

cubohemioctahedron = Lar.cubohemioctahedron()

V,EV,FV,CV = Lar.struct2lar(cubohemioctahedron)
GL.VIEW([ GL.GLGrid(V,FV, GL.Point4d(1,1,1,0.1)) ,GL.GLFrame]);

cop_EV = Lar.coboundary_0(EV::Lar.Cells);
cop_FE = Lar.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);
W = convert(Lar.Points, V');

V, copEV, copFE, copCF = Lar.Arrangement.spatial_arrangement(W::Lar.Points, cop_EV::Lar.ChainOp, cop_FE::Lar.ChainOp);


using Test
using GEMPIC

include("test_maxwell_2d_fem.jl")
include("test_vm_1d1v_spin.jl")
include("test_vm_1d2v.jl")
include("test_mesh.jl")
include("test_spline_pp.jl")
include("test_maxwell_1d_fem.jl")
include("test_sampling.jl")
include("test_sampling_spin.jl")
include("test_particle_mesh_coupling_spline_1d.jl")
include("test_particle_mesh_coupling_spline_2d.jl")

if Threads.nthreads() <= 2
    include("test_hamiltonian_splitting.jl")
end

include("test_hamiltonian_splitting_spin.jl")
include("test_hamiltonian_splitting_boris.jl")
include("test_save_particles.jl")


using Test
using LinearAlgebra
import GPLinearODEMaker
GLOM = GPLinearODEMaker

println("Testing...")

@testset "cholesky factorizations" begin
    A = [4. 12 -16; 12 37 -43; -16 -43 98]
    chol_A = GLOM.ridge_chol(A)
    @test chol_A==cholesky(A)==GLOM.symmetric_A(A, chol=true)
    @test chol_A.L==LowerTriangular([2. 0 0;6 1 0;-8 5 3])
    A = [4. 12 -16; 12 37 -43; -16 -43 88]
    @test_logs (:warn, "added a ridge")  GLOM.ridge_chol(A)
    println()
end

@testset "nlogL derivatives" begin
    kernel, n_kern_hyper = include("../src/kernels/se_kernel.jl")

    n = 100
    xs = 20 .* sort(rand(n))
    y1 = sin.(xs) .+ 0.1 .* randn(n)
    noise1 = 0.1 .* ones(n)
    y2 = cos.(xs) .+ 0.2 .* randn(n)
    noise2 = 0.2 .* ones(n)

    ys = collect(Iterators.flatten(zip(y1, y2)))
    noise = collect(Iterators.flatten(zip(noise1, noise2)))

    prob_def = GLOM.GLO(kernel, n_kern_hyper, 2, 2, xs, ys; noise = noise, a0=[[1. 0.1];[0.1 1]])

    @test GLOM.est_dΣdθ(prob_def, 1 .+ rand(n_kern_hyper); return_bool=true, print_stuff=false)
    @test GLOM.test_∇nlogL_GLOM(prob_def, 1 .+ rand(n_kern_hyper), print_stuff=false)
    @test GLOM.test_∇∇nlogL_GLOM(prob_def, 1 .+ rand(n_kern_hyper), print_stuff=false)
    println()
end


"""
    ac_resample(ac, ratio[; filter])

Resample three-dimensional autocorrelation function `ac` by a ratio of
`ratio` times along the third axis, i.e. `ac` of shape `(x,y,z)`
becomes an array of shape `(x,y,ratio * z)`.
"""
function ac_resample(ac     :: AbstractArray{T, 3},
                     ratio  :: Union{Integer, Rational};
                     filter :: VoF = ratio |> filter_length |> filter_coeffs) where T <: AF
    x, y, z = size(ac)

    # Upsampling coefficient
    n = numerator(ratio)

    # Upsampled autocorrelation function
    ac_up = zeros(T, (x, y, n*z))

    ## To preserve "energy" of a signal
    filter = n * filter

    # Upsample
    for k in 1:z
        ac_up[:,:,n*(k-1)+1] = ac[:,:,k]
    end

    # Filter
    for k in 1:n:n*z
        cureven  = ac_up[:,:,mod1(k + 0, n*z)]
        nexteven = ac_up[:,:,mod1(k + n, n*z)]

        for l in 1:n-1
            ac_up[:,:,k + l] = filter[l]*cureven + filter[n-l]*nexteven
        end
    end

    # Downsample the result
    slices = (ac_up[:,:,k] for k in 1:denominator(ratio):n*z)
    return reduce((x, y) -> cat(x, y; dims = (3)), slices)
end


# Creating a test problem
n = 100;
p = 5;
U = randn(n,p);
V = randn(n,p);
K = DiaSymSemiseparable(U,V,ones(n));
x = randn(n);
Kfull = tril(U*V') + triu(V*U',1) + I;

# Testing multiplication
@test isapprox(K*x, Kfull*x, atol = 1e-6)

#  Testing multiplication with the adjoint operator
@test isapprox(K'*x, Kfull'*x, atol = 1e-6)


using Manifolds, Manopt, Test, ManifoldsBase

@testset "Differentials (on Sphere(2))" begin
    # The Adjoint Differentials test using the same variables as the differentials
    # test
    p = [1.0, 0.0, 0.0]
    q = [0.0, 1.0, 0.0]
    M = Sphere(2)
    X = log(M, p, q)
    # Tept differentials (1) Dp of Log_pq
    Y = similar(X)
    @test adjoint_differential_log_basepoint(M, p, p, X) == -X
    adjoint_differential_log_basepoint!(M, Y, p, p, X)
    @test Y == -X
    @test adjoint_differential_log_argument(M, p, p, X) == X
    adjoint_differential_log_argument!(M, Y, p, p, X)
    @test Y == X
    @test adjoint_differential_exp_basepoint(M, p, zero_vector(M, p), X) == X
    adjoint_differential_exp_basepoint!(M, Y, p, zero_vector(M, p), X)
    @test Y == X
    @test adjoint_differential_exp_argument(M, p, zero_vector(M, p), X) == X
    adjoint_differential_exp_argument!(M, Y, p, zero_vector(M, p), X)
    @test Y == X
    for t in [0, 0.15, 0.33, 0.66, 0.9]
        @test adjoint_differential_geodesic_startpoint(M, p, p, t, X) == (1 - t) * X
        adjoint_differential_geodesic_startpoint!(M, Y, p, p, t, X)
        @test Y == (1 - t) * X
        @test norm(M, p, adjoint_differential_geodesic_endpoint(M, p, p, t, X) - t * X) ≈ 0 atol =
            10.0^(-16)
        adjoint_differential_geodesic_endpoint!(M, Y, p, p, t, X)
        @test norm(M, p, Y - t * X) ≈ 0 atol = 10.0^(-16)
    end
    Mp = PowerManifold(M, NestedPowerRepresentation(), 3)
    pP = [p, q, p]
    qP = [p, p, q]
    XP = [X, zero_vector(M, p), -X]
    YP = similar.(XP)
    @test norm(
        Mp,
        pP,
        adjoint_differential_forward_logs(Mp, pP, XP) -
        [-X, zero_vector(M, p), zero_vector(M, p)],
    ) ≈ 0 atol = 4 * 10.0^(-16)
    adjoint_differential_forward_logs!(Mp, YP, pP, XP)
    @test norm(Mp, pP, YP - [-X, zero_vector(M, p), zero_vector(M, p)]) ≈ 0 atol =
        4 * 10.0^(-16)
    ZP = [[0.0, π / 2, 0.0], [0.0, 0.0, 0.0], [π / 2, 0.0, 0.0]]
    @test adjoint_differential_log_argument(Mp, pP, qP, XP) == ZP
    adjoint_differential_log_argument!(Mp, YP, pP, qP, XP)
    @test ZP == YP
end


#=
    Interpolate using weighted essentially non-oscillatory (WENO) techniques.
=#

export WENO4Interpolation

"""
    WENOInterpolationType

Abstract class for all types of WENO interpolation.
"""
abstract type WENOInterpolationType <: InterpolationType end

"""
    WENO4Interpolation

Fourth-order WENO interpolation. Based on Janett et al. (2019),
"A novel fourth-order WENO interpolation technique. A possible new tool designed for radiative transfer",
https://ui.adsabs.harvard.edu/abs/2019A&A...624A.104J
"""
struct WENO4Interpolation{T} <: WENOInterpolationType where T <: Number
    nodes::Vector{T}
    α2::Vector{T}  # coefficients for the α2 and α3 weights
    α3::Vector{T}  # coefficients for the α2 and α3 weights
    q2::Array{T, 2}  # coefficients for the quadratic Lagrange polynomials q2
    q3::Array{T, 2}  # coefficients for the quadratic Lagrange polynomials q3
end

function Base.:(==)(o1::WENO4Interpolation, o2::WENO4Interpolation)
    o1.nodes == o2.nodes &&
    o1.α2 == o2.α2 &&
    o1.α3 == o2.α3 &&
    o1.q2 == o2.q2 &&
    o1.q3 == o2.q3
end

lbounds(itp::WENO4Interpolation) = (first(itp.nodes),)
ubounds(itp::WENO4Interpolation) = (last(itp.nodes),)

function interpolate(nodes::AbstractVector{<:Number}, A::AbstractVector{<:Number})
    n = length(nodes)
    α2 = Vector{eltype(nodes)}(undef, n)
    α3 = Vector{eltype(nodes)}(undef, n)
    q2 = Array{eltype(nodes)}(undef, 3, n)
    q3 = Array{eltype(nodes)}(undef, 3, n)

    for i ∈ 1:n
        xim, xi, xip, xipp = get_stencil_nodes(i, nodes)
        yim, yi, yip, yipp = get_stencil_values(i, A)
        q2[:, i], q3[:, i] = weno4_q(xim, xi, xip, xipp, yim, yi, yip, yipp)
        α2[i], α3[i] = weno4_α(xim, xi, xip, xipp, yim, yi, yip, yipp)
    end
    WENO4Interpolation(nodes, α2, α3, q2, q3)
end

function (itp::WENO4Interpolation)(x::Number)
    @boundscheck (checkbounds(Bool, itp, x) || Base.throw_boundserror(itp, (x,)))
    n = length(itp.nodes)
    i = binary_search(itp.nodes, x)
    xim, xi, xip, xipp = get_stencil_nodes(i, itp.nodes)
    Δim = x - xim
    Δi = x - xi
    Δip = x - xip
    Δipp = x - xipp
    q2 = itp.q2[1, i] * Δi * Δip + itp.q2[2, i] * Δim * Δip +  itp.q2[3, i] * Δim * Δi
    q3 = itp.q3[1, i] * Δip * Δipp + itp.q3[2, i] * Δi * Δipp +  itp.q3[3, i] * Δi * Δip
    if i == 1
        result = q3
    elseif i == n - 1
        result = q2
    else
        α2 = Δipp * itp.α2[i]
        α3 = Δip * itp.α3[i]
        ω2 = α2 / (α2 + α3)
        ω3 = α3 / (α2 + α3)
        result = ω2 * q2 + ω3 * q3
    end
    return result
end

function get_stencil_nodes(index::Int, nodes::AbstractVector{<:Number})
    n = length(nodes)
    if index == n
        index -= 1
    end
    xi = nodes[index]
    xip = nodes[index+1]
    if index == 1
        xim = zero(eltype(nodes))
        xipp = nodes[index+2]
    elseif index == n - 1
        xim = nodes[index-1]
        xipp = zero(eltype(nodes))
    else
        xim = nodes[index-1]
        xipp = nodes[index+2]
    end
    return (xim, xi, xip, xipp)
end

function get_stencil_values(index::Int, A::AbstractVector{<:Number})
    n = length(A)
    if index == n
        index -= 1
    end
    yi = A[index]
    yip = A[index+1]
    if index == 1
        yim = zero(eltype(A))
        yipp = A[index+2]
    elseif index == n - 1
        yim = A[index-1]
        yipp = zero(eltype(A))
    else
        yim= A[index-1]
        yipp = A[index+2]
    end
    return (yim, yi, yip, yipp)
end

"""
    binary_search(a::AbstractVector{<:Number}, x)

Return the index of the last value in `a` less than or equal to `x`.
`a` is assumed to be sorted.
"""
function binary_search(a::AbstractVector{<:Number}, x)
    n = length(a)
    ileft = 2
    i = n - 1
    while ileft <= i
        mid = (ileft + i) >> 1
        if @inbounds x < a[mid]
            i = mid - 1
        else
            ileft = mid + 1
        end
    end
    return i
end

"""
Compute the smoothness indicators β2 and β3 for 4th order WENO interpolation.
Arguments are the x and y values for the 4-point stencil.
"""
function weno4_α(xim, xi, xip, xipp, yim, yi, yip, yipp)
    ε = 1.0f-6
    him = xi - xim
    hi = xip - xi
    hip = xipp - xip
    H = him + hi + hip
    # Derivatives
    yyim = -((2 * him + hi) * H + him * (him + hi)) / (him * (him + hi) * H) * yim
    yyim += ((him + hi) * H) / (him * hi * (hi + hip)) * yi
    yyim -= (him * H) / ((him + hi) * hi * hip) * yip
    yyim += (him * (him + hi)) / ((hi + hip) * hip * H) * yipp
    yyi = -(hi * (hi + hip)) / (him * (him + hi) * H) * yim
    yyi += (hi * (hi + hip) - him * (2 * hi + hip)) / (him * hi * (hi + hip)) * yi
    yyi += (him * (hi + hip)) / ((him + hi) * hi * hip) * yip
    yyi -= (him * hi) / ((hi + hip) * hip * H) * yipp
    yyip = (hi * hip) / (him * (him + hi) * H) * yim
    yyip -= (hip * (him + hi)) / (him * hi * (hi + hip)) * yi
    yyip += ((him + 2 * hi) * hip - (him + hi) * hi) / ((him + hi) * hi * hip) * yip
    yyip += ((him + hi) * hi) / ((hi + hip) * hip * H) * yipp
    yyipp = -((hi + hip) * hip) / (him * (him + hi) * H) * yim
    yyipp += (hip * H) / (him * hi * (hi + hip)) * yi
    yyipp -= ((hi + hip) * H) / ((him + hi) * hi * hip) * yip
    yyipp += ((2 * hip + hi) * H + hip * (hi + hip)) / ((hi + hip) * hip * H) * yipp
    # Smoothness indicators
    β2 = (hi + hip)^2 * (abs(yyip - yyi) / hi - abs(yyi - yyim) / him)^2
    β3 = (him + hi)^2 * (abs(yyipp - yyip) / hip - abs(yyip - yyi) / hi)^2
    α2 = - one(typeof(xim)) / ((xipp - xim) * (ε + β2))
    α3 = one(typeof(xim)) / ((xipp - xim) * (ε + β3))
    return (α2, α3)
end

"""
Compute the q2 and q3 Lagrange interpolation factors for 4th order WENO interpolation.
Arguments are the x and y values for the 4-point stencil.
"""
function weno4_q(xim, xi, xip, xipp, yim, yi, yip, yipp)
    him = xi - xim
    hi = xip - xi
    hip = xipp - xip
    q2 = Vector{eltype(xim)}(undef, 3)
    q3 = Vector{eltype(xim)}(undef, 3)
    q2[1] = yim / (him * (him + hi))
    q2[2] = -yi / (him * hi)
    q2[3] = yip / ((him + hi) * hi)
    q3[1] = yi / (hi * (hi + hip))
    q3[2] = -yip / (hi * hip)
    q3[3] = yipp / ((hi + hip) * hip)
    return (q2, q3)
end


using Test
using LinearAlgebra
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

polyorder = 1
numqp = required_quadrature_order(polyorder)
basis = TensorProductBasis(2,polyorder)
levelset = InterpolatingPolynomial(1,basis)
levelsetcoeffs = [-1.,-1.,1.,1.]
update!(levelset,levelsetcoeffs)

xL,xR = [-1.,-1.],[1.,1.]
pquad = area_quadrature(levelset,+1,xL,xR,numqp)
nquad = area_quadrature(levelset,-1,xL,xR,numqp)
squad = surface_quadrature(levelset,xL,xR,numqp)

lambda,mu = 1.,2.
penalty = 1.0
dx = 0.1
e11 = dx / 2.0
e22 = -lambda / (lambda + 2mu) * e11
dy = e22
penalty = 1.
stiffness = CutCell.plane_strain_voigt_hooke_matrix(lambda,mu)
cellmap = CutCell.CellMap([0.,0.],[2.,1.])

pbf = CutCell.bilinear_form(basis,pquad,stiffness,cellmap)
nbf = CutCell.bilinear_form(basis,nquad,stiffness,cellmap)

normals = repeat([1.,0.],1,numqp)
top = CutCell.coherent_traction_operator(basis,squad,squad,normals,stiffness,cellmap)

facescale = CutCell.scale_area(cellmap,normals)
mm = penalty*CutCell.interface_mass_matrix(basis,squad,squad,facescale)

eta = -1.0
BF = [pbf         zeros(8,8)
      zeros(8,8)  nbf       ]
T1 = -0.5*[-top  -top
           +top  +top]
T2 = eta*T1'

MM = [+mm  -mm
      -mm  +mm]

K = BF+T1+T2+MM
@test norm(K-K') > 1.0

rhs = zeros(16)

CutCell.apply_dirichlet_bc!(K,rhs,[5,6],1,0.,2)
CutCell.apply_dirichlet_bc!(K,rhs,[5],2,0.,2)
CutCell.apply_dirichlet_bc!(K,rhs,[3,4],1,dx,2)

sol = K\rhs
disp = reshape(sol,2,:)

testdisp = [0.0 0.0 dx dx 0.0 0.0 dx dx
            0.0 dy  0. dy 0.  dy  0. dy]

@test allapprox(disp,testdisp,1e2*eps())

eta = 1.0
BF = [pbf         zeros(8,8)
      zeros(8,8)  nbf       ]
T1 = -0.5*[-top  -top
           +top  +top]
T2 = eta*T1'

MM = [+mm  -mm
      -mm  +mm]

K = BF+T1+T2+MM
@test isapprox(norm(K-K'),0.0,atol=1e1eps())

rhs = zeros(16)

CutCell.apply_dirichlet_bc!(K,rhs,[5,6],1,0.,2)
CutCell.apply_dirichlet_bc!(K,rhs,[5],2,0.,2)
CutCell.apply_dirichlet_bc!(K,rhs,[3,4],1,dx,2)

sol = K\rhs
disp = reshape(sol,2,:)

testdisp = [0.0 0.0 dx dx 0.0 0.0 dx dx
            0.0 dy  0. dy 0.  dy  0. dy]

@test allapprox(disp,testdisp,1e2*eps())

eta = 0.0
BF = [pbf         zeros(8,8)
      zeros(8,8)  nbf       ]
T1 = -0.5*[-top  -top
           +top  +top]
T2 = eta*T1'

MM = [+mm  -mm
      -mm  +mm]

K = BF+T1+T2+MM
@test norm(K-K') > 1.0

rhs = zeros(16)

CutCell.apply_dirichlet_bc!(K,rhs,[5,6],1,0.,2)
CutCell.apply_dirichlet_bc!(K,rhs,[5],2,0.,2)
CutCell.apply_dirichlet_bc!(K,rhs,[3,4],1,dx,2)

sol = K\rhs
disp = reshape(sol,2,:)

testdisp = [0.0 0.0 dx dx 0.0 0.0 dx dx
            0.0 dy  0. dy 0.  dy  0. dy]

@test allapprox(disp,testdisp,1e2*eps())


struct Beam
    λ::Real; z::Real; n::Real
    x::Real; k::Real; q::Complex
end
Beam(λ::Real, w0::Real, n0::Real=1) = Beam(λ, 0, n0, 0, 0, 1im*π*n0*w0^2/λ)

η(e::Interface) = e.η
η(e::Union{Tan,Sag}) = dη(e.e)
η(e::Element) = 1

dz(e::FreeSpace) = e.L
dz(e::Union{Tan,Sag}) = dz(e.e)
dz(e::Element) = 0

function transform(e::Element, Γ::Beam)
    M = RTM(e)
    return Beam(Γ.λ, Γ.z+dz(e), Γ.n/η(e), M*[Γ.x,Γ.k]..., /(M*[Γ.q,1]...))
end

function beamtrace(elems::Vector{<:Element}, Γ0::Beam)
    Γs = Vector{Beam}(undef, length(elems)+1); Γs[1] = Γ0
    for (ind, elem) in enumerate(elems)
        Γs[ind+1] = transform(elem, Γs[ind])
    end
    return Γs
end

spotsize(Γ::Beam) = /(-Γ.λ, π*Γ.n*imag(1/Γ.q)) |> sqrt
location(Γ::Beam) = Γ.z

"""

Return a function that calculates the beam's waist radius as a
function of the propagation distance.

Expects an optical system (e.g. a vector or collection of
[`Element`](@ref)) as first argument, and a [`Beam`](@ref) as second
argument. The optional keyword argument `outside` can be used to
define a return value for beam positions outside those covered by
the system; the default is to throw a DomainError.

"""
function waistradiusfunc(elements, beam::Beam; outside=nothing)
    beams = Vector{Beam}(undef, length(elements)+1)
    beams[1] = beam
    for (i, el) in enumerate(elements)
        beam = transform(el, beam)
        beams[i+1] = beam
    end
    return function(z)
        if z < beams[1].z || z > beams[length(beams)].z
            if isnothing(outside)
                throw(DomainError(string(
                    "system does not cover ",
                    "the requested beam position"
                )))
            end
            return outside
        end
        # find last element before requested beam position z; note
        # that this could be accelerated by a more intelligent
        # algorithm since the beams[i].z should be sorted (and this
        # algorithm already depends on this expected ordering).
        i = 0
        while beams[i+1].z < z
            i += 1
        end
        i = max(i, 1)
        # add the effect of a FreeSpace element to reach the requested
        # beam position z
        beam = transform(FreeSpace(z - beams[i].z), beams[i])
        return spotsize(beam)
    end
end

discretize(e::FreeSpace, N::Int) = fill(FreeSpace(e.L/N), N)
discretize(e::Element, N::Int) = e
discretize(els::Vector{<:Element}, N::Int) = vcat(discretize.(els,N)...)


export LayerNorm
using Statistics
using Knet.Train20: param

# TODO: this is slow, need a kernel, maybe https://github.com/tensorflow/tensorflow/pull/6205/files
# TODO: support other axes

"""
    LayerNorm(size::Integer; eps=1e-5)
    LayerNorm(γ, β, ϵ)

References:
* [Ba, Kiros and Hinton 2016](https://arxiv.org/abs/1607.06450) Layer Normalization
* torch.nn.LayerNorm
* tf.keras.layers.LayerNormalization
"""
struct LayerNorm; γ; β; ϵ; end

function LayerNorm(dmodel; eps=1e-5)
    γ = param(dmodel; init=ones)
    β = param(dmodel; init=zeros)
    LayerNorm(γ, β, eps)
end

function (l::LayerNorm)(x, o...)
    μ = mean(x,dims=1)
    σ = std(x,mean=μ,dims=1)
    ϵ = eltype(x)(l.ϵ)
    l.γ .* (x .- μ) ./ (σ .+ ϵ) .+ l.β # TODO: doing x .- μ twice?
end

# function (l::LayerNorm)(x::MaskedArray, o...)
#     MaskedArray(l(x.array), x.mask) # TODO: shouldn't normalization ignore masked values?
# end


# # Linear regression

# Import the necessary packages:

using Distributions
using MLJBase
using Soss
using SossMLJ
using Statistics

# In this example, we fit a Bayesian linear regression model with the
# canonical link function.

# Suppose that we are given a matrix of features `X` and a column vector of
# labels `y`. `X` has `n` rows and `p` columns. `y` has `n` elements. We assume
# that our observation vector `y` is a realization of a random variable `Y`.
# We define `μ` (mu) as the expected value of `Y`, i.e. `μ := E[Y]`. Our model
# comprises three components:
#
# 1. The probability distribution of `Y`: for linear regression, we assume that each `Yᵢ` follows a normal distribution with mean `μᵢ` and variance `σ²`.
# 2. The systematic component, which consists of linear predictor `η` (eta), which we define as `η := Xβ`, where `β` is the column vector of `p` coefficients.
# 3. The link function `g`, which provides the following relationship: `g(E[Y]) = g(μ) = η = Xβ`. It follows that `μ = g⁻¹(η)`, where `g⁻¹` denotes the inverse of `g`. For linear regression, the canonical link function is the identity function. Therefore, when using the canonical link function, `μ = g⁻¹(η) = η`.
#
# In this model, the parameters that we want to estimate are `β` and `σ`.
# We need to select prior distributions for these parameters. For each `βᵢ`
# we choose a normal distribution with zero mean and variance `s²`. Here, `βᵢ`
# denotes the `i`th component of `β`. For `σ`, we will choose a half-normal
# distribution with variance `t²`. `s` and `t` are hyperparameters that we will
# need to choose.

# We define this model using the Soss probabilistic programming library:

m = @model X, s, t begin
    p = size(X, 2) # number of features
    β ~ Normal(0, s) |> iid(p) # coefficients
    σ ~ HalfNormal(t) # dispersion
    η = X * β # linear predictor
    μ = η # `μ = g⁻¹(η) = η`
    y ~ For(eachindex(μ)) do j
        Normal(μ[j], σ) # `Yᵢ ~ Normal(mean=μᵢ, variance=σ²)`
    end
end;

# Generate some synthetic features. Let us generate two continuous
# features and two binary categorical features.

num_rows = 1_000
x1 = randn(num_rows)
x2 = randn(num_rows)
x3 = Int.(rand(num_rows) .> 0.5)
x4 = Int.(rand(num_rows) .> 0.5)
X = (x1 = x1, x2 = x2, x3 = x3, x4 = x4)

# Define the hyperparameters of our prior distributions:

hyperparams = (s=0.1, t=0.1)

# Convert the Soss model into a `SossMLJModel`:

model = SossMLJModel(;
    model       = m,
    hyperparams = hyperparams,
    infer       = dynamicHMC,
    response    = :y,
);

# Generate some synthetic labels:

args = merge(model.transform(X), hyperparams)
truth = rand(m(args))

# Create an MLJ machine for fitting our model:

mach = MLJBase.machine(model, X, truth.y)

# Fit the machine. This may take several minutes.

fit!(mach)

# Construct the posterior distribution and the joint posterior predictive distribution:

predictor_joint = predict_joint(mach, X)
typeof(predictor_joint)

# Draw a single sample from the joint posterior predictive distribution:

single_sample = rand(predictor_joint; response = :y)

# Evaluate the logpdf of the joint posterior predictive distribution at this sample:

logpdf(predictor_joint, single_sample)

# True `β`:

truth.β

# Posterior distribution of `β`

predict_particles(mach, X; response = :β)

# Difference between the posterior distribution of `β` to the true values:

truth.β - predict_particles(mach, X; response = :β)

# Compare the joint posterior predictive distribution of `μ` to the true values:

truth.μ - predict_particles(mach, X; response = :μ)

# Compare the joint posterior predictive distribution of `y` to the true values:

truth.y - predict_particles(mach, X; response = :y)

# Construct each of the marginal posterior predictive distributions:

predictor_marginal = MLJBase.predict(mach, X)
typeof(predictor_marginal)

# `predictor_marginal` has one element for each row in `X`

size(predictor_marginal)

# Draw a single sample from each of the marginal posterior predictive distributions:

only.(rand.(predictor_marginal))

# Use cross-validation to evaluate the model with respect to the expected value of the root mean square error (RMSE)

evaluate!(mach, resampling=CV(; nfolds = 6, shuffle = true), measure=rms_expected, operation=predict_particles)

# Use cross-validation to evaluate the model with respect to the median of the root mean square error (RMSE)

evaluate!(mach, resampling=CV(; nfolds = 6, shuffle = true), measure=rms_median, operation=predict_particles)


"""
# Hooke's atom

Find the total energy and density of the
[Hooke's atom](https://en.wikipedia.org/wiki/Hooke's_atom) using LDA DFT.
"""

using AtomEnergyLevels
using SpecialFunctions, PyPlot
pygui(true);

x = -30.0:0.1:5;
r  = exp.(x);

# exact density of the ground state, see
# S. Kais et al // Density functionals and dimensional
# renormalization for an exactly solvable model, JCP 99, 417 (1993);
# http://dx.doi.org/10.1063/1.465765

N² = 1 / (π^(3/2) * (8 + 5 * sqrt(π)));
ρₑ = @. 2N² * exp(-1/2 * r^2) * (sqrt(π/2) * (7/4 + 1/4 * r^2 + (r + 1/r) * erf(r/sqrt(2))) + exp(-1/2 * r^2));

xalpha = lda(2, x, conf = c"[He]", Vex = r -> 1/8 * r^2,
             xc = xc_func(xc_type_lda, false, ρ -> LDA_X(ρ, α = 0.798), ρ -> 0));
svwn5  = lda(2, x, conf = c"[He]", Vex = r -> 1/8 * r^2);

begin
    title("Hooke's atom density")
    xlabel("r, Bohr")
    ylabel("ρ(r)")

    ax1 = PyPlot.axes()
    ax1.set_xlim([0,5])

    plot(r, ρₑ, label = "Exact")
    plot(r, xalpha.density ./ r, label = "Xα, α=0.798")
    plot(r, svwn5.density ./ r, label = "LDA")

    legend(loc = "upper right", fancybox = "true")
    grid("on")
end


using LinearAlgebra
using SparseArrays
using PartitionedArrays
using Test
using IterativeSolvers

function test_fem_sa(parts)

  u(x) = 1.0 #x[1]+x[2]
  lx = 2.0
  ls = (lx,lx)
  nx = 10
  ns = (nx,nx)
  n = prod(ns)
  h = lx/nx
  D = length(ns)
  Ae = (h/6)*[
    4.0  -1.0  -1.0  -2.0
   -1.0   4.0  -2.0  -1.0
   -1.0  -2.0   4.0  -1.0
   -2.0  -1.0  -1.0   4.0
  ]

  cis_gcells = CartesianIndices(ns)
  cis_gnodes = CartesianIndices(ns.+1)
  cis_enodes = CartesianIndices(ntuple(i->2,Val(D)))
  lis_enodes = LinearIndices(cis_enodes)
  lis_gnodes = LinearIndices(cis_gnodes)
  c1 = CartesianIndex(ntuple(i->1,Val(D)))

  cells = PRange(parts,ns)

  # Loop over owned cells, and fill the coo-vectors
  # Note that during the process we will touch remote rows and cols
  # This will be fixed later with an assembly.
  I,J,V = map_parts(cells.partition) do cells
    I = Int[]
    J = Int[]
    V = Float64[]
    for ocell in cells.oid_to_lid
      gcell = cells.lid_to_gid[ocell]
      ci_gcell = cis_gcells[gcell]
      for ci_erow in cis_enodes
        ci_grow = ci_gcell + (ci_erow - c1)
        grow = lis_gnodes[ci_grow]
        erow = lis_enodes[ci_erow]
        boundary = any(s->(1==s||s==(nx+1)),Tuple(ci_grow))
        if boundary
          push!(I,grow)
          push!(J,grow)
          push!(V,1)
        else
          for ci_ecol in cis_enodes
            ci_gcol = ci_gcell + (ci_ecol - c1)
            gcol = lis_gnodes[ci_gcol]
            ecol = lis_enodes[ci_ecol]
            push!(I,grow)
            push!(J,gcol)
            push!(V,Ae[erow,ecol])
          end
        end
      end
    end
    I,J,V
  end

  # Create rows and cols without ghost layer
  rows = PRange(parts,ns.+1)
  cols = PRange(parts,ns.+1)

  # Add remote row gids to the rows ghost layer
  add_gids!(rows,I)

  # Start an assembly of the COO Vectors,
  # i.e., send and add the triplets (i,j,v) to the part that owns i.
  t = async_assemble!(I,J,V,rows)

  # Meanwhile we can fill the rhs.
  # Allocate it.
  b = PVector(0.0,rows)

  # Fill it.
  # we use `global_view` in order to be able to index
  # b with global ids from within the parts.
  # Note that we touch non owned rows.
  map_parts(global_view(b,b.rows),cells.partition) do b, cells
    for ocell in cells.oid_to_lid
      gcell = cells.lid_to_gid[ocell]
      ci_gcell = cis_gcells[gcell]
      for ci_erow in cis_enodes
        ci_grow = ci_gcell + (ci_erow - c1)
        grow = lis_gnodes[ci_grow]
        erow = lis_enodes[ci_erow]
        boundary = any(s->(1==s||s==(nx+1)),Tuple(ci_grow))
        if boundary
          x = (Tuple(ci_grow) .- 1) .* h
          b[grow] += u(x)
        end
      end
    end
  end

  # Wait the assembly of the coo vectors to finish
  map_parts(wait∘schedule,t)

  # Now we can add off processor col ids to the ghost layer of cols.
  add_gids!(cols,J)

  # Compress the coo vectors and build the matrix
  A = PSparseMatrix(I,J,V,rows,cols,ids=:global)

  # When filling b we have touched remote rows.
  # Send and add their contribution to the owner part
  # NOTE: calling this assembly before the previous one finishes can be
  # problematic since we can have wrongly matching snd/rcv in MPI.
  assemble!(b)

  # Setup initial and exact solution
  x = PVector(0.0,cols)
  x̂ = similar(x)
  map_parts(x.values,x̂.values,x.rows.partition) do x,x̂,rows
    for lid in rows.oid_to_lid
      grow = rows.lid_to_gid[lid]
      ci_grow = cis_gnodes[grow]
      xi = (Tuple(ci_grow) .- 1) .* h
      x̂[lid] = u(xi)
      boundary = any(s->(1==s||s==(nx+1)),Tuple(ci_grow))
      if boundary
        x[lid] = u(xi)
      end
    end
  end

  # Solve!
  IterativeSolvers.cg!(x,A,b,verbose=i_am_main(parts))

  @test norm(x-x̂) < 1.0e-5

end


include("blackpaper.jl")
using DelimitedFiles

# predefined hyperparameters for solution

# center location

const x_00 = 0.0
const y_00 = 0.0

# velocity field radius

const r_1 = 1.0

# center movement hyperparameters

const u_00 = 1.0 # transfer speed for center
const r_00 = 0.1 # radius of the circle where the center of the initial condition travels
const T_0  = 1.0  # time of the single rotation of the initial condition

# hyperparameters for transfer speeds

const T_1 = 0.5 # time of the single rotation around the center
const T_2 = 0.5 # Decaying factor

# hyperparameters for initial conditions

const H_0   = 1.0 # height of the cone
const r_0   = 0.5 # radius of the cone
#angles defining the slice of the cone
const phi_0 = pi / 4.0
const phi_1 = pi / 2.0

# define new fuctions from templates
# to build the solution, we need a single sample of
# initial conditions, transfer speeds, and center point movement functions

# radius-vector from central point defined above
function r(x, y)
	return r_template(x, y, x_00, y_00)
end

# x coordinate for center point
function x0_linear(t)
	return x0_linear_template(t, u_00, x_00)
end

# y coordinate for center point
function y0_linear(t)
	return x0_linear_template(t, u_00, y_00)
end

function a1_linear(r)
	return r
end

function a2_constant(t)
	return a2_constant_template(t, T_1)
end

# initial conditions in polar coordinates
function cone_polar(r, phi)
	return cone_template(r, phi, r_0, H_0)
end

function cone_slice_polar(r, phi)
	return cone_slice_template(r, phi, r_0, H_0, phi_0, phi_1)
end

# initial conditions in cartesian coordinates
function cone_cartesian(xs, ys)

	grid = [(x, y) for x in xs, y in ys]

	radius = inflate(
			(x, y) -> ((x .- x_00) ^ 2 + (y .- y_00) ^ 2) ^ 0.5,
			xs, ys
	)

	rlh = [
		r != 0.0 ?
		acos((x .- x_00) / r) :
		0.5 * pi
		for (r, (x, y)) in zip(radius, grid)
	]

	rrh = [
		r != 0.0 ?
		(2.0 * pi .- acos((x .- x_00) / r)) :
		1.5 * pi
		for  (r, (x, y)) in zip(radius, grid)
	]

	phi     = zeros((length(xs), length(ys)))
	mask    = [y .<  y_00 for (x, y) in grid]
	notmask	= [y .>= y_00 for (x, y) in grid]
	phi[mask]    = rrh[mask]
	phi[notmask] = rlh[notmask]

	periods = floor.(phi / (2.0 * pi))
	phi     = phi - periods * 2.0 * pi

	return cone_template(radius, phi, r_0, H_0)

function cone_slice_cartesian(xs, ys)

	grid = [(x, y) for x in xs, y in ys]

	radius = inflate(
			(x, y) -> ((x .- x_00) ^ 2 + (y .- y_00) ^ 2) ^ 0.5,
			xs, ys
	)

	rlh = [
		r != 0.0 ?
		acos((x .- x_00) / r) :
		0.5 * pi
		for (r, (x, y)) in zip(radius, grid)
	]

	rrh = [
		r != 0.0 ?
		(2.0 * pi .- acos((x .- x_00) / r)) :
		1.5 * pi
		for  (r, (x, y)) in zip(radius, grid)
	]

	phi     = zeros((length(xs), length(ys)))
	mask    = [y .<  y_00 for (x, y) in grid]
	notmask	= [y .>= y_00 for (x, y) in grid]
	phi[mask]    = rrh[mask]
	phi[notmask] = rlh[notmask]

	periods = floor.(phi / (2.0 * pi))
	phi     = phi - periods * 2.0 * pi

	return cone_slice_template(r, phi, r_0, H_0, phi_0, phi_1)

# transfer velocities
#function v1(x, y, t)
#	return v1_template(x, y, t, )

function solution(x, y, t)

	return solution_template(
		x, y, t,
		cone_slice,
		a1_linear,
		a2_constant,
		x0_linear,
		y0_linear
	)

end

# setup the domain

x = -1.0:0.01:1.0
y = -1.0:0.01:1.0
t = 0.0:0.1:10

res = solution(x, y, t)

# save as .csv file

writedlm("result.csv", res, ',')


"""
    function elo_ranking(
        players::Vector{Player},
        ranking_scheme_to_update::String,
        last_tournament_id::String
    )

Update the player's rankings using Elo Ranking system.
Computed after the completion of a tournament.
A player's ranking is not updated after the result of each match.
The ranking for each team is computed as average of individual player rankings.
K factor is calculated by dividing 800 by the number of games a player's rating is based on
plus the number of games the player completed in a tournament.

Arguments
=========
- `players`:
    an array of players in the simulation.
- `ranking_scheme_to_update`:
    the name of the ranking scheme to update for each Player.
- `last_tournament_id`:
    the last tournament that was completed.
"""

function elo_ranking(
    players::Vector{Player},
    ranking_scheme_to_update::String,
    last_tournament_id::String
)
    updatedRankings = Dict{Int64, Int64}();
    for player in players
        if haskey(player.doubles_results, last_tournament_id)
            expected_score = 0
            actual_score = 0
            teammate_id = player.doubles_results[last_tournament_id].partner_player_id
            teammate_ranking = search_elo_ranking(teammate_id, players)
            for match in player.doubles_results[last_tournament_id].matches
                RA = (player.rankings["Elo"] + teammate_ranking) / 2
                RB = search_elo_ranking(match.opponent1_player_id, players) + search_elo_ranking(match.opponent2_player_id, players)
                EA = 1 / (1 + 10^((RB-RA)/400))
                SA = match.result == "Won" ? 1 : 0
                expected_score += EA
                actual_score += SA
            end
            k_factor = 800 / (player.matches_played + length(player.doubles_results[last_tournament_id].matches))
            updatedRankings[player.id] = player.rankings["Elo"] + k_factor*(actual_score-expected_score)
        else
            # If player didn't participate in tournament, their ranking doesn't change
            updatedRankings[player.id] = player.rankings["Elo"];
        end
    end
    map(player -> player.rankings[ranking_scheme_to_update] = updatedRankings[player.id], players)
    return nothing
end

"""
    function search_elo_ranking(
        player_id::Int64,
        players::Vector{Player}
    )::Float64

Search for a player's current Elo Ranking

Arguments
=========
- `player_id`:
    the player whose ranking is under search.
- `players`:
    an array of players in the simulation.
"""

function search_elo_ranking(
    player_id::Int64,
    players::Vector{Player}
)::Float64
    player_ranking = filter(p -> p.id == player_id, players)[1].rankings["Elo"]
    return player_ranking
end


using Test
using CartesianMesh
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell

function displacement(alpha, x::V) where {V<:AbstractVector}
    u1 = alpha * x[2] * sin(pi * x[1])
    u2 = alpha * (x[1]^3 + cos(pi * x[2]))
    return [u1, u2]
end

function displacement(alpha, x::M) where {M<:AbstractMatrix}
    npts = size(x)[2]
    return hcat([displacement(alpha, x[:, i]) for i = 1:npts]...)
end

function body_force(lambda, mu, alpha, x)
    b1 = alpha * (lambda + 2mu) * pi^2 * x[2] * sin(pi * x[1])
    b2 =
        -alpha * (6mu * x[1] + (lambda + mu) * pi * cos(pi * x[1])) +
        alpha * (lambda + 2mu) * pi^2 * cos(pi * x[2])
    return [b1, b2]
end

function displacement_nodeids(mesh)
    bn = CutCell.bottom_boundary_node_ids(mesh)
    tn = CutCell.top_boundary_node_ids(mesh)
    ln = CutCell.left_boundary_node_ids(mesh)
    return unique!(vcat(bn, tn, ln))
end

function is_traction_boundary(x)
    if x[1] ≈ 1.0
        return true
    else
        return false
    end
end

function right_face_traction_force(lambda, mu, alpha, x)
    s11 =
        (lambda + 2mu) * alpha * pi * x[2] * cos(pi * x[1]) -
        lambda * alpha * pi * sin(pi * x[2])
    s12 = alpha * mu * (3x[1]^2 + sin(pi * x[1]))

    return [s11, s12]
end

function add_cell_error_squared!(
    err,
    interpolater,
    exactsolution,
    cellmap,
    quad,
    detjac,
)
    for (p, w) in quad
        numsol = interpolater(p)
        exsol = exactsolution(cellmap(p))
        err .+= (numsol - exsol) .^ 2 * detjac * w
    end
end

function mesh_L2_error(
    nodalsolutions,
    exactsolution,
    nodalconnectivity,
    cellmaps,
    basis,
    errorquad,
)

    ndofs = size(nodalsolutions)[1]
    detjac = CutCell.determinant_jacobian(cellmaps[1])
    err = zeros(ndofs)
    interpolater = InterpolatingPolynomial(ndofs, basis)
    for (cellid, cellmap) in enumerate(cellmaps)
        nodeids = nodalconnectivity[:, cellid]
        elementsolutions = nodalsolutions[:, nodeids]
        update!(interpolater, elementsolutions)
        add_cell_error_squared!(
            err,
            interpolater,
            exactsolution,
            cellmap,
            errorquad,
            detjac,
        )
    end
    return sqrt.(err)
end

function apply_displacement_boundary_condition!(
    matrix,
    rhs,
    displacement_function,
    femesh,
)

    boundarynodeids = displacement_nodeids(femesh)
    nodalcoordinates = CutCell.nodal_coordinates(femesh)
    boundarynodecoordinates = nodalcoordinates[:, boundarynodeids]
    boundarydisplacement = displacement_function(boundarynodecoordinates)
    CutCell.apply_dirichlet_bc!(
        matrix,
        rhs,
        boundarynodeids,
        boundarydisplacement,
    )
end

function linear_system(
    basis,
    quad,
    facequads,
    stiffness,
    femesh,
    bodyforcefunc,
    tractionfunc,
    istractionboundary,
)

    cellmaps = CutCell.cell_maps(femesh)
    nodalconnectivity = CutCell.nodal_connectivity(femesh)
    cellconnectivity = CutCell.cell_connectivity(femesh)

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    cellmatrix = CutCell.bilinear_form(basis, quad, stiffness, cellmaps[1])
    CutCell.assemble_bilinear_form!(sysmatrix, cellmatrix, nodalconnectivity, 2)
    CutCell.assemble_body_force_linear_form!(
        sysrhs,
        bodyforcefunc,
        basis,
        quad,
        cellmaps,
        nodalconnectivity,
    )
    CutCell.assemble_traction_force_linear_form!(
        sysrhs,
        tractionfunc,
        basis,
        facequads,
        cellmaps,
        nodalconnectivity,
        cellconnectivity,
        istractionboundary,
    )

    K = CutCell.make_sparse(sysmatrix, femesh)
    R = CutCell.rhs(sysrhs, femesh)

    return K, R
end

function error_for_num_elements(numelements, polyorder, numqp)
    x0 = [0.0, 0.0]
    widths = [1.0, 1.0]
    nelements = [numelements, numelements]
    mesh = UniformMesh(x0, widths, nelements)
    stiffness = plane_strain_voigt_hooke_matrix(lambda, mu)

    basis = TensorProductBasis(2, polyorder)
    quad = tensor_product_quadrature(2, numqp)
    facequads = CutCell.face_quadratures(numqp)
    errorquad = tensor_product_quadrature(2, numqp + 2)

    femesh = CutCell.Mesh(mesh, basis)

    matrix, rhs = linear_system(
        basis,
        quad,
        facequads,
        stiffness,
        femesh,
        x -> body_force(lambda, mu, alpha, x),
        x -> right_face_traction_force(lambda, mu, alpha, x),
        is_traction_boundary,
    )
    apply_displacement_boundary_condition!(
        matrix,
        rhs,
        x -> displacement(alpha, x),
        femesh,
    )

    sol = matrix \ rhs
    nodalsolutions = reshape(sol, 2, :)

    nodalconnectivity = CutCell.nodal_connectivity(femesh)
    cellmaps = CutCell.cell_maps(femesh)

    err = mesh_L2_error(
        nodalsolutions,
        x -> displacement(alpha, x),
        nodalconnectivity,
        cellmaps,
        basis,
        errorquad,
    )
    return err
end

function required_quadrature_order(polyorder)
    ceil(Int, 0.5 * (2polyorder + 1))
end

function convergence(nelements, polyorder)
    numqp = required_quadrature_order(polyorder)
    err = zeros(length(nelements), 2)
    for (idx, nelement) in enumerate(nelements)
        err[idx, :] = error_for_num_elements(nelement, polyorder, numqp)
    end
    dx = 1.0 ./ nelements
    u1err = err[:,1]
    u2err = err[:,2]

    u1rate = sum(diff(log.(u1err)) ./ diff(log.(dx)))/(length(nelements)-1)
    u2rate = sum(diff(log.(u2err)) ./ diff(log.(dx)))/(length(nelements)-1)

    return round(u1rate,digits=3),round(u2rate,digits=3)
end

const lambda = 1.0
const mu = 2.0
const alpha = 0.01

numelements = [2,4,8,16,32,64]
u1rate, u2rate = convergence(numelements,1)
# println("Convergence of linear elements : ", u1rate, "    ", u2rate)
@test isapprox(u1rate,2.0,atol=0.05)
@test isapprox(u2rate,2.0,atol=0.05)

u1rate, u2rate = convergence(numelements,2)
# println("Convergence of quadratic elements : ", u1rate, "    ", u2rate)
@test isapprox(u1rate,3.0,atol=0.05)
@test isapprox(u2rate,3.0,atol=0.05)

# u1rate, u2rate = convergence(numelements,3)
# # println("Convergence of cubic elements : ", u1rate, "    ", u2rate)
# @test isapprox(u1rate,4.0,atol=0.05)
# @test isapprox(u2rate,4.0,atol=0.05)
#
#
# u1rate, u2rate = convergence(numelements,4)
# # println("Convergence of quartic elements : ", u1rate, "    ", u2rate)
# @test isapprox(u1rate,5.0,atol=0.1)
# @test isapprox(u2rate,5.0,atol=0.1)


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export LeftStringWeight

"""
`LeftStringWeight(x)`

| Set                     |           ``\\oplus``  |  ``\\otimes``  | ``\\bar{0}`` | ``\\bar{1}`` |
|:-----------------------:|:----------------------:|:--------------:|:------------:|:------------:|
|``L^*\\cup\\{\\infty\\}``|  longest common prefix |     ``\\cdot`` |``\\infty``   |``\\epsilon`` |

where ``L^*`` is Kleene closure of the set of characters ``L`` and ``\\epsilon`` the empty string.
"""
struct LeftStringWeight <: Semiring
  x::String
  iszero::Bool
end
LeftStringWeight(x::String) = LeftStringWeight(x,false)
reversetype(::Type{S}) where {S <: LeftStringWeight} = RightStringWeight

zero(::Type{LeftStringWeight}) = LeftStringWeight("",true)
one(::Type{LeftStringWeight})  = LeftStringWeight("",false)

# longest common prefix
function lcp(a::LeftStringWeight,b::LeftStringWeight)
  if a.iszero
    return b
  elseif b.iszero
    return a
  else
    cnt = 0
    for x in zip(a.x,b.x)
      if x[1] == x[2]
        cnt +=1
      else
        break
      end
    end
    return LeftStringWeight(a.x[1:cnt])
  end
end

*(a::T, b::T) where {T<:LeftStringWeight}=
((a.iszero) | (b.iszero)) ? zero(T) : T(a.x * b.x)
+(a::LeftStringWeight, b::LeftStringWeight) = lcp(a,b)
function /(a::LeftStringWeight, b::LeftStringWeight)
  if b == zero(LeftStringWeight)
    throw(ErrorException("Cannot divide by zero"))
  elseif a == zero(LeftStringWeight)
    return zero(LeftStringWeight)
  else
    str = ""
    for i=1:length(a.x)
      if i > length(b.x)
        str *= a.x[i]
      end
    end
    return LeftStringWeight(str)
  end
end
reverse(a::LeftStringWeight) = RightStringWeight(reverse(a.x),a.iszero)

#properties
isidempotent(::Type{W}) where {W <: LeftStringWeight} = true
isleft(::Type{W}) where {W <: LeftStringWeight}= true
isweaklydivisible(::Type{W}) where {W <: LeftStringWeight}= true


#
#----------------------------------------------------------------------
export bdfExtK, bdfExtK!
#----------------------------------------------------------------------
""" k'th Order Backward Difference Formula with Extrapolation for integrating

 du = f(t) + g(t)           (1)
 dt

 where f(t) is rapidly changing, and g(t) is considerably slower.
 For example, we treat diffusion implicity, and convection explicityly
 while time-stepping Navier-Stokes as diffusion timescales are an order
 of magnitude faster (pressure is also treated implicitly but that is
 another story) This allows for larger time-steps based on the intertial term.

 Implementation:

 input: t: array of timesteps [t(i+1), t(i), ..., t(i-k+1)] where
           k is the order of the extrapolation

 output: a: extrapolation coefficients for f(t) (k)
         b: interpolation coefficients for du/dt (k+1)

 (1) ⟺   u' - f(t) = g(t)

 where du/dt(t=i+1) = u(i+1)b(1) + ... + u(i-k+1)*b(k+1),
           f(t=i+1) = f(i+1)
           g(t=i+1) = g(i)a(1) + ... + g(i-k+1)*a(k)

"""
function bdfExtK(t;k=3)
    t  = unique(t)
    kk = length(t) - 1
    t1 = t[1]
    t0 = t[2:end]

    a = interpMat(t1,t0)
    b = derivMat(t)[1,:]'

    if(kk < k)
        global a = hcat(a,zeros(1,k-kk))
        global b = hcat(b,zeros(1,k-kk))
    else
        a = a[1:k]
        b = b[1:k+1]
    end

    if(kk==0) # steady state
        global a[1] = 1.0
    end

    return reshape(a,k), reshape(b,k+1)
end
#--------------------------------------#
function bdfExtK!(a::Vector
                 ,b::Vector
                 ,t::Vector)

    k = length(t) - 1
    aa,bb = bdfExtK(t;k=k)

    a .= aa
    b .= bb

    return
end
#----------------------------------------------------------------------
export TimeStepper
#----------------------------------------------------------------------
struct TimeStepper{T,U}
    time::Vector{T}
    bdfA::Vector{T}
    bdfB::Vector{T}

    istep::Array{U,1} # current step
    fstep::Array{U,1} # final step

    dt::Array{T,1} # time step size
    Ti::Array{T,1} # start time
    Tf::Array{T,1} # end time

end
#--------------------------------------#
function TimeStepper(Ti,Tf,dt,k=3) # add flag for variable dt

    time = [Ti] .* ones(k+1)
    bdfA,bdfB = bdfExtK(time)

    istep = [0]
    fstep = [0] # steady

    if(dt != 0)
        global fstep = [Int64(ceil(Tf/dt))]
    end

    return TimeStepper{Float64,Int64}(time,bdfA,bdfB
                                     ,istep,fstep
                                     ,[dt],[Ti],[Tf])
end
#--------------------------------------#
function updateHist!(tstep::TimeStepper)

    @unpack time, istep, dt, bdfA, bdfB = tstep

    updateHist!(time)

    istep  .+= 1
    time[1] += dt[1]
    bdfExtK!(bdfA,bdfB,time)

    return
end
#----------------------------------------------------------------------


@defcomp grosseconomy begin
    regions = Index()

    YGROSS = Variable(index=[time, regions]) # Gross world product GROSS of abatement and damages (trillions 2005 USD per year)
    K = Variable(index=[time, regions]) # Capital stock (trillions 2005 US dollars)

    al = Parameter(index=[time, regions]) # Level of total factor productivity
    l = Parameter(index=[time, regions]) # Level of population and labor
    I = Parameter(index=[time, regions]) # Investment (trillions 2005 USD per year)
    gama = Parameter() # Capital elasticity in production function
    dk = Parameter(index=[regions]) # Depreciation rate on capital (per year)
    k0 = Parameter(index=[regions]) # Initial capital value (trill 2005 USD)

    # TODO remove this, just a temporary output trick
    L = Variable(index=[time, regions])

    function run_timestep(p, v, d, t)
        #Define function for K
        for r in d.regions
            if is_first(t)
                v.K[t,r] = p.k0[r]
            else
                v.K[t,r] = (1 - p.dk[r])^10 * v.K[t-1,r] + 10 * p.I[t-1,r]
            end
        end

        #Define function for YGROSS
        for r in d.regions
            v.YGROSS[t,r] = (p.al[t,r] * (p.l[t,r]/1000)^(1-p.gama)) * (v.K[t,r]^p.gama)
        end

        # TODO remove this, just a temporary output trick
        for r in d.regions
            v.L[t,r] = p.l[t,r]
        end
    end
end


"""
    ⊠(V₁::VectorSpace, V₂::VectorSpace)

Given two vector spaces `V₁` and `V₂` (`ElementarySpace` or `ProductSpace`), or thus,
objects of corresponding fusion categories ``C₁`` and ``C₂``, ``V₁ ⊠ V₂`` constructs the
Deligne tensor product, an object in ``C₁ ⊠ C₂`` which is the natural tensor product of
those categories. In particular, the corresponding type of sectors (simple objects) is
given by `sectortype(V₁ ⊠ V₂) == sectortype(V₁) ⊠ sectortype(V₂)` and can be thought of as
a tuple of the individual sectors.

The Deligne tensor product also works in the type domain and for sectors and tensors. For
group representations, we have `Rep[G₁] ⊠ Rep[G₂] == Rep[G₁ × G₂]`, i.e. these are the
natural representation spaces of the direct product of two groups.
"""
⊠(V1::VectorSpace, V2::VectorSpace) = (V1 ⊠ one(V2)) ⊗ (one(V1) ⊠ V2)

# define deligne products with empty tensor product: just add a trivial sector of the type of the empty space to each of the sectors in the non-empty space
function ⊠(V::GradedSpace, P0::ProductSpace{<:EuclideanSpace{ℂ},0})
    I1 = sectortype(V)
    I2 = sectortype(P0)
    return Vect[I1 ⊠ I2](ifelse(isdual(V), dual(c), c) ⊠ one(I2) => dim(V, c) for c in sectors(V); dual = isdual(V))
end

function ⊠(P0::ProductSpace{<:EuclideanSpace{ℂ},0}, V::GradedSpace)
    I1 = sectortype(P0)
    I2 = sectortype(V)
    return Vect[I1 ⊠ I2](one(I1) ⊠ ifelse(isdual(V), dual(c), c) => dim(V, c) for c in sectors(V); dual = isdual(V))
end

function ⊠(V::ComplexSpace, P0::ProductSpace{<:EuclideanSpace{ℂ},0})
    I2 = sectortype(P0)
    return Vect[I2](one(I2) => dim(V); dual = isdual(V))
end

function ⊠(P0::ProductSpace{<:EuclideanSpace{ℂ},0}, V::ComplexSpace)
    I1 = sectortype(P0)
    return Vect[I1](one(I1) => dim(V); dual = isdual(V))
end

function ⊠(P::ProductSpace{<:EuclideanSpace{ℂ},0}, P0::ProductSpace{<:EuclideanSpace{ℂ},0})
    I1 = sectortype(P)
    I2 = sectortype(P0)
    return one(Vect[I1 ⊠ I2])
end

function ⊠(P::ProductSpace{<:EuclideanSpace{ℂ}}, P0::ProductSpace{<:EuclideanSpace{ℂ},0})
    I1 = sectortype(P)
    I2 = sectortype(P0)
    S = Vect[I1 ⊠ I2]
    N = length(P)
    return ProductSpace{S,N}(map(V->V ⊠ P0, tuple(P...)))
end

function ⊠(P0::ProductSpace{<:EuclideanSpace{ℂ},0}, P::ProductSpace{<:EuclideanSpace{ℂ}})
    I1 = sectortype(P0)
    I2 = sectortype(P)
    S = Vect[I1 ⊠ I2]
    N = length(P)
    return ProductSpace{S,N}(map(V->P0 ⊠ V, tuple(P...)))
end


using Revise
using ADCME
using LinearAlgebra

xc = rand(10)
yc = rand(10)
e = rand(10)
c = rand(10)
d = rand(3)
r = RBF2D(xc, yc; c=c, eps=e, d=d, kind = 0)

x = rand(5)
y = rand(5)
o = r(x, y)

O = zeros(5)
for i = 1:5
    for j = 1:10
        d = sqrt((x[i]-xc[j])^2 + (y[i]-yc[j])^2)
        O[i] += c[j] * exp(-(e[j]*d)^2)
    end
    O[i] += d[1] + d[2] * x[i] + d[3] * y[i]
end

sess = Session(); init(sess)
@test norm(run(sess, o)-O)<1e-5


# a rather literal translation of the pseudocode at https://en.wikipedia.org/wiki/MD5

const s = UInt32[7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22,
                 5, 9, 14, 20, 5, 9, 14, 20, 5, 9, 14, 20, 5, 9, 14, 20,
                 4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23,  4, 11, 16, 23,
                 6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21]

const K = UInt32[0xd76aa478, 0xe8c7b756, 0x242070db, 0xc1bdceee, 0xf57c0faf,
    0x4787c62a, 0xa8304613, 0xfd469501, 0x698098d8, 0x8b44f7af, 0xffff5bb1,
    0x895cd7be, 0x6b901122, 0xfd987193, 0xa679438e, 0x49b40821, 0xf61e2562,
    0xc040b340, 0x265e5a51, 0xe9b6c7aa, 0xd62f105d, 0x02441453, 0xd8a1e681,
    0xe7d3fbc8, 0x21e1cde6, 0xc33707d6, 0xf4d50d87, 0x455a14ed, 0xa9e3e905,
    0xfcefa3f8, 0x676f02d9, 0x8d2a4c8a, 0xfffa3942, 0x8771f681, 0x6d9d6122,
    0xfde5380c, 0xa4beea44, 0x4bdecfa9, 0xf6bb4b60, 0xbebfbc70, 0x289b7ec6,
    0xeaa127fa, 0xd4ef3085, 0x04881d05, 0xd9d4d039, 0xe6db99e5, 0x1fa27cf8,
    0xc4ac5665, 0xf4292244, 0x432aff97, 0xab9423a7, 0xfc93a039, 0x655b59c3,
    0x8f0ccc92, 0xffeff47d, 0x85845dd1, 0x6fa87e4f, 0xfe2ce6e0, 0xa3014314,
    0x4e0811a1, 0xf7537e82, 0xbd3af235, 0x2ad7d2bb, 0xeb86d391]

function md5(msgbytes)
    a0::UInt32 = 0x67452301  # A
    b0::UInt32 = 0xefcdab89  # B
    c0::UInt32 = 0x98badcfe  # C
    d0::UInt32 = 0x10325476  # D

    oldlen = length(msgbytes)
    umsg = push!([UInt8(b) for b in msgbytes], UInt8(0x80))
    while length(umsg) % 64 != 56
        push!(umsg, UInt8(0))
    end
    append!(umsg, reinterpret(UInt8, [htol(UInt64(oldlen) * 8)]))

    for j in 1:64:length(umsg)-1
        arr = view(umsg, j:j+63)
        M = [reinterpret(UInt32, arr[k:k+3])[1] for k in 1:4:62]
        A = a0
        B = b0
        C = c0
        D = d0

        for i in 0:63
            if 0 ≤ i ≤ 15
                F = D ⊻ (B & (C ⊻ D))
                g = i
            elseif 16 ≤ i ≤ 31
                F = C ⊻ (D & (B ⊻ C))
                g = (5 * i + 1) % 16
            elseif 32 ≤ i ≤ 47
                F = B ⊻ C ⊻ D
                g = (3 * i + 5) % 16
            elseif 48 ≤ i ≤ 63
                F = C ⊻ (B | (~D))
                g = (7 * i) % 16
            end
            F += A + K[i+1] + M[g+1]
            A = D
            D = C
            C = B
            B += ((F) << s[i+1]) | (F >> (32 - s[i+1]))
        end

        a0 += A
        b0 += B
        c0 += C
        d0 += D
    end
    digest = join(map(x -> lpad(string(x, base=16), 2, '0'), reinterpret(UInt8, [a0, b0, c0, d0])), "") # Output is in little-endian
end

for pair in [0xd41d8cd98f00b204e9800998ecf8427e => "", 0x0cc175b9c0f1b6a831c399e269772661 => "a",
   0x900150983cd24fb0d6963f7d28e17f72 => "abc", 0xf96b697d7cb7938d525a2f31aaf161d0 => "message digest",
   0xc3fcd3d76192e4007dfb496cca67e13b => "abcdefghijklmnopqrstuvwxyz",
   0xd174ab98d277d9f5a5611c2c9f419d9f => "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789",
   0x57edf4a22be3c955ac49da2e2107b67a => "12345678901234567890123456789012345678901234567890123456789012345678901234567890"]
   println("MD5 of $(pair[2]) is $(md5(pair[2])), which checks with $(string(pair[1], base=16)).")
end


"""
    dic(loglike::AbstractVector{<:Real})

Computes Deviance Information Criterion (DIC).

# Arguments
* `log_lik::AbstractArray`: A vector of posterior log likelihoods

# Returns
* `dic::Real`: DIC value
"""
function dic(loglike::AbstractVector{<:Real})
    D = deviance.(loglike)
    return mean(D) + 0.5 * var(D)
end

deviance(loglikelihood::Real) = -2 * loglikelihood


module VirusLoadCurve

    export fitVLF,
           LogVirusLoadFunction,
           fitVTM,
           LogViralTargetModel,
           VirusLoadData,
           Boxplots
    using OrdinaryDiffEq
    using Optim, LsqFit
    using Plots, LaTeXStrings
    using DelimitedFiles, DataFrames
    using Statistics
    using Turing
    using Printf
    using ProgressMeter

    import Base.summary

    struct VirusLoadData{T<:Real}
        t::Vector{T}
        v::Vector{T}
    end

    @recipe function f(data::VirusLoadData)
        tmin, tmax = extrema(data.t)
        vmin, vmax = extrema(data.v)
        x := data.t
        y := data.v
        seriestype := :scatter
        xaxis := ("Time (days)", (0.0, tmax), font(14))
        yaxis := (L"\log\,V(t)", font(14))
        label := "Data"
        grid := :none
        ()
    end

    include("virus_load_function.jl")
    include("virus_target_model.jl")

end #module


function run_distributed(num_sims::Integer) :: Float64
    @assert num_sims > 0

    acum = @distributed (+) for _ in 1:num_sims
        x = rand(Float64)
        y = rand(Float64)

        r = sqrt(x^2 + y^2)
        if r <= 1.0
            1.0
        else
            0.0
        end
    end

    area = acum / num_sims
    return 4*area
end


using Example
if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

# @test hello("Julia") == "Hello, Julia"
# @test domath(2.0) ≈ 7.0

### SDM
K = 5
α = ones(K)
β = 100.0
d = SparseDirMix(α, β)
rand(d)

### SBM
a, b = shape_Dir2GenDir(α)

nsim = 10e3
xD = hcat([rDirichlet(α) for i = 1:nsim ]...)
xGD = hcat([rGenDirichlet(a,b)[1] for i = 1:nsim ]...)

using Statistics
sum(abs.(mean(xD, dims=2) .- mean(xGD, dims=2)))

η = 1.0e3
π_small = 0.5
π_large = 0.1
γ = 1.5
δ = 1.5
p = SBMprior(K, η, π_small, π_large, γ, δ)

rand(p)

x = [0, 0, 2, 1, 0]

pp = SBM_multinom_post(p, x)

rand(pp)


#############################################################################
#############################################################################
#
# This file defines the polynomial type with several operations
#
#############################################################################
#############################################################################

####################################
# Polynomial type and construction #
####################################

"""
A Polynomial type - designed to be for polynomials with integer coefficients.
"""
struct Polynomial
    terms::Array{Any,1}
        #The terms in the stack need to satisfy:
            # Will never have terms with 0 coefficient
            # Will never have two terms with same coefficient
        #An empty terms stack means that the polynomial is zero
    Polynomial() = new(Array{Any,1}[])
    #Inner constructor
    Polynomial(h::Array{Any,1}) = new(h)
end

struct PolynomialModP
    terms::Polynomial
    prime::Int

    #Constructor for the case where only a prime is given, initialize an empty polynomial and enforce mod(prime) conditions
    PolynomialModP(prime::Int) = new(Polynomial(), prime)
    #Inner constructor
    PolynomialModP(poly::Polynomial, prime::Int) = new(poly, prime)
end

"""
Construct a polynomial with a single term.
"""
function Polynomial(t::Term)
    terms = Array{Any,1}([])
    t.coeff != 0 && push!(terms, t)
    return Polynomial(terms)
end

"""
Construct a polynomial with a vector of terms.
"""
function Polynomial(tv::Vector{Term})
    terms = Array{Any,1}([])
    for t in tv
        t.coeff != 0 && push!(terms,t)
    end
    return Polynomial(terms)
end

"""
Construct a polynomial of the form x^p-x.
"""
cyclotonic_polynomial(p::Int) = Polynomial([Term(1,p), Term(-1,0)])

"""
Construct a polynomial of the form x-n.
"""
linear_monic_polynomial(n::Int) = Polynomial([Term(1,1), Term(-n,0)])

"""
Construct a polynomial of the form x.
"""
x_poly() = Polynomial(Term(1,1))

"""
Creates the zero polynomial.
"""
zero(::Type{Polynomial})::Polynomial = Polynomial()

"""
Creates the unit polynomial.
"""
one(::Type{Polynomial})::Polynomial = Polynomial(one(Term))
one(p::Polynomial) = one(typeof(p))

"""
Generates a random polynomial.
"""
function rand(::Type{Polynomial} ;
                degree::Int = -1,
                terms::Int = -1,
                max_coeff::Int = 100,
                mean_degree::Float64 = 5.0,
                prob_term::Float64  = 0.7,
                monic = false,
                condition = (p)->true)

    while true
        _degree = degree == -1 ? rand(Poisson(mean_degree)) : degree
        _terms = terms == -1 ? rand(Binomial(_degree,prob_term)) : terms
        degrees = vcat(sort(sample(0:_degree-1,_terms,replace = false)),_degree)
        coeffs = rand(1:max_coeff,_terms+1)
        monic && (coeffs[end] = 1)
        p = Polynomial( [Term(coeffs[i],degrees[i]) for i in 1:length(degrees)] )
        condition(p) && return p
    end
end

###########
# Display #
###########

"""
Show a polynomial.
"""
function show(io::IO, p::Polynomial)
    p = deepcopy(p)
    if iszero(p)
        print(io,"0")
    else
        n = length(p.terms)
        for (i,t) in enumerate(p.terms)
            print(io, t, i != n ? " + " : "")
        end
    end
end

"""
Show a polynomial Modulo P.
"""
function show(io::IO, p::PolynomialModP)
    input_poly = deepcopy(p.terms)
    if iszero(input_poly)
        print(io."0")
    else
        n = length(input_poly.terms)
        for (i,t) in enumerate(input_poly.terms)
            print(io, t, i != n ? " + " : "")
        end
        print(" mod ", p.prime)
    end
end

##############################################
# Iteration over the terms of the polynomial #
##############################################

"""
Allows to do iteration over the terms of the polynomial. The iteration is in an arbitrary order.
"""
iterate(p::Polynomial, state=1) = iterate(p.terms, state)

##############################
# Queries about a polynomial #
##############################

"""
The number of terms of the polynomial.
"""
length(p::Polynomial) = length(p.terms)

"""
The leading term of the polynomial.
"""
#Assuming every array of polynomials given to the struct is in ascending order of the term degrees.
leading(p::Polynomial)::Term = isempty(p.terms) ? zero(Term) : last(p.terms)

"""
The trailing term of the polynomial.
"""
#Assuming every array of polynomials given to the struct is in ascending order of the term degrees.
trailing(p::Polynomial)::Term = isempty(p.terms) ? zero(Term) : first(p.terms)

"""
Returns the coefficients of the polynomial.
"""
coeffs(p::Polynomial)::Vector{Int} = [t.coeff for t in p]

"""
The degree of the polynomial.
"""
degree(p::Polynomial)::Int = leading(p).degree

"""
The content of the polynomial is the GCD of its coefficients.
"""
content(p::Polynomial)::Int = euclid_alg(coeffs(p))

"""
Evaluate the polynomial at a point `x`.
"""
evaluate(f::Polynomial, x::T) where T <: Number = sum(evaluate(t,x) for t in f)

################################
# Pushing and popping of terms #
################################

"""
Push a new term into the polynomial.
"""
#Note that ideally this would throw and error if pushing another term of degree that is already in the polynomial
function push!(p::Polynomial, t::Term)
    iszero(t) && return #don't push a zero
    push!(p.terms,t)
end

"""
Pop the trailing term out of the polynomial.
"""
pop!(p::Polynomial)::Term = pop!(p.terms)

"""
Pop the leading term out of the polynomial.
"""
popfirst!(p::Polynomial)::Term = popfirst!(p.terms)

"""
Check if the polynomial is zero.
"""
iszero(p::Polynomial)::Bool = isempty(p.terms)

#################################################################
# Transformation of the polynomial to create another polynomial #
#################################################################

"""
The negative of a polynomial.
"""
-(p::Polynomial) = Polynomial(map((pt)->-pt, p.terms))

"""
Create a new polynomial which is the derivative of the polynomial.
"""
function derivative(p::Polynomial)::Polynomial
    der_p = Polynomial()
    for term in p
        der_term = derivative(term)
        !iszero(der_term) && push!(der_p,der_term)
    end
    return der_p
end

"""
The prim part (multiply a polynomial by the inverse of its content).
"""
prim_part(p::Polynomial) = p ÷ content(p)

"""
A square free polynomial.
"""
square_free(p::Polynomial, prime::Int)::Polynomial = (PolynomialModP(p, prime) ÷ PolynomialModP(gcd(PolynomialModP(p, prime),PolynomialModP(derivative(p), prime)), prime))

#################################
# Queries about two polynomials #
#################################

"""
Check if two polynomials are the same
"""
==(p1::Polynomial, p2::Polynomial)::Bool = p1.terms == p2.terms

"""
Check if a polynomial is equal to 0.
"""
#Note that in principle there is a problem here. E.g The polynomial 3 will return true to equalling the integer 2.
==(p::Polynomial, n::T) where T <: Real = iszero(p) == iszero(n)

##################################################################
# Operations with two objects where at least one is a polynomial #
##################################################################

"""
Subtraction of two polynomials.
"""
-(p1::Polynomial, p2::Polynomial)::Polynomial = p1 + (-p2)

"""
Multiplication of polynomial and term.
"""
*(t::Term,p1::Polynomial)::Polynomial = iszero(t) ? Polynomial() : Polynomial(map((pt)->t*pt, p1.terms))
*(p1::Polynomial, t::Term)::Polynomial = t*p1

"""
Multiplication of polynomial and an integer.
"""
*(n::Int,p::Polynomial)::Polynomial = p*Term(n,0)
*(p::Polynomial,n::Int)::Polynomial = n*p

"""
Integer division of a polynomial by an integer.

Warning this may not make sense if n does not divide all the coefficients of p.
"""
÷(p::Polynomial,n::Int) = (prime)->Polynomial(map((pt)->((pt ÷ n)(prime)), p.terms))

"""
Integer division of a polynomial mod(P) by an integer.

Warning this may not make sense if n does not divide all the coefficients of p.
"""
÷(p::PolynomialModP,n::Int) = PolynomialModP(Polynomial(map((pt)->((pt ÷ n)(p.prime)), p.terms)), p.prime)

"""
Take the smod of a polynomial with an integer.
"""
function mod(f::Polynomial, p::Int)::Polynomial
    p_out = Polynomial()
    for t in f
        push!(p_out, mod(t, p)) #if coeff reduced to zero, push! will handle it
    end
    return p_out
end

"""
Power of a polynomial mod prime.
"""
# function pow_mod(p::Polynomial, n::Int, prime::Int)
#     n < 0 && error("No negative power")
#     out = one(p)
#     for _ in 1:n
#         out *= p
#         out = mod(out, prime)
#     end
#     return out
# end

"""
Power of a polynomial mod prime.
"""
function ^(p::PolynomialModP, n::Int)
    n < 0 && error("No negative power")
    out = one(p.terms)
    for _ in 1:n
        out *= p.terms
        out = mod(out, p.prime)
    end
    return out
end


@info "Starting SINDy tests"
@testset "SINDy" begin
    # Create a nonlinear pendulum
    function pendulum(u, p, t)
        x = u[2]
        y = -9.81sin(u[1]) - 0.1u[2]^3 -0.2*cos(u[1])
        return [x;y]
    end

    u0 = [0.99π; -1.0]
    tspan = (0.0, 20.0)
    dt = 0.3
    prob = ODEProblem(pendulum, u0, tspan)
    sol = solve(prob, Tsit5(), saveat = dt)

    # Create the differential data
    DX = similar(sol[:,:])
    for (i, xi) in enumerate(eachcol(sol[:,:]))
        DX[:,i] = pendulum(xi, [], 0.0)
    end

    # Create a basis
    @variables u[1:2]
    # Lots of polynomials
    polys = Operation[1]
    for i ∈ 1:5
        push!(polys, u.^i...)
        for j ∈ 1:i-1
            push!(polys, u[1]^i*u[2]^j)
        end
    end

    # And some other stuff
    h = [cos(u[1]); sin(u[1]); u[1]*u[2]; u[1]*sin(u[2]); u[2]*cos(u[2]); polys]

    basis = Basis(h, u)

    opt = STRRidge(1e-2)
    Ψ = SINDy(sol[:,:], DX, basis, opt, maxiter = 2000)
    @test_nowarn set_threshold!(opt, 1e-2)

    @test length(Ψ) == 2
    @test size(Ψ) == (2, )
    @test parameters(Ψ) ≈ [1.0; -0.2; -9.81; -0.1]
    W = get_coefficients(Ψ)
    @test size(W) == (length(basis), 2)
    @test_nowarn get_aicc(Ψ)
    @test sum(get_sparsity(Ψ)) == 4
    @test sum(get_error(Ψ)) < 1e-10
    @test Ψ(u0) ≈ pendulum(u0, nothing, nothing)

    # Simulate
    estimator = ODEProblem(dynamics(Ψ), u0, tspan, parameters(Ψ))
    sol_ = solve(estimator,Tsit5(), saveat = dt)
    @test sol[:,:] ≈ sol_[:,:]

    opt = ADMM(1e-2, 0.7)
    Ψ = SINDy(sol[:,:], DX, basis, opt, maxiter = 5000)
    @test_nowarn set_threshold!(opt, 1e-2)

    # Simulate
    estimator = ODEProblem(dynamics(Ψ), u0, tspan, parameters(Ψ))
    sol_2 = solve(estimator,Tsit5(), saveat = dt)
    @test norm(sol[:,:] - sol_2[:,:], 2) < 2e-1

    opt = SR3(1e-2, 1.0)
    Ψ = SINDy(sol[:,:], DX, basis, opt, maxiter = 5000)
    @test_nowarn set_threshold!(opt, 0.1)

    # Simulate
    estimator = ODEProblem(dynamics(Ψ), u0, tspan, parameters(Ψ))
    sol_3 = solve(estimator,Tsit5(), saveat = dt)
    @test norm(sol[:,:] - sol_3[:,:], 2) < 1e-1

    # Now use the threshold adaptation
    opt = STRRidge(1e-2)
    λs = exp10.(-7:0.1:-1)
    Ψ = SINDy(sol[:,:], DX[:, :], basis, λs, opt, maxiter = 100)
    estimator = ODEProblem(dynamics(Ψ), u0, tspan, parameters(Ψ))
    sol_4 = solve(estimator, Tsit5(), saveat = dt)
    @test norm(sol[:,:] - sol_4[:,:], 2) < 1e-1

    # Check for errors
    @test_nowarn SINDy(sol[:,:], DX[1,:], basis, λs, opt, maxiter = 1)
    @test_nowarn SINDy(sol[:, :], DX[1, :], basis, λs, opt, maxiter = 1, denoise = true, normalize = true)

    # Check with noise
    X = sol[:, :] + 1e-3*randn(size(sol[:,:])...)
    opt = SR3(1e-1, 10.0)
    Ψ = SINDy(X, DX, basis, λs, opt, maxiter = 10000, denoise = true, normalize = true)

    estimator = ODEProblem(dynamics(Ψ), u0, tspan, parameters(Ψ))
    sol_5 = solve(estimator,Tsit5(), saveat = dt)
    @test norm(sol[:,:] - sol_5[:,:], 2) < 5e-1

    # Check sparse_regression
    X .= Array(sol)
    opt = SR3(1e-1, 1.0)
    maxiter = 5000
    θ = basis(X)
    Ξ1, iters_1 = sparse_regression(X, DX, basis, parameters(basis), [], maxiter, opt, true, true, eps())
    Ξ2 = similar(Ξ1)
    Ξ3 = similar(Ξ1)
    iters_2 = sparse_regression!(Ξ2, X, DX, basis, parameters(basis), [], maxiter, opt, true, true, eps())
    iters_3 = sparse_regression!(Ξ3, θ, DX, maxiter, opt, true, true, eps())

    @test iters_1 == iters_2 == iters_3
    @test Ξ1 ≈ Ξ2 ≈ Ξ3
    @test isapprox(norm(Ξ1'*θ - DX,2), 10.18; atol = 1e-1)
end


const LOWERA = codepoint('a')
const LOWERZ = codepoint('z')
const UPPERA = codepoint('A')
const UPPERZ = codepoint('Z')

function rotate(n::Int, s::String)
    String(wraprot.(n,codeunits(s)))
end

# Prettier, but 25% slower for big strings:
# rotate(n::Int, s::String) = map(c->rotate(n,c),s)

rotate(n,c::Char) = Char(wraprot(n,codepoint(c)))

function wraprot(n::Int,c::T)::UInt8 where T <: Unsigned
    if LOWERA <= c <= LOWERZ
        rem(c - T(LOWERA) + T(n), T(26)) + T(LOWERA)
    elseif UPPERA <= c <= UPPERZ
        rem(c - T(UPPERA) + T(n), T(26)) + T(UPPERA)
    else
        c
    end
end

# Shamelessly cribbed from base/reducedim.jl
for n in 1:26
    name = Symbol("R$(n)_str")
    @eval begin
         macro $name(s)
            rotate($n,s)
        end
    end
end

# For testing:
# using BenchmarkTools
# @benchmark rotate(23,s) setup=(s=rand('a':'z',100_000)|>join)


module MyTestRepo

using Plots

function lsq_fit(x, y)
    @assert length(x) == length(y)
    N = length(x)
    X = zeros(N, 2)
    for j=1:N
        X[j,1] = 1.0
        X[j,2] = x[j]
    end
    C = X'*X
    alpha = inv(C)*X'*y
    return alpha
end

function plot_fit(x, y, fn)
    # x = [0.0, 2.0, 4.0]
    # y = [0.0, 3.0, 5.0]
    α0, α1 = lsq_fit(x, y)
    scatter(x, y, color=:black, label="Data")
    plot!(x, α0 .+ α1*x, color=:black, lw=3)
    savefig(fn)
end

export lsq_fit, plot_fit

#=
comment
long comment
=#

end


using Zygote, NNlib, Test
using Zygote: gradient
using NNlib: conv
import Random

function ngradient(f, xs::AbstractArray...)
  grads = zero.(xs)
  for (x, Δ) in zip(xs, grads), i in 1:length(x)
    δ = sqrt(eps())
    tmp = x[i]
    x[i] = tmp - δ/2
    y1 = f(xs...)
    x[i] = tmp + δ/2
    y2 = f(xs...)
    x[i] = tmp
    Δ[i] = (y2-y1)/δ
  end
  return grads
end

gradcheck(f, xs...) =
  all(isapprox.(ngradient(f, xs...),
                gradient(f, xs...), rtol = 1e-5, atol = 1e-5))

gradtest(f, xs::AbstractArray...) = gradcheck((xs...) -> sum(sin.(f(xs...))), xs...)
gradtest(f, dims...) = gradtest(f, rand.(Float64, dims)...)

Random.seed!(0)

@test gradtest((x, W, b) -> σ.(W*x .+ b), 5, (2,5), 2)
@test gradtest((x, W, b) -> σ.(W*x .+ b), (5,3), (2,5), 2)
@test gradtest((x, W, b) -> logσ.(W*x .+ b), 5, (2,5), 2)
@test gradtest((x, W, b) -> logσ.(W*x .+ b), (5,3), (2,5), 2)

@test gradtest((w, x) -> w'*x, randn(10, 2), randn(10))
@test gradtest((w, x) -> transpose(w)*x, randn(5,5), randn(5,5))

@test gradtest(x -> sum(x, dims = (2, 3)), (3,4,5))
@test gradtest(x -> prod(x, dims = (2, 3)), (3,4,5))
@test gradtest(x -> prod(x), (3,4,5))

@test gradtest(x -> softmax(x).*(1:3), 3)
@test gradtest(x -> softmax(x).*(1:3), (3,5))
@test gradtest(x -> logsoftmax(x).*(1:3), 3)
@test gradtest(x -> logsoftmax(x).*(1:3), (3,5))

@test gradtest(conv, rand(10, 3, 2), randn(Float64,2, 3, 2))
@test gradtest(conv, rand(10, 10, 3, 2), randn(Float64,2, 2, 3, 2))
@test gradtest(conv, rand(10, 10, 10, 3, 2), randn(Float64,2, 2, 2, 3, 2))

@test gradtest(x -> maxpool(x, (2,2)), rand(10, 10, 3, 2))
@test gradtest(x -> maxpool(x, (2,2,2)), rand(10, 10, 10, 3, 2))

@test gradtest(x -> meanpool(x, (2,2)), rand(10, 10, 3, 2))
@test gradtest(x -> meanpool(x, (2,2,2)), rand(5, 5, 5, 3, 2))

@test gradtest(x -> permutedims(x, [3,1,2]), rand(4,5,6))

@test gradtest(x -> repeat(x; inner=2), rand(5))
@test gradtest(x -> repeat(x; inner=2, outer=3), rand(5))
@test gradtest(x -> repeat(x; inner=(2,2,1), outer=(1,1,3)), rand(5,4,3))

@test gradtest(kron, rand(5), rand(3))
@test gradtest(kron, rand(5), rand(3), rand(8))
@test gradtest(kron, rand(5,1), rand(3,1))
@test gradtest(kron, rand(5,1), rand(3,1), rand(8,1))
@test gradtest(kron, rand(5,2), rand(3,2), rand(8,2))

@testset "maximum" begin
  @test gradtest(maximum, rand(2, 3))

  @test gradtest(x -> maximum(x, dims=1), rand(2, 3))
  @test gradtest(x -> maximum(x, dims=2), rand(2, 3))
  @test gradtest(x -> maximum(x, dims=3), rand(2, 3, 4))

  @test gradtest(x -> maximum(x, dims=[1, 2]), rand(2, 3, 4))
end

@testset "minimum" begin
  @test gradtest(minimum, rand(2, 3))

  @test gradtest(x -> minimum(x, dims=1), rand(2, 3))
  @test gradtest(x -> minimum(x, dims=2), rand(2, 3))
  @test gradtest(x -> minimum(x, dims=3), rand(2, 3, 4))

  @test gradtest(x -> minimum(x, dims=[1, 2]), rand(2, 3, 4))
end


using SamplingHS3

struct gm_test
    gm::SamplingHS3.GaussMixture
    gm_p::Vector{SamplingHS3.GaussParam}

    G::Array{Float64,2}
    h::Array{Float64,1}
    vol::Float64
end

function test_gauss_mixture( means , sigmas , bounds )
    gmx = Vector{SamplingHS3.GaussParam}(length(means))
    for zi=1:length(means)
        gmx[zi] = SamplingHS3.GaussParam(means[zi],sigmas[zi])
    end
    bb1 = SamplingHS3.get_bounding_box(bounds[:,1],bounds[:,2])
    G   = bb1[1]
    h   = bb1[2]
    vol = prod( bounds[:,2]-bounds[:,1] )
    return gm_test( SamplingHS3.GaussMixture(gmx,ones(length(gmx))),gmx,G,h , vol)
end

function run_gm_test( gmt , n_live , n_steps_decorr , sfactor , time_s )
    d = size(gmt.G,2)
    f_cost_gm = (x) -> -log(SamplingHS3.pdf_gaussmixture_nonnormalized(gmt.gm,x)+1.e-100)
    X_init = SamplingHS3.sample_unif_hr( gmt.G , gmt.h , zeros(d) , n_live , 2000 )
    ns_result = SamplingHS3.nested_sampling( X_init , f_cost_gm , gmt.G , gmt.h , SamplingHS3.NSConfig(n_steps_decorr,sfactor,SamplingHS3.StopCriterion(max_time=time_s),SamplingHS3.MCMCConfig()) )

    ns_x_sequence = reduce( (x,y) -> [x y] , map( x -> x.X_sorted_out , ns_result.ns_iter_data ) )

    # test reconstruction of live samples:
    ns_x_livesamples = SamplingHS3.reconstruct_live_samples(ns_result.ns_iter_data)
    # test resampling:
    rs_result = SamplingHS3.resample_ns_data( ns_result.ns_iter_data )

    #scatter( ns_x_sequence[1,:] , ns_x_sequence[2,:] , ns_x_sequence[3,:] , ma = 0.75 , zcolor=(1:size(ns_x_sequence,2)) , m=2, msw=0 , c=:lightrainbow , lab="grad" )
    ei_result = SamplingHS3.eval_integral(gmt.vol,ns_result.ns_iter_data)
    return ei_result[1]
end

gm_test_a = test_gauss_mixture( [[4;4;4],[-4;-4;-4]] , [eye(3),eye(3)] , repmat([-10 10],3,1) )
gm_test_a_results = Vector{Float64}(10)
for zi=1:length(gm_test_a_results)
    ival = run_gm_test( gm_test_a , 200 , 6 , 0.8 , 2. )
    gm_test_a_results[zi] = ival
end


# Utilities we need. unzip is being contributed to ZipFile,
# the more memory-efficient lda is being offered to MultivariateStats.
function unzip(inputfilename, outputpath=pwd())
    r = ZipFile.Reader(inputfilename)
    for f in r.files
        outpath = joinpath(outputpath, f.name)
        if isdirpath(outpath)
            mkpath(outpath)
        else
            open(outpath, "w") do io
                write(io, read(f))
            end
        end
    end
    nothing
end

## Memory-efficient LDA
# These are large vectors, and storing the full covariance matrix is
# problematic. This algorithm works in terms of a projection onto a
# subspace spanning the means, which also serves to regularize. See
#   Yu, Hua, and Jie Yang. "A direct LDA algorithm for
#      high-dimensional data—with application to face recognition."
#      Pattern recognition 34.10 (2001): 2067-2070.
function lda(X::AbstractMatrix{T}, group::Vector{Int}) where {T}
    nd = size(X,1)
    dmeans, dX = lda_prepare(X, group)
    UB, SB, _ = svd(dmeans, thin=true)
    ikeep = SB .>= sqrt(eps(T))*maximum(SB)
    Y = UB[:,ikeep]
    Z = scale(Y, 1 ./ SB[ikeep])
    projW = Z'*dX
    UW, SW, _ = svd(projW, thin=true)
    eigenval = 1 ./ SW.^2
    eigenvec = Z*scale(UW, 1 ./ SW)
    # Normalize
    for j = 1:size(eigenvec,2)
        evn = zero(T)
        for i = 1:nd
            evn += eigenvec[i,j]^2
        end
        evn = sqrt(evn)
        for i = 1:nd
            eigenvec[i,j] /= evn
        end
    end
    # Start with largest eigenvalues first
    return eigenvec[:,end:-1:1], eigenval[end:-1:1]
end

function lda_prepare(X::AbstractMatrix{T}, group::Vector{Int}) where {T}
    nd = size(X,1)
    npoints = size(X,2)
    xbar = zeros(T,nd)
    ngroups = maximum(group)
    if ngroups > npoints
        error("The group indices must be consecutive, starting at 1")
    end
    ngroups > 1 || error("Must have more than one group")
    # Calculate the group means
    means = zeros(T,nd,ngroups)
    npergroup = zeros(Int,ngroups)
    for j = 1:npoints
        k = group[j]
        for i = 1:nd
            means[i,k] += X[i,j]
        end
        npergroup[k] += 1
    end
    keep = npergroup .> 0
    if !all(keep)
        npergroup = npergroup[keep]
        means = means[:,keep]
    end
    gindex = cumsum(keep)
    ngroups = length(npergroup)
    for j = 1:ngroups
        n = npergroup[j]
        for i = 1:nd
            xbar[i] += means[i,j]
            means[i,j] /= n
        end
    end
    for i = 1:nd
        xbar[i] /= npoints
    end
    # Compute differences from the means
    dX = Array(T, nd, npoints)
    for j = 1:npoints
        k = gindex[group[j]]
        for i = 1:nd
            dX[i,j] = X[i,j] - means[i,k]
        end
    end
    dmeans = means .- xbar
    dmeans, dX
end


using Test, LinearAlgebra
using OrdinaryDiffEq, DiffEqBase

const ALGOS = (
    (SymplecticEuler, true, 1),
    (VelocityVerlet, false, 2),
    (VerletLeapfrog, true, 2),
    (PseudoVerletLeapfrog, true, 2),
    (McAte2, true, 2),
    (Ruth3, true, 3),
    (McAte3, true, 3),
    (CandyRoz4, true, 4),
    (McAte4, true, 4),
    (CalvoSanz4, true, 4),
    # (McAte42, true, 4), known to be broken
    (McAte5, true, 5),
    (Yoshida6, true, 6),
    (KahanLi6, true, 6),
    (McAte8, true, 8),
    (KahanLi8, true, 8),
    (SofSpa10, true, 10),
)

function dp(p, q, pa, t)
    0q .+ pa[2]
end

function dq(p, q, pa, t)
    p .* pa[1]
end

dp(res, p, q, pa, t) = (res .= dp(p, q, pa, t))
dq(res, p, q, pa, t) = (res .= dq(p, q, pa, t))

dynode(iip, dp, dq) = DynamicalODEFunction{iip}(dp, dq)

# [0:1] used in dp, dq; [3:4] start values for p0, q0
const PARAMS = ((1.0, 0.1, 1.0, 0.0), (0.1, 1.0, 1.0, -1.0))
const IIPS = (true, false)
const TSPAN = (0.0, 1.0)

solution(t, w) = (w[2] * t + w[3], (w[2] / 2 * t + w[3]) * w[1] * t + w[4])
apa(iip::Bool, x) = iip ? vcat.(x) : x
errorbound(dt, d, x) = 100 * abs(dt)^d + 1000 * eps(norm(x))
function printerrors(text, calc, solution, pa, t1)
    print(text, ": ")
    print(norm(calc[1] - solution(t1, pa)[1]), " ")
    print(norm(calc[2] - solution(t1, pa)[2]))
    println()
end

@testset "symplectic $alg-$iip-$pa" for (alg, x, d) in ALGOS, iip in IIPS, pa = PARAMS

    dt = 0.01
    tspan = TSPAN
    t0, t1 = tspan
    dynfun = dynode(iip, dp, dq)
    p0, q0 = apa(iip, solution(t0, pa))
    prob = DynamicalODEProblem(dynfun, p0, q0, tspan, pa)

    if x || pa[1] == 1
        sol = solve(prob, alg(); dt = dt)
        calc = sol(t1)
        # printerrors("$alg-$iip-$pa", calc, solution, pa, t1)
        @test calc[1] ≈ solution(t1, pa)[1] rtol = errorbound(dt, d, calc[1])
        @test calc[2] ≈ solution(t1, pa)[2] rtol = errorbound(dt, d, calc[2])
    else
        @test_throws ArgumentError solve(prob, alg(); dt = dt)
    end
end


# Solver for homogeneous sites
@doc raw"""
    sumBassHomogeneousSolver(p,q,m,c0,alpha,F,T,n,Pi)

This function maximizes the sum of homogeneous Bass functions and returns the maximizer
and maximum. Note that any permutation of the maximizer is also a maximizer.
"""
function sumBassHomogeneousSolver(p::Real,q::Real,m::Real,c0::Real,alpha::Real,F::Real,T::Real,n::Int,Pi::Real)
    if n*F > Pi
        throw(DomainError(Pi,"the problem is unfeasible because the capacity is too small."))
    end
    if n*T <= Pi # Trivial solution
        x = fill(float(T),n)
        v = alpha*n*Bass(T,p,q,m,c0)
        return x,v
    end

    # Case 1
    vstar = 0.0
    xstar = 0
    for i in 1:n # This loop ensures that case does not violate the constraint x<T
        if F+(Pi-F*n)/i < T
            objective = [j * (Bass(F + (Pi-F*n)/j,p,q,m,c0) - Bass(F,p,q,m,c0)) for j in i:n]
            xstar = argmax(objective) + i - 1
            vstar = alpha*(xstar*Bass(F + (Pi-F*n)/xstar,p,q,m,c0) + (n-xstar)*Bass(F,p,q,m,c0))
            break
        end
    end

    # Case 2
    xstar2 = floor(Int,(Pi-n*F)/(T-F))
    vstar2 = alpha*(xstar2*Bass(T,p,q,m,c0) + Bass(Pi-xstar2*T-(n-xstar2-1)*F,p,q,m,c0) +
        (n-xstar2-1)*Bass(F,p,q,m,c0))

    x = Vector{Float64}(undef,n)
    if vstar >= vstar2
        x[1:xstar] .= F + (Pi-F*n)/xstar
        x[(xstar+1):n] .= F
        return x,vstar
    else
        x[1:xstar2] .= T
        x[xstar2+1] = Pi-xstar2*T-(n-xstar2-1)*F
        x[xstar2+2:n] .= F
        return x,vstar2
    end

end


@testset "irreducibles $p" for p = (2, 3)
    Z = ZZ/p
    n = 5
    irrn = collect(irreducibles(Z[:x], n))
    @test irreducible(Z[:x], n) == irrn[1]
    @test length(irrn) == necklace(p, n)
    m = 7
    irrm = collect(irreducibles(Z[:x], m))
    f = prod(irrn[2:4]) * prod(irrm[3:4]) * irrn[2]
    @test isreducible(f)
    ff = factor(f)
    @test length(ff) == 5
    @test prod(ff) == f

    @test length(irrn) == num_irreducibles(Z[:x], n)
    @test length(irrn) == necklace(p, n)
    @test length(irrm) == necklace(p, m)

    redn = collect(reducibles(Z[:x], n))
    @test length(redn) + length(irrn) == p^n
end

@testset "factor non-monic" begin
    Z = ZZ/3
    x = monom(Z[:x])
    p = (-x+1)^9
    @test prod(factor(p)) == p
    p = zero(p)
    @test factor(p) == [p => 1]
    p = one(p)
    @test factor(p) == [p => 1]
    p = x + 1
    @test factor(p) == [p => 1]
    p = -x + 1
    @test factor(p) == [-one(x) => 1, x - 1 => 1]
    p = -(x+1)^3
    @test factor(p) == [-one(x) => 1, x + 1 => 3]
end


# # Conway's game of life

# ![](game_of_life.gif)

# [Game of life on wikipedia](https://en.wikipedia.org/wiki/Conway%27s_Game_of_Life).

# It is also available from the `Models` module as [`Models.game_of_life`](@ref).

using Agents, Plots
using Random # hide
gr(); # hide
nothing # hide

# ## 1. Define the rules

# Rules of Conway's game of life: DSRO (Death, Survival, Reproduction, Overpopulation).
# Cells die if the number of their living neighbors is <D or >O,
# survive if the number of their living neighbors is ≤S,
# come to life if their living neighbors are  ≥R and ≤O.
rules = (2, 3, 3, 3)
nothing # hide

# ## 2. Build the model

# First, define an agent type. It needs to have the compulsary `id` and `pos` fields, as well as an `status` field that is `true` for cells that are alive and `false` otherwise.

mutable struct Cell <: AbstractAgent
    id::Int
    pos::Dims{2}
    status::Bool
end

# The following function builds a 2D cellular automaton. `rules` is of type `Tuple{Int,Int,Int,Int}` representing DSRO.

# `dims` is a tuple of integers determining the width and height of the grid environment.
# `metric` specifies whether cells connect to their diagonal neighbors.

# This function creates a model where all cells are "off".

function build_model(; rules::Tuple, dims = (100, 100), metric = :chebyshev, seed = 120)
    space = GridSpace(dims; metric)
    properties = Dict(:rules => rules)
    model = ABM(Cell, space; properties, rng = MersenneTwister(seed))
    idx = 1
    for x in 1:dims[1]
        for y in 1:dims[2]
            add_agent_pos!(Cell(idx, (x, y), false), model)
            idx += 1
        end
    end
    return model
end
nothing # hide

# Now we define a stepping function for the model to apply the rules to agents.

function ca_step!(model)
    new_status = fill(false, nagents(model))
    for agent in allagents(model)
        nlive = nlive_neighbors(agent, model)
        if agent.status == true && (nlive ≤ model.rules[4] && nlive ≥ model.rules[1])
            new_status[agent.id] = true
        elseif agent.status == false && (nlive ≥ model.rules[3] && nlive ≤ model.rules[4])
            new_status[agent.id] = true
        end
    end

    for k in keys(model.agents)
        model.agents[k].status = new_status[k]
    end
end

function nlive_neighbors(agent, model)
    neighbor_positions = nearby_positions(agent, model)
    all_neighbors =
        Iterators.flatten(ids_in_position(np, model) for np in neighbor_positions)
    sum(model[i].status == true for i in all_neighbors)
end
nothing # hide

# now we can instantiate the model:
model = build_model(rules = rules, dims = (50, 50))

# Let's make some random cells on
for i in 1:nagents(model)
    if rand(model.rng) < 0.2
        model.agents[i].status = true
    end
end

# ## 3. Animate the model

# We use the [`plotabm`](@ref) function to create an animation.

cd(@__DIR__) #src
ac(x) = x.status == true ? :black : :white
anim = @animate for i in 0:100
    i > 0 && step!(model, dummystep, ca_step!, 1)
    p1 = plotabm(model; ac = ac, as = 3, am = :square, showaxis = false)
end
nothing # hide

# We can now save the animation to a gif.

gif(anim, "game_of_life.gif", fps = 5)