import Test: @test, @testset
import Luna: Fields, FFTW, Grid, Maths, PhysData, Processing, Modes, Tools, Maths
import Statistics: mean, std
import Random: MersenneTwister

# note that most of the Fields.jl code is tested in many other modules

function getceo(t, Et, It, ω0)
    Δt = t[argmax(It)] - t[argmax(Et)]
    Δt*ω0
end

@testset "Wavelength" begin
    # real
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    ϕ = [0.0, 0.0]
    grid = Grid.RealGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{Float64}(undef, length(grid.t))
    FT = FFTW.plan_rfft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    @test isapprox(PhysData.wlfreq(grid.ω[argmax(abs2.(Eω))]), λ0, rtol=3e-4)
    λ0 = 320e-9
    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    @test isapprox(PhysData.wlfreq(grid.ω[argmax(abs2.(Eω))]), λ0, rtol=3e-4)
    λ0 = 800e-9
    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    @test isapprox(PhysData.wlfreq(grid.ω[argmax(abs2.(Eω))]), λ0, rtol=3e-4)
    λ0 = 320e-9
    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    @test isapprox(PhysData.wlfreq(grid.ω[argmax(abs2.(Eω))]), λ0, rtol=3e-4)

    # Envelope
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    ϕ = [0.0, 0.0]
    grid = Grid.EnvGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{ComplexF64}(undef, length(grid.t))
    FT = FFTW.plan_fft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    @test isapprox(PhysData.wlfreq(grid.ω[argmax(abs2.(Eω))]), λ0, rtol=3e-4)
    λ0 = 320e-9
    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    @test isapprox(PhysData.wlfreq(grid.ω[argmax(abs2.(Eω))]), λ0, rtol=3e-4)
    λ0 = 800e-9
    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    @test isapprox(PhysData.wlfreq(grid.ω[argmax(abs2.(Eω))]), λ0, rtol=3e-4)
    λ0 = 320e-9
    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    @test isapprox(PhysData.wlfreq(grid.ω[argmax(abs2.(Eω))]), λ0, rtol=3e-4)
end

@testset "Energy" begin
    # real
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    ϕ = [0.0, 0.0]
    grid = Grid.RealGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{Float64}(undef, length(grid.t))
    FT = FFTW.plan_rfft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    @test isapprox(energy_t(Et), energy, rtol=1e-14)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    @test isapprox(energy_t(Et), energy, rtol=1e-14)

    # Envelope
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    ϕ = [0.0, 0.0]
    grid = Grid.EnvGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{ComplexF64}(undef, length(grid.t))
    FT = FFTW.plan_fft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    @test isapprox(energy_t(Et), energy, rtol=1e-14)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    @test isapprox(energy_t(Et), energy, rtol=1e-14)
end

@testset "Duration" begin
    # real
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    ϕ = [0.0, 0.0]
    grid = Grid.RealGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{Float64}(undef, length(grid.t))
    FT = FFTW.plan_rfft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Maths.hilbert(Et))
    @test isapprox(Maths.fwhm(grid.t, It), τfwhm, rtol=1e-5)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Maths.hilbert(Et))
    @test isapprox(Maths.fwhm(grid.t, It), τfwhm, rtol=1e-5)

    # Envelope
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    ϕ = [0.0, 0.0]
    grid = Grid.EnvGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{ComplexF64}(undef, length(grid.t))
    FT = FFTW.plan_fft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Et)
    @test isapprox(Maths.fwhm(grid.t, It), τfwhm, rtol=2e-5)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Et)
    @test isapprox(Maths.fwhm(grid.t, It), τfwhm, rtol=3e-5)
end

@testset "Position" begin
    # real
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    τ0 = 5e-15
    # elements of ϕ are [CEP, group delay, GDD, TOD, ...]
    # so [0.0, τ0] is a delay by τ0
    ϕ = [0.0, τ0]
    grid = Grid.RealGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{Float64}(undef, length(grid.t))
    FT = FFTW.plan_rfft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Maths.hilbert(Et))
    @test isapprox(grid.t[argmax(It)], τ0, rtol=1e-15, atol=1e-15)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Maths.hilbert(Et))
    @test isapprox(grid.t[argmax(It)], τ0, rtol=1e-15, atol=1e-15)

    # Envelope
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    τ0 = 5e-15
    # elements of ϕ are [CEP, group delay, GDD, TOD, ...]
    # so [0.0, τ0] is a delay by τ0
    ϕ = [0.0, τ0]
    grid = Grid.EnvGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{ComplexF64}(undef, length(grid.t))
    FT = FFTW.plan_fft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Et)
    @test isapprox(grid.t[argmax(It)], τ0, rtol=1e-15, atol=1e-15)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Et)
    @test isapprox(grid.t[argmax(It)], τ0, rtol=1e-15, atol=1e-15)

    # non zero
    τ0 = -564e-15

    #real
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    τ0 = 5e-15
    # elements of ϕ are [CEP, group delay, GDD, TOD, ...]
    # so [0.0, τ0] is a delay by τ0
    ϕ = [0.0, τ0]
    grid = Grid.RealGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{Float64}(undef, length(grid.t))
    FT = FFTW.plan_rfft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Maths.hilbert(Et))
    @test isapprox(grid.t[argmax(It)], τ0, rtol=1e-15, atol=1e-15)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Maths.hilbert(Et))
    @test isapprox(grid.t[argmax(It)], τ0, rtol=1e-15, atol=1e-15)

    # Envelope
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    τ0 = 5e-15
    # elements of ϕ are [CEP, group delay, GDD, TOD, ...]
    # so [0.0, τ0] is a delay by τ0
    ϕ = [0.0, τ0]
    grid = Grid.EnvGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{ComplexF64}(undef, length(grid.t))
    FT = FFTW.plan_fft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Et)
    @test isapprox(grid.t[argmax(It)], τ0, rtol=1e-15, atol=1e-15)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=ϕ)
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Et)
    @test isapprox(grid.t[argmax(It)], τ0, rtol=1e-15, atol=1e-15)
end

@testset "CEO" begin
    # real
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    ϕCEO = 0.0
    grid = Grid.RealGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{Float64}(undef, length(grid.t))
    FT = FFTW.plan_rfft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=[ϕCEO])
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Maths.hilbert(Et))
    @test isapprox(getceo(grid.t, Et, It, PhysData.wlfreq(λ0)), ϕCEO, rtol=1e-15, atol=1e-15)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=[ϕCEO])
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Maths.hilbert(Et))
    @test isapprox(getceo(grid.t, Et, It, PhysData.wlfreq(λ0)), ϕCEO, rtol=1e-15, atol=1e-15)

    # Envelope
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    ϕCEO = 0.0
    grid = Grid.EnvGrid(1.0, λ0, (160e-9, 3000e-9), 10e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{ComplexF64}(undef, length(grid.t))
    FT = FFTW.plan_fft(x, 1)

    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=[ϕCEO])
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Et)
    @test isapprox(getceo(grid.t, real(Et.*exp.(im .* grid.ω0 .* grid.t)), It, PhysData.wlfreq(λ0)), ϕCEO, rtol=1e-15, atol=1e-15)

    input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=[ϕCEO])
    Eω = input(grid, FT)
    Et = FT \ Eω
    It = abs2.(Et)
    @test isapprox(getceo(grid.t, real(Et.*exp.(im .* grid.ω0 .* grid.t)), It, PhysData.wlfreq(λ0)), ϕCEO, rtol=1e-15, atol=1e-15)

    # non zero

    #real
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    τ0 = 0.0
    grid = Grid.RealGrid(1.0, λ0, (100e-9, 3000e-9), 1e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{Float64}(undef, length(grid.t))
    FT = FFTW.plan_rfft(x, 1)

    # Make CEO exact multiple of one grid point to avoid issues with argmax() in getceo()
    δt = grid.t[2] - grid.t[1]
    for i = 1:10
        ϕCEO = i*δt*PhysData.wlfreq(λ0)

        input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=[ϕCEO])
        Eω = input(grid, FT)
        Et = FT \ Eω
        It = abs2.(Maths.hilbert(Et))
        @test isapprox(abs(getceo(grid.t, Et, It, PhysData.wlfreq(λ0))), ϕCEO, rtol=1e-10)

        input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=[ϕCEO])
        Eω = input(grid, FT)
        Et = FT \ Eω
        It = abs2.(Maths.hilbert(Et))
        @test isapprox(abs(getceo(grid.t, Et, It, PhysData.wlfreq(λ0))), ϕCEO, rtol=1e-10)
    end

    # Envelope
    τfwhm = 30e-15
    λ0 = 800e-9
    energy = 1e-6
    τ0 = 0.0
    grid = Grid.EnvGrid(1.0, λ0, (100e-9, 3000e-9), 1e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{ComplexF64}(undef, length(grid.t))
    FT = FFTW.plan_fft(x, 1)

    # Make CEO exact multiple of one grid point to avoid issues with argmax() in getceo()
    δt = grid.t[2] - grid.t[1]

    for i = 1:10
        ϕCEO = i*δt*PhysData.wlfreq(λ0)

        input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=[ϕCEO])
        Eω = input(grid, FT)
        Et = FT \ Eω
        It = abs2.(Et)
        @test isapprox(
            abs(getceo(grid.t, real(Et.*exp.(im .* grid.ω0 .* grid.t)), It, PhysData.wlfreq(λ0))),
            ϕCEO,
            rtol=1e-10)

        input = Fields.SechField(λ0=λ0, τfwhm=τfwhm, energy=energy, ϕ=[ϕCEO])
        Eω = input(grid, FT)
        Et = FT \ Eω
        It = abs2.(Et)
        @test isapprox(
            abs(getceo(grid.t,real(Et.*exp.(im .* grid.ω0 .* grid.t)), It, PhysData.wlfreq(λ0))),
            ϕCEO,
            rtol=1e-10)
    end
end

@testset "CW fields" begin
    λ0 = 1064e-9
    Pavg = 20.0
    Δλ = 4e-9
    grid = Grid.EnvGrid(1.0, λ0, (980e-9, 1160e-9), 500e-12)
    energy_t = Fields.energyfuncs(grid)[1]
    x = Array{ComplexF64}(undef, length(grid.t))
    FT = FFTW.plan_fft(x, 1)
    input = Fields.CWSech(λ0=λ0, Pavg=Pavg, Δλ=Δλ, rng=MersenneTwister(0))
    Eω = input(grid, FT)
    Et = FT \ Eω
    I = Fields.It(Et, grid)
    istart = findfirst(isequal(1.0), grid.twin)
    iend = findlast(isequal(1.0), grid.twin)
    # test average power
    @test isapprox(mean(I[istart:iend]), Pavg, rtol=5e-16)
    # test coherence time
    @test isapprox(Processing.coherence_time(grid, Et), 3.35/(PhysData.c*(Δλ)/λ0^2*2π), rtol=1e-2)
    idcs = sortperm(PhysData.wlfreq.(grid.ω))
    # test spectral width
    @test isapprox(Maths.fwhm(PhysData.wlfreq.(grid.ω)[idcs], abs2.(Eω[idcs])), Δλ, rtol=3e-3)
    # now do the same for a number of realisations
    Eωs = hcat([Fields.CWSech(λ0=λ0, Pavg=Pavg, Δλ=Δλ, rng=MersenneTwister(i))(grid, FT) for i = 1:5]...)
    Iωs = abs2.(Eωs)
    Iωav = mean(Iωs, dims=2)[:,1]
    idcs = sortperm(PhysData.wlfreq.(grid.ω))
    # test average spectral width
    @test isapprox(Maths.fwhm(PhysData.wlfreq.(grid.ω)[idcs], Iωav[idcs], minmax=:max), Δλ, rtol=6e-4)
    Ets = FFTW.ifft(Eωs, 1)
    Its = abs2.(Ets)
    Itav = mean(Its[istart:iend,:])
    # test average power
    @test isapprox(Itav, Pavg, rtol=5e-16)
    # test diversity of power fluctuations
    @test mean(std(Its[istart:iend,:], dims=2)[:,1]) > 10
end

