using LinearAlgebra

import PotentialFlow.Utils: derivative, extract_derivative, value, partials,
          Dual,ComplexDual, ComplexGradientConfig,
          dz_partials, gradient, chunk_mode_gradient, chunk_mode_jacobian

import PotentialFlow.Elements: gradient_position, gradient_strength,
          jacobian_position, jacobian_strength, jacobian_param

const DELTA=1e-6
const BIGEPS = 1e-8
const TOL=5e-6
const BIGTOL=1e-4
const BIGGESTTOL=1e-3

safenorm(a) = norm(filter(x -> ~isnan(x),a))

@testset "Complex Automatic Differentiation" begin

  @testset "Basic derivatives" begin

    z = rand(ComplexF64)

    dz1, dzstar1 = derivative(z -> log(sqrt(z)),z)

    dzex = 1/(2*z)

    @test isapprox(abs(dz1-dzex),0,atol=BIGEPS)
    @test isapprox(abs(dzstar1),0,atol=BIGEPS)

    dz1, dzstar1 = derivative(z -> conj(z),z)
    @test dzstar1 == one(z) && dz1 == zero(z)

    f(z) = 0.5im/(π*conj(z))

    dz1, dzstar1 = derivative(f,z)
    dzstarex = -0.5im/(pi*conj(z)^2)

    @test isapprox(abs(dz1),0,atol=BIGEPS)
    @test isapprox(abs(dzstar1-dzstarex),0,atol=BIGEPS)

    z0 = rand(ComplexF64)
    dz1, dzstar1 = derivative(z -> f(z-z0),z)

    dzstarex = -0.5im/(pi*conj(z-z0)^2)

    @test isapprox(abs(dz1),0,atol=BIGEPS)
    @test isapprox(abs(dzstar1-dzstarex),0,atol=BIGEPS)

    dz1, dzstar1 = derivative(abs,z)

    dzex = 0.5*conj(z)/abs(z)
    dzstarex = 0.5*z/abs(z)
    @test isapprox(abs(dz1-dzex),0,atol=BIGEPS)
    @test isapprox(abs(dzstar1-dzstarex),0,atol=BIGEPS)

    dz1, dzstar1 = derivative(z -> z + sqrt(log(z)) - 1/z^2,z)

    dzex = 1 + 1/(2*sqrt(log(z))*z) + 2/z^3
    dzstarex = 0.0

    @test isapprox(abs(dz1-dzex),0,atol=BIGEPS)
    @test isapprox(abs(dzstar1-dzstarex),0,atol=BIGEPS)

    dz1, dzstar1 = derivative(z -> real(z - √(z - 1)*√(z + 1)),z)

    dzex = 0.5*(1 - 0.5*sqrt((z+1)/(z-1)) - 0.5*sqrt((z-1)/(z+1)))
    dzstarex = 0.5*conj(1 - 0.5*sqrt((z+1)/(z-1)) - 0.5*sqrt((z-1)/(z+1)))

    @test isapprox(abs(dz1-dzex),0,atol=BIGEPS)
    @test isapprox(abs(dzstar1-dzstarex),0,atol=BIGEPS)

  end

  @testset "Nested derivatives" begin

    z0 = rand(ComplexF64)
    y0 = rand(ComplexF64)
    dyz, dystarz, dyzstar, dystarzstar = derivative(y -> derivative(z -> 1/(z-y),z0),y0)

    dyzex = -2/(z0-y0)^3
    @test isapprox(abs(dyz-dyzex),0,atol=BIGEPS)
    @test isapprox(abs(dystarz),0,atol=BIGEPS)
    @test isapprox(abs(dyzstar),0,atol=BIGEPS)
    @test isapprox(abs(dystarzstar),0,atol=BIGEPS)

    z0 = rand(ComplexF64)
    y0 = rand(ComplexF64)
    dyz, dystarz, dyzstar, dystarzstar = derivative(y -> derivative(z -> 1/(z-conj(y)),z0),y0)

    dystarzex = -2/(z0-conj(y0))^3
    @test isapprox(abs(dyz),0,atol=BIGEPS)
    @test isapprox(abs(dystarz-dystarzex),0,atol=BIGEPS)
    @test isapprox(abs(dyzstar),0,atol=BIGEPS)
    @test isapprox(abs(dystarzstar),0,atol=BIGEPS)

    z0 = rand(ComplexF64)
    y0 = rand(ComplexF64)
    dyz, dystarz, dyzstar, dystarzstar = derivative(y -> derivative(z -> 1/(conj(z)-y),z0),y0)

    dyzstarex = -2/(conj(z0)-y0)^3
    @test isapprox(abs(dyz),0,atol=BIGEPS)
    @test isapprox(abs(dystarz),0,atol=BIGEPS)
    @test isapprox(abs(dyzstar-dyzstarex),0,atol=BIGEPS)
    @test isapprox(abs(dystarzstar),0,atol=BIGEPS)

    z0 = rand(ComplexF64)
    y0 = rand(ComplexF64)
    dyz, dystarz, dyzstar, dystarzstar = derivative(y -> derivative(z -> 1/conj(z-y),z0),y0)

    dystarzstarex = -2/conj(z0-y0)^3
    @test isapprox(abs(dyz),0,atol=BIGEPS)
    @test isapprox(abs(dystarz),0,atol=BIGEPS)
    @test isapprox(abs(dyzstar),0,atol=BIGEPS)
    @test isapprox(abs(dystarzstar-dystarzstarex),0,atol=BIGEPS)

  end

  @testset "Chunk mode derivatives" begin

    n = 20
    z = rand(ComplexF64,n)
    fcn(z) = sum(log.(sqrt.(z)))

    cfg = ComplexGradientConfig(fcn,z)
    dz, dzstar = gradient(fcn,z,cfg)

    ichk = rand(1:n)
    @test dz[ichk] ≈ 1.0/(2*z[ichk])

    N = rand(1:n)
    cfg2 = ComplexGradientConfig(fcn,z,PotentialFlow.Utils.Chunk{N}())
    dz_chunk, dzstar_chunk = gradient(fcn,z,cfg2)
    @test dz == dz_chunk
    @test dzstar == dzstar_chunk

    n = 20
    z = rand(ComplexF64,n)
    fcn2(z) = log.(sqrt.(z))

    cfg = ComplexGradientConfig(fcn2,z)
    dz, dzstar = PotentialFlow.Utils.jacobian(fcn2,z,cfg)

    ichk = rand(1:n)
    @test dz[ichk,ichk] ≈ 1.0/(2*z[ichk])

    N = rand(1:n)
    cfg2 = ComplexGradientConfig(fcn2,z,PotentialFlow.Utils.Chunk{N}())

    dz_chunk, dzstar_chunk = PotentialFlow.Utils.jacobian(fcn2,z,cfg2)

    @test dz_chunk == dz
    @test dzstar_chunk == dzstar

  end

  nblob = 5
  pos = rand(ComplexF64,nblob)
  str = rand(length(pos))

  σ = 1e-2
  blobs = Vortex.Blob.(pos,str,σ)
  z = rand(ComplexF64)

  i = rand(1:nblob)

  # For finite difference approximations
  dz = zeros(ComplexF64,length(blobs))
  dz[i] = DELTA
  dΓ = zeros(Float64,length(blobs))
  dΓ[i] = DELTA
  blobsx⁺ = Vortex.Blob.(Elements.position(blobs).+dz,Elements.circulation.(blobs),σ)
  blobsy⁺ = Vortex.Blob.(Elements.position(blobs).+im*dz,Elements.circulation.(blobs),σ)
  blobsΓ⁺ = Vortex.Blob.(Elements.position(blobs),Elements.circulation.(blobs).+dΓ,σ)

  @testset "Basic operations with duals" begin

    dualpos = one(ComplexDual{Nothing},Elements.position(blobs)[i])

    @test value(dualpos) == Elements.position(blobs)[i]
    pr, pi = partials(dualpos)
    @test pr == [1.0,0.0] && pi == [0.0,1.0]

    newblob = Vortex.Blob(dualpos,Elements.circulation(blobs[i]),σ)
    @test value(Elements.position(newblob)) == Elements.position(blobs)[i]

    dwdz,dwdzstar  = extract_derivative(Nothing,induce_velocity(z,newblob,0.0))

    w2 = Elements.circulation(newblob)*PotentialFlow.Blobs.blob_kernel(z - Elements.position(newblob),σ)
    dwdz2, dwdzstar2 = extract_derivative(Nothing,w2)

    @test isapprox(abs(dwdz-dwdz2),0.0,atol=BIGEPS)
    @test isapprox(abs(dwdzstar-dwdzstar2),0.0,atol=BIGEPS)

end

@testset "Induced velocity at a point" begin

    compute_velocity(v) = induce_velocity(z,v,0.0)

    # Finite difference approximations
    w_fd = compute_velocity(blobs)
    wx⁺_fd = compute_velocity(blobsx⁺)
    wy⁺_fd = compute_velocity(blobsy⁺)
    wΓ⁺_fd = compute_velocity(blobsΓ⁺)
    dwdx_fd = (wx⁺_fd - w_fd)/dz[i]
    dwdy_fd = (wy⁺_fd - w_fd)/dz[i]
    dwdz_fd = 0.5*(dwdx_fd - im*dwdy_fd)
    dwdzstar_fd = 0.5*(dwdx_fd + im*dwdy_fd)
    dwdΓ_fd = (wΓ⁺_fd - w_fd)/dΓ[i]

    #=
    dwdx_fd = (induce_velocity(z,blobsx⁺,0.0) - induce_velocity(z,blobs,0.0))/dz[i]
    dwdy_fd =(induce_velocity(z,blobsy⁺,0.0) - induce_velocity(z,blobs,0.0))/dz[i]
    dwdz_fd = 0.5*(dwdx_fd - im*dwdy_fd)
    dwdzstar_fd = 0.5*(dwdx_fd + im*dwdy_fd)
    dwdΓ_fd = (induce_velocity(z,blobsΓ⁺,0.0) - induce_velocity(z,blobs,0.0))/dΓ[i]
    =#

    # Auto differentation
    cfg = ComplexGradientConfig(z -> (),blobs)
    newblobs = Vortex.seed_position(blobs,cfg)
    dwdz2, dwdzstar2 = dz_partials(induce_velocity(z,newblobs,0.0))

    # chunk mode
    nchunk = 3
    cfg2 = PotentialFlow.Utils.ComplexGradientConfig(z -> (),blobs,PotentialFlow.Utils.Chunk{nchunk}())

    # Auto diff with the API
    dwdz, dwdzstar = gradient_position(compute_velocity,blobs)
    @test dwdz == dwdz2 && dwdzstar == dwdzstar2

    @test_skip abs(dwdz[i]-dwdz_fd) ≈ 0.0 atol=TOL
    @test_skip abs(dwdzstar[i]-dwdzstar_fd) ≈ 0.0 atol=TOL

    dwdz_chunk, dwdzstar_chunk = gradient_position(compute_velocity,blobs,cfg2)
    @test dwdz == dwdz_chunk && dwdzstar == dwdzstar_chunk

    newblobs = Vortex.seed_strength(blobs,cfg)

    @test sum(value.(Vortex.circulation.(newblobs))) -
              value(Vortex.circulation(newblobs)) == 0

    dwdz_tmp, dwdzstar_tmp = dz_partials(induce_velocity(z,newblobs,0.0))
    dwdΓ2 = dwdz_tmp+dwdzstar_tmp

    # Auto diff with the API
    dwdΓ = gradient_strength(compute_velocity,blobs)
    @test dwdΓ == dwdΓ2

    @test_skip abs(dwdΓ[i]-dwdΓ_fd) ≈ 0.0 atol=TOL

    dwdΓ_chunk = gradient_strength(compute_velocity,blobs,cfg2)

    @test dwdΓ  == dwdΓ_chunk

end

N = 7
L = 2.0
c = complex(0)
α = 0.0
ċ = complex(0.0)
α̇ = 0.0
p = PotentialFlow.Plate(N,L,c,α)
motion = PotentialFlow.RigidBodyMotion(ċ,α̇)

Plates.enforce_no_flow_through!(p, motion, blobs, 0.0)
px⁺ = deepcopy(p)
Plates.enforce_no_flow_through!(px⁺, motion, blobsx⁺, 0.0)
py⁺ = deepcopy(p)
Plates.enforce_no_flow_through!(py⁺, motion, blobsy⁺, 0.0)
pΓ⁺ = deepcopy(p)
Plates.enforce_no_flow_through!(pΓ⁺, motion, blobsΓ⁺, 0.0)

@testset "Induced velocity at plate points" begin

    function plate_velocity(v)
      w = zeros(Complex{Elements.property_type(eltype(v))},N)
      induce_velocity!(w,p,v,0.0)
      return w
    end

    C = plate_velocity(blobs)
    Cx⁺ = plate_velocity(blobsx⁺)
    Cy⁺ = plate_velocity(blobsy⁺)
    CΓ⁺ = plate_velocity(blobsΓ⁺)

    dwdx_fd = (Cx⁺ - C)/dz[i]
    dwdy_fd = (Cy⁺ - C)/dz[i]
    dwdz_fd = 0.5*(dwdx_fd - im*dwdy_fd)
    dwdzstar_fd = 0.5*(dwdx_fd + im*dwdy_fd)

    # Auto diff
    dwdz, dwdzstar = jacobian_position(plate_velocity,blobs)

    # test that the induced velocities and their derivatives match
    #@test isapprox(norm(value.(C2) - C),0.0,atol=BIGEPS)
    @test_skip norm(dwdz[:,i] - dwdz_fd) ≈ 0.0 atol=TOL
    @test_skip norm(dwdzstar[:,i] - dwdzstar_fd) ≈ 0.0 atol=TOL

    dchebt! = Plates.Chebyshev.plan_transform!(Plates._dct_data(Float64,N))

    # finite diff
    dchebt! * C
    dchebt! * Cx⁺
    dchebt! * Cy⁺
    dchebt! * CΓ⁺
    dCdx_fd = (Cx⁺ - C)/dz[i]
    dCdy_fd = (Cy⁺ - C)/dz[i]
    dCdΓ_fd = (CΓ⁺ - C)/dΓ[i]
    dCdz_fd = 0.5*(dCdx_fd - im*dCdy_fd)
    dCdzstar_fd = 0.5*(dCdx_fd + im*dCdy_fd)

    # auto diff
    cfg = ComplexGradientConfig(z -> (),blobs)
    newblobs = Vortex.seed_position(blobs,cfg)
    C2 = zeros(eltype(cfg.duals),N)
    induce_velocity!(C2,p,newblobs,0.0)
    dchebt! = Plates.Chebyshev.plan_transform!(Plates._dct_data(Dual{Nothing,Float64,2*nblob},N))
    dchebt! * C2
    dCdz, dCdzstar = dz_partials(C2,i)

    @test_skip norm(dCdz - dCdz_fd) ≈ 0.0 atol=TOL
    @test_skip norm(dCdzstar - dCdzstar_fd) ≈ 0.0 atol=TOL

    # diff wrt strength
    newblobs = Vortex.seed_strength(blobs,cfg)
    C2 = zeros(eltype(cfg.duals),N)
    induce_velocity!(C2,p,newblobs,0.0)
    dchebt! * C2
    dCdz_tmp, dCdzstar_tmp = dz_partials(C2,i)
    dCdΓ = dCdz_tmp+dCdzstar_tmp

    @test_skip norm(dCdΓ - dCdΓ_fd) ≈ 0.0 atol=TOL

    # Now with enforce_no_flow_through
    newblobs = Vortex.seed_position(blobs,cfg)
    pdual = PotentialFlow.Plate{Elements.property_type(eltype(newblobs))}(N,L,c,α)
    Plates.enforce_no_flow_through!(pdual, motion, newblobs, 0.0)
    dCdz, dCdzstar = dz_partials(pdual.C,i)

    @test_skip norm(dCdz - dCdz_fd) ≈ 0.0 atol=TOL
    @test_skip norm(dCdzstar - dCdzstar_fd) ≈ 0.0 atol=TOL

    n = rand(0:N-1)
    @test p.A[n] ≈ value(pdual.A[n]) atol=BIGEPS

    # note that we need to wrap A in complex to ensure it gets dispatched
    # to the correct extract_derivative.
    dAdz, dAdzstar = dz_partials(complex(pdual.A[n]),i)
    @test dAdz == -0.5im*(dCdz[n+1] - conj(dCdzstar[n+1]))
    @test dAdzstar == conj(dAdz)

    dΓdz, dΓdzstar = dz_partials(complex(pdual.Γ))
    @test all(dΓdz .== complex(0.0))

    # with dualized strength
    newblobs = Vortex.seed_strength(blobs,cfg)
    pdual = PotentialFlow.Plate{Elements.property_type(eltype(newblobs))}(N,L,c,α)
    Plates.enforce_no_flow_through!(pdual, motion, newblobs, 0.0)
    dCdz_tmp, dCdzstar_tmp = dz_partials(pdual.C,i)
    dCdΓ = dCdz_tmp+dCdzstar_tmp

    @test_skip norm(dCdΓ - dCdΓ_fd) ≈ 0.0 atol=TOL

    n = rand(0:N-1)
    @test p.A[n] ≈ value(pdual.A[n]) atol=BIGEPS

    dΓdz, dΓdzstar = dz_partials(complex(pdual.Γ))
    dΓdΓ = dΓdz+dΓdzstar
    @test all(dΓdΓ .== complex(-1.0))

    @test pdual.A[n] == imag(pdual.C[n+1])
