using ParametricSurfaces # for generating some simple point distributions
using LoopVectorization  # for automatically vectorizing some typical kernels

# some point distributions for testing
function cube_uniform_surface_mesh(dx)
    geo = ParametricSurfaces.Cube() # a cube with low corner at (0,0,0) and high corner (2,2,2)
    range = -1+dx/2:dx:1-dx/2
    pts   = Vector{Geometry.Point3D}()
    for ent in boundary(geo)
        for u in range, v in range
            push!(pts,ent((u,v)))
        end
    end
    return pts
end

function spheroid(N,r,center,zstretch)
    pts   = Vector{Geometry.Point3D}(undef,N)
    phi   = π*(3-sqrt(5)) # golden angle in radians
    for i = 1:N
        ytmp   = 1 - ((i-1)/(N-1))*2
        radius = sqrt(1-ytmp^2)
        theta    = phi*i
        x = cos(theta) * radius * r + center[1]
        y = ytmp * r + center[2]
        z = zstretch * sin(theta) * radius * r + center[3]
        pts[i] = SVector(x,y,z)
    end
    # Output
    return pts
end

sphere_uniform(N,r,center=(0,0,0)) = spheroid(N,r,center,1)

function cube_nonuniform_surface_mesh(dx,p=2)
    geo = ParametricSurfaces.Cube() # a cube with low corner at (0,0,0) and high corner (2,2,2)
    range = -1+dx/2:dx:1-dx/2
    pts   = Vector{Geometry.Point3D}()
    # degenerate change of variables mapping [0,1] onto itself with p
    # derivatives vanishing at each endpoint
    cov   = Integration.KressP(order=p)
    # map cov to act on [-1,1]
    χ     = (u) -> begin
        û = 0.5*(u+1)
        ŝ = cov(û)
        s = 2*ŝ - 1
    end
    for ent in boundary(geo)
        for u in range, v in range
            push!(pts,ent((χ(u),χ(v))))
        end
    end
    return pts
end

function cube_uniform_volume_mesh(dx)
    range = dx/2:dx:1-dx/2
    [SVector(a,b,c) for a in range, b in range, c in range] |> vec
end

function sphere_uniform_surface_mesh(dx)
    geo   = ParametricSurfaces.Sphere(radius=1) # a sphere of radius 1 centered at (0,0,0)
    range = -1+dx/2:dx:1-dx/2
    pts   = Vector{Geometry.Point3D}()
    for ent in boundary(geo)
        for u in range, v in range
            push!(pts,ent((u,v)))
        end
    end
    return pts
end

# define some typical kernels
struct HelmholtzKernel
    k::Float64
end

function (K::HelmholtzKernel)(x,y)
    d = norm(x-y)
    v = exp(im*K.k*d)/(4*π*d)
    return (!iszero(d))*v
end

IFGF.wavenumber(K::HelmholtzKernel) = K.k

function IFGF.near_interaction!(C,K::HelmholtzKernel,X,Y,σ,I,J)
    Tx = eltype(X)
    Ty = eltype(Y)
    @assert Tx <: SVector && Ty <: SVector
    Xm = reshape(reinterpret(eltype(Tx),X),3,:)
    Ym = reshape(reinterpret(eltype(Ty),Y),3,:)
    @views helmholtz3d_sl_vec!(C[I],Xm[:,I],Ym[:,J],σ[J],K.k)
end

function IFGF.transfer_factor(K::HelmholtzKernel,x,Y)
    yc  = IFGF.center(Y)
    yp  = IFGF.center(IFGF.parent(Y))
    d   = norm(x-yc)
    dp  = norm(x-yp)
    exp(im*K.k*(d-dp))*dp/d
end

function helmholtz3d_sl_vec!(C,X,Y,σ,k)
    m,n = size(X,2), size(Y,2)
    C_T = reinterpret(Float64, C)
    C_r = @views C_T[1:2:end,:]
    C_i = @views C_T[2:2:end,:]
    σ_T = reinterpret(Float64, σ)
    σ_r = @views σ_T[1:2:end,:]
    σ_i = @views σ_T[2:2:end,:]
    @turbo for j in 1:n
        for i in 1:m
            d2 = (X[1,i] - Y[1,j])^2
            d2 += (X[2,i] - Y[2,j])^2
            d2 += (X[3,i] - Y[3,j])^2
            d  = sqrt(d2)
            s, c = sincos(k * d)
            zr = inv(4π*d) * c
            zi = inv(4π*d) * s
            C_r[i] += (!iszero(d))*(zr*σ_r[j] - zi*σ_i[j])
            C_i[i] += (!iszero(d))*(zi*σ_r[j] + zr*σ_i[j])
        end
    end
    return C
end

struct LaplaceKernel
end

function (K::LaplaceKernel)(x,y)
    d = norm(x-y)
    v = 1/(4*π*d)
    return (!iszero(d))*v
end

function IFGF.near_interaction!(C,K::LaplaceKernel,X,Y,σ,I,J)
    Xm = reshape(reinterpret(Float64,X),3,:)
    Ym = reshape(reinterpret(Float64,Y),3,:)
    @views laplace3d_sl_vec!(C[I],Xm[:,I],Ym[:,J],σ[J])
end

function laplace3d_sl_vec!(C,X,Y,σ)
    m,n = size(X,2), size(Y,2)
    @turbo for j in 1:n
        for i in 1:m
            d2 = (X[1,i] - Y[1,j])^2
            d2 += (X[2,i] - Y[2,j])^2
            d2 += (X[3,i] - Y[3,j])^2
            # fast invsqrt code taken from here
            # https://benchmarksgame-team.pages.debian.net/benchmarksgame/program/nbody-julia-8.html
            invd    = @fastmath Float64(1 / sqrt(Float32(d2)))
            invd = 1.5invd - 0.5d2 * invd * (invd * invd)
            C[i] += (!iszero(d2))*(inv(4π)*invd*σ[j])
            # C[i] += inv(4π*sqrt(d2))*σ[j] # significalty slower
        end
    end
    return C
end

struct MaxwellKernel
    k::Float64
end

function (K::MaxwellKernel)(x,y)
    r   = x - y
    d   = norm(r)
    # helmholtz greens function
    g   = exp(im*K.k*d)/(4π*d)
    gp  = im*K.k*g - g/d
    gpp = im*K.k*gp - gp/d + g/d^2
    RRT = r*transpose(r) # rvec ⊗ rvecᵗ
    G   = g*LinearAlgebra.I + 1/K.k^2*(gp/d*LinearAlgebra.I + (gpp/d^2 - gp/d^3)*RRT)
    return (!iszero(d))*G
end

IFGF.wavenumber(K::MaxwellKernel) = K.k

# function IFGF.centered_factor(K::MaxwellKernel,x,Y)
#     yc = IFGF.center(Y)
#     r  = x-yc
#     d = norm(r)
#     g   = exp(im*k*d)/(4π*d)
#     return g    # works fine!
# end

# function IFGF.transfer_factor(K::HelmholtzKernel,x,Y)
#     yc  = IFGF.center(Y)
#     yp  = IFGF.center(IFGF.parent(Y))
#     d   = norm(x-yc)
#     dp  = norm(x-yp)
#     exp(im*K.k*(d-dp))*dp/d
# end

# function helmholtz3d_sl_vec!(C,X,Y,σ,k)
#     m,n = size(X,2), size(Y,2)
#     C_T = reinterpret(Float64, C)
#     C_r = @views C_T[1:2:end,:]
#     C_i = @views C_T[2:2:end,:]
#     σ_T = reinterpret(Float64, σ)
#     σ_r = @views σ_T[1:2:end,:]
#     σ_i = @views σ_T[2:2:end,:]
#     @turbo for j in 1:n
#         for i in 1:m
#             d2 = (X[1,i] - Y[1,j])^2
#             d2 += (X[2,i] - Y[2,j])^2
#             d2 += (X[3,i] - Y[3,j])^2
#             d  = sqrt(d2)
#             s, c = sincos(k * d)
#             zr = inv(4π*d) * c
#             zi = inv(4π*d) * s
#             C_r[i] += (!iszero(d))*(zr*σ_r[j] - zi*σ_i[j])
#             C_i[i] += (!iszero(d))*(zi*σ_r[j] + zr*σ_i[j])
#         end
#     end
#     return C
# end


__precompile__(true)

module TransformUtils

using LinearAlgebra
using ManifoldsBase
using Manifolds
using StaticArrays
using DocStringExtensions

import Base: convert, promote_rule, *, \, vec
import LinearAlgebra: transpose, adjoint, normalize, normalize!
import ManifoldsBase: vee, vee!, hat, hat!

export
  Quaternion,
  AngleAxis,
  AxisAngle,
  SO3,
  so3,
  SE3,
  Euler,
  skew,
  vee!,
  vee,
  vec,
  setso3!,
  *,
  transpose,
  matrix,
  matrix!,
  inverse,
  compare,
  normalize!,
  normalize,
  q_conj,
  q_conj!,
  convert,
  convert!,
  rotate!,
  rotate,
  wrapRad,
  logmap,
  logmap_SO2,
  rightJacExmap,
  rightJacExmapinv,
  deltaso3vee,
  # Should be good to go
  veeAngleAxis,
  veeQuaternion,
  veeEuler,
  A_invB,
  ominus,
  oplus,
  ⊖,
  ⊕,
  \,
  doinversetransform!,

  # type aliases
  FloatInt,
  VectorFloatInt,

  # basic Taylor exported to try against Pade version in expm
  expmOwn,
  expmOwn1,
  expmOwnT,
  expmOwn2,
  expmOwn3,
  expmOwn4,

  # TODO -- needs refactoring
  R,
  SE2,
  se2vee,
  se2vee!,
  TU,
  veePose3,
  veePose

const TU = TransformUtils
const LA = LinearAlgebra

include("LazyLieManifolds.jl")
include("Legacy.jl")
include("LegacySpecialEuclidean.jl")

end #module