@testset "Propagation" begin
    λ0 = 800e-9
    τfwhm = 2.5e-15
    grid = Grid.RealGrid(1, λ0, (400e-9, 1200e-9), 500e-15)
    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=1e-6)
    x = Array{Float64}(undef, length(grid.t))
    FT = FFTW.plan_rfft(x, 1)
    Eω = input(grid, FT)
    Eωβ1 = Fields.prop_taylor(Eω, grid, [0, 10e-15], λ0)
    Et = FT \ Eωβ1
    @test isapprox(Maths.moment(grid.t, abs2.(Maths.hilbert(Et))), 10e-15, rtol=1e-6)

    # Test sign of dispersion
    Eωβ2 = Fields.prop_taylor(Eω, grid, [0, 0, 15e-30], λ0) # positive chirp
    Et = FT \ Eωβ2
    gab = Maths.gabor(grid.t, Et, [-10e-15, 10e-15], 3e-15) # spectrogram
    ω0 = Maths.moment(grid.ω, abs2.(gab))
    @test ω0[1] < ω0[2] # mean frequency at earlier time should be lower (upchirp)

    # Test pulse stretching for Gaussian pulse
    τfwhm = 30e-15
    τ0 = Tools.τfw_to_τ0(τfwhm, :gauss)
    input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=1e-6)
    Eω = input(grid, FT)
    Eωβ2 = Fields.prop_taylor(Eω, grid, [0, 0, τ0^2], λ0) # should lead to √2 increase
    Et = FT \ Eωβ2
    τfwβ2 = Maths.fwhm(grid.t, abs2.(Maths.hilbert(Et)); method=:spline)
    τ0β2 = Tools.τfw_to_τ0(τfwβ2, :gauss)
    @test τ0β2 ≈ √2 * τ0

    # Test pulse stretching and sign of the dispersion for modal propagation
    ω0 = PhysData.wlfreq(λ0)
    # Artificial mode with τ0^2 2nd order dispersion over 2 m and α=0.1
    β(ω; z=0) = ω0/PhysData.c + 1/(0.999*PhysData.c)*(ω-ω0) + τ0^2/4*(ω-ω0)^2
    α(ω; z=0) = 0.1
    m = Modes.arbitrary(neff=Modes.neff_from_αβ(α, β))
    Eωm = copy(Eω)
    Fields.prop_mode!(Eωm, grid.ω, m, 2, λ0)
    Et = FT \ Eωm
    τfwm = Maths.fwhm(grid.t, abs2.(Maths.hilbert(Et)); method=:spline)
    τ0m = Tools.τfw_to_τ0(τfwm, :gauss)
    @test τ0m ≈ √2 * τ0
    # Check signs are correct:
    # α = 0.1 should give loss
    et, eω = Fields.energyfuncs(grid)
    @test eω(Eω)*exp(-0.2) ≈ eω(Eωm)
    # β2 > 0 should give positive chirp:
    gab = Maths.gabor(grid.t, Et, [-10e-15, 10e-15], 3e-15)
    ω0 = Maths.moment(grid.ω, abs2.(gab))
    @test ω0[1] < ω0[2]

    # Test sign of dispersion for glass
    Eωglass = Fields.prop_material(Eω, grid, :SiO2, 0.5e-3, λ0)
    Et = FT \ Eωglass
    gab = Maths.gabor(grid.t, Et, [-10e-15, 10e-15], 3e-15)
    ω0 = Maths.moment(grid.ω, abs2.(gab))
    @test ω0[1] < ω0[2]

    # Test sign of dispersion for chirped mirrors
    for mirror in (:PC70, :ThorlabsUMC)
        Eωmirr = Fields.prop_mirror(Eω, grid, mirror, 2) # one pair
        Et = FT \ Eωmirr
        gab = Maths.gabor(grid.t, Et, [-10e-15, 10e-15], 3e-15)
        ω0 = Maths.moment(grid.ω, abs2.(gab))
        @test ω0[1] > ω0[2] # negative chirp, so frequency should go down with time
    end
end

@testset "Compression" begin
# Short pulse with 100 fs^2
λ0 = 800e-9
τfwhm = 10e-15
grid = Grid.RealGrid(1, λ0, (400e-9, 1200e-9), 500e-15)
x = Array{Float64}(undef, length(grid.t))
FT = FFTW.plan_rfft(x, 1)
input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=1e-6)
Eω = input(grid, FT)
Et = FT \ Eω
Eωβ2 = Fields.prop_taylor(Eω, grid, [0, 0, 100e-30], λ0)
ϕs, Eωcomp = Fields.optcomp_taylor(Eωβ2, grid, λ0)
Etcomp = FT \ Eωcomp
@test ϕs[3] ≈ -100e-30
@test isapprox(Maths.fwhm(grid.t, abs2.(Maths.hilbert(Etcomp))), τfwhm; rtol=1e-3)

# Long pulse with 40000 fs^2 (stretches 220 fs to ~5 ps)
λ0 = 1030e-9
τfwhm = 220e-15
grid = Grid.RealGrid(1, λ0, (980e-9, 1080e-9), 20e-12)
x = Array{Float64}(undef, length(grid.t))
FT = FFTW.plan_rfft(x, 1)
input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=1e-6)
Eω = input(grid, FT)
Et = FT \ Eω
Eωβ2 = Fields.prop_taylor(Eω, grid, [0, 0, 4e-25], λ0)
ϕs, Eωcomp = Fields.optcomp_taylor(Eωβ2, grid, λ0)
Etcomp = FT \ Eωcomp
@test ϕs[3] ≈ -4e-25
@test isapprox(Maths.fwhm(grid.t, abs2.(Maths.hilbert(Etcomp))), τfwhm; rtol=1e-3)

# Short pulse with GDD and TOD
λ0 = 800e-9
τfwhm = 10e-15
grid = Grid.RealGrid(1, λ0, (400e-9, 1200e-9), 500e-15)
x = Array{Float64}(undef, length(grid.t))
FT = FFTW.plan_rfft(x, 1)
input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=1e-6)
Eω = input(grid, FT)
Et = FT \ Eω
Eωβ2 = Fields.prop_taylor(Eω, grid, [0, 0, 100e-30, 800e-45], λ0)
ϕs, Eωcomp = Fields.optcomp_taylor(Eωβ2, grid, λ0; order=3)
Etcomp = FT \ Eωcomp
@test all(ϕs .≈ [0, 0, -100e-30, -800e-45])
@test isapprox(Maths.fwhm(grid.t, abs2.(Maths.hilbert(Etcomp))), τfwhm; rtol=1e-3)

# Material insertion
λ0 = 800e-9
τfwhm = 10e-15
grid = Grid.RealGrid(1, λ0, (400e-9, 1200e-9), 500e-15)
x = Array{Float64}(undef, length(grid.t))
FT = FFTW.plan_rfft(x, 1)
input = Fields.GaussField(λ0=λ0, τfwhm=τfwhm, energy=1e-6)
Eω = input(grid, FT)
Et = FT \ Eω
EωFS = Fields.prop_material(Eω, grid, :SiO2, 2e-3, λ0)
d, Eωcomp = Fields.optcomp_material(EωFS, grid, :SiO2, λ0, -1e-2, 1e-2)
Etcomp = FT \ Eωcomp
@test d ≈ -2e-3
@test isapprox(Maths.fwhm(grid.t, abs2.(Maths.hilbert(Etcomp))), τfwhm; rtol=1e-3)
end

@testset "Gaussian beam initialisation" begin
    a = 16e-6
    gas = :Kr
    pres = 17.2
    τfwhm = 230e-15
    λ0 = 1030e-9
    energy = 5.2e-6
    modes = (
        Capillary.MarcatiliMode(a, gas, pres, n=1, m=1, kind=:HE, ϕ=0.0, loss=false),
        Capillary.MarcatiliMode(a, gas, pres, n=1, m=2, kind=:HE, ϕ=0.0, loss=false),
        Capillary.MarcatiliMode(a, gas, pres, n=1, m=3, kind=:HE, ϕ=0.0, loss=false),
        Capillary.MarcatiliMode(a, gas, pres, n=1, m=4, kind=:HE, ϕ=0.0, loss=false),
        Capillary.MarcatiliMode(a, gas, pres, n=2, m=1, kind=:HE, ϕ=0.0, loss=false),
        Capillary.MarcatiliMode(a, gas, pres, n=3, m=1, kind=:HE, ϕ=0.0, loss=false),
        Capillary.MarcatiliMode(a, gas, pres, n=0, m=1, kind=:TE, ϕ=0.0, loss=false),
        Capillary.MarcatiliMode(a, gas, pres, n=0, m=1, kind=:TM, ϕ=0.0, loss=false)
    )
    inputs = Fields.gauss_beam_init(modes, 2π/λ0, a*0.64, Fields.GaussField, λ0=λ0, τfwhm=τfwhm, energy=energy)
    inputs = (inputs..., ((mode=i, fields=(Fields.ShotNoise(),)) for i=1:length(modes))...)
    @test inputs[1].fields[1].energy/energy ≈ 0.9807131210817726
    @test inputs[2].fields[1].energy/energy ≈ 0.006182621678046407
    @test inputs[3].fields[1].energy/energy ≈ 0.0013567813790567626
    @test inputs[4].fields[1].energy/energy ≈ 0.0008447236094573648
    @test inputs[5].fields[1].energy/energy < 1e-20
    @test inputs[6].fields[1].energy/energy < 2e-20
    @test inputs[7].fields[1].energy/energy < 1e-20
    @test inputs[8].fields[1].energy/energy < 1e-20
end

@testset "DataField" begin
    # create a Gaussian pulse with a delay in the frequency domain, write it to a file,
    # load it back in a check it produces the correct pulse
    import DelimitedFiles: writedlm, readdlm
    import FFTW
    λlims = (200e-9, 4e-6)
    trange = 200e-15
    τfwhm = 10e-15
    τ0 = 20e-15
    λ0 = 800e-9
    f0 = PhysData.c/λ0
    σt = Maths.fwhm_to_σ(τfwhm)
    σf = 1/(4π*σt)
    f = collect(range(PhysData.c/λlims[2], PhysData.c/λlims[1]; length=2048))
    If = Maths.gauss.(f, σf; x0=f0)
    ϕ = @. -2π*τ0*(f-f0) # Fourier transform in the maths convention here--pos. delay = neg. slope
    dat = [f If ϕ]
    grid = Grid.RealGrid(1, λ0, λlims, trange)
    FT = FFTW.plan_rfft(copy(grid.t), 1)
    field = mktempdir() do td
        tf = joinpath(td, tempname())
        writedlm(tf, dat, ' ')

        @test readdlm(tf, ' ') == dat

        Fields.DataField(tf; energy=1e-6)
    end
    Eω = field(grid, FT)
    t, Et = Processing.getEt(grid, Eω)
    @test isapprox(Maths.fwhm(t, abs2.(Et)), τfwhm, rtol=1e-5)
    @test isapprox(Maths.moment(t, abs2.(Et)), τ0, rtol=1e-5)
end


module HWunconstrained

	# imports: which packages are we going to use in this module?
	using Distributions, Optim, Plots, DataFrames
	using Random
	using Statistics
	using LinearAlgebra

    export maximize_like_grad, makeData

	# methods/functions
	# -----------------

	# data creator
	# should/could return a dict with beta,numobs,X,y,norm)
	# true coeff vector, number of obs, data matrix X, response vector y, and a type of parametric distribution for G.
	function makeData(n=10_000)

 		# your turn!

		return Dict("beta"=>beta,"n"=>numobs,"X"=>X,"y"=>y,"dist"=>norm)
	end

	# log likelihood function at x
	function loglik(betas::Vector,d::Dict)

	end

	# gradient of the likelihood at x
	function grad!(storage::Vector,betas::Vector,d)

	end

	# hessian of the likelihood at x
	function hessian!(storage::Matrix,betas::Vector,d)

	end

	function info_mat(betas::Vector,d)

	end

	function inv_Info(betas::Vector,d)

	end

	"""
	inverse of observed information matrix
	"""
	function inv_observedInfo(betas::Vector,d)

	end

	"""
	standard errors
	"""
	function se(betas::Vector,d::Dict)

	end

	# function that maximizes the log likelihood without the gradient
	# with a call to `optimize` and returns the result
	function maximize_like(x0=[0.8,1.0,-0.1],meth=NelderMead())

	end
	function maximize_like_helpNM(x0=[ 1; 1.5; -0.5 ],meth=NelderMead())

	end

	# function that maximizes the log likelihood with the gradient
	# with a call to `optimize` and returns the result
	function maximize_like_grad(x0=[0.8,1.0,-0.1],meth=BFGS())

	end

	function maximize_like_grad_hess(x0=[0.8,1.0,-0.1],meth=Newton())

	end

	function maximize_like_grad_se(x0=[0.8,1.0,-0.1],meth=BFGS())

	end

	# visual diagnostics
	# ------------------

	# function that plots the likelihood
	# we are looking for a figure with 3 subplots, where each subplot
	# varies one of the parameters, holding the others fixed at the true value
	# we want to see whether there is a global minimum of the likelihood at the the true value
	function plotLike()

	end
	function plotGrad()

	end

	function saveplots()
		p1 = plotLike()
		savefig(p1,joinpath(dirname(@__FILE__),"..","likelihood.png"))
		p2 = plotGrad()
		savefig(p2,joinpath(dirname(@__FILE__),"..","gradient.png"))
		@info("saved both plots.")
	end

end