end

@testset "Induced velocity with bc enforced" begin

    # Now apply the full differentiation to evaluate sensitivity of induced velocity
    # First, by finite difference

    w_fd =  induce_velocity(z,(p,blobs),0.0)
    wx⁺_fd = induce_velocity(z,(px⁺,blobsx⁺),0.0)
    wy⁺_fd = induce_velocity(z,(py⁺,blobsy⁺),0.0)
    wΓ⁺_fd = induce_velocity(z,(pΓ⁺,blobsΓ⁺),0.0)

    dwdx_fd = (wx⁺_fd - w_fd)/dz[i]
    dwdy_fd = (wy⁺_fd - w_fd)/dz[i]
    dwdz_fd = 0.5*(dwdx_fd - im*dwdy_fd)
    dwdzstar_fd = 0.5*(dwdx_fd + im*dwdy_fd)
    dwdΓ_fd = (wΓ⁺_fd - w_fd)/dΓ[i]

    # now autodiff
    function f(v)
      ptmp = PotentialFlow.Plate{Elements.property_type(eltype(v))}(N,L,c,α)
      Plates.enforce_no_flow_through!(ptmp, motion, v, 0.0)
      out = induce_velocity(z,(ptmp,v),0.0)
      return out
    end

    dwdz, dwdzstar = gradient_position(f,blobs)

    @test_skip abs(dwdz[i]-dwdz_fd) ≈ 0.0 atol=TOL
    @test_skip abs(dwdzstar[i]-dwdzstar_fd) ≈ 0.0 atol=TOL

    dwdΓ = gradient_strength(f,blobs)

    @test_skip abs(dwdΓ[i]-dwdΓ_fd) ≈ 0.0 atol=TOL
end

@testset "Self-induced velocity with bcs enforced" begin

    function self_velocity(v)
      ptmp = PotentialFlow.Plate{Elements.property_type(eltype(v))}(N,L,c,α)
      Plates.enforce_no_flow_through!(ptmp, motion, v, 0.0)
      wself = zeros(Complex{Elements.property_type(eltype(v))},length(v))
      self_induce_velocity!(wself,v, 0.0)
      induce_velocity!(wself, v, ptmp, 0.0)
      return wself
    end

    # chunk mode
    nchunk = 3
    cfg2 = PotentialFlow.Utils.ComplexGradientConfig(z -> (),blobs,PotentialFlow.Utils.Chunk{nchunk}())

    wself_fd = self_velocity(blobs)
    wselfx⁺_fd = self_velocity(blobsx⁺)
    wselfy⁺_fd = self_velocity(blobsy⁺)
    wselfΓ⁺_fd = self_velocity(blobsΓ⁺)
    dwdx_fd = (wselfx⁺_fd - wself_fd)/dz[i]
    dwdy_fd = (wselfy⁺_fd - wself_fd)/dz[i]
    dwdz_fd = 0.5*(dwdx_fd - im*dwdy_fd)
    dwdzstar_fd = 0.5*(dwdx_fd + im*dwdy_fd)
    dwdΓ_fd = (wselfΓ⁺_fd - wself_fd)/dΓ[i]

    dwdz, dwdzstar = jacobian_position(self_velocity,blobs)

    @test_skip norm(dwdz[:,i]-dwdz_fd) ≈ 0.0 atol=TOL
    @test_skip norm(dwdzstar[:,i]-dwdzstar_fd) ≈ 0.0 atol=TOL

    dwdz_chunk, dwdzstar_chunk = jacobian_position(self_velocity,blobs,cfg2)
    @test dwdz == dwdz_chunk
    @test dwdzstar == dwdzstar_chunk

    dwdΓ = jacobian_strength(self_velocity,blobs)

    @test_skip norm(dwdΓ[:,i]-dwdΓ_fd) ≈ 0.0 atol=TOL

    dwdΓ_chunk = jacobian_strength(self_velocity,blobs,cfg2)
    @test dwdΓ == dwdΓ_chunk

end

@testset "Acceleration induced on plate" begin

    targvel = fill(ċ, length(p))

    function compute_Ċ(v)
      ptmp = PotentialFlow.Plate{Elements.property_type(eltype(v))}(N,L,c,α)
      Plates.enforce_no_flow_through!(ptmp, motion, v, 0.0)
      wself = zeros(Complex{Elements.property_type(eltype(v))},length(v))
      self_induce_velocity!(wself,v, 0.0)
      induce_velocity!(wself, v, ptmp, 0.0)
      Ċ = zeros(Complex{Elements.property_type(eltype(v))},length(ptmp))
      Plates.induce_acc!(Ċ, ptmp.zs, targvel, v, wself)
    end

    Ċ_fd = compute_Ċ(blobs)
    Ċx⁺_fd = compute_Ċ(blobsx⁺)
    Ċy⁺_fd = compute_Ċ(blobsy⁺)
    ĊΓ⁺_fd = compute_Ċ(blobsΓ⁺)

    dĊdx_fd = (Ċx⁺_fd - Ċ_fd)/dz[i]
    dĊdy_fd = (Ċy⁺_fd - Ċ_fd)/dz[i]
    dĊdz_fd = 0.5*(dĊdx_fd - im*dĊdy_fd)
    dĊdzstar_fd = 0.5*(dĊdx_fd + im*dĊdy_fd)
    dĊdΓ_fd = (ĊΓ⁺_fd - Ċ_fd)/dΓ[i]

    dĊdz, dĊdzstar = jacobian_position(compute_Ċ,blobs)

    @test_skip norm(dĊdz[:,i]-dĊdz_fd) ≈ 0.0 atol=BIGTOL
    @test_skip norm(dĊdzstar[:,i]-dĊdzstar_fd) ≈ 0.0 atol=BIGTOL

    dĊdΓ = jacobian_strength(compute_Ċ,blobs)

    @test_skip norm(dĊdΓ[:,i]-dĊdΓ_fd) ≈ 0.0 atol=BIGTOL

  end

  Δt = 0.01
  lesp = 0.2
  tesp = 0.0
  z₊ = rand()
  z₋ = rand()

  plesp⁺ = deepcopy(p)
  dlesp = DELTA

  @testset "Pressure" begin

    press_fd = Plates.surface_pressure_inst(p,motion,blobs,(z₊,z₋),0.0,Δt,lesp,tesp)
    pressx⁺_fd = Plates.surface_pressure_inst(px⁺,motion,blobsx⁺,(z₊,z₋),0.0,Δt,lesp,tesp)
    pressy⁺_fd = Plates.surface_pressure_inst(py⁺,motion,blobsy⁺,(z₊,z₋),0.0,Δt,lesp,tesp)
    pressΓ⁺_fd = Plates.surface_pressure_inst(pΓ⁺,motion,blobsΓ⁺,(z₊,z₋),0.0,Δt,lesp,tesp)

    dpdx_fd = (pressx⁺_fd - press_fd)/dz[i]
    dpdy_fd = (pressy⁺_fd - press_fd)/dz[i]
    dpdz_fd = 0.5*(dpdx_fd - im*dpdy_fd)
    dpdzstar_fd = 0.5*(dpdx_fd + im*dpdy_fd)
    dpdΓ_fd = (pressΓ⁺_fd - press_fd)/dΓ[i]

    function compute_pressure(v)
      ptmp = PotentialFlow.Plate{Elements.property_type(eltype(v))}(N,L,c,α)
      return complex(Plates.surface_pressure_inst(ptmp,motion,v,(z₊,z₋),0.0,Δt,lesp,tesp))
    end

    dpdz, dpdzstar = jacobian_position(compute_pressure,blobs)

    @test safenorm(dpdz[:,i]-dpdz_fd)/safenorm(dpdz_fd) ≈ 0.0 atol=BIGTOL
    @test safenorm(dpdzstar[:,i]-dpdzstar_fd)/safenorm(dpdzstar_fd) ≈ 0.0 atol=BIGTOL

    dpdΓ = real(jacobian_strength(compute_pressure,blobs))

    @test safenorm(dpdΓ[:,i]-dpdΓ_fd)/safenorm(dpdΓ_fd) ≈ 0.0 atol=BIGTOL

    presslesp⁺_fd = Plates.surface_pressure_inst(plesp⁺,motion,blobs,(z₊,z₋),0.0,Δt,lesp+dlesp,tesp)
    dpdlesp_fd = (presslesp⁺_fd - press_fd)/dlesp

    function lesp_to_pressure(v,lesp)
      ptmp = PotentialFlow.Plate{Elements.property_type(eltype(v))}(N,L,c,α)
      press = Plates.surface_pressure_inst(ptmp,motion,v,(z₊,z₋),0.0,Δt,lesp,tesp)
      return complex(press)
    end

    dpdlesp = jacobian_param(lesp_to_pressure,(blobs,lesp))

    @test safenorm(dpdlesp-dpdlesp_fd)/safenorm(dpdlesp) ≈ 0.0 atol=BIGTOL

  end

  ## Larger number of blobs

  nblob = 20
  pos = rand(ComplexF64,nblob)
  str = rand(length(pos))

  σ = 1e-2
  blobs = Vortex.Blob.(pos,str,σ)
  z = rand(ComplexF64)

  i = rand(1:nblob)

  # For finite difference approximations
  dz = zeros(ComplexF64,length(blobs))
  dz[i] = DELTA
  dΓ = zeros(Float64,length(blobs))
  dΓ[i] = DELTA
  blobsx⁺ = Vortex.Blob.(Elements.position(blobs).+dz,Elements.circulation.(blobs),σ)
  blobsy⁺ = Vortex.Blob.(Elements.position(blobs).+im*dz,Elements.circulation.(blobs),σ)
  blobsΓ⁺ = Vortex.Blob.(Elements.position(blobs),Elements.circulation.(blobs).+dΓ,σ)

  N = 7
  L = 2.0
  c = complex(0)
  α = 0.0
  ċ = complex(0.0)
  α̇ = 0.0
  p = PotentialFlow.Plate(N,L,c,α)
  motion = PotentialFlow.RigidBodyMotion(ċ,α̇)

  Plates.enforce_no_flow_through!(p, motion, blobs, 0.0)
  px⁺ = deepcopy(p)
  Plates.enforce_no_flow_through!(px⁺, motion, blobsx⁺, 0.0)
  py⁺ = deepcopy(p)
  Plates.enforce_no_flow_through!(py⁺, motion, blobsy⁺, 0.0)
  pΓ⁺ = deepcopy(p)
  Plates.enforce_no_flow_through!(pΓ⁺, motion, blobsΓ⁺, 0.0)

  @testset "Acceleration induced on plate" begin

      targvel = fill(ċ, length(p))

      function compute_Ċ(v)
        ptmp = PotentialFlow.Plate{Elements.property_type(eltype(v))}(N,L,c,α)
        Plates.enforce_no_flow_through!(ptmp, motion, v, 0.0)
        wself = zeros(Complex{Elements.property_type(eltype(v))},length(v))
        self_induce_velocity!(wself,v, 0.0)
        induce_velocity!(wself, v, ptmp, 0.0)
        Ċ = zeros(Complex{Elements.property_type(eltype(v))},length(ptmp))
        Plates.induce_acc!(Ċ, ptmp.zs, targvel, v, wself)
      end

      Ċ_fd = compute_Ċ(blobs)
      Ċx⁺_fd = compute_Ċ(blobsx⁺)
      Ċy⁺_fd = compute_Ċ(blobsy⁺)
      ĊΓ⁺_fd = compute_Ċ(blobsΓ⁺)

      dĊdx_fd = (Ċx⁺_fd - Ċ_fd)/dz[i]
      dĊdy_fd = (Ċy⁺_fd - Ċ_fd)/dz[i]
      dĊdz_fd = 0.5*(dĊdx_fd - im*dĊdy_fd)
      dĊdzstar_fd = 0.5*(dĊdx_fd + im*dĊdy_fd)
      dĊdΓ_fd = (ĊΓ⁺_fd - Ċ_fd)/dΓ[i]

      dĊdz, dĊdzstar = jacobian_position(compute_Ċ,blobs)

      @test_skip norm(dĊdz[:,i]-dĊdz_fd) ≈ 0.0 atol=BIGTOL
      @test_skip norm(dĊdzstar[:,i]-dĊdzstar_fd) ≈ 0.0 atol=BIGTOL

      dĊdΓ = jacobian_strength(compute_Ċ,blobs)

      @test_skip norm(dĊdΓ[:,i]-dĊdΓ_fd) ≈ 0.0 atol=BIGTOL

    end

end


#!/usr/bin/env julia

p(n) = mapreduce(x -> (4 / (1 + (x * x))), +, (0.5 / n):(1 / n):1) / n