@testset "points.jl" begin
    p1 = Point(80.0, 190.0)
    p2 = Point(40.0, -10.0)

    mp = midpoint(p1, p2)
    @test mp.ϕ ≈ 69.58472222222 atol = 0.001
    @test mp.λ ≈ -15.62638888888 atol = 0.001
    @test midpoint(Arc(p1, p2)).ϕ ≈ 69.58472222222 atol = 0.001

    ip = intermediate_point(p1, p2, 0.5)
    @test ip.ϕ ≈ 69.58472222222 atol = 0.001
    @test ip.λ ≈ -15.62638888888 atol = 0.001

    p3 = Point(0.0, -10.0)
    p4 = Point(0.0, 0.0)
    p5 = Point(0.0, 10.0)
    lss = Arcs([p3, p4, p5])
    ip_lss = intermediate_point(lss, 0.6)
    @test ip_lss.ϕ ≈ 0.0 atol = 0.01
    @test ip_lss.λ ≈ 2.0 atol = 0.01
    @test intermediate_point(lss, 0.0).λ ≈ -10.0 atol = 0.01
    @test intermediate_point(lss, 1.0).λ ≈ 10.0 atol = 0.01
    @test midpoint(lss).λ ≈ 0.0 atol = 0.01

    ad = 89.9321605919
    @test SphericalGeometry.normalize(destination_point(p1, ad, 45.0).λ, -180.0, 180.0) ≈ -35.44694444444 atol = 0.001
    @test destination_point(Line(p1, 45.0), ad).ϕ ≈ 7.12027777777 atol = 0.001

    intp = intersection_point(p4, p2, 45.0, 20.0)
    @test intp.ϕ ≈ 27.43166666666 atol = 0.001
    @test intp.λ ≈ 148.73138888888 atol = 0.001
    intp2 = intersection_point(Line(p4, 45.0), Line(p2, 20.0))
    @test intp2.ϕ ≈ 27.43166666666 atol = 0.001
    @test intp2.λ ≈ 148.73138888888 atol = 0.001

    p1ni = Point(50.0, 10.0)
    p2ni = Point(50.0, 30.0)
    p3ni = Point(40.0, 20.0)
    p4ni = Point(60.0, 20.0)

    @test intersection_point(p1ni, p2ni, p3ni, p4ni).ϕ ≈ 50.43250 atol = 0.01
    @test intersection_point(p1ni, p2ni, p3ni, p4ni).λ ≈ 20.0 atol = 0.0001
    @test intersection_point(p1ni, p2ni, p3ni, p2ni).λ ≈ 30.0 atol = 0.0001
    @test intersection_point(Arc(p1ni, p2ni), Arc(p3ni, p2ni)).λ ≈ 30.0 atol = 0.0001

    pa11 = Point(20.0, 10.0)
    pa12 = Point(90.0, 60.0)
    pa21 = Point(10.0, 50.0)
    pa22 = Point(80.0, 5.0)

    @test intersection_point(pa11, pa12, pa21, pa22).λ ≈ 10.0 atol = 0.1
    @test intersection_point(pa11, pa12, pa21, pa22).ϕ ≈ 79.06854 atol = 0.01

    p4ni2 = Point(49.0, 20.0)
    @test isnan(intersection_point(p1ni, p2ni, p3ni, p4ni2).λ)

    # p1a = Point(10.0374, 0.0)
    # p2a = Point(-90.0, 0.0)
    # p3a = Point(0.0, 0.0)
    # @test isinf(intersection_point(p3a, p2a, p1a, p3a).ϕ)

    p1ip1 = Point(5.0, -10.0)
    p2ip1 = Point(5.0, -5.0)
    p3ip1 = Point(5.0, 5.0)
    p4ip1 = Point(5.0, 6.0)
    p5ip1 = Point(5.0, 11.0)
    p6ip1 = Point(-1.0, 5.0)
    p1ip2 = Point(0.0, 0.0)
    p2ip2 = Point(0.0, 10.0)
    p3ip2 = Point(10.0, 10.0)
    p4ip2 = Point(10.0, 0.0)
    p5ip2 = Point(1.0, 0.0)

    @test length(intersection_points(Arcs([p1ip1, p2ip1, p3ip1, p4ip1, p5ip1,
    p6ip1]), Arcs([p1ip2, p2ip2, p3ip2, p4ip2, p5ip2]))) == 4
    @test length(intersection_points(Arcs([p1ip1, p2ip1, p3ip1, p4ip1, p5ip1,
    p6ip1]), Polygon(p1ip2, [p1ip2, p2ip2, p3ip2, p4ip2, p5ip2]))) == 4
    @test length(intersection_points(Polygon(p1ip2, [p1ip1, p2ip1, p3ip1, p4ip1, p5ip1,
    p6ip1]), Polygon(p1ip2, [p1ip2, p2ip2, p3ip2, p4ip2, p5ip2]))) == 6

    #concave polygon
    pc1 = Point(64.0, 27.0)
    pc2 = Point(47.0, 0.0)
    pc3 = Point(42.0, 25.0)
    pc4 = intermediate_point(pc2,pc3,0.33)
    pc5 = Point(52.0, 16.0)
    pc6 = intermediate_point(pc2,pc3,0.67)
    pc7 = Point(55.0, -7.0)
    pc8 = intermediate_point(pc2,pc1,0.33)
    pc9 = Point(62.0, 4.0)
    pc10 = intermediate_point(pc2,pc1,0.67)
    pc11 = Point(56.0, 21.0)

    polygon_concave = Polygon(pc11,
    [pc1,pc3,pc6,pc5,pc4,pc2,pc7,pc8,pc9,pc10])

    @test length(intersection_points(Arc(Point(63.0, 27.0), Point(65.0, 27.0)), polygon_concave)) == 1
    @test length(intersection_points(Arc(Point(63.0, 28.0), Point(65.0, 28.0)), polygon_concave)) == 0
    @test length(intersection_points(Arc(Point(62.0, 4.0), Point(55.0, -7.0)), polygon_concave)) == 2
    @test length(intersection_points(Arc(Point(56.0, 21.0), Point(55.0, 21.0)), polygon_concave)) == 0
    @test length(intersection_points(Arc(Point(64.0, 26.0), Point(47.0, -1.0)), polygon_concave)) == 4
    @test length(intersection_points(Arc(Point(56.0, 21.0), Point(52.0, 16.0)), polygon_concave)) == 1
    @test length(intersection_points(Arc(Point(52.0, 16.0), Point(56.0, 21.0)), polygon_concave)) == 1
    @test length(intersection_points(Arc(Point(56.0, 21.0), Point(56.0, 30.0)), polygon_concave)) == 1

    pt1 = Point(20.0, -5.0)
    pt2 = Point(21.0, 5.0)
    pt3 = Point(22.0, 15.0)
    pa = Point(40.0, 5.0)
    pb = Point(30.0, 5.0)
    pc = Point(10.0, 5.0)
    pd = Point(0.0, 5.0)
    arcs123 = Arcs([pt1, pt2, pt3])

    z = intersection_points(arcs123, Arcs([pt2,pc,pd]))

    @test length(intersection_points(arcs123, Arcs([pt1,pc,pd]))) == 1
    @test length(intersection_points(arcs123, Arcs([pt2,pc,pd]))) == 1
    @test length(intersection_points(arcs123, Arcs([pt3,pc,pd]))) == 1
    @test length(intersection_points(arcs123, Arcs([pb,pt1,pc]))) == 1
    @test length(intersection_points(arcs123, Arcs([pb,pt2,pc]))) == 1
    @test length(intersection_points(arcs123, Arcs([pb,pt3,pc]))) == 1
    @test length(intersection_points(arcs123, Arcs([pa,pb,pt1]))) == 1
    @test length(intersection_points(arcs123, Arcs([pa,pb,pt2]))) == 1
    @test length(intersection_points(arcs123, Arcs([pa,pb,pt3]))) == 1

    @test length(intersection_points(arcs123, Arcs([pb,Point(10.0, 4.0),pd]))) == 1
    @test length(intersection_points(arcs123, Arcs([pa,pb,pc]))) == 1

    @test !ison(p1ip1, Arc(p2ip1, p3ip1))
    @test ison(Point(0.0, 5.0), Arc(p1ip2, p2ip2))
    @test ison(Point(0.0, 5.0), Arcs([p1ip2, p2ip2, p3ip2]))
    @test ison(Point(0.0, 5.0), Polygon(p5ip2, [p1ip2, p2ip2, p3ip2]))

    p1i = Point(0.0, 0.0)
    p2i = Point(0.0, 10.0)
    p3i = Point(10.0, 10.0)
    p4i = Point(5.0, 5.0)
    p5i = Point(10.0, 0.0)
    px_inside = Point(7.0, 8.0)

    p6i_out = Point(15.0, 0.0)
    p7i_out = Point(6.0, 5.0)
    p8i_out = Point(6.0, -2.0)
    p9i_in = Point(6.0, 2.0)
    p10i_in = Point(1.0, 6.0)
    p11i_in = Point(1.0, 1.0)

    poly_inside = Polygon(px_inside, [p1i, p2i, p3i, p4i, p5i])
    @test !isinside(p6i_out, poly_inside)
    @test !isinside(p7i_out, poly_inside)
    @test !isinside(p8i_out, poly_inside)
    @test isinside(p9i_in, poly_inside)
    @test isinside(p10i_in, poly_inside)

    @test isinside(Arc(p9i_in, p10i_in), poly_inside)
    @test !isinside(Arc(p8i_out, p10i_in), poly_inside)
    @test !isinside(Arc(p10i_in, p8i_out), poly_inside)
    @test !isinside(Arc(p6i_out, p8i_out), poly_inside)

    @test isinside(Arcs([p9i_in, p10i_in]), poly_inside)
    @test !isinside(Arcs([p8i_out, p10i_in]), poly_inside)
    @test !isinside(Arcs([p10i_in, p8i_out]), poly_inside)
    @test !isinside(Arcs([p6i_out, p8i_out]), poly_inside)

    poly2 = Polygon(Point(2.0, 2.0), [p9i_in, p11i_in, p10i_in])
    @test isinside(poly2, poly_inside)

    poly3 = Polygon(px_inside, [p1i, p2i, p5i, p3i])
    arcs3 = Arcs([p1i, p3i, p2i, p5i])
    @test !isselfintersecting(poly2)
    @test isselfintersecting(poly3)
    @test isselfintersecting(arcs3)
    @test !isselfintersecting(Arcs([p1i, p2i, p3i, p4i, p5i]))

    # https://www.onboardintelligence.com/GC for generating test (sphere)
    pz1 = Point(50.0, 70.0)
    pz2 = Point(50.0, -70.0)
    azi = azimuth(pz1, pz2)
    @test highest_latitude(pz1, azi) ≈ 73.9871 atol = 0.0001
    @test lowest_latitude(pz1, azi) ≈ -73.9871 atol = 0.0001
    @test highest_latitude_point(pz1, pz2).ϕ ≈ 73.9871 atol = 0.0001
    @test lowest_latitude_point(pz1, pz2).ϕ ≈ 50.0 atol = 0.001
    @test highest_latitude_point(pz1, Point(60.0, 70.0)).ϕ ≈ 60.0 atol = 0.001
    @test highest_latitude_point(Point(60.0, 70.0), pz1).ϕ ≈ 60.0 atol = 0.001
    @test lowest_latitude_point(pz1, Point(60.0, 70.0)).ϕ ≈ 50.0 atol = 0.001
    @test lowest_latitude_point(Point(60.0, 70.0), pz1).ϕ ≈ 50.0 atol = 0.001
    @test highest_latitude_point(pz1, Point(60.0, 80.0)).ϕ ≈ 60.0 atol = 0.0001

    azi2 = azimuth(Point(30.0, 89.0), Point(50.0, -60.0))
    @test highest_latitude_point(Point(30.0, 89.0), azi2).ϕ ≈ 73.2626 atol=0.0001
    @test highest_latitude_point(Point(30.0, 89.0), azi2).λ ≈ 8.9986 atol=0.0001
    @test lowest_latitude_point(Point(30.0, 89.0), azi2).λ ≈ -171.0014 atol=0.0001

    @test highest_latitude_point(pz1, azi).ϕ ≈ 73.9871 atol = 0.0001
    @test lowest_latitude_point(pz1, azi).ϕ ≈ -73.9871 atol = 0.0001
    @test highest_latitude_point(pz1, azi).λ ≈ 0.0 atol = 0.0001

    @test highest_latitude_point(Point(30.0, 90.0), 0.0).ϕ ≈ 90.0 atol = 0.1

end


using Test
using Algorithm101
const ALGO = Algorithm101

@testset "Monte Carlo Integral" begin
    f = x -> x
    @test isapprox(ALGO.mc_integral(f, 0:0.00001:1), 0.5, atol=1e-2)

    f = x -> x^2
    @test isapprox(ALGO.mc_integral(f, 0:0.00001:1), 0.333, atol=1e-2)

    f = x -> x * (atan(x))^2
    @test isapprox(ALGO.mc_integral(f, 0:0.00001:1), 0.173, atol=1e-2)

    f = x ->  1 / (1 + (tan(x))^4)
    @test isapprox(ALGO.mc_integral(f, 0:0.00001:1), 0.769, atol=1e-2)
end


"""
    A version of E7 for testing linearization
"""
module E7A

using ModelBaseEcon

model = Model()
model.options.substitutions = false

@parameters model begin
    delta = 0.1000000000000000
    p_dlc_ss = 0.0040000000000000
    p_dlinv_ss = 0.0040000000000000
    p_growth = 0.0040000000000000
end

@variables model begin
    dlc; dlinv; dly; lc; linv;
    lk; ly;
end

@shocks model begin
    dlc_shk; dlinv_shk;
end

@autoexogenize model begin
    lc = dlc_shk
    linv = dlinv_shk
end

@equations model begin
    dlc[t] = (1 - 0.2 - 0.2) * p_dlc_ss + 0.2 * dlc[t - 1] + 0.2 * dlc[t + 1] + dlc_shk[t]
    dlinv[t] = (1 - 0.5) * p_dlinv_ss + 0.1 * dlinv[t - 2] + 0.1 * dlinv[t - 1] + 0.1 * dlinv[t + 1] + 0.1 * dlinv[t + 2] + 0.1 * dlinv[t + 3] + dlinv_shk[t]
    lc[t] = lc[t - 1] + dlc[t]
    linv[t] = linv[t - 1] + dlinv[t]
    ly[t] = log(exp(lc[t]) + exp(linv[t]))
    dly[t] = ly[t] - ly[t - 1]
    lk[t] = log((1 - delta) * exp(lk[t - 1]) + exp(linv[t]))
end

@initialize model

@steadystate model linv = lc - 7;
@steadystate model lc = 14;

end


function LinearAlgebra.cross(a::T1, b::T2) where {T1 <: AbstractPICDataStructure, T2 <: AbstractPICDataStructure}
    @assert size(a) == size(b)
    @assert scalarness(T1) === scalarness(T2) === VectorQuantity() "The cross product is only defined for Vector quantities"
    @assert propertynames(a) == propertynames(b)
    @assert dimensionality(a) == dimensionality(b) == 3 "The cross product is only defined in 3D"

    x_data = @. a.y * b.z - a.z * b.y
    y_data = @. a.z * b.x - a.x * b.z
    z_data = @. a.x * b.y - a.y * b.x

    data = StructArray((x_data, y_data, z_data); names=propertynames(a))
    grid = getdomain(a)
    name = nameof(a) * " × " * nameof(b)
    newstruct(a, data, grid; name)
end


# Set model up
m = Model990(testing = true)
m <= Setting(:alternative_policy, DSGE.taylor99())
Γ0, Γ1, C, Ψ, Π = DSGE.taylor99_eqcond(m)

# Check policy rule
eq = m.equilibrium_conditions
endo = m.endogenous_states
@testset "Check Taylor99 rule has the right coefficients." begin
    @test Γ0[eq[:eq_mp], endo[:R_t]] == 1
    @test Γ0[eq[:eq_mp], endo[:π_a_t]] == -1.5/4
    @test Γ0[eq[:eq_mp], endo[:y_t]]  == -1/4
    @test Γ0[eq[:eq_mp], endo[:y_f_t]] == 1/4
end


using PyPlot
using Roots

ntime = 100
time = collect(linspace(0,10,ntime))
gaussian=exp(-time.^2/2.)

function psd_func(x,omega)
psd = 0.
for k=1:nexp
  a1 = x[1+(k-1)*4]
  a2 = x[2+(k-1)*4]
  b  = x[3+(k-1)*4]
  om = x[4+(k-1)*4]
  psdt = (a2*om*(b^2 + om^2 - omega^2) + a1*b*(b^2 + om^2 + omega^2))/
         (b^4 + (om^2 - omega^2)^2 + 2*b^2*(om^2 + omega^2))
  psd += psdt
end
return psd
end

nexp = 4
x = [-0.5247762751816804,0.2828915570634826,1.3837656415003445,1.8323467099030555, 63.120917132667905,0.0,1.478130410600365,0.0, -32.17183668366479,0.0,1.5929285243364661,0.0, -29.423970266940803,0.0,1.38693874263686,0.0]

nfreq = 1000
omega = logspace(-1,1,nfreq)
psd = zeros(nfreq)
for i=1:nfreq
  psd[i]=psd_func(x,omega[i])
end

# model = gauss_model(atrial)
function dpsd_domega(omega)
# Let the kernel components be:
# a1*Exp[-b*t]*Cos[omega_0*t] + a2*Exp[-b*t]*Sin[omega_0*t]
# x = [{a1,a2,b,omega_0}_1; {a1,a2,b,omega_0}_2; {a1,a2,b,omega_0}_3; ... ]
#
dpsd = zero(Real)
for i=1:nexp
  a1 = x[(i-1)*4+1]
  a2 = x[(i-1)*4+2]
  b  = x[(i-1)*4+3]
  om0 = x[(i-1)*4+4]
  dpsd += (-2*sqrt(2/pi)*omega*(a1*b*(b^4 - 3*om0^4 + 2*om0^2*omega^2 + omega^4 - 2*b^2*(om0^2 - omega^2)) -
     a2*om0*(-3*b^4 + (om0^2 - omega^2)^2 - 2*b^2*(om0^2 + omega^2))))/
     ((b^2 + (om0 - omega)^2)^2*(b^2 + (om0 + omega)^2)^2)
