struct DummySolver <: SemialgebraicSets.AbstractAlgebraicSolver end

@testset "Basic semialgebraic set" begin
    Mod.@polyvar x y
    @test isa(FullSpace(), FullSpace)
    V = @set x * y == 1
    @test V isa AlgebraicSet
    @test_throws ArgumentError addinequality!(V, x*y)
    @testset "Basic" begin
        S = @set x - y == 0 && x^2*y >= 1
        addequality!(S, x^2 - y)
        addinequality!(S, x + y - 1)
        # Algebraic set forces `Float64`
        @test S isa BasicSemialgebraicSet{Float64}
        @test sprint(show, S) == "{ (x, y) | x - y = 0, x^2 - y = 0, x^2*y - 1.0 ≥ 0, x + y - 1.0 ≥ 0 }"
        @test sprint(show, MIME"text/plain"(), S) == "Basic semialgebraic Set defined by 2 equalities\n x - y = 0\n x^2 - y = 0\n2 inequalities\n x^2*y - 1.0 ≥ 0\n x + y - 1.0 ≥ 0\n"
        @test S.V isa AlgebraicSet{Float64}
        @test sprint(show, S.V) == "{ (x, y) | x - y = 0, x^2 - y = 0 }"
        @test sprint(show, MIME"text/plain"(), S.V) == "Algebraic Set defined by 2 equalities\n x - y = 0\n x^2 - y = 0\n"
        @test S === MultivariatePolynomials.changecoefficienttype(S, Float64)
        @test S.V === MultivariatePolynomials.changecoefficienttype(S.V, Float64)
        @test S.V.I === convert(typeof(S.V.I), S.V.I)
        @test BasicSemialgebraicSet{Int, polynomialtype(x, Int)}() isa BasicSemialgebraicSet{Float64, polynomialtype(x, Float64)}
        @test Int32(2)*x^2*y isa MultivariatePolynomials.AbstractTerm{Int32}
        Sf = MultivariatePolynomials.changecoefficienttype(S, Float32)
        @test Sf isa BasicSemialgebraicSet{Float32}
        @test Sf.V isa AlgebraicSet{Float32}

        @testset "Mixed types" begin
            S = (@set Int32(2)*x^2*y == 0 && 1.0x^2*y >= 0 && (6//3)*x^2*y == -y && 1.5x+y >= 0)
            S2 = S ∩ V
            S3 = V ∩ S
            @test inequalities(S2) == inequalities(S3) == S.p
            @test equalities(S2) == S3.V.I.p
        end

        T = (@set x*y^2 == -1 && x^2 + y^2 <= 1)
        V2 = @set T.V && V && x + y == 2.0
        @test V2 isa AlgebraicSet
        @test V2.I.p == [equalities(T); equalities(V); x + y - 2.0]
        S4 = @set S && T
        @test S4.p == [S.p; inequalities(T)]
        @test equalities(S4) == [S.V.I.p; T.V.I.p]

        @testset "Different variables" begin
            T = (@set x == x^2 && y <= y^2)
            @test sprint(show, T) == "{ (x, y) | -x^2 + x = 0, y^2 - y ≥ 0 }"
            @test sprint(show, MIME"text/plain"(), T) == "Basic semialgebraic Set defined by 1 equalitty\n -x^2 + x = 0\n1 inequalitty\n y^2 - y ≥ 0\n"
        end
    end
    @testset "Basic with no equality" begin
        S = @set x + y ≥ 1 && x ≤ y
        @test sprint(show, S) == "{ (x, y) | x + y - 1 ≥ 0, -x + y ≥ 0 }"
        @test sprint(show, MIME"text/plain"(), S) == """
Basic semialgebraic Set defined by no equality
2 inequalities
 x + y - 1 ≥ 0
 -x + y ≥ 0
"""
        @test sprint(show, S.V) == "R^n"
        @test sprint(show, MIME"text/plain"(), S.V) == """
Algebraic Set defined by no equality
"""
        @test S isa BasicSemialgebraicSet{Int}
        @test S === MultivariatePolynomials.changecoefficienttype(S, Int)
        @test S.V isa FullSpace
        Sf = MultivariatePolynomials.changecoefficienttype(S, Float64)
        @test Sf isa BasicSemialgebraicSet{Float64}
        @test Sf.V isa FullSpace

        @test S ∩ FullSpace() === S
        @test S.V ∩ FullSpace() === S.V
        @test FullSpace() ∩ S  === S
        @test FullSpace() ∩ S.V  === S.V
        @test S.V === MultivariatePolynomials.changecoefficienttype(S.V, Float64)
    end
    @testset "Basic with fixed variables" begin
        S = @set x == 1 && x ≤ x^2
        @test sprint(show, S) == "{ (x) | x - 1 = 0, x^2 - x ≥ 0 }"
        @test sprint(show, MIME"text/plain"(), S) == "Basic semialgebraic Set defined by 1 equalitty\n x - 1 = 0\n1 inequalitty\n x^2 - x ≥ 0\n"
        @test sprint(show, S.V) == "{ (x) | x - 1 = 0 }"
        @test sprint(show, MIME"text/plain"(), S.V) == "Algebraic Set defined by 1 equalitty\n x - 1 = 0\n"

        S = @set x == 1 && x ≤ y && 2 == y
        @test S isa BasicSemialgebraicSet{Int}
        @test S.V isa FixedVariablesSet{<:AbstractVariable, Int}
        @test rem(x + y, ideal(S.V)) == 3
        @test S === MultivariatePolynomials.changecoefficienttype(S, Int)
        @test S.V === MultivariatePolynomials.changecoefficienttype(S.V, Int)
        @test S.V.ideal === convert(typeof(S.V.ideal), S.V.ideal)
        Sf = MultivariatePolynomials.changecoefficienttype(S, Float64)
        @test Sf isa BasicSemialgebraicSet{Float64}
        @test Sf.V isa FixedVariablesSet{<:AbstractVariable, Float64}
        @test rem(x + y, ideal(Sf.V)) == 3

        S = @set x == 1 && x ≤ y && 2 == y && 1 == x
        @test S isa BasicSemialgebraicSet{Int}
        @test S.V isa FixedVariablesSet{<:AbstractVariable, Int}
        @test rem(x + y, ideal(S.V)) == 3
        Sf = MultivariatePolynomials.changecoefficienttype(S, Float64)
        @test Sf isa BasicSemialgebraicSet{Float64}
        @test Sf.V isa FixedVariablesSet{<:AbstractVariable, Float64}
        @test rem(x + y, ideal(Sf.V)) == 3

        @testset "Empty" begin
            for S in (@set(x == 1 && x ≤ y && 2 == y && 2 == x),
                      @set(x == 1 && x ≤ y && 2 == x),
                      @set(x == 1 && x ≤ y && 2 == x && y == 3),
                      @set(x ≤ y && y == 3 && x == 1 && y == 1),
                      @set(x ≤ y && y == 3 && x == 1 && y == 3 && x == 2))
                @test S isa BasicSemialgebraicSet{Int}
                @test S.V isa FixedVariablesSet{<:AbstractVariable, Int}
                @test isempty(S.V)
                @test iszero(length(S.V))
                @test isempty(collect(S.V))
                @test iszero(rem(x + y, ideal(S.V)))
            end
        end
    end
    @testset "Basic mixed fixed variables and equalities" begin
        for S in [@set(x == 1 && x ≤ y && 2 + y == x),
                  @set(x == 1 && 2 + y == x && x ≤ y),
                  @set(x ≤ y && 2 + y == x && x == 1),
                  @set(x ≤ y && x == 1 && 2 + y == x),
                  @set(2 + y == x && x ≤ y && x == 1),
                  @set(2 + y == x && x ≤ y && x == 1)]
            @test S isa BasicSemialgebraicSet{Float64}
            @test S.V isa AlgebraicSet{Float64}
        end

        solver = DummySolver()
        for S in [@set(x == 1 && x ≤ y && 2 + y == x, solver),
                  @set(x == 1 && 2 + y == x && x ≤ y, solver),
                  @set(x ≤ y && 2 + y == x && x == 1, solver),
                  @set(x ≤ y && x == 1 && 2 + y == x, solver),
                  @set(2 + y == x && x ≤ y && x == 1, solver),
                  @set(2 + y == x && x ≤ y && x == 1, solver)]
            @test S isa BasicSemialgebraicSet{Float64}
            @test S.V isa AlgebraicSet{Float64}
            @test S.V.solver === solver
        end
    end
end


using Test
using DataDeps
using ParameterEstimocean

using Statistics

@testset "Transformations tests" begin
    data_path = datadep"two_day_suite_2m/free_convection_instantaneous_statistics.jld2";
    field_names = :u

    raw_observations = SyntheticObservations(data_path; field_names = field_names)
    Nz = raw_observations.grid.Nz
    times = raw_observations.times

    map = ConcatenatedOutputMap()

    @info "  Test space and time slicing..."
    time_indices = 101
    z_indices = [collect(20:2:Int(Nz/2)); Nz-4]

    time_transformation = TimeIndices(t=time_indices)
    space_transformation = SpaceIndices(y=:, z=z_indices)

    transformation = Transformation(space = space_transformation,
                                    time = time_transformation,
                                    normalization = nothing)

    observations = SyntheticObservations(data_path; field_names = field_names, transformation)

    sliced_raw_observations = raw_observations.field_time_serieses.u[:, :, z_indices, time_indices][1, 1, :, :]
    ny, nz = size(sliced_raw_observations)
    sliced_raw_observations = reshape(sliced_raw_observations, (ny * nz, 1))

    @test parent(observation_map(map, observations)) ≈ parent(sliced_raw_observations)

    time_indices = [100, 200]
    z_indices = 20:40

    time_transformation = TimeIndices(t=time_indices)
    space_transformation = SpaceIndices(y=:, z=z_indices)

    transformation = Transformation(space = space_transformation,
                                    time = time_transformation,
                                    normalization = ZScore())

    observations = SyntheticObservations(data_path; field_names = field_names, transformation)

    sliced_raw_observations = raw_observations.field_time_serieses.u[:, :, z_indices, time_indices][1, 1, :, :]
    ny, nz = size(sliced_raw_observations)
    sliced_raw_observations = reshape(sliced_raw_observations, (ny * nz, 1))

    # normalize
    μ, σ =  mean(sliced_raw_observations), sqrt(cov(sliced_raw_observations; corrected=false))
    @. sliced_raw_observations = (sliced_raw_observations - μ) / σ

    @test parent(observation_map(map, observations)) ≈ parent(sliced_raw_observations)

    @info "  Test applying space and time weights..."

    time_scale_factor = 0.55
    time_weights = time_scale_factor * ones(size(times))

    depth_scale = 100
    z = raw_observations.grid.zᵃᵃᶜ[1:Nz]
    z_weights = @. exp(z / depth_scale)
    space_weights = reshape(z_weights, 1, 1, Nz)

    transformation = Transformation(space = space_weights,
                                    time = time_weights,
                                    normalization = nothing)

    observations = SyntheticObservations(data_path; field_names = field_names, transformation)

    weighted_raw_observations = raw_observations.field_time_serieses.u[1, 1, 1:Nz, :]
    @. weighted_raw_observations *= z_weights * time_scale_factor
    ny, nz = size(weighted_raw_observations)
    weighted_raw_observations = reshape(weighted_raw_observations, (ny * nz, 1))

    @test parent(observation_map(map, observations)) ≈ parent(weighted_raw_observations)
end


# This is just a bit of code to set up a temporary package environment since not everyone
# will have the same environment set up; this is largely unnecessary when working only on
# your own code.
# See https://bkamins.github.io/julialang/2020/06/28/automatic-project-environments.html
using Pkg: Pkg
# Change to a newly created temporary directory
cd(mktempdir()) do
    Pkg.activate(".")  # activate a new project in the temporary directory
    Pkg.UPDATED_REGISTRY_THIS_SESSION[] = true  # don't update the registry; on Windows it's slow...
    # Add the required packages
    Pkg.add([Pkg.PackageSpec(name="StochasticDiffEq"),
        Pkg.PackageSpec(name="OrdinaryDiffEq"),
        Pkg.PackageSpec(name="StaticArrays"),
        Pkg.PackageSpec(name="Plots"),])
end

# Load required packages
using StochasticDiffEq  # SDE solvers
using OrdinaryDiffEq  # ODE solvers
using StaticArrays  # SVectors for speed
using Plots  # Plotting

# Note on StaticArrays: Ordinary vectors (like Matlab) are created with square brackets,
# e.g., [1, 2, 3], and are good for most purposes. However, since they are variable size the
# compiler cannot optimise as much as it could if they were a fixed size. StaticArrays
# provides fixed-size arrays which allow the compiler to do lots of optimisations and so
# often work much faster for small arrays (typically where there are fewer than 20 elements
# in the array).

"""
    NTMD

A nonlinear tuned-mass-damper system with stochastic forcing of the form

```math
m₁x₁″ + c₁(x₁′ - x₂′) + k₁x₁ + k₂(x₁ - x₂) + μ(x₁ - x₂)^3 = f
m₂x₂″ + c₂(x₂′ - x₁′) + k₂(x₂ - x₁) + μ(x₂ - x₁)^3 = 0
```

where

```math
f = F₀cos(ωt) + ξ(t)
```

and `ξ(t)` is a Brownian motion with variance `σ`.

"""
Base.@kwdef mutable struct NTMD
    m₁::Float64 = 1.0
    c₁::Float64 = 0.02
    k₁::Float64 = 1.0
    m₂::Float64 = 0.5
    c₂::Float64 = 0.05
    k₂::Float64 = 0.95
    μ::Float64 = 0.1
    σ::Float64 = 0.5
    F₀::Float64 = 1.0
    ω::Float64 = 1.0
end

function ntmd_dt(u, p::NTMD, t)
    # Deterministic part
    x₁ = u[1]
    x₁′ = u[2]
    x₂ = u[3]
    x₂′ = u[4]
    f = p.F₀*cos(p.ω*t)
    return SVector(
        x₁′,
        (f - p.c₁*(x₁′ - x₂′) - p.k₁*x₁ - p.k₂*(x₁ - x₂) - p.μ*(x₁ - x₂)^3)/p.m₁,
        x₂′,
        (-p.c₂*(x₂′ - x₁′) - p.k₂*(x₂ - x₁) - p.μ*(x₂ - x₁)^3)/p.m₂,
    )
end

function ntmd_dW(u, p::NTMD, t)
    # Stochastic part
    return SVector(0.0, p.σ/p.m₁, 0.0, 0.0)
end

function ntmd_integrate_sde(u₀, p::NTMD, t)
    # Solve the full SDE
    sde = SDEProblem(ntmd_dt, ntmd_dW, SVector{4, Float64}(u₀), (0.0, Float64(t)), p)
    return solve(sde, SOSRA())
end

function ntmd_integrate_ode(u₀, p::NTMD, t)
    # Solve only the deterministic part of the SDE
    ode = ODEProblem(ntmd_dt, SVector{4, Float64}(u₀), (0.0, Float64(t)), p)
    return solve(ode, Tsit5())
end

function ntmd_example()
    p = NTMD(ω=1.1)  # NTMD with default parameters except ω
    sol = ntmd_integrate_ode(ones(4), p, 2π/p.ω*1000)  # integrate to remove the transient
    u₀ = sol[:, end]  # use the end point for more simulations
    sol = ntmd_integrate_ode(u₀, p, 2π/p.ω*20)  # integrate the deterministic part only (for comparison purposes)
    sol_sde = ntmd_integrate_sde(u₀, p, 2π/p.ω*20)  # integrate the full stochastic differential equation
    plot(sol.t, sol[1, :])
    plot!(sol_sde.t, sol_sde[1, :])
    gui()
    return sol_sde
end


using JOLI
using WAVEFORM
using PyPlot

n = 50*[1;1;1]
d = 12.5*[1;1;1]
o = 0.0*[1;1;1]
v0 = 2000
freq = 8.0
λ = v0/freq

t0 = 0.0
f0 = 10.0
unit = "m/s"
xsrc = [(n.*d)[1]/2]
ysrc = [(n.*d)[2]/2]
zsrc = [100.0]
xrec = range(0.0,stop=(n.*d)[1]::Float64,length=n[1])
yrec = range(0.0,stop=(n.*d)[2]::Float64,length=n[2])
zrec = [100.0]
freqs = [freq]
nsrc = length(xsrc)*length(ysrc)*length(zsrc)
nfreq = length(freqs)
model = Model{Int64,Float64}(n,d,o,t0,f0,unit,freqs,xsrc,ysrc,zsrc,xrec,yrec,zrec)

comp_n = n
comp_d = d
comp_o = o
npml = convert(Int,λ/minimum(comp_d))*ones(Int64,(2,3))
npml = 10*ones(Int64,(2,3))
scheme = WAVEFORM.helm3d_operto27
cut_pml = true
implicit_matrix = true
srcfreqmask = trues(nsrc,nfreq)
misfit = WAVEFORM.least_squares
lsopts = LinSolveOpts(solver=:fgmres,maxinnerit=5,tol=1e-6)
lsopts.precond = :mlgmres
opts = PDEopts{Int64,Float64}(scheme,comp_n,comp_d,comp_o,cut_pml,implicit_matrix,npml,misfit,srcfreqmask,lsopts)

v = v0*ones(Float64,n...)
v = vec(v)
opts.implicit_matrix = true

(H,comp_grid,T,DT_adj,P) = helmholtz_system(v,model,freq,opts)
nt = comp_grid.comp_n
q = zeros(eltype(H),tuple(nt...))
q[div.(nt,2)...] = 1.0
q = vec(q)
u = H\q
u = reshape(u,nt...)

clim = 0.5
imshow(real(u[:,:,div(nt[3],2)]),vmin=-clim,vmax=clim)


using AstroAngles
using Formatting
using StableRNGs
using Test

rng = StableRNG(206265)

randdegree(rng, Ns...) = (rand(rng, Ns...) .- 0.5) .* 720
randrad(rng, Ns...) = (rand(rng, Ns...) .- 0.5) .* 4π
randha(rng, Ns...) = (rand(rng, Ns...) .- 0.5) .* 48
randdms(rng, Ns...) = randdegree(rng, Ns...) .|> deg2dms
randhms(rng, Ns...) = randdegree(rng, Ns...) .|> deg2hms

@testset "conversions" begin include("conversions.jl") end
@testset "parsing" begin include("parsing.jl") end
@testset "printing" begin include("printing.jl") end


# Discrete Measure
#
# Represents a generic discrete distribution with finitely many atoms

# Returns a weighted sample with one row for each distinct row in X and an
# associated weight equaling the sum of all weights.  Rows are sorted in
# ascending order if X is univariate.
function compresssample(X::AbstractArray{T}, q::AbstractArray{Float64}) where T
    n = size(X,1);
    p = size(X,2);
    # Add each row of X with its weight to an accumulator
    a = DataStructures.Accumulator{Array{T,2},Float64}()
    for i in 1:n
        push!(a,X[i:i,:],q[i])
    end
    # Collect distinct points
    nnew = length(a);
    Xnew = zeros(Float64,nnew,p);
    qnew = zeros(Float64,nnew,1);
    for (i,x) in enumerate(keys(a))
        Xnew[i,:] = x;
        qnew[i] = a[x];
    end
    # Sort points if univariate
    if p == 1
        order = sortperm(Xnew[:,1]);
        Xnew = Xnew[order,:];
        qnew = qnew[order,:];
    end
    (Xnew, qnew)
end

mutable struct SteinDiscrete
    # n x p matrix of support points
    support::Array{Float64,2}
    # n x 1 vector of weights associated with support points
    weights::Array{Float64,2}
    # Internal constructor collapses X to its unique sorted elements and
    # adjusts q accordingly; sorts X if univariate.
    SteinDiscrete(support, weights) = (
        (support_,weights_) = compresssample(support,weights);
        return new(support_, weights_);
    )
end

# Constructor with uniform weights q_i = 1/n
SteinDiscrete(X) = SteinDiscrete(
    X,
    [1/size(X,1) for i=1:size(X,1)]
);


# basic_spaces.jl

##############################
# Definitions of basic spaces
##############################

"A geometric space with integer type."
const IntegerSpace{T <: Integer} = GeometricSpace{T}

"A geometric space with rational types."
const RationalSpace{T} = GeometricSpace{Rational{T}}

subeltype(::Type{RationalSpace{T}}) where {T} = T
similar_space(::Type{RationalSpace{T}}, ::Type{S}) where {T,S} = RationalSpace{S}

"A geometric space with floating point type."
const RealSpace{T <: AbstractFloat} = GeometricSpace{T}

"A geometric space with complex type."
const ComplexSpace{T} = GeometricSpace{Complex{T}}

subeltype(::Type{ComplexSpace{T}}) where {T} = T
similar_space(::Type{ComplexSpace{T}}, ::Type{S}) where {T,S} = ComplexSpace{S}

"""
A Euclidean space with static vector types. It is the space of all vectors
of fixed length `N`, with entries of type `T`.
"""
const VectorSpace{N,T} = GeometricSpace{SVector{N,T}}
const EuclideanSpace{N,T} = VectorSpace{N,T}

subeltype(::Type{VectorSpace{N,T}}) where {N,T} = T
similar_space(::Type{VectorSpace{N,T}}, ::Type{S}) where {N,T,S} = VectorSpace{N,S}

"""
A general Array space, with arrays of dimension `N` and element type `T`.
Note that the arrays can have any size. Thus, `ArraySpace{2,Float64}` contains
all possible matrices of size `m × n` for any combination of `m,n ∈ ℕ`.
"""
const ArraySpace{N,T} = GeometricSpace{Array{T,N}}

subeltype(::Type{ArraySpace{N,T}}) where {N,T} = T
similar_space(::Type{ArraySpace{N,T}}, ::Type{S}) where {N,T,S} = ArraySpace{N,S}

## Convenience functions

^(::Type{GeometricSpace{T}}, ::Type{Val{N}}) where {T <: Number,N} = VectorSpace{N,T}

## Some standard spaces

"The set of integers of type Int (ℤ = \\BbbZ)."
const ℤ = IntegerSpace{Int}
"The set of rational numbers of type Rational{Int} (ℚ = \\BbbQ)."
const ℚ = RationalSpace{Int}
"The set of reals of type Float64 (ℝ = \\BbbR)."
const ℝ = RealSpace{Float64}
"The complex plane with Float64 real and imaginar parts (ℂ = \\BbbC)."
const ℂ = ComplexSpace{Float64}

"The space ℝ^1"
const ℝ1 = ℝ^Val{1}
"The space ℝ^2"
const ℝ2 = ℝ^Val{2}
"The space ℝ^3"
const ℝ3 = ℝ^Val{3}
"The space ℝ^4"
const ℝ4 = ℝ^Val{4}

##############################
# Embeddings and conversions
##############################

## Isomorphism between scalars and 1D vectors

isomorphism_reduction(::Type{GeometricSpace{T}}, ::Type{VectorSpace{1,S}}) where {T <: Number,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(::Type{GeometricSpace{T}}, x::SVector{1,T}) where {T} = x[1]
convert_space(::Type{VectorSpace{1,T}}, x::T) where {T} = SVector(x)

## Isomorphism between the complex plane and ℝ^2

isomorphism_reduction(::Type{VectorSpace{2,T}}, ::Type{ComplexSpace{S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(::Type{ComplexSpace{T}}, x::SVector{2,T}) where {T} = x[1] + im*x[2]
convert_space(::Type{VectorSpace{2,T}}, x::Complex{T}) where {T} = SVector(real(x), imag(x))

## Isomorphism between vector spaces of the same dimension

# Normally, this line never gets called when T equals S, because that case is
# caught higher up in the algorithms.
isomorphism_reduction(::Type{VectorSpace{N,T}}, ::Type{VectorSpace{N,S}}) where {N,T,S} =
    (GeometricSpace{T},GeometricSpace{S})

# Just in case, we make sure if it ever happens that we notice
isomorphism_reduction(::Type{VectorSpace{N,T}}, ::Type{VectorSpace{N,T}}) where {N,T} =
    error("isomorphism_reduction was called with identical types")

convert_space(A::Type{VectorSpace{N,T}}, x::SVector{N,T}) where {N,T} = x
convert_space(A::Type{VectorSpace{N,T}}, x::SVector{N,S}) where {N,T,S} =
    SVector{N,T}(convert_space.(GeometricSpace{T}, x))

restrict_space(A::Type{VectorSpace{N,T}}, x::SVector{N,T}) where {N,T} = x
restrict_space(A::Type{VectorSpace{N,T}}, x::SVector{N,S}) where {N,T,S} =
    SVector{N,T}(restrict_space.(GeometricSpace{T}, x))

## For completeness, we have an isomorphism rule for ComplexSpace and RationalSpace
## to themselves as well. These are probably not useful.

isomorphism_reduction(::Type{ComplexSpace{T}}, ::Type{ComplexSpace{S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(A::Type{ComplexSpace{T}}, x::Complex{T}) where {T} = x
convert_space(A::Type{ComplexSpace{T}}, x::Complex{S}) where {T,S} =
    convert_space(GeometricSpace{T}, real(x)) + im * convert_space(GeometricSpace{T}, imag(x))

restrict_space(A::Type{ComplexSpace{T}}, x::Complex{T}) where {T} = x
restrict_space(A::Type{ComplexSpace{T}}, x::Complex{S}) where {T,S} =
    restrict_space(GeometricSpace{T}, real(x)) + im * restrict_space(GeometricSpace{T}, imag(x))

isomorphism_reduction(::Type{RationalSpace{T}}, ::Type{RationalSpace{S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(A::Type{RationalSpace{T}}, x::Rational{T}) where {T} = x
convert_space(A::Type{RationalSpace{T}}, x::Rational{S}) where {T,S} =
    convert_space(GeometricSpace{T}, numerator(x)) // convert_space(GeometricSpace{T}, denominator(x))

restrict_space(A::Type{RationalSpace{T}}, x::Rational{T}) where {T} = x
restrict_space(A::Type{RationalSpace{T}}, x::Rational{S}) where {T,S} =
    restrict_space(GeometricSpace{T}, numerator(x)) // restrict_space(GeometricSpace{T}, denominator(x))

## Embedding of R^N in R^(N+1)

# We only go up to dimension four
embedding_reduction(::Type{VectorSpace{1,T}}, ::Type{VectorSpace{2,S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})
embedding_reduction(::Type{VectorSpace{2,T}}, ::Type{VectorSpace{3,S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})
embedding_reduction(::Type{VectorSpace{3,T}}, ::Type{VectorSpace{4,S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(::Type{VectorSpace{2,T}}, x::SVector{1,T}) where {T} = SVector(x[1], 0)
convert_space(::Type{VectorSpace{3,T}}, x::SVector{2,T}) where {T} = SVector(x[1], x[2], 0)
convert_space(::Type{VectorSpace{4,T}}, x::SVector{3,T}) where {T} = SVector(x[1], x[2], x[3], 0)

# Since these are true embeddings, we also provide a left inverse
restrict_space(::Type{VectorSpace{1,T}}, x::SVector{2,T}) where {T} = SVector(x[1])
restrict_space(::Type{VectorSpace{2,T}}, x::SVector{3,T}) where {T} = SVector(x[1], x[2])
restrict_space(::Type{VectorSpace{3,T}}, x::SVector{4,T}) where {T} = SVector(x[1], x[2], x[3])

# Superspaces of a vector space have one higher dimension
superspace(::Type{VectorSpace{1,T}}) where {T} = VectorSpace{2,T}
superspace(::Type{VectorSpace{2,T}}) where {T} = VectorSpace{3,T}
superspace(::Type{VectorSpace{3,T}}) where {T} = VectorSpace{4,T}


module Visualizations
export
    plot_tfbdry,
    wiggle,
    wiggle!

using
    Plots,
    Statistics
using
    ..Utils

"""
    treenodes_matrix(x)

Given a `BitVector` of nodes in a binary tree, output the matrix representation of the
nodes.

# Arguments
- `x::BitVector`: BitVector representing a binary tree, where an input is 1 if the
  corresponding node exists and has children, and 0 if the corresponding node does not
  exist/does not have children.

# Returns
`::BitMatrix`: BitMatrix representation of the tree, where each column corresponds to a
level of the binary tree. The inputs in the matrix are 1 if the corresponding node exists
and has children, and 0 if the corresponding node does not exist/does not have children.

# Visualization
```
# Binary tree
      x
     / \\
    x   o
   / \\
  x   o
 / \\
o   o

# BitVector representation
[1,1,0,1,0,0,0]

# BitMatrix representation
[1 1 1 1;
 1 1 0 0;
 1 0 0 0]
```

# Examples
```@repl
using Wavelets, WaveletsExt

tree = maketree(8, 3, :dwt)
WaveletsExt.Visualizations.treenodes_matrix(tree)
```
"""
function treenodes_matrix(x::BitVector)
    n = (length(x) + 1) >> 1
    L = floor(Integer, log2(n))

    result = falses(L+1, n)
    i = 1
    for lvl in 0:L
        c = 1 << lvl                  # number of nodes at current level
        n₀ = nodelength(n, lvl)       # length of nodes at current level
        for node in 0:(c - 1)
            rng = (node * n₀ + 1):((node + 1) * n₀)
            result[lvl+1, rng] = repeat([x[i]], n₀)
            i += 1
        end
    end

    return result
end

"""
    plot_tfbdry(tree[; start, nd_col, ln_col, bg_col])

Given a tree, output a visual representation of the leaf nodes, user will have the option to
start the node count of each level with 0 or 1.

# Arguments
- `tree::BitVector`: Tree for plotting the leaf nodes. Comes in the form of a `BitVector`.
- `depth::Integer`: (Default: `log2(length(tree)+1)-1 |> Int`) Maximum depth to be
  displayed.

# Keyword Arguments
- `start::Integer`: (Default: `0`) Whether to zero-index or one-index the root of the tree.
- `nd_col::Symbol`: (Default: `:white`) Color of the leaf nodes.
- `ln_col::Symbol`: (Default: `:white`) Color of lines in plot.
- `bg_col::Symbol`: (Default: `:black`) Color of background.

# Returns
`::Plots.Plot`: Plot object with the visual representation of the leaf nodes.

# Examples
```julia
using Wavelets, WaveletsExt

# Build a tree using Wavelets `maketree`
tree = maketree(128, 7, :dwt)

# Plot the leaf nodes
plot_tfbdry(tree)
```
"""
function plot_tfbdry(tree::BitVector,
                     depth::Integer = log2(length(tree)+1)-1 |> Int;
                     start::Integer = 0,
                     nd_col::Symbol = :white,
                     ln_col::Symbol = :white,
                     bg_col::Symbol = :black)

    @assert 0 <= start <= 1
    leaf = getleaf(tree,:binary)

    ncol = 1 << depth
    nrow = depth+1
    mat = treenodes_matrix(leaf[1:(1<<(depth+1)-1)])

    p = heatmap(0:(ncol-1), start:(nrow+start-1), mat, color = [bg_col, nd_col],
                legend = false, background_color = bg_col)

    plot!(p, xlims = (start-0.5, ncol+start-0.5), ylims = (-0.5, nrow-0.5), yticks = 0:nrow)
    # plot horizontal lines
    for j in 0:(nrow-1)
        x_rng = (start-0.5):(ncol+start-0.5)
        y_val = (j-0.5)*ones(ncol+1)
        plot!(p, x_rng, y_val, color = ln_col, legend = false)
    end

    # plot vertical lines
    for j in 1:(nrow-1)
        for jj in 1:2^(j-1)
            x_val = (ncol/2^j)*(2*jj-1)+start-0.5;
            y_rng = j-0.5:nrow-0.5
            plot!(p, x_val*ones(nrow-j+1), y_rng, color = ln_col)
        end
    end
    plot!(p, (ncol+start-0.5)*ones(nrow+1), -0.5:nrow-0.5, color = ln_col)
    plot!(p, (start-0.5)*ones(nrow+1), -0.5:nrow-0.5, color = ln_col)
    plot!(p, yaxis = :flip)

    return p
end

"""
    wiggle(wav[; taxis, zaxis, sc, EdgeColor, FaceColor, Orient, Overlap, ZDir])
    wiggle(plt, wav[; taxis, zaxis, sc, EdgeColor, FaceColor, Orient, Overlap, ZDir])

Plots a set of shaded wiggles.

# Arguments
- `plt::Plots.Plot`: Input plot to plot shaded wiggles.
- `wav::AbstractArray{<:Number,2}`: Matrix of waveform columns.

# Keyword Arguments
- `taxis::AbstractVector`: (Default: `1:size(wav,1)`) Time axis vector
- `zaxis::AbstractVector`: (Default: `1:size(wav,2)`) Space axis vector
- `sc::Real`: (Default: `1`) Scale factor/magnification.
- `EdgeColor::Symbol`: (Default: `:black`) Sets edge of wiggles color.
- `FaceColor::Symbol`: (Default: `:black`) Sets shading color of wiggles.
- `Overlap::Bool`: (Default: `true`) How signals are scaled.
    - `true`  - Signals overlap (default);
    - `false` - Signals are scaled so they do not overlap.
- `Orient::Symbol`: (Default: `:across`) Controls orientation of wiggles.
    - `:across` - from left to right
    - `:down`   - from top to down
- `ZDir::Symbol`: (Default: `:normal`) Direction of space axis.
    - `:normal`  - First signal at bottom (default)
    - `:reverse` - First signal at top.

# Returns
`::Plots.Plot`: Shaded wiggles on top of current plot object.

# Examples
```julia
using Plots, WaveletsExt

# Generate random signals
x = randn(16, 5)

# ----- Build wiggles -----
# Method 1
wiggle(x)

# Method 2
p = Plot()
wiggle(p, x)
```

Translated by Nicholas Hausch -- MATLAB file provided by Naoki Saito. The previous MATLAB
version contributors are Anthony K. Booer (SLB) and Bradley Marchand (NSWC-PC).

Revised by Naoki Saito, Feb. 05, 2018. Maintained by Zeng Fung Liew for newest Julia version
compatibility.

**See also:** [`wiggle!`](@ref)
"""
wiggle(args...; kwargs...) = wiggle!(Plots.Plot(), args...; kwargs...)
wiggle(plt::Plots.Plot, args...; kwargs...) = wiggle!(deepcopy(plt), args...; kwargs...)

"""
    wiggle!(wav[; taxis, zaxis, sc, EdgeColor, FaceColor, Orient, Overlap, ZDir])
    wiggle!(plt, wav[; taxis, zaxis, sc, EdgeColor, FaceColor, Orient, Overlap, ZDir])

Plot a set of shaded wiggles on the current displayed graphics or on top of `plt`. If there
are no displayed graphics currently available, a new `Plots.Plot` object is generated to
plot the shaded wiggles.

# Arguments
- `plt::Plots.Plot`: Input plot to plot shaded wiggles.
- `wav::AbstractArray{<:Number,2}`: Matrix of waveform columns.

# Keyword Arguments
- `taxis::AbstractVector`: (Default: `1:size(wav,1)`) Time axis vector
- `zaxis::AbstractVector`: (Default: `1:size(wav,2)`) Space axis vector
- `sc::Real`: (Default: `1`) Scale factor/magnification.
- `EdgeColor::Symbol`: (Default: `:black`) Sets edge of wiggles color.
- `FaceColor::Symbol`: (Default: `:black`) Sets shading color of wiggles.
- `Overlap::Bool`: (Default: `true`) How signals are scaled.
    - `true`  - Signals overlap (default);
    - `false` - Signals are scaled so they do not overlap.
- `Orient::Symbol`: (Default: `:across`) Controls orientation of wiggles.
    - `:across` - from left to right
    - `:down`   - from top to down
- `ZDir::Symbol`: (Default: `:normal`) Direction of space axis.
    - `:normal`  - First signal at bottom (default)
    - `:reverse` - First signal at top.

# Returns
`::Plots.Plot`: Shaded wiggles on top of current plot object.

# Examples
```julia
using Plots, WaveletsExt

# Generate random signals
x = randn(16, 5)

# ----- Build wiggles -----
# Build onto existing plot
plt = plot()
wiggle!(x)

# Build onto a specified plot
wiggle!(plt, x)
```

Translated by Nicholas Hausch -- MATLAB file provided by Naoki Saito. The previous MATLAB
version contributors are Anthony K. Booer (SLB) and Bradley Marchand (NSWC-PC).

Revised by Naoki Saito, Feb. 05, 2018. Maintained by Zeng Fung Liew for newest Julia version
compatibility.

**See also:** [`wiggle`](@ref)
"""
function wiggle!(args...; kwargs...)
    @nospecialize
    local plt
    try
        plt = current()
    catch
        return wiggle(args..., kwargs...)
    end
    wiggle!(current(), args...; kwargs...)
end

function wiggle!(plt::Plots.Plot,
                 wav::AbstractArray{T,2};
                 taxis::AbstractVector=1:size(wav,1),
                 zaxis::AbstractVector=1:size(wav,2),
                 sc::Real=1,
                 EdgeColor::Symbol=:black,
                 FaceColor::Symbol=:black,
                 Overlap::Bool=true,
                 Orient::Symbol=:across,
                 ZDir::Symbol=:normal) where T<:Number
    # Set axes
    (n, m) = size(wav)

    # Sanity check
    @assert Orient ∈ [:across, :down]
    @assert ZDir ∈ [:normal, :reverse]
    if length(taxis) != n
        error("Inconsistent taxis dimension!")
    end
    if length(zaxis) != m
        error("Inconsistent zaxis dimension!")
    end

    # For calculation purposes
    maxrow = zeros(m)
    minrow = zeros(m)
    for (i, wavᵢ) in enumerate(eachcol(wav))
        maxrow[i] = maximum(wavᵢ)
        minrow[i] = minimum(wavᵢ)
    end

    # Scale the data for plotting
    dz = mean(diff(zaxis))
    if Overlap
        wamp = 2 * dz * (sc/maximum(maxrow-minrow)) * wav
    else
        wmax = maximum(maxrow) <= 0 ? 0 : maximum(maxrow)
        wmin = minimum(minrow) >= 0 ? 0 : minimum(minrow)
        wamp = sc*wav/(wmax-wmin)
    end

    # Set initial plot
    t0 = minimum(taxis)
    t1 = maximum(taxis)
    z0 = minimum(zaxis)
    z1 = maximum(zaxis)
    if Orient == :down
        plot!(plt, xlims=(z0-dz,z1+dz), ylims=(t0,t1), yflip=true, legend=:none)
    else
        plot!(plt, xlims=(t0,t1), ylims=(z0-dz,z1+dz), legend=:none)
    end
    if ZDir == :reverse
        wamp = reverse(wamp, dims=2)
    end

    # Plot each wavelet
    for (i, wampᵢ) in enumerate(eachcol(wamp))
        t = deepcopy(taxis)
        w_sign = sign.(wampᵢ)
        for j in 1:(n-1)
            if (w_sign[j]!=w_sign[j+1] && w_sign[j]!=0 && w_sign[j+1]!=0)
                wampᵢ = [wampᵢ; 0]
                t = [t; t[j]-wampᵢ[j]*(t[j+1]-t[j])/(wampᵢ[j+1]-wampᵢ[j])]
            end
        end
        ix = sortperm(t)
        t = t[ix]
        wampᵢ = wampᵢ[ix]
        len = length(t)
        indperm = [1:len; len:-1:1]
        inputx = t[indperm]
        inputy = zaxis[i] .+ [wampᵢ; min.(wampᵢ[len:-1:1],0)]
        # In the plot! functions below, theoretically speaking, either fillrange = zaxis[k]
        # or fillrange=[zaxis[k], zaxis[k]+dz] should be used. However, those do not
        # generate the desired plots as of O3/19/2018. Somehow, the relative value of 0,
        # i.e., fillrange=0, works well, which is used temporarily.
        if Orient == :down
            plot!(plt, inputy, inputx, fillrange=0, fillalpha=0.75,
                  fillcolor=FaceColor, linecolor=EdgeColor, orientation=:v)
        else
            plot!(plt, inputx, inputy, fillrange=0, fillalpha=0.75,
                  fillcolor=FaceColor, linecolor=EdgeColor)
        end
    end
    return plt
end

end # end module


struct PINningSparseSynapseParameter
end

@snn_kw mutable struct PINningSparseSynapse{VIT=Vector{Int32},VFT=Vector{Float32}}
    param::PINningSparseSynapseParameter = PINningSparseSynapseParameter()
    colptr::VIT # column pointer of sparse W
    I::VIT      # postsynaptic index of W
    W::VFT  # synaptic weight
    rI::VFT # postsynaptic rate
    rJ::VFT # presynaptic rate
    g::VFT  # postsynaptic conductance
    P::VFT  # <rᵢrⱼ>⁻¹
    q::VFT  # P * r
    f::VFT  # postsynaptic traget
    records::Dict = Dict()
end

"""
[PINing Sparse Synapse](https://www.ncbi.nlm.nih.gov/pubmed/26971945)
"""
PINningSparseSynapse

function PINningSparseSynapse(pre, post; σ = 1.5, p = 0.0, α = 1, kwargs...)
    w = σ / √(p * pre.N) * sprandn(post.N, pre.N, p)
    rowptr, colptr, I, J, index, W = dsparse(w)
    rI, rJ, g = post.r, pre.r, post.g
    P = α .* (I .== J)
    f, q = zeros(post.N), zeros(post.N)
    PINningSparseSynapse(;@symdict(colptr, I, W, rI, rJ, g, P, q, f)..., kwargs...)
end

function forward!(c::PINningSparseSynapse, param::PINningSparseSynapseParameter)
    @unpack colptr, I, W, rJ, g, P, q = c
    fill!(q, zero(Float32))
    fill!(g, zero(Float32))
    @inbounds for j in 1:(length(colptr) - 1)
        rJj = rJ[j]
        for s = colptr[j]:(colptr[j+1] - 1)
            i = I[s]
            g[i] += W[s] * rJj
            q[i] += P[s] * rJj
        end
    end
end

function plasticity!(c::PINningSparseSynapse, param::PINningSparseSynapseParameter, dt::Float32, t::Float32)
    @unpack colptr, I, W, rI, g, P, q, f = c
    C = 1 / (1 + dot(q, rI))
    @inbounds for j in 1:(length(colptr) - 1)
        for s in colptr[j]:(colptr[j+1] - 1)
            i = I[s]
            P[s] += -C * q[i] * q[j]
            W[s] += C * (f[i] - g[i]) * q[j]
        end
    end
end


module test_SolveOneNonlinearEquation

import ..ModiaMath: solveOneNonlinearEquation
using Test
@eval using Printf

@testset "ModiaMath.NonlinearEquations: test solveOneNonlinearEquation" begin
    # Test from Modelica.Math.Nonlinear.Examples.solveNonlinearEquations1
    y          = zeros(3)
    y_analytic = zeros(3)

    fun1(u)       = u^2 - 1
    y[1]          = solveOneNonlinearEquation(fun1, -0.5, 10.0)
    y_analytic[1] = 1.0

    fun2(u;w=1.0) = 3*u - sin(w*u) - 1
    y[2]          = solveOneNonlinearEquation(u->fun2(u; w=3.0), 0.0, 5.0)
    y_analytic[2] = 0.6448544035840080891877538

    fun3(u; p=zeros(2), m=0.0) = p[1] + log(p[2]*u) - m*u
    y[3]          = solveOneNonlinearEquation(u->fun3(u; p=[5.0,1.0], m=1.0), 1.0, 8.0)
    y_analytic[3] = 6.9368474072202187221643182

    for i in 1:length(y)
        println("\n... Results of Solve_SingleNonlinearEquations:")
        println("fun", string(i), ":")
        @printf("   analytical zero     = %20.16e\n", y_analytic[i])
        @printf("   numerical zero      = %20.16e\n", y[i])
        @printf("   absolute difference = %20.16e\n", y_analytic[i] - y[i])
        @test isapprox(y[i], y_analytic[i],  atol=1e-12)
    end

    # Test error case
    errorOccured = false
    try
       yy = solveOneNonlinearEquation(u->fun2(u; w=3.0), 0.0, 0.5)
    catch
       errorOccured = true
    end
    @test errorOccured

    # Test tolerance
    yy = solveOneNonlinearEquation(fun1, -0.5, 10.0; tolerance=1e-4)
    @test isapprox(yy, 1.0,  atol=1e-4)

    # Tests from https://github.com/JuliaMath/Roots.jl/test/test_find_zero.jl
    fns = [(x -> x^2 - exp(x) - 3x + 2,   0.257530285439860, -1, 0.5),
           (x -> x^3 + 4x^2 -10,          1.365230013414097, 0.5, 2.0),
           (x -> exp(x)*sin(x) - 2x - 5, -2.523245230732555, -3.0, -1.0),
           (x -> log(x) + sqrt(x) - 5,    8.309432694231571,  7, 9),
           (x -> sqrt(x) - 1/x - 3,       9.633595562832696,  8, 10),
           (x -> exp(-x) - cos(x), 1.292695719373398, 1.0, 2.0),
           (x -> cos(x)^2 - x/5, 1.085982678007472, 0.5, 2.0),
           (x -> x^10 - x^3 - x - 1, -0.674177935277052, -1.0, 0.0),
           (x -> sin(x) - x + 1, 1.934563210752024, 1.0, 20.0),
           (x -> exp(-x^2 + x + 2) - cos(x) + x^3 + 1, -1.0899423334391694, -2, 0),
           (x -> asin(x^2 -1) - x/2 + 1, 0.594810968398369, 0.3, 0.9),
           (x -> tanh(x) - tan(x), 7.068582745628732, 5.5, 7.1)
    ]

    for (i, (f, u_exact, u_min, u_max)) in enumerate(fns)
        u = solveOneNonlinearEquation(f, u_min, u_max)
        @test isapprox(u, u_exact,  atol=1e-12)
    end

    # Some more tests from https://github.com/JuliaMath/Roots.jl/test/test_find_zero.jl
    # based on http://ir.igsnrr.ac.cn/bitstream/311030/8840/1/%E4%BE%AF%E9%BA%9F%E7%A7%91(SCI)2.pdf
    # which have multiple roots
    multiplicity_tests = [
         (x -> (x - sqrt(5))^4 / ((x-1)^2 + 2),    2.236067977499790, 0, 3)
         (x -> (sin(x)^2 - 2x + 1)^5,              0.71483582544138924, 0, 1)
         (x -> (8x*exp(-x^2) -2x - 3)^8,           -1.7903531791589544, -2, -1)
         (x -> (2x*cos(x) + x^2 - 3)^10/(x^2 + 1),  2.9806452794385368, 0, 3)
         (x -> (exp(-x^2 + x + 3) - x + 2)^9,       2.4905398276083051, 0, 4)

         (x -> (exp(-x) + 2sin(x))^4,              3.1627488709263654, 2, 4)
         (x -> (log(x^2 + 3x + 5) - 2x + 7)^8,     5.4690123359101421, 3, 6)
         (x -> (sqrt(x^2 + 2x + 5) - 2sin(x) - x^2 + 3)^5,   2.3319676558839640, 0, 4)
         (x -> (x-2)^4 / ( (x-1)^2 + 1),           2.0000000000000000, 0, 4)
         (x -> abs(x - 2.5)^(15/4) * exp(x),       2.5, 0, 4)

         (x -> (sqrt(x) - 1/x - 1)^7,              2.147899035704787, 0, 4)
         (x -> (log(x) + sqrt(x) - 5)^3,           8.309432694231572, 0, 10)
         (x -> (sin(x)*cos(x) - x^3 + 1)^9,        1.117078770687451, 0, 2)
         (x -> ((x-3)*exp(x))^5,                   3.0000000000000000, 0, 5)
         (x -> (log(x) + sqrt(x^4 + 1) -2)^7,      1.222813963628973, 1, 3)
         ]

    for (i, (f, u_exact, u_min, u_max)) in enumerate(fns)
        u = solveOneNonlinearEquation(f, u_min, u_max)
        @test isapprox(u, u_exact,  atol=1e-12)
    end
end

end


export NLfunForward_test1

"""
weighted forward function
"""
function NLfunForward_test1(A, x, g, params)

    e = params["numr"]*params["nr"]
    w = params["w"]

    e1 = params["numr"]*params["nr1"]
    W_data = params["weight"]
    U_t = copy(W_data[1:e1,1])
    V_t = copy(W_data[e1+1:end,1])
    U = reshape(copy(U_t),params["numr"], params["nr1"])
    V = reshape(copy(V_t),params["numc"], params["nr1"])

    PU = U*U'
    PV = V*V'
    m_U::Int,n_U::Int = size(PU)
    m_V::Int,n_V::Int = size(PV)

    PU_org = Matrix{Complex{Float64}}(I,m_U,n_U)-copy(PU)
    PV_org = Matrix{Complex{Float64}}(I,m_V,n_V)-copy(PV)

    QU = (1/w)*PU+PU_org
    QV = (1/w)*PV+PV_org

    L = x[1:e]
    R = x[(e+1):end]
    L = reshape(L, params["numr"], params["nr"])
    R = reshape(R, params["numc"], params["nr"])

    if isempty(g)
        f1 = params["afun"](QU*L*(QV*R)', params)
        f2 = 0.
    else
        fp = params["afunT"](g)
        f1 = [vec(QU*fp*(QV*R)); vec(QV*fp'*(QU*L))]
        f2 = vec(QU*fp*QV)
    end

    return f1,f2
end


"""
    Duality(optimizer)

Duality uses Lagrangian relaxation to compute over-approximated bounds for a network

# Problem requirement
1. Network: any depth, any activation function that is monotone
2. Input: hyperrectangle
3. Output: halfspace

# Return
`BasicResult`

# Method
Lagrangian relaxation.
Default `optimizer` is `GLPKSolverMIP()`.

# Property
Sound but not complete.

# Reference
[K. Dvijotham, R. Stanforth, S. Gowal, T. Mann, and P. Kohli,
"A Dual Approach to Scalable Verification of Deep Networks,"
*ArXiv Preprint ArXiv:1803.06567*, 2018.](https://arxiv.org/abs/1803.06567)
"""
@with_kw struct Duality <: Solver
    optimizer = GLPK.Optimizer
end

function solve(solver::Duality, problem::Problem)

    @assert is_halfspace_equivalent(problem.output) "Duality only accepts HalfSpace output sets. Got a $(length(d))-d polytope."

    model = Model(solver)
    c, d = tosimplehrep(problem.output)
    d = d[1]

    λ = init_vars(model, problem.network, :λ)
    μ = init_vars(model, problem.network, :μ)
    o = dual_value(solver, problem, model, λ, μ)
    @constraint(model, last(λ) .== -c)
    optimize!(model)

    # Interpret result:
    if termination_status(model) != OPTIMAL
        return BasicResult(:unknown)
    elseif value(o) - d <= 0.0
        return BasicResult(:holds)
    else
        return BasicResult(:violated)
    end
end

function dual_value(solver::Duality,
                    problem::Problem,
                    model::Model,
                    λ::Vector{Vector{VariableRef}},
                    μ::Vector{Vector{VariableRef}})
    bounds = get_bounds(problem)
    layers = problem.network.layers

    # prepend an array of 0s to λ so the λᵢ
    # term cancels in the input layer
    λ = [zeros(dim(bounds[1])), λ...]

    o = 0
    for i in 1:length(layers)
        Lᵢ, Bᵢ = layers[i], bounds[i]
        W, b = Lᵢ.weights, Lᵢ.bias
        c, r = Bᵢ.center, Bᵢ.radius
        B̂ᵢ₊₁ = approximate_affine_map(Lᵢ, Bᵢ)

        # layer value
        o += λ[i]'*c - μ[i]'*(W*c + b) + sum(symbolic_abs.(λ[i] .- W'*μ[i]) .* r)

        # activation value
        o += activation_value(Lᵢ.activation, μ[i], λ[i+1], low(B̂ᵢ₊₁), high(B̂ᵢ₊₁))
    end

    @objective(model, Min, o)
    return o
end

function activation_value(σ::ReLU, μᵢ, λᵢ, l̂ᵢ, ûᵢ)

    gᵢl̂ᵢ = @. μᵢ*l̂ᵢ - λᵢ*σ(l̂ᵢ)
    gᵢûᵢ = @. μᵢ*ûᵢ - λᵢ*σ(ûᵢ)

    max = symbolic_max

    # NOTE: ifelse evaluates all of its arguments in advance.
    # This may seem like a bug since symbolic_max mutates the model
    # and introduces new variables, but it is fine. Only those
    # variables that are part of the objective will affect the problem,
    # and those will only the the ones summed *after* taking the ifelse
    return sum(@. ifelse(l̂ᵢ < 0 < ûᵢ,
                         max(gᵢl̂ᵢ, gᵢûᵢ, 0),
                         max(gᵢl̂ᵢ, gᵢûᵢ)))
end

function activation_value(σ::Id, μᵢ, λᵢ, l̂ᵢ, ûᵢ)
    max = symbolic_max
    sum(@. max(μᵢ*l̂ᵢ - λᵢ*l̂ᵢ, μᵢ*ûᵢ - λᵢ*ûᵢ))
end

function activation_value(σ::Any, μᵢ, λᵢ, l̂ᵢ, ûᵢ)
    max = symbolic_max
    sum(@. max(μᵢ*l̂ᵢ, μᵢ*ûᵢ) + max(-λᵢ*σ(l̂ᵢ), -λᵢ*σ(ûᵢ)))
end


# test basic forward convolve, see IIF issue #477

# using Revise

using Test
using IncrementalInference
using Statistics

@testset "Test basic convolution result..." begin

function forwardConvolve(X0::Array{Float64,2}, model)

  fg = initfg()

  addVariable!(fg, :x0, ContinuousScalar)
  initManual!(fg, :x0, X0)
  addVariable!(fg, :x1, ContinuousScalar)

  addFactor!(fg, [:x0;:x1], model)

  ## TODO -- dont use name here, add API to just use z2 here
  return approxConv(fg, :x0x1f1, :x1)
end

## Start

# first numerical values -- samples from the marginal of X0
z1 = Normal(0,0.1)
X0 = rand(z1, 1,100)

## predict -- project / conv
# 0 -> 1 seconds
# make approx function
z2 = Normal(11,1.0) # odo
statemodel = LinearRelative( z2 )
X1_ = forwardConvolve(X0, statemodel)

## measure -- product of beliefs, using `ApproxManifoldProducts.jl`

predX1 = manikde!(X1_, ContinuousScalar)
z3 = Normal(9.5,0.75)
measX1 = manikde!(reshape(rand(z3,100),1,:), ContinuousScalar)

# do actual product
posterioriX1 = predX1 * measX1
X1 = getPoints(posterioriX1)

## predict, 1->2 seconds
z4 = Normal(8,2.0) # odo
statemodel = LinearRelative( z4 )
X2_ = forwardConvolve(X1, statemodel)

@test size(X2_) == (1,100)
@test 15 < Statistics.mean(X2_) < 25

end


"""
    SGD(::Vector{XVariable}; lr=1e-4, L1decay=0.001, L2decay=0.01)

Implements stochastic gradient descent
"""
mutable struct SGD <: Optimizer
    xparams::Vector{XVariable}
    lr::AbstractFloat
    L1decay::AbstractFloat
    L2decay::AbstractFloat
    name::String
    function SGD(xparams::Vector{XVariable}; lr=1e-4, L1decay=0.001, L2decay=0.01)
        new(xparams, lr, L1decay, L2decay, "SGD")
    end
end

# pretty printing
function Base.show(io::IO, O::SGD)
    print("SGD(lr=$(O.lr), L1decay=$(O.L1decay), L2decay=$(O.L2decay))")
end

function update!(O::SGD; clipfn::Function=LPInfNormClip, clipvalue=10.0)
    μ = - O.lr
    λ₁ = O.L1decay
    λ₂ = O.L2decay

    for i = 1:length(O.xparams)
        c , θ = O.xparams[i]
        ∇ = clipfn(setNanInfZero(δ(θ)), clipvalue)
        𝒗 = ᵛ(θ)
        if c == 'w'
            if λ₁==0 && λ₂==0
                @. 𝒗 += μ * ∇
            elseif λ₁==0 && λ₂!=0
                @. 𝒗 += μ * (∇ + λ₂ * 𝒗)
            elseif λ₁!=0 && λ₂==0
                @. 𝒗 += μ * (∇ + λ₁ * sign(𝒗))
            else  # λ₁!=0 && λ₂!=0
                @. 𝒗 += μ * (∇ + λ₁ * sign(𝒗) + λ₂ * 𝒗)
            end
        else
            @. 𝒗 += μ * ∇
        end
    end
end


"""
    block_jackknife(Y::JArray{Float64,2}, subsample::Float64)

Generate block jackknife samples as in Kunsch (1989).

This technique subsamples a time series dataset by removing, in turn, all the blocks of consecutive observations
with a given size.

# Arguments
- `Y`: observed measurements (`nxT`), where `n` and `T` are the number of series and observations.
- `subsample`: block size as a percentage of the original sample size. It is bounded between 0 and 1.

# References
Kunsch (1989) and Pellegrino (2020)
"""
function block_jackknife(Y::JArray{Float64,2}, subsample::Float64)

    # Check inputs
    check_bounds(subsample, 0, 1);

    # Dimensions
    n, T = size(Y);

    # Block size
    block_size = Int64(ceil(subsample*T));
    if block_size == 0
        error("subsample is too small!");
    end

    # Number of block jackknifes samples - as in Kunsch (1989)
    samples = T-block_size+1;

    # Initialise jackknife_data
    jackknife_data = JArray{Float64,3}(undef, n, T, samples);

    # Loop over j=1, ..., samples
    for j=1:samples

        # Index of missings
        ind_j = collect(j:j+block_size-1);

        # Block jackknife data
        jackknife_data[:, :, j] = Y;
        jackknife_data[:, ind_j, j] .= missing;
    end

    # Return jackknife_data
    return jackknife_data;
end

"""
    optimal_d(n::Int64, T::Int64)

Select the optimal value for d.

# Arguments
- `n`: Number of series
- `T`: Number of observations
"""
function optimal_d(n::Int64, T::Int64)
    objfun_array = zeros(n*T);

    for d=1:fld(n*T,2)
        objfun_array[d] = objfun_optimal_d(n, T, d);
    end

    return argmax(objfun_array);
end

"""
    objfun_optimal_d(n::Int64, T::Int64, d::Int64)

Objective function to select the optimal value for d.

# Arguments
- `n`: Number of series
- `T`: Number of observations
- `d`: Candidate d
"""
function objfun_optimal_d(n::Int64, T::Int64, d::Int64)
    fun = no_combinations(n*T, d) - (d>=n).*no_combinations(n*T-n, d-n)*T;

    for i=2:fld(d, n)
        fun -= (-1^(i-1))*no_combinations(T, i)*no_combinations(n*T-i*n, d-i*n);
    end

    return fun;
end

"""
    artificial_jackknife(Y::JArray{Float64,2}, subsample::Float64, max_samples::Int64)

Generate artificial jackknife samples as in Pellegrino (2020).

The artificial delete-d jackknife is an extension of the delete-d jackknife for dependent data problems.
This technique replaces the actual data removal step with a fictitious deletion, which consists of
imposing `d`-dimensional (artificial) patterns of missing observations to the data. This approach
does not alter the data order nor destroy the correlation structure.

# Arguments
- `Y`: observed measurements (`nxT`), where `n` and `T` are the number of series and observations.
- `subsample`: `d` as a percentage of the original sample size. It is bounded between 0 and 1.
- `max_samples`: if `C(n*T,d)` is large, artificial_jackknife would generate `max_samples` jackknife samples.

# References
Pellegrino (2020)
"""
function artificial_jackknife(Y::JArray{Float64,2}, subsample::Float64, max_samples::Int64)

    # Dimensions
    n, T = size(Y);
    nT = n*T;

    # Check inputs
    check_bounds(subsample, 0, 1);

    # Set d using subsample
    d = Int64(ceil(subsample*nT));

    # Error management
    if d == 0
        error("subsample is too small!");
    end

    # Warning
    if subsample > 0.5
        @warn "this algorithm might be unstable for `subsample` larger than 0.5!";
    end

    # Get vec(Y)
    vec_Y = convert(JArray{Float64}, Y[:]);

    # Initialise loop (controls)
    C_nT_d = no_combinations(nT, d);
    samples = convert(Int64, min(C_nT_d, max_samples));
    zeros_vec = zeros(d);

    # Initialise loop (output)
    ind_missings = Array{Int64}(zeros(d, samples));
    jackknife_data = JArray{Float64,3}(undef, n, T, samples);

    # Loop over j=1, ..., samples
    for j=1:samples

        # First draw
        if j == 1
            if samples == C_nT_d
                ind_missings[:,j] = rand_without_replacement(nT, d);
            else
                ind_missings[:,j] = rand_without_replacement(n, T, d);
            end

        # Iterates until ind_missings[:,j] is neither a vector of zeros, nor already included in ind_missings
        else
            while ind_missings[:,j] == zeros_vec || is_vector_in_matrix(ind_missings[:,j], ind_missings[:, 1:j-1])
                if samples == C_nT_d
                    ind_missings[:,j] = rand_without_replacement(nT, d);
                else
                    ind_missings[:,j] = rand_without_replacement(n, T, d);
                end
            end
        end

        # Add (artificial) missing observations
        vec_Y_j = copy(vec_Y);
        vec_Y_j[ind_missings[:,j]] .= missing;

        # Store data
        jackknife_data[:, :, j] = reshape(vec_Y_j, n, T);
    end

    # Return jackknife_data
    return jackknife_data;
end


#derivative logic

#ForwardDiff compiles one method per function, so separating the
#differentiation logic from the property logic allows the differentials
#to be compiled only once

"""
    ∂f∂t(model,v,t,z=SA[1.0])

returns `f` and `∂f/∂T` at constant total volume and composition, where f is the total helmholtz energy, given by `eos(model,v,T,z)`

"""
function ∂f∂t(model,v,t,z=SA[1.0])
    return ForwardDiff.derivative(∂t -> eos(model,v,∂t,z),t)
end

"""
    ∂f∂v(model,v,t,z=SA[1.0])

returns `f` and `∂f/∂v` at constant temperature and composition, where f is the total helmholtz energy, given by `eos(model,v,T,z)`, and v is the total volume

"""
function ∂f∂v(model,v,t,z)
    return ForwardDiff.derivative(∂v -> eos(model,∂v,t,z),v)
end

#returns a tuple of the form ([∂f∂v,∂f∂t],f),using the least amount of computation
"""
    ∂f(model,v,t,z)

returns zeroth order (value) and first order derivative information of the total helmholtz energy (given by `eos(model,v,T,z)`).
the result is given in two values:

```
grad_f,fval = ∂2f(model,v,t,z)
```

where:
```
fval   = f(v,T) = eos(model,v,T,z)

grad_f = [ ∂f/∂v; ∂f/∂T]

 ```

Where `v` is the total volume, `T` is the temperature and `f` is the total helmholtz energy.
"""
function ∂f(model,v,t,z)
    f(w) = eos(model,first(w),last(w),z)
    v,t = promote(v,t)
    vt_vec = SVector(v,t)
    ∂result = DiffResults.GradientResult(vt_vec)
    res_∂f =  ForwardDiff.gradient!(∂result, f,vt_vec)
    _f =  DiffResults.value(res_∂f)
    _∂f = DiffResults.gradient(res_∂f)
    return (_∂f,_f)
end

#returns p and ∂p∂v at constant t
#it doesnt do a pass over temperature, so its
#faster that d2f when only requiring d2fdv2

"""
    p∂p∂v(model,v,t,z=SA[1.0])

returns `p` and `∂p/∂v` at constant temperature, where p is the pressure = `pressure(model,v,T,z)` and `v` is the total volume.

"""
function p∂p∂v(model,v,t,z=SA[1.0])
    v_vec =   SVector(v)
    f(∂v) = pressure(model,only(∂v),t,z)
    ∂result = DiffResults.GradientResult(v_vec)
    res_∂f =  ForwardDiff.gradient!(∂result, f,v_vec)
    _p =  DiffResults.value(res_∂f)
    _∂p∂v = only(DiffResults.gradient(res_∂f))
    return _p,_∂p∂v
end

"""
    ∂2f(model,v,t,z)

returns zeroth order (value), first order and second order derivative information of the total helmholtz energy (given by `eos(model,v,T,z)`).
the result is given in three values:

```
hess_f,grad_f,fval = ∂2f(model,v,t,z)
```

where:
```
fval   = f(v,T) = eos(model,v,T,z)

grad_f = [ ∂f/∂v; ∂f/∂T]

hess_f = [ ∂²f/∂v²; ∂²f/∂v∂T
          ∂²f/∂v∂T; ∂²f/∂v²]
 ```

Where `v` is the total volume, `T` is the temperature and `f` is the total helmholtz energy.
"""
function ∂2f(model,v,t,z)
    f(w) = eos(model,first(w),last(w),z)
    v,t = promote(v,t)
    vt_vec =   SVector(v,t)
    ∂result = DiffResults.HessianResult(vt_vec)
    res_∂f =  ForwardDiff.hessian!(∂result, f,vt_vec)
    _f =  DiffResults.value(res_∂f)
    _∂f = DiffResults.gradient(res_∂f)
    _∂2f = DiffResults.hessian(res_∂f)
    return (_∂2f,_∂f,_f)
end

"""
    ∂2p(model,v,t,z)

returns zeroth order (value), first order and second order derivative information of the pressure.
the result is given in three values:

```
hess_p,grad_p,pval = ∂2p(model,v,t,z)
```

where:
```
pval   = p(v,T) = pressure(model,v,T,z)

grad_p = [ ∂p/∂v; ∂p/∂T]

hess_p = [ ∂²p/∂v²; ∂²p/∂v∂T
          ∂²p/∂v∂T; ∂²p/∂v²]
 ```

Where `v` is the total volume, `T` is the temperature and `p` is the pressure.
"""
function ∂2p(model,v,t,z)
    f(w) = pressure(model,first(w),last(w),z)
    v,t = promote(v,t)
    vt_vec =   SVector(v,t)
    ∂result = DiffResults.HessianResult(vt_vec)
    res_∂f =  ForwardDiff.hessian!(∂result, f,vt_vec)
    _f =  DiffResults.value(res_∂f)
    _∂f = DiffResults.gradient(res_∂f)
    _∂2f = DiffResults.hessian(res_∂f)
    return (_∂2f,_∂f,_f)
end

"""
    f_hess(model,v,t,z)

returns the second order volume (`v`) and temperature (`T`) derivatives of the total helmholtz energy (given by `eos(model,v,T,z)`). the result is given in a 2x2 `SMatrix`, in the form:

```
[ ∂²f/∂v²  ∂²f/∂v∂T
 ∂²f/∂v∂T  ∂²f/∂v²]
 ```

use this instead of the ∂2f if you only need second order information. ∂2f also gives zeroth and first order derivative information, but due to a bug in the used AD, it allocates more than necessary.
"""
function f_hess(model,v,t,z)
    f(w) = eos(model,first(w),last(w),z)
    v,t = promote(v,t)
    vt_vec = SVector(v,t)
    return ForwardDiff.hessian(f,vt_vec)
end

"""
    ∂p2∂p3(model,v,t,z=SA[1.0])

returns `∂²p/∂v²` and `∂³p/∂v³`, in a single ForwardDiff pass. used mainly in `crit_pure` objective function

"""
function ∂p2∂p3(model,v,t,z=SA[1.0])
    v_vec =   SVector(v)
    f(∂v) = ForwardDiff.derivative(∂2v -> pressure(model,only(∂2v),t,z),only(∂v))
    ∂result = DiffResults.GradientResult(v_vec)
    res_∂f =  ForwardDiff.gradient!(∂result, f,v_vec)
    _p =  DiffResults.value(res_∂f)
    _∂p∂v = only(DiffResults.gradient(res_∂f))
    return _p,_∂p∂v
end


#apply_weights(A::Array, W::Array) = sum(A .* W)
apply_weights(A::Array, W::Missing) = sum(A .* (1/length(A)))

@inline function orthogonal_sine(t::Real, j::Real, block_num::Integer, block_width::Real)
    # Note for simplicity and speed I am going to assume that we will never call this after (k+1)h_n. This is due to the way it will be calculated in the loop.
    start_of_block = block_num*block_width
    return  t > start_of_block ? ((sqrt(2 * block_width))/(j*pi)) *sin((j * pi * (t-start_of_block))/block_width) : 0.0
end

function spectral_lmm_array(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts); regularisation::Union{Missing,Symbol} = :covariance_default, regularisation_params::Dict = Dict(),
                          only_regulise_if_not_PSD::Bool = false, numJ::Integer = 100, num_blocks::Integer = 10, block_width::Real = (maximum(ts.df[:,ts.time])) / num_blocks,
                          microstructure_noise_var::Dict{Symbol,<:Real} = two_scales_volatility(ts, assets)[2])
    numAssets = length(assets)
    SS = zeros(num_blocks, numJ, numAssets)
    indices = Dict(assets .=> collect(1:numAssets))
    lasttimes = zeros(numAssets)
    lastvalues = zeros(numAssets)
    H_kn = zeros(num_blocks, numAssets)
    for i in 1:nrow(ts.df)
        asset = ts.df[i, ts.grouping]
        if (asset in assets) == false continue end
        assetslot = indices[asset]
        newtime = ts.df[i, ts.time]
        newval  = ts.df[i, ts.value]
        if lasttimes[assetslot] < eps()
            lasttimes[assetslot]  = newtime
            lastvalues[assetslot] = newval
            continue
        end
        timegap = (newtime - lasttimes[assetslot])
        centertime = lasttimes[assetslot] + 0.5*timegap
        ret = simple_differencing([newval], [lastvalues[assetslot]])[1]
        lasttimes[assetslot]  = newtime
        lastvalues[assetslot] = newval

        block_num = min(num_blocks-1,Int(floor(newtime/block_width)))

        # Looking over j values for the S Array.
        for j in 1:numJ
             SS[block_num+1, j, assetslot] += (pi * j/block_width) * ret * orthogonal_sine(centertime, j, block_num, block_width)
        end
        # Now doing the H_k^n matrix which is independent of j.
        H_kn[block_num+1, assetslot] += (timegap^2) * (microstructure_noise_var[asset]/block_width)
    end
    # Now we assemble the covariance matrices for every j,k.
    R = eltype(SS)
    uncorrected_covar_matrices = Array{Array{Hermitian{R},1},1}(undef,num_blocks)
    for block_num in 1:num_blocks
        hermitian_array = Array{Hermitian{R},1}(undef, numJ)
        for j in 1:numJ
            hermitian_array[j] = Hermitian( SS[block_num,j,:] * transpose(SS[block_num,j,:])  )
        end
        uncorrected_covar_matrices[block_num] = hermitian_array
    end

    # The efficient weights (see page 1324 of Bibinger, Hautsch, Malec & Reiss 2014) are not implemented here. It is highly error proone as
    # many inversions are required which is 1) computationally intensive and 2) can cause singularity errors as there is no guarantee of psd.
    # So going to fall back to equal weighting.
    weights = repeat([missing],num_blocks)

    corrected_matrices = Array{CovarianceMatrix{R},1}(undef,num_blocks)
    for block_num in 1:num_blocks
        weights_for_block = weights[block_num]
        block_H = Diagonal(H_kn[block_num,:])
        adjustments = ((pi .*  collect(1:numJ)) ./block_width).^2
        block_adjustments = Array{Diagonal,1}(undef,numJ)
        for nn in 1:numJ block_adjustments[nn] = adjustments[nn] * block_H end
        corrected_mats = uncorrected_covar_matrices[block_num] .- block_adjustments
        mat = Hermitian(apply_weights(corrected_mats, weights_for_block))

        # Regularisation
        dont_regulise = ismissing(regularisation) || (only_regulise_if_not_PSD && is_psd_matrix(mat))
        mat = dont_regulise ? mat : regularise(mat, ts, assets, regularisation; regularisation_params... )

        # In some cases we get negative terms on the diagonal with this algorithm.
        negative_diagonals = findall(diag(mat) .< eps())
        covar = make_nan_covariance_matrix(assets, ts.time_period_per_unit)
        if length(negative_diagonals) < 1
            corr, vols = cov2cor_and_vol(mat, 1) # No adjustment for duration as we already have the spot matrix.
            covar.correlation = corr
            covar.volatility = vols
         end
        corrected_matrices[block_num] = covar
    end
    return corrected_matrices
end

"""
    spectral_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                        regularisation::Union{Missing,Symbol} = :covariance_default,
                        regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                        numJ::Integer = 100, num_blocks::Integer = 10,
                        block_width::Real = (maximum(ts.df[:,ts.time]) - minimum(ts.df[:,ts.time])) / num_blocks,
                        microstructure_noise_var::Dict{Symbol,<:Real} = two_scales_volatility(ts, assets)[2])

Estimation of a CovarianceMatrix using the spectral covariance method.
### Inputs
* `ts` - The tick data.
* `assets` - The assets you want to estimate volatilities for.
* `regularisation` - A symbol representing what regularisation technique should be used. If missing no regularisation is performed.
* `regularisation_params` - keyword arguments to be consumed in the regularisation algorithm.
* `only_regulise_if_not_PSD` - Should regularisation only be attempted if the matrix is not psd already.
* `numJ` - The number of J values. See the paper for details.
* `num_blocks` - The number of blocks to split the time frame into. See the paper for details.
* `block_width` - The width of each block to split the time frame into.
* `microstructure_noise_var` - Estimates of microstructure noise variance for each asset.
### Returns
* A `CovarianceMatrix`.

### References
Bibinger M, Hautsch N, Malec P, Reiss M (2014). “Estimating the quadratic covariation matrix from noisy observations: Local method of moments and efficiency.” The Annals of Statistics, 42(4), 1312–1346. doi:10.1214/14-AOS1224.
"""
function spectral_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                             regularisation::Union{Missing,Symbol} = :covariance_default,
                             regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                             numJ::Integer = 100, num_blocks::Integer = 10,
                             block_width::Real = (maximum(ts.df[:,ts.time]) - minimum(ts.df[:,ts.time])) / num_blocks,
                             microstructure_noise_var::Dict{Symbol,<:Real} = two_scales_volatility(ts, assets)[2])
    corrected_matrices = spectral_lmm_array(ts, assets; regularisation = regularisation, regularisation_params = regularisation_params, only_regulise_if_not_PSD = only_regulise_if_not_PSD, numJ = numJ, num_blocks = num_blocks,
                                            block_width = block_width, microstructure_noise_var = microstructure_noise_var)
    cor_weights = repeat([1.0], length(corrected_matrices))
    spectral_estimate = combine_covariance_matrices(corrected_matrices, cor_weights)
    return spectral_estimate
end


## Exponential moving average
function EWMA(x)
    ewmaVar = zeros(length(x),1)
    ewmaVar[1] = var(x)
    for i in 2:length(x)
         ewmaVar[i] = 0.94*ewmaVar[i-1]+ (1.0 - 0.94)*x[i-1]^2.0
    end
    ewmaVol = sqrt.(ewmaVar);
    return ewmaVol
end

## OLS estimate
function OLSestimator(y,x)
    estimate = (transpose(x)*x) \ (transpose(x)*y)
    return estimate
end

## Function to generate polynoms
function delta3j(js)
    alphaj =  collect(linspace(0.05, 0.95, js))
    delta0 = collect(linspace(1.0, 1.0, length(alphaj)))
    delta1 = 2.0*(alphaj - 0.5)
    delta2 = (2.0^2.0)*((alphaj - 0.5).^2.0)
    delta3 = (2.0^3.0)*((alphaj - 0.5).^3.0)
    delta4 = (2.0^4.0)*((alphaj - 0.5).^4.0)
    delta5 = (2.0^5.0)*((alphaj - 0.5).^5.0)
    return [delta0 delta1 delta2 delta3 delta4 delta5]
end

## Fucntion to obtain initial estimates
function getinitlogit(rets, js,p1,p2)
    n = length(rets)
    x2 = log.(1 + abs.(rets))[1:(n-1)];
    alphaj = collect(linspace(0.05, 0.95, js))
    coefdf = zeros(js, 3)
    estlog = 0.0
    for j in collect(1:(js))
        indx1 = (rets .< (EWMA(rets).*quantile.(Normal(0,1), alphaj[j])) ).*1.0;
        data1 = DataFrame([indx1[2:n], indx1[1:(n-1)], x2], [:y, :x1, :x2]);

        estlog = glm(@formula(y ~ x1 + x2), data1, Binomial(), LogitLink());
        coefdf[j,:] = coef(estlog)
    end

    ## Estimation of kappas from deltas
    delta = delta3j(js)
    kapa1 = OLSestimator(coefdf[:, 2], delta[:,1:(p1+1)]);
    kapa2 = OLSestimator(coefdf[:, 3], delta[:,1:(p2+1)]);
    kapainit = [coefdf[:, 1]; kapa1; kapa2]
    return kapainit
end

## LL function
function loglikeFinalstd(x, param0,p1,p2)
    ##
    rets = x
    n = length(rets)
    m = zeros(n)
    theta = zeros(n)
    n_alpha = length(param0) - (p1+p2+2)
    prob = zeros(n, n_alpha)

    alphaj = collect(linspace(0.05, 0.95, n_alpha))

    cuts=zeros(n,n_alpha);
    for i in collect(1:n_alpha)
        cuts[:,i]=EWMA(rets)*quantile.(Normal(0,1), alphaj[i])
    end
    ##
    powers = [p1; p2]
    ## looping over time=j, and quantiles=i
    for j in collect(1:n)
        for i in collect(1:n_alpha)
            ##
            theta[j] = param0[1 + (i-1)*1]
            poly_1 = param0[1 + n_alpha]
            poly_2 = param0[2 + n_alpha + powers[1]]
            ## first polynoms in delta1
            for l in collect(1:powers[1])
                poly_1 = poly_1 + param0[1 + n_alpha + l] * (2.0 * alphaj[i] - 1.0)^l
            end
            ## second polynoms in delta2
            for l in collect(1:powers[2])
                poly_2 = poly_2 + param0[2 + n_alpha + powers[1]+l] * (2.0 * alphaj[i] - 1.0)^l
            end
            ##
            if j==1
                theta[j]=theta[j] + poly_1*mean(rets.<cuts[:,i]) + poly_2*mean(log.(1.0+abs.(rets)))
            else
                theta[j]=theta[j] + poly_1*(rets[j-1]<cuts[j-1,i]) + poly_2*log.(1.0+abs.(rets[j-1]))
            end
            ##
            prob[j, i] = 1.0 / (1.0 + exp.(-theta[j]))
        end
        m[j] = (rets[j] <= cuts[j,1]) * log.(prob[j, 1])
        for i in collect(2:n_alpha)
            if (prob[j, i] - prob[j, i-1]) <= 0.0
                prob[j, i] = prob[j, i-1] + 0.00001
            end

            m[j] = m[j] + ((rets[j] <= cuts[j,i]) * (rets[j] > cuts[j,i-1])) * log.(prob[j, i] - prob[j, i-1])
        end
            if (1.0 - prob[j, n_alpha]) <= 0.0
                prob[j, n_alpha] = 0.99999
            end
        m[j] = m[j] + (rets[j] > cuts[j,n_alpha]) * log.(1.0 - prob[j, n_alpha])
    end
    ##
    return m
end

function loglikeFinal(x, param0,p1,p2)

    m = loglikeFinalstd(x, param0,p1,p2)

    return -sum(m)
end

# Estimate parameters
function OrderedLogitparameters(x,q,p1,p2)

    kapainit = getinitlogit(x,q,p1,p2)
    estone = optimize(kapainit -> loglikeFinal(x, kapainit,p1,p2), kapainit, NelderMead(), Optim.Options(g_tol=1e-5, iterations=100))

    par = estone.minimizer

    return par
end

# Parameters and inference
function OrderedLogit(x,q,p1,p2)

    kapainit = getinitlogit(x,q,p1,p2)
    estone = optimize(kapainit -> loglikeFinal(x, kapainit,p1,p2), kapainit, NelderMead(), Optim.Options(g_tol=1e-5, iterations=100))

    par = estone.minimizer

    n=length(x)
    b=par'
    k=length(b)
    frac = 0.00001;
    h=frac*b

    e = eye(k)

    DfDp=zeros(n,k)
    for i in 1:k
        DfDp[:,i] = (loglikeFinalstd(x, b'+h'.*e[i,:],p1,p2)-loglikeFinalstd(x, b'-h'.*e[i,:],p1,p2))./h[i]/2.0
    end

    DfDp2=zeros(k,k);
    for i in 1:k
        for j in 1:i
            DfDp2[i,j] = mean(DfDp[:,i].*DfDp[:,j])
            DfDp2[j,i] = DfDp2[i,j];
        end
    end

    Df2Dp=zeros(k,k);
    for i in 1:k
    for j in 1:i
        if i==j;
            Df2Dp[i,i] = mean((loglikeFinalstd(x, b'+h'.*e[i,:],p1,p2)-2.0*loglikeFinalstd(x, b,p1,p2)+loglikeFinalstd(x, b'-h'.*e[i,:],p1,p2))./(h[i]^2.0));
        else
            Df2Dp[i,j] = mean((loglikeFinalstd(x, b'+h'.*(e[i,:]+e[j,:]),p1,p2)-loglikeFinalstd(x, b'+h'.*(e[i,:]-e[j,:]),p1,p2)-loglikeFinalstd(x, b'+h'.*(e[j,:]-e[i,:]),p1,p2)+loglikeFinalstd(x, b'-h'.*(e[i,:]+e[j,:]),p1,p2))./(4.0*h[i]*h[j]));
            Df2Dp[j,i] = Df2Dp[i,j];
        end
    end
    end
    Df2Dp=-Df2Dp;
    J = inv(Df2Dp)

    sterr=sqrt.(diag(J*DfDp2*J)/n)
    t=b'./sterr
    LL=estone.minimum

    AIC = -2LL + 2k + 2k*(k+1)/(n-k-1)
    BIC = -2LL + k*log(n)

    return [par,sterr,t,LL,AIC,BIC]
end

function forecastProbs(x, y, q,p1,p2)

    kapainit = getinitlogit(x,q,p1,p2)

    estone = optimize(kapainit -> loglikeFinal(x, kapainit,p1,p2), kapainit, NelderMead(), Optim.Options(g_tol=1e-5, iterations=100))
    param0 = estone.minimizer

    rets0 = x
    rets = y
    n = length(rets)
    m = zeros(n)
    theta = zeros(n)
    n_alpha = length(param0) - (p1+p2+2)
    prob = zeros(n, n_alpha)
    alphaj = collect(linspace(0.05, 0.95, n_alpha))
    cuts=EWMA(rets)[length(rets)].*quantile.(Normal(0,1), alphaj)
    powers = [p1; p2]
    wrong=zeros(n)
    ## looping over time=j, and quantiles=i
    for j in collect(2:n)
        for i in collect(1:n_alpha)
            ##
            theta[j] = param0[1 + (i-1)*1]
            poly_1 = param0[1 + n_alpha]
            poly_2 = param0[2 + n_alpha + powers[1]]
            ## first polynoms in delta1
            for l in collect(1:powers[1])
                poly_1 = poly_1 + param0[1 + n_alpha + l] * (2.0 * alphaj[i] - 1.0)^l
            end
            ## second polynoms in delta2
            for l in collect(1:powers[2])
                poly_2 = poly_2 + param0[2 + n_alpha + powers[1]+l] * (2.0 * alphaj[i] - 1.0)^l
            end

            theta[j]=theta[j] + poly_1*(rets[j-1]<cuts[i]) + poly_2*log.(1.0+abs.(rets[j-1]))

            prob[j, i] = 1.0 / (1.0 + exp.(-theta[j]))
        end
        for i in collect(2:n_alpha)
            if (prob[j, i] - prob[j, i-1]) <= 0.0
                    prob[j, i] = prob[j, i-1] + 0.00001
                    #wrong[j] = 1
            end
        end
    end

    return prob[2:n,:]
end

## ------------------------------------------
## Statistical evaluation of forecasts
## ------------------------------------------

function Interpol(xs,ys,x)
    n = length(xs)
    dxs = xs[2:n]-xs[1:n-1]
    dys = ys[2:n]-ys[1:n-1]
    Ds = dys./dxs
    ms = [Ds[1];(Ds[1:n-2]+Ds[2:n-1])/2;Ds[n-1]]
    as = ms[1:n-1]./Ds[1:n-1]
    bs = ms[2:n]./Ds[1:n-1]
    cir = (as.^2)+(bs.^2)
    ex = cir.>9
    tau = 3./sqrt.(cir).*ex + 1-ex
    ms = ms.*[tau;1].*[1;tau]
    if findfirst(xs.>x) == 0
        ind_u = 2
    elseif findfirst(xs.>x) == 1
        ind_u = 2
    else
        ind_u = findfirst(xs.>x)
    end
    x_u = xs[ind_u]
    x_l = xs[ind_u-1]
    h = x_u-x_l
    t = (x-x_l)./h
    fx = ys[ind_u-1]*(2*t^3-3*t^2+1) + h*ms[ind_u-1]*(t^3-2*t^2+t) + ys[ind_u]*(-2*t^3+3*t^2) + h*ms[ind_u]*(t^3-t^2)

    return fx
end

function Berkowitz(pit)
    Zb = zeros(length(pit))
    pit[pit.>=1] =0.9999999
    pit[pit.<=0] =0.0000001
    for t in collect(1:length(pit))
        Zb[t] = norminvcdf(abs.(pit[t]))
    end
    naux = length(Zb)-1
    mu_hat = sum(Zb)/naux

    Xaux = Zb[1:(length(Zb)-1)]-mu_hat
    Yaux = Zb[2:length(Zb)]-mu_hat

    rho_hat = Xaux'Yaux/(Xaux'Xaux)
    eps_hat = Yaux-rho_hat*Xaux
    s2_hat = sum(eps_hat.^2)/naux

    LLu = -1/2*naux*log(s2_hat) - 1/2/s2_hat*eps_hat'eps_hat
    LLr = - 1/2*Yaux'Yaux
    LR = -2*(LLr-LLu)
    return [LR,ccdf(Chisq(3), LR)]
end

function GonzalezRivera(pit,alpha_c,alpha_main,k_main)
    # alpha_main    -- alpha-contour for k-aggregated test
    # k_main        -- lag for alpha-aggregated test

    n_a = length(alpha_c)   # collection of alpha-contours
    K_c=collect(1:n_a)
    n_k = length(K_c)                   # collection of lags for alpha-aggregated test
    Lambda_a = zeros(n_k,n_k)
    L_a =  zeros(n_k,1)
    report1=zeros(n_a,2,n_k)
    report2=zeros(1,2)
    report3=zeros(1,2)
    for k in 1:n_k
        alpha_contour_hat = zeros(length(alpha_c),1)
        PITk = [pit[1:(length(pit)-k)] pit[(1+k):length(pit)]]
        Tk = size(PITk)[1]
        Ind = zeros(Tk,n_a)
        Omega_k = zeros(n_a,n_a)

        for i in 1:n_a
            Indi = 1*(PITk.<sqrt(alpha_c[i]))
            Ind[:,i] = Indi[:,1].*Indi[:,2];
            alpha_contour_hat[i] = sum(Ind[:,i])/Tk;

            for j in 1:(i-1)
                    Omega_k[i,j] = alpha_c[j].*(1-alpha_c[i])+2*alpha_c[j]*alpha_c[i]^0.5.*(1-alpha_c[i]^0.5)
            end
            Omega_k = Omega_k + Omega_k'
            Omega_k[i,i] = alpha_c[i].*(1-alpha_c[i])+2*alpha_c[i]^1.5.*(1-alpha_c[i]^0.5)
        end

        for j in 1:n_k
            Lambda_a[j,k] = 4*alpha_main^1.5*(1-alpha_main^0.5)
        end
        Lambda_a[k,k] = alpha_main.*(1-alpha_main)+2*alpha_main^1.5.*(1-alpha_main^0.5)
        C_k = sqrt(Tk)*(alpha_contour_hat-alpha_c)
        t_c = C_k./sqrt.(alpha_c.*(1-alpha_c)+2*alpha_c.^1.5.*(1-alpha_c.^0.5))
        C_k = C_k'inv(Omega_k)*C_k

        Indi = 1*(PITk.<sqrt(alpha_main))
        Ind = Indi[:,1].*Indi[:,2]
        alpha_main_hat = sum(Ind)/Tk

        L_a[k] = sqrt(Tk)*(alpha_main_hat-alpha_main)

        # report individual k stats and p-vals
        report1[:,:,k]=[t_c ccdf.(Chisq(1),(t_c.^2))]

        # report alpha-aggregated
        if k == k_main
            report2=[C_k ccdf.(Chisq(n_a), C_k)]
        end

    end
    # report lag-aggregated
    L_a = L_a'inv(Lambda_a)*L_a;

    report3=[L_a ccdf.(Chisq(n_k), L_a)]

    return [report1,report2,report3]
end

## NUMERICAL INTEGRATION ROUTINES

function ComputeInt(r,rets,cuts,ext_probs,n)

    a = 2*minimum(rets)
    b = 2*maximum(rets)
    x = cos.(pi*(2*collect(1:1:n)-1)./2./n)
    S = 0

    for i in 1:n
        y=(x[i]+1)*(b-a)/2+a
        CDFy = Interpol([2*minimum(rets);cuts;2*maximum(rets)],ext_probs,y)
        S = S + ((CDFy-1*(y.>=r)).^2)*sqrt(1-x[i]^2)
    end

    return S*pi*(b-a)/2/n
end


using DataFrames
using Gadfly
using Mamba
using Distributions
using StatsBase

dat = readtable("test.txt.gz", separator = '\t', header=false, names=[:Length, :DeltaG])

hex=plot(dat, x=:Length, y=:DeltaG, Geom.hexbin(xbincount=35),
         Coord.cartesian(xmin=0, ymax=15),
         Guide.xlabel("Duplex Length"),
         Guide.ylabel("Gibbs Free Energy Change (ΔG)"))

dens=plot(dat[dat[:Length] .>= 5,:], color=:Length, x=:DeltaG, Geom.density, Guide.colorkey("Duplex Length"),
          Guide.xlabel("Gibbs Free Energy Change (ΔG)"))

dd = DuplexDist( dat )
model=plotcdf( dd, dat[:Length], dat[:DeltaG] )

draw(PDF("deltaG.pdf", 8inch, 8inch), vstack(hstack(hex,dens), model))

type DuplexDist
   length::Distributions.Poisson
   deltag::Vector{Distributions.Normal}
end

# G[L] ~ Normal(mu[L], s[L])
# L ~ Poisson(lambda)
function DuplexDist( df::DataFrame )
   # Fit Model
   dgdat = Vector{DataArray}()
   for i in 1:maximum(df[:Length])
      push!( dgdat, dat[dat[:,1] .== i, 2] )
   end
   lendat = df[:Length]

   len = fit(Poisson, lendat)
   dg  = Vector{Distributions.Normal}(length(dgdat))
   for i in 1:N
      dg[i] = fit(Normal, dgdat[i])
   end

   DuplexDist( len, dg )
end

function cdf( dd::DuplexDist, len::Int, dg::Float64 )
   prob = 0.0
   for i in len:length(dd.deltag)
      prob += pdf( dd.length, i ) * cdf( dd.deltag[i], dg )
   end
   prob += ccdf( dd.length, length(dd.deltag)+1 )
   prob
end

# P(G,L) = P(G|L)P(L)
pdf( dd::DuplexDist, len::Int, dg::Float64 ) = pdf( dd.length, len ) * pdf( dd.deltag[len], dg )

function plotcdf( dd::DuplexDist, lenvec, dgvec )
   l = Vector{Int}()
   d = Vector{Float64}()
   c = Vector{Float64}()
   p = Vector{Float64}()
   irange = 1:maximum(lenvec)
   jrange = minimum(dgvec):1.0:maximum(dgvec)
   for i in irange
      for j in jrange
         push!( l, i )
         push!( d, j )
         push!( c, cdf( dd, i, j ) )
         push!( p, pdf( dd, i, j ) )
      end
   end
   z = reshape( p, (length(jrange), length(irange)))

   pval = plot(layer(x=l, y=d, color=c, Geom.rectbin),
         Coord.cartesian(xmin=0, ymax=15),
         Guide.xlabel("Duplex Length"),
         Guide.ylabel("Gibbs Free Energy Change (ΔG)"),
         Guide.colorkey("p-value"), Scale.color_log10)

   prob =plot(layer(x=l, y=d, color=p, Geom.rectbin),
         Coord.cartesian(xmin=0, ymax=15),
         Guide.xlabel("Duplex Length"),
         Guide.ylabel("Gibbs Free Energy Change (ΔG)"),
         Guide.colorkey("Pdf(x)"))

   hstack(prob, pval)
end

#= MAMBA

fitdat = Dict{Symbol, Any}(
  :len => len,
  :dg => dg
)
N = length(dg)

model = Model(
   dg = Stochastic(1, (len, mu, s2) -> UnivariateDistribution[ Normal(mu[i], sqrt(s2[i])) for i in 1:N ], false),

   mu = Stochastic(1, () -> Normal(0, 1)),
   s2 = Stochastic(1, () -> InverseGamma(0.01, 0.01)),

   len = Stochastic(alpha -> begin
                               lambda = max(0, 1/alpha)
                               Poisson(lambda)
                             end, false),
   alpha = Stochastic(() -> Exponential(1/10)),
)

inits = [
  Dict{Symbol, Any}(
     :alpha => rand(Exponential(1/10)),
     :len => mean(fitdat[:len]),
     :s2 => [rand(InverseGamma(0.01, 0.01)) for i in 1:N],
     :mu => [mean(dg[i]) for i in 1:N],
     :dg => [mean(dg[i]) for i in 1:N]
  )
for j in 1:5
]

scheme = [Slice(:alpha, 1.0, Univariate),
          Slice(:len, 1.0, Univariate),
          NUTS([:mu, :s2])]

setinputs!(model, fitdat)
setinits!(model, inits[1])
setsamplers!(model, scheme)

sim = mcmc(model, fitdat, inits, 10000, burnin=250, thin=2, chains=1)

=#


using NPZ
using JuMP
using Ipopt
using LinearAlgebra
using .Threads
using WoodburyMatrices
using Statistics
using UnicodePlots

function main(B)
    X1 = npzread("Kernels/G1.npy")
    X2 = npzread("Kernels/G2.npy")
    X3 = npzread("Kernels/G3.npy")
    nfactor = npzread("Testing/nfactor.npy")

    n = size(X3, 1)
    ndas = size(X1, 1)

    T = Float64
    ϵ = 1.0
    λ = 100.0
    α = 100.0

    X = [X3 zeros(T, size(X3));
        zeros(T, size(X3)) X3;
        nfactor*X1 nfactor*X2]
    Xn = maximum(sqrt.(sum(X.^2, dims=2)))
    #Xn = sqrt.(sum(X.^2, dims=2)) #normalize all rows to be 1?
    X ./= Xn

    v, V = eigen(X'*X + ϵ*I)

    #eigtol = 5#quantile(v, 0.99)

    n_eigs = 5#sum(v.>=eigtol)

    Γ = V[:, (end-n_eigs+1):end]

    μ_init = ones(n) ./ n .* B
    μn = sum(μ_init)/B
    μ_init ./= μn
    U_init = copy(Γ)

    function Ml(μ, X, ϵ)
        M = ϵ*I + X'*Diagonal(μ)*X
        Mc = cholesky(Hermitian(M))
        return (M, Mc)
    end

    Ml(μ) = Ml(μ, X, ϵ)

    function subproblem_sμ(μ_init, u, X, λ, B, maxiter=1000)

        function f(μ, u, λ)
            M, Mc = Ml(μ)
            return -logdet(M) + λ*tr(u'*(Mc\u))
        end

        f(μ) = f(μ, u, λ)

        jump_wrapper_f(x...) = f(vcat([x[i] for i in 1:length(x)], [x[i] for i in 1:length(x)], ones(T, ndas)))

        function gi(i, Mc, μ, u, λ)
            ld_term = @views -X[i,:]'*(Mc\X[i,:])
            tmp_A = SymWoodbury(Mc, X[i,:], -μ[i])
            tmp_b = @views tmp_A \ X[i,:]
            #tmp_1 = (u'*tmp_b)^2
            tmp_1 = sum((u'*tmp_b).^2) # think that this should be correct for tr(U'MU) for matrix U but need to check
            tmp_d = @views X[i,:]'*tmp_b
            tmp_den = (1+μ[i]*tmp_d)^2
            eig_term = ((1+μ[i]*tmp_d)*tmp_1 - μ[i]*tmp_1*tmp_d) / tmp_den
            return ld_term - λ*eig_term #apparently need a -ve sign here, work out why
        end

        function g!(gvec, μ, X, u, λ, n)
            M, Mc = Ml(μ)
            @threads for i = 1:n
                g1 = gi(i, Mc, μ, u, λ) # in respect to x component
                g2 = gi(i+n, Mc, μ, u, λ) # in respect to y component
                gvec[i] = g1 + g2
            end
        end

        g!(gvec, μ) = g!(gvec, μ, X, u, λ, n)

        jump_wrapper_g!(gvec, x...) = g!(gvec, vcat([x[i] for i in 1:length(x)], [x[i] for i in 1:length(x)], ones(T, ndas)))

        model = Model(optimizer_with_attributes(Ipopt.Optimizer, "max_iter"=>maxiter, "tol"=>1e-4/B))
        JuMP.register(model, :f, n, jump_wrapper_f, jump_wrapper_g!)

        @variable(model, 0.0<=x[1:n]<= 1.0)

        set_start_value.(x, μ_init)

        @NLconstraint(model, +(x...) <= B)

        @NLobjective(model, Min, f(x...))

        optimize!(model)

        return (value.(x), vcat(value.(x), value.(x), ones(T, ndas)), objective_value(model))

    end

    function subproblem_sO(u_init, μ, γ, X, λ, α, maxiter=1000)
        _, Mc = Ml(μ)

        f(u) = α*sum((u.-γ).^2) + λ*tr(u'*(Mc\u))
        function g!(gvec, u)
            gvec .= 2*α*(u-γ) + 2*λ*(Mc\u)
        end
        jump_wrapper_f(x...) = f([x[i] for i in 1:length(x)])
        jump_wrapper_g!(gvec, x...) = g!(gvec, [x[i] for i in 1:length(x)])

        sumsquares(x...) = sum(x[i]^2 for i in 1:length(x))
        function sumsquares_grad!(gvec, x...)
            for i in 1:length(x)
                gvec[i] = 2*x[i]
            end
        end

        model = Model(optimizer_with_attributes(Ipopt.Optimizer, "max_iter"=>maxiter))
        JuMP.register(model, :f, length(γ), jump_wrapper_f, jump_wrapper_g!)
        JuMP.register(model, :sumsquares, length(γ), sumsquares, sumsquares_grad!)

        @variable(model, x[1:length(γ)])

        set_start_value.(x, u_init)

        @NLconstraint(model, sumsquares(x...) == 1)

        @NLobjective(model, Min, f(x...))

        optimize!(model)

        return (value.(x), objective_value(model))

    end

    function subproblem_sO(U_init, μ, Γ, X, λ, α, m, n, maxiter=1000)
            _, Mc = Ml(μ)

            f(U) = α*sum((U.-Γ).^2) + λ*tr(U'*(Mc\U))

            function g!(gvec, U, m, n)
                @threads for i = 1:n
                    @views gvec[((i-1)*m+1):i*m] .= 2*α*(U[1:m, i].-Γ[1:m, i]) .+ 2*λ*(Mc\U[1:m, i])
                end
            end

            L = LinearIndices((m, n))
            jump_wrapper_f(x...) = f([x[L[i,j]] for i in 1:m, j in 1:n])
            jump_wrapper_g!(gvec, x...) = g!(gvec, [x[L[i,j]] for i in 1:m, j in 1:n], m, n)

            sumprod(x...) = sum(x[i]*x[i+m] for i in 1:m)
            function sumprod_grad!(gvec, x...)
                for i in 1:m
                    gvec[i] = x[i+m]
                    gvec[i+m] = x[i]
                end
            end

            model = Model(optimizer_with_attributes(Ipopt.Optimizer, "max_iter"=>maxiter))
            JuMP.register(model, :f, m*n, jump_wrapper_f, jump_wrapper_g!)
            JuMP.register(model, :sumprod, 2*m, sumprod, sumprod_grad!)

            @variable(model, x[1:m, 1:n])

            set_start_value.(x, U_init)

            for i = 1:n
                for j = i:n
                    x_i = x[1:m, i]
                    x_j = x[1:m, j]
                    @NLconstraint(model, sumprod(x_i..., x_j...) == ((i==j) ? 1 : 0))            end
            end

            @NLobjective(model, Min, f(x...))

            optimize!(model)

            return (value.(x), objective_value(model)) # want to have all of the x here as a matrix
    end

    function optimize_spectral(μ_init, u_init, iters, γ, X, λ, α, B)
        μ = copy(μ_init)
        u = copy(u_init)
        ob1l = Float64[]
        ob2l = Float64[]
        for iter in 1:iters
            μ, μl, ob1 = subproblem_sμ(μ, u, X, λ, B)
            u, ob2 = subproblem_sO(u, μl, γ, X, λ, α)
            push!(ob1l, ob1)
            push!(ob2l, ob2)
        end
        μ, μl, ob1 = subproblem_sμ(μ, u, X, λ, B)
        push!(ob1l, ob1)

        return (μ, u, ob1l, ob2l)
    end

    function optimize_spectral(μ_init, U_init, iters, Γ, X, λ, α, B, m, n)
        μ = copy(μ_init)
        U = copy(U_init)
        ob1l = Float64[]
        ob2l = Float64[]
        for iter in 1:iters
            μ, μl, ob1 = subproblem_sμ(μ, U, X, λ, B)
            U, ob2 = subproblem_sO(U, μl, Γ, X, λ, α, m, n)
            push!(ob1l, ob1)
            push!(ob2l, ob2)
        end
        μ, μl, ob1 = subproblem_sμ(μ, U, X, λ, B)
        push!(ob1l, ob1)

        return (μ, U, ob1l, ob2l)
    end

    μ, U_res, ob1l, ob2l = optimize_spectral(μ_init, U_init, 5, Γ, X, λ, α, B, size(Γ)...)

    M, _ = Ml(vcat(μ,μ, ones(ndas)))
    vμ, Vμ = eigen(M)

    npzwrite("SpectralResults/res_spectral_$B.npy", μ)

    lineplot(ob1l)

    lineplot(ob2l)
end

for i in [4,8,16,32,64,128,256]
    main(i)
end

#yeh


using Printf

include("alias.jl")
include("utils.jl")
include("PGBF.jl")
include("CGBF.jl")
include("overlap.jl")
include("kinetic.jl")

function test_kinetic()

  s = PGBF(1.0)
  c = init_CGBF( 0.0, 0.0, 0.0 )
  push!(c,1.0,1.0)

  @assert isapprox(kinetic(1.0, (0.0,0.0,0.0), (0,0,0),
                           1.0, (0.0,0.0,0.0), (0,0,0)), 2.9530518648229536)
  @assert isapprox( kinetic(s,s), 1.5 )
  #@assert isapprox(kinetic(c,c),1.5)
  @printf "test_kinetic is passed\n"
end

test_kinetic()


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Julia 1.3.1
#     language: julia
#     name: julia-1.3
# ---

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# # Lagrangian particle tracking
#
# Material particles following ocean currents can be analyzed in terms of trajectories. These can simply be computed by integrating velocities through time within a [Lagrangian framework](https://en.wikipedia.org/wiki/Lagrangian_and_Eulerian_specification_of_the_flow_field). In `Julia` this is easily done using the [IndividualDisplacements.jl](https://JuliaClimate.github.io/IndividualDisplacements.jl/dev/) and [OrdinaryDiffEq.jl](https://docs.juliadiffeq.org/latest) packages.

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Import software, grid, and velocity fields
#
# 1. pre-requisites
# 2. read variables

# + {"slideshow": {"slide_type": "subslide"}}
#]add MITgcmTools#master; add OrdinaryDiffEq; add IndividualDisplacements#master;

# + {"slideshow": {"slide_type": "-"}, "cell_style": "center"}
using IndividualDisplacements, MeshArrays, OrdinaryDiffEq
using Plots, Statistics, MITgcmTools, DataFrames

include("helper_functions.jl")
get_grid_if_needed()
get_velocity_if_needed()

Γ=read_llc90_grid();

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Set Gridded Variables
#
# 1. Select depth
# 2. time average & normalize to grid units
# 3. apply exchange function to `u,v,lon,lat`
# 4. store everything in `uv_etc` dictionary

# + {"slideshow": {"slide_type": "-"}}
uv_etc=read_uv_etc(20,Γ);

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Initialize Trajectory Computation
#
# 1. set `duComp` to interpolation function
# 2. set `u0` initial location array

# + {"slideshow": {"slide_type": "-"}}
du_dt=IndividualDisplacements.VelComp!
(u0,du)=initialize_locations(uv_etc,10);

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Compute Trajectories
#
# _Note: `ODEProblem` and `solve` settings can still be refined_
# -

tspan = (0.0,uv_etc["t1"]-uv_etc["t0"])
prob = ODEProblem(du_dt,u0,tspan,uv_etc)
sol = solve(prob,Euler(),dt=uv_etc["dt"])
size(sol)

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Post-Process Trajectories
#
# 1. Copy `sol` to a `DataFrame`
# 2. Map position to lon,lat coordinates

# + {"slideshow": {"slide_type": "subslide"}}
df=postprocess_ODESolution(sol,uv_etc);

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Plot Trajectories
#
#
# For example, you can:
#
# - use either `Makie.jl` via `PlotMakie`
# - or use `PyPlot.jl` via `PlotMapProj`
#

# + {"slideshow": {"slide_type": "-"}}
p=dirname(pathof(IndividualDisplacements));

# + {"slideshow": {"slide_type": "slide"}}
#include(joinpath(p,"plot_pyplot.jl"));
#PyPlot.figure(); PlotMapProj(df,5000)

#include(joinpath(p,"plot_makie.jl")); AbstractPlotting.inline!(true); #for Juno, set to false
#scene=PlotMakie(df,5000,180.0) #Makie.save("LatLonCap300mDepth.png", scene)

include(joinpath(p,"plot_plots.jl"));
plt=PlotBasic(df,1000,180.0)


"""
    ErrorMap

A nonlinear mapping between the space of unit quaternions and three-dimensional
rotation errors.

These mappings are extremely useful for converting from globally nonsingular 3D rotation
representations such as `UnitQuaternion` or `RotMatrix3` to a three-parameter error that can
be efficiently used in gradient-based optimization methods that optimize deviations about
a current iterate using first or second-order information.

# Usage
    errmap(v::AbstractVector)  # "forward" map from a 3D error to a `UnitQuaternion`
    inv(errmap)(R::Rotation)   # "inverse" map from a rotation (via `UnitQuaternion`) to a 3D error

where `errmap <: ErrorMap`

# Implemented Maps
- `CayleyMap`:  Uses `RodriguesParam` as the error representation (default).
    Goes singular at 180° and does not have a sign ambiguity.
- `ExponentialMap`: Uses the canonical exponential map from Lie Group theory. Computationally
    expensive to compute. Exhibits kinematic singularities.
- `MRPMap`: Uses a scaled `MRP` as the error representation. Singular at 360° but has
    a sign ambiguity (with the "shadow" MRP set).
- `QuatVecMap`: Uses the vector part of the quaternions. Cheapest map to compute, but
    goes singular at 180° and suffers from sign ambiguity.
- `IdentityMap`: Maps values through directly. Only works with three-parameter rotation
    representations with the following methods: `R(::SVector{3})` and `SVector(::R)::SVector{3}`
"""
abstract type ErrorMap end
struct CayleyMap <: ErrorMap end
struct ExponentialMap <: ErrorMap end
struct MRPMap <: ErrorMap end
struct QuatVecMap <: ErrorMap end
struct IdentityMap <: ErrorMap end

"""
    InvErrorMap

The nonlinear mapping from unit quaternions to a three-dimensional error state. Obtained by
inverting an `ErrorMap`, i.e.

    InvCayleyMap() = inv(CayleyMap())

# Usage
    imap(R::Rotation)             # "inverse" map from a rotation to a 3D error
    inv(imap)(v::AbstractVector)  # "forward" map from a 3D error to a `UnitQuaternion`

where `imap <: InvErrorMap`.

See `ErrorMap` for documentation on the implemented maps.
"""
abstract type InvErrorMap end
struct InvCayleyMap <: InvErrorMap end
struct InvExponentialMap <: InvErrorMap end
struct InvMRPMap <: InvErrorMap end
struct InvQuatVecMap <: InvErrorMap end

inv(::CayleyMap) = InvCayleyMap()
inv(::ExponentialMap) = InvExponentialMap()
inv(::QuatVecMap) = InvQuatVecMap()
inv(::MRPMap) = InvMRPMap()
inv(::IdentityMap) = IdentityMap()

inv(::InvCayleyMap) = CayleyMap()
inv(::InvExponentialMap) = ExponentialMap()
inv(::InvQuatVecMap) = QuatVecMap()
inv(::InvMRPMap) = MRPMap()

# Scalings
@inline scaling(::Type{ExponentialMap}) = 0.5
@inline scaling(::Type{QuatVecMap}) = 1.0
@inline scaling(::Type{CayleyMap}) = 1.0
@inline scaling(::Type{MRPMap}) = 2.0
scaling(m::M) where M <: ErrorMap = scaling(M)

# Quaternion Maps
(::ExponentialMap)(ϕ::AbstractVector) = expm(ϕ/scaling(ExponentialMap))

function (::QuatVecMap)(v::AbstractVector)
    check_length(v, 3)
    μ = 1/scaling(QuatVecMap)
    UnitQuaternion(sqrt(1-μ^2*v'v), μ*v[1], μ*v[2], μ*v[3])
end

function (::CayleyMap)(g::AbstractVector)
    check_length(g, 3)
    g /= scaling(CayleyMap)
    M = 1/sqrt(1+g'g)
    UnitQuaternion(M, M*g[1], M*g[2], M*g[3])
end

function (::MRPMap)(p::AbstractVector)
    check_length(p, 3)
    p /= scaling(MRPMap)
    n2 = p'p
    M = 2/(1+n2)
    UnitQuaternion((1-n2)/(1+n2), M*p[1], M*p[2], M*p[3])
end

function (::IdentityMap)(q::AbstractVector)
    check_length(q, 3)
    UnitQuaternion(q[1], q[2], q[3], q[4])
end

# Quaternion Map Jacobians
"""
    jacobian(::ErrorMap, ϕ)

Jacobian of the quaternion map that takes a three-dimensional vector `ϕ` and returns a
    unit quaternion.
Returns a 4x3 Static Matrix

For all the maps (except the `IdentityMap`)
`jacobian(::ErrorMap, zeros(3)) = [0; I] = Hmat()'`
"""
function jacobian(::ExponentialMap, ϕ, eps=1e-5)
    μ = 1/scaling(ExponentialMap)
    θ = norm(ϕ)
    cθ = cos(μ*θ/2)
    sincθ = sinc(μ*θ/2π)
    if θ < eps
        0.5*μ*[-0.5*μ*sincθ*ϕ'; sincθ*I + (cθ - sincθ)*ϕ*ϕ']
    else
        0.5*μ*[-0.5*μ*sincθ*ϕ'; sincθ*I + (cθ - sincθ)*ϕ*ϕ'/(ϕ'ϕ)]
    end
end

function jacobian(::QuatVecMap, v)
    μ = 1/scaling(QuatVecMap)
    μ2 = μ*μ
    M = -μ2/sqrt(1-μ2*v'v)
    @SMatrix [v[1]*M v[2]*M v[3]*M;
              μ 0 0;
              0 μ 0;
              0 0 μ]
end

function jacobian(::CayleyMap, g)
    μ = 1/scaling(CayleyMap)
    μ2 = μ*μ
    n = 1+μ2*g'g
    ni = 1/n
    μ*[-μ*g'; -μ2*g*g' + I*n]*ni*sqrt(ni)
end

function jacobian(::MRPMap, p)
    μ = 1/scaling(MRPMap)
    μ2 = μ*μ
    n = 1+μ2*p'p
    2*[-2*μ2*p'; I*μ*n - 2*μ*μ2*p*p']/n^2
end

jacobian(::IdentityMap, q) = I

############################################################################################
#                             INVERSE RETRACTION MAPS
############################################################################################
# (emap::ErrorMap)(R::Rotation) =
#     throw(ArgumentError("Must use the inverse map to convert a rotation to an error state, e.g. inv(emap)(q)"))
# (emap::InvErrorMap)(R::Rotation) = emap(UnitQuaternion(R))  # automatically call the inverse map for rotations
# NOTE: Julia v1.0 doesn't allow these methods on abstract types. Leave out for now.

for imap in (InvExponentialMap,InvCayleyMap,InvMRPMap,InvQuatVecMap)
    @eval begin
        @inline (imap::$imap)(R::Rotation) = imap(UnitQuaternion(R))
    end
end

(::InvExponentialMap)(q::UnitQuaternion) = scaling(ExponentialMap)*logm(q)
(::InvCayleyMap)(q::UnitQuaternion) = scaling(CayleyMap) * vector(q)/q.w
(::InvMRPMap)(q::UnitQuaternion) = scaling(MRPMap)*vector(q)/(1+q.w)
(::InvQuatVecMap)(q::UnitQuaternion) = scaling(QuatVecMap)*vector(q) * sign(q.w)

# ~~~~~~~~~~~~~~~ Inverse map Jacobians ~~~~~~~~~~~~~~~ #
"""
    jacobian(::InvErrorMap, q::UnitQuaternion)

Jacobian of the inverse quaternion map, returning a 3x4 matrix.
For all maps: `jacobian(::InvErrorMap, UnitQuaternion(I)) = [0 I] = Hmat()'`
"""
function jacobian(::InvExponentialMap, q::UnitQuaternion, eps=1e-5)
    μ = scaling(ExponentialMap)
    s = scalar(q)
    v = vector(q)
    θ2 = v'v
    θ = sqrt(θ2)
    datan = 1/(θ2 + s^2)
    ds = -datan*v

    if θ < eps
        return 2*μ*[ds (v*v' + I)/s]
    else
        atanθ = atan(θ,s)
        dv = ((s*datan - atanθ/θ)v*v'/θ + atanθ*I )/θ
        d0 = ((s*datan - atanθ/θ)v*v'/θ^2 + atanθ/θ*I )
        d1 = (s*datan - atanθ/θ)
        d2 = v*v'/θ2
        d3 = atanθ/θ * I
        return 2*μ*[ds dv]
    end
end

function jacobian(::InvQuatVecMap, q::UnitQuaternion)
    μ = scaling(QuatVecMap)
    return sign(q.w) * SA[
                0. μ 0 0;
                0. 0 μ 0;
                0. 0 0 μ]
end

function jacobian(::InvCayleyMap, q::UnitQuaternion)
    μ = scaling(CayleyMap)
    si = 1/q.w
    return μ*@SMatrix [-si^2*q.x si 0 0;
                       -si^2*q.y 0 si 0;
                       -si^2*q.z 0 0 si]
end

function jacobian(::InvMRPMap, q::UnitQuaternion)
    μ = scaling(MRPMap)
    si = 1/(1+q.w)
    return μ*@SMatrix [-si^2*q.x si 0 0;
                       -si^2*q.y 0 si 0;
                       -si^2*q.z 0 0 si]
end

jacobian(::IdentityMap, q::UnitQuaternion) = I

# ~~~~~~~~~~~~~~~ Inverse map Jacobian derivative ~~~~~~~~~~~~~~~ #
"""
    ∇jacobian(::InvErrorMap, q::UnitQuaternion, b::SVector{3})

Jacobian of G(q)'b, where G(q) = jacobian(::InvErrorMap, q),
    b is a 3-element vector
"""
function ∇jacobian(::InvExponentialMap, q::UnitQuaternion, b::SVector{3}, eps=1e-5)
    μ = scaling(ExponentialMap)
    s = scalar(q)
    v = vector(q)
    θ2 = v'v
    θ = sqrt(θ2)
    datan = 1/(θ2 + s^2)
    ds = -datan*v

    if θ < eps
        # return 2*μ*[b'ds; (v*v'b + b)/s]
        return 2*μ*[b'*(datan^2*2s*v) -b'datan*I;
                    -(v*v'b +b)/s^2 (I*(v'b) + v*b')/s]
    else
        dsds = 2b's*datan^2*v
        dsdv = b'*(-datan*I + 2datan^2*v*v')

        atanθ = atan(θ,s)
        d1 = (s*datan - atanθ/θ)
        d2 = v*v'b/θ2
        d3 = atanθ/θ*b
        d1ds = (datan - 2s^2*datan^2 + datan)
        dvds = d1ds*d2 - datan*b

        d1dv =  (-2s*datan^2*v' - s*datan*v'/θ^2 + atanθ/θ^3*v')
        d2dv = (I*(v'b) + v*b')/θ2 - 2(v*v'b)/θ^4 * v'
        d3dv = b*(s*datan*v'/θ^2 - atanθ/θ^3*v')
        dvdv = d2*d1dv + d1*d2dv + d3dv

        # return 2*μ*[ds'b; dv'b]
        return 2*μ*@SMatrix [
            dsds    dsdv[1] dsdv[2] dsdv[3];
            dvds[1] dvdv[1] dvdv[4] dvdv[7];
            dvds[2] dvdv[2] dvdv[5] dvdv[8];
            dvds[3] dvdv[3] dvdv[6] dvdv[9];
        ]
    end
end

function ∇jacobian(::InvCayleyMap, q::UnitQuaternion, b::SVector{3})
    μ = scaling(CayleyMap)
    si = 1/q.w
    v = vector(q)
    μ*@SMatrix [
        2*si^3*(v'b) -si^2*b[1] -si^2*b[2] -si^2*b[3];
       -si^2*b[1] 0 0 0;
       -si^2*b[2] 0 0 0;
       -si^2*b[3] 0 0 0;
    ]
end

function ∇jacobian(::InvMRPMap, q::UnitQuaternion, b::SVector{3})
    μ = scaling(MRPMap)
    si = 1/(1+q.w)
    v = vector(q)
    μ * @SMatrix [
        2*si^3*(v'b) -si^2*b[1] -si^2*b[2] -si^2*b[3];
       -si^2*b[1] 0 0 0;
       -si^2*b[2] 0 0 0;
       -si^2*b[3] 0 0 0;
    ]
end

function ∇jacobian(::InvQuatVecMap, q::UnitQuaternion, b::SVector{3})
    μ = scaling(QuatVecMap)
    @SMatrix zeros(4,4)
end


using CompScienceMeshes

coarse = meshsphere(1.0, 0.45)
fine = CompScienceMeshes.lineofsight_refinement(m)
vtoc, nc = CompScienceMeshes.vertextocellmap(fine)

colors = zeros(numcells(fine))
for i in 1:numvertices(coarse)
    c = rand()
    for j in 1 : (nc[i])
        colors[vtoc[i,j]] = c
    end
end

import PlotlyJS
PlotlyJS.plot(patch(fine.mesh, colors))