function picalc(numsteps)

  # Get number of processes.
  # np = nworkers()
  np = 1

  println("Calculating PI using:")
  println("  ", numsteps, " slices")
  println("  ", np, " worker(s)")
  start = time()

  mypi = p(numsteps)

  stop = time()

  elapsed = (stop - start)

  println("Obtained value of PI: ", mypi)
  println("Time taken: ", round(elapsed; digits=3), " seconds")

end

const numsteps = if length(ARGS) > 0
    parse(Int, ARGS[1])
else
    1_000_000_000
end

# Warm up kernel
print("  Warming up...")
warms = time()
p(10)
warmt = time() - warms
println("done. [", round(warmt; digits=3), "s]\n")

# Run the full example
picalc(numsteps)


# #
# Reproduce the figures in
# [V. Duchêne and C. Klein](http://dx.doi.org/10.3934/dcdsb.2021300)
# on the Serre-Green-Naghdi and its fully dispersive counterpart.
# #
export PlotSolitaryWaveWGN1,PlotSolitaryWaveWGN2,PlotSolitaryWaveWGN3,PlotJacobianWGN,IntegrateSolitaryWaveWGN,StabilitySolitaryWaveWGN,IntegrateWGN
using WaterWaves1D,FFTW,Plots,LinearAlgebra,ProgressMeter;
#using JLD #(uncomment if using function save)

#---- Figures 1 and 2
"""
	PlotSolitaryWaveWGN1(;kwargs)

First method to compute the WGN solitary wave.
Use GMRES-iterative method for the inversion of the Jacobian.

Arguments are all optional:
- `c` the velocity,
- `N` the number of collocation points,
- `L` the half-length of the mesh,
- `verbose` provides details on residuals if `true` (default is `false`)
- `name`: a string used to save the figure as a name.pdf.

Return `(η,u,v,mesh)`, where `mesh.x` are the collocation points.
"""
function PlotSolitaryWaveWGN1(;c=2,N=2^10,L=10*π,verbose=false,name=nothing)
	param = ( μ  = 1,
			ϵ  = 1,
        	N  = N,
            L  = L,
			c  = c
				)

	(η,u,v,mesh) = SolitaryWaveWhithamGreenNaghdi(
				param;
				method=2, α = 0,
				tol =  1e-14, max_iter=10,
				ktol =0*1e-11, gtol = 1e-16,
				iterative = true, q=1,
				verbose = verbose)

	(ηGN,uGN) = SolitaryWaveSerreGreenNaghdi(param)

	plt = plot(layout=(1,2))
	plot!(plt[1,1], mesh.x, [u uGN];
	  title=string("c=",c),
	  xlabel = "x",
	  ylabel = "u",
	  label=["WGN" "SGN"])
	plot!(plt[1,2], fftshift(mesh.k),
	  [log10.(abs.(fftshift(fft(u)))) log10.(abs.(fftshift(fft(uGN))))];
	  title="frequency",
	  label=["WGN" "SGN"])
	display(plt)
	if name != nothing savefig(string(name,".pdf")); end
	return (η,u,v,mesh)
end

#---- Figure 3
"""
	PlotSolitaryWaveWGN2(;kwargs)

Second method to compute the WGN solitary wave. Same as `PlotSolitaryWaveWGN1`, but
use non-iterative method for the inversion of the Jacobian.
To be used with higher values of the velocity (default is `c=20`).
"""
function PlotSolitaryWaveWGN2(;c=20,L=10*π,N=2^10,verbose=false,name=nothing)
	param = ( μ  = 1,
		ϵ  = 1,
      	N  = N,
        L  = L,
		c = c )

	(η,u,v,mesh) = SolitaryWaveWhithamGreenNaghdi(
				param;
				method=2, α = 1, #ici α = 1 évite des oscillations important si c = 3 ou c = 20
				tol =  1e-14, max_iter=15,
				iterative = false, q=1,
				verbose = verbose)

	(ηGN,uGN) = SolitaryWaveSerreGreenNaghdi(param)

	plt = plot(layout=(1,2))
	plot!(plt[1,1], mesh.x, [u uGN];
	  title=string("c=",c),
	  xlabel = "u",
	  ylabel = "u/c",
	  label=["WGN" "SGN"])
	plot!(plt[1,2], fftshift(mesh.k),
	  [log10.(abs.(fftshift(fft(u)))) log10.(abs.(fftshift(fft(uGN))))];
	  title="frequency",
	  label=["WGN" "SGN"])
    display(plt)
	if name != nothing savefig(string(name,".pdf")); end
	return (η,u,v,mesh)
end
#---- Figure 4
"""
	PlotSolitaryWaveWGN3(;kwargs)

Third method to compute the WGN solitary wave. Same as `PlotSolitaryWaveWGN1`, but
- use non-iterative method for the inversion of the Jacobian; and
- use a rescaled equation.
To be used with highest values of the velocity (default is `c=100`).
"""
function PlotSolitaryWaveWGN3(;c=100,L=10*π,N=2^10,verbose=false,name=nothing)
	param = ( μ  = 1,
		ϵ  = 1,
      	N  = N,
        L  = L,
		c = c )

	(η,u,v,mesh) = SolitaryWaveWhithamGreenNaghdi(
				param;
				method=3, α = 1,
				tol =  1e-10, max_iter=10,
				iterative = false, q=1,
				verbose = verbose)
	(ηGN,uGN) = SolitaryWaveSerreGreenNaghdi(param)

	plt = plot(layout=(1,2))
	plot!(plt[1,1], mesh.x, [u/c uGN/c];
	  title=string("c=",c),
	  xlabel = "x",
	  ylabel = "u/c",
	  label=["WGN" "SGN"])
	plot!(plt[1,2], fftshift(mesh.k),
	  [log10.(abs.(fftshift(fft(u/c)))) log10.(abs.(fftshift(fft(uGN/c)))) ];
	  title="frequency",
	  label=["WGN" "SGN"])
	display(plt)
	if name != nothing savefig(string(name,".pdf")); end
	return (η,u,v,mesh)
end