end
return dpsd
end

# Now, solve for roots of PSD:

#omega_min = newton(dpsd_domega,5.0)
omega_min = fzero(dpsd_domega,[1e-4,1e4])
fmin = psd_func(x,omega_min)

plot(omega,psd)
scatter(omega_min,fmin)


# ODE SIR/SICR models
# variables are concentrations, X/N
# Z = 1-S

# SICRLQM ODE
# Cases are infective, mitigation applied
# 7+1
function sicrqm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q,mbase,mlocation = p
    beta *= mbase + (1-mbase)/(1 + exp(mrate*(t - mlocation)))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

# SICR ODE
# Cases are noninfective
# 4+1 parameters
function sicr!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar = p
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I)*(1-Z)
end

# SICRq ODE
# Cases are quarantined with effectiveness q
# 5+1 parameters
function sicrq!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q = p
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

# SICRf ODE
# Cases are quarantined with effectiveness q
# 5+1 parameters
function sicrf!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,f = p
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I)*(1-Z) - sigmac*I - f*sigma*I
	du[5] = dZ = beta*(I)*(1-Z)
end

# SICRqf ODE
# Cases are quarantined with factor q
# Cases recover or die with factor f
# 6+1 parameters
function sicrqf!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q,f = p
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - f*sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

# Mitigation applied
#6+1
function sicrm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,mc,ml = p
    beta *= mc + (1-mc)/(1 + exp(t - 35. - ml))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I)*(1-Z)
end

# Mitigation applied
# 7+1
function sicrfm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,f,mc,ml = p
    beta *= mc + (1-mc)/(1 + exp(t - 35. - ml))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I)*(1-Z) - sigmac*I - f*sigma*I
	du[5] = dZ = beta*(I)*(1-Z)

end

# 7+1
function sicrqm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q,mc,ml = p
    beta *= mc + (1-mc)/(1 + exp(t - 35. - ml))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

# 8+1
function sicrqfm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q,f,mc,ml = p
    beta *= mc + (1-mc)/(1 + exp(t - 35. - ml))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - f*sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

function sicrlqfm!(du,u,p,t)
	I,C = u
	beta,sigmac,sigma,q,f,mbase,mlocation,mrate,cmax,c50 = p
    beta *= mbase + (1-mbase)/(1 + exp(mrate*(t - mlocation)))
    sigmac *= 1 + cmax*t/(c50+t)
    du[1] = dI = beta*(I+q*C) - sigmac*I - f*sigma*I
	du[2] = dC = sigmac*I - sigma*C
end

function linearrd!(du,u,p,t)
	I,C = u
	beta,sigmac,sigmar,sigmad,q,f,mbase,mlocation,mrate,cmax,c50 = p
    beta *= mbase + (1-mbase)/(1 + exp(mrate*(t - mlocation)))
    sigmac *= 1 + cmax*t/(c50+t)
    du[1] = dI = beta*(I+q*C) - sigmac*I - f*(sigmar+sigmad)*I
	du[2] = dC = sigmac*I - (sigmar+sigmad)*C
end


#####
##### two-dimensional fluxes
#####
"
function euler_fluxes_2D(rhoL,uL,vL,betaL,rhoR,uR,vR,betaR,
                         rhologL,betalogL,rhologR,betalogR)
assumes primitive variables ordering: UL = (rhoL,uL,...,betaL),
                                      UR = (rhoR,uR,...,betaR)
"
function euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
                         rhoR,uR,vR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))
    vavg   = (@. .5*(vL+vR))

    unorm = (@. uL*uR + vL*vR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. FxS1*vavg)
    FxS4 = (@. f4aux*uavg)

    FyS1 = (@. rholog*vavg)
    FyS2 = (@. FxS3)
    FyS3 = (@. FyS1*vavg + pa)
    FyS4 = (@. f4aux*vavg)
    return (FxS1,FxS2,FxS3,FxS4),(FyS1,FyS2,FyS3,FyS4)
end

function euler_fluxes_2D_x(rhoL,uL,vL,betaL,rhoR,uR,vR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_2D_x(rhoL,uL,vL,betaL,rhologL,betalogL,
                             rhoR,uR,vR,betaR,rhologR,betalogR)
end
function euler_fluxes_2D_y(rhoL,uL,vL,betaL,rhoR,uR,vR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_2D_y(rhoL,uL,vL,betaL,rhologL,betalogL,
                             rhoR,uR,vR,betaR,rhologR,betalogR)
end

function euler_fluxes_2D_x(rhoL,uL,vL,betaL,rhologL,betalogL,
                           rhoR,uR,vR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))
    vavg   = (@. .5*(vL+vR))

    unorm = (@. uL*uR + vL*vR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. FxS1*vavg)
    FxS4 = (@. f4aux*uavg)

    return (FxS1,FxS2,FxS3,FxS4)
end

function euler_fluxes_2D_y(rhoL,uL,vL,betaL,rhologL,betalogL,
                           rhoR,uR,vR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))
    vavg   = (@. .5*(vL+vR))

    unorm = (@. uL*uR + vL*vR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FyS1 = (@. rholog*vavg)
    FyS2 = (@. FyS1*uavg)
    FyS3 = (@. FyS1*vavg + pa)
    FyS4 = (@. f4aux*vavg)
    return (FyS1,FyS2,FyS3,FyS4)
end

#####
##### shim functions
#####

# dispatch to n-dimensional constitutive routines, with optional entropy scaling
function primitive_to_conservative(rho,u,v,p)
   rho,rhoU,E = primitive_to_conservative_nd(rho,(u,v),p)
   return rho,rhoU...,E
end
function v_ufun(rho,rhou,rhov,E)
    v1,vU,vE = v_ufun_nd(rho,(rhou,rhov),E)
    return scale_entropy_output(v1,vU...,vE)
end
function u_vfun(v1,vU1,vU2,vE)
    v1,vU1,vU2,vE = scale_entropy_input(v1,vU1,vU2,vE)
    rho,rhoU,E = u_vfun_nd(v1,(vU1,vU2),vE)
    return rho,rhoU...,E
end
function conservative_to_primitive_beta(rho,rhou,rhov,E)
    rho,U,beta = conservative_to_primitive_beta_nd(rho,(rhou,rhov),E)
    return rho,U...,beta
end

Sfun(rho,rhou,rhov,E) = Sfun_nd(rho,(rhou,rhov),E)
pfun(rho,rhou,rhov,E) = pfun_nd(rho,(rhou,rhov),E)
betafun(rho,rhou,rhov,E) = betafun_nd(rho,(rhou,rhov),E)

"function euler_fluxes_2D(rhoL,uL,vL,betaL,rhoR,uR,vR,betaR)
assumes primitive variables ordering: UL = (rhoL,uL,...,betaL),
                                      UR = (rhoR,uR,...,betaR)"
function euler_fluxes(rhoL,uL,vL,betaL,rhoR,uR,vR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
                           rhoR,uR,vR,betaR,rhologR,betalogR)
end

euler_fluxes(rhoL,uL,vL,betaL,rhologL,betalogL,
             rhoR,uR,vR,betaR,rhologR,betalogR) =
     euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
                     rhoR,uR,vR,betaR,rhologR,betalogR)

euler_fluxes(varL,varR) = euler_fluxes(varL...,varR...)
euler_fluxes(varL,varR,logL,logR) = euler_fluxes(varL...,logL...,varR...,logR...)


"""
    x_ray_emission(n_cm3::Real, T_K::Real;
                   Emin::Real=5.e4, Emax::Real=1.e10,
                   xH::Real=0.76)

X-Ray emission of a particle with number density `n_cm3` in ``1/cm^3`` and temperature `T` in ``K``.
`Emin` and `Emax` give the minimum and maximum energy of the oberservation.
`xH` gives the hydrogen fraction used in the simulation.

# Returns
X-Ray surface brightness contribution in units of [erg/cm^2/s/Hz].
"""
function x_ray_emission(n_cm3::Real, T_K::Real;
                        Emin::Real=5.e4, Emax::Real=1.e10,
                        xH::Real=0.76)

    prefac = 4.0 * 2.42e-24 / (1 + xH)
    T_eV   = T_K * cgs2eV

    return prefac * sqrt( k_B * T_eV * 1.e-3 ) * n_cm3^2 *
            ( exp( -Emin / (k_B * T_eV) ) - exp( -Emax / (k_B * T_eV) ) )
end


# Lattice generation

module MyLattice

export Lattice, coord, flavor, nns, nnns, num_sites, num_uc, adj
export bravais_full