function solve_tracer(wf::WaterFlooding; phase::Symbol=:water)
    si = wf.Si
    sj = wf.Sj
    s̃ = wf.S̃

    fw(s) = fractional_flow(s, wf.kr, wf.μw, wf.μo)
    dfw(s) = fw_derivative(s, wf.kr, wf.μw, wf.μo)
    Δfw(s) = (fw(s) - fw(si)) / (s - si)

    # vw(s) = s < s̃ ? Δfw(s̃) : dfw(s)
    vc(s) = fw(s) ./ s

    sc = fzero(s -> vc.(s) - dfw(s), 0.9*sj)

    return Tracer(:water, sc, vc(sc))
end

function solve_tracer(pf::PolymerFlooding; phase::Symbol=:water)
    sb = pf.Sb
    f(s) = fractional_flow(s, pf.kr1, pf.μw, pf.μo)
    vc = f(sb) / sb

    return Tracer(:water, sb, vc)
end

function tracer_profile(wf::WaterFlooding, c::Tracer, time::Real)
    si = wf.Si
    sj = wf.Sj
    s̃ = wf.S̃
    sc = c.sc

    fw(sw) = fractional_flow(sw, wf.kr, wf.μw, wf.μo)
    dfw(sw) = fw_derivative(sw, wf.kr, wf.μw, wf.μo)
    Δfw(sw) = (fw(sw) - fw(si)) / (sw - si)

    s = vcat(collect(sc:0.01:sj), [sj])
    x = Vector{Float64}(undef, length(s))
    x[1] = c.v * time
    x[2:end - 1] = dfw(s[2:end - 1]) * time
    x[end] = 0.0

    return x, s
end

function tracer_profile(pf::PolymerFlooding, c::Tracer, time::Real)
    si = pf.Si
    sj = pf.Sj
    sb = pf.Sb
    ŝ = pf.S̃
    sc = c.sc
    vc = c.v

    df(s) = fw_derivative(s, pf.kr2, pf.μw, pf.μo)

    s = vcat([sc, sb ,sb, ŝ], collect(ŝ:0.01:sj))
    x = zeros(length(s))
    x[1] = vc * time
    x[2:3] .= pf.VΔc * time
    x[4:end] = df(s[4:end]) * time

    return x, s
end


export potential,CorrugatedMorse,calc_VG

using SpecialFunctions

struct CorrugatedMorse{T <: Abstract2DLattice}
    D
    κ
    h
    lattice::T
end

function get_w(z,ch::ScatteringChannels,ki,conf::CorrugatedMorse)
    get_V(z,ch,conf) + Diagonal(ch.kz²)
end

function get_V(z,ch::ScatteringChannels,conf::CorrugatedMorse)
    @unpack gi = ch
    l = length(gi)
    V = Array{Float64,2}(undef,l,l)
    for i in 1:l, j in 1:l
        V[i,j] = get_VG(z,gi[i]-gi[j],conf)
    end
    V
end

function get_VG(z,g,conf::CorrugatedMorse{HexGrid})
    @unpack κ,h,D = conf
    if g[1]==0 && g[2]==0
        return D*(exp(-2*κ*z)+2*exp(-κ*z))
    end

    V0 = calc_VG_factor([0,0],conf,1)
    VG = calc_VG_factor(g,conf,V0)

    VG *= exp(-2*κ*z)*D

    return VG
end

function calc_VG_factor(g,conf::CorrugatedMorse{HexGrid},div)
    @unpack κ,h,D = conf
    @unpack a = conf.lattice

    kmax = 10
    α = 2*κ*h*a/3
    VG = 0.0
    for k in -kmax:kmax
        VG += besseli(k,α)*(besseli(k+g[2],α)*besseli(k+g[1],α) + besseli(k-g[2],α)*besseli(k-g[1],α))
    end
    VG / div
end