#------ Figure 6
"""
	PlotJacobianWGN(;kwargs)

Compute the Jacobian matrix to be inverted in SGN or WGN equation.

Arguments are all optional:
- `c` the velocity,
- `N` the number of collocation points,
- `L` the half-length of the mesh,
- `SGN`: use SGN if `true`, and WGN if `false` (default is `false`),
- `name`: a string used to save the figure as a name.pdf.

"""
function PlotJacobianWGN(;c=20,L=10*π,N=2^10,SGN=false,verbose=false,name=nothing)
	ϵ,μ,α=1,1,0

	(η,u,v,mesh) = SolitaryWaveWhithamGreenNaghdi(
				(c=c,L=L,N=N,μ=μ,ϵ=ϵ); SGN = SGN,
				method=2, α = 1, #ici α = 1 évite des oscillations important si c = 3 ou c = 20
				tol =  1e-12, max_iter=15,
				iterative = false, q=1,
				verbose = verbose)

	k,x=mesh.k,mesh.x
	if SGN == true
		F₀ = sqrt(μ)*1im * k
	else
		F₁ = tanh.(sqrt(μ)*abs.(k))./(sqrt(μ)*abs.(k))
		F₁[1] = 1
		F₀ = 1im * sqrt.(3*(1 ./F₁ .- 1)).*sign.(k)
	end
	x₀ = mesh.x[1]
	FFT = exp.(-1im*k*(x.-x₀)');
	IFFT = exp.(1im*k*(x.-x₀)')/length(x);
	M₀ = IFFT * Diagonal( F₀ )* FFT
	M(v) = Diagonal( v )

	dxu = real.(ifft(1im*k.*fft(u)))
	dxu ./= norm(dxu,2)
	hu = c ./(c .- ϵ*u)
	Fu = hu.* real.(ifft(F₀.*fft(u)))
	F2u = real.(ifft(F₀.*fft(hu.^2 .* Fu ) ))
	Du = c ./hu .+ 2*ϵ/3 * F2u ./ hu .+ ϵ^2/c * hu .* Fu.^2 .- (hu.^2)/c
	Jac = (-1/3 *M(1 ./ hu.^2 )* M₀ * M(hu.^3)* (c*M₀ .+ 3*ϵ * M( Fu ))
					.+ ϵ * M( hu .* Fu ) * M₀
					.+ M( Du ) .+ α*M( hu.^2 )*dxu*dxu' *M(1 ./ hu.^2 ) )
	Jacstar = -1/3 *M(1 ./ hu.^2 )* M₀ * M(hu.^3)* c*M₀

	plt = plot(layout=(1,2))
	surface!(plt[1,1],fftshift(k),fftshift(k)[N:-1:1],
		log10.(abs.(FFT*Jac*IFFT)),
		title = "Jacobian")
	surface!(plt[1,2],fftshift(k),fftshift(k)[N:-1:1],
		log10.(abs.(FFT*Jacstar*IFFT)),
		title = "non-diagonal part")
	display(plt)

	if name != nothing savefig(string(name,".pdf")); end
	return (Jac,Jacstar,FFT,IFFT)
end

#---- Figures 7 and 8
"""
	IntegrateSolitaryWaveWGN(;kwargs)

Integrate in time the SGN or WGN solitary wave.

Arguments are all optional:
- `c` the velocity,
- `N` the number of collocation points,
- `L` the half-length of the mesh,
- `T` the final time of integration,
- `dt` the timestep,
- `SGN`: use SGN if `true`, and WGN if `false` (default is `false`),
- `name`: a string used to save raw data and the figure.

Return `problem` of type `Problem`, containing all the information.
"""
function IntegrateSolitaryWaveWGN(;SGN=false,c=2,N=2^10,L=10*π,T=1,dt=1/2000,name=nothing)
	if name != nothing ns=floor(Int,max(1,T/dt/100)) else ns=1 end

	param = ( μ  = 1, ϵ  = 1, c = c,
				N  = N,
	            L  = L,
	            T  = T,
	            dt = dt, ns=ns)
	if SGN == true
		(ηinit,uinit,vinit,mesh) = SolitaryWaveSerreGreenNaghdi(param)
		(ηfin,ufin,vfin) = SolitaryWaveSerreGreenNaghdi(param; x₀ = c*T)
	else
		(ηinit,uinit,vinit,mesh) = SolitaryWaveWhithamGreenNaghdi(
				param; method = 3, max_iter=10,α=1,verbose=true, tol=1e-12)
	  	(ηfin,ufin,vfin) = SolitaryWaveWhithamGreenNaghdi(
				param; x₀ = c*T, method = 3, max_iter=10,α=1,verbose=true, tol=1e-12)
	end
	init     = Init(mesh,ηinit,vinit)
	model = WhithamGreenNaghdi(param;SGN=SGN, ktol=0, iterate=true, precond = true)
	problem = Problem(model, init, param)
	solve!( problem )

	(ηcomp,vcomp,ucomp)   =  model.mapfrofull(last(problem.data.U))

	E(η,u,v) = sum(η.^2 .+ (1 .+ param.ϵ*η).*u.*v)
	dE(η1,u1,v1,η2,u2,v2) = sum(η1.^2-η2.^2) + sum((1 .+ param.ϵ*η1).*u1.*v1 - (1 .+ param.ϵ*η2).*u2.*v2)

	print(string("absolute energy difference: ",dE(ηinit,uinit,vinit,ηcomp,ucomp,vcomp),"\n"))
	print(string("relative energy difference: ",dE(ηinit,uinit,vinit,ηcomp,ucomp,vcomp)/E(ηinit,uinit,vinit),"\n"))

	plt = plot(layout=(1,2))
	plot!(plt[1,1],fftshift(mesh.k),fftshift(log10.(abs.(fft(uinit))));
			title = "frequency",
			label = "initial",
			xlabel="x",
			ylabel="u")
	plot!(plt[1,1],fftshift(mesh.k),fftshift(log10.(abs.(fft(ucomp))));
			label = "final")
	plot!(plt[1,2],mesh.x,ufin-ucomp;
			title = string("error at time t=",problem.times.tfin),
			label="difference",
			xlabel="x",
			ylabel="δu")
	display(plt)

	if name != nothing
		savefig(string(name,".pdf"));
		create_animation(problem;name=string(name,"-anim.pdf"))
		plot_solution(problem)
		savefig(string(name,"-final.pdf"));
		#save(problem,name);
	end
	return problem
end

#------ Figures 9 to 15

"""
	StabilitySolitaryWaveWGN(;kwargs)

Integrate in time a pertubed SGN or WGN solitary wave.

Arguments are all optional:
- `p` (1,2, or a real) is the type of perturbation:
	1. If `p=1`, then `u` is multiplied by `λ = 0.99`
	2. If `p=2`, then `u` is multiplied by `λ = 1.01`
	3. Otherwise, one adds `p exp(-x^2)` to `u`
- `c` the velocity,
- `N` the number of collocation points,
- `L` the half-length of the mesh,
- `T` the final time of integration,
- `dt` the timestep,
- `SGN`: use SGN if `true`, and WGN if `false` (default is `false`),
- `iterate`: use GMRES if `true`, and LU decomposition otherwise (default is `true`),
- `precond` gives some choice in the preconditioner for GMRES,
- `dealias` use 2/3 dealiasing rule if `1`, (default is `0`, i.e. no dealiasing),
- `name`: a string used to save raw data and the figures.

Return `problem` of type `Problem`, containing all the information.
"""
function StabilitySolitaryWaveWGN(;p=2,c=2,N=2^10,L=10*π,T=10,dt=10/10^4,SGN=false,precond=true,iterate=true,dealias=0,name=nothing)
	if name != nothing ns=floor(Int,max(1,T/dt/100)) else ns=1 end
	if p == 1
		λ = 0.99
	elseif p == 2
		λ = 1.01
	else
		λ = p
	end

	param = ( μ = 1, ϵ = 1, c = c, λ = λ,
				N  = N,
	            L  = L,
	            T  = T,
	            dt = dt, ns=ns,
				dealias = dealias,
				SGN=SGN,
				ktol=0*1e-6,
				gtol=1e-12,
				iterate=iterate)

	if SGN == true
		(η,u,v,mesh) = SolitaryWaveSerreGreenNaghdi(param)
	else
		(η,u,v,mesh) = SolitaryWaveWhithamGreenNaghdi(
				param; method = 3, tol=1e-14, max_iter=10,α=1,verbose=true)
	end
	k = mesh.k

	if precond > 0
		precond = Diagonal( 1 .+ μ/3*(precond^2*k).^2 )
	elseif precond < 0
		precond = Diagonal( (1 .+ μ/3000*k.^8)  )
	end
	if p == 1 || p == 2
		u= λ*u
	else
		u .+= λ*exp.(-mesh.x.^2)
	end
	if SGN == true
		F₀ = sqrt(param.μ)*1im * k
	else
		F₁ = tanh.(sqrt(param.μ)*abs.(k))./(sqrt(param.μ)*abs.(k))
		F₁[1] = 1
		F₀ = 1im * sqrt.(3*(1 ./F₁ .- 1)).*sign.(k)
	end
	DxF(v) = real.(ifft(F₀ .* fft(v)))
	h = 1 .+ param.ϵ*η
	v = u - 1/3 ./h .* (DxF(h.^3 .*DxF(u)))

	init = Init(mesh,η,v)
	model = WhithamGreenNaghdi(param;SGN=param.SGN, dealias = param.dealias, ktol=param.ktol, gtol=param.gtol, iterate=param.iterate, precond = precond)
	problem = Problem(model, init, param)
	solve!( problem )

	#if name != nothing save(problem,name); end

	(ηfin,vfin,ufin)   =  model.mapfrofull(last(problem.data.U))

	E(η,u,v) = sum(η.^2 .+ (1 .+ param.ϵ*η).*u.*v)
	dE(η1,u1,v1,η2,u2,v2) = sum(η1.^2-η2.^2) + sum((1 .+ param.ϵ*η1).*u1.*v1 - (1 .+ param.ϵ*η2).*u2.*v2)

	print(string("normalized error: ",dE(η,u,v,ηfin,ufin,vfin)/E(η,u,v),"\n"))

	plt = plot(layout=(1,2))
	plot!(plt[1,1],fftshift(mesh.k),fftshift(log10.(abs.(fft(η))));
			title = "frequency",
			label = "initial")
	plot!(plt[1,1],fftshift(mesh.k),fftshift(log10.(abs.(fft(ηfin))));
			label = "final")
	plot!(plt[1,2],mesh.x,[η ηfin];
			title = string("at time t=",problem.times.tfin),
			label=["zeta initial" "zeta final"])
	plot!(plt[1,2],mesh.x,[u ufin];
			label=["u initial" "u final"])
	display(plt)
	if name != nothing savefig(string(name,"-final.pdf")); end

	ts = problem.times.ts
	x = mesh.x
	k = mesh.k
	X=interpolate(mesh,real.(ifft(problem.data.U[1][:,1])))[1].x
	zs=zeros(length(X),length(ts));
	for i in 1:length(ts)
		zs[:,i].=interpolate(mesh,real.(ifft(problem.data.U[i][:,1])))[2]
	end

	plt = plot()
	scatter!(plt,ts,maximum(abs.(zs),dims=1)',
			title="maximum of surface deformation",
			label="",
			xlabel="time t")
	display(plt)

	if name != nothing
		savefig(string(name,"-znorm.pdf"));

		us=zeros(length(X),length(ts));
		@showprogress 1 "Computing v..." for i in 1:length(ts)
			us[:,i].=interpolate(mesh,model.mapfro(problem.data.U[i])[2])[2]
		end
		plt = plot()
		scatter!(plt,ts,maximum(abs.(us),dims=1)',
				title="maximum of velocity",
				label="",
				xlabel="time t")
		display(plt)
		savefig(string(name,"-vnorm.pdf"));

		plt=plot()
		my_cg = cgrad([:blue,:green])
		#surface!(plt,X,ts,zs',view_angle=(20,30), color = my_cg)
		display(plt)
		savefig(string(name,"-evol.pdf"));

		create_animation(problem;name=string(name,"-anim"))
	end
	return problem
end

#---- Figures 16 to end
"""
	IntegrateWGN(scenario;kwargs)

Integrate in time SGN or WGN with an initial data depending on a given `scenario`
- If `scenario = 1`, produces a dispersive shock wave for a unidirectional wave constructed from the Saint-Venant model.
- If `scenario = 2`, produces a dispersive shock wave for a unidirectional wave constructed from the Camassa-Holm model.
- If `scenario = 3`, the initial data challenges the non-cavitation assumption.

Other arguments are optional:
- `δ` is either
	- the square root of the shallowness parameter (the nonlinear parameter `ϵ=1`) if `scenario=1` or `2`;
	- the minimal depth if `scenario=3`
- `N` the number of collocation points,
- `L` the half-length of the mesh,
- `T` the final time of integration,
- `dt` the timestep,
- `SGN`: use SGN if `true`, and WGN if `false` (default is `false`),
- `iterate`: use GMRES if `true`, and LU decomposition otherwise (default is `true`),
- `precond` gives some choice in the preconditioner for GMRES,
- `dealias`: dealiasing with Orlicz rule `1-dealias/(dealias+2)` (default is `0`, i.e. no dealiasing);
- `name`: a string used to save raw data and the figures.

Return `(problem,plt)` where `problem` contains all the information and `plt` a plot of the final time solution.
"""
function IntegrateWGN(scenario;δ=0.1,N=2^11,L=3*π,x₀=-3,T= 5,dt = 5/10^4,SGN=false,dealias=0,iterate=true,precond=true,name=nothing)

	if name != nothing ns=floor(Int,max(1,T/dt/100)) else ns=1 end
	if scenario == 1 || scenario == 2
		μ  = δ^2
	else
		μ = 1
	end
	param = ( μ  = μ, ϵ  = 1,
				N  = N, L  = L,
	            T  = T, dt = dt,
				ns=ns )

	mesh=Mesh(param)
	k=mesh.k
	if precond > 0
		precond = Diagonal( 1 .+ μ/3*(precond^2*k).^2 )
	elseif precond < 0
		precond = Diagonal( (1 .+ μ/3000*k.^8)  )
	end

	if scenario == 1
		η= exp.(-(mesh.x .-x₀).^2)
		u= 2*sqrt.(1 .+ param.ϵ*η) .-2
	elseif scenario == 2
		function krasny!(v)
			v[abs.(v).<1e-14].=0
			return v
		end
		Dx(v) = real.(ifft( 1im*mesh.k.* krasny!(fft(v))))
		Dx2(v) = Dx(Dx(v))
		ϵ = param.ϵ;μ=param.μ;
		w = - (mesh.x.-x₀).* exp.(-(mesh.x.-x₀).^2)
		u = w .+ μ/12 *Dx2(w) .+ μ*ϵ/6* w.*Dx2(w)
		η = u .+ ϵ/4*u.^2 .- μ/6* Dx2( u+3*ϵ/4* u.^2) .- μ*ϵ/6 * u .* Dx2(u) .- 5*μ*ϵ/48 * Dx(u).^2
	elseif scenario == 3
		η = -(1-δ)*exp.(-mesh.x.^2)
		u = -mesh.x.*exp.(-mesh.x.^2)
	else
		error("the first argument must be 1, 2 or 3")
	end
	if SGN == true
		F₀ = sqrt(param.μ)*1im*mesh.k
	else
		F₁ 	= tanh.(sqrt(param.μ)*abs.(mesh.k))./(sqrt(param.μ)*abs.(mesh.k))
		F₁[1]= 1                 # Differentiation
		F₀   = 1im * sqrt.(3*(1 ./F₁ .- 1)).*sign.(mesh.k)
	end
	DxF(v) = real.(ifft(F₀ .* fft(v)))
	h=1 .+ param.ϵ*η
	v= u - 1/3 ./h .* (DxF(h.^3 .*DxF(u)))

	init     = Init(mesh,η,v)
	model = WhithamGreenNaghdi(param;SGN=SGN, ktol=0, gtol=1e-12, iterate=iterate, precond = precond, dealias = dealias)
	problem = Problem(model, init, param)
	solve!( problem )

	(ηfin,vfin,ufin)   =  model.mapfrofull(last(problem.data.U))
	E(η,u,v) = sum(η.^2 .+ (1 .+ param.ϵ*η).*u.*v)
	dE(η1,u1,v1,η2,u2,v2) = sum(η1.^2-η2.^2) + sum((1 .+ param.ϵ*η1).*u1.*v1 - (1 .+ param.ϵ*η2).*u2.*v2)

	print(string("normalized error: ",dE(η,u,v,ηfin,ufin,vfin)/E(η,u,v),"\n"))

	fftηfin=last(problem.data.U)[:,1]

	plt = plot(layout=(1,2))
	plot!(plt[1,1],fftshift(mesh.k),fftshift(log10.(abs.(fftηfin)));
			title = "frequency",label="")
	plot!(plt[1,2],mesh.x,real.(ifft(fftηfin));
			title = string("surface deformation at time t=",problem.times.tfin),
			label="")
	display(plt)

	if name != nothing
		savefig(string(name,".pdf"));
		#save(problem,name);
		ts = problem.times.ts
		x = mesh.x
		k = mesh.k
		zs=zeros(param.N,length(ts));
		@showprogress 1 for i in 1:length(ts)
			zs[:,i].=real.(ifft(problem.data.U[i][:,1]))
		end
		plt3=plot()
		my_cg = cgrad([:blue,:green])
		surface!(plt3,x,ts,zs',view_angle=(20,30), color = my_cg)
		display(plt3)
		savefig(string(name,"-evol.pdf"));
		create_animation(problem;ylims=false,name=string(name,"-anim"))

	end
	display(plt)
	return problem,plt
end
nothing


function pmatmul(A::SparseMatrixCSC, v::Vector{Float64}, slc::Vector{UnitRange{Int64}})

    x = @distributed (+) for s in slc
        A[:,s] * v[s]
    end

    return x
end

function pdot(u::Vector{Float64}, v::Vector{Float64}, slc::Vector{UnitRange{Int64}})

    x = @distributed (+) for s in slc
        dot(u[s], v[s])
    end

    return x
end

function pnorm(u::Vector{Float64}, slc::Vector{UnitRange{Int64}})

    return sqrt(pdot(u, u, slc))
end

function pgmres(A::SparseMatrixCSC, x::Vector{Float64}, r::Vector{Float64}, m::Int64, slc::Vector{UnitRange{Int64}})

    n = length(x0)
    Q = zeros(Float64, (n,m+1))
    H = zeros(Float64, (m+2,m+1))
    sn = zeros(Float64, m)
    cn = zeros(Float64, m)
    e1 = zeros(Float64, m+1)
    e1[1] = 1.0

    rnorm = pnorm(r, slc)
    beta = rnorm* e1
    Q[:,1] = r / rnorm

    # begin iteration
    nk = 0
    for k=1:m
        nk += 1

        # do arnoldi iteration
        q = pmatmul(A, Q[:,k], slc)
        for i=1:k
            H[i,k] = pdot(q, Q[:,i], slc)
            q = q - H[i,k] * Q[:,i]
        end
        qnorm = pnorm(q, slc)
        H[k+1,k] = qnorm
        Q[:,k+1] = q / qnorm

        # apply givens rotation
        for i=1:k-1
            temp = cn[i] * H[i,k] + sn[i] * H[i+1,k]
            H[i+1,k] = -sn[i] * H[i,k] + cn[i] * H[i+1,k]
            H[i,k] = temp
        end

        # update sn and cn
        t = sqrt(H[k,k]^2 + H[k+1,k]^2)
        cn[k] = H[k,k] / t
        sn[k] = H[k+1,k] / t

        # eliminate H[k+1, k]
        H[k,k] = cn[k] * H[k,k] + sn[k] * H[k+1,k]
        H[k+1,k] = 0.0

        # update residual
        beta[k+1] = -sn[k]*beta[k]
        beta[k] = cn[k]*beta[k]
    end

    # compute correction to x
    y = H[1:nk,1:nk] \ beta[1:nk]
    return x + Q[:,1:nk] * y
end


function get_nnd(s)
    nnd = []
    clusters = unique(s)
    for c in clusters
        g_c = findall(x-> x==c, s)
        forwd = diff(g_c)
        backwd = map(x->abs(x),diff(reverse(g_c)))
        nnd_s = map((x,y)-> minimum([x,y]),forwd,backwd)
        push!(nnd,nnd_s)
    end
    return vcat(nnd...)
end
function get_H(s::Array{Int64,1})
    T = Dict{Any,Float64}()
    n = length(s)
    for i = 1:length(s)
        T[s[i]] = get(T, s[i],0) + 1/n
    end
    p = collect(values(T))
    return mapreduce(x-> -x*log(x),+,p)
end
function get_freq(s)
    T = Dict{Int64,Int64}()
    n = length(s)
    for i = 1:n
        T[s[i]] = get(T, s[i],0) + 1
    end
    return T
end
function get_nnd_distro(s,dmax::Int64)
    P = zeros(dmax)
    n = length(s)
    for i in 1:dmax
        ndata = findall(x-> x==i, s)
        if !isempty(ndata)
            P[i] = (length(ndata) + 1/dmax)/(n + 1/dmax)
        else
            P[i] = (1/dmax) / (n+1/dmax)
        end
    end
    return P
end
KS_d(x,y) = maximum(abs.(cumsum(x) - cumsum(y)))
function get_gene_cor(gene_exp)
    pcor = Float64[]
    g_di = Float64[]
    lg = size(gene_exp)[1]
    for i in 1:lg-1
        for j in i+1:lg
            #try
            push!(pcor, cor(gene_exp[i,2:end-1], gene_exp[j,2:end-1]))
            push!(g_di, abs(gene_exp[i,end]- gene_exp[j,end]))
        end
    end
    g_ord = sortperm(g_di)
    return [g_di[g_ord] pcor[g_ord]]
end
function get_gene_null(gene_exp)
    pcor = Float64[]
    g_di = Float64[]
    lg = size(gene_exp)[1]
    for i in 1:lg-1
        for j in i+1:lg
            #try
            push!(pcor, cor(shuffle(gene_exp[i,2:end-1]), gene_exp[j,2:end-1]))
            push!(g_di, abs(gene_exp[i,end]- gene_exp[j,end]))
        end
    end
    g_ord = sortperm(g_di)
    return [g_di[g_ord] pcor[g_ord]]
end
function get_smooth_cor(gene_cor; num_w=1000, q_1=0.25, q_2=0.75)
    w_s = round((gene_cor[:,1][end] - gene_cor[:,1][1] )/ num_w)
    med_dis = Float64[]
    med_cor = Float64[]
    q_cor = []
    for nw in 0:num_w-1
        ini = nw*w_s
        fin = nw*w_s + w_s
        ix_dw = findall(x-> ini<x<fin, gene_cor[:,1])
        if isempty(ix_dw); continue; end
        push!(med_dis, median(gene_cor[ix_dw,1]))
        m_cor = median(gene_cor[ix_dw,2])
        push!(med_cor, m_cor)
        qs = quantile(gene_cor[ix_dw,2],[q_1,q_2])
        push!(q_cor,[m_cor-qs[1] qs[2]-m_cor])
    end
    q_cor =vcat(q_cor...)
    return med_dis, med_cor, q_cor
end
moving_average(vs,n) = [sum(@view vs[i:(i+n-1)])/n for i in 1:(length(vs)-(n-1))]
function get_plot_cor(gene_cor,gene_null)
    dis_cor, val_cor, q_cor = get_smooth_cor(gene_cor)
    dis_null, val_null, q_null = get_smooth_cor(gene_null)
    mv_w = 11
    dis_cor_mv = moving_average(dis_cor,mv_w)
    val_cor_mv = moving_average(val_cor,mv_w)
    q_cor_mv = [moving_average(q_cor[:,1],mv_w) moving_average(q_cor[:,2],mv_w)]
    dis_null_mv = moving_average(dis_null,mv_w)
    val_null_mv = moving_average(val_null,mv_w)
    q_null_mv = [moving_average(q_null[:,1],mv_w) moving_average(q_null[:,2],mv_w)]
    pt = plot(leg=false, #title="$gtitle",
        size=(1100,1000)
    )
    scatter!(pt,gene_cor[:,1]/10^8,gene_cor[:,2],
        alpha=0.08,
        ylims=(-0.3,1.0),
        color="#36983B"
    )
    scatter!(pt,gene_null[:,1]/10^8,gene_null[:,2],
        alpha=0.08,
        color="#EA5108"
    )
    plot!(pt,dis_cor_mv/10^8, val_cor_mv,ribbon=q_cor_mv, color="#36983B",linewidth=2)
    plot!(pt,dis_null_mv/10^8, val_null_mv,ribbon=q_null_mv, color="#EA5108",linewidth=2)
    return pt
end
function get_genes_ch(gene_tags, genes_data)
    g_ch = []
    g_st = []
    ix_x = []
    genes= []
    for g in gene_tags
        ix = findfirst(x-> x==g,genes_data[:,3])
        if isnothing(ix)
            #println("gene $g not found")
            continue
        end
        push!(ix_x, ix)
        push!(genes, g)
        push!(g_ch, genes_data[ix,5])
        push!(g_st, genes_data[ix,4])
    end
    return [genes g_st g_ch]
end


# ====================================================================
# Component structure
mutable struct OffsetSlope_1D <: AbstractComponent
    offset::Parameter
    x0::Parameter
    slope::Parameter

    function OffsetSlope_1D(offset::Number, x0::Number, slope::Number)
        out = new(Parameter(offset), Parameter(x0), Parameter(slope))
        out.x0.fixed = true
        return out
    end
end

mutable struct OffsetSlope_2D <: AbstractComponent
    offset::Parameter
    x0::Parameter
    y0::Parameter
    slopeX::Parameter
    slopeY::Parameter

    function OffsetSlope_2D(offset::Number, x0::Number, y0::Number, slopeX::Number, slopeY::Number)
        out = new(Parameter(offset), Parameter(x0), Parameter(y0), Parameter(slopeX), Parameter(slopeY))
        out.x0.fixed = true
        out.y0.fixed = true
        return out
    end
end

OffsetSlope(offset, x0, slope) = OffsetSlope_1D(offset, x0, slope)
OffsetSlope(offset, x0, y0, slopeX, slopeY) = OffsetSlope_2D(offset, x0, y0, slopeX, slopeY)

# ====================================================================
# Evaluate component
function evaluate!(buffer, comp::OffsetSlope_1D, x::AbstractDomain{1},
                   offset, x0, slope)
    @. (buffer = slope * (x - x0) + offset)
end

function evaluate!(buffer, comp::OffsetSlope_2D, domain::AbstractDomain{2},
                   offset, x0, y0, slopeX, slopeY)
    x = domain[1]
    y = domain[2]
    @. (buffer =
        slopeX * (x - x0) +
        slopeY * (y - y0) +
        offset)
end


using EconFunctions
using ContinuousTransformations
using Base.Test
using Compat: ∘
import ForwardDiff.derivative

@testset "quadrature general" begin
    @test_throws ArgumentError Quadrature(1..2, collect(1:3), collect(1:4))
end

@testset "normal quadrature" begin
    q = quadrature_standard_normal(10)
    @test sprint(show, q) == "Quadrature of 10 nodes on " * sprint(show, ℝ)
    @test length(q) == 10
    @test domain(q) == ℝ
    @test q(one) ≈ 1
    @test isapprox(q(identity), 0, atol=eps()) # FIXME replace with new test syntax
    @test q(x->x^2) ≈ 1
    μ = 2.0
    σ = 3.0
    a = Affine(σ, μ)
    @test q(a) ≈ μ
    @test q((x->x^2) ∘ a) ≈ μ^2 + σ^2
end

@testset "uniform quadrature" begin
    q = quadrature_standard_uniform(9)
    @test sprint(show, q) == "Quadrature of 9 nodes on " * sprint(show, 𝕀)
    @test length(q) == 9
    @test domain(q) == 𝕀
    @test q(one) ≈ 1
    @test q(identity) ≈ 0.5
    @test q(x->(x-0.5)^2) ≈ 1/12
end

@testset "CRRA utility" begin
    u = CrraUtility(2.0)
    @test u(0) == -Inf
    @test u(1) < u(2)
    @test derivative(u, 3) ≈ abs(crra_u′(2.0)(3))
    c = 2.0
    σ = 3.1
    @test derivative(σ -> CrraUtility(σ)(c), σ) ≈
        derivative(σ -> (c^(1-σ)-1)/(1-σ), σ)
end


##############################################################################
##
## Graveyard of Alternative Functions
##
## The proper setting of epsilon depends on whether we're doing central
##  differencing or forward differencing. See, for example, section 5.7 of
##  Numerical Recipes.
##
##############################################################################

##############################################################################
##
## Derivative of f: C -> R
##
##############################################################################

macro forwardrule(x, e)
    x, e = esc(x), esc(e)
    quote
        $e = sqrt(eps(eltype($x))) * max(one(eltype($x)), abs($x))
    end
end

macro centralrule(x, e)
    x, e = esc(x), esc(e)
    quote
        $e = cbrt(eps(eltype($x))) * max(one(eltype($x)), abs($x))
    end
end

macro hessianrule(x, e)
    x, e = esc(x), esc(e)
    quote
        $e = eps(eltype($x))^(1/4) * max(one(eltype($x)), abs($x))
    end
end

macro complexrule(x, e)
    x, e = esc(x), esc(e)
    quote
        $e = eps($x)
    end
end

function finite_difference{T <: Number}(f::Function,
                                        x::T,
                                        dtype::Symbol = :central)
    if dtype == :forward
        @forwardrule x epsilon
        xplusdx = x + epsilon
        return (f(xplusdx) - f(x)) / epsilon
    elseif dtype == :central
        @centralrule x epsilon
        xplusdx, xminusdx = x + epsilon, x - epsilon
        return (f(xplusdx) - f(xminusdx)) / (epsilon + epsilon)
    elseif dtype == :complex
        @complexrule x epsilon
        xplusdx = x + epsilon * im
        return imag(f(xplusdx)) / epsilon
    else
        error("dtype must be :forward, :central or :complex")
    end
end

##############################################################################
##
## Complex Step Finite Differentiation Tools
##
## Martins, Sturdza, and Alonso (2003) suggest the only non-analytic
##  fuction of which complex step finite difference approximation
##  will fail and finite difference will not is abs().
##  They suggest redefining as follows for z = x + im*y
##
##  if x < 0
##      -x - im * y
##  else
##      x + im * y
##
## This is provided below as complex_differentiable_abs (renaming encouraged!)
##
## Also, if your fuctions has control flow using < or >, you must compare
## real(z) for your control flow.
##
##############################################################################

function complex_differentiable_abs{T <: Complex}(z::T)
    if real(z) < 0
        return -real(z) - im * imag(z)
    else
        return real(z) + im * imag(z)
    end
end

##############################################################################
##
## Gradient of f: R^n -> R
##
##############################################################################

function finite_difference!{S <: Number, T <: Number}(f::Function,
                                                      x::Vector{S},
                                                      g::Vector{T},
                                                      dtype::Symbol)
    # What is the dimension of x?
    n = length(x)

    # Iterate over each dimension of the gradient separately.
    # Use xplusdx to store x + dx instead of creating a new vector on each pass.
    # Use xminusdx to store x - dx instead of creating a new vector on each pass.
    if dtype == :forward
        # Establish a baseline value of f(x).
        f_x = f(x)
        for i = 1:n
            @forwardrule x[i] epsilon
            oldx = x[i]
            x[i] = oldx + epsilon
            f_xplusdx = f(x)
            x[i] = oldx
            g[i] = (f_xplusdx - f_x) / epsilon
        end
    elseif dtype == :central
        for i = 1:n
            @centralrule x[i] epsilon
            oldx = x[i]
            x[i] = oldx + epsilon
            f_xplusdx = f(x)
            x[i] = oldx - epsilon
            f_xminusdx = f(x)
            x[i] = oldx
            g[i] = (f_xplusdx - f_xminusdx) / (epsilon + epsilon)
        end
    else
        error("dtype must be :forward or :central")
    end

    return
end
function finite_difference{T <: Number}(f::Function,
                                        x::Vector{T},
                                        dtype::Symbol = :central)
    # Allocate memory for gradient
    g = Array(Float64, length(x))

    # Mutate allocated gradient
    finite_difference!(f, float(x), g, dtype)

    # Return mutated gradient
    return g
end

##############################################################################
##
## Jacobian derivative of f: R^n -> R^m
##
##############################################################################

function finite_difference_jacobian!{R <: Number,
                                     S <: Number,
                                     T <: Number}(f::Function,
                                                  x::Vector{R},
                                                  f_x::Vector{S},
                                                  J::Array{T},
                                                  dtype::Symbol = :central)
    # What is the dimension of x?
    m, n = size(J)

    # Iterate over each dimension of the gradient separately.
    if dtype == :forward
        for i = 1:n
            @forwardrule x[i] epsilon
            oldx = x[i]
            x[i] = oldx + epsilon
            f_xplusdx = f(x)
            x[i] = oldx
            J[:, i] = (f_xplusdx - f_x) / epsilon
        end
    elseif dtype == :central
        for i = 1:n
            @centralrule x[i] epsilon
            oldx = x[i]
            x[i] = oldx + epsilon
            f_xplusdx = f(x)
            x[i] = oldx - epsilon
            f_xminusdx = f(x)
            x[i] = oldx
            J[:, i] = (f_xplusdx - f_xminusdx) / (epsilon + epsilon)
        end
    else
        error("dtype must :forward or :central")
    end

    return
end
function finite_difference_jacobian{T <: Number}(f::Function,
                                                 x::Vector{T},
                                                 dtype::Symbol = :central)
    # Establish a baseline for f_x
    f_x = f(x)

    # Allocate space for the Jacobian matrix
    J = zeros(length(f_x), length(x))

    # Compute Jacobian inside allocated matrix
    finite_difference_jacobian!(f, x, f_x, J, dtype)

    # Return Jacobian
    return J
end

##############################################################################
##
## Second derivative of f: R -> R
##
##############################################################################

function finite_difference_hessian{T <: Number}(f::Function,
                                                x::T)
    @hessianrule x epsilon
    (f(x + epsilon) - 2*f(x) + f(x - epsilon))/epsilon^2
end
function finite_difference_hessian(f::Function,
                                   g::Function,
                                   x::Number,
                                   dtype::Symbol = :central)
    finite_difference(g, x, dtype)
end

##############################################################################
##
## Hessian of f: R^n -> R
##
##############################################################################

function finite_difference_hessian!{S <: Number,
                                    T <: Number}(f::Function,
                                                 x::Vector{S},
                                                 H::Array{T})
    # What is the dimension of x?
    n = length(x)

    epsilon = NaN
    # TODO: Remove all these copies
    xpp, xpm, xmp, xmm = copy(x), copy(x), copy(x), copy(x)
    fx = f(x)
    for i = 1:n
        xi = x[i]
        @hessianrule x[i] epsilon
        xpp[i], xmm[i] = xi + epsilon, xi - epsilon
        H[i, i] = (f(xpp) - 2*fx + f(xmm)) / epsilon^2
        @centralrule x[i] epsiloni
        xp = xi + epsiloni
        xm = xi - epsiloni
        xpp[i], xpm[i], xmp[i], xmm[i] = xp, xp, xm, xm
        for j = i+1:n
            xj = x[j]
            @centralrule x[j] epsilonj
            xp = xj + epsilonj
            xm = xj - epsilonj
            xpp[j], xpm[j], xmp[j], xmm[j] = xp, xm, xp, xm
            H[i, j] = (f(xpp) - f(xpm) - f(xmp) + f(xmm))/(4*epsiloni*epsilonj)
            xpp[j], xpm[j], xmp[j], xmm[j] = xj, xj, xj, xj
        end
        xpp[i], xpm[i], xmp[i], xmm[i] = xi, xi, xi, xi
    end
    Base.LinAlg.copytri!(H,'U')
end
function finite_difference_hessian{T <: Number}(f::Function,
                                                x::Vector{T})
    # What is the dimension of x?
    n = length(x)

    # Allocate an empty Hessian
    H = Array(Float64, n, n)

    # Mutate the allocated Hessian
    finite_difference_hessian!(f, x, H)

    # Return the Hessian
    return H
end
function finite_difference_hessian{T <: Number}(f::Function,
                                                g::Function,
                                                x::Vector{T},
                                                dtype::Symbol = :central)
    finite_difference_jacobian(g, x, dtype)
end

##############################################################################
##
## Taylor Series based estimates of first and second derivatives
##
## TODO: Fill in multivariate equivalents
##
##############################################################################

# Higher precise finite difference method based on Taylor series approximation.
# h is the stepsize
function taylor_finite_difference(f::Function,
                                  x::Real,
                                  dtype::Symbol = :central,
                                  h::Real = 10e-4)
    if dtype == :forward
        f_x = f(x)
        d = 2^3 * (2^2 * (f(x + h) - f_x) - (f(x + 2 * h) - f_x))
        d += - (2^2 * (f(x + 2 * h) - f_x) - (f(x + 4 * h) - f_x))
        d /= 3 * 2^2 * h
    elseif dtype == :central
        d = 4^5 * (2^3 * (f(x + h) - f(x - h)) - (f(x + 2 * h) - f(x - 2 * h)))
        d -= 2^3 * (f(x + 4 * h) - f(x - 4 * h)) - (f(x + 8 * h) - f(x - 8 * h))
        d /= (4^5 * (2^4 - 2^2) - (2^6 - 2^4)) * h
    else
        error("dtype must be :forward or :central")
    end
    return d
end

function taylor_finite_difference_hessian(f::Function,
                                          x::Real,
                                          h::Real)
    f_x = f(x)
    d = 4^6 * (2^4 * (f(x + h) + f(x - h) - 2 * f_x) - (f(x + 2 * h) + f(x - 2 * h) - 2 * f_x))
    d += - (2^4 * (f(x + 4 * h) + f(x - 4 * h) - 2 * f_x) - (f(x + 8 * h) + f(x - 8 * h) - 2 * f_x))
    return d / (3 * 2^6 * (2^8 - 1) * h^2)
end

##############################################################################
##
## TODO: Implement directional_derivative()
##
##############################################################################

# The function "dirderivative" calculates directional derivatives in the direction v.
# The function supplied must have the form Vector{Float64} -> Float64
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, h::Float64, twoside::Bool)
#     derivative(t::Float64 -> f(x0 + v*t) / norm(v), 0.0, h, twoside)
# end
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, h::Float64)
#     dirderivative(f, v, x0, h, true)
# end
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, )
#     derivative(f, v, x0, 0.0001)
# end
# function dirderivative(f::Function, v::Vector{Float64})
#     x -> dirderivative(f, v, x)
# end


"""
A linear Gaussian CPD, always returns a Normal

	Assumes that target and all parents can be converted to Float64 (ie, are numeric)

	P(x|parents(x)) = Normal(μ=a×parents(x) + b, σ)
"""
mutable struct LinearGaussianCPD <: CPD{Normal}
    target::NodeName
    parents::NodeNames
    a::Vector{Float64}
    b::Float64
    σ::Float64
end
LinearGaussianCPD(target::NodeName, μ::Float64, σ::Float64) = LinearGaussianCPD(target, NodeName[], Float64[], μ, σ)

name(cpd::LinearGaussianCPD) = cpd.target
parents(cpd::LinearGaussianCPD) = cpd.parents
nparams(cpd::LinearGaussianCPD) = length(cpd.a) + 2
function (cpd::LinearGaussianCPD)(a::Assignment)

    # compute A⋅v + b
    μ = cpd.b
    for (i, p) in enumerate(cpd.parents)
        μ += a[p]*cpd.a[i]
    end

    Normal(μ, cpd.σ)
end
(cpd::LinearGaussianCPD)() = (cpd)(Assignment()) # cpd()
(cpd::LinearGaussianCPD)(pair::Pair{NodeName}...) = (cpd)(Assignment(pair)) # cpd(:A=>1)

function Distributions.fit(::Type{LinearGaussianCPD},
    data::DataFrame,
    target::NodeName;
    min_stdev::Float64=0.0, # an optional minimum on the standard deviation
    )

    # no parents

    arr = data[!,target]
    eltype(arr) <: Real || error("fit LinearGaussianCPD requrires target to be numeric")

    μ = convert(Float64, mean(arr))
    σ = convert(Float64, stdm(arr, μ))
    σ = max(σ, min_stdev)

    LinearGaussianCPD(target, NodeName[], Float64[], μ, σ)
end
function Distributions.fit(::Type{LinearGaussianCPD},
    data::DataFrame,
    target::NodeName,
    parents::NodeNames;
    min_stdev::Float64=0.0, # an optional minimum on the standard deviation
    )

    if isempty(parents)
        return fit(LinearGaussianCPD, data, target, min_stdev=min_stdev)
    end

    # ---------------------
    # pull parental dataset
    # 1st row is all of the data for the 1st parent
    # 2nd row is all of the data for the 2nd parent, etc.

    nparents = length(parents)
    X = Array{Float64}(undef, nrow(data), nparents+1)
    for (i,p) in enumerate(parents)
        arr = data[!,p]
    	for j in 1 : nrow(data)
            X[j,i] = convert(Float64, arr[j])
    	end
    end
    X[:,end] .= 1.0

    y = convert(Vector{Float64}, data[!,target])

    # --------------------
    # solve the regression problem
    #   β = (XᵀX)⁻¹Xᵀy
    #
    #     X is the [nsamples × nparents+1] data matrix
    #     where the last column is 1.0
    #
    #     y is the [nsamples] vector of target values
    #
    # NOTE: this will fail if X is not full rank

    β = (X'*X)\(X'*y)

    a = β[1:nparents]
    b = β[end]
    σ = max(std(y), min_stdev)

    LinearGaussianCPD(target, parents, a, b, σ)
end


using LinearAlgebra, SparseArrays

function simple_matrix(T::Type, nrow::Int, ncol::Int; symmetric::Bool = false)
  symmetric && @assert nrow == ncol
  if nrow == ncol == 1
    return ones(T, nrow, ncol)
  end
  U, _ = qr(rand(T, nrow, nrow))
  local V
  if symmetric
    V = U
  else
    V, _ = qr(rand(T, ncol, ncol))
  end
  S = T[(1 + (i - 1) / (nrow - 1)) * (i == j) for i = 1:nrow, j = 1:ncol]
  return U * S * V'
end

function simple_sparse_matrix(T::Type, nrow::Int, ncol::Int; symmetric::Bool = false)
  m = min(nrow, ncol)
  p = floor(Int, sqrt(m))
  q = div(m, p)
  M = kron(Matrix(I, q, q), simple_matrix(T, p, p, symmetric = symmetric))
  pq = p * q
  R = if m > pq
    simple_matrix(T, nrow - pq, ncol - pq, symmetric = symmetric)
  else
    spzeros(nrow - pq, ncol - pq)
  end
  A = [sparse(M) spzeros(pq, ncol - pq); spzeros(nrow - pq, pq) R]
end

simple_vector(T::Type, n::Int) = T[-(-one(T))^i for i = 1:n]


"""
    Normal(μ,σ)

The *Normal distribution* with mean `μ` and standard deviation `σ≥0` has probability density function

```math
f(x; \\mu, \\sigma) = \\frac{1}{\\sqrt{2 \\pi \\sigma^2}}
\\exp \\left( - \\frac{(x - \\mu)^2}{2 \\sigma^2} \\right)
```

Note that if `σ == 0`, then the distribution is a point mass concentrated at `μ`.
Though not technically a continuous distribution, it is allowed so as to account for cases where `σ` may have underflowed,
and the functions are defined by taking the pointwise limit as ``σ → 0``.

```julia
Normal()          # standard Normal distribution with zero mean and unit variance
Normal(mu)        # Normal distribution with mean mu and unit variance
Normal(mu, sig)   # Normal distribution with mean mu and variance sig^2

params(d)         # Get the parameters, i.e. (mu, sig)
mean(d)           # Get the mean, i.e. mu
std(d)            # Get the standard deviation, i.e. sig
```

External links

* [Normal distribution on Wikipedia](http://en.wikipedia.org/wiki/Normal_distribution)

"""
struct Normal{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    σ::T
    Normal{T}(µ::T, σ::T) where {T<:Real} = new{T}(µ, σ)
end

function Normal(μ::T, σ::T; check_args=true) where {T <: Real}
    check_args && @check_args(Normal, σ >= zero(σ))
    return Normal{T}(μ, σ)
end

#### Outer constructors
Normal(μ::Real, σ::Real) = Normal(promote(μ, σ)...)
Normal(μ::Integer, σ::Integer) = Normal(float(μ), float(σ))
Normal(μ::T) where {T <: Real} = Normal(μ, one(T))
Normal() = Normal(0.0, 1.0, check_args=false)

const Gaussian = Normal

# #### Conversions
convert(::Type{Normal{T}}, μ::S, σ::S) where {T <: Real, S <: Real} = Normal(T(μ), T(σ))
convert(::Type{Normal{T}}, d::Normal{S}) where {T <: Real, S <: Real} = Normal(T(d.μ), T(d.σ), check_args=false)

@distr_support Normal -Inf Inf

#### Parameters

params(d::Normal) = (d.μ, d.σ)
@inline partype(d::Normal{T}) where {T<:Real} = T

location(d::Normal) = d.μ
scale(d::Normal) = d.σ

eltype(::Normal{T}) where {T} = T

#### Statistics

mean(d::Normal) = d.μ
median(d::Normal) = d.μ
mode(d::Normal) = d.μ

var(d::Normal) = abs2(d.σ)
std(d::Normal) = d.σ
skewness(d::Normal{T}) where {T<:Real} = zero(T)
kurtosis(d::Normal{T}) where {T<:Real} = zero(T)

entropy(d::Normal) = (log2π + 1)/2 + log(d.σ)

#### Evaluation

# Helpers
"""
    xval(d::Normal, z::Real)

Computes the x-value based on a Normal distribution and a z-value.
"""
function xval(d::Normal, z::Real)
    if isinf(z) && iszero(d.σ)
        d.μ + one(d.σ) * z
    else
        d.μ + d.σ * z
    end
end
"""
    zval(d::Normal, x::Real)

Computes the z-value based on a Normal distribution and a x-value.
"""
zval(d::Normal, x::Real) = (x - d.μ) / d.σ

gradlogpdf(d::Normal, x::Real) = -zval(d, x) / d.σ
# logpdf
function logpdf(d::Normal, x::Real)
    if iszero(d.σ)
        d.μ == x ? Inf : -Inf
    else
        -(zval(d, x)^2 + log2π) / 2 - log(d.σ)
    end
end
# pdf
function pdf(d::Normal, x::Real)
    if iszero(d.σ)
        d.μ == x ? Inf : 0.0
    else
        exp(-zval(d, x)^2 / 2) * invsqrt2π / d.σ
    end
end
# logcdf
function logcdf(d::Normal, x::Real)
    if iszero(d.σ)
        d.μ ≤ x ? 0.0 : -Inf
    else
        z = zval(d, x)
        if z < -1.0
            log(erfcx(-z * invsqrt2)/2) - abs2(z)/2
        else
            log1p(-erfc(z * invsqrt2)/2)
        end
    end
end
# logccdf
function logccdf(d::Normal, x::Real)
    if iszero(d.σ) && x == d.μ
        z = zval(Normal(zero(d.μ), d.σ), one(x))
    else
        z = zval(d, x)
    end
    z > 1.0 ?
    log(erfcx(z * invsqrt2) / 2) - z^2 / 2 :
    log1p(-erfc(-z * invsqrt2) / 2)
end
# cdf
function cdf(d::Normal{T}, x::S) where {T, S <: Real}
    if iszero(d.σ)
        convert(promote_type(T, S), d.μ ≤ x)
    else
        erfc(-zval(d, x) * invsqrt2) / 2
    end
end
# ccdf
function ccdf(d::Normal, x::Real)
    z = iszero(d.σ) && x == d.μ ?
        zval(Normal(zero(d.μ), d.σ), one(x)) :
        zval(d, x)
    erfc(z * invsqrt2) / 2
end
# invlogcdf
"""
    norminvlogcdf(lp::Real)

Helper function that calls `_norminvlogcdf_impl` used for `invlogccdf` with the Normal distributions.
"""
norminvlogcdf(lp::Real) = _norminvlogcdf_impl(convert(Float64, lp))
norminvlogcdf(lp::Union{Float16,Float32}) = convert(typeof(lp), _norminvlogcdf_impl(convert(Float64, lp)))
invlogcdf(d::Normal, lp::Real) = xval(d, norminvlogcdf(lp))
# invlogccdf
invlogccdf(d::Normal, lp::Real) = xval(d, -norminvlogcdf(lp))
# quantile
function quantile(d::Normal, p::Real)
    if iszero(d.σ)
        if iszero(p)
            -Inf
        elseif isone(p)
            Inf
        else
            0.5
        end
    end
    xval(d, -erfcinv(2p) * sqrt2)
end
# cquantile
function cquantile(d::Normal, q::Real)
    if iszero(d.σ)
        if iszero(q)
            Inf
        elseif isone(q)
            -Inf
        else
            0.5
        end
    end
    xval(d, erfcinv(2q) * sqrt2)
end

# norminvcdf & norminvlogcdf implementation
"""
    _norminvlogcdf_impl(lp::Float64)

Uses `_qnorm_ker1` and `_qnorm_ker2` to obtain the ational approximations for the inverse cdf and its logarithm, from:
Wichura, M.J. (1988) Algorithm AS 241: The Percentage Points of the Normal Distribution
  Journal of the Royal Statistical Society. Series C (Applied Statistics), Vol. 37, No. 3, pp. 477-484
"""
function _norminvlogcdf_impl(lp::Float64)
    if isfinite(lp) && lp < 0.0
        q = exp(lp) - 0.5
        # qnorm_kernel(lp, q, true)
        if abs(q) ≤ 0.425
            _qnorm_ker1(q)
        else
            r = sqrt(q < 0 ? -lp : -log1mexp(lp))
            return copysign(_qnorm_ker2(r), q)
        end
    elseif lp ≥ 0.0
        iszero(lp) ? Inf : NaN
    else # lp is -Inf or NaN
        lp
    end
end
"""
    _qnorm_ker1(r::Float64)

Along with `_qnorm_ker2` these helpers enable the rational approximations for the inverse cdf and its logarithm
"""
function _qnorm_ker1(q::Float64)
    # pre-condition: abs(q) ≤ 0.425
    r = 0.180625 - q^2
    return q * @horner(r,
                       3.38713_28727_96366_6080e0,
                       1.33141_66789_17843_7745e2,
                       1.97159_09503_06551_4427e3,
                       1.37316_93765_50946_1125e4,
                       4.59219_53931_54987_1457e4,
                       6.72657_70927_00870_0853e4,
                       3.34305_75583_58812_8105e4,
                       2.50908_09287_30122_6727e3) /
    @horner(r,
            1.0,
            4.23133_30701_60091_1252e1,
            6.87187_00749_20579_0830e2,
            5.39419_60214_24751_1077e3,
            2.12137_94301_58659_5867e4,
            3.93078_95800_09271_0610e4,
            2.87290_85735_72194_2674e4,
            5.22649_52788_52854_5610e3)
end
"""
    _qnorm_ker2(r::Float64)

Along with `_qnorm_ker1` these helpers enable the rational approximations for the inverse cdf and its logarithm
"""
function _qnorm_ker2(r::Float64)
    if r < 5.0
        r -= 1.6
        @horner(r,
                1.42343_71107_49683_57734e0,
                4.63033_78461_56545_29590e0,
                5.76949_72214_60691_40550e0,
                3.64784_83247_63204_60504e0,
                1.27045_82524_52368_38258e0,
                2.41780_72517_74506_11770e-1,
                2.27238_44989_26918_45833e-2,
                7.74545_01427_83414_07640e-4) /
        @horner(r,
                1.0,
                2.05319_16266_37758_82187e0,
                1.67638_48301_83803_84940e0,
                6.89767_33498_51000_04550e-1,
                1.48103_97642_74800_74590e-1,
                1.51986_66563_61645_71966e-2,
                5.47593_80849_95344_94600e-4,
                1.05075_00716_44416_84324e-9)
    else
        r -= 5.0
        @horner(r,
                6.65790_46435_01103_77720e0,
                5.46378_49111_64114_36990e0,
                1.78482_65399_17291_33580e0,
                2.96560_57182_85048_91230e-1,
                2.65321_89526_57612_30930e-2,
                1.24266_09473_88078_43860e-3,
                2.71155_55687_43487_57815e-5,
                2.01033_43992_92288_13265e-7) /
        @horner(r,
                1.0,
                5.99832_20655_58879_37690e-1,
                1.36929_88092_27358_05310e-1,
                1.48753_61290_85061_48525e-2,
                7.86869_13114_56132_59100e-4,
                1.84631_83175_10054_68180e-5,
                1.42151_17583_16445_88870e-7,
                2.04426_31033_89939_78564e-15)
    end
end

mgf(d::Normal, t::Real) = exp(t * d.μ + d.σ^2 / 2 * t^2)
cf(d::Normal, t::Real) = exp(im * t * d.μ - d.σ^2 / 2 * t^2)

#### Sampling

rand(rng::AbstractRNG, d::Normal{T}) where {T} = d.μ + d.σ * randn(rng, T)

#### Fitting

struct NormalStats <: SufficientStats
    s::Float64    # (weighted) sum of x
    m::Float64    # (weighted) mean of x
    s2::Float64   # (weighted) sum of (x - μ)^2
    tw::Float64    # total sample weight
end

function suffstats(::Type{<:Normal}, x::AbstractArray{T}) where T<:Real
    n = length(x)

    # compute s
    s = x[1]
    for i = 2:n
        @inbounds s += x[i]
    end
    m = s / n

    # compute s2
    s2 = abs2(x[1] - m)
    for i = 2:n
        @inbounds s2 += abs2(x[i] - m)
    end

    NormalStats(s, m, s2, n)
end

function suffstats(::Type{<:Normal}, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
    n = length(x)

    # compute s
    tw = w[1]
    s = w[1] * x[1]
    for i = 2:n
        @inbounds wi = w[i]
        @inbounds s += wi * x[i]
        tw += wi
    end
    m = s / tw

    # compute s2
    s2 = w[1] * abs2(x[1] - m)
    for i = 2:n
        @inbounds s2 += w[i] * abs2(x[i] - m)
    end

    NormalStats(s, m, s2, tw)
end

# Cases where μ or σ is known

struct NormalKnownMu <: IncompleteDistribution
    μ::Float64
end

struct NormalKnownMuStats <: SufficientStats
    μ::Float64      # known mean
    s2::Float64     # (weighted) sum of (x - μ)^2
    tw::Float64     # total sample weight
end

function suffstats(g::NormalKnownMu, x::AbstractArray{T}) where T<:Real
    μ = g.μ
    s2 = abs2(x[1] - μ)
    for i = 2:length(x)
        @inbounds s2 += abs2(x[i] - μ)
    end
    NormalKnownMuStats(g.μ, s2, length(x))
end

function suffstats(g::NormalKnownMu, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
    μ = g.μ
    s2 = abs2(x[1] - μ) * w[1]
    tw = w[1]
    for i = 2:length(x)
        @inbounds wi = w[i]
        @inbounds s2 += abs2(x[i] - μ) * wi
        tw += wi
    end
    NormalKnownMuStats(g.μ, s2, tw)
end

struct NormalKnownSigma <: IncompleteDistribution
    σ::Float64

    function NormalKnownSigma(σ::Float64)
        σ > 0 || throw(ArgumentError("σ must be a positive value."))
        new(σ)
    end
end

struct NormalKnownSigmaStats <: SufficientStats
    σ::Float64      # known std.dev
    sx::Float64      # (weighted) sum of x
    tw::Float64     # total sample weight
end

function suffstats(g::NormalKnownSigma, x::AbstractArray{T}) where T<:Real
    NormalKnownSigmaStats(g.σ, sum(x), Float64(length(x)))
end

function suffstats(g::NormalKnownSigma, x::AbstractArray{T}, w::AbstractArray{T}) where T<:Real
    NormalKnownSigmaStats(g.σ, dot(x, w), sum(w))
end

# fit_mle based on sufficient statistics

fit_mle(::Type{<:Normal}, ss::NormalStats) = Normal(ss.m, sqrt(ss.s2 / ss.tw))
fit_mle(g::NormalKnownMu, ss::NormalKnownMuStats) = Normal(g.μ, sqrt(ss.s2 / ss.tw))
fit_mle(g::NormalKnownSigma, ss::NormalKnownSigmaStats) = Normal(ss.sx / ss.tw, g.σ)

# generic fit_mle methods

function fit_mle(::Type{<:Normal}, x::AbstractArray{T}; mu::Float64=NaN, sigma::Float64=NaN) where T<:Real
    if isnan(mu)
        if isnan(sigma)
            fit_mle(Normal, suffstats(Normal, x))
        else
            g = NormalKnownSigma(sigma)
            fit_mle(g, suffstats(g, x))
        end
    else
        if isnan(sigma)
            g = NormalKnownMu(mu)
            fit_mle(g, suffstats(g, x))
        else
            Normal(mu, sigma)
        end
    end
end

function fit_mle(::Type{<:Normal}, x::AbstractArray{T}, w::AbstractArray{Float64}; mu::Float64=NaN, sigma::Float64=NaN) where T<:Real
    if isnan(mu)
        if isnan(sigma)
            fit_mle(Normal, suffstats(Normal, x, w))
        else
            g = NormalKnownSigma(sigma)
            fit_mle(g, suffstats(g, x, w))
        end
    else
        if isnan(sigma)
            g = NormalKnownMu(mu)
            fit_mle(g, suffstats(g, x, w))
        else
            Normal(mu, sigma)
        end
    end
end


"""
    matB(X, y; fun = mean)
Compute the between covariance matrix ("B") of `X`.
* `X` : X-data (n, p).
* `y` : A vector (n,) defing the class memberships.
"""
matB = function(X, y)
    X = ensure_mat(X)
    y = vec(y)
    z = aggstat(X; group = y, fun = mean)
    B = matcov(z.res, mweights(z.ni))
    (B = B, ct = z.res, lev = z.lev, ni = z.ni)
end

"""
    matW(X, y; fun = mean)
Compute the within covariance matrix ("W") of `X`.
* `X` : X-data (n, p).
* `y` : A vector (n,) defing the class memberships.
"""
matW = function(X, y)
    X = ensure_mat(X)
    y = vec(y)
    ztab = tab(y)
    lev = ztab.keys
    nlev = length(lev)
    ni = collect(values(ztab))
    # Case with y(s) with only 1 obs
    sum(ni .== 1) > 0 ? sigma_1obs = matcov(X) : nothing
    # End
    w = mweights(ni)
    Wi = list(nlev, Matrix{Float64})
    W = zeros(1, 1)
    @inbounds for i in 1:nlev
        if ni[i] == 1
            Wi[i] = sigma_1obs
        else
            s = findall(y .== lev[i])
            Wi[i] = matcov(X[s, :])
        end
        if i == 1
            W = w[i] * Wi[i]
        else
            W = W + w[i] * Wi[i]
            # Alternative: Could give weight=0 to the class(es) with 1 obs
        end
    end
    (W = W, Wi = Wi, lev = lev, ni = ni)
end


module grids

export Grid,createFVgrid, createAdvecDiffOperator

type Grid
   lft::Array{Float64,1}
   ctr::Array{Float64,1}
   rgt::Array{Float64,1}
   BC::Array{Int64,1}
end

function createDifferentiationMatrix(vgrid::Array{Float64,1})
   Nv = length(vgrid)
   for i in 1:Nv
     ddv[i,2:end-1] = 1.0/(vgrid[3:end]-vgrid[1:end-2])
     ddv[i,1] = 1.0/(vgrid[2]-vgrid[1])
     ddv[i,end] = 1.0/(vgrid[end]-vgrid[end-1])
   end
end

function createFVgrid(Nv::Int64,v0::Float64,vf::Float64,BCs::Array{Int64,1})
   leftFaces = Float64[]
   centerPts = Float64[]
   rightFaces = Float64[]
   dv = (vf-v0)/Nv
   for i in 1:Nv
      push!(leftFaces, (i-1)*dv)
   end
   rightFaces = leftFaces[2:end]
   push!(rightFaces,leftFaces[end] + (leftFaces[end]-leftFaces[end-1]))

   centerPts = 0.5*( leftFaces[1:end] + rightFaces[1:end] )

   return Grid(leftFaces,centerPts,rightFaces,BCs)
end

function createAdvecDiffOperator(grid::Grid,jacobian::Function,advection::Function,diffusion::Function,boundaryVals::Array{Float64,1},source::Array{Float64,1})

   N = length(grid.ctr)
   a = zeros(Float64,N)
   b = zeros(Float64,N)
   c = zeros(Float64,N)
   operator = spzeros(N,N)

   # See ../docs/operator.pdf for description of this
   # Interior points
   denom = jacobian(grid.ctr[2:end-1]).*(grid.rgt[2:end-1]-grid.lft[2:end-1])
   A_left = jacobian(grid.lft[2:end-1]).*advection(grid.lft[2:end-1])
   A_right = jacobian(grid.rgt[2:end-1]).*advection(grid.rgt[2:end-1])
   D_left = (jacobian(grid.lft[2:end-1]).*diffusion(grid.lft[2:end-1]))./(grid.ctr[2:end-1] - grid.ctr[1:end-2])
   D_right = (jacobian(grid.rgt[2:end-1]).*diffusion(grid.rgt[2:end-1]))./(grid.ctr[3:end] - grid.ctr[2:end-1])

   a[1:end-2] = (-0.5*A_left + D_left)./denom
   b[2:end-1] = (0.5*A_right - 0.5*A_left - D_right - D_left)./denom
   c[3:end] = (0.5*A_right + D_right)./denom

   # i=1 boundary condition
   if grid.BC[1] == 0
      # Dirichlet boundary conditions
      scaleFactor = b[2]
      b[1] = scaleFactor
      source[1] = boundaryVals[1]*scaleFactor
   elseif grid.BC[1] == 1
      # Flux-specified boundary
      denom = jacobian(grid.ctr[1])*(grid.rgt[1]-grid.lft[1])
      b[1] = (0.5*jacobian(grid.rgt[1]).*advection(grid.rgt[1]) - (jacobian(grid.rgt[1]).*diffusion(grid.rgt[1]))./(grid.ctr[2] - grid.ctr[1]) )/denom
      c[2] = (0.5*jacobian(grid.rgt[1]).*advection(grid.rgt[1]) + (jacobian(grid.rgt[1]).*diffusion(grid.rgt[1]))./(grid.ctr[2] - grid.ctr[1]) )/denom
      source[1] = source[1] + jacobian(grid.lft[1])*boundaryVals[1] / (jacobian(grid.ctr[1])*(grid.rgt[1]-grid.lft[1]))
   else
         error("Invalid boundary condition for i = 1")
   end

   # i=N boundary condition
   if grid.BC[2] == 0
      # Dirichlet boundary conditions
      scaleFactor = b[N-1]
      b[N] = scaleFactor
      source[N] = boundaryVals[2]*scaleFactor
   elseif grid.BC[2] == 1
      # Flux-specified boundary
      denom = jacobian(grid.ctr[N])*(grid.rgt[N]-grid.lft[N])
      a[N-1] = (-0.5*jacobian(grid.lft[N]).*advection(grid.lft[N]) + (jacobian(grid.lft[N]).*diffusion(grid.lft[N]))./(grid.ctr[N] - grid.ctr[N-1]) )/denom
      b[N] = (- 0.5*jacobian(grid.lft[N]).*advection(grid.lft[N]) - (jacobian(grid.lft[1]).*diffusion(grid.lft[1]))./(grid.ctr[N] - grid.ctr[N-1]) )/denom
      source[N] = source[N] - jacobian(grid.rgt[N])*boundaryVals[2] / (jacobian(grid.ctr[N])*(grid.rgt[N]-grid.lft[N]))
   else
         error("Invalid boundary condition for i = N")
   end

   for icol in 2:N-1
      operator[icol-1:icol+1,icol] = [a[icol],b[icol],c[icol]]
   end
   operator[1,1] = b[1]
   operator[1,2] = c[2]
   operator[N,N-1] = a[N-1]
   operator[N,N] = b[N]

   return operator

end

end


"""
    LIF

A leaky-integrate-fire neuron described by the following differential equation

``\\frac{\\mathrm{d}v}{\\mathrm{d}t} = \\frac{R}{\\tau} I - \\lambda``

# Fields
- `voltage::VT`: membrane potential
- `current::VT`: injected (unprocessed) current
- `lastt::IT`: the last time this neuron processed a spike
- `τm::VT`: membrane time constant
- `vreset::VT`: reset voltage potential
- `R::VT`: resistive constant (typically = 1)
"""
mutable struct LIF{VT<:Real, IT<:Integer} <: AbstractCell
    # required fields
    voltage::VT
    current::VT

    # model specific fields
    lastt::IT
    τm::VT
    vreset::VT
    R::VT
end

Base.show(io::IO, ::MIME"text/plain", neuron::LIF) =
    print(io, """LIF:
                     voltage: $(neuron.voltage)
                     current: $(neuron.current)
                     τm:      $(neuron.τm)
                     vreset:  $(neuron.vreset)
                     R:       $(neuron.R)""")
Base.show(io::IO, neuron::LIF) =
    print(io, "LIF(voltage = $(neuron.voltage), current = $(neuron.current))")

"""
    LIF(τm, vreset, vth, R = 1.0)

Create a LIF neuron with zero initial voltage and empty current queue.
"""
LIF(τm::Real, vreset::Real, R::Real = 1.0) = LIF{Float32, Int}(vreset, 0, 0, τm, vreset, R)

isactive(neuron::LIF, t::Integer; dt::Real = 1.0) = (neuron.current > 0)
getvoltage(neuron::LIF) = neuron.voltage

"""
    excite!(neuron::LIF, current)
    excite!(neurons::AbstractArray{<:LIF}, current)

Excite a `neuron` with external `current`.
"""
excite!(neuron::LIF, current) = (neuron.current += current)
excite!(neurons::T, current) where T<:AbstractArray{<:LIF} = (neurons.current .+= current)

"""
    spike!(neuron::LIF, t::Integer; dt::Real = 1.0)
    spike!(neurons::AbstractArray{<:LIF}, spikes; dt::Real = 1.0)

Record a output spike from the threshold function with the `neuron` body.
Sets `neuron.lastt`.
"""
spike!(neuron::LIF, t::Integer; dt::Real = 1.0) = (t > 0) && (neuron.voltage = neuron.vreset)
function spike!(neurons::T, spikes; dt::Real = 1.0) where T<:AbstractArray{<:LIF}
    map!((x, y, s) -> (s > 0) ? y : x, neurons.voltage, neurons.voltage, neurons.vreset, spikes)
end

"""
    lif(t::Real, I, V; R, tau)
    lif!(t::AbstractArray{<:Real}, I::AbstractArray{<:Real}, V::AbstractArray{<:Real}; R::AbstractArray{<:Real}, tau::AbstractArray{<:Real})
    lif!(t::CuVector{<:Real}, I::CuVector{<:Real}, V::CuVector{<:Real}; R::CuVector{<:Real}, tau::CuVector{<:Real})

Evaluate a leaky integrate-and-fire neuron.
Use `CuVector` instead of `Vector` to evaluate on GPU.

# Fields
- `t`: time since last evaluation in seconds
- `I`: external current
- `V`: current membrane potential
- `vrest`: resting membrane potential
- `R`: resistance constant
- `tau`: time constant
"""
function lif(t::Real, I, V; vrest, R, tau)
    V = V * exp(-t / tau) + vrest
    V += I * (R / tau)

    return V
end
function lif!(t::AbstractArray{<:Real}, I::AbstractArray{<:Real}, V::AbstractArray{<:Real}; vrest::AbstractArray{<:Real}, R::AbstractArray{<:Real}, tau::AbstractArray{<:Real})
    # account for leakage
    @. V = V * exp(-t / tau) + vrest

    # apply update step
    @. V += I * (R / tau)

    return V
end
function lif!(t::CuVector{<:Real}, I::CuVector{<:Real}, V::CuVector{<:Real}; vrest::CuVector{<:Real}, R::CuVector{<:Real}, tau::CuVector{<:Real})
    # account for leakage
    @. V = V * exp(-t / tau) + vrest

    # apply update step
    @. V += I * (R / tau)

    return V
end

"""
    evaluate!(neuron::LIF, t::Integer; dt::Real = 1.0)
    (neuron::LIF)(t::Integer; dt::Real = 1.0)
    evaluate!(neurons::AbstractArray{<:LIF}, t::Integer; dt::Real = 1.0)

Evaluate the neuron model at time `t`. Return the resulting membrane potential.
"""
function evaluate!(neuron::LIF, t::Integer; dt::Real = 1.0)
    neuron.voltage = lif((t - neuron.lastt) * dt, neuron.current, neuron.voltage; vrest = 0, R = neuron.R, tau = neuron.τm)
    neuron.lastt = t
    neuron.current = 0

    return neuron.voltage
end
(neuron::LIF)(t::Integer; dt::Real = 1.0) = evaluate!(neuron, t; dt = dt)
function evaluate!(neurons::T, t::Integer; dt::Real = 1.0) where T<:AbstractArray{<:LIF}
    lif!((t .- neurons.lastt) .* dt, neurons.current, neurons.voltage;
         vrest = zero(neurons.voltage), R = neurons.R, tau = neurons.τm)
    neurons.lastt .= t
    neurons.current .= 0

    return neurons.voltage
end

"""
    reset!(neuron::LIF)
    reset!(neurons::AbstractArray{<:LIF})

Reset `neuron` by setting the membrane potential to `neuron.vreset`.
"""
function reset!(neuron::LIF)
    neuron.lastt = 0
    neuron.voltage = neuron.vreset
    neuron.current = 0
end
function reset!(neurons::T) where T<:AbstractArray{<:LIF}
    neurons.lastt .= 0
    neurons.voltage .= neurons.vreset
    neurons.current .= 0
end


using Test, OrdinaryDiffEq, HyperbolicDiffEq, Roots

function calc_error(balance_law, uₐₙₐ, tmin, tmax, fnum, reconstruction, N, parallel)
    u₀ = x -> uₐₙₐ(tspan[1], x)

    tspan = (tmin, tmax)
    mesh = UniformPeriodicMesh1D(0., 2., N)
    fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, reconstruction, parallel)
    ode = semidiscretise(fv, u₀, tspan)
    if order(reconstruction) <= 4
        sol = solve(ode, SSPRK104(), dt=0.5*max_dt(tspan[1], ode.u0, fv),
                    save_everystep=false, dense=false)
    else
        sol = solve(ode, SSPRK104(), dt=0.01*max_dt(tspan[1], ode.u0, fv),
                    save_everystep=false, dense=false)
    end
    uana = compute_coefficients(x->uₐₙₐ(tspan[end],x), mesh, order(reconstruction))

    N, integrate(u->sum(abs.(u)), sol[end]-uana, mesh)
end

function calc_order_estimate(Ns_and_errors)
    orders = zeros(length(Ns_and_errors)-1)
    for i in 1:length(orders)
        orders[i] = -log(Ns_and_errors[i+1][2]/Ns_and_errors[i][2]) /
                     log(Ns_and_errors[i+1][1]/Ns_and_errors[i][1])
    end
    mean(orders)
end

function calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, reconstruction, Ns, parallel=Val{:serial}())
    calc_order_estimate(calc_error.(balance_law, uₐₙₐ, tspan..., fnum, reconstruction, Ns, parallel))
end

# test convergence orders
Ns_lo = 10 .* 2 .^ (2:6)
Ns_hi =  4 .* 2 .^ (2:5)

balance_law = ConstantLinearAdvection()
tspan = (0., 1.)
fnum = LocalLaxFriedrichsFlux()
uₐₙₐ = (t,x) -> sinpi(x-t)

println("CentralReconstruction")
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, CentralReconstruction(1), Ns_lo, Val{:serial}()) > 0.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, CentralReconstruction(3), Ns_lo, Val{:serial}()) > 2.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, CentralReconstruction(5), Ns_lo, Val{:serial}()) > 4.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, CentralReconstruction(7), Ns_hi, Val{:serial}()) > 6.8