# This will not be accessible outside the module
function bravais_int(Xvecs)
    # Xvecs should be a d by d matrix of integers giving winding vectors as columns
    # This function returns a unique set of integer coordinates of points defining a cluster given the integer winding numbers

    const nmax = 40
    dim = size(Xvecs,1)
    if dim == 3
    	newxvecs = Xvecs
    elseif dim == 2
    	newxvecs = zeros(Int64,3,3)
    	newxvecs[1:2,1:2] = Xvecs
    	newxvecs[3,3] = 1
    else
        error("Lattice must have dimension 2 or 3")
    end
    tp = dot(newxvecs[:,1],cross(newxvecs[:,2],newxvecs[:,3]))

    numpts = abs(tp)
    function reduced_pt(pt)
        # gives a unique integer point inside a supercell defined by Xvecs
        # Reciprocal lattice for the supercell

        dim = length(pt)
        if dim == 2
            pt3d = [pt; 0]
        else
            pt3d = pt
        end
        recip = Array{Float64}(3,3)
        for a in 1:3
            recip[:,a] = cross(newxvecs[:,mod(a,3)+1],newxvecs[:,mod(a+1,3)+1])/tp
        end
        fracpt = mod(round(recip'*float(pt3d),3),1)
        newpt = round(Int64,newxvecs*fracpt)[1:dim]
        return newpt
    end

    myset=Set{Vector{Int64}}()
    skip = false
    if dim == 3
        for i in 0:nmax
            for j in 0:nmax
                for k in 0:nmax
                    thispt = [i,j,k]
                    modpt = reduced_pt(thispt)
                    push!(myset,modpt)
                    length(myset)==numpts && (skip = true; break)
                end
                skip && break
            end
            skip && break
        end
    else
    	for i in 0:nmax
    		for j in 0:nmax
    			thispt = [i,j]
    			modpt = reduced_pt(thispt)
    			push!(myset,modpt)
    			length(myset)==numpts && (skip = true; break)
    		end
    		skip && break
    	end
    end

    intpts = Array{Int64}(dim,numpts)
    for (i,vec) in enumerate(myset)
        intpts[:,i] = vec
    end

    return intpts
end

# Lattice type.  Stores all sorts of good stuff

type Lattice
    sites::Matrix{Float64}
    flavor::Vector{Int64}
    firstneighbors::Matrix{Int64}
    secondneighbors::Matrix{Int64}
    adjacency::Matrix{Int64}
    nsites::Int64
    nuc::Int64
    numnns::Int64
    numnnns::Int64
end

coord(lat::Lattice,i) = lat.sites[:,i]
flavor(lat::Lattice,i) = lat.flavor[i]
nns(lat::Lattice,i) = lat.firstneighbors[:,i]
nnns(lat::Lattice,i) = lat.secondneighbors[:,i]
num_sites(lat::Lattice) = lat.nsites
num_uc(lat::Lattice) = lat.nuc
adj(lat::Lattice,i,j) = lat.adjacency[i,j]

function bravais_full(avecs,basis,Xvecs)
    # Will construct a Lattice object describing a cluster of lattice defined by these parameters
    # avecs should be a d by d matrix of primitive lattice vectors arranged in columns
    # basis should be a d by n matrix of basis vectors for n sites in a unit cell.
    # Xvecs should be a d by d matrix of integers giving winding vectors as columns
    # This function should return a full Lattice object

    dim = size(avecs,1)
    int_pts = bravais_int(Xvecs) # integer coordinates of distinct lattice points
    nuc = size(int_pts,2) # number of unit cells
    nbasis = size(basis,2) # number of sites in the primitive unit cell
    nsites = nuc*nbasis # total number of sites in the cluster

    # form real space winding vectors
    wind = avecs*Xvecs

    # form reciprocal space vectors for supercell
    if dim == 3
    	newwind = wind
    elseif dim == 2
    	newwind = zeros(3,3)
    	newwind[1:2,1:2] = wind
    	newwind[3,3] = 1.0
    else
        error("Lattice must have dimension 2 or 3")
    end
    tp = dot(newwind[:,1],cross(newwind[:,2],newwind[:,3]))
    bigrecip = Array{Float64}(3,3)
    for a in 1:3
        bigrecip[:,a] = cross(newwind[:,mod(a,3)+1],newwind[:,mod(a+1,3)+1])/tp
    end
    if dim == 3
        recip = bigrecip
    elseif dim == 2
        recip = bigrecip[1:2,1:2]
    end

    function reduced_pt(pt)
        # gives a unique point inside a supercell defined by wind
        fracpt = mod(recip'*pt,1)
        newpt = wind*fracpt
        return newpt
    end

    function distance(pt1,pt2)
        # finds closest distance between two points taking into account periodicity
        # round to 3 digits
        diff = pt1-pt2
        fracdiff = recip'*diff
        normfrac = mod(fracdiff+0.5*ones(dim),1)-0.5*ones(dim)
        newdiff = wind*normfrac
        return round(norm(newdiff),3)
    end

    # Compute all the sites in the unit cell
    bravais_points = avecs*int_pts
    sites = Array{Float64}(dim,nuc*nbasis) # real space coordinate of site i
    flavors = Array{Int64}(nuc*nbasis) # basis type of site i
    for i in 1:nuc
        for a in 1:nbasis
            sites[:,nuc*(a-1)+i] = reduced_pt(bravais_points[:,i]+basis[:,a])
            flavors[nuc*(a-1)+i] = a
        end
    end

    # Compute neighbors
    # For the moment, make simple assumptions:
    # Find only first and second neighbors
    # Assume the same number of these for every site
    # Store them in nns[a,i] and nnns[a,i], here i is the site, and a runs over the nns and nnns of that site

    # First figure out how many there are and at what distance
    alldist = Array{Float64}(0)
    for i in 2:nsites
        push!(alldist,distance(sites[:,1],sites[:,i]))
    end
    dists = unique(alldist)
    sort!(dists)
    nndist = dists[1]
    nnndist = dists[2]
    numnns = length(find(x->x==nndist,alldist))
    numnnns = length(find(x->x==nnndist,alldist))

    # Now construct all of them and store in two 2d arrays
    thenns = Array{Int64}(numnns,nsites)
    thennns = Array{Int64}(numnnns,nsites)
    for i in 1:nsites
        dists = [distance(sites[:,i],sites[:,j]) for j in 1:nsites]
        thenns[:,i] = find(x->x==nndist,dists)
        thennns[:,i] = find(x->x==nnndist,dists)
    end

    # Also compute adjacency matrix: entries = 1 if NN, 2 if NNN, 0 otherwise
    adjacency = Array{Int64}(nsites,nsites)
    for i in 1:nsites
        for j in 1:nsites
            dist = distance(sites[:,i],sites[:,j])
            if dist == nndist
                adjacency[i,j]=1
            elseif dist == nnndist
                adjacency[i,j]=2
            else
                adjacency[i,j]=0
            end
        end
    end

    Lattice(sites,flavors,thenns,thennns,adjacency,nsites,nuc,numnns,numnnns)
end

end


@testset "FBM" begin
    rng = MersenneTwister(42)
    h = 0.3
    k = FBMKernel(; h=h)
    v1 = rand(rng, 3)
    v2 = rand(rng, 3)
    @test k(v1, v2) ≈
          (
        sqeuclidean(v1, zero(v1))^h + sqeuclidean(v2, zero(v2))^h -
        sqeuclidean(v1 - v2, zero(v1 - v2))^h
    ) / 2 atol = 1e-5
    @test repr(k) == "Fractional Brownian Motion Kernel (h = $(h))"

    test_interface(k)
    @test repr(k) == "Fractional Brownian Motion Kernel (h = $(h))"
    test_ADs(FBMKernel; ADs=[:ReverseDiff, :Zygote])
    @test_broken "Tests failing for kernelmatrix(k, x) for ForwardDiff"
    test_params(k, ([h],))
end


#=
1D advection using FV
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("FVadvection1d");

useLog("FVadvection1dlog", level=3)

# Configuration setup
domain(1)
solverType(FV)
timeStepper(RK4)

# Mesh
n = 60 # number of elements
mesh(LINEMESH, elsperdim=n, bids=2)

# Set the order for flux
# finiteVolumeOrder(3);

# Variables and BCs
u = variable("u", location=CELL)
v = variable("v", location=CELL)
w = variable("w", location=CELL)
# boundary(u, 1, FLUX, "t<0.2 ? sin(pi*t/0.2)^2 : 0")
# boundary(u, 2, NO_BC)
# boundary(v, 1, FLUX, "t<0.2 ? sin(pi*t/0.2)^2 : 0")
# boundary(v, 2, NO_BC)
# boundary(w, 1, FLUX, "t<0.2 ? sin(pi*t/0.2)^2 : 0")
# boundary(w, 2, NO_BC)

boundary(u, 1, FLUX, "t<0.2 ? sin(pi*t/0.2)^2 : 0")
boundary(v, 1, DIRICHLET, "t<0.2 ? sin(pi*t/0.2)^2 : 0")
boundary(w, 1, DIRICHLET, 1)
boundary(u, 2, NO_BC)
boundary(v, 2, NO_BC)
boundary(w, 2, DIRICHLET, 0)

# Time interval and initial condition
T = 0.5;
timeInterval(T)
initial(u, "0")
initial(v, "0")
initial(w, "0")

# The flux and source terms of the conservation equation
# F and S in the following equation:
# Dt(int(u dx)/V) = int(S dx) - int(F ds)
coefficient("a", 1) # advection velocity
# The "upwind" function applies upwinding to the term (a.n)*f with flow velocity a.
# The optional third parameter is for tuning. Default upwind = 0, central = 1. Choose something between these.
flux([u, v, w], ["upwind(a,u)", "upwind(a,v)", "upwind(a,w)"])
# Note that there is no source() for this problem.

# @exportCode("fvad1dcode") # uncomment to export generated code to a file
# @importCode("advec1dcode") # uncomment to import code from a file

solve([u,v,w])

finalize_finch()

##### Uncomment below to compare to exact solution

# # The exact solution with constant velocity v
# a = 1;
# n = size(Finch.fv_info.cellCenters,2);
# exact = zeros(n);
# x = Finch.fv_info.cellCenters[:]
# for i=1:n
#     xt = x[i] - a*T;
#     if xt < 0
#         exact[i] = xt > -0.2 ? sin(pi*-xt/0.2)^2 : 0
#     end
# end

# using Plots
# pyplot();
# display(plot([x x x x], [exact u.values[:] v.values[:] w.values[:]], markershape=:circle, label=["exact" "u" "v" "w"]))


using Latexify
using Test
using Markdown

str = "2*x^2 - y/c_2"
ex = :(2*x^2 - y/c_2)

desired_output = "2 \\cdot x^{2} - \\frac{y}{c_{2}}"

@test latexify(:(a = [x / y; 3; 4])) == replace(
raw"$a = \left[
\begin{array}{c}
\frac{x}{y} \\
3 \\
4 \\
\end{array}
\right]$", "\r\n"=>"\n")

@test latexify(:(a = [x / y 2 3 4])) == replace(
raw"$a = \left[
\begin{array}{cccc}
\frac{x}{y} & 2 & 3 & 4 \\
\end{array}
\right]$", "\r\n"=>"\n")

@test latexify(:(a = [x / y 2; 3 4])) == replace(
raw"$a = \left[
\begin{array}{cc}
\frac{x}{y} & 2 \\
3 & 4 \\
\end{array}
\right]$", "\r\n"=>"\n")

@test latexraw(str) == latexraw(ex)
@test latexraw(ex) == desired_output

@test latexify(:(u[1, 2]); index = :bracket) == raw"$u\left[1, 2\right]$"
@test latexify(:(u[1, 2]); index = :subscript) == raw"$u_{1,2}$"

array_test = [ex, str]
@test all(latexraw.(array_test) .== desired_output)

@test latexraw(:(@__dot__(x / y))) == raw"\frac{x}{y}"
@test latexraw(:(@. x / y)) == raw"\frac{x}{y}"
@test latexraw(:(eps())) == raw"\mathrm{eps}()"

@test latexraw(:y_c_a) == "y_{c\\_a}"
@test latexraw(1.0) == "1.0"
@test latexraw(1.00) == "1.0"
@test latexraw(1) == "1"
@test latexraw(:(log10(x))) == "\\log_{10}\\left( x \\right)"
@test latexraw(:(sin(x))) ==  "\\sin\\left( x \\right)"
@test latexraw(:(sin(x)^2)) ==  "\\sin^{2}\\left( x \\right)"
@test latexraw(:(asin(x))) ==  "\\arcsin\\left( x \\right)"
@test latexraw(:(asinh(x))) ==  "\\mathrm{arcsinh}\\left( x \\right)"
@test latexraw(:(sinc(x))) ==  "\\mathrm{sinc}\\left( x \\right)"
@test latexraw(:(acos(x))) ==  "\\arccos\\left( x \\right)"
@test latexraw(:(acosh(x))) ==  "\\mathrm{arccosh}\\left( x \\right)"
@test latexraw(:(cosc(x))) ==  "\\mathrm{cosc}\\left( x \\right)"
@test latexraw(:(atan(x))) ==  "\\arctan\\left( x \\right)"
@test latexraw(:(atan2(x))) ==  "\\arctan\\left( x \\right)"
@test latexraw(:(atanh(x))) ==  "\\mathrm{arctanh}\\left( x \\right)"
@test latexraw(:(acot(x))) ==  "\\mathrm{arccot}\\left( x \\right)"
@test latexraw(:(acoth(x))) ==  "\\mathrm{arccoth}\\left( x \\right)"
@test latexraw(:(asec(x))) ==  "\\mathrm{arcsec}\\left( x \\right)"
@test latexraw(:(sech(x))) ==  "\\mathrm{sech}\\left( x \\right)"
@test latexraw(:(asech(x))) ==  "\\mathrm{arcsech}\\left( x \\right)"
@test latexraw(:(acsc(x))) ==  "\\mathrm{arccsc}\\left( x \\right)"
@test latexraw(:(csch(x))) ==  "\\mathrm{csch}\\left( x \\right)"
@test latexraw(:(acsch(x))) ==  "\\mathrm{arccsch}\\left( x \\right)"
@test latexraw(:(x ± y)) == "x \\pm y"
@test latexraw(:(f(x))) ==  "f\\left( x \\right)"
@test latexraw("x = 4*y") == "x = 4 \\cdot y"
@test latexraw(:(sqrt(x))) == "\\sqrt{x}"
@test latexraw(complex(1,-1)) == "1-1\\textit{i}"
@test latexraw(1//2) == "\\frac{1}{2}"
@test latexraw(missing) == "\\textrm{NA}"
@test latexraw("x[2]") == raw"x\left[2\right]"
@test latexraw("x[2, 3]") == raw"x\left[2, 3\right]"
@test latexraw("α") == raw"\alpha"
@test latexraw("α + 1") == raw"\alpha + 1"
@test latexraw("α₁") == raw"\alpha{_1}"
@test latexraw("γ³") == raw"\gamma{^3}"
@test latexraw("β₃_hello") == raw"\beta{_3}_{hello}"
@test latexraw("β₃₄") == raw"\beta{_{34}}"
@test latexraw("β₃₄¹⁴") == raw"\beta{_{34}^{14}}"
@test latexraw("β₃¹⁴") == raw"\beta{_3^{14}}"
@test latexraw("β¹⁴₃") == raw"\beta{^{14}_3}"
@test latexraw("β¹⁴") == raw"\beta{^{14}}"
@test latexraw("β⁴") == raw"\beta{^4}"
@test latexraw(Inf) == raw"\infty"
@test latexraw(:Inf) == raw"\infty"
@test latexraw("Inf") == raw"\infty"
@test latexraw(-Inf) == raw"-\infty"
@test latexraw(:($(3+4im)*a)) == raw"\left( 3+4\textit{i} \right) \cdot a"
@test latexraw(:(a*$(3+4im))) == raw"a \cdot \left( 3+4\textit{i} \right)"

@test latexify(:((-1) ^ 2)) == replace(
raw"$\left( -1 \right)^{2}$", "\r\n"=>"\n")
@test latexify(:($(1 + 2im) ^ 2)) == replace(
raw"$\left( 1+2\textit{i} \right)^{2}$", "\r\n"=>"\n")
@test latexify(:($(3 // 2) ^ 2)) == replace(
raw"$\left( \frac{3}{2} \right)^{2}$", "\r\n"=>"\n")

### Test broadcasting
@test latexraw(:(sum.((a, b)))) == raw"\mathrm{sum}\left( a, b \right)"

### Test for correct signs in nested sums/differences.
@test latexraw("-(-1)") == raw" + 1"
@test latexraw("+(-1)") == raw"-1"
@test latexraw("-(+1)") == raw" - 1"
@test latexraw("-(1+1)") == raw" - \left( 1 + 1 \right)"
@test latexraw("1-(-2)") == raw"1 + 2"
@test latexraw("1 + (-(2))") == raw"1 - 2"
@test latexraw("1 + (-2 -3 -4)") == raw"1 -2 - 3 - 4"
@test latexraw("1 - 2 - (- 3 - 4)") == raw"1 - 2 - \left(  - 3 - 4 \right)"
@test latexraw("1 - 2 - (- 3 -(2) + 4)") == raw"1 - 2 - \left(  - 3 - 2 + 4 \right)"
@test latexraw("1 - 2 - (- 3 -(2 - 8) + 4)") == raw"1 - 2 - \left(  - 3 - \left( 2 - 8 \right) + 4 \right)"
@test latexraw(:(-$(3+5im))) == raw" - \left( 3+5\textit{i} \right)"
@test latexraw(:($(3+4im)-$(3+5im))) == raw"3+4\textit{i} - \left( 3+5\textit{i} \right)"

# @test_throws ErrorException latexify("x/y"; env=:raw, bad_kwarg="should error")

@test latexraw(:(3 * (a .< b .<= c < d <= e > f <= g .<= h .< i == j .== k != l .!= m))) ==
raw"3 \cdot \left( a < b \leq c < d \leq e > f \leq g \leq h < i = j = k \neq l \neq m \right)"

#### Test the fmt keyword option
@test latexify([32894823 1.232212 :P_1; :(x / y) 1.0e10 1289.1]; env=:align, fmt="%.2e") == replace(
raw"\begin{align}
3.29e+07 =& 1.23e+00 =& P_{1} \\
\frac{x}{y} =& 1.00e+10 =& 1.29e+03
\end{align}
", "\r\n"=>"\n")

@test latexify([32894823 1.232212 :P_1; :(x / y) 1.0e10 1289.1]; env=:equation, fmt="%.2e") == replace(
raw"\begin{equation}
\left[
\begin{array}{ccc}
3.29e+07 & 1.23e+00 & P_{1} \\
\frac{x}{y} & 1.00e+10 & 1.29e+03 \\
\end{array}
\right]
\end{equation}
", "\r\n"=>"\n")

@test latexify([32894823 1.232212 :P_1; :(x / y) 1.0e10 1289.1]; env=:table, fmt="%.2e") == replace(
raw"\begin{tabular}{ccc}
$3.29e+07$ & $1.23e+00$ & $P_{1}$\\
$\frac{x}{y}$ & $1.00e+10$ & $1.29e+03$\\
\end{tabular}
", "\r\n"=>"\n")

@test latexify([32894823 1.232212 :P_1; :(x / y) 1.0e10 1289.1]; env=:mdtable, fmt="%.2e") ==
Markdown.md"|    $3.29e+07$ | $1.23e+00$ |    $P_{1}$ |
| -------------:| ----------:| ----------:|
| $\frac{x}{y}$ | $1.00e+10$ | $1.29e+03$ |
"

@test latexify(1234.2234; env=:inline, fmt="%.2e") ==
raw"$1.23e+03$"

@test latexify(1234.2234; env=:raw, fmt="%.2e") ==
raw"1.23e+03"

test_functions = [:sinh, :alpha, :Theta, :cosc, :acoth, :acot, :asech, :lambda,
                  :asinh, :sinc, :eta, :kappa, :nu, :asin, :epsilon, :sigma,
                  :upsilon, :phi, :tanh, :iota, :Psi, :acosh, :log, :zeta, :mu,
                  :csc, :xi, :tau, :beta, :Lambda, :Xi, :Phi, :acsc, :atan,
                  :sech, :atanh, :gamma, :Delta, :rho, :sec, :log10, :delta,
                  :pi, :cot, :log2, :cos, :Omega, :psi, :atan2, :Gamma, :cosh,
                  :acos, :Pi, :Upsilon, :omega, :coth, :chi, :tan, :csch,
                  :acsch, :theta, :asec, :Sigma, :sin]

@test latexify(["3*$(func)(x)^2/4 -1" for func = test_functions]) == replace(
raw"\begin{equation}
\left[
\begin{array}{c}
\frac{3 \cdot \sinh^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \alpha\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \Theta\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \mathrm{cosc}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \mathrm{arccoth}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \mathrm{arccot}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \mathrm{arcsech}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \lambda\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \mathrm{arcsinh}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \mathrm{sinc}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \eta\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \kappa\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \nu\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \arcsin^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \epsilon\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \sigma\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \upsilon\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \phi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \tanh^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \iota\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \Psi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \mathrm{arccosh}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \log\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \zeta\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \mu\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \csc^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \xi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \tau\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \beta\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \Lambda\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \Xi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \Phi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \mathrm{arccsc}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \arctan^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \mathrm{sech}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \mathrm{arctanh}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \Gamma\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \Delta\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \rho\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \sec^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \log_{10}\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \delta\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \pi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \cot^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \log_{2}\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \cos^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \Omega\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \psi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \arctan^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \Gamma\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \cosh^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \arccos^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \Pi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \Upsilon\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \left( \omega\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \coth^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \chi\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \tan^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \mathrm{csch}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \mathrm{arccsch}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \theta\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \mathrm{arcsec}^{2}\left( x \right)}{4} - 1 \\
\frac{3 \cdot \left( \Sigma\left( x \right) \right)^{2}}{4} - 1 \\
\frac{3 \cdot \sin^{2}\left( x \right)}{4} - 1 \\
\end{array}
\right]
\end{equation}
", "\r\n"=>"\n")

## Test logical operators
@test latexraw(:(x && y)) == "x \\wedge y"
@test latexraw(:(x || y)) == "x \\vee y"
@test latexraw(:(x || !y)) == "x \\vee \\neg y"

## Test {cases} enviroment
@test latexraw(:(R(p,e,d) = e ? 0 : log(p) - d)) == replace(
raw"R\left( p, e, d \right) = \begin{cases}
0 & \text{if } e\\
\log\left( p \right) - d & \text{otherwise}
\end{cases}", "\r\n"=>"\n")

@test latexraw(:(R(p,e,d,t) = if (t && e); 0 elseif (t && !e); d else log(p) end)) == replace(
raw"R\left( p, e, d, t \right) = \begin{cases}
0 & \text{if } t \wedge e\\
d & \text{if } t \wedge \neg e\\
\log\left( p \right) & \text{otherwise}
\end{cases}", "\r\n"=>"\n")

@test latexraw(:(function reward(p,e,d,t)
    if t && e
        return 0
    elseif t && !e
        return -1d
    elseif 2t && e
        return -2d
    elseif 3t && e
        return -3d
    else
        return log(p)
    end
end)) == replace(
raw"\mathrm{reward}\left( p, e, d, t \right) = \begin{cases}
0 & \text{if } t \wedge e\\
-1 \cdot d & \text{if } t \wedge \neg e\\
-2 \cdot d & \text{if } 2 \cdot t \wedge e\\
-3 \cdot d & \text{if } 3 \cdot t \wedge e\\
\log\left( p \right) & \text{otherwise}
\end{cases}", "\r\n"=>"\n")


"""
`UnivariateFeatureSelector` has the following fields:

* `method::Function` (required)- Method to calculate feature importance. The
  method chosen will determine the scoring. Below is the scring with available
  statistical method to obtain them.

    * Correlation - higher score means more important

        * `pearson_correlation`

    * P-value - lower score means more important

        * `f_test`
        * `chisq_test`

* `k::Union{Int64,Nothing}` - Select top `k` features with the highest correlation to target
  variable. You could ignore this by specifying k == nothing. This defaults to nothing.
* `threshold::Union{Float64,Nothing}` - Select features with correlation more than or equal
  to threshold. To ignore, simply set threshold to nothing (default behavior).
"""
mutable struct UnivariateFeatureSelector
    method::Function
    k::Union{Int64,Nothing}
    threshold::Union{Float64,Nothing}
end

function UnivariateFeatureSelector(;
    method::Function,
    k::Union{Int64,Nothing} = nothing,
    threshold::Union{Float64,Nothing} = nothing,
)
    UnivariateFeatureSelector(method, k, threshold)
end

"""
    select_features(selector,
                    X::DataFrame,
                    y::Vector;
                    verbose::Bool=false)

Select features based on the importance, which is defined by `selector.method`
to target `y`. if `verbose` is true, logs will be printed - this defaults to
false. This function will return only the feature names of selected features.

If you have feature `X_data` as matrix and feature names `X_features` as a
Vector, you can replace `X` with `X_data` and `X_features` (in this order).

# Example

```jldoctest
julia> using RDatasets, FeatureSelectors, DataFrames

julia> boston = dataset("MASS", "Boston");

julia> selector = UnivariateFeatureSelector(method=pearson_correlation, k=5)
UnivariateFeatureSelector(FeatureSelectors.pearson_correlation, 5, nothing)

julia> select_features(
           selector,
           boston[:, Not(:MedV)],
           boston.MedV
       )
5-element Vector{String}:
 "LStat"
 "Rm"
 "PTRatio"
 "Indus"
 "Tax"

```
"""
function select_features(
    selector::UnivariateFeatureSelector,
    X_data::Matrix,
    X_features::Vector,
    y::Vector;
    verbose::Bool = false,
)
    score_arr = selector.method(X_data, y)
    # Correlation should be sorted descending
    if selector.method == pearson_correlation
        rev = true
        filter_operator = >=
        if verbose
            @info "method is correlation score, result will be " * "descendingly sorted" rev
        end
    else
        rev = false
        filter_operator = <=
        if verbose
            @info "method is p-value, result will be " * "ascendingly sorted" rev
        end
    end
    sorted_tup =
        sort([x for x in zip(X_features, score_arr)], by = v -> abs(v[2]), rev = rev)
    if !isnothing(selector.k)
        if selector.k < 1
            @warn "k cannot be less than 1. Resetting value to 1"
            selector.k = 1
        end
        verbose && @info "Filtering top k features" selector.k
        sorted_tup = sorted_tup[1:selector.k]
    end
    if !isnothing(selector.threshold)
        verbose && @info "Filtering by threshold" selector.threshold
        sorted_tup = filter(v -> filter_operator(abs(v[2]), selector.threshold), sorted_tup)
    end
    first.(sorted_tup)
end

select_features(
    selector::UnivariateFeatureSelector,
    X::DataFrame,
    y::Vector;
    verbose::Bool = false,
) = select_features(
    selector,
    Matrix(X),
    names(X),
    y;
    verbose = verbose,
)

"""
    pearson_correlation(X_data::Matrix, y::Vector)

Calculate pearson's correlation on `X_data` to `y`.
"""
pearson_correlation(X_data::Matrix, y::Vector) = cor(X_data, y)

"""
    f_test(X_data::Matrix, y::Vector)

Calculate p-value using f-test method.
"""
function f_test(X_data::Matrix, y::Vector)
    pvals = [pvalue(VarianceFTest(X_col, y)) for X_col in eachcol(X_data)]
end

"""
    chisq_test(X_data::Matrix, y::Vector)

Calculate p-value using chi-square test.
"""
function chisq_test(X_data::Matrix, y::Vector)
    # DataFrame for easy aggregation
    data_df = DataFrame(hcat(X_data, y), :auto)
    # Rename cols to prevent name clash with count aggregation below
    rename!(data_df, [Symbol("input$i") for i = 1:size(data_df)[2]])
    y_name = names(data_df[!, [end]])[1]
    pvals = Vector{Float64}()
    for col_name in names(data_df[:, 1:end-1])
        # Hacky method to generate frequency count and put 0 for missing
        # combination. `by` function is not able to achieve this. See:
        # https://github.com/JuliaData/DataFrames.jl/issues/2136

        # Feature + target to process and put count = 1 for each row
        Xy_df = hcat(data_df[:, [col_name, y_name]], Int.(ones(size(data_df)[1])))
        rename!(Xy_df, :x1 => :Count)
        # This nested join will append missing in Count to missing combination
        # of value in col_name and y_name
        Xy_df = outerjoin(
            crossjoin(
                DataFrame(; Symbol(col_name) => unique(Xy_df[:, col_name])),
                DataFrame(; Symbol(y_name) => unique(Xy_df[:, y_name])),
            ),
            Xy_df,
            on = [col_name, y_name],
        )
        # sum count and recode missing to 0
        _y = replace(
            sort(combine(r -> sum(r.Count), groupby(Xy_df, [col_name, y_name]))).x1,
            missing => 0,
        )
        _y = reshape(
            _y,
            (length(unique(Xy_df[:, col_name])), length(unique(Xy_df[:, y_name]))),
        )
        push!(pvals, pvalue(ChisqTest(_y)))
    end
    pvals
end


using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using Agents, LightGraphs, DataFrames
# import GLMakie
# using GraphMakie.NetworkLayout

n0 = 100
seed = 42
s = 80 # steps to run
er = n -> erdos_renyi(n, 200; seed)
ba = n -> barabasi_albert(n, 3, 2; complete=true, seed)
ws = n -> watts_strogatz(n, 4, 0.1; seed)

graphs = [g(n0) for g in (er, ba, ws)]

# g = graphs[i](n0)

# %% First version: one agent per node, no moving

# model parameters:
τ = 0.1  # transmission probability
# d = 0.02 # death probability? Should I?
p = 14   # period to recover

# Create an ABM on a graph `g`
@agent Human GraphAgent begin
    status::Symbol # :S, :I, :R
    days_infected::Int
end
function init_sir_abm(g; τ = 0.1, p = 14, d = 0.02)
    abm = ABM(Human, GraphSpace(g); properties = (τ = τ, p = p, d = d))
    # Add one agent to each node, with infected agent at node 1
    starting_status(i) = i == 1 ? (:I, 1) : (:S, 0)
    fill_space!(abm, starting_status)
    return abm
end
# Define rules of the ABM simulation
function agent_step!(human, abm) # for each agent at each step
    human.status ≠ :I && return # do nothing if not infected
    # Infect all susceptible neighbors with probability τ
    for friend in nearby_agents(human, abm)
        if friend.status == :S && rand() < abm.τ
            friend.status = :I
        end
    end
    # Check if human recovers or dies
    human.days_infected += 1
    if human.days_infected ≥ abm.p
        human.status = :R
    end
    # elseif (rand() < abm.d)
    #     kill_agent!(human, abm)
    # end
end
# run simulation
g = graphs[1]
abm = init_sir_abm(g)
infected(abm) = count(a.status == :I for a in allagents(abm))
recovered(abm) = count(a.status == :R for a in allagents(abm))
susceptible(abm) = count(a.status == :S for a in allagents(abm))
mdata = [infected, recovered, susceptible]

# infected(x) = count(i == :I for i in x)
# recovered(x) = count(i == :R for i in x)
# adata = [(:status, infected), (:status, recovered)]

names = ("Erdős–Rényi ", "Barabási–Albert", "Watts-Strogatz")
names = ("ER ", "BA", "WS")

fig = figure(figsize = (figx/2, figy))
for (i, g) in enumerate(graphs)
    abm = init_sir_abm(g)
    _, mdf = run!(abm, agent_step!, 100; mdata)
    plot(mdf[!, :infected]; color = "C$(i-1)", label = "$(names[i])")
    plot(mdf[!, :recovered]; color = "C$(i-1)", lw = 1.5, ls = "--")
    plot(mdf[!, :susceptible]; color = "C$(i-1)", lw = 1.5, ls = ":")
end
legend(handlelength=1)
xlabel("time")
fig.tight_layout(pad=0.3)

# %% Ensemble version
using Statistics
fig = figure(figsize = (figx/2, figy))
ax = gca()
for (i, g) in enumerate(graphs)
    _, mdf = ensemblerun!(seed -> init_sir_abm(g), agent_step!, dummystep, 100; mdata, ensemble = 10)
    μdf = combine(groupby(mdf, :step), [:infected, :susceptible, :recovered] .=> mean .=> [:I, :S, :R])
    σdf = combine(groupby(mdf, :step), [:infected, :susceptible, :recovered] .=> std .=> [:I, :S, :R])

    # gdf = groupby(mdf, :ensemble)
    # I = sum(df[!, :infected] for df in gdf) ./ 5
    # R = sum(df[!, :recovered] for df in gdf) ./ 5
    # S = sum(df[!, :susceptible] for df in gdf) ./ 5
    I = μdf[!, :I]
    Iσ = σdf[!, :I]
    ax.plot(I; color = "C$(i-1)", label = "$(names[i])")
    ax.fill_between(0:length(I)-1, I .- Iσ, I .+ Iσ; color = "C$(i-1)", alpha = 0.25)
    # plot(R; color = "C$(i-1)", lw = 1.5, ls = "--")
    # plot(S; color = "C$(i-1)", lw = 1.5, ls = ":")
end
legend(handlelength=1)
xlabel("time (days)")
ylabel("infected agents")
fig.tight_layout(pad=0.3)

wsave(plotsdir("10", "sir_network"), fig)

# %% Second version: pretty much what we have online: many agents per node (city)
# movement between cities and infecting only within node.

nt = 50
seed = 42
s = 80 # steps to run
er = erdos_renyi(nt, 100; seed)
ba = barabasi_albert(nt, 3, 2; complete=true, seed)
ws = watts_strogatz(nt, 4, 0.1; seed)

graphs = (er, ba, ws)

function init_sir_abm_transport(g; τ=0.05, p=14, m = 0.5)
    tprobs = transition_probabilities(g)
    properties = @ntuple tprobs τ p m
    abm = ABM(Human, GraphSpace(g); properties)
    # Add a bunch of agents at each node, set agent0 as infected.
    for i in positions(abm, :random)
        for j in 1:10
            add_agent!(i, abm, :S, 0)
        end
    end
    abm[1].status = :I
    return abm
end

function transition_probabilities(g::AbstractGraph)
    degrees = degree(g)
    distributions = Vector{DiscreteNonParametric}(undef, nv(g))
    for i in 1:nv(g)
        n = neighbors(g, i)
        if isempty(n)
            vals = [i]; probs = [1.0]
        else
            vals = n
            outdtot = sum(degrees[j] for j in n)
            probs = [degrees[j]/outdtot for j in n]
        end
        distributions[i] = DiscreteNonParametric(vals, probs)
    end
    return distributions
end

function agent_step!(human, abm)
    # Travelling
    if rand(abm.rng) < abm.m
        newnode = rand(abm.rng, abm.tprobs[human.pos])
        move_agent!(human, newnode, abm)
    end
    # Infecting
    human.status ≠ :I && return # do nothing if not infected
    for person in agents_in_position(human, abm) # nearby_agents(human, abm, 0)
        if person.status == :S && rand() < abm.τ
            person.status = :I
        end
    end
    # Check if human recovers or dies
    human.days_infected += 1
    if human.days_infected ≥ abm.p
        human.status = :R
    end
end

fig = figure(figsize = (figx/2, figy))
ax = gca()
for (i, g) in enumerate(graphs)
    _, mdf = ensemblerun!(seed -> init_sir_abm_transport(g), agent_step!, dummystep, 100; mdata, ensemble = 10)
    μdf = combine(groupby(mdf, :step), [:infected, :susceptible, :recovered] .=> mean .=> [:I, :S, :R])
    σdf = combine(groupby(mdf, :step), [:infected, :susceptible, :recovered] .=> std .=> [:I, :S, :R])

    # gdf = groupby(mdf, :ensemble)
    # I = sum(df[!, :infected] for df in gdf) ./ 5
    # R = sum(df[!, :recovered] for df in gdf) ./ 5
    # S = sum(df[!, :susceptible] for df in gdf) ./ 5
    I = μdf[!, :I]
    Iσ = σdf[!, :I]
    ax.plot(I; color = "C$(i-1)", label = "$(names[i])")
    ax.fill_between(0:length(I)-1, I .- Iσ, I .+ Iσ; color = "C$(i-1)", alpha = 0.25)
    # plot(R; color = "C$(i-1)", lw = 1.5, ls = "--")
    # plot(S; color = "C$(i-1)", lw = 1.5, ls = ":")
end
legend(handlelength=1)
xlabel("time (days)")
ylabel("infected agents")
# fig.tight_layout(pad=0.3)


# number of visible keypoints = 13
# mask = [1 1 1 1;
#         1 0 1 1;
#         1 0 1 1;
#         1 1 1 0]
viewpoint_indices = [1, 2, 3, 4, 1, 3, 4, 1, 3, 4, 1, 2, 3]
point_indices     = [1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4]
indices = SBA.Indices(viewpoint_indices, point_indices)
N = length(viewpoint_indices)
A = randn(Float64, 2, 4, N)  # n_pose_params = 4
B = randn(Float64, 2, 3, N)  # n_point_params = 3
x_true = x_pred = randn(Float64, 2, N)

# n_rows(J) = 2 * n_visible_keypoints
#           = 2 * 16 = 26
# n_cols(J) = n _point_params * n_points + n_pose_params * n_viewpoints
#           = 3 * 4 + 4 * 4 = 12 + 16 = 28
# n_rows(J) < n_cols(J)
# J' * J can not be invertible. ArgumentError should be thrown
@test_throws ArgumentError SBA.sba(indices, x_true, x_pred, A, B)

# number of visible keypoints = 14
# mask = [1 1 1 1;
#         1 1 1 1;
#         1 1 1 0;
#         1 1 1 0]
viewpoint_indices = [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 1, 2, 3]
point_indices     = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 4, 4]
indices = SBA.Indices(viewpoint_indices, point_indices)
N = length(viewpoint_indices)
A = randn(Float64, 2, 4, N)  # n_pose_params = 4
B = randn(Float64, 2, 3, N)  # n_point_params = 3
x_true = x_pred = randn(Float64, 2, N)

# n_rows(J) = 2 * n_visible_keypoints
#           = 2 * 14 = 28
# n_cols(J) = n _point_params * n_points + n_pose_params * n_viewpoints
#           = 3 * 4 + 4 * 4 = 12 + 16 = 28
# n_rows(J) = n_cols(J)
# J' * J can be invertible. Nothing should be thrown
SBA.sba(indices, x_true, x_pred, A, B)


##############################################################################
##
## REF: http://optlab-server.sce.carleton.ca/POAnimations2007/NonLinear7.html
##
## Initial points
## N + 1 points of dimension N
## Store in N + 1 by N array
## Compute centroid of non-maximal points
## Iteratively update p_h
## Parameters: a, g, b
## y_l is current minimum
## y_h is current maximum
##
##############################################################################

function nelder_mead_trace!(tr::OptimizationTrace,
                            p::Matrix,
                            y::Vector,
                            n::Integer,
                            f::Function,
                            iteration::Integer,
                            store_trace::Bool,
                            show_trace::Bool)
    d = Dict()
    d["Standard Deviation over Simplex"] = std(y)
    os = OptimizationState(centroid(p), f(centroid(p)), iteration, d)
    if store_trace
        push!(tr, os)
    end
    if show_trace
        println(os)
    end
end

function nelder_mead(f::Function,
                     initial_x::Vector;
                     a::Real = 1.0,
                     g::Real = 2.0,
                     b::Real = 0.5,
                     tolerance::Real = 1e-8,
                     iterations::Integer = 1_000,
                     store_trace::Bool = false,
                     show_trace::Bool = false)

    # Set up a simplex of points around starting value
    m = length(initial_x)
    n = m + 1
    p = repmat(initial_x, 1, n)
    for i in 1:length(p)
        p[i] += randn()
    end

    # Count function calls
    f_calls = 0

    # Maintain a record of the value of f() at n points
    y = Array(Float64, n)
    for i in 1:n
        y[i] = f(p[:, i])
    end
    f_calls += n

    # Count iterations
    iteration = 0

    # Maintain a trace
    tr = OptimizationTrace()
    if store_trace || show_trace
        nelder_mead_trace!(tr, p, y, n, f, iteration, store_trace, show_trace)
    end

    # Monitor convergence
    converged = false

    # Cache p_bar, y_bar, p_star and p_star_star
    p_bar = Array(Float64, m)
    y_bar = Array(Float64, m)
    p_star = Array(Float64, m)
    p_star_star = Array(Float64, m)

    # Iterate until convergence or exhaustion
    while !converged && iteration < iterations
        # Augment the iteration counter
        iteration = iteration + 1

        # Find p_l and p_h, the minimum and maximum values of f() among p
        # Always take the first min or max if many exist
        y_l, l = findmin(y)
        p_l = p[:, l]
        y_h, h = findmax(y)
        p_h = p[:, h]

        # Compute the centroid of the non-maximal points
        # Cache function values of all non-maximal points
        fill!(p_bar, 0.0)
        let
            tmpindex = 0
            for i in 1:n
                if i != h
                    tmpindex += 1
                    y_bar[tmpindex] = y[i]
                    p_bar[:] += p[:, i]
                end
            end
        end
        for j in 1:m
            p_bar[j] /= m
        end

        # Compute a reflection
        for j in 1:m
            p_star[j] = (1 + a) * p_bar[j] - a * p_h[j]
        end
        y_star = f(p_star)
        f_calls += 1

        if y_star < y_l
            # Compute an expansion
            for j in 1:m
                p_star_star[j] = g * p_star[j] + (1 - g) * p_bar[j]
            end
            y_star_star = f(p_star_star)
            f_calls += 1

            if y_star_star < y_l
                p_h[:] = p_star_star
                p[:, h] = p_star_star
                y[h] = y_star_star
            else
                p_h = p_star
                p[:, h] = p_star
                y[h] = y_star
            end
        else
            if all(y_star .> y_bar)
                if y_star < y_h
                    p_h[:] = p_star
                    p[:, h] = p_h
                    y[h] = y_star
                end

                # Compute a contraction
                for j in 1:m
                    p_star_star[j] = b * p_h[j] + (1 - b) * p_bar[j]
                end
                y_star_star = f(p_star_star)
                f_calls += 1

                if y_star_star > y_h
                    for i = 1:n
                        for j in 1:m
                            p[j, i] = (p[j, i] + p_l[j]) / 2.0
                        end
                        y[i] = f(p[:, i])
                    end
                else
                    p_h[:] = p_star_star
                    p[:, h] = p_h
                    y[h] = y_star_star
                end
            else
                p_h[:] = p_star
                p[:, h] = p_h
                y[h] = y_star
            end
        end

        v = sqrt(var(y) * (m / n))

        if store_trace || show_trace
            nelder_mead_trace!(tr, p, y, n, f, iteration, store_trace, show_trace)
        end

        if v <= tolerance
            converged = true
        end
    end

    OptimizationResults("Nelder-Mead",
                        initial_x,
                        centroid(p),
                        f(centroid(p)),
                        iteration,
                        converged,
                        tr,
                        f_calls,
                        0)
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
g = (k^2 for k in -3:3)

# %%
collect(g)

# %%
v = [k^2 for k in -3:3]

# %%
Tuple(g)

# %%
t = Tuple(k^2 for k in -3:3)

# %%
Set(g)

# %%
s = Set(k^2 for k in -3:3)

# %%
d = Dict(k => k^2 for k in -3:3)

# %%
:([k^2 for k in -3:3]) |> Meta.show_sexpr

# %%
:(Tuple(k^2 for k in -3:3)) |> Meta.show_sexpr

# %%
:(Set(k^2 for k in -3:3)) |> Meta.show_sexpr

# %%
:(Dict(k => k^2 for k in -3:3)) |> Meta.show_sexpr

# %%
f(x, y) = (d = x^2+y^2; d == 0 ? d : x^2*y/d)
X = Y = range(-1, 1; length=201)

z = [f(x, y) for y in Y, x in X]
Z = f.(X', Y)
z == Z

# %%
using Plots
surface(X, Y, Z; camera=(40, 70))

# %%
surface(X, Y, f; camera=(40, 70))

# %%


using LinearAlgebra
using BaryRational
using Test
using Random

include("test_aaa.jl")
include("test_fh.jl")

@testset "BaryRational.jl" begin
    @testset "aaa_rational" begin
        @test test_aaa_gamma_poles()
        @test test_aaa_spiral()
        @test test_aaa_exp()
        @test test_aaa_length()
        @test test_aaa_tan()
        @test test_aaa_gamma()
        @test test_aaa_ranlog()
        @test test_aaa_infinp()
        @test test_aaa_naninp()
        @test test_aaa_residuals()
        #@test test_aaa_case2()
        #@test test_aaa_scale_invar()
    end
    @testset "FH_rational_interpolation" begin
        @test test_fh_runge()
        @test test_abs_x()
    end
end


@testset "pairwise" begin
    rng = MersenneTwister(123456)
    d = SqEuclidean()
    Ns = (4, 5)
    D = 3
    x = [randn(rng, D) for _ in 1:Ns[1]]
    y = [randn(rng, D) for _ in 1:Ns[2]]
    X = hcat(x...)
    Y = hcat(y...)
    K = zeros(Ns)

    @test KernelFunctions.pairwise(d, x, y) ≈ pairwise(d, X, Y, dims=2)
    @test KernelFunctions.pairwise(d, x) ≈ pairwise(d, X, dims=2)
    KernelFunctions.pairwise!(K, d, x, y)
    @test K ≈ pairwise(d, X, Y, dims=2)
    K = zeros(Ns[1], Ns[1])
    KernelFunctions.pairwise!(K, d, x)
    @test K ≈ pairwise(d, X, dims=2)

    x = randn(rng, 10)
    X = reshape(x, :, 1)
    y = randn(rng, 11)
    Y = reshape(y, :, 1)
    K = zeros(10, 11)
    @test KernelFunctions.pairwise(d, x, y) ≈ pairwise(d, X, Y; dims=1)
    @test KernelFunctions.pairwise(d, x) ≈ pairwise(d, X; dims=1)
    KernelFunctions.pairwise!(K, d, x, y)
    @test K ≈ pairwise(d, X, Y, dims=1)
end


# ### semilog x

#md # ![FILE_NAME.png](images/FILE_NAME.png)

plot(log, 0+0.00001, 5, xscale=:log10, title="semilogx")


using FrameFun
using BasisFunctions, Test, DomainSets
using BasisFunctions.Test: test_generic_dict_interface

import BasisFunctions.Test: point_outside_domain
using BasisFunctions.Test: generic_test_discrete_measure, generic_test_measure

point_outside_domain(::ExtensionFrame) = 1.

@testset "ExtensionFrame" begin
    d1 = extensionframe(Fourier(10),0.0..0.5)
    test_generic_dict_interface(d1)
    d2 = extensionframe(Fourier(10)^2,(0.0..0.5)^2)
    @test d1 == extensionframe(0.0..0.5, Fourier(10))
    @test d2 isa TensorProductDict
    @test d2 isa ExtensionFrameTensor
    @test basis(d1) == Fourier(10)
    @test basis(d2) == Fourier(10)^2
    @test support(d1) == 0.0..0.5
    @test support(d2) == (0.0..0.5)^2
    @test measure(d1) isa SubWeight
    @test measure(d2) isa ProductWeight
    @test iscompatible(extensiondual(d1, measure(d1)),component(extensiondual(d2, measure(d2)), 1))
end

@testset "SubWeight" begin
    g = PeriodicEquispacedGrid(3,-1,1)^2
    d = UnitDisk()
    sg = subgrid(g,d)
    μ = discretemeasure(sg)
    @test μ isa DiscreteSubWeight
    generic_test_discrete_measure(μ)
    @test supermeasure(μ) == discretemeasure(g)
    @test !(isnormalized(μ))
    @test restrict(supermeasure(μ),d)≈μ

    g = PeriodicEquispacedGrid(3,-1,1)×PeriodicEquispacedGrid(4,-1,1)
    d = UnitInterval()^2
    @assert subgrid(g,d) isa ProductSubGrid
    sg = subgrid(g,d)
    μ = discretemeasure(sg)
    generic_test_discrete_measure(μ )
    @test μ isa DiscreteTensorSubWeight
    @test supermeasure(μ) ≈ discretemeasure(g)
    @test components(μ) == (component(μ,1),component(μ,2))
    @test component(μ,1) == restrict(component(supermeasure(μ),1),UnitInterval())

    m = JacobiWeight(rand(),rand())
    μ = restrict(m, UnitInterval())
    generic_test_measure(μ)
    @test supermeasure(μ) == m
    μ = submeasure(m, UnitInterval())
    generic_test_measure(μ)
    @test supermeasure(μ) == m
end


@testset verbose=true "Base" begin
    @testset "Numbers $T" for T in Types
        # Note that, because `Num` from Symbolics is a weird type, we have to
        # be a little more explicit below than we normally would be.  Also,
        # because of signed zeros in the float types, we have to take the
        # absolute value of the difference before comparing to zero.

        # Define basis elements
        u = Quaternion{T}(1)
        i = Quaternion{T}(𝐢)
        j = Quaternion{T}(𝐣)
        k = Quaternion{T}(𝐤)
        basis = [u, i, j, k]

        # Check basis elements
        for (index, element) in enumerate(basis)
            for component in 1:4
                if component == index
                    @test element[component] == one(T)
                else
                    @test element[component] == zero(T)
                end
            end
        end
        @test u.w==one(T) && u.x==zero(T) && u.y==zero(T) && u.z==zero(T)
        @test i.w==zero(T) && i.x==one(T) && i.y==zero(T) && i.z==zero(T)
        @test j.w==zero(T) && j.x==zero(T) && j.y==one(T) && j.z==zero(T)
        @test k.w==zero(T) && k.x==zero(T) && k.y==zero(T) && k.z==one(T)

        # Check equality with constants; note that these are *equal*, but not *identical*
        @test u == one(T) + zero(T)*𝐢 == one(T)
        @test i == 𝐢 == imx
        @test j == 𝐣 == imy
        @test k == 𝐤 == imz

        @test Quaternion{T}(1, 2, 3, 4) == Quaternion{T}(SVector{4, T}(1, 2, 3, 4))
        @test Quaternion{T}(0, 2, 3, 4) == Quaternion{T}(2, 3, 4)
        @test Quaternion{T}(1, 0, 0, 0) == Quaternion{T}(1)
        @test Quaternion(T.([1, 2, 3, 4])...) == Quaternion(SVector{4, T}(1, 2, 3, 4))
        @test Quaternion(T.([0, 2, 3, 4])...) == Quaternion(T(2), T(3), T(4))
        @test Quaternion(T.([1, 0, 0, 0])...) == Quaternion(T(1))
        if !(T<:Integer)
            @test Rotor{T}(1, 2, 3, 4) == Rotor{T}(SVector{4, T}(1, 2, 3, 4)/√T(30))
            @test Rotor{T}(0, 2, 3, 4) == Rotor{T}(2, 3, 4)
            @test Rotor{T}(1, 0, 0, 0) == Rotor{T}(1)
            @test Rotor(T.([1, 2, 3, 4])...) == Rotor(SVector{4, T}(1, 2, 3, 4)/√T(30))
            @test Rotor(T.([0, 2, 3, 4])...) == Rotor(T(2), T(3), T(4))
            @test Rotor(T.([1, 0, 0, 0])...) == Rotor(T(1))
        end
        @test QuatVec{T}(1, 2, 3, 4) == QuatVec{T}(SVector{4, T}(0, 2, 3, 4))
        @test QuatVec{T}(0, 2, 3, 4) == QuatVec{T}(2, 3, 4)
        @test QuatVec{T}(1, 0, 0, 0) == QuatVec{T}(1)
        @test QuatVec(T.([0, 2, 3, 4])...) == QuatVec(SVector{4, T}(0, 2, 3, 4))
        @test QuatVec(T.([0, 2, 3, 4])...) == QuatVec(T(2), T(3), T(4))
        @test QuatVec(T.([0, 0, 0, 0])...) == QuatVec(T(0))

        # Test indexing
        q = Quaternion(T(1), T(2), T(3), T(4))
        @test q[[1, 2]] == [T(1), T(2)]
        @test q[[3, 4]] == [T(3), T(4)]
        @test q[[4, 2]] == [T(4), T(2)]
        @test q[[4, 2, 3]] == [T(4), T(2), T(3)]
        @test q[[4, 2, 3, 1]] == [T(4), T(2), T(3), T(1)]

        # Test copy constructor and self-equality
        @test Quaternion(u) == Quaternion{T}(u) == u
        @test Quaternion(i) == Quaternion{T}(i) == i
        @test Quaternion(j) == Quaternion{T}(j) == j
        @test Quaternion(k) == Quaternion{T}(k) == k
        @test Quaternionic.wrapper(Quaternion(u)) == Quaternion
        @test Quaternionic.wrapper(Rotor(u)) == Rotor
        @test Quaternionic.wrapper(QuatVec(u)) == QuatVec
        @test u == one(T)
        @test one(T) == u
        @test i != one(T)
        @test one(T) != i
        @test j != one(T)
        @test one(T) != j
        @test k != one(T)
        @test one(T) != k
        @test isequal(u, u)
        @test !isequal(u, i)
        @test !isequal(u, j)
        @test !isequal(u, k)
        @test !isequal(i, u)
        @test isequal(i, i)
        @test !isequal(i, j)
        @test !isequal(i, k)
        @test !isequal(j, u)
        @test !isequal(j, i)
        @test isequal(j, j)
        @test !isequal(j, k)
        @test !isequal(k, u)
        @test !isequal(k, i)
        @test !isequal(k, j)
        @test isequal(k, k)

        # Check isapprox
        @test u ≈ one(T)
        @test one(T) ≈ u
        if T !== Num  # Num(1) ≉ Num(2) doesn't work
            @test i ≉ one(T)
            @test one(T) ≉ i
            @test j ≉ one(T)
            @test one(T) ≉ j
            @test k ≉ one(T)
            @test one(T) ≉ k
        end

        # Check non-sensical basis elements are not allowed
        @test_throws DomainError one(QuatVec)
        @test_throws DomainError one(QuatVec{T})
        @test_throws DomainError one(QuatVec{T}(1, 2, 3))
        @test_throws DomainError zero(Rotor)
        @test_throws DomainError zero(Rotor{T})
        @test_throws DomainError zero(Rotor{T}(1))

        # Check "real" part
        @test real(u) == one(T)
        @test real(i) == zero(T)
        @test real(j) == zero(T)
        @test real(k) == zero(T)

        # Check "imag" part
        @test imag(u) == [zero(T), zero(T), zero(T)]
        @test imag(i) == [one(T), zero(T), zero(T)]
        @test imag(j) == [zero(T), one(T), zero(T)]
        @test imag(k) == [zero(T), zero(T), one(T)]

        # Check "isreal"
        @test isreal(u)
        @test !isreal(i)
        @test !isreal(j)
        @test !isreal(k)

        # Check "isinteger"
        if T != Num
            @test isinteger(u)
            @test !isinteger(1.2u)
        end
        @test !isinteger(i)
        @test !isinteger(j)
        @test !isinteger(k)

        if T<:AbstractFloat
            # Check "isnan"
            @test !isnan(u)
            @test !isnan(i)
            @test !isnan(j)
            @test !isnan(k)
            @test isnan(T(NaN) + 0imx)
            @test isnan(T(NaN)imx)
            @test isnan(T(NaN)imy)
            @test isnan(T(NaN)imz)

            # Check "isinf"
            @test !isinf(u)
            @test !isinf(i)
            @test !isinf(j)
            @test !isinf(k)
            @test isinf(T(Inf) + 0imx)
            @test isinf(T(Inf)imx)
            @test isinf(T(Inf)imy)
            @test isinf(T(Inf)imz)
        end

        # Check "isone"
        @test isone(u)
        @test !isone(2u)
        @test !isone(i)
        @test !isone(j)
        @test !isone(k)

        # Check "flipsign"
        @test flipsign(u, 1) == u
        @test flipsign(u, -1) == -u
        @test flipsign(i, 1) == i
        @test flipsign(i, -1) == -i
        @test flipsign(j, 1) == j
        @test flipsign(j, -1) == -j
        @test flipsign(k, 1) == k
        @test flipsign(k, -1) == -k

        if T != Num
            # Check "in"
            @test u ∈ 0:2
            @test u ∉ 2:4
            @test i ∉ 0:2
            @test j ∉ 0:2
            @test k ∉ 0:2
        end

    end

    @testset "bswap" begin
        @test bswap(Quaternion(1)) == Quaternion(2^56)
        @test bswap(1imx) == (2^56)imx
        @test bswap(1imy) == (2^56)imy
        @test bswap(1imz) == (2^56)imz
        @test bswap(bswap(Quaternion(1))) == Quaternion(1)
        @test bswap(bswap(1imx)) == 1imx
        @test bswap(bswap(1imy)) == 1imy
        @test bswap(bswap(1imz)) == 1imz
    end

    @testset "hash" begin
        for T1 in [FloatTypes...; IntTypes...]
            for T2 in [FloatTypes...; IntTypes...]
                q1 = Quaternion{T1}(1, 2, 3, 4)
                q2 = Quaternion{T2}(1, 2, 3, 4)
                @test isequal(q1, q2) && hash(q1)==hash(q2)
            end
        end

        for T1 in FloatTypes
            for T2 in FloatTypes
                q1 = Quaternion(T1(0.0))
                q2 = Quaternion(T2(-0.0))
                @test !isequal(q1, q2) && hash(q1)!=hash(q2)
                q1 = Quaternion(T1(NaN))
                q2 = Quaternion(T2(NaN))
                @test isequal(q1, q2) && hash(q1)==hash(q2)
            end
        end
    end

    @testset "io" begin
        io = IOBuffer()

        Base.show(io, MIME("text/plain"), Quaternion{Float64}(1, 2, 3, 4))
        @test String(take!(io)) == "1.0 + 2.0𝐢 + 3.0𝐣 + 4.0𝐤"
        Base.show(io, MIME("text/plain"), Quaternion{Int64}(1, 2, 3, 4))
        @test String(take!(io)) == "1 + 2𝐢 + 3𝐣 + 4𝐤"
        Base.show(io, MIME("text/plain"), Quaternion(a-b, b*c, c/d, d+e))
        @test String(take!(io)) == "a - b + b*c𝐢 + (c / d)𝐣 + (d + e)𝐤"
        Base.show(io, MIME("text/latex"), Quaternion{Float64}(1, 2, 3, 4))
        @test String(take!(io)) == "\$1.0 + 2.0\\,\\mathbf{i} + 3.0\\,\\mathbf{j} + 4.0\\,\\mathbf{k}\$"
        Base.show(io, MIME("text/latex"), Quaternion{Int64}(1, 2, 3, 4))
        @test String(take!(io)) == "\$1 + 2\\,\\mathbf{i} + 3\\,\\mathbf{j} + 4\\,\\mathbf{k}\$"
        Base.show(io, MIME("text/latex"), Quaternion(a-b, b*c, c/d, d+e))
        @test String(take!(io)) == "\$a - b + b c\\,\\mathbf{i} + \\frac{c}{d}\\,\\mathbf{j} + \\left(d + e\\right)\\,\\mathbf{k}\$"

        for T in PrimitiveTypes
            io = IOBuffer()
            q = Quaternion{T}(1, 2, 3, 4)
            write(io, q)
            seekstart(io)
            p = read(io, typeof(q))
            @test q == p
        end
    end

end


################################################
# Approximate info-reward functions
################################################

function negentropy(p::AbstractArray{T}) where T<:Real
    s = zero(T)
    z = zero(T)
    for i = 1:length(p)
        @inbounds pi = p[i]
        if pi > z
            s += pi * log2(pi)
        end
    end
    return log2(length(p)) + s
end

using Distances, Plots; gr()
x = 0:0.01:1
y1 = [negentropy([x[i], 1 - x[i]]) for i = 1:length(x)]
y2 = [cityblock([x[i], 1 - x[i]], [0.5, 0.5]) for i = 1:length(x)]
y3 = [euclidean([x[i], 1 - x[i]], [0.5, 0.5]) for i = 1:length(x)]
y4 = [norm([x[i], 1 - x[i]] - [0.5, 0.5], Inf) for i = 1:length(x)]
plot(x, y1, label = "Negative Entropy", legend = :top,
    xlab = "Belief, P(State = 1)",
    ylab = "Belief Reward",
    linestyle = :solid,
    linewidth = 2)
plot!(x, y2, label = "1-Norm", linestyle = :dash,
    linecolor = :red, linealpha = 1, linewidth = 2)
plot!(x, y3, label = "2-Norm", linestyle = :dot,
    linecolor = :red, linealpha = 1, linewidth = 2)
plot!(x, y4, label = "Infinity-Norm", linestyle = :dashdot,
    linecolor = :red, linealpha = 0.6, linewidth = 2)
fn = string("belief_reward_function.pdf")
path = joinpath(homedir(),".julia\\v0.6\\RobustInfoPOMDP\\data\\figures\\",fn)
savefig(path)

x = 1:0.1:1
y = 1:0.1:1
f(x,y) = negentropy([x - y])
X = repmat(x',length(y),1)
Y = repmat(y,1,length(x))
Z = map(f,X,Y)
p1 = contour(x,y,f,fill=true)
p2 = contour(x,y,Z)
plot(p1,p2)

using ForwardDiff
dne(x) = ForwardDiff.gradient(negentropy, x)

function ev(i::Int, len::Int)
    v = zeros(len)
    v[i] = 1.0
    return v
end

e1 = ev(1,4)
negentropy(e1)
negentropy(c)
negentropy([0.2, 0.25, 0.25, 0.3])
dne([0.2, 0.25, 0.25, 0.3])

ints = 0.1:0.1:0.9
n = length(ints)
alphas = Vector{Vector{Float64}}(n)
for (i, int) in enumerate(ints)
    x = [int, 1 - int]
    m = dne(x)
    y = negentropy(x)
    b = y - dot(m,x)
    a1 = dot(m, [0,1]) + b
    a2 = dot(m, [1,0]) + b
    alphas[i] = [a1, a2]
end

plot(x, y1, linecolor = :black, legend = :bottomright)
plot([0,1], alphas, linealpha = 0.2, linecolor = :black)

y

dot(m,[0,1]) + b
dot(m,[1,0]) + b

ax = [-2.32, 0.848]
plot!([0,1],ax)

using RPOMDPModels
ip = CyberIPOMDP()

rmax = -Inf
b = fill(1/27, 27)
for α in ip.inforeward
    rmax = max(rmax, dot(α, b))
end
rmax

ip.inforeward

# complicated info reward function
ra21 = [1.0, 0.0]
ra22 = [0.9, 0.1]
ra23 = [0.8, 0.2]
ra24 = [0.7, 0.3]
ralphas2 = [ra21, ra22, ra23, ra24, 1-ra24, 1-ra23, 1-ra22, 1-ra21]
plot([0,1], ralphas2)

using Plots; gr()
using Distances

function tinfo(rinfo::Vector{Vector{Float64}}, b::Vector{Float64})
    rmax = -Inf
    for α in rinfo
        rmax = max(rmax, dot(α, b))
    end
    rmax
end

# simple info reward function
irs = [[1.0, -1/3, -1/3, -1/3],
         [-1/3, 1.0, -1/3, -1/3],
         [-1/3, -1/3, 1.0, -1/3],
         [-1/3, -1/3, -1/3, 1.0]]

c = fill(1/4, 4)
e1 = ev(1,4)
tinfo(irs, c)

# complicated info reward function
vhi = -10.0
vlo = 0.1
irc = [[1.0, vhi, vhi, vhi],
        [vhi, 1.0, vhi, vhi],
        [vhi, vhi, 1.0, vhi],
        [vhi, vhi, vhi, 1.0],
        [vlo, -vlo/3, -vlo/3, -vlo/3],
        [-vlo/3, vlo, -vlo/3, -vlo/3],
        [-vlo/3, -vlo/3, vlo, -vlo/3],
        [-vlo/3, -vlo/3, -vlo/3, vlo]]

tinfo(irc, c)
tinfo(irc, e1)
tinfo(irs, [0.5, 0.5, 0.0, 0.0])
tinfo(irc, [0.5, 0.5, 0.0, 0.0])
negentropy([0.5, 0.5, 0.0, 0.0])
norm([0.5, 0.5, 0.0, 0.0] - c, 1)
x = 0:0.001:1.0
y0 = [norm([xi, 1-xi, 0.0, 0.0] - c, 1) for xi in x]
y1 = [negentropy([xi, 1-xi, 0.0, 0.0]) for xi in x]
y2 = [tinfo(irs, [xi, 1-xi, 0.0, 0.0]) for xi in x]
y3 = [tinfo(irc, [xi, 1-xi, 0.0, 0.0]) for xi in x]
plot(x, y2, label = "Simple", legend = :bottomleft,
    xlab = "Belief, P(State = Bad | Position = 1)",
    ylab = "Expected Total Discounted Reward")
plot!(x, y3, label = "Complex")
plot!(x, y2)
plot!(x, y3)


"""
$(TYPEDEF)

A system of partial differential equations.

# Fields
$(FIELDS)

# Example

```julia
using ModelingToolkit

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

#2D PDE
C=1
eq  = Dtt(u(t,x)) ~ C^2*Dxx(u(t,x))

# Initial and boundary conditions
bcs = [u(t,0) ~ 0.,# for all t > 0
       u(t,1) ~ 0.,# for all t > 0
       u(0,x) ~ x*(1. - x), #for all 0 < x < 1
       Dt(u(0,x)) ~ 0. ] #for all  0 < x < 1]

# Space and time domains
domains = [t ∈ (0.0,1.0),
           x ∈ (0.0,1.0)]

@named pde_system = PDESystem(eq,bcs,domains,[t,x],[u])
```
"""
struct PDESystem <: ModelingToolkit.AbstractMultivariateSystem
    "The equations which define the PDE"
    eqs::Any
    "The boundary conditions"
    bcs::Any
    "The domain for the independent variables."
    domain::Any
    "The independent variables"
    ivs::Any
    "The dependent variables"
    dvs::Any
    "The parameters"
    ps::Any
    """
    defaults: The default values to use when initial conditions and/or
    parameters are not supplied in `ODEProblem`.
    """
    defaults::Dict
    """
    type: type of the system
    """
    connector_type::Any
    """
    systems: The internal systems. These are required to have unique names.
    """
    systems::Vector
    """
    name: the name of the system
    """
    name::Symbol
    @add_kwonly function PDESystem(eqs, bcs, domain, ivs, dvs,
                                   ps = SciMLBase.NullParameters();
                                   defaults = Dict(),
                                   systems = [],
                                   connector_type = nothing,
                                   checks::Bool = true,
                                   name)
        if checks
            all_dimensionless([dvs; ivs; ps]) || check_units(eqs)
        end
        eqs = eqs isa Vector ? eqs : [eqs]
        new(eqs, bcs, domain, ivs, dvs, ps, defaults, connector_type, systems, name)
    end
end

function Base.getproperty(x::PDESystem, sym::Symbol)
    if sym == :indvars
        return getfield(x, :ivs)
        Base.depwarn("`sys.indvars` is deprecated, please use `get_ivs(sys)`", :getproperty,
                     force = true)

    elseif sym == :depvars
        return getfield(x, :dvs)
        Base.depwarn("`sys.depvars` is deprecated, please use `get_dvs(sys)`", :getproperty,
                     force = true)

    else
        return getfield(x, sym)
    end
end

Base.summary(prob::PDESystem) = string(nameof(typeof(prob)))
function Base.show(io::IO, ::MIME"text/plain", sys::PDESystem)
    println(io, summary(sys))
    println(io, "Equations: ", get_eqs(sys))
    println(io, "Boundary Conditions: ", get_bcs(sys))
    println(io, "Domain: ", get_domain(sys))
    println(io, "Dependent Variables: ", get_dvs(sys))
    println(io, "Independent Variables: ", get_ivs(sys))
    println(io, "Parameters: ", get_ps(sys))
    print(io, "Default Parameter Values", get_defaults(sys))
    return nothing
end


# y: data
# yhat_generator: function that samples data given the parameters
# algo parameters: Namedtuple of algorithm parameters
# max_time: max time allowed for the algo
# N_samples: number of samples desired
function ABC_SMC(y,yhat_generator,algo_parameters,N_samples)
    CPUtic()
    resample_method = algo_parameters[:resample_method]
    kernel = algo_parameters[:kernel]
    kernel_density = algo_parameters[:kernel_density]
    sd = algo_parameters[:sd]
    eta_y = algo_parameters[:eta](y)
    size_y = size(y)
    q = length(algo_parameters[:prior])
    eps_list = algo_parameters[:eps_list]
    time_final = algo_parameters[:time_final]
    verbose = algo_parameters[:verbose]

    ess_list = zeros(time_final)
    ess_threshold = Int(N_samples*0.6)
    result = zeros(N_samples,q)
    weights = ones(N_samples)./N_samples
    theta = zeros(N_samples,q)
    t = 1
    # sample from the prior initially (same as rejection ABC)
    nParticle=1
    while nParticle < N_samples
        p = rand.(algo_parameters[:prior])
        yhat = yhat_generator(p)
        if algo_parameters[:d](eta_y,algo_parameters[:eta](yhat)) < eps_list[t]
            result[nParticle,:].=p
            nParticle +=1
        end
    end
    ess_list[t] = 1/(sum(weights.^2))
    t += 1

    result_prev = result
    weights_prev = weights
    while (t <= time_final)
        nParticle=1
        p_cand = Tuple(zeros(q))
        while nParticle <= N_samples
            # sample the index with with weights
            sampled_index = wsample(1:N_samples,weights_prev,1)
            p_star = Tuple(result[sampled_index,:])
            # propose candidate parameter from the kernel
            valid = true
            while valid
                p_cand = kernel(p_star, sd)
                valid = any(p_cand .< 0)
            end
            # if the density  of the prior at proposed param is 0, start again
            prior_cand = prod(pdf.(algo_parameters[:prior],p_cand))
            if prior_cand == 0
                continue
            end
            yhat = yhat_generator(p_cand)
            if (algo_parameters[:d](eta_y, algo_parameters[:eta](yhat)) < eps_list[t])
                # accept the proposed parameter
                result[nParticle,:] .= p_cand
                # Compute the weight of the current particle
                weight_denom = 0
                for j=1:N_samples
                    weight_denom += weights_prev[j]*prod((kernel_density.(Tuple(result_prev[j,:]), p_cand, sd) ))
                end
                weights[nParticle] = prior_cand/weight_denom
                nParticle+=1
            end

        end
        # Normalize weights
        weights = weights./sum(weights)
        # Compute effective sample size
        ess = 1/sum(weights.^2)
        ess_list[t] = ess
        # resample if the effective sample size is below the threshold
        if ess < ess_threshold
            idx = resample_method(weights)
            result = result[idx,:]
        end
        result_prev = result
        t += 1
    end
    t_end = CPUtoc_modified(false)
    if verbose
        @show ess_list
        @printf("Done SMC\n")
    end
    return result, t_end
end


module TestSparse
using LinearAlgebra
using CovarianceFunctions
using SparseArrays
using Test

@testset "sparse" begin
    n = 128
    d = 32
    X = randn(d, n)
    k = CovarianceFunctions.EQ()
    δ = 1e-6
    G = gramian(k, X)
    S = sparse(G, δ)
    @test isapprox(G, S, atol = δ, rtol = δ)
end

end # module


using StochasticDiffEq,  DiffEqNoiseProcess, Test

f1(u, p, t) = 1.01u
g1(u, p, t) = 1.01u
dt = 1//2^(4)
prob1 = SDEProblem(f1,g1,1.0,(0.0,1.0))
sol1 = solve(prob1,EM(),dt=dt,save_noise = true)

W2 = NoiseWrapper(sol1.W)
prob2 = SDEProblem(f1,g1,1.0,(0.0,1.0),noise=W2)
dt = 1//2^(5)
sol2 = solve(prob2,EM(),dt=dt)

W3 = NoiseWrapper(sol1.W)
prob3 = SDEProblem(f1,g1,1.0,(0.0,1.0))
dt = 1//2^(5)
sol3 = solve(prob3,EM(),dt=dt,noise=W3)

using Plots
plotly()
plot(sol1, label="reference")
plot!(sol2, label="good coarse reconstruction")
plot!(sol3, label="bad coarse reconstruction")


# Adapted from:
# SOSDEMO5 --- Upper bound for the structured singular value mu
# Section 3.5 of SOSTOOLS User's Manual

using SparseArrays

function sosdemo5_test(optimizer, config::MOIT.TestConfig, feasible::Bool, γ)
    @polyvar x[1:8]

    # The matrix under consideration
    α = 3 + sqrt(3);
    β = sqrt(3) - 1;
    a = sqrt(2/α);
    b = 1/sqrt(α);
    c = b;
    d = -sqrt(β/α);
    f = (1 + im)*sqrt(1/(α*β));
    U = [a 0; b b; c im*c; d f];
    V = [0 a; b -b; c -im*c; -im*f -d];
    M = U*V';

    Z = monomials(x, 1)

    function build_A(i)
        H = M[i,:]*M[i,:]' - (γ^2)*sparse([i],[i],[1],4,4)
        H = [real(H) -imag(H); imag(H) real(H)]
        dot(Z, H*Z)
    end
    A = build_A.(1:4)

    model = _model(optimizer)
    PolyJuMP.setpolymodule!(model, SumOfSquares)

    # -- Q(x)'s -- : sums of squares
    # Monomial vector: [x1; ... x8]
    Q = Vector{GramMatrix{JuMP.VariableRef, monomialtype(x[1]),
                          monovectype(x[1])}}(undef, 4)
    @variable model Q[1:4] SOSPoly(Z)

    # -- r's -- : constant sum of squares
    Z = monomials(x, 0)
    #r = Matrix{GramMatrix{JuMP.VariableRef}}(4,4) # FIXME doesn't work with 1x1 SDP matrix :(
    @variable model r[i=1:4,j=(i+1):4] >= 0

    # Constraint : -sum(Qi(x)*Ai(x)) - sum(rij*Ai(x)*Aj(x)) + I(x) >= 0
    expr = 0
    # Adding term
    for i = 1:4
        expr -= A[i]*Q[i]
    end
    for i = 1:4
        for j = (i+1):4
            expr -= A[i]*A[j]*r[i,j]
        end
    end
    # Constant term: I(x) = -(x1^4 + ... + x8^4)
    I = -sum(x.^4)
    expr = expr + I

    @constraint model expr >= 0

    JuMP.optimize!(model)

    # Program is feasible, thus 0.8724 is an upper bound for mu.
    if feasible
        @test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
    else
        @test JuMP.dual_status(model) == MOI.INFEASIBILITY_CERTIFICATE
    end
end
const γ_opt = 0.8723
sosdemo5_infeasible_test(optimizer, config, ε = 1e-2) = sosdemo5_test(optimizer, config, false, γ_opt - ε)
sosdemo5_feasible_test(optimizer, config, ε = 1e-2) = sosdemo5_test(optimizer, config, true, γ_opt + ε)
sd_tests["sosdemo5_infeasible"] = sosdemo5_infeasible_test
sd_tests["sosdemo5_feasible"] = sosdemo5_feasible_test