function calc_VG(g,l::Generic2DLattice,conf::CorrugatedMorse,z,steps=100)
    VG = 0+0im
    rng = range(0.0,length=steps,1)
    for i in rng , j in rng
        r1 = l.a1.*i
        r2 = l.a2.*j
        r = r1+r2
        VG += potential(conf,r[1],r[2],z)*exp(-1im*(g'*r))
    end
    VG *= 2π/(steps^2)
end

function potential(conf::CorrugatedMorse,x,y,z)
    @unpack D,κ,h,lattice = conf
    @unpack a1,a2 = lattice
    R = [x,y]

    # TODO: only for 60 deg(a1,a2)
    ξ = h*(cos(2π*(a1'*R)/norm(a1)^2)+
           cos(2π*(a2'*R)/norm(a2)^2)+
           cos(2π*((a1-a2)'*R)/norm(a1-a2)^2))
    @. D*(exp(-2*(z-ξ)/κ)-2*exp(-z/κ))

end


@testset "Interpolation" begin
    @testset "T = $(T)" for (T,o,γ) in [(Float64,x -> one(Float64),1.0),(ComplexF64,x -> one(ComplexF64),exp(im*π/4))]
        L = 10.0
        f = x -> γ * sin(2π*x/L)

        N = 100

        rr = range(0, stop=L, length=1000)

        @testset "$(label)" for (label,R,tol) in [
            ("Finite-differences", FiniteDifferences(N, L/(N+1)), 0.01),
            ("Uniform, staggered finite-differences", StaggeredFiniteDifferences(N, L/(N-1/2)), 0.04),
            ("Non-uniform, staggered finite-differences", StaggeredFiniteDifferences(0.1, 0.2, 0.01, L), 0.04),
            ("Implicit finite-differences", ImplicitFiniteDifferences(N, L/(N+1)), 0.01),
            ("FE-DVR", FEDVR(range(0, stop=L, length=N), 4), 5e-6),
            ("B-splines", BSpline(LinearKnotSet(5, 0, L, N)), 5e-9)
        ]
            x = axes(R,1)
            fc = R \ f.(x)

            χ = R[rr,:]
            @test χ*fc ≈ f.(rr) atol=tol
            if !(R isa AbstractFiniteDifferences)
                # We don't bother testing the various
                # finite-differences for interpolations of constants,
                # since they do not yet support non-zero Dirichlet
                # boundary conditions.
                oc = R \ o.(x)
                @test norm(χ*oc .- 1) < 5e-14
            end
        end
    end
end


using MyGraph

G = @buildgraph false 0,1,4 0,7,8 1,2,8 1,7,11 2,3,7 2,8,2 2,5,4 3,4,9 3,5,14 4,5,10 5,6,2 6,8,6 6,7,1 7,8,7
T = mstkruskal(G)


@doc """ multi-class logarithmic loss
""" ->
function mlogloss{T<:AbstractFloat}(pd::Array{T,2}, gt::Array{T,2})
    # T can't be Real since raised an integer x to a negative power -n.
    function max_min(x::T)
        max(min(x,1-10.0^(-15)),10.0^(-15))
    end
    pd = map(max_min, pd)
    log_pd = log(pd)
    -sum(sum(gt .* log_pd, 1))
end


export RBM

struct RBM{T} <: KetNeuralNetwork
    a::Vector{T}
    b::Vector{T}
    W::Matrix{T}
end

"""
    RBMSplit([T=Complex{STD_REAL_PREC}], N, α, [initW, initb])

Constructs a Restricted Bolzmann Machine to encode a wavefunction,
with weights of type `T` (Defaults to ComplexF32), `N` input neurons,
N⋅α hidden neurons.
This is the Neural Quantum State (NQS) Ansatz.

`N` must match the size of the lattice.

The initial parameters of the neurons are initialized with a rescaled normal
distribution of width 0.01 for the coupling matrix and 0.05 for the local
biases. The default initializers can be overriden by specifying

initW=(dims...)->rescaled_normal(T, 0.01, dims...)
initb=(dims...)->rescaled_normal(T, 0.05, dims...)
inita=(dims...)->rescaled_normal(T, 0.01, dims...)

Refs:
    https://arxiv.org/abs/1606.02318
"""
RBM(in, α, args...) = RBM(ComplexF32, in, α, args...)
RBM(T::Type, in, α,
    initW=(dims...)->rescaled_normal(T, 0.01, dims...),
    initb=(dims...)->rescaled_normal(T, 0.05, dims...),
    inita=(dims...)->rescaled_normal(T, 0.01, dims...)) =
    RBM(inita(in), initb(convert(Int,α*in)),
        initW(convert(Int,α*in), in))

input_type(net::RBM{T}) where T = real(T)
weight_type(net::RBM) = out_type(net)
out_type(net::RBM{T}) where T = T
input_shape(net::RBM) = length(net.a)
random_input_state(net::RBM{T}) where T = T.([rand(0:1) for i=1:length(net.a)])
is_analytic(net::RBM) = true

(net::RBM)(σ::State) = net(config(σ))
(net::RBM{T})(σ) where T = transpose(σ)*net.a .+ sum(logℒ.(net.b .+ net.W*σ))

function Base.show(io::IO, m::RBM{T}) where T
    print(io, "RBM{$T}, n=$(length(m.a)), n_hid=$(length(m.b)) => α=$(length(m.b)/length(m.a)))")
end

# Cached version
mutable struct RBMCache{T} <: NNCache{RBM{T}}
    θ::Vector{T}
    logℒθ::Vector{T}
    σ::Vector{T}
    valid::Bool # = false
end

cache(net::RBM{T}) where T =
    RBMCache(Vector{T}(undef,length(net.b)),
             Vector{T}(undef,length(net.b)),
             Vector{T}(undef,length(net.a)),
                  false)

function (net::RBM{T})(c::RBMCache, σ) where T
    θ = c.θ
    logℒθ = c.logℒθ

    #θ .= net.b .+ net.W * σ
    copyto!(c.σ, σ)
    copyto!(θ, net.b)
    BLAS.gemv!('N', T(1.0), net.W, c.σ, T(1.0), θ)

    logℒθ .= logℒ.(θ)
    logψ = dot(σ,net.a) + sum(logℒθ)
    return logψ
end

function logψ_and_∇logψ!(∇logψ, net::RBM{T}, c::RBMCache, σ) where T
    θ = c.θ
    logℒθ = c.logℒθ

    #θ .= net.b .+ net.W * σ
    copyto!(c.σ, σ)
    copyto!(θ, net.b)
    BLAS.gemv!('N', T(1.0), net.W, c.σ, T(1.0), θ)

    logℒθ .= logℒ.(θ)
    logψ = dot(σ,net.a) + sum(logℒθ)

    ∇logψ.a .= σ
    ∇logψ.b .= ∂logℒ.(θ)
    ∇logψ.W .= ∇logψ.b  .* transpose(σ)
    return logψ
end


function createnodes(p::Array{Float64,2},t::Array{Int,2},porder::Int)
# %CREATEDGNODES Computes the Coordinates of the DG nodes.
# %   DGNODES=CREATENODES(MESH,FD,FPARAMS)
# %
# %      MESH:      Mesh Data Structure
# %      FD:        Distance Function d(x,y)
# %      FPARAMS:   Additional parameters passed to FD
# %      DGNODES:   Triangle indices (NPL,2,NT). The nodes on
# %                 the curved boundaries are projected to the
# %                 true boundary using the distance function FD
# %

# % npv : number of nodes per volume element
# % nfv : number of faces per volume element
# % npf : number of nodes per face element

nve,ne = size(t);
nd=size(p,1);

elemtype = 0;
if (nd==2) && (nve==4)
    elemtype=1;
end
if (nd==3) && (nve==8)
    elemtype=1;
end

plocal,~,~,~,perm = masternodes(porder,nd,elemtype);
npl = size(plocal,1);

philocal = zeros(Float64, npl, nve);

if nd==1
    xi  = plocal[:,1];
    philocal[:,1] = 1.0 .- xi;
    philocal[:,2] = xi;
elseif nd==2 && nve==3 # tri
    xi  = plocal[:,1];
    eta = plocal[:,2];
    philocal[:,1] = 1.0 .- xi .- eta;
    philocal[:,2] = xi;
    philocal[:,3] = eta;
elseif nd==2 && nve==4 # quad
    xi  = plocal[:,1];
    eta = plocal[:,2];
    philocal[:,1] = (1.0 .- xi).*(1.0 .- eta);
    philocal[:,2] = xi.*(1.0 .- eta);
    philocal[:,3] = xi.*eta;
    philocal[:,4] = (1.0 .- xi).*eta;
elseif nd==3 && npv==4 # tet
    xi   = plocal[:,1];
    eta  = plocal[:,2];
    zeta = plocal[:,3];
    philocal[:,1] = 1.0 - xi - eta - zeta;
    philocal[:,2] = xi;
    philocal[:,3] = eta;
    philocal[:,4] = zeta;
elseif nd==3 && npv==8 # hex
    xi   = plocal[:,1];
    eta  = plocal[:,2];
    zeta = plocal[:,3];
    philocal[:,1] = (1.0 .-xi).*(1.0 .-eta).*(1.0 .-zeta);
    philocal[:,2] = xi.*(1.0 .-eta).*(1.0 .-zeta);
    philocal[:,3] = xi.*eta.*(1.0 .-zeta);
    philocal[:,4] = (1.0 .-xi).*eta.*(1.0 .-zeta);
    philocal[:,5] = (1.0 .-xi).*(1.0 .-eta).*(zeta);
    philocal[:,6] = xi.*(1.0 .-eta).*(zeta);
    philocal[:,7] = xi.*eta.*(zeta);
    philocal[:,8] = (1.0 .-xi).*eta.*(zeta);
end

# Allocate nodes
dgnodes = zeros(Float64,npl,ne,nd);
for dim=1:nd
  for node=1:nve
    dp = reshape(philocal[:,node],npl,1)*reshape(p[dim,t[node,:]],1,ne);
    dgnodes[:,:,dim] = dgnodes[:,:,dim] + reshape(dp,npl,ne);
  end
end

dgnodes = permutedims(dgnodes,[1, 3, 2]);

return dgnodes, elemtype, perm

end


export domaintype,
    codomaintype,
    coefficienttype,
    prectype

"""
A `Dictionary{S,T}` is an ordered family of functions, in which each function
maps a variable of type `S` to a variable of type `T`. The dictionary can be
thought of as an array, where the elements are functions that map `S` to `T`.

A `Dictionary{S,T}` has domain type `S` and codomain type `T`. The domain type
corresponds to the type of a domain in the `DomainSets.jl` package, and it is the
type of the expected argument to the elements of the dictionary. The
codomain type is the type of the output.

Each dictionary is ordered via its index set: the ordering is determined by the
iterator of the index set. A dictionary `d` can be indexed in several ways:
- the linear index is a positive natural number between `1` and `length(d)`
- the native index is an index that more closely corresponds to the conventional
  mathematical notation of the dictionary, if that differs from linear indexing.
  For example, polynomials may have degree ranging from `0` to `length(d)-1`.
  Fourier series may have negative frequencies.

Some dictionaries define other types of indices, such as multilinear and product
indices. See `CompositeDictionary` and `ProductDictionary`. There is always a
bijection between the different types of index sets supported by a dictionary.

A dictionary is the basic building block of the package. They can support several
operations, such as interpolation, evaluation, differentiation and others. This
functionality is made available using a specific interface. See the documentation
of the features for a description of that interface and the syntax.
"""
abstract type Dictionary{S,T}
end

# Useful abstraction for special cases
const Dictionary1d{S <: Number,T} = Dictionary{S,T}
# Warning: these are shaky definitions of multivariate functions
const Dictionary2d{S <: Number,T} = Union{Dictionary{NTuple{2,S},T},Dictionary{SVector{2,S},T}}
const Dictionary3d{S <: Number,T} = Union{Dictionary{NTuple{3,S},T},Dictionary{SVector{3,S},T}}
const Dictionary4d{S <: Number,T} = Union{Dictionary{NTuple{4,S},T},Dictionary{SVector{4,S},T}}

"The type of the elements of the domain of the dictionary."
domaintype(::Type{<:Dictionary{S,T}}) where {S,T} = S
domaintype(dict::Dictionary) = domaintype(typeof(dict))

"The type of the elements of the codomain of the dictionary."
codomaintype(::Type{<:Dictionary{S,T}}) where {S,T} = T
codomaintype(dict::Dictionary) = codomaintype(typeof(dict))

"The type of the expansion coefficients in a dictionary."
# By default we set it equal to the codomaintype
coefficienttype(D::Type{<:Dictionary{S,T}}) where {S,T} = codomaintype(D)
coefficienttype(dict::Dictionary) = coefficienttype(typeof(dict))

prectype(D::Type{<:Dictionary}) = prectype(domaintype(D), codomaintype(D))
numtype(D::Type{<:Dictionary}) = numtype(domaintype(D), codomaintype(D))

# The dimension of a function set is the dimension of its domain type
dimension(dict::Dictionary) = dimension(domaintype(dict))

dimension(dict::Dictionary, i) = dimension(component(dict, i))

"Are the functions in the dictionary real-valued?"
isreal(d::Dictionary) = isreal(codomaintype(d))

"Is the dictionary a (truncation of a) basis?"
isbasis(d::Dictionary) = false

"Is the dictionary orthogonal (with respect to the given measure)?"
isorthogonal(d::Dictionary) = hasmeasure(d) && isorthogonal(d, measure(d))
isorthogonal(d::Dictionary, μ::Measure) = isorthonormal(d, μ)

"Is the dictionary orthonormal (with respect to the given measure)?"
isorthonormal(d::Dictionary) = hasmeasure(d) && isorthonormal(d, measure(d))
isorthonormal(d::Dictionary, μ::Measure) = false

isorthogonal(d::Dictionary, grid::AbstractGrid) = isorthogonal(d, discretemeasure(grid))
isorthonormal(d::Dictionary, grid::AbstractGrid) = isorthonormal(d, discretemeasure(grid))

"Is the dictionary biorthogonal (with respect to the given measure)?"
isbiorthogonal(d::Dictionary) = hasmeasure(d) && isbiorthogonal(d, measure(d))
isbiorthogonal(d::Dictionary, μ::Measure) = isorthogonal(d, μ)

size(d::Dictionary, j) = size(d)[j]

"""
The length and size of a dictionary may not be enough to uniquely determine the
size of the dictionary, if it has additional internal structure. The output of
dimensions is such that `resize(dict, dimensions(dict)) == dict`.
"""
dimensions(d::Dictionary1d) = length(d)
dimensions(d::Dictionary) = size(d)

length(d::Dictionary) = prod(size(d))

firstindex(d::Dictionary) = first(eachindex(d))
lastindex(d::Dictionary) = last(eachindex(d))

#############################
# Domain and codomain type
#############################

similar(s::Dictionary, ::Type{T}) where {T} = similar(s, T, size(s))

similar(s::Dictionary, size::Int...) = similar(s, domaintype(s), size...)
similar(s::Dictionary, dims::Base.Dims) = similar(s, domaintype(s), dims...)

similar(s::Dictionary, ::Type{T}, dims::Base.Dims) where {T} = similar(s, T, dims...)

# This is a default routine that ony appies when nothing changes.
function similar(dict::Dictionary{S,T}, ::Type{S}, n::Int) where {S,T}
    @assert n == length(dict)
    dict
end

resize(s::Dictionary, dims...) = similar(s, domaintype(s), dims...)

widen(d::Dictionary) = similar(d, widen(domaintype(d)))

"Return a set of zero coefficients in the native format of the set."
zeros(s::Dictionary) = zeros(coefficienttype(s), s)
ones(s::Dictionary) = ones(coefficienttype(s), s)

# By default we assume that the native format corresponds to an array of the
# same size as the set. This is not true, e.g., for multidicts.
zeros(::Type{T}, s::Dictionary) where {T} = zeros(T, size(s))
ones(::Type{T}, s::Dictionary) where {T} = ones(T, size(s))

"What is the type of the coefficient vector of the dictionary?"
containertype(d::Dictionary) = typeof(zeros(d))

"Transform `a` to a coefficient vector for the given dictionary."
tocoefficientformat(a, d::Dictionary) = reshape(a, size(d))

function rand(dict::Dictionary)
    c = zeros(dict)
    T = coefficienttype(dict)
    for i in eachindex(c)
        c[i] = rand(T)
    end
    c
end

###########
# Indexing
###########

IndexStyle(d::Dictionary) = IndexLinear()

"""
Dictionaries are ordered lists. Their ordering is defined by the way their
index sets are ordered.

The `ordering` of a dictionary returns a list-like object that can be indexed
with integers between `1` and `length(dict)`. This operation returns the
corresponding native index. This defines the ordering of the native index set
of the dictionary.
"""
ordering(dict::Dictionary) = Base.OneTo(length(dict))

eachindex(d::Dictionary) = eachindex(IndexStyle(d), d)
eachindex(::IndexLinear, d::Dictionary) = axes1(d)
eachindex(::IndexCartesian, d::Dictionary) = CartesianIndices(axes(d))

axes1(d::Dictionary) = Base.OneTo(length(d))
axes(d::Dictionary) = map(Base.OneTo, size(d))

"Compute the native index corresponding to the given index."
native_index(dict::Dictionary, idx) = _native_index(dict, idx)
# We redirect to a fallback _native_index in case the concrete dictionary
# did not implement native_index.  We explicitly convert a linear index using the ordering.
# Anything else we throw an error because the index looks invalid
_native_index(dict::Dictionary, idx::Int) = ordering(dict)[idx]
_native_index(dict::Dictionary, idx::NativeIndex) = idx
_native_index(dict::Dictionary, idx::AbstractShiftedIndex) = idx
_native_index(dict::Dictionary, idx) = throw(ArgumentError("invalid index: $idx"))

"Compute the linear index corresponding to the given index."
linear_index(dict::Dictionary, idx) = _linear_index(dict, idx)
# We can accept an integer unchanged, anything else we pass to the ordering
_linear_index(dict::Dictionary, idx::Int) = idx
_linear_index(dict::Dictionary, idxn) = ordering(dict)[idxn]

##################################################
## Conversion between coefficient representations
##################################################

"""
Convert the set of coefficients in the native format of the dictionary
to a linear list in a vector.
"""
function linearize_coefficients(dict::Dictionary, coef_native)
    coef_linear = zeros(eltype(coef_native), length(dict))
    linearize_coefficients!(dict, coef_linear, coef_native)
end

linearize_coefficients!(dict::Dictionary, coef_linear::Vector, coef_native) =
    copyto!(coef_linear, coef_native)
# Note that copyto! is defined in util/common.jl

"""
Convert a linear set of coefficients back to the native representation of the dictionary.
"""
function delinearize_coefficients(dict::Dictionary, coef_linear::Vector)
    coef_native = zeros(eltype(coef_linear), dict)
    delinearize_coefficients!(dict, coef_native, coef_linear)
end

delinearize_coefficients!(dict::Dictionary, coef_native, coef_linear::Vector) =
    copyto!(coef_native, coef_linear)

"Promote the given coefficients to the native representation of the dictionary."
native_coefficients(dict::Dictionary, coef) = _native_coefficients(dict, coef)
# TODO: we create an unnecessary copy here if the native type is a vector
_native_coefficients(dict::Dictionary, coef::Vector) = delinearize_coefficients(dict, coef)
_native_coefficients(dict::Dictionary, coef) = coef

# Sets have a native size and a linear size. However, there is not necessarily a
# bijection between the two. You can always convert a native size to a linear size,
# but the other direction can be done in general only approximately.
# For example, a 2D tensor product set can only support sizes of the form n1 * n2. Its native size may be
# (n1,n2) and its linear size n1*n2, but not any integer n maps to a native size tuple.
# By convention, we denote a native size variable by size_n.
"Compute the native size best corresponding to the given linear size."
approximate_native_size(d::Dictionary, size_l) = size_l

"Compute the linear size corresponding to the given native size."
linear_size(d::Dictionary, size_n) = size_n

"Suggest a suitable size, close to `n`, to resize the given dictionary."
approx_length(d::Dictionary, n::Int) = n
approx_length(d::Dictionary, n::Real) = approx_length(d, round(Int,n))

###############################
## Properties of function sets
###############################

"""
Does the dictionary implement a differentiation operator?
An optional second argument may specify an exact order of the derivative.
"""
hasderivative(Φ::Dictionary) = false

function hasderivative(Φ::Dictionary, order)
    # We have to be correct for zero order derivatives, because even if the
    # dictionary itself does not support derivatives, it could be part of a
    # composition where other elements do
    if orderiszero(order)
        true
    else
        order == 1 ? hasderivative(Φ) : false
    end
end

"Does the dictionary implement an antiderivative?"
hasantiderivative(Φ::Dictionary) = false

function hasantiderivative(Φ::Dictionary, order)
    if orderiszero(order)
        true
    else
        order == 1 ? hasantiderivative(Φ) : false
    end
end

"Does the dictionary have an associated interpolation grid?"
hasinterpolationgrid(d::Dictionary) = false

"Does the dictionary have a transform associated with some space?"
hastransform(d1::Dictionary, d2) = false

# Convenience functions: default grid, and conversion from grid to space
hastransform(d::Dictionary) = hasinterpolationgrid(d) && hastransform(d, interpolation_grid(d))
hastransform(d::Dictionary, grid::AbstractGrid) =
    hastransform(d, GridBasis{coefficienttype(d)}(grid))

"Does the dictionary support extension and restriction operators?"
hasextension(d::Dictionary) = false

# A concrete Dictionary may also override extension_set and restriction_set
# The default is simply to resize.
extension_set(d::Dictionary, n) = resize(d, n)
restriction_set(d::Dictionary, n) = resize(d, n)

###############################
## Iterating over dictionaries
###############################

# Default iterator over sets of functions: based on underlying index iterator.
function iterate(d::Dictionary)
    iter = eachindex(d)
    first_item, first_state = iterate(iter)
    (d[first_item], (iter, (first_item, first_state)))
end

function iterate(d::Dictionary, state)
    iter, iter_tuple = state
    iter_item, iter_state = iter_tuple
    next_tuple = iterate(iter, iter_state)
    if next_tuple != nothing
        next_item, next_state = next_tuple
        (d[next_item], (iter,next_tuple))
    end
end

include("dict_evaluation.jl")
include("dict_moments.jl")


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Geometry{Dim,T}

A geometry embedded in a `Dim`-dimensional space with coordinates of type `T`.
"""
abstract type Geometry{Dim,T} end

"""
    embeddim(geometry)

Return the number of dimensions of the space where the `geometry` is embedded.
"""
embeddim(::Type{<:Geometry{Dim,T}}) where {Dim,T} = Dim
embeddim(g::Geometry) = embeddim(typeof(g))

"""
    paramdim(geometry)

Return the number of parametric dimensions of the `geometry`. For example, a
sphere embedded in 3D has 2 parametric dimension (polar and azimuthal angles).
"""
paramdim(g::Geometry) = paramdim(typeof(g))

"""
    coordtype(geometry)

Return the machine type of each coordinate used to describe the `geometry`.
"""
coordtype(::Type{<:Geometry{Dim,T}}) where {Dim,T} = T
coordtype(g::Geometry) = coordtype(typeof(g))

"""
    measure(geometry)

Return the measure of the `geometry`, i.e. the length, area, or volume.
"""
function measure end

"""
    boundary(geometry)

Return the boundary of the `geometry`.
"""
function boundary end

# ----------------
# IMPLEMENTATIONS
# ----------------
include("primitives.jl")
include("polytopes.jl")

"""
    GeometrySet(geometries)

A set of `geometries` seen as a single geometry.

## Example

In a geographic map, countries can be described with
multiple polygons (a.k.a. MultiPolygon).
"""
struct GeometrySet{Dim,T,G<:Geometry{Dim,T}} <: Geometry{Dim,T}
  geometries::Vector{G}
end


#using GLOQ
#x = rand(Float64,8)
#X = [0 x[1]+x[2]*im x[3]+x[4]*im;
#     0 x[5]         x[6]+x[7]*im;
#     0 0            x[8]]
#H = 0.5*(X+adjoint(X))
HK = [ 0.0 0.0 1.0;
       0.0 0.0 0.0;
       1.0 0.0 0.0]
HS = [ 0.0 1.0 0.0;
      -1.0 0.0 0.0;
       0.0 0.0 0.0]
H = HK + 1.0im.*HS
L1 = [0.0 0.25 0.0;
      0.0 0.0  0.5;
      0.0 0.0  0.0]
L2 = [0.0 0.0 0.0;
      0.0 0.5 0.0;
      0.0 0.0 1.0]
L_list = (L1,L2)
LH,LD = GLOQ.make_lindblad_operator_complex(H,L_list)
LK,LS,LD2 = GLOQ.make_lindblad_operator(HK,HS,L_list)
L = LH+LD
println("Test whether the real-valued and complex-valued operator match each other")
@test (LD==LD2)&&( LH==-1.0im*(LK+1.0im*LS) )

println("\n Lindblad operator assembled\n")

tspan_length = 21
t_span = range(0.0,pi,length=tspan_length)
t_span2 = zeros(Float64,tspan_length)
dt = pi/(tspan_length-1);
for i = 1:tspan_length
      t_span2[i] = (i-1)*dt
end
u1 = [1.0;0;0]
rho1_history = GLOQ.exponential_solver_complex(u1,L,t_span;initial_type = "states")
u2 = [0;1.0;0]
rho2 = (u2*u2')[:]
rho2_history = GLOQ.exponential_solver_complex(rho2,L,t_span2;initial_type = "density")
u3 = [1.0/sqrt(2.0);-im/sqrt(2.0);0]
rho3_history = GLOQ.exponential_solver_complex(u3,L,t_span;initial_type = "states")
u4 = [0.0;1.0/sqrt(2.0);-im/sqrt(2.0)]
rho4_history = GLOQ.exponential_solver_complex(u4,L,t_span;initial_type = "states")
u5 = [1.0/sqrt(2.0);0.0;-im/sqrt(2.0)]
rho5_history = GLOQ.exponential_solver_complex(u5,L,t_span;initial_type = "states")

population1_history = GLOQ.get_population(rho1_history)
population2_history = GLOQ.get_population(rho2_history)
population3_history = GLOQ.get_population(rho3_history)
population4_history = GLOQ.get_population(rho4_history)
population5_history = GLOQ.get_population(rho5_history)

#################################################
# test real solvers
rho1_history_u,rho1_history_v = GLOQ.exponential_solver(real(u1),-imag(u1),LK,LS,LD2,t_span;initial_type = "states")
rho2_u,rho2_v = GLOQ.convert_state_to_density(real(u2),-imag(u2))
rho2_history_u,rho2_history_v = GLOQ.exponential_solver(rho2_u,rho2_v,LK,LS,LD2,t_span;initial_type = "density")
rho3_history_u,rho3_history_v = GLOQ.exponential_solver(real(u3),-imag(u3),-imag(LH),real(LH),LD,t_span;initial_type = "states")
rho4_u,rho4_v = GLOQ.convert_state_to_density(real(u4),-imag(u4))
rho4_history_u,rho4_history_v = GLOQ.exponential_solver(rho4_u,rho4_v,LK,LS,LD2,t_span;initial_type = "density")
rho5_history_u,rho5_history_v = GLOQ.exponential_solver(real(u5),-imag(u5),LK,LS,LD2,t_span;initial_type = "states")

@test ( (norm(rho1_history_u-1.0im*rho1_history_v-rho1_history,Inf)<1e-14)&&
        (norm(rho2_history_u-1.0im*rho2_history_v-rho2_history,Inf)<1e-14)&&
        (norm(rho3_history_u-1.0im*rho3_history_v-rho3_history,Inf)<1e-14)&&
        (norm(rho4_history_u-1.0im*rho4_history_v-rho4_history,Inf)<1e-14)&&
        (norm(rho5_history_u-1.0im*rho5_history_v-rho5_history,Inf)<1e-14) )

population1_history2 = GLOQ.get_population(rho1_history_u)
population2_history2 = GLOQ.get_population(rho2_history_u)
population3_history2 = GLOQ.get_population(rho3_history_u)
population4_history2 = GLOQ.get_population(rho4_history_u)
population5_history2 = GLOQ.get_population(rho5_history_u)

@test ( (norm(population1_history-population1_history2,Inf)<1e-14)&&
        (norm(population2_history-population2_history2,Inf)<1e-14)&&
        (norm(population3_history-population3_history2,Inf)<1e-14)&&
        (norm(population4_history-population4_history2,Inf)<1e-14)&&
        (norm(population5_history-population5_history2,Inf)<1e-14) )

# check the interface with DifferentialEquations

rho3 = (u3*u3')[:]
problem3 = GLOQ.LindbladODEProblemComplex(rho3,L,t_span[end];initial_type="density")
rho3_history_diffeq = solve(problem3,saveat=t_span)
difference = maximum(abs.(rho3_history-rho3_history_diffeq))
@test(difference<5e-7)

N_density = length(u3)^2;
problem3_real = GLOQ.LindbladODEProblem(real(u3),-imag(u3),LK,LS,LD,t_span[end];initial_type="states")
rho3_history_diffeq_real = solve(problem3_real,saveat=t_span)

println(maximum(abs.(real(rho3_history)-rho3_history_diffeq_real[1:N_density,:])))
println(maximum(abs.(imag(rho3_history)+rho3_history_diffeq_real[N_density+1:end,:])))
@test( (maximum(abs.(real(rho3_history)-rho3_history_diffeq_real[1:N_density,:]))<5e-5)&&
       (maximum(abs.(imag(rho3_history)+rho3_history_diffeq_real[N_density+1:end,:]))<5e-5) )


import FiniteDiff
import ForwardDiff
import ReverseDiff
# import Zygote

using Test

grads = let
    F(x, θ, deg) = (θ[1] - x[1])^deg + θ[2] * (x[2] - x[1]^deg)^deg

    F_idx_val(ca) = F(ca[Val(:x)], ca[Val(:θ)], ca[Val(:deg)])
    F_idx_sym(ca) = F(ca[:x], ca[:θ], ca[:deg])
    F_prop(ca) = F(ca.x, ca.θ, ca.deg)

    ca = ComponentArray(x = [1, 2], θ = [1.0, 100.0], deg = 2)

    (
        truth = [-400, 200],
        finite = FiniteDiff.finite_difference_gradient(ca -> F_prop(ca), ca).x,
        forward = ForwardDiff.gradient(ca -> F_prop(ca), ca).x,
        reverse = ReverseDiff.gradient(ca -> F_prop(ca), ca).x,
        # zygote = Zygote.gradient(ca -> F(ca), ca)[1].x,
    )
end

@test grads.finite ≈ grads.truth
@test grads.forward ≈ grads.truth
@test grads.reverse ≈ grads.truth
# @test grads.zygote ≈ grads.truth


##Operators

for op in (:Derivative,:Integral)
    @eval begin
        function ($op){T<:IntervalDomain}(d::Vector{T})
            n=length(d)
            R=zeros(Operator{mapreduce(eltype,promote_type,d)},n,n)
            for k=1:n
                R[k,k]=$op(d[k])
            end

            R
        end
    end
end

function Evaluation{T<:IntervalDomain}(d::Vector{T},x...)
    n=length(d)
    R=zeros(Operator{mapreduce(eltype,promote_type,d)},n,n)
    for k=1:n
        R[k,k]=Evaluation(d[k],x...)
    end

    R
end

## Construction

function Base.diagm{T<:Operator}(d::Vector{T})
    D=zeros(Operator{mapreduce(eltype,promote_type,d)},length(d),length(d))
    for k=1:length(d)
        D[k,k]=d[k]
    end
    D
end

##TODO: unify with other blkdiag
function Base.blkdiag{T<:Operator}(d1::Vector{T},d2::Vector{T})
  if isempty(d1)&&isempty(d2)
    error("Empty blkdiag")
  end
  if isempty(d1)
    TT=mapreduce(eltype,promote_type,d2)
  elseif isempty(d2)
    TT=mapreduce(eltype,promote_type,d1)
  else
     TT=promote_type(mapreduce(eltype,promote_type,d1),
                    mapreduce(eltype,promote_type,d2))
  end

  D=zeros(Operator{TT},length(d1)+length(d2),2)
  D[1:length(d1),1]=d1
  D[length(d1)+1:end,2]=d2
  D
end

Base.blkdiag(a::Operator,b::Operator) = blkdiag(Operator{promote_type(eltype(a),eltype(b))}[a],
                                                Operator{promote_type(eltype(a),eltype(b))}[b])

## broadcase

.*{N<:Number}(A::Array{N},D::Operator)=Operator{promote_type(N,eltype(D))}[A[k,j]*D for k=1:size(A,1),j=1:size(A,2)]
.*{N<:Number}(D::Operator,A::Array{N})=A.*D


export elasticqrfact!, elasticqrfact, updateelasticqrfact!

# We cannot define subtypes of concrete types in Julia
# struct ElasticQR{T, S<:AbstractMatrix{T}} <: Factorization{T}
#     factors::ElasticArray{T,2,1,Array{T,1}}
#     τ::Array{T,1}
#     # cache::Array{Float64,1}
# end
# ElasticQR(factors::AbstractMatrix{T}, τ::AbstractVector{T}) where {T} = ElasticQR{T}(ElasticMatrix(factors), τ)

function elasticqrfact!(A::AbstractMatrix{T}) where {T}
    LinearAlgebra.require_one_based_indexing(A)
    m, n = size(A)
    τ = zeros(T, min(m,n))
    @inbounds for k = 1:min(m - 1 + !(T<:Real), n)
        x = view(A, k:m, k)
        τk = LinearAlgebra.reflector!(x)
        τ[k] = τk
        LinearAlgebra.reflectorApply!(x, τk, view(A, k:m, k + 1:n))
    end
    QR(ElasticMatrix(A), τ)
end

function elasticqrfact(A::AbstractMatrix{T}, arg...; kwargs...) where T
    LinearAlgebra.require_one_based_indexing(A)
    AA = similar(A, LinearAlgebra._qreltype(T), size(A))
    copyto!(AA, A)
    return elasticqrfact!(AA, arg...; kwargs...)
end

function updateelasticqrfact!(S::QR{T,ElasticArray{T,2,1,Array{T,1}}}, a::AbstractVector{T}, cache::AbstractVector{T}) where {T}
    m, n = size(S.factors)
    mul!(cache, S.Q', a)
    @inbounds append!(S.factors, cache)
    x = view(S.factors, n+1:m, n+1)
    τk = LinearAlgebra.reflector!(x)
    @inbounds push!(S.τ, τk)
    nothing
end

updateelasticqrfact!(S::QR{T,ElasticArray{T,2,1,Array{T,1}}}, a::AbstractVector{T}) where {T} = updateelasticqrfact!(S, a, zero(a))

function updateelasticqrfact!(S::QR{T,ElasticArray{T,2,1,Array{T,1}}}, A::AbstractMatrix{T}, cache) where {T}
    m, n = size(S.factors)
    mA, nA = size(A)
    @assert mA == m "Length of the columns do not match"

    @inbounds mul!(cache, S.Q', A)
    @inbounds append!(S.factors, cache)
    # @for k=n+1:n+nA
    #     append!(S.factors, S.Q'

    @inbounds for k = n+1:min(m - 1 + !(T<:Real), n+nA)
        x = view(S.factors, k:m, k)
        τk = LinearAlgebra.reflector!(x)
        push!(S.τ, τk)
        LinearAlgebra.reflectorApply!(x, τk, view(S.factors, k:m, k + 1:n + nA))
    end
end

updateelasticqrfact!(S::QR{T,ElasticArray{T,2,1,Array{T,1}}}, A::AbstractMatrix{T}) where {T} = updateelasticqrfact!(S, A, zero(A))


module Perceptron

using MetadataTools,  DocStringExtensions

export PerceptronClassifier, fit!, predict

type PerceptronClassifier{T}
    W::AbstractMatrix{T}
    b::AbstractVector{T}
    n_classes::Int
    n_features::Int
end

function Base.show{T}(io::IO, p::PerceptronClassifier{T})
    n_classes  = p.n_classes
    n_features = p.n_features
    print(io, "PerceptronClassifier{$T}(n_classes=$n_classes, n_features=$n_features)")
end

PerceptronClassifier(T::Type, n_classes::Int, n_features::Int) = PerceptronClassifier{T}(rand(T, n_features, n_classes),
                                                                                zeros(T, n_classes),
                                                                                n_classes,
                                                                                n_features)

"""
$(SIGNATURES)

Compute the accuracy betwwen `y` and `y_hat`.
"""
function accuracy(y::AbstractVector, y_hat::AbstractVector)
    acc = 0.
    @fastmath for k = 1:length(y)
            @inbounds  acc += y[k] == y_hat[k]
    end
    return acc/length(y_hat)
end

"""
$(SIGNATURES)

Predicts the class for a given input in a `PerceptronClassifier`.
The placeholder is used to avoid allocating memory for each matrix-vector multiplication.

- Returns the predicted class.
"""
function predict(h::PerceptronClassifier, x::AbstractVector, class_placeholder::AbstractVector)
    @fastmath class_placeholder .= At_mul_B!(class_placeholder, h.W, x) .+ h.b
    return indmax(class_placeholder)
end

"""
$(SIGNATURES)

Function to predict the class for a given example.

- Returns the predicted class.
"""
function predict(h::PerceptronClassifier, x::AbstractVector)
    score = h.W' * x .+ h.b
    return @fastmath indmax(score), score
end

"""
$(SIGNATURES)

Function to predict the class for a given input batch.
- Returns the predicted class.
"""
function predict(h::PerceptronClassifier, X::AbstractMatrix)
    predictions = zeros(Int32, size(X, 2))
    class_placeholder = zeros(eltype(h.W), h.n_classes)

    @inbounds for m in 1:length(predictions)
        predictions[m] = predict(h, view(X,:,m), class_placeholder)
    end
    return predictions
end

"""
$(SIGNATURES)

>    fit!(h::PerceptronClassifier,
>         X::Array,
>         y::Array;
>         n_epochs=50,
>         learning_rate=0.1,
>         print_flag=false,
>         compute_accuracy=true,
>         seed=srand(1234),
>         pocket=false,
>         shuffle_data=false)

##### Arguments

- **`h`**, (PerceptronClassifier{T} type), Multiclass perceptron.
- **`X`**, (Array{T,2} type), data contained in the columns of X.
- **`y`**, (Vector{T} type), class labels (as integers from 1 to n_classes).

##### Keyword arguments

- **`n_epochs`**, (Integer type), number of passes (epochs) through the data.
- **`learning_rate`**, (Float type), learning rate (The standard perceptron is with learning_rate=1.)
- **`compute_accuracy`**, (Bool type), if `true` the accuracy is computed at the end of every epoch.
- **`print_flag`**, (Bool type), if `true` the accuracy is printed at the end of every epoch.
- **`seed`**, (MersenneTwister type), seed for the permutation of the datapoints in case there the data is shuffled.
- **`pocket`** , (Bool type), if `true` the best weights are saved (in the pocket) during learning.
- **`shuffle_data`**, (Bool type),  if `true` the data is shuffled at every epoch (in reality we only shuffle indicies for performance).

"""
function fit!(h::PerceptronClassifier, X::AbstractArray, y::AbstractVector, scores::Array;
              n_epochs=50, learning_rate=1., print_flag=false,
              compute_accuracy=false, seed=srand(1234), pocket=false,
              shuffle_data=false)

    n_features, n_samples = size(X)
    @assert length(y) == n_samples

    T = eltype(X)
    learning_rate = T(learning_rate)

    class_placeholder = zeros(T, h.n_classes)
    y_preds = zeros(Int64, n_samples)
    data_indices = Array(1:n_samples)
    max_acc = zero(T)

    if pocket
        W_hat = zeros(T, h.n_features, h.n_classes)
        b_hat = zeros(T, h.n_classes)
    end

    @fastmath for epoch in 1:n_epochs

        n_mistakes = 0
        if shuffle_data
            shuffle!(seed, data_indices)
        end

        @inbounds for m in data_indices
            x = view(X, :, m)

            y_hat = predict(h, x, class_placeholder)

            if y[m] != y_hat
                n_mistakes += 1
                ####  wij ← wij − η (yj −tj) · xi
                h.W[:, y[m]]  .= h.W[:, y[m]]  .+ learning_rate .* x
                h.b[y[m]]     .= h.b[y[m]]     .+ learning_rate
                h.W[:, y_hat] .= h.W[:, y_hat] .- learning_rate .* x
                h.b[y_hat]    .= h.b[y_hat]    .- learning_rate
            end
        end

        if compute_accuracy
             @inbounds for m in  data_indices
                 y_preds[m] = predict(h, view(X, :, m), class_placeholder)
            end
            acc = accuracy(y, y_preds)
            push!(scores, acc)
        else
            acc = (n_samples - n_mistakes)/n_samples
            push!(scores, acc)
        end

        if pocket
            if acc > max_acc
                max_acc = acc
                copy!(W_hat, h.W)
                copy!(b_hat, h.b)
            end
        end

        if print_flag
            print("Epoch: $(epoch) \t Accuracy: $(round(acc,3))\r")
            flush(STDOUT)
        end
    end

end

end # module


using MultivariateFunctions
using DataFrames
using Dates
using Random
using GLM

tol = 10*eps()

Random.seed!(1)
obs = 1000
X = rand(obs)
y = X .+ rand(Normal(),obs) .+ 7

# Basic use case with 2 degrees
lm1 = fit(LinearModel, hcat(ones(obs), X, X .^ 2), y)
glm_preds = predict(lm1, hcat(ones(obs), X, X .^ 2))
package_approximation, reg1 = create_ols_approximation(y, X, 2; intercept = true, base_x = 0.0)
package_predictions = evaluate.(package_approximation, X)
sum(abs.(glm_preds .- package_predictions)) < 1e-10

# Degree of 1 with no intercept
lm1 = fit(LinearModel, hcat(X), y)
glm_preds = predict(lm1, hcat(X))
package_approximation, reg2 = create_ols_approximation(y, X, 1; intercept = false)
package_predictions = evaluate.(package_approximation, X)
sum(abs.(glm_preds .- package_predictions)) < 1e-10

# Degree of 1 with no intercept
lm1 = fit(LinearModel, hcat(ones(obs)), y)
glm_preds = predict(lm1, hcat(ones(obs)))
package_approximation, reg3 = create_ols_approximation(y, X, 0; intercept = true)
package_predictions = evaluate.(package_approximation, X)
sum(abs.(glm_preds .- package_predictions)) < 1e-10

# With applying a base
Xmin = X .- 5
lm1 = fit(LinearModel, hcat(ones(obs), Xmin, Xmin .^ 2), y)
glm_preds = predict(lm1, hcat(ones(obs), Xmin, Xmin .^ 2))
package_approximation, reg4 = create_ols_approximation(y, X, 2; intercept = true, base_x = 5.0)
package_predictions = evaluate.(package_approximation, X)
sum(abs.(glm_preds .- package_predictions)) < 1e-10

# With x as dates and a base.
X = Array{Date}(undef, obs)
baseDate =  Date(2016, 7, 21)
StartDate = Date(2018, 7, 21)
for i in 1:obs
    X[i] = StartDate +Dates.Day(2* (i-1))
end
XConverted = years_between.(X, baseDate)
lm1 = fit(LinearModel, hcat(ones(obs), XConverted, XConverted .^ 2), y)
glm_preds = predict(lm1, hcat(ones(obs), XConverted, XConverted .^ 2))
package_approximation, reg5 = create_ols_approximation(y, X, 2; intercept = true, base_date = baseDate)
package_predictions = evaluate.(package_approximation, X)
sum(abs.(glm_preds .- package_predictions)) < 1e-10

##### Multivariate Tests

Random.seed!(1992)
nObs = 1000
dd = DataFrame()
dd[:x] = rand( Normal(),nObs) + 0.1 .* rand( Normal(),nObs)
dd[:z] = rand( Normal(),nObs) + 0.1 .* rand( Normal(),nObs)
dd[:w] = (0.5 .* rand( Normal(),nObs)) .+ 0.7.*(dd[:z] .- dd[:x]) + 0.1 .* rand( Normal(),nObs)
dd[:y] = (dd[:x] .*dd[:w] ) .* (dd[:z] .- dd[:w]) .+ dd[:x] + rand( Normal(),nObs)
dd[7,:y] = 1.0

# Creating basic PE_Units
linear    = PE_Unit(0.0,0.0,1)
quadratic = PE_Unit(0.0,0.0,2)
# Creating Model
constant_term = PE_Function(1.0, Dict{Symbol,PE_Unit}())
x_lin         = PE_Function(1.0, Dict{Symbol,PE_Unit}(:x => linear))
z_lin         = PE_Function(1.0, Dict{Symbol,PE_Unit}(:z => linear))
w_lin         = PE_Function(1.0, Dict{Symbol,PE_Unit}(:w => linear))
w_quad        = PE_Function(1.0, Dict{Symbol,PE_Unit}(:w => quadratic))
x_lin_z_quad  = PE_Function(1.0, Dict{Symbol,PE_Unit}([:x, :z] .=> [linear, quadratic]))
model = constant_term + x_lin + z_lin + w_lin + w_quad + x_lin_z_quad
# Regression and testing
mod_1, reg_1 = create_ols_approximation(dd, :y, model)
dd[:predicted_y] = evaluate(mod_1, dd)
sum(abs.(dd[:predicted_y] - reg_1.rr.mu)) < 1e-12

# Now we look at a saturated model.
y = :y
x_variables = [:x,:z]
degree = 2
intercept = true
univariate_dim_name = :default
mod_2, reg_2 = create_saturated_ols_approximation(dd, y, x_variables, degree)
dd[:predicted_y_2] = evaluate(mod_2, dd)
sum(abs.(dd[:predicted_y_2] - reg_2.rr.mu)) < 1e-12

# Now we look at a saturated model.
y = :y
x_variables = [:x,:y,:z]
degree = 3
intercept = true
univariate_dim_name = :default
mod_3, reg_3 = create_saturated_ols_approximation(dd, y, x_variables, degree)
dd[:predicted_y_3] = evaluate(mod_3, dd)
sum(abs.(dd[:predicted_y_3] - reg_3.rr.mu)) < 1e-11

# Now the second model should have higher RSS than the first because it has fewer variables and terms. We can test this:
sum((dd[:predicted_y_2] .- dd[:y]).^2) > sum((dd[:predicted_y] .- dd[:y]).^2)
# The third model should have lower RSS than the first because it has more terms. Testing this:
sum((dd[:predicted_y_3] .- dd[:y]).^2) < sum((dd[:predicted_y] .- dd[:y]).^2)


unstructured_index(i::T, j::T, nx::T) where {T<:Integer} = (j - 1) * nx + i

function unstructured_grid(coords::T, corners::T) where {T<:AbstractArray{<:Real,3}}
    nx = size(coords, 1)
    ny = size(coords, 2)
    ne = nx * ny
    np = (nx + 1) * (ny + 1)

    points = zeros(np, 2)
    for i = 1:nx+1, j = 1:ny+1
        idx = unstructured_index(i, j, nx + 1)
        points[idx, :] .= corners[i, j, :]
    end

    cells = zeros(Int, ne, 4)
    for i = 1:nx, j = 1:ny
        idx = unstructured_index(i, j, nx)
        cells[idx, 1] = unstructured_index(i, j, nx + 1)
        cells[idx, 2] = unstructured_index(i + 1, j, nx + 1)
        cells[idx, 3] = unstructured_index(i + 1, j + 1, nx + 1)
        cells[idx, 4] = unstructured_index(i, j + 1, nx + 1)
    end

    return cells, points
end

function unstructured_grid(x::T, y::T, x0, x1, y0, y1) where {T<:AbstractMatrix}
    Δx = zero(x[:, 1])
    Δx[1] = 2.0 * (x[1, 1] - x0)
    for i = 2:length(Δx)-1
        Δx[i] = 2.0 * (x[i, 1] - x[i-1, 1] - Δx[i-1])
    end
    Δx[end] = 2.0 * (x1 - x[end, 1])

    Δy = zero(y[1, :])
    Δy[1] = 2.0 * (y[1, 1] - y0)
    for i = 2:length(Δy)-1
        Δy[i] = 2.0 * (y[1, i] - y[1, i-1] - Δy[i-1])
    end
    Δy[end] = 2.0 * (y1 - y[1, end])

    points = zeros(length(Δx) + 1, length(Δy) + 1, 2)
    for i in axes(points, 1), j in axes(points, 2)
        if i <= length(Δx) && j <= length(Δy)
            points[i, j, 1] = x[i, j] - Δx[i] / 2
            points[i, j, 2] = y[i, j] - Δy[j] / 2
        elseif i == length(Δx) + 1 && j <= length(Δy)
            points[i, j, 1] = x[i-1, j] + Δx[i-1] / 2
            points[i, j, 2] = y[i-1, j] - Δy[j] / 2
        elseif i <= length(Δx) && j == length(Δy) + 1
            points[i, j, 1] = x[i, j-1] - Δx[i] / 2
            points[i, j, 2] = y[i, j-1] + Δy[j-1] / 2
        elseif i == length(Δx) + 1 && j == length(Δy) + 1
            points[i, j, 1] = x[i-1, j-1] + Δx[i-1] / 2
            points[i, j, 2] = y[i-1, j-1] + Δy[j-1] / 2
        end
    end

    unstructured_grid(cat(x, y, dims = 3), points)
end

function unstructured_grid(coords::T, xrange, yrange) where {T<:AbstractArray{<:Real,3}}
    xcoords = coords[:, :, 1]
    ycoords = coords[:, :, 2]
    x0, x1 = xrange
    y0, y1 = yrange

    unstructured_grid(xcoords, ycoords, x0, x1, y0, y1)
end

"""
    triangulate(points::T) where {T<:AbstractMatrix}

Split rectangular mesh into triangles

- @arg points: structured node coordinates with shape (npoints, ndim)
"""
function triangulate(points::T) where {T<:AbstractMatrix}
    sp = pyimport("scipy.spatial")
    tri = sp.Delaunay(points)

    return tri.simplices .+ 1 # python is zero-indexed
end


# instantiate environment
using Pkg; Pkg.instantiate()

using GeoStats
using SpectralGaussianSimulation
using DensityRatioEstimation
using CategoricalArrays
using LossFunctions
using ProgressMeter
using DataFrames
using MLJ, CSV
using LinearAlgebra
using Random

# reproducible results
Random.seed!(2020)

# generate random images with given spatial mean, range and sill
function generator(nimgs=1; mean=0., range=1., sill=1., size=(100,100))
  γ = GaussianVariogram(range=range, sill=sill)
  p = SimulationProblem(RegularGrid{Float64}(size), (:Z1=>Float64,:Z2=>Float64), nimgs)
  s = SpecGaussSim(:Z1=>(mean=mean, variogram=γ), :Z2=>(mean=mean, variogram=γ))
  solve(p, s)
end

# generate a covariate shift configuration with given δ and τ
function covariateshift(δ, τ, r; ns=1, nt=1, size=(100,100))
  μ₁, σ₁ = 0.0, 1.0
  μ₂, σ₂ = (3*√2*σ₁)δ + μ₁, τ*σ₁
  simgs = generator(ns, mean=μ₁, sill=σ₁^2, range=r, size=size)
  timgs = generator(nt, mean=μ₂, sill=σ₂^2, range=r, size=size)
  simgs, timgs
end

# geostatistical learning problem with given covariate shift configuration
function problem(; δ=0.0, τ=1.0, r=10.0, size=(100,100))
  # covariate shift
  simgs, timgs = covariateshift(δ, τ, r, nt=101, size=size)

  # sine-norm labeling function
  label(z, p=1) = sin(4*norm(z, p)) < 0 ? 1 : 0
  function f(Γ)
    labels = [label(z) for z in view(Γ, [:Z1,:Z2])]
    ldict  = OrderedDict(:LABEL => categorical(labels))
    georef(ldict, domain(Γ))
  end

  # add labels to all samples
  Ωss = [join(Ω, f(Ω)) for Ω in simgs]
  Ωts = [join(Ω, f(Ω)) for Ω in timgs]

  # geostatistical learning problem
  t = ClassificationTask((:Z1,:Z2), :LABEL)
  p = LearningProblem(Ωss[1], Ωts[1], t)

  # return problem and other Ωt samples
  p, Ωts[2:end]
end

# estimators of generalization error
error_cv( m, p, k, ℒ) = error(PointwiseLearn(m), p, CrossValidation(k, loss=ℒ))
error_bcv(m, p, r, ℒ) = error(PointwiseLearn(m), p, BlockCrossValidation(r, loss=ℒ))
error_drv(m, p, k, ℒ) = error(PointwiseLearn(m), p, DensityRatioValidation(k, loss=ℒ, estimator=LSIF(σ=2.0,b=10)))

# true error (empirical approximation)
function error_empirical(m, p, Ωts, ℒ)
  # train on source data
  lm = learn(task(p), sourcedata(p), m)

  # test on various samples of target data
  ϵs = map(Ωts) do Ωt
    Ω̂t = perform(task(p), Ωt, lm)
    y  = vec(Ωt[:LABEL])
    ŷ  = vec(Ω̂t[:LABEL])
    LossFunctions.value(ℒ[:LABEL], y, ŷ, AggMode.Mean())
  end

  # return mean loss
  mean(ϵs)
end

function experiment(m, δ, τ, r, ℒ)
  # sample a problem
  p, Ωts = problem(δ=δ, τ=τ, r=r)

  # parameters for validation methods
  rᵦ = 20.
  s  = size(sourcedata(p))
  k  = round(Int, prod(s ./ rᵦ))

  # try different error estimates
  cv  = error_cv( m, p, k,  ℒ)[:LABEL]
  bcv = error_bcv(m, p, rᵦ, ℒ)[:LABEL]
  drv = error_drv(m, p, k,  ℒ)[:LABEL]

  # actual error (empirical estimate)
  actual = error_empirical(m, p, Ωts, ℒ)

  # model name (without "Classifier" suffix)
  model = chop(info(m).name, tail=10)

  (δ=δ, τ=τ, r=r, CV=cv, BCV=bcv, DRV=drv, ACTUAL=actual, MODEL=model)
end

# -------------
# MAIN SCRIPT
# -------------

# learning models
@load KNeighborsClassifier pkg="ScikitLearn"
@load DecisionTreeClassifier pkg="DecisionTree"

# misclassification loss
ℒ = Dict(:LABEL => MisclassLoss())

# parameter ranges
mrange = [KNeighborsClassifier(), DecisionTreeClassifier()]
δrange = 0.0:0.1:0.7
τrange = 0.3:0.1:0.9
rrange = [1e-4,1e+1,2e+1]

# experiment iterator and progress
iterator = Iterators.product(mrange, δrange, τrange, rrange)
progress = Progress(length(iterator), "Gaussian experiment:")

# return missing in case of failure
skip = e -> (println("Skipped: $e"); missing)

# perform experiments
results = progress_pmap(iterator, progress=progress,
                        on_error=skip) do (m, δ, τ, r)
  experiment(m, δ, τ, r, ℒ)
end

# save results
fname = joinpath(@__DIR__,"results","gaussian.csv")
CSV.write(fname, DataFrame(skipmissing(results)))


using IntelVectorMath
using ForwardDiff
using BenchmarkTools

function IVM.atan(xd::Vector{ForwardDiff.Dual{T, V, 1}},
                  yd::Vector{ForwardDiff.Dual{T, V, 1}}) where {T, V}

    x = ForwardDiff.value.(xd)
    y = ForwardDiff.value.(yd)
    z = IVM.atan(x, y)

    δx = ForwardDiff.partials.(xd, 1)
    δy = ForwardDiff.partials.(yd, 1)
    δz = @. (δy*x - δx*y)/(x^2 + y^2)

    ForwardDiff.Dual{T}.(z, δz)
end

struct fast_atan_dual_cache{T}
    x::Vector{T}
    y::Vector{T}
    z::Vector{T}
    δ::Matrix{T}
end

function allocate_fast_atan_dual_cache(T::Type, N::Int)
    x = Vector{T}(undef, N)
    y = Vector{T}(undef, N)
    z = Vector{T}(undef, N)
    δ = Matrix{T}(undef, N, 3)
    fast_atan_dual_cache(x, y, z, δ)
end

function IVM.atan!(zd::Vector{ForwardDiff.Dual{T, V, 1}},
                   xd::Vector{ForwardDiff.Dual{T, V, 1}},
                   yd::Vector{ForwardDiff.Dual{T, V, 1}},
                   cache::fast_atan_dual_cache{V}) where {T, V}

   x, y, z = cache.x, cache.y, cache.z
   δx, δy, δz = view(cache.δ,:,1), view(cache.δ,:,2), view(cache.δ,:,3)

    x .= ForwardDiff.value.(xd)
    y .= ForwardDiff.value.(yd)
    IVM.atan!(z, x, y)

    δx .= ForwardDiff.partials.(xd, 1)
    δy .= ForwardDiff.partials.(yd, 1)
    δz .= @. (δy*x - δx*y)/(x^2 + y^2)

    zd .= ForwardDiff.Dual{T}.(z, δz)
end

N = 500
x = rand(N)
y = rand(N)
δx = rand(N)
δy = rand(N)
xd = ForwardDiff.Dual{DERTag}.(x, δx)
yd = ForwardDiff.Dual{DERTag}.(y, δy)
zd = IVM.atan(xd,yd)

cache_atan = allocate_fast_atan_dual_cache(Float64, N)
@btime IVM.atan!(zd, xd, yd, cache_atan)
@btime $zd .= atan.($xd, $yd)

z = rand(N)
@btime IVM.atan!(z, x, y)
@btime z .= atan.(x, y)

struct DERTag end
v = rand(3)
p1 = rand(3)
p2 = rand(3)
p = [p1, p2]
d = ForwardDiff.Dual{DERTag}.(v, p...)

ForwardDiff.partials.(d, 1)
ForwardDiff.partials.(d, 2)
ForwardDiff.value.(d)


"""
    is_leap_year(year)

Return `true` if `year` is a leap year in the gregorian calendar.
    on every year that is evenly divisible by 4
    except every year that is evenly divisible by 100
      unless the year is also evenly divisible by 400
"""
function is_leap_year(year)
    if year % 4 == 0 && year % 100 != 0
		return true
	elseif year % 4 == 0 && year % 400 == 0
        return true
	else
		return false
    end

end


abstract type AbstractSection end

update!(sh::AbstractSection) = error("Not yet implemented. You can use the dummy function `sh->true`.")

####################################################################################################
@views function sectionShooting(x::AbstractVector, T, normal::AbstractVector, center::AbstractVector)
	N = length(center)
	# we only constrain the first point to lie on a specific hyperplane
	# this avoids the temporary xc - centers
	return (dot(x, normal) - dot(center, normal)) * T
end

# section for Standard Shooting
"""
$(TYPEDEF)

This composite type (named for Section Standard Shooting) encodes a type of section implemented by a hyperplane. It can be used in conjunction with [`ShootingProblem`](@ref). The hyperplane is defined by a point `center` and a `normal`.

$(TYPEDFIELDS)

"""
struct SectionSS{Tn, Tc}  <: AbstractSection
	"Normal to define hyperplane"
	normal::Tn

	"Representative point on hyperplane"
	center::Tc
end

(sect::SectionSS)(u, T) = sectionShooting(u, T, sect.normal, sect.center)

isEmpty(sect::SectionSS{Tn, Tc}) where {Tn, Tc} = (Tn == Nothing) || (Tn == Nothing)

# we update the field of Section, useful during continuation procedure for updating the section
function update!(sect::SectionSS, normal, center)
	copyto!(sect.normal, normal)
	copyto!(sect.center, center)
	sect
end

####################################################################################################
# Poincare shooting based on Sánchez, J., M. Net, B. Garcı́a-Archilla, and C. Simó. “Newton–Krylov Continuation of Periodic Orbits for Navier–Stokes Flows.” Journal of Computational Physics 201, no. 1 (November 20, 2004): 13–33. https://doi.org/10.1016/j.jcp.2004.04.018.

function sectionHyp!(out, x, normals, centers)
	for ii = 1:length(normals)
		out[ii] = dot(normals[ii], x) - dot(normals[ii], centers[ii])
	end
	out
end

"""
$(TYPEDEF)

This composite type (named for SectionPoincaréShooting) encodes a type of Poincaré sections implemented by hyperplanes. It can be used in conjunction with [`PoincareShootingProblem`](@ref). Each hyperplane is defined par a point (one example in `centers`) and a normal (one example in `normals`).

$(TYPEDFIELDS)

# Constructor(s)
    SectionPS(normals::Vector{Tv}, centers::Vector{Tv})

"""
struct SectionPS{Tn, Tc, Tnb, Tcb} <: AbstractSection
	M::Int64				# number of hyperplanes
	normals::Tn 			# normals to define hyperplanes
	centers::Tc 			# representative point on each hyperplane
	indices::Vector{Int64} 	# indices to be removed in the operator Ek

	normals_bar::Tnb
	centers_bar::Tcb

	function SectionPS(normals, centers)
		@assert length(normals) == length(centers)
		M = length(normals)
		indices = zeros(Int64, M)
		for ii in 1:M
			indices[ii] = argmax(abs.(normals[ii]))
		end
		nbar = [R(normals[ii], indices[ii]) for ii=1:M]
		cbar = [R(centers[ii], indices[ii]) for ii=1:M]

		return new{typeof(normals), typeof(centers), typeof(nbar), typeof(cbar)}(M, normals, centers, indices, nbar, cbar)
	end

	SectionPS(M = 0) = new{Nothing, Nothing, Nothing, Nothing}(M, nothing, nothing, Int64[], nothing, nothing)
end

(hyp::SectionPS)(out, u) = sectionHyp!(out, u, hyp.normals, hyp.centers)
isEmpty(sect::SectionPS{Tn, Tc, Tnb, Tcb}) where {Tn, Tc, Tnb, Tcb} = (Tn == Nothing) || (Tn == Nothing)

"""
	update!(hyp::SectionPS, normals, centers)

Update the hyperplanes saved in `hyp`.
"""
function update!(hyp::SectionPS, normals, centers)
	M = hyp.M
	@assert length(normals) == M "Wrong number of normals"
	@assert length(centers) == M "Wrong number of centers"
	for ii in 1:M
		hyp.normals[ii] .= normals[ii]
		hyp.centers[ii] .= centers[ii]
		hyp.indices[ii] = argmax(abs.(normals[ii]))
		R!(hyp.normals_bar[ii], normals[ii], hyp.indices[ii])
		R!(hyp.centers_bar[ii], centers[ii], hyp.indices[ii])
	end
	return hyp
end

# Operateur Rk from the paper above
@views function R!(out, x::AbstractVector, k::Int)
	out[1:k-1] .= x[1:k-1]
	out[k:end] .= x[k+1:end]
	return out
end

R!(hyp::SectionPS, out, x::AbstractVector, k::Int) = R!(out, x, hyp.indices[k])
R(x::AbstractVector, k::Int) = R!(similar(x, length(x) - 1), x, k)
R(hyp::SectionPS, x::AbstractVector, k::Int) = R!(hyp, similar(x, length(x) - 1), x, k)

# differential of R
dR!(hyp::SectionPS, out, dx::AbstractVector, k::Int) = R!(hyp, out, dx, k)

# Operateur Ek from the paper above
function E!(hyp::SectionPS, out, xbar::AbstractVector, ii::Int)
	@assert length(xbar) == length(hyp.normals[1]) - 1 "Wrong size for the projector / expansion operators, length(xbar) = $(length(xbar)) and length(normal) = $(length(hyp.normals[1]))"
	k = hyp.indices[ii]
	nbar = hyp.normals_bar[ii]
	xcbar = hyp.centers_bar[ii]
	coord_k = hyp.centers[ii][k] - (dot(nbar, xbar) - dot(nbar, xcbar)) / hyp.normals[ii][k]

	@views out[1:k-1] .= xbar[1:k-1]
	@views out[k+1:end] .= xbar[k:end]
	out[k] = coord_k
	return out
end

function E(hyp::SectionPS, xbar::AbstractVector, ii::Int)
	out = similar(xbar, length(xbar) + 1)
	E!(hyp, out, xbar, ii)
end

# differential of E!
function dE!(hyp::SectionPS, out, dxbar::AbstractVector, ii::Int)
	k = hyp.indices[ii]
	nbar = hyp.normals_bar[ii]
	xcbar = hyp.centers_bar[ii]
	coord_k = - dot(nbar, dxbar) / hyp.normals[ii][k]

	@views out[1:k-1]   .= dxbar[1:k-1]
	@views out[k+1:end] .= dxbar[k:end]
	out[k] = coord_k
	return out
end

function dE(hyp::SectionPS, dxbar::AbstractVector, ii::Int)
	out = similar(dxbar, length(dxbar) + 1)
	dE!(hyp, out, dxbar, ii)
end


using Random

include("input/fMultiSinGen.jl")

function gensignalARX(θ_true, τ_true; M1=2, M2=2, fMin=0.8, fMax=1.0, fs=1.0, T=100)

    # Orders
    M = M1+1+M2

    # Parameters
    θ_true = θ_scale .*(rand(M,) .- 0.5)

    # Input frequency and amplitude
    input, _ = fMultiSinGen(T, 1, 1, fMin=fMin, fMax=fMax, fs=fs, type_signal="odd")

    # Observation array
    output = zeros(T,)
    errors = zeros(T,)

    for k = 1:T

        # Input
        input[k] = A[k]*cos(ω*k)

        # Generate noise
        errors[k] = sqrt(inv(τ_true))*randn(1)[1]

        # Output
        if k < M
            output[k] = input[k] + errors[k]
        else
            # Update history vectors
            x_kmin1 = output[k-1:-1:k-M1]
            z_kmin1 = input[k-1:-1:k-M2]

            # Compute output
            output[k] = θ_true'*[x_kmin1; input[k]; z_kmin1] + errors[k]
        end
    end

    return input, output
end

function gensignalNARX(ϕ, θ_scale, τ_true; M1=2, M2=2, degree=1, fMin=0.8, fMax=1.0, fs=1.0, T=100)

    # Orders
    M = (M1+1+M2)*degree + 1

    # Parameters
    θ_true = θ_scale .*(rand(M,) .- 0.5)
    θ_true[end] = 0.

    # Input frequency and amplitude
    input, _ = fMultiSinGen(T, 1, 1, fMin=fMin, fMax=fMax, fs=fs, type_signal="odd")

    # Observation array
    output = zeros(T,)
    errors = zeros(T,)

    for k = 1:T

        # Input
        input[k] = A[k]*cos(ω*k)

        # Generate noise
        errors[k] = sqrt(inv(τ_true))*randn(1)[1]

        # Output
        if k < M
            output[k] = input[k] + errors[k]
        else
            # Update history vectors
            x_kmin1 = output[k-1:-1:k-M1]
            z_kmin1 = input[k-1:-1:k-M2]

            # Compute output
            output[k] = θ_true'*ϕ([x_kmin1; input[k]; z_kmin1]) + errors[k]
        end
    end

    return input, output
end

function gensignalNARMAX(ϕ, θ_scale, τ_true; M1=2, M2=2, M3=2, degree=1, fMin=0.8, fMax=1.0, fs=1.0, uStd=1., T=100, num_periods=1, points_period=1)

    # Orders
    M = (M1+1+M2+M3)*degree + 1

    # Parameters
    θ_true = θ_scale .*(rand(M,) .- 0.5)
    θ_true[end] = 0.0

    # Input frequency and amplitude
    input, _ = fMultiSinGen(points_period, num_periods, 1, fMin=fMin, fMax=fMax, fs=fs, type_signal="full", uStd=uStd)

    # Observation array
    output = zeros(T,)
    errors = zeros(T,)

    for k = 1:T

        # Generate noise
        errors[k] = sqrt(inv(τ_true))*randn(1)[1]

        # Output
        if k < (maximum([M1,M2,M3])+1)
            output[k] = input[k] + errors[k]
        else
            # Update history vectors
            x_kmin1 = output[k-1:-1:k-M1]
            z_kmin1 = input[k-1:-1:k-M2]
            r_kmin1 = errors[k-1:-1:k-M3]

            # Compute output
            output[k] = θ_true'*ϕ([x_kmin1; input[k]; z_kmin1; r_kmin1]) + errors[k]
        end
    end

    return input, output
end

function split_data(T, split_index; start_index=10)

    # Check inputs
    if split_index > T
        error("Split index "*string(split_index)*"must be lower than end time T="*string(T))
    end

    # Select training set
    ix_trn = Array{Int64,1}(start_index:split_index)

    # Select validation set
    ix_val = Array{Int64,1}(split_index+1:T)

    return ix_trn, ix_val
end


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using SparseArrays
Random.seed!(233)

function sparse_to_dense(A)
    A = constant(A)
    ij = A.o.indices
    vv = values(A)
    m, n = size(A)
    sparse_to_dense_ = load_op_and_grad("./build/libSparseToDense","sparse_to_dense_ad")
    m_, n_ = convert_to_tensor(Any[m,n], [Int64,Int64])
    out = sparse_to_dense_(ij, vv, m_,n_)
    set_shape(out, (m, n))
end

A = sprand(10,10,0.3)
# TODO: specify your input parameters
u = sparse_to_dense(A)
sess = Session(); init(sess)
@show run(sess, u) - Array(A)

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(x)
    B = SparseTensor(ij[:,1]+1, ij[:,2]+1, x, m, n)
    return sum(sparse_to_dense(B)^2)
end

# TODO: change `m_` and `v_` to appropriate values
A = constant(A)
m, n = size(A)
K = length(values(A))
ij = A.o.indices
m_ = constant(rand(K))
v_ = rand(K)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


using PermPlain
using Test
using PermPlain: permcycles, cycstoperm
using Random

Random.seed!(1234)

makec() = PCycles([[1, 6, 14, 7, 3, 11], [2, 8, 12, 9, 10, 5, 15, 4]], 15)
makep() = [3, 7, 2, 4, 8, 10, 1, 6, 9, 5]

@testset "isperm, cyclelengths, numcycles" begin
    ntrials = 50
    for n in (3, 4, 5, 10, 15, 30)
        for trial in 1:ntrials
            p = randperm(n)
            sp = permsparse(p)
            @test permlist(permcycles(permmatrix(sp))) == p
            cl = cyclelengths(p)
            ncl = numcycles(p)
            for prep in (p, permcycles(p), permmatrix(p), sp)
                @test cl == cyclelengths(prep)
                @test ncl == numcycles(prep)
                @test isperm(prep)
            end
            @test permsgn(permcompose(p, p)) == 1
            @test permcommute(p, p)
        end
    end
end

@testset "permpower canoncycles" begin
    np = 3
    c = makec()
    p = permlist(c)
    pp = permpower(p, np);
    pp1 = permpower2(p, np);
    c1 = permpower(c,np);
    p2 = permlist(c1);
    @test pp == p2
    @test pp1 == p2
    @test permcycles(pp) == canoncycles(c1)
    @test cyc_pow_perm(c, np) == pp
end

@testset "permsparse" begin
    p = makep()
    sp = permsparse(permmatrix(p))
    @test permlist(permcycles(sp)) == p

    sp = permsparse(permcycles(p))
    @test permlist(permmatrix(sp)) == p
end

@testset "permcompose" begin
    ntrials = 100
    for n in (3, 4, 10)
        for trial in 1:ntrials
            p = randperm(n);
            q = randperm(n);
            pm = permmatrix(p);
            qm = permmatrix(q);
            @test permcompose(p, q) == permlist(qm * pm)
        end
    end
end

@testset "permkron permcommute" begin
    ntrials = 100
    nbad = 0
    ngood = 0
#    for n in (3, 5, 10)
    for n in (3, 4)
        for trial in 1:ntrials
            p = randperm(n);
            q = randperm(n);
            pq = permkron(p, q);
            pm = permmatrix(p);
            qm = permmatrix(q);
            pqm = kron(pm, qm);
            @test pq == permlist(pqm)
            @test isperm(pq)
            if permcommute(p, q) != permcommute(pm, qm)
                nbad += 1
                # println(p)
                # println(q)
                # println()
            else
                ngood += 1
            end
        end
    end
    @info "$nbad / $(nbad + ngood) permcommute tests failed"
end