@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, CentralReconstruction(1), Ns_lo, Val{:threads}()) > 0.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, CentralReconstruction(3), Ns_lo, Val{:threads}()) > 2.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, CentralReconstruction(5), Ns_lo, Val{:threads}()) > 4.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, CentralReconstruction(7), Ns_hi, Val{:threads}()) > 6.8

println("ModifiedENO")
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(1), Ns_lo, Val{:serial}()) > 0.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(2), Ns_lo, Val{:serial}()) > 1.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(3), Ns_lo, Val{:serial}()) > 2.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(4), Ns_lo, Val{:serial}()) > 3.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(5), Ns_lo, Val{:serial}()) > 4.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(6), Ns_hi, Val{:serial}()) > 5.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(7), Ns_hi, Val{:serial}()) > 6.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(8), Ns_hi, Val{:serial}()) > 7.5

@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(1), Ns_lo, Val{:threads}()) > 0.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(2), Ns_lo, Val{:threads}()) > 1.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(3), Ns_lo, Val{:threads}()) > 2.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(4), Ns_lo, Val{:threads}()) > 3.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(5), Ns_lo, Val{:threads}()) > 4.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(6), Ns_hi, Val{:threads}()) > 5.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(7), Ns_hi, Val{:threads}()) > 6.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(8), Ns_hi, Val{:threads}()) > 7.5

println("MinL2Choice")
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{1}(), MinL2Choice()), Ns_lo, Val{:serial}()) > 0.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{2}(), MinL2Choice()), Ns_lo, Val{:serial}()) > 1.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{3}(), MinL2Choice()), Ns_lo, Val{:serial}()) > 2.8
@test 0.2 < calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{4}(), MinL2Choice()), Ns_lo, Val{:serial}()) < 0.4
@test 0.2 < calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{5}(), MinL2Choice()), Ns_lo, Val{:serial}()) < 0.4
@test 0.5 < calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{6}(), MinL2Choice()), Ns_hi, Val{:serial}()) < 1.1
@test 0.5 < calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{7}(), MinL2Choice()), Ns_hi, Val{:serial}()) < 1.1
@test 0.5 < calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{8}(), MinL2Choice()), Ns_hi, Val{:serial}()) < 1.1

println("BiasedMinL2Choice")
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{1}(), BiasedMinL2Choice()), Ns_lo, Val{:serial}()) > 0.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{2}(), BiasedMinL2Choice()), Ns_lo, Val{:serial}()) > 1.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{3}(), BiasedMinL2Choice()), Ns_lo, Val{:serial}()) > 2.8
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{4}(), BiasedMinL2Choice()), Ns_lo, Val{:serial}()) > 3.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{5}(), BiasedMinL2Choice()), Ns_lo, Val{:serial}()) > 4.8

# Order reduction observed by Rogerson & Meiburg and Shu (1990).
balance_law = ConstantLinearAdvection()
tspan = (0., 1.)
fnum = LocalLaxFriedrichsFlux()
uₐₙₐ = (t,x) -> (sinpi(x-t)^2)^2
println("ClassicalChoiceENO")
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{2}(), ClassicalChoiceENO()), Ns_lo, Val{:serial}()) > 1.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{3}(), ClassicalChoiceENO()), Ns_lo, Val{:serial}()) > 2.2
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{4}(), ClassicalChoiceENO()), Ns_lo, Val{:serial}()) < 2.2
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{5}(), ClassicalChoiceENO()), Ns_lo, Val{:serial}()) < 1.9

println("BiasedENOChoice")
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{2}(), BiasedENOChoice()), Ns_lo, Val{:serial}()) > 1.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{3}(), BiasedENOChoice()), Ns_lo, Val{:serial}()) > 2.6
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{4}(), BiasedENOChoice()), Ns_lo, Val{:serial}()) > 3.4
@test calc_order_estimate(balance_law, uₐₙₐ, tspan, fnum, ModifiedENO(Val{5}(), BiasedENOChoice()), Ns_lo, Val{:serial}()) > 4.6

# test evaluation
mesh = UniformPeriodicMesh1D(0., 2., 50)

K = 1
fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, ModifiedENO(K), Val{:serial})
u = compute_coefficients(sinpi, mesh, K)
xplot, uplot = evaluate_coefficients(u, fv)
@test norm(sinpi.(xplot) - uplot) < 5.e-1
K = 2
fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, ModifiedENO(K), Val{:serial})
u = compute_coefficients(sinpi, mesh, K)
xplot, uplot = evaluate_coefficients(u, fv)
@test norm(sinpi.(xplot) - uplot) < 4.e-2
K = 3
fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, ModifiedENO(K), Val{:serial})
u = compute_coefficients(sinpi, mesh, K)
xplot, uplot = evaluate_coefficients(u, fv)
@test norm(sinpi.(xplot) - uplot) < 2.e-3
K = 4
fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, ModifiedENO(K), Val{:serial})
u = compute_coefficients(sinpi, mesh, K)
xplot, uplot = evaluate_coefficients(u, fv)
@test norm(sinpi.(xplot) - uplot) < 2.e-4
K = 5
fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, ModifiedENO(K), Val{:serial})
u = compute_coefficients(sinpi, mesh, K)
xplot, uplot = evaluate_coefficients(u, fv)
@test norm(sinpi.(xplot) - uplot) < 7.e-6
K = 6
fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, ModifiedENO(K), Val{:serial})
u = compute_coefficients(sinpi, mesh, K)
xplot, uplot = evaluate_coefficients(u, fv)
@test norm(sinpi.(xplot) - uplot) < 5.e-7
K = 7
fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, ModifiedENO(K), Val{:serial})
u = compute_coefficients(sinpi, mesh, K)
xplot, uplot = evaluate_coefficients(u, fv)
@test norm(sinpi.(xplot) - uplot) < 3.e-8
K = 8
fv = UniformPeriodicReconstructedFV1D(balance_law, mesh, fnum, ModifiedENO(K), Val{:serial})
u = compute_coefficients(sinpi, mesh, K)
xplot, uplot = evaluate_coefficients(u, fv)
@test norm(sinpi.(xplot) - uplot) < 2.e-9


function solve_perspective_relaxation(U, mu, lambda, k_sparse, k_rank;
    zero_indices=[], one_indices=[], solver_output=0)
    """
    This function exactly solves the perspective reformulation relaxation of the
    problem of interest.

    :param U: An arbitrary n-by-n matrix.
    :param mu: The regularization parameter for the sparse matrix penalty
                   (Float64).
    :param lambda: The regularization parameter for the low rank matrix penalty
                   (Float64).
    :param k_sparse: The maximum sparsity of the sparse matrix (Int64).
    :param k_rank: The maximum rank of the low rank matrix (Int64).
    :param zero_indices: List of 2-tuples of Int64 where each entry consists of
                         an index (i, j) such that Y_ij is constrained to take
                         value 0.
    :param one_indices: List of 2-tuples of Int64 where each entry consists of
                        an index (i, j) such that S_ij is constrained to take
                        value 1 where S is the binary matrix denoting the
                        sparsity pattern of Y.
    :param solver_output: Solver_output param to be passed to Mosek (Int64).

    :return: This function returns two values. The first value is a tuple of 5
             n-by-n arrays that correspond to the solution of the perspective
             relaxation (the first element in the tuple is the matrix X, the
             second element is the matrix Y, the third element is the matrix Z,
             the fourth element is the matrix P and the fifth element is the
             matrix Theta). The second value is the optimal objective value of
             the relaxation (Float64).
    """

    n = size(U)[1]

    num_zeros = size(zero_indices)[1]
    num_ones = size(one_indices)[1]

    @assert num_ones <= k_sparse
    @assert num_zeros <= n * n - k_sparse

    # Build perspective formulation relaxation using JuMP
    m = Model(Mosek.Optimizer)
    set_optimizer_attribute(m, "MSK_IPAR_LOG", solver_output)

    @variable(m, X[i=1:n, j=1:n])
    @variable(m, Y[i=1:n, j=1:n])
    @variable(m, Z[i=1:n, j=1:n] >= 0)
    @variable(m, P[i=1:n, j=1:n])
    @variable(m, P_hat[i=1:n, j=1:n])
    @variable(m, alpha[i=1:n, j=1:n])
    @variable(m, theta[i=1:n, j=1:n] >= 0)

    @constraint(m, [i=1:n, j=1:n], Z[i, j] <= 1)
    @SDconstraint(m, [i=1:n, j=1:n],
        [theta[i, j] Y[i, j]; Y[i, j] Z[i, j]] >= 0)
    for (i, j) in zero_indices
        @constraint(m, Z[i, j] == 0)
        @constraint(m, Y[i, j] == 0)
        @constraint(m, theta[i, j] == 0)
    end
    for (i, j) in one_indices
        @constraint(m, Z[i, j] == 1)
        @constraint(m, Y[i, j] == (U[i, j] - X[i, j]) / (1 + mu))
    end
    @constraint(m, sum(Z[i, j] for i=1:n, j=1:n) <= k_sparse)

    @constraint(m, sum(P[i, i] for i = 1:n) <= k_rank)
    @constraint(m, 1*Matrix(I, n, n) - P in PSDCone())
    @constraint(m, [alpha X; X' P] in PSDCone())

    @objective(m, Min, sum((U[i, j]-X[i, j]-Y[i, j])^2 for i=1:n for j=1:n) +
        lambda * sum(alpha[i, i] for i=1:n) +
        mu * sum(theta[i, j] for i=1:n for j=1:n))

    # Solve perspective relaxation
    optimize!(m)

    X_sol = value.(X)
    Y_sol = value.(Y)

    P_sol = value.(P)
    Z_sol = value.(Z)
    alpha_sol = value.(alpha)
    theta_sol = value.(theta)

    return (X_sol, Y_sol, Z_sol, P_sol, alpha_sol), objective_value(m)

end;


export Qubit, Q0, Q1, RandomQubit

"""
A `Qubit` represents a single qubit.

* `Qubit()` creates the `|0>` value; equivalent to `Qubit(1,0)`.
* `Qubit(0)` and `Qubit(1)` create the `|0>` and `|1>` values, respecively. They are also
  available as `Q0` and `Q1`.
* `Qubit(w::Complex,z::Complex)` creates a qubit whose state is `[w;z]` (normalized)
"""
mutable struct Qubit <: QuantumState
    vec::Vector{CC}
    function Qubit(w, z)
        v = [w; z]
        v /= _norm(v)
        new(v)
    end
end

Qubit(b::Integer = 0) = b == 0 ? Qubit(1, 0) : Qubit(0, 1)

"""
`Q0` is the value `Qubit(1,0)` which is a quantum zero.
"""
const Q0 = Qubit(0)

"""
`Q1` is the value `Qubit(0,1)` which is a quantum one.
"""
const Q1 = Qubit(1)

function show(io::IO, q::Qubit)
    x = q.vec[1]
    y = q.vec[2]
    print(io, "Qubit($x,$y)")
end

"""
`measure!(q::Qubit)` performs a quantum measurement on the qubit `q`
returning `0` or `1` and overwriting `q` to be either `Q0` or `Q1`
respectively.
"""
function measure!(q::Qubit)::Int
    p = abs(q.vec[2])^2
    val::Int = 0
    if rand() < p
        val = 1
        q.vec = Q1.vec
    else
        q.vec = Q0.vec
    end
    return val
end

"""
`RandomQubit()` returns a random qubit.
"""
function RandomQubit()
    a = randn() + randn() * im
    b = randn() + randn() * im
    return Qubit(a, b)
end

length(Q::Qubit) = 1