using BrownianDynamics
using OffsetArrays
using DifferentialEquations
using DiffEqDiffTools
using DSP
using FFTW
using Test
using SparseArrays
using LinearAlgebra

# write your own tests here
@testset "conv_mat" begin
    vv = [1.0;1.0+im;2.0;1.0*im;1.0;3.0+im]
    T = Complex{Float64}
    nu = length(vv)

    n = 2nu-1
    np2 = n > 1024 ? nextprod([2,3,5],n) : nextpow(2,n)
    upad = [vv;zeros(T,np2-nu)]

    p = plan_fft!(upad)
    ip = plan_ifft!(upad)

    @test conv_mat(vv) == conv_mat(p,ip,vv)
    @test isapprox(conv_mat(p,ip,vv)*vv,conv(vv,vv)[4:9])
    u = rand(6)
    v = rand(6)
    @test isapprox(conv_mat(u)*v,conv(u,v)[4:9])
end

# @testset "flux!" begin
#     # Test that the Jacobian is correct.
# end

## Initialize a spectral system.
nn = 64
nd2 = round(Int,nn/2)
p = OffsetArray{Complex{Float64}}(undef,-round(Int,nn/2):round(Int,nn/2-1))
p[:] = zeros(Complex{Float64},nn)
p[-2:2] = [0.0 0.0 1.0 0.0 0.0]
vv = OffsetArray{Complex{Float64}}(undef,-round(Int,nn/2):round(Int,nn/2-1))
vv[:] = zeros(Complex{Float64},nn)
# vv[-2:2] = [0.25 0.0 1.0 0.0 0.25]
vv[-2:2] = 2*[0.0 0.5im 0.0 -0.5im 0.0]
tt = OffsetArray{Complex{Float64}}(undef,-round(Int, nn/2):round(Int,nn/2-1))
tt[:] = zeros(Complex{Float64},nn)
tt[-1:1] = [0.0 1.0 0.0]

p = collect(p)
vv = collect(vv)
tt = collect(tt)
ϕ0 = [p;tt]

T = Complex{Float64}
nu = length(vv)
nv = length(vv)
n = nu+nv-1
np2 = n > 1024 ? nextprod([2,3,5],n) : nextpow(2,n)
upad = [vv;zeros(T,np2-nu)]
vpad = [vv;zeros(T,np2-nv)]

plan = plan_fft!(upad)
iplan = plan_ifft!(upad)
y = similar(upad)

A = 4.0
B = [1.0 100.0]
force = 4.0
T0 = 1.0

params = SpectralParameters{Float64}(plan,iplan,A,B |> vec,force,vv,T0)
tspan = (0.0,0.5)
prob = ODEProblem(spectral_rhs!,ϕ0,tspan,params)

@testset "spectral_rhs!" begin
    jac_diff = DiffEqDiffTools.finite_difference_jacobian((dϕ,ϕ) -> spectral_rhs!(dϕ,ϕ,params,0),ϕ0)
    jac = spzeros(Complex{Float64},2nn,2nn)
    jac = spectral_rhs!(Val{:jac},jac,ϕ0,params,0)
    @test norm(jac*ϕ0 - jac_diff*ϕ0)/nn < 1e-9

    sol = solve(prob,TRBDF2(autodiff=false),abstol=1e-10,reltol=1e-10)

    jac_diff = DiffEqDiffTools.finite_difference_jacobian((dϕ,ϕ) -> spectral_rhs!(dϕ,ϕ,params,0),sol[end])
    jac = spzeros(Complex{Float64},2nn,2nn)
    jac = spectral_rhs!(Val{:jac},jac,sol[end],params,0)
    @test norm(jac*sol[end] - jac_diff*sol[end])/nn < 1e-9
end

@testset "solve_steady_state" begin
    integrator = init(prob,TRBDF2(autodiff=false))
    u = solve_steady_state(integrator,params,steadytol=1e-10,maxiters=5)
    @test norm([spectral_rhs!(copy(ϕ0),u,params,0);u[nd2+1]-1])/nn < 1e-10
end


nload{T<:Bus}(buses::Vector{T}) = length([b for b in buses if isa(b, LoadBus)])

ngen{T<:Bus}(buses::Vector{T}) = length([b for b in buses if isa(b, GenBus)])

count_bus_type(grid::Grid) = (nload(buses(grid)), ngen(buses(grid)))

#######################
### CHECK FUNCTIONS ###
#######################
# Everything here uses connectivty matrices as input instead of grids because
# all functions are supposed to work both for buses and for substations,
# effectively giving us two representations of the network (fine and coarse)
"""
    laplacian(con_mat::AbstractMatrix{<:Real})

Return the Laplacian matrix for a graph defined by adjacency matrix `con_mat`.
"""
function laplacian(con_mat::AbstractMatrix{<:Real})
    lap = zeros(Int, size(con_mat))
    conmat = con_mat .> 0 # adjusting for the case of substation connectivty matrix
    degree = sum(conmat, 1) # this gives the node degrees
    n = length(degree)
    for j in 1:(n - 1), i in (j + 1):n
        if conmat[i,j] > 0
            lap[i,j] = -1
        end
        lap[j,j] = degree[j]
    end
    lap[n,n] = degree[n]
    return Symmetric(lap, :L)
end

"""
    total_links(con_mat::AbstractMatrix{<:Real})

Return the total number of links in a system with connectivity matrix `con_mat`.
"""
function total_links(con_mat::AbstractMatrix{<:Real})
    return (1 / 2) * trace(laplacian(con_mat))
end

"""
    mean_node_deg(con_mat::AbstractMatrix{<:Real})

Return the average nodal degree of a system with connectivity matrix `con_mat`.
"""
function mean_node_deg{T}(con_mat::AbstractArray{T,2})
    return (1 / size(con_mat, 1)) * trace(laplacian(con_mat))
end

"""
    cluster_coeff(con_mat::AbstractMatrix{<:Real})

Return the clustering coefficient of a system with connectivity matrix `con_mat`.
"""
function cluster_coeff(con_mat::AbstractMatrix{<:Real})
    conmat = con_mat .> 0 # adjusting for the case of substation connectivty matrix
    degree = sum(conmat, 1) # this gives the node degrees
    clust_coeff = Vector{Float64}(length(degree))
    for i in 1:length(degree)
        maxedges = degree[i] * (degree[i] - 1) / 2
        cumul = 0
        for j in 1:length(degree)
            if conmat[i,j]
                cumul += dot(conmat[i,:],conmat[:,j])
            end
        end
        push!(clust_coeff, maxedges > 0 ? cumul/maxedges : 0)
    end
    return mean(clust_coeff)
end

"""
    cluster_coeff_degw(con_mat::AbstractMatrix{<:Real})

Return the degree-weighted clustering coefficient of a system with connectivity
matrix `con_mat`.
"""
function cluster_coeff_degw(con_mat::AbstractMatrix{<:Real})
    conmat = con_mat .> 0 # adjusting for the case of substation connectivty matrix
    degree = sum(conmat, 1) # this gives the node degrees
    clust_coeff = Vector{Float64}(length(degree))
    for i in 1:length(degree)
        maxedges = degree[i] * (degree[i] - 1) / 2
        cumul = 0
        for j in 1:length(degree)
            if conmat[i, j]
                cumul += dot(conmat[i,:], conmat[:,j])
            end
        end
        push!(clust_coeff, maxedges > 0 ? degree[i] * cumul / maxedges : 0)
    end
    return mean(clust_coeff) / mean(degree)
end

"""
    test_connectivity(con_mat::AbstractMatrix{<:Real}, verb=true)

Return `true` if the system with connectivity matrix `con_mat`is connected and `false`
otherwise. If `verb=true` the result will be printed onscreen together with the Fiedler
eigenvalue.
"""
function test_connectivity(con_mat::AbstractMatrix{<:Real}, verb=true)
    lap = laplacian(con_mat)
    evals = eigvals(lap)
    unique_vals = [v for v in Set(evals)]
    sort!(unique_vals)
    if length(unique_vals) > 1 && unique_vals[2] > 0
        fiedler = unique_vals[2]
        if verb
            println("Fiedler eigenvalue: $fiedler. System is connected.")
        end
        return true
    else
        if verb
            println("Fiedler eigenvalue equals zero. Network is NOT connected!")
        end
        return false
    end
end

function net_adj_mat(con_mat::AbstractMatrix{<:Real})
    lap = laplacian(con_mat)
    return Diagonal(lap) - lap
end

"""
    dijkstra(con_mat::AbstractMatrix{<:Real}, v1)

Compute minimum distances to each node in the network with connectivity matrix
con_mat, starting from node v1, using Dijkstra's algorithm.
"""
function dijkstra(con_mat::AbstractMatrix{<:Real}, v1)
    n = size(con_mat, 1)
    c = v1
    visited = [v1]
    dist = [i == v1 ? 0 : Inf for i in 1:n]
    r_dist = [Inf for d in dist]
    while length(visited) < n
        neigh = [i for i in 1:n if con_mat[c, i] > 0 && !(i in visited)]
        for i in neigh
            if dist[i] > (dist[c] + con_mat[c, i])
                dist[i] = (dist[c] + con_mat[c, i])
                r_dist[i] = dist[i]
            end
        end
        c = indmin(r_dist)
        push!(visited, c)
        r_dist[c] = Inf
    end
    return dist
end

"""
    mean_shortest_path(conmat::AbstractMatrix{<:Real})

Compute the mean shortest distance between nodes for a network represented by `conmat`.
The `conmat` matrix should specify the distance between each connected node and be equal to
zero whenever two nodes are not connected. Distances are computed using Dijkstra's algorithm.
"""
function mean_shortest_path(conmat::AbstractMatrix{<:Real})
    n = size(conmat, 1)
    mdist = [mean(dijkstra(conmat, i)) for i in 1:n]
    return mean(mdist)
end

"""
    mean_shortest_path(conmat::AbstractMatrix{<:Real}, distmat::AbstractMatrix{<:Real})

Compute the mean shortest distance between nodes for a network with connectivity matrix
`conmat` and distances between nodes given by `distmat`. Distances are computed using
Dijkstra's algorithm.
"""
function mean_shortest_path(
        conmat::AbstractMatrix{<:Real},
        distmat::AbstractMatrix{<:Real}
    )
    return mean_shortest_path(conmat.*distmat)
end

"""
    robustness_line(con_mat::AbstractMatrix{<:Real}, n=1000)

Return the average number of transmission lines that have to be randomly removed from a
system with connectivity matrix `con_mat` before it becomes disconnected. Average is
computed over n iterations.
"""
function robustness_line(con_mat::AbstractMatrix{<:Real}, n=1000)
    s = size(con_mat, 1)
    conns = []
    for i in 1:(s - 1), j in (i + 1):s
        con_mat[j, i] > 0 ? push!(conns, (i, j)) : nothing
    end
    total = 0
    for rep in 1:n
        connscopy = copy(conns)
        conmat = con_mat .> 0 # adjusting for the case of substation connectivty matrix
        step = 0
        while test_connectivity(conmat, false)
            k = rand(1:length(connscopy))
            ij = connscopy[k]
            conmat[ij[1],ij[2]] = false
            conmat[ij[2],ij[1]] = false
            deleteat!(connscopy, k)
            step += 1
        end
        total += step
    end
    return total / n
end

"""
    robustness_node(con_mat::AbstractMatrix{<:Real}, n=1000)

Return the average number of nodes that have to be randomly removed from a system with
connectivity matrix `con_mat` before it becomes disconnected. Average is computed over n
iterations.
"""
function robustness_node(con_mat::AbstractMatrix{<:Real}, n=1000)
    s = size(con_mat, 1)
    total = 0
    for rep in 1:n
        nodes = collect(1:s)
        conmat = con_mat .> 0 # adjusting for the case of substation connectivty matrix
        step = 0
        while test_connectivity(conmat, false)
            k = rand(1:length(nodes))
            i = nodes[k]
            for j in 1:s
                conmat[i,j] = false
                conmat[j,i] = false
            end
            deleteat!(nodes, k)
            step += 1
        end
        total += step
    end
    return total / n
end


# An example of the Burgers code for the
# propagation of Burgers equation in time
# Creates a colormap showing u(x, t) of Burgers equation
# initialized with the sine function used in lecture

using Burgers
using Printf
using PyPlot
include("Utils.jl")
using .Utils

Ttot   = 2.0   # total duration
Nt     = 1024  # number of time divisions
Ncells = 256   # number of position divisions
Nghost = 2     # padding on position divisions for boundary conditions

res   = Resolution{Float64, Int64}(Ttot, Nt, Ncells, Nghost)
sc    = minmod::SlopeCalc
u2fun = get_evolved(get_sinewave, res, sc)

# u(x, t) can be plotted as a colour map
xlabel("t")
ylabel("x")
pcolormesh(u2fun.tdat, u2fun.xdat, u2fun.u2dat)
colorbar()
savefig("images/colormap2.png", dpi = 300)


"""
Generate coefficients related to PML boundaries, e.g., damping profiles.
This function outputs a, b, and k variables in eq. 25-26 from Komatitsch 2007 (Geophysics).

	TODO
* At the moment, K=1, and is dummy, unlike in the case of EM when it was useful. Removing K and K_inv might save us memory when using GPUs.
"""
function update_pml!(
    pml::NamedVector{T},
    mgrid, # 1D grid
    flags::Vector{Bool},
    δt::Float64,
    np::Int64,
    velavg::Float64,
    freqpeak::Float64, # dominant frequency in Hz
) where {T<:Data.Array}

    δx = step(mgrid)
    nx = length(mgrid)

    NPOWER = 2.e0
    K_MAX_PML = 1.e0 # from Gedney page 8.11
    ALPHA_MAX_PML = pi * freqpeak # from Festa and Vilotte

    # thickness of the PML layer in meters
    thickness_PML = np * δx

    "reflection coefficient (INRIA report section 6.1) http://hal.inria.fr/docs/00/07/32/19/PDF/RR-3471.pdf"
    Rcoef = 0.001e0

    #! check that NPOWER is okay
    #if(NPOWER < 1) stop "NPOWER must be greater than 1"

    # compute d0 from INRIA report section 6.1 http://hal.inria.fr/docs/00/07/32/19/PDF/RR-3471.pdf
    d0 = -(NPOWER + 1) * (velavg) * log(Rcoef) / (2.e0 * thickness_PML)

    k = zeros(nx)
    k_half = zero(k)
    fill!(k, 1)
    fill!(k_half, 1)
    d = zeros(nx)
    d_half = zero(d)
    alpha = zero(d)
    alpha_half = zero(d)
    a=zero(d)
    b=zero(d)
    a_half=zero(d)
    b_half=zero(d)

    "damping in the X direction"
    "origin of the PML layer (position of right edge minus thickness, in meters)"
    xoriginleft = thickness_PML
    xoriginright = (nx - 1) * δx - thickness_PML

    for ix = 1:nx
        # abscissa of current grid point along the damping profile
        xval = δx * real(ix - 1)

        #---------- left edge
        if (flags[1])

            # define damping profile at the grid points
            abscissa_in_PML = xoriginleft - xval
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML
                d[ix] = d0 * abscissa_normalized .^ NPOWER
                # this taken from Gedney page 8.2
                k[ix] = 1.e0 + (K_MAX_PML - 1.e0) * abscissa_normalized .^ NPOWER
                alpha[ix] =
                    ALPHA_MAX_PML * (1.e0 - abscissa_normalized) + 0.01e0 * ALPHA_MAX_PML
            end

            # define damping profile at half the grid points
            abscissa_in_PML = xoriginleft - (xval + δx / 2.e0)
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML
                d_half[ix] = d0 * abscissa_normalized .^ NPOWER
                # this taken from Gedney page 8.2
                k_half[ix] = 1.e0 + (K_MAX_PML - 1.e0) * abscissa_normalized .^ NPOWER
                alpha_half[ix] =
                    ALPHA_MAX_PML * (1.e0 - abscissa_normalized) + 0.01e0 * ALPHA_MAX_PML
            end
        end

        #---------- right edge
        if (flags[2])

            # define damping profile at the grid points
            abscissa_in_PML = xval - xoriginright
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML
                d[ix] = d0 * abscissa_normalized^NPOWER
                # this taken from Gedney page 8.2
                k[ix] = 1.e0 + (K_MAX_PML - 1.e0) * abscissa_normalized^NPOWER
                alpha[ix] =
                    ALPHA_MAX_PML * (1.e0 - abscissa_normalized) + 0.01e0 * ALPHA_MAX_PML
            end

            # define damping profile at half the grid points
            abscissa_in_PML = xval + δx / 2.e0 - xoriginright
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML
                d_half[ix] = d0 * abscissa_normalized^NPOWER
                # this taken from Gedney page 8.2
                k_half[ix] = 1.e0 + (K_MAX_PML - 1.e0) * abscissa_normalized^NPOWER
                alpha_half[ix] =
                    ALPHA_MAX_PML * (1.e0 - abscissa_normalized) + 0.01e0 * ALPHA_MAX_PML
            end

        end
        #
        # just in case, for -5 at the end
        (alpha[ix] < 0.0) ? alpha[ix] = 0.0 : nothing
        (alpha_half[ix] < 0.0) ? alpha_half[ix] = 0.0 : nothing

        # see equation 25 Komatitsch, 2007, get b and a (need k, alpha, and d before this )
        b[ix] = exp(-(d[ix] / k[ix] + alpha[ix]) * δt)
        b_half[ix] = exp(-(d_half[ix] / k_half[ix] + alpha_half[ix]) * δt)

        # this to avoid division by zero outside the PML
        (abs(d[ix]) > 1.e-6) ?
        a[ix] = d[ix] * (b[ix] - 1.e0) / (k[ix] * (d[ix] + k[ix] * alpha[ix])) : nothing
        (abs(d_half[ix]) > 1.e-6) ?
        a_half[ix] =
            d_half[ix] * (b_half[ix] - 1.e0) /
            (k_half[ix] * (d_half[ix] + k_half[ix] * alpha_half[ix])) : nothing

    end
    copyto!(pml[:kI], inv.(k))
    copyto!(pml[:k_halfI], inv.(k_half))
    copyto!(pml[:a],a)
    copyto!(pml[:b],b)
    copyto!(pml[:a_half],a_half)
    copyto!(pml[:b_half],b_half)

end

"""
Generate a NamedArray with PML coefficients for all the dimensions that are then stored in the FDTD structs.
"""
function get_pml(mgrid)
    dnames = dim_names(length(mgrid))
    pnames = [:a, :b, :kI, :a_half, :b_half, :k_halfI]
    return NamedArray(
        [
            NamedArray(
                Data.Array.([
                    zeros(ni),
                    zeros(ni),
                    ones(ni),
                    zeros(ni),
                    zeros(ni),
                    ones(ni),
                ]),
                pnames,
            ) for ni in length.(mgrid)
        ],
        dnames,
    )
end

"""
Just a loop over dims of pml
"""
function update_pml!(
    pml::NamedVector{T},
    mgrid,
    pml_edges::Vector{Symbol},
    args...,
) where {T<:NamedVector}
    for (i, dim) in enumerate(dim_names(length(mgrid)))
        update_pml!(
            pml[dim],
            mgrid[i],
            [
                any(pml_edges .== Symbol(string(dim), "min")),
                any(pml_edges .== Symbol(string(dim), "max")),
            ],
            args...,
        )
    end
end

function update_pml!(pac)
    update_pml!(
        pac.pml,
        pac.exmedium.mgrid,
        pac.pml_edges,
        pac.fc[:dt],
        npml-3,
        Statistics.mean(pac.exmedium.bounds[:vp]),
        pac.fc[:freqpeak],
    )
end

for dimnames in [zip([:1, :2, :3], dim_names(3)), zip([:1, :2], dim_names(2))]
    is = broadcast(x -> Symbol(string("i", x)), getindex.(collect(dimnames), 2))
    ist = Meta.parse(string("(", [string(s, ",") for s in is]..., ")"))
    N = Meta.parse(string(length(is)))

    for (idim, dim) in dimnames
        i = Symbol("i", string(dim))
        ismoff = replace(is, i => :($i + moff))
        isdoff = replace(is, i => :(doff + $i))
        for (fname, fnamenp, idoff) in zip(
            [Symbol("memory", string(dim), "!"), Symbol("memory1", string(dim), "!")],
            [Symbol("memorynp", string(dim), "!"), Symbol("memorynp1", string(dim), "!")],
            [:($i + doff), :($i + doff + 1)],
        )
            @eval @parallel_indices(
                $ist,
                function $fnamenp(memory::Data.Array{$N}, d, a, b, kI, moff, doff)
                    memory[$(ismoff...)] =
                        b[$idoff] * memory[$(ismoff...)] + a[$idoff] * d[$(isdoff...)]
                    d[$(isdoff...)] = d[$(isdoff...)] * kI[$idoff] + memory[$(ismoff...)]
                    return
                end
            )
            @eval function $fname(memory::Data.Array{$N}, d, a, b, kI)
                sm = collect(size(memory))
                setindex!(sm, npml, $idim)
                # first npml points
                @parallel map(x -> (:)(1, x), Tuple(sm)) $fnamenp(memory, d, a, b, kI, 0, 0)
                # last npml points independent of d
                @parallel map(x -> (:)(1, x), Tuple(sm)) $fnamenp(
                    memory,
                    d,
                    a,
                    b,
                    kI,
                    npml,
                    getindex(size(d), $idim) - npml,
                )
            end
        end
    end
end


"""
    Double pendulum
"""
struct Impact end
struct Nominal end

struct DoublePendulum{T,X} <: Model{T}
    nq::Int
    nu::Int
    nw::Int
    nc::Int

    m1::T    # mass link 1
    J1::T    # inertia link 1
    l1::T    # length link 1
    lc1::T   # length to COM link 1

    m2::T    # mass link 2
    J2::T    # inertia link 2
    l2::T    # length link 2
    lc2::T   # length to COM link 2

    g::T     # gravity

    b1::T    # joint friction
    b2::T
end

lagrangian(model::DoublePendulum, q, q̇) = 0.0

function kinematics(model::DoublePendulum, x)
    [model.l1 * sin(x[1]) + model.l2 * sin(x[1] + x[2]),
     -1.0 * model.l1 * cos(x[1]) - model.l2 * cos(x[1] + x[2])]
end

function kinematics_elbow(model::DoublePendulum, x)
    [model.l1 * sin(x[1]),
     -1.0 * model.l1 * cos(x[1])]
end

function M_func(model::DoublePendulum, x)
    a = (model.J1 + model.J2 + model.m2 * model.l1 * model.l1
         + 2.0 * model.m2 * model.l1 * model.lc2 * cos(x[2]))

    b = model.J2 + model.m2 * model.l1 * model.lc2 * cos(x[2])

    c = model.J2

    [a b;
     b c]
end

function τ_func(model::DoublePendulum, x)
    a = (-1.0 * model.m1 * model.g * model.lc1 * sin(x[1])
         - model.m2 * model.g * (model.l1 * sin(x[1])
         + model.lc2 * sin(x[1] + x[2])))

    b = -1.0 * model.m2 * model.g * model.lc2 * sin(x[1] + x[2])

    [a, b]
end

function c_func(model::DoublePendulum, q, q̇)
    a = -2.0 * model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[2]
    b = -1.0 * model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[2]
    c = model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[1]
    d = 0.0

    [a b;
     c d]
end

function B_func(model::DoublePendulum, x)
    [0.0; 1.0]
end

function C_func(model::DoublePendulum, q, q̇)
    c_func(model, q, q̇) * q̇ - τ_func(model, q)
end

function ϕ_func(model, q)
    [0.5 * π - q[2], q[2] + 0.5 * π]
end

function P_func(model, q)
    ϕ = ϕ_func(model, q)
    Symbolics.jacobian(ϕ, q)
end

function dynamics(model::DoublePendulum{T,Impact}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2)

	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ B_func(model, qm2) * u1[1]
        + transpose(P_func(model, q2)) * λ1
        - h[1] * 0.5 .* vm2) # damping
end

function dynamics(model::DoublePendulum{T,Nominal}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2)

	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ B_func(model, qm2) * u1[1]
        - h[1] * 0.5 .* vm2) # damping
end

function residual(model::DoublePendulum{T,Impact}, z, θ, κ) where T
    nq = model.nq
    nu = model.nu
    nc = model.nc

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    q2 = z[1:nq]
    λ1 = z[nq .+ (1:nc)]
    s1 = z[nq + nc .+ (1:nc)]

    [
     dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b),
        h, q0, q1, u1, zeros(model.nw), λ1, q2);
     s1 .- ϕ_func(model, q2);
     λ1 .* s1 .- κ;
    ]

end

function residual(model::DoublePendulum{T,Nominal}, z, θ, κ) where T
    nq = model.nq
    nu = model.nu
    nc = model.nc

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    q2 = z[1:nq]
    return dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b),
        h, q0, q1, u1, zeros(model.nw), zeros(model.nc), q2)
end

acrobot_impact = DoublePendulum{Float64,Impact}(2, 1, 0, 2,
    1.0, 0.333, 1.0, 0.5, 1.0, 0.333, 1.0, 0.5, 9.81, 0.0, 0.0)

acrobot_nominal = DoublePendulum{Float64,Nominal}(2, 1, 0, 0,
    1.0, 0.333, 1.0, 0.5, 1.0, 0.333, 1.0, 0.5, 9.81, 0.0, 0.0)


data = map(x->split(x, " -> "),split(strip(read(joinpath(@__DIR__, "./5.txt"), String)), '\n'))
data = map(y->map(x->parse.(Int,split(x,",")),y),data)
data = map(z->reduce(vcat,z),data)

line_plot = zeros(Int,1000,1000)

function points_on_line(x1,y1,x2,y2)
  if x1 > x2
    x1, y1, x2, y2 = x2, y2, x1, y1
  end
  slope = (y2-y1)/(x2-x1)
  y_intercept = y1-slope*x1
  [(x,Int(x*slope+y_intercept)) for x in range(x1, x2) if Int(x*slope) == x*slope]
end

function plotline(x1,y1,x2,y2)
  global line_plot

  if x1 == x2
    if y1 < y2
      line_plot[x1,y1:y2] .+= 1
    else
      line_plot[x1,y2:y1] .+= 1
    end
  elseif y1 == y2
    if x1 < x2
      line_plot[x1:x2,y1] .+= 1
    else
      line_plot[x2:x1,y1] .+= 1
    end
  else
    # this is for part2 of the problem
    points = points_on_line(x2,y2,x1,y1)
    for (x,y) in points
      line_plot[x,y] += 1
    end
  end
end

function solve(data)
  global line_plot

  for d in data
    plotline(d...)
  end

  length(filter(i->i>=2,line_plot))
end

solve(data)


"""
    mechanical_energy(mechanism, storage)

    mechanism's total mechanical energy

    mechanism: Mechanism
    storage: Storage
"""
function mechanical_energy(mechanism::Mechanism, storage::Storage)
    ke = kinetic_energy(mechanism, storage)
    pe = potential_energy(mechanism, storage)
    me = ke + pe
    return me
end

"""
    kinetic_energy(mechanism, storage)

    mechanism's kinetic energy from linear and angular velocity

    mechanism: Mechanism
    storage: Storage
"""
function kinetic_energy(mechanism::Mechanism, storage::Storage{T,N}) where {T,N}
    ke = zeros(T,N)
    for i = 1:N
        ke[i] = kinetic_energy(mechanism, storage, i)
    end
    return ke
end

function kinetic_energy(mechanism::Mechanism, storage::Storage{T,N}, t::Int) where {T,N}
    ke = 0.0
    for (i,body) in enumerate(mechanism.bodies)
        vl = storage.vl[i][t]
        ωl = storage.ωl[i][t]
        ke += 0.5 * body.mass * vl' * vl
        ke += 0.5 * ωl' * body.inertia * ωl
    end
    return ke
end

"""
    potential_energy(mechanism, storage)

    mechanism's potential energy from gravity and springs

    mechanism: Mechanism
    storage: Storage
"""
function potential_energy(mechanism::Mechanism, storage::Storage{T,N}) where {T,N}
    pe = zeros(T,N)
    for i = 1:N
        pe[i] = potential_energy(mechanism, storage, i)
    end
    return pe
end

function potential_energy(mechanism::Mechanism{T,Nn,Ne,Nb}, storage::Storage{T,Ns}, t::Int) where {T,Nn,Ne,Nb,Ns}
    pe = 0.0

    # Gravity
    for (i,body) in enumerate(mechanism.bodies)
        pe -= body.mass * dot(mechanism.gravity, storage.x[i][t]) # TODO: confirm this is correct
    end

    # Springs
    for joint in mechanism.joints
        parent_id = joint.parent_id
        child_id = joint.child_id
        if joint.spring
            for element in (joint.translational, joint.rotational)
                if element.spring > 0
                    xb = storage.x[child_id - Ne][t] # TODO this is sketchy way to get the correct index
                    qb = storage.q[child_id - Ne][t] # TODO this is sketchy way to get the correct index

                    if parent_id != 0
                        xa = storage.x[parent_id - Ne][t] # TODO this is sketchy way to get the correct index
                        qa = storage.q[parent_id - Ne][t] # TODO this is sketchy way to get the correct index
                    else
                        xa, qa = current_configuration(mechanism.origin.state)
                    end

                    force = spring_force(:parent, element, xa, qa, xb, qb)
                    pe += 0.5 * dot(force, force) ./ element.spring
                end
            end
        end
    end
    return pe
end


using Base.Test, Convex, SCS, SchattenNorms

set_default_solver(SCSSolver(verbose=0))

x,y,z = [0 1; 1 0], [0 -1im; 1im 0], [1 0; 0 -1]

liou(x) = kron(conj(x),x)
liou(x,y) = kron(transpose(y),x)

function rand_unitary(d)
    rm = randn(d,d)+1im*randn(d,d)
    u,_,_ = svd(rm)
    return u
end

cnot = [1. 0. 0. 0.; 0. 1. 0. 0.; 0. 0. 0. 1.; 0. 0. 1. 0.]
cnotL = liou(cnot)

#Check the diamond norm of randomly chosen Pauli channels
p1 = rand(4); p1 = p1/sum(p1)
p2 = rand(4); p2 = p2/sum(p2)
pauli_dnorm = sum(abs(p1-p2))
pauli = Matrix{Complex128}[ [1 0; 0 1],
                            [0 1; 1 0],
                            [0 -im; im 0],
                            [1 0; 0 -1] ]

u = liou(rand_unitary(2))
v = liou(rand_unitary(2))

#ee = rand_cp_map(2)
#ff = rand_cp_map(2)

@test_approx_eq_eps 2.0 dnorm(cnotL-eye(16)) 1e-5
@test_approx_eq_eps 2.0 dnorm(eye(4)-liou(x,x)) 1e-5
@test_approx_eq_eps 2.0 dnorm(eye(4)-liou(z,z)) 1e-5
@test_approx_eq_eps 2.0 dnorm(eye(4)-liou(y,y)) 1e-5
@test_approx_eq_eps 2.0 dnorm(liou(x,x)-liou(y,y)) 1e-5
@test_approx_eq_eps 2.0 dnorm(liou(x,x)-liou(z,z)) 1e-5
@test_approx_eq_eps 2.0 dnorm(liou(z,z)-liou(y,y)) 1e-5

@test_approx_eq_eps 2.0 dnorm(u*eye(4)-u*liou(x,x)) 1e-5
@test_approx_eq_eps 2.0 dnorm(u*eye(4)-u*liou(z,z)) 1e-5
@test_approx_eq_eps 2.0 dnorm(u*eye(4)-u*liou(y,y)) 1e-5
@test_approx_eq_eps 2.0 dnorm(u*liou(x,x)-u*liou(y,y)) 1e-5
@test_approx_eq_eps 2.0 dnorm(u*liou(x,x)-u*liou(z,z)) 1e-5
@test_approx_eq_eps 2.0 dnorm(u*liou(z,z)-u*liou(y,y)) 1e-5

pc1 = sum([ p1[ii+1]*liou(pauli[ii+1]) for ii in 0:3])
pc2 = sum([ p2[ii+1]*liou(pauli[ii+1]) for ii in 0:3])
calc_dnorm1 = dnorm(pc1-pc2)

@test_approx_eq_eps pauli_dnorm calc_dnorm1  1e-5

duv  = dnorm(u-v)
@time begin
    duv2 = dnorm(u*v'-eye(4))
    duv3 = dnorm(v'u-eye(4))
    duv4 = dnorm(eye(4)-u'*v)
    duv5 = dnorm(eye(4)-v*u')
end

@test_approx_eq_eps duv duv2 1e-5
@test_approx_eq_eps duv duv3 1e-5
@test_approx_eq_eps duv duv4 1e-5
@test_approx_eq_eps duv duv5 1e-5

#def = dnorm(ee-ff)
#@time begin
#    def2 = dnorm(u*(ee-ff)*v)
#    def3 = dnorm(u*(ee-ff))
#    def4 = dnorm((ee-ff)*v)
#    def5 = dnorm(u'*(ee-ff))
#end

#println(def)
#println(def2)
#println(def3)
#println(def4)
#println(def5)


# SCHEME
abstract type Scheme end

# By default, scheme definition can ignore time step
(scheme::Scheme)(flux, mesh, w, i_face, dt) = scheme(flux, mesh, w, i_face)

####################################
struct Centered <: Scheme end
####################################

function (::Centered)(flux::LinearAdvectionFlux, mesh, w, i_face)
    n = normal_vector(mesh, i_face)
    i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, i_face)
    return 0.5*(flux(w[i_cell_1], n) + flux(w[i_cell_2], n))
end

####################################
struct Downwind <: Scheme end
####################################

function (::Downwind)(flux::LinearAdvectionFlux, mesh, w, i_face)
    n = normal_vector(mesh, i_face)
    i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, i_face)
    λ = eigvals(flux, (w[i_cell_1] + w[i_cell_2])/2, n)
    return λ < zero(λ) ? flux(w[i_cell_1], n) : flux(w[i_cell_2], n)
end

##################################
struct Upwind <: Scheme end
##################################

# Single-wave problems
function (::Upwind)(flux::LinearAdvectionFlux, mesh, w, i_face)
    n = normal_vector(mesh, i_face)
    i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, i_face)
    λ = eigvals(flux, (w[i_cell_1] + w[i_cell_2])/2, n)
    return λ > zero(λ) ? flux(w[i_cell_1], n) : flux(w[i_cell_2], n)
end

# Multi-wave problems
function (::Upwind)(flux, mesh, w, i_face)
    n = normal_vector(mesh, i_face)
    i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, i_face)
    w_mean = (w[i_cell_1] + w[i_cell_2])/2
    λ, R = eigen(flux, w_mean, n)
    left_to_right = λ .> zero(λ)
    if all(left_to_right)
        return flux(w[i_cell_1], n)
    elseif !any(left_to_right)
        return flux(w[i_cell_2], n)
    else
        f₁ = flux(w[i_cell_1], n)
        f₂ = flux(w[i_cell_2], n)
        decentrement = R \ (f₁ - f₂)
        return (f₁ + f₂ + R*diagm(sign.(λ))*decentrement)/2
        # L = inv(R)
        # return (f₁ + f₂ + R*diagm(sign.(λ))*L*(f₁ - f₂))/2
    end
end

###################################
abstract type BoundaryCondition end
###################################

(scheme::BoundaryCondition)(model, mesh, w, i_face, dt) = scheme(model, mesh, w, i_face)

#########################################
struct NeumannBC <: BoundaryCondition end
#########################################

function (::NeumannBC)(flux, mesh, w, i_face)
    i_cell_1 = cell_next_to_boundary_face(mesh, i_face)
    return flux(w[i_cell_1], normal_vector(mesh, i_face))
end

################################################################################
#                                     Div                                      #
################################################################################

numerical_flux(flux, mesh, w, Φ, i_face, dt) = Φ(flux, mesh, w, i_face, dt)

function div!(Δw, flux, mesh, w, scheme::Scheme, dt=0.0)
    @inbounds for i_face in inner_faces(mesh)
        ϕ = numerical_flux(flux, mesh, w, scheme, i_face, dt)
        i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, i_face)
        Δw[i_cell_1] += ϕ * face_area(mesh, i_face) / cell_volume(mesh, i_cell_1)
        Δw[i_cell_2] -= ϕ * face_area(mesh, i_face) / cell_volume(mesh, i_cell_2)
    end
end

function div!(Δw, flux, mesh, w, scheme::BoundaryCondition, dt=0.0)
    @inbounds for i_face in boundary_faces(mesh)
        ϕ = numerical_flux(flux, mesh, w, scheme, i_face, dt)
        i_cell = cell_next_to_boundary_face(mesh, i_face)
        Δw[i_cell] += ϕ * face_area(mesh, i_face) / cell_volume(mesh, i_cell)
    end
end

function div!(Δw, flux, mesh, w, schemes::NTuple{N, Union{Scheme, BoundaryCondition}}, dt=0.0) where N
    for scheme in schemes
        div!(Δw, flux, mesh, w, scheme, dt)
    end
end

flux_in_cell(flux, mesh, w, scheme, dt, i_face, i_cell) = numerical_flux(flux, mesh, w, scheme, i_face, dt) * face_area(mesh, i_face) / cell_volume(mesh, i_cell)
function Δw_type(flux, mesh, w, schemes::NTuple{N, Union{Scheme, BoundaryCondition}}, dt) where N
    T = Base.return_types(flux_in_cell, typeof.((flux, mesh, w, schemes[1], dt, first(inner_faces(mesh)), first(all_cells(mesh)))))[1]
    if T in (Union{}, Any)
        return eltype(w)
    else
        return T
    end
end
Δw_type(f, m, w, s::Union{Scheme, BoundaryCondition}, dt) = Δw_type(f, m, w, (s,), dt)

function div(flux, mesh, w, schemes::NTuple{N, Union{Scheme, BoundaryCondition}}, dt=0.0) where N
    Δw = zeros(Δw_type(flux, mesh, w, schemes, dt), size(w))
    div!(Δw, flux, mesh, w, schemes, dt)
    return Δw
end

div(flux, mesh, w, scheme::Union{Scheme, BoundaryCondition}, dt=0.0) = div(flux, mesh, w, (scheme,), dt)


# ---------------------------------------------------------------------------- #
#
#   basisFuncLineLag.jl
#
#   Get Legendre basis functions for Line
#
#   λυτέος
#   Fall 2017
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

# function basisFuncLineLeg( ::Type{Val{0}}, s::Array{Float64} )
#
#   dim = 1
#   nx  = length( s )
#
#   # Constant
#   nphi = 1
#   phi  = Array{Float64}(nx, nphi)
#   dphi = Array{Float64}(nx, nphi, dim)
#
#   # value
#   phi[:,1] = 1.0
#
#   # derivative in x
#   dphi[:,1,1] = 0.0
#
#   # transpose
#   phi   = phi'
#   dphi  = dphi[:,:,1]'
#
#   return phi, dphi
#
# end

function basisFuncLineLeg( P::P1, s::Array{Float64} )

  dim = 1
  nx  = length( s )

  # Linear
  nphi = 2
  phi  = Array{Float64}(nx, nphi)
  dphi = Array{Float64}(nx, nphi, dim)

  t = 1 - s

  phi[:,1] = 1
  phi[:,2] = sqrt(3).*(s - t)

  dphi[:,1,1] =  0
  dphi[:,2,1] =  2.*sqrt(3)

  # transpose
  phi   = phi'
  dphi  = dphi[:,:,1]'

  return phi, dphi

end

function basisFuncLineLeg( P::P2, s::Array{Float64} )

  dim = 1
  nx  = length( s )

  # Quadratric
  nphi = 3
  phi  = Array{Float64}(nx, nphi)
  dphi = Array{Float64}(nx, nphi, dim)

  t = 1 - s

  phi[:,1] = 1
  phi[:,2] = sqrt(3).*(s - t)
  phi[:,3] = sqrt(5).*(1 - 6.*s.*t)

  dphi[:,1,1] =  0
  dphi[:,2,1] =  2.*sqrt(3)
  dphi[:,3,1] =  6.*sqrt(5).*(s - t)

  # transpose
  phi   = phi'
  dphi  = dphi[:,:,1]'

  return phi, dphi

end

function basisFuncLineLeg( P::P3, s::Array{Float64} )

  dim = 1
  nx  = length( s )

  # Cubic
  nphi = 4
  phi  = Array{Float64}(nx, nphi)
  dphi = Array{Float64}(nx, nphi, dim)

  t = 1 - s

  phi[:,1] = 1
  phi[:,2] = sqrt(3).*(s - t)
  phi[:,3] = sqrt(5).*(1 - 6.*s.*t)
  phi[:,4] = sqrt(7).*(1 - 10.*s.*t).*(s - t)

  dphi[:,1,1] =  0
  dphi[:,2,1] =  2.*sqrt(3)
  dphi[:,3,1] =  6.*sqrt(5).*(s - t)
  dphi[:,4,1] = 12.*sqrt(7).*(1 - 5.*s.*t)

  # transpose
  phi   = phi'
  dphi  = dphi[:,:,1]'

  return phi, dphi

end

function basisFuncLineLeg( P::P4, s::Array{Float64} )

  dim = 1
  nx  = length( s )

  # Quartic
  nphi = 5
  phi  = Array{Float64}(nx, nphi)
  dphi = Array{Float64}(nx, nphi, dim)

  t = 1 - s

  phi[:,1] = 1
  phi[:,2] = sqrt(3).*(s - t)
  phi[:,3] = sqrt(5).*(1 - 6.*s.*t)
  phi[:,4] = sqrt(7).*(1 - 10.*s.*t).*(s - t)
  phi[:,5] = 3.*(1 - 10.*s.*t.*(2 - 7.*s.*t))

  dphi[:,1,1] =  0
  dphi[:,2,1] =  2.*sqrt(3)
  dphi[:,3,1] =  6.*sqrt(5).*(s - t)
  dphi[:,4,1] = 12.*sqrt(7).*(1 - 5.*s.*t)
  dphi[:,5,1] = 60.*(1 - 7.*s.*t).*(s - t)

  # transpose
  phi   = phi'
  dphi  = dphi[:,:,1]'

  return phi, dphi

end

# function basisFuncLineLeg( P::P5, s::Array{Float64} )
#
#   dim = 1
#   nx  = length( s )
#
#   # Quintic
#   nphi = 6
#   phi  = Array{Float64}(nx, nphi)
#   dphi = Array{Float64}(nx, nphi, dim)
#
#   t = 1 - s
#
#   phi[:,1] = 1;
#   phi[:,2] = sqrt(3).*(s - t);
#   phi[:,3] = sqrt(5).*(1 - 6.*s.*t);
#   phi[:,4] = sqrt(7).*(1 - 10.*s.*t).*(s - t);
#   phi[:,5] = 3.*(1 - 10.*s.*t.*(2 - 7.*s.*t));
#   phi[:,6] = sqrt(11).*(1 - 14.*s.*t.*(s - 2.*t).*(2.*s - t)).*(s - t);
#
#
#   dphi[:,1,1] =  0
#   dphi[:,2,1] =  2.*sqrt(3)
#   dphi[:,3,1] =  6.*sqrt(5).*(s - t)
#   dphi[:,4,1] = 12.*sqrt(7).*(1 - 5.*s.*t)
#   dphi[:,5,1] = 60.*(1 - 7.*s.*t).*(s - t)
#   dphi[:,6,1] = 30.*sqrt(11).*(1 - 14.*s.*t.*(1 - 3.*s.*t))
#
#   # transpose
#   phi   = phi'
#   dphi  = dphi[:,:,1]'
#
#   return phi, dphi
#
# end


"""
    simulate_forward(eventHistory, maxT, startT, bg, kappa, kern)

Simulate a Hawkes process given we have already observed some events. This can be used to forecast future events.

# Arguments
* `eventHistory` The events that have ocured so far.
* `maxT` Maximum time value to simulate forward to.
* `startT` Starting time from which to simulate forward.
* `bg` Background value
* `kappa` Kappa value
* `kern` Kernel function or distribution

"""
function simulate_forward(eventHistory, maxT::Number, startT::Number, bg::Number, kappa::Number, kern::Function)
    bgF(x) = intensity(collect(x .+ startT), eventHistory, bg, kappa, kern)
    events = simulate(bgF, kappa, kern, maxT)
    events .+ startT
end

function simulate_forward(eventHistory, maxT::Number, startT::Number, bg::Number, kappa::Number, kern::Distribution)
    kernD, kern_cdf = generate_kernel_functions(kern)
    simulate_forward(eventHistory, maxT, startT, bg, kappa, kernD)
end


includet("simulation_model.jl")
using StatsBase, GLM, DataFrames, LinearAlgebra, Plots

function estimateOptimalPolicy()
    players = SCTypeData(10000)
    d = 10000
    βs = rand(Uniform(0.5, 1.1), d)

    μs = zeros(d, 2)
    σs = zeros(d, 4)

    for i in 1:d
        data = ObservedData(players, [0, βs[i]])
        σs[i, :] = reshape(cov([data.y data.x]), (1, 4))
        μs[i, 1] = mean(data.y)
        μs[i, 2] = mean(data.x)
     end

    display(plot(βs, σs[:, 4], seriestype=:scatter))
    μ_coef = [coef(lm(@formula(μ ~ β + β^2), DataFrame(μ=μs[:, i], β=βs))) for i in 1:2]
    σ_coef = [coef(lm(@formula(σ ~ β + β^2), DataFrame(σ=σs[:, i], β=βs))) for i in 1:4]
    function loss(β)
        μb = [ dot(μ_coef[i], [1, β, β^2]) for i in 1:2]
        σb = reshape([ dot(σ_coef[i],[1, β, β^2]) for i in 1:4], (2,2))
        samples = rand(MvNormal(μb, σb), 1000000)
        β0 = mean(samples[1, :] .- samples[2, :].*β)
        return mean((samples[1, :] .- β0 .- samples[2, :]*β).^2)
    end

    βstar = Optim.optimize(loss, 0.0, 1.0).minimizer
    return βstar
end

#then we estimate a linear function for mean and variance of the x variables

# then we optimize the loss function when we integrate over a normal distribution


using ContinuousWavelets, Plots, Wavelets, FFTW, Logging
global_logger(SimpleLogger(min_level = Logging.Error))
global_logger(Logging.SimpleLogger(stderr, Logging.Error))
n = 2047;
t = range(0, n / 2000, length = n); # 2kHz sampling rate
f = testfunction(n, "Doppler");
p1 = plot(t, f, legend = false, title = "Doppler", xticks = false, linewidth = 2)
c = wavelet(Morlet(π), β = 2);
res = ContinuousWavelets.cwt(f, c)
# plotting
freqs = getMeanFreq(ContinuousWavelets.computeWavelets(n, c)[1])
freqs[1] = 0
p2 = heatmap(t, freqs, abs.(res)', xlabel = "time (s)", ylabel = "frequency (Hz)", colorbar = false, c = :viridis)
l = @layout [a{0.3h}; b{0.7h}]
plot(p1, p2, layout = l)
savefig("../docs/doppler.svg");#hide
DEFAULT_BOUNDARY;
supertype(typeof(morl));
a, b = computeWavelets(n, c)
heatmap(a)
using LaTeXStrings
dRate = 4
waveType = Morlet()
Ψ1 = wavelet(waveType, s = 8, β = dRate, averagingLength = 2)
# sketch of how the frequencies are chosen
pyplot()
locs = ContinuousWavelets.polySpacing(8, Ψ1);
#Figure 3.1
scatter(1:length(locs), locs, legend = :bottomright, label = "mean log frequency", xlabel = "Wavelet Index (x)", ylabel = "log-Frequency (y)", color = :black)
scatter!(length(locs):length(locs), locs[end:end], markersize = 10, markershape = :x, color = :black, label = :none)
firstW, lastW, stepS = ContinuousWavelets.genSamplePoints(8, Ψ1)
b = (dRate / Ψ1.Q)^(1 ./ dRate) * (8 + Ψ1.averagingLength)^((dRate - 1) / dRate)
t = range(1, stop = length(locs), step = 0.1)
curve = b .* (range(firstW, stop = (locs[end] / b)^dRate, length = length(t))) .^ (1 / dRate)
plot!(t, curve, color = :blue, line = :dash, label = L"y=a(mx+x_0)^{^1/_\beta}", legend = :bottomright, legendfontsize = 12, xrange = (0, length(locs) + 3), xticks = [1; 5:5:1+length(locs)...], yrange = (minimum(locs) - 1, maximum(locs) + 1), yticks = (range(floor(Int, minimum(locs)), ceil(Int, maximum(locs)), step = 2), [L"\alpha", (4:2:8)..., "N.Octaves"]))
x = range(15, stop = 28, step = 0.5)
ycord(x) = locs[end] .+ b / dRate * 24 .^ (1 / dRate - 1) .* (x .- length(locs))
plot!(x, ycord(x), c = :black, line = 2, label = :none)
annotate!(length(locs) - 1 / 8, locs[end] + 7.5 / 16, Plots.text(L"\frac{dy}{dx}=^{1}/_{Q}", 11, :black, :center))
savefig("plotOfLogCentralFrequencies.svg")

n = 2047
Ψ1 = wavelet(morl, s = 8, β = 1)
d1, ξ = computeWavelets(n, Ψ1)
Ψ2 = wavelet(morl, s = 8, β = 2)
d2, ξ = Wavelets.computeWavelets(n, Ψ2)
Ψ4 = wavelet(morl, s = 8, β = 4)
d4, ξ = Wavelets.computeWavelets(n, Ψ4)
matchingLimits = (minimum([d1 d2 d4]), maximum([d1 d2 d4]))# for
plot(heatmap(1:size(d1, 2), ξ, d1, color = :Greys,
        yaxis = (L"\omega",),
        xaxis = ("wavelet index",),
        title = L"\beta=1" * " (" * L"\Psi1" * ")", colorbar = false,
        clims = matchingLimits),
    heatmap(1:size(d2, 2), ξ, d2, color = :Greys,
        yticks = [],
        xaxis = ("wavelet index",),
        title = L"\beta=2" * " (" * L"\Psi2" * ")", colorbar = false,
        clims = matchingLimits),
    heatmap(1:size(d4, 2), ξ, d4, color = :Greys, yticks = [],
        xaxis = ("wavelet index",),
        title = L"\beta=4" * " (" * L"\Psi4" * ")"),
    layout = (1, 3), clims = matchingLimits,
    colorbar_title = L"\widehat{\psi_i}")#hide

savefig("changeBeta.png")
aasdf = randn(10, 10)
plot(aasdf, aspect_ratio = 1.0)
plot(aasdf)
l = @layout [grid(4, 3) [a{0.5w}; b{0.5w}; c{0.5w}; d{0.5w}]]
ps = [heatmap(i * j * rand(10, 10), colorbar = false) for i = 1:4, j = 1:4]
plot(ps..., layout = l, size = (1100, 800))
Q = 8
eval(:Q)

using ContinuousWavelets, Plots, Wavelets, FFTW, Logging
using Plots;
gr();
Plots.reset_defaults()
global_logger(Logging.SimpleLogger(stderr, Logging.Error))
n = 2047;
function mapTo(waveType, isReal = true, window = 1:2047; d = 1, kwargs...)
    if isReal
        c = wavelet(waveType; β = d, kwargs...)
        waves, ω = computeWavelets(n, c)
        return circshift(irfft(waves, 2 * n, 1), (1024, 0))[window, :]
    else
        c = wavelet(waveType; β = d, kwargs...)
        waves, ω = computeWavelets(n, c)
        waves = cat(waves, zeros(2047, size(waves, 2)), dims = 1)
        return circshift(ifft(waves, 1), (1024, 0))[window, :]
    end
end
tmp = mapTo(Morlet(π), false; averagingLength = -0.2)[:, 2]
p1 = plot([real.(tmp) imag.(tmp)], title = "Morlet", labels = ["real" "imaginary"], ticks = nothing, linewidth = 5)
tmp = mapTo(paul2, false, averagingLength = -.5)[:, 2]
p2 = plot([real.(tmp) imag.(tmp)], title = "Paul 2", labels = ["real" "imaginary"], ticks = nothing, linewidth = 5)
p3 = plot(mapTo(dog2; averagingLength = -1.5)[:, 2], title = "derivative of gaussians (dog2)", legend = false, ticks = nothing, linewidth = 5)
p4 = plot(mapTo(cHaar, true; averagingLength = 1)[:, 2], title = "Haar", legend = false, ticks = nothing, linewidth = 5)
p5 = plot(mapTo(cBeyl, true; d = 1, averagingLength = -0)[:, 2], title = "Beylkyin", legend = false, ticks = nothing, linewidth = 5)
p6 = plot(mapTo(cVaid, true; d = 1, averagingLength = -0)[:, 2], title = "Vaidyanthan", legend = false, ticks = nothing, linewidth = 5)
p7 = plot(mapTo(cDb2; d = 1, averagingLength = -0)[:, 2], title = "Daubhechies 2", legend = false, ticks = nothing, linewidth = 5)
p8 = plot(mapTo(cCoif2, true; d = 1, averagingLength = -0)[:, 2], title = "Coiflet 2", legend = false, ticks = nothing, linewidth = 5)
p9 = plot(mapTo(cSym4, true; d = 1, averagingLength = -0)[:, 2], title = "Symlet 4", legend = false, ticks = nothing, linewidth = 5)
k = 0600;
p10 = plot(mapTo(cBatt4, true, 1024-k:1024+k; d = 1, averagingLength = -1)[:, 2], title = "Battle-Lemarie, 4", legend = false, ticks = nothing, linewidth = 5);
plot(p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, layout = (2, 5), size = 300 .* (5, 2.2))
savefig("mothers.svg")#hide
2047 / 2

c = wavelet(Morlet(π), averagingType = NoAve(), β = 2);
exs = cat(testfunction(n, "Doppler"), testfunction(n, "Blocks"), testfunction(n, "Bumps"), testfunction(n, "HeaviSine"), dims = 2)
p1 = plot(plot(exs[:, 1], legend = false, title = "Doppler", yticks = [], xticks = [], linewidth = 2), plot(exs[:, 2], legend = false, title = "Blocks", yticks = [], xticks = [], linewidth = 2), plot(exs[:, 3], legend = false, title = "Bumps", yticks = [], xticks = [], linewidth = 2), plot(exs[:, 4], legend = false, title = "HeaviSine", yticks = [], xticks = [], linewidth = 2), layout = (1, 4))
c = wavelet(cDb2, β = 2, extraOctaves = -0);
res = circshift(ContinuousWavelets.cwt(exs, c), (0, 1, 0))
p2 = plot(heatmap(identity.(res[:, :, 1])', xticks = false, yticks = [], c = :viridis, colorbar = false), heatmap(identity.(res[:, :, 2])', xticks = false, yticks = [], c = :viridis, colorbar = false), heatmap(identity.(res[:, :, 3])', xticks = false, yticks = [], c = :viridis, colorbar = false), heatmap(identity.(res[:, :, 4])', xticks = false, yticks = [], c = :viridis, colorbar = false), layout = (1, 4))
l = @layout [a{0.3h}; b{0.7h}]
plot(p1, p2, layout = l)
savefig("multiEx.svg")
ContinuousWavelets.getNOctaves(n, wavelet(cCoif2))
log2(n)
n / 2
c = wavelet(cCoif2, β = 1, Q = 1)
wave, ω = ContinuousWavelets.computeWavelets(n, wavelet(cCoif4, β = 1, Q = 1); space = true)
i = 1;
count(abs.(wave[:, i]) / norm(wave[:, i]) .> 1e-7);
plot(wave[1000:1090, end])
ContinuousWavelets.calcDepth(c, n)
log2(n)
log2(length(qmf(c.waveType)))
log2(n)

gr()
n = 2047
Ψ1 = wavelet(morl, s = 8, β = 1)
d1, ξ = computeWavelets(n, Ψ1)
Ψ2 = wavelet(morl, s = 8, β = 2)
d2, ξ = Wavelets.computeWavelets(n, Ψ2)
Ψ4 = wavelet(morl, s = 8, β = 4)
d4, ξ = Wavelets.computeWavelets(n, Ψ4)
matchingLimits = (minimum([d1 d2 d4]), maximum([d1 d2 d4]))# for
plot(heatmap(1:size(d1, 2), ξ, d1, color = :Greys,
        yaxis = (L"\omega",),
        xaxis = ("wavelet index",),
        title = L"\beta=1" * " (" * L"\Psi1" * ")", colorbar = false,
        clims = matchingLimits),
    heatmap(1:size(d2, 2), ξ, d2, color = :Greys,
        yticks = [],
        xaxis = ("wavelet index",),
        title = L"\beta=2" * " (" * L"\Psi2" * ")", colorbar = false,
        clims = matchingLimits),
    heatmap(1:size(d4, 2), ξ, d4, color = :Greys, yticks = [],
        xaxis = ("wavelet index",),
        title = L"\beta=4" * " (" * L"\Psi4" * ")"),
    layout = (1, 3), clims = matchingLimits,
    colorbar_title = L"\widehat{\psi_i}")

# readme examples
using ContinuousWavelets, Plots, Wavelets
n = 2047;
f = testfunction(n, "Doppler");
p1 = plot(f, legend = false, title = "Doppler", xlims = (0, 2000))
c = wavelet(Morlet(π), β = 2);
res = cwt(f, c)
p2 = heatmap(abs.(res)', xlabel = "time index",
    ylabel = "frequency index", colorbar = false)
l = @layout [a{0.3h}; b{0.7h}]
plot(p1, p2, layout = l)

f = testfunction(n, "Bumps");
p1 = plot(f, legend = false, title = "Bumps", xlims = (0, 2000), linewidth = 2)
c = wavelet(dog2, β = 2);
res = ContinuousWavelets.cwt(f, c)
# dropping the middle peaks
res[620:1100, :] .= 0
# and smoothing the remaining peaks
res[:, 10:end] .= 0
freqs = ContinuousWavelets.getMeanFreq(f, c)
p2 = heatmap(1:n, freqs, abs.(res)', xlabel = "time (ms)", ylabel = "Frequency (Hz)", colorbar = false, c = :viridis)
dropped = ContinuousWavelets.icwt(res, c, DualFrames())
p1 = plot(f, legend = false, title = "Smoothing and dropping bumps", linewidth = 2)
plot!(dropped, linewidth = 3)
l = @layout [a{0.3h}; b{0.7h}]
plot(p1, p2, layout = l)
savefig("../docs/bumps.svg")

n = 2047
f = testfunction(n, "Blocks") + 0.10randn(n)
f = testfunction(n, "HeaviSine")
wave = wavelet(morl)
fCWT = cwt(f, wave)
fRecon = ContinuousWavelets.icwt(fCWT, wave, DualFrames())
plot([real.(fRecon) f], labels = ["reconstructed" "original"])
plot(fRecon)
deNoiseCWT = copy(fCWT)
deNoiseCWT[abs.(fCWT)/norm(fCWT, Inf).<0.07] .= 0
heatmap((abs.(fCWT) / norm(fCWT, Inf))' .< 0.01)
heatmap(fCWT', c = :viridis)
heatmap(deNoiseCWT', c = :viridis)
heatmap(fCWT' - deNoiseCWT')
fRecon = ContinuousWavelets.icwt(deNoiseCWT, wave, DualFrames())
plot([fRecon f], labels = ["reconstructed" "original"])
plot(plot(f, title = "with noise"), plot(fRecon, title = "threshold and inverted"), layout = (2, 1), legend = false)
fRecon = ContinuousWavelets.icwt(fCWT, wave, DualFrames())
plot([real.(fRecon) f])

c = wavelet(morl, β = 2);
dualframe, dfNorm = getDualCoverage(n, c, DualFrames())
dualPenrose, dfNorm = getDualCoverage(n, c, PenroseDelta())
dualNaive, dfNorm = getDualCoverage(n, c, NaiveDelta())
plot([dualframe dualPenrose dualNaive], labels = ["Canonical" "Delta with least squares weights" "delta with naive weights"], legend = :topleft)

pyplot()
n = 2047
Ψ1 = wavelet(morl, s = 8, β = 1)
d1, ξ = computeWavelets(n, Ψ1)
Ψ2 = wavelet(morl, s = 8, β = 2)
d2, ξ = computeWavelets(n, Ψ2)
Ψ4 = wavelet(morl, s = 8, β = 4)
d4, ξ = computeWavelets(n, Ψ4)
matchingLimits = (minimum([d1 d2 d4]), maximum([d1 d2 d4]))#hide
plot(heatmap(1:size(d1, 2), ξ, d1, color = :Greys, yaxis = (L"\omega",), xaxis = ("wavelet index",), title = L"\beta=1" * " (" * L"\Psi1" * ")", colorbar = false, clims = matchingLimits), heatmap(1:size(d2, 2), ξ, d2, color = :Greys, yticks = [], xaxis = ("wavelet index",), title = L"\beta=2" * " (" * L"\Psi2" * ")", colorbar = false, clims = matchingLimits), heatmap(1:size(d4, 2), ξ, d4, color = :Greys, yticks = [], xticks = [1, 5, 10, 14, 18], xaxis = ("wavelet index",), title = L"\beta=4" * " (" * L"\Psi4" * ")"), layout = (1, 3), clims = matchingLimits, colorbar_title = L"\widehat{\psi_i}")
using LaTeXStrings
d4


export matching_poly

"""
`matching_poly(G)` returns the matching polynomial of the
`SimpleGraph` `G`.
"""
function matching_poly(G::SimpleGraph, cache_flag::Bool=true)
    if cache_flag && cache_check(G,:matching_poly)
        return cache_recall(G,:matching_poly)
    end
    if NE(G)==0
        p = Poly([0,1])^NV(G)
        if cache_flag
            SimpleGraphs.cache_save(G,:matching_poly,p)
        end
        return p
    end

    if is_connected(G)
        e = first(G.E)  # get any edge
        v,w = e

        G1 = deepcopy(G)
        SimpleGraphs.delete!(G1,v,w)
        p1 = matching_poly(G1,false)

        SimpleGraphs.delete!(G1,v)
        SimpleGraphs.delete!(G1,w)
        p2 = matching_poly(G1,false)

        p = p1 - p2
        if cache_flag
            cache_save(G,:matching_poly,p)
        end
        return p
    end

    comps = parts(components(G))
    p = Poly([1])
    for S in comps
        H = induce(G,S)
        pH = matching_poly(H,false)
        p *= pH
    end
    if cache_flag
        cache_save(G,:matching_poly,p)
    end
    return p
end


# Arrays

glorot_uniform(dims...) = glorot_uniform(STD_REAL_PREC, dims...)
glorot_normal(dims...) = glorot_normal(STD_REAL_PREC, dims...)
rescaled_normal(scale::Real, dims...) = rescaled_normal(STD_REAL_PREC, scale, dims...)

glorot_uniform(T::Type, dims...) = (rand(T, dims...) .- T(0.5)) .* sqrt(T(24)/sum(dims))
glorot_normal(T::Type, dims...) = randn(T, dims...) .* sqrt(T(2)/sum(dims))
rescaled_normal(T::Type, scale::Real, dims::Integer...) = randn(T, dims...) .* T(scale) .* sqrt(T(24)/sum(dims))

# Utils

@inline ℒ2(x) = 2one(x)*cosh(x)
@inline ∂logℒ2(x) = tanh(x)
@inline logℒ2(x::T) where T<:Real =
    log(2one(x))+ (abs(x)<= T(12.0) ? log(cosh(x)) : abs(x) - log(2one(X)))
@inline logℒ2(x::Complex{T}) where T<:Real =
    log(T(2.0))+ (abs(x)<= T(12.0) ? log(cosh(x)) : abs(x) - log(2one(x)))

@inline ℒ(x) = one(x) + exp(x)
#@inline ∂logℒ(x) = one(x)/(one(x)+exp(-x))
const ∂logℒ = NNlib.σ
NNlib.σ(x::Complex) = one(x)/(one(x)+exp(-x))
#@inline logℒ(x) = log1p(exp(x))#log(one(x) + exp(x))
const logℒ = NNlib.softplus
NNlib.softplus(x::Complex) = log1p(exp(x))#log(one(x) + exp(x))


using Printf

const DEFAULT_ARMIJO_α₀ = 1.0
const DEFAULT_ARMIJO_σ₀ = 0.1
const DEFAULT_ARMIJO_σ₁ = 0.5
const DEFAULT_ARMIJO_p  = 0.5

struct BacktrackingState{OBJ,OPT,T} <: LinesearchState where {OBJ <: AbstractObjective, OPT <: Options, T <: Number}
    objective::OBJ
    config::OPT

    α₀::T
    ϵ::T
    p::T

    function BacktrackingState(objective; config = Options(),
                    α₀::T = DEFAULT_ARMIJO_α₀,
                    ϵ::T = DEFAULT_WOLFE_ϵ,
                    p::T = DEFAULT_ARMIJO_p) where {T}
        new{typeof(objective), typeof(config), T}(objective, config, α₀, ϵ, p)
    end
end

function BacktrackingState(F::Callable, x::Number = DEFAULT_ARMIJO_α₀; D = nothing, kwargs...)
    objective = UnivariateObjective(F, D, x)
    BacktrackingState(objective; kwargs...)
end

function BacktrackingState(F::Callable, x::AbstractVector; D = nothing, kwargs...)
    objective = MultivariateObjective(F, D, x)
    BacktrackingState(objective; kwargs...)
end

Base.show(io::IO, ls::BacktrackingState) = print(io, "Backtracking")

LinesearchState(algorithm::Backtracking, objective; kwargs...) = BacktrackingState(objective; kwargs...)

function (ls::BacktrackingState{<:UnivariateObjective})()
    local α = ls.α₀
    local y₀ = value!(ls.objective, zero(α))
    local d₀ = derivative!(ls.objective, zero(α))

    for _ in 1:ls.config.max_iterations
        if value!(ls.objective, α) ≥ y₀ + ls.ϵ * α * d₀
            α *= ls.p
        else
            break
        end
    end

    return α
end

backtracking(f, x; kwargs...) = BacktrackingState(f, x; kwargs...)()


using LinearAlgebra, Krylov, LinearOperators, SparseArrays, Stopping

#Krylov @kdot
macro kdot(n, x, y)
  return esc(:(Krylov.krylov_dot($n, $x, 1, $y, 1)))
end

using SolverTools, Logging
"""
Randomized coordinate descent

Sect. 3.7 in Gower, R. M., & Richtárik, P. (2015).
Randomized iterative methods for linear systems.
SIAM Journal on Matrix Analysis and Applications, 36(4), 1660-1690.
"""
function RandomizedCD(
  A::AbstractMatrix,
  b::AbstractVector{T};
  is_zero_start::Bool = true,
  x0::AbstractVector{T} = zeros(T, size(A, 2)),
  atol::AbstractFloat = 1e-7,
  rtol::AbstractFloat = 1e-15,
  max_iter::Int = size(A, 2)^2,
  max_time::Float64 = 60.0,
  max_cntrs = Stopping._init_max_counters_linear_operators(quick = 20000),
  verbose::Int = 100,
  kwargs...,
) where {T <: AbstractFloat}
  m, n = size(A)
  x = copy(x0)
  res = is_zero_start ? b : b - A * x
  nrm0 = norm(res)

  time_init = time()
  elapsed_time = time_init

  cntrs = LACounters()
  max_f = false

  OK = nrm0 <= atol
  k = 0

  #@info log_header([:iter, :un, :time], [Int, T, T])
  #@info log_row(Any[0, res[1], elapsed_time])
  while !OK && (k <= max_iter) && (elapsed_time - time_init <= max_time) && !max_f

    #rand a number between 1 and n
    #224.662 ns (4 allocations: 79 bytes) - independent of the n
    i = mod(k, n) + 1#Int(floor(rand() * n) + 1)
    Ai = A[:, i]

    #ei = zeros(n); ei[i] = 1.0 #unit vector in R^n
    #xk  = Ai == 0 ? x0 : x0 - dot(Ai,res)/norm(Ai,2)^2 * ei
    Aires = @kdot(m, Ai, res)
    nAi = @kdot(m, Ai, Ai)
    x[i] -= Aires / nAi

    #res = b - A*x
    res += Ai * Aires / nAi
    cntrs.nprod += 1
    nrm = norm(res, Inf)
    OK = nrm <= atol + nrm0 * rtol

    sum, max_f = 0, false
    for f in [:nprod, :ntprod, :nctprod]
      ff = getfield(cntrs, f)
      max_f = max_f || (ff > max_cntrs[f])
      sum += ff
    end
    max_f = max_f || sum > max_cntrs[:neval_sum]
    k += 1
    elapsed_time = time()
    if mod(k, verbose) == 0 #print every 20 iterations
      # @info log_row(Any[k, res[1], elapsed_time])
    end
  end

  return x, OK, k
end


function get_preaveraged_prices(ts::SortedDataFrame, asset::Symbol, k_n::Real, gs::Vector)
   prices = Array(ts.df[ts.groupingrows[asset],ts.value])
   times = Array(ts.df[ts.groupingrows[asset],ts.time])
   durations = Array(times[2:end] .- times[1:(end-1)])
   times = times[2:end]
   diffs = simple_differencing(Array(prices[2:end]), Array(prices[1:(end-1)]))
   lastrow = length(diffs) - k_n + 1
   if lastrow < 1 return missing, missing end
   preaveraged_returns = Array{eltype(diffs),1}(undef, lastrow)
   for i in 1:lastrow
       preaveraged_returns[i] = sum(diffs[i:(i+k_n-2)] .* gs)
   end
   return preaveraged_returns, times
end

function HY_n(A::Tuple{Vector{<:Real},Vector{<:Real}}, B::Tuple{Vector{<:Real},Vector{<:Real}}, k_n::Integer, psi::Real)
   # See top of page 62 of Christensen et al.
   cov_sum = 0.0
   for i in 1:(length(A[1]) - k_n + 2)
      Y_i = A[1][i]
      i_start = A[2][i]
      i_end = A[2][i+k_n]
      j_index_i_start = max(1, searchsortedfirst(B[2], i_start) - k_n) # Because an observation k_n ticks ago for the other asset also overlaps.
      j_index_i_end   = min(length(B[1]),searchsortedlast(B[2], i_end))
      for j in j_index_i_start:j_index_i_end
          cov_sum += Y_i * B[1][j]
      end
   end
   return cov_sum / ((psi * k_n)^2)
end

# This is the 3rd equation in the aymptotic theory section of the paper.
univariate_HYn(vect::Vector, k_n::Real, psi::Real) = sum(map(i -> vect[i] * sum( map(j -> vect[i-j], (1-k_n):(k_n-1)  ) ), k_n:(length(vect)-2k_n+1)))/((psi*k_n)^2)

g = (f = x-> min(x, 1-x), psi = 0.25)

"""
    preaveraged_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,
                           theta::Real = 0.15, g::NamedTuple = g)

Estimation of the CovarianceMatrix using preaveraging 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.
* `theta` - A theta value. See paper for details.
* `g` - A tuple containing a preaveraging method (with name "f") and a ψ value. See paper for details.
### Returns
* A `CovarianceMatrix`.

### References
Christensen K, Podolskij M, Vetter M (2013). “On covariation estimation for multivariate continuous Itô semimartingales with noise in non-synchronous observation schemes.” Journal of Multivariate Analysis, 120, 59–84. doi:10.1016/j.jmva.2013.05.002.
"""
function preaveraged_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,
                                theta::Real = 0.15, g::NamedTuple = g)
   # The defaults are from the paper (Christensen et al 2013). theta and the formula for k_n is from halfway down page 67. g is from page 64.
   number_of_ticks = nrow(ts.df)
   k_n = Int(ceil(min(number_of_ticks/2,theta * sqrt(number_of_ticks))))
   gs = g.f.( collect(1:1:(k_n-1)) ./ k_n )
   prev_prices = get_preaveraged_prices.(Ref(ts), assets, k_n, Ref(gs))

  lens = findall(map(i -> ismissing(prev_prices[i][1]), 1:length(prev_prices)))
  if length(lens) > 0
     @warn string("Cannot estimate the correlation matrix with ", number_of_ticks, " ticks. There are insufficient ticks for ", assets[lens])
     return make_nan_covariance_matrix(assets, ts.time_period_per_unit)
  end

   N = length(assets)
   # First doing HY_n from equation 2.5
   HYn = zeros(N,N)
   for i in 1:N
      A = prev_prices[i]
      for j in i:N
         if i ==j
            HYn[i,i] = univariate_HYn(A[1], k_n, g.psi)
         else
            B = prev_prices[j]
            HYn[i,j] = HY_n(A, B, k_n, g.psi)
         end
      end
   end
   HYn = Hermitian(HYn)
   # Regularisation
   dont_regulise = ismissing(regularisation) || (only_regulise_if_not_PSD && is_psd_matrix(HYn))
   HYn = dont_regulise ? HYn : regularise(HYn, ts, assets, regularisation; regularisation_params... )

   # In some cases we get negative terms on the diagonal with this algorithm.
   negative_diagonals = findall(diag(HYn) .< eps())
   covar = make_nan_covariance_matrix(assets, ts.time_period_per_unit)
   if length(negative_diagonals) == 0
       corr, _ = cov2cor_and_vol(HYn, ts.time_period_per_unit, ts.time_period_per_unit)
       covar.correlation = corr
   end

   # We can use this to get the correlation matrix but the variances are too low - as a result of preveraging.
   # We will instead use the two scales vol of Zhang, mykland, Ait-Sahalia 2005.
   voldict = two_scales_volatility(ts, assets)[1]
   vols = map(a -> voldict[a], assets)
   covar.volatility = vols
   return covar
end


using Gadfly

plot(sin, 0, 25, Coord.cartesian(aspect_ratio=1.0))


#module SSEquilibrium

#export ModelParams, utility, TauchenApprox, HHBellmanMap, SolveHHBellman, UpdateWage, SolveWage

using Parameters
using LinearAlgebra
using Distributions
using Optim
using QuantEcon
using SpecialFunctions
using SparseArrays
using ThreadTools

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Initialize parameters & asset grids

@doc """
    ModelParams

A structure containing all model parameters, including asset and matching productivity shock grids.
"""
@with_kw mutable struct ModelParams
    A::Float64              = 1.0                           # output productivity factor
    r::Float64              = 0.06                          # annual interest rate
    r_mon::Float64          = (1.0 + r)^(1.0/12.0) - 1.0    # monthly interest rate
    worker_surplus::Float64 = 0.5                           # worker's share of surplus
    β::Float64              = 0.9948                        # discount factor
    β_mon::Float64          = 1.0/(1.0 + r_mon)             # monthly depreciation rate
    γ::Float64              = 1.0                           # curvature of utility function
    λ::Float64              = 0.02                          # exogenous separation rate
    b::Float64              = 0.4                           # UI benefits
    B::Float64              = 0.15                          # utility from leisure
    κ::Float64              = 0.522                         # vacancy posting cost with log utility
    θ::Float64              = 1.0                           # labor market tightness, v/u, normalized to 1
    α::Float64              = 0.5                           # power term in the matching function
    η::Float64              = 0.3133                        # scale parameter in the matching function
    p_θ::Float64            = η*θ^α                         # unemployed worker matching rate
    q_θ::Float64            = η*θ^(α-1)                     # vacancy matching probability
    amin::Float64           = -6.0                          # minimum asset
    amax::Float64           = 120.0                         # maximum asset
    agrid::Vector{Float64}  = LinRange(amin, amax, 50)      # asset grid w/ 50 nodes
    ρ_x::Float64            = 0.97                          # persistence of idiosyncratic productivity ln(x)
    σ_x::Float64            = 0.13                          # std. dev. of innovation to ln(x)
    π_x                     = TauchenApprox(9, ρ_x, σ_x).distribution           # probability weights on ln(x)
    xgrid                   = exp.(TauchenApprox(9, ρ_x, σ_x).state_values)     # match productivity grid
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    utility(para::ModelParams, c::Real)

Household Bernoulli utility function.
"""
function utility(para::ModelParams, c::Real)

    γ = para.γ # Store risk-aversion parameter

    # Define the utility mapping
    # condition on the input `c`
    if c > 0
        if γ == 1
          return log(c) # if σ = 1, CRRA utility is log
        else
          return (c^(1-γ)-1)/(1-γ) # if σ != 1, regular specification
        end
    else
        return -Inf # infinite disutility if c∈(-∞,0]
    end
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    utility(para::ModelParams, cvec)

Household Bernoulli utility function applied to a vector.
"""
function utility(para::ModelParams, cvec)

  u(c) = utility(para,c)

  return u.(cvec)
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    TauchenApprox(N, ρ, σ, μ, n_std)

Tauchen's (1996) method for approximating AR(1) process with finite markov chain.

**Output:**
⋅ transition    -- Estimated transition matrix
⋅ distribution  -- Estimated stationary distribution (vector)
⋅ state_values  -- Etimated state values (vector)
"""
function TauchenApprox(N::Integer, ρ::T1, σ::T2, μ=zero(promote_type(T1, T2)), n_std::Integer=3) where {T1 <: Real, T2 <: Real}

    std_norm_cdf(x::T) where {T <: Real} = 0.5 * SpecialFunctions.erfc(-x/sqrt(2))
    std_norm_cdf(x::Array{T}) where {T <: Real} = 0.5 .* SpecialFunctions.erfc(-x./sqrt(2))

    # Get discretized space
    a_bar = n_std * sqrt(σ^2 / (1 - ρ^2))
    y = range(-a_bar, stop=a_bar, length=N)
    d = y[2] - y[1]

    # Get transition probabilities
    Π = zeros(promote_type(T1, T2), N, N)
    Threads.@threads for row = 1:N
        # Do end points first
        Π[row, 1] = std_norm_cdf((y[1] - ρ*y[row] + d/2) / σ)
        Π[row, N] = 1 - std_norm_cdf((y[N] - ρ*y[row] - d/2) / σ)

        # fill in the middle columns
        for col = 2:N-1
            Π[row, col] = (std_norm_cdf((y[col] - ρ*y[row] + d/2) / σ) -
                           std_norm_cdf((y[col] - ρ*y[row] - d/2) / σ))
        end
    end

    # state values
    ȳ = y .+ μ / (1 - ρ) # center process around its mean (wbar / (1 - rho)) in new variable

    # renormalize transition matrix
    Π = Π./sum(Π, dims = 2)

    # Generate stationary distribution from Π
    N = size(Π)[1]
    π0 = ones(1,N)/N
    diff = 1.
    while diff > 1e-10
        π1 = π0*Π
        diff = norm(π1-π0,Inf)
        π0 = π1
    end
    π0 = copy(π0')

    return (transition = Π, distribution = π0, state_values = ȳ)
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    HHBellmanMap(para::ModelParams, wage, W_old, U_old)

Iterates on the bellman equation using continuation value function.

**Output:**
⋅ W_new         -- Updated value of employment (matrix)
⋅ U_new         -- Updated value of unemployment (vector)
⋅ emp_policy    -- Asset policy function in case of employment (matrix)
⋅ unemp_policy  -- Asset policy function in case of unemployment (vector)
"""
function HHBellmanMap(para::ModelParams, wage, W_old, U_old)

    @unpack agrid, xgrid, π_x, β, b, B, p_θ, r   = para

    u(c) = utility(para,c) #shorthand for utility

    W_new           = similar(W_old)
    U_new           = similar(U_old)
    emp_policy      = similar(W_new, Int)
    unemp_policy    = similar(U_new, Int)

    N_x = length(xgrid)
    N_a = length(agrid)

    Threads.@threads for x_i  in 1:N_x

        if x_i == 5
            for a_i in 1:N_a
                cvec_emp = (1+r)*agrid[a_i] + wage[a_i,x_i] .- agrid
                obj_emp     = u(cvec_emp) .+ β * sum( π_x[i]*max.(W_old[:,i], U_old) for i in 1:N_x)
                obj_emp     = vec(obj_emp)
                W_new[a_i,x_i], emp_policy[a_i,x_i] = findmax(obj_emp)

                cvec_unemp = (1+r)*agrid[a_i] + b .- agrid
                obj_unemp   = u(cvec_unemp) .+ B .+ β*(1-p_θ)*U_old .+ β*(p_θ)*W_old[:,5]
                obj_unemp   = vec(obj_unemp)
                U_new[a_i], unemp_policy[a_i] = findmax(obj_unemp)
            end
        else
            for a_i in 1:N_a
                cvec_emp = (1+r)*agrid[a_i] + wage[a_i,x_i] .- agrid
                obj_emp     = u(cvec_emp) .+ β * sum( π_x[i]*max.(W_old[:,i], U_old) for i in 1:N_x)
                obj_emp     = vec(obj_emp)
                W_new[a_i,x_i], emp_policy[a_i,x_i] = findmax(obj_emp)
            end
        end
    end

    return W_new, U_new, emp_policy, unemp_policy
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    SolveHHBellman(para::ModelParams, wage, W0, U0, ϵ=1e-6)

Iterates on the Bellman map until convergence.

**Output:**
⋅ W             -- Approximated value of employment (matrix)
⋅ U             -- Approximated value of unemployment (vector)
⋅ emp_policy    -- Approximated asset policy function in case of employment (matrix)
⋅ unemp_policy  -- Approximated asset policy function in case of unemployment (vector)
"""
function SolveHHBellman(para::ModelParams, wage, W0, U0, ϵ=1e-5)

    W_old = W0
    U_old = U0
    emp_policy      = similar(W_old, Int)
    unemp_policy    = similar(U_old, Int)

    diff = 1.

    while diff > ϵ
        W_new, U_new = HHBellmanMap(para, wage, W_old, U_old)
        diff = norm((U_new-U_old)[:],Inf)
        W_old = W_new
        U_old = U_new
    end

    W, U, emp_policy, unemp_policy = HHBellmanMap(para, wage, W_old, U_old)

    return W, U, emp_policy, unemp_policy
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    UpdateWage(para::ModelParams, W, U, emp_policy, wage_old)

Applies step #3(c) of the Bils, Chang, and Kim (2011) steady state equilibrium algorithm to update the wage, w¹(a,x;θ⁰), object.

**Output:**
⋅ wage_new  -- New iteration of the wage mapping (matrix)
⋅ J         -- New iteration of the value of a matched job (matrix)
"""
function UpdateWage(para::ModelParams, W, U, emp_policy, wage_old)

    @unpack agrid, xgrid, β, λ, α, r, π_x = para

    N_x = length(xgrid)
    N_a = length(agrid)

    J = zeros(N_a, N_x)
    wage_new = similar(wage_old)

    Threads.@threads for x_i in 1:N_x
        for a_i in 1:N_a
            c_e = (1+r)*agrid[a_i] + wage_old[a_i,x_i] - emp_policy[a_i,x_i]
            J[a_i,x_i] = ((1-α)/α)*(W[a_i,x_i] - U[a_i])*c_e
        end
    end

    Threads.@threads for x_i in 1:N_x
        for a_i in 1:N_a
            wage_new[a_i,x_i] = 1*xgrid[x_i] - J[a_i,x_i] + β*(1-λ)*sum( π_x[i]*max.( J[emp_policy[a_i,x_i],i] , 0.0 ) for i in 1:N_x)
        end
    end

    return wage_new, J
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Initialize value functions and wage
@doc """
    InitializeFunctions(para::ModelParams)

Initializes W(a,x), U(a), and wage(a,x).
"""
function InitializeFunctions(para::ModelParams)
    @unpack agrid, xgrid, β, r, b, B = para
    N_a = length(agrid)
    N_x = length(xgrid)
    W       = ones(N_a,N_x)
    U       = ones(N_a)
    wage    = ones(N_a,N_x)

    Threads.@threads for x_i  in 1:N_x

        if x_i == 5
            for a_i in 1:N_a
                wage[a_i,x_i] = xgrid[x_i]*2
                W[a_i,x_i] = log((1+r)*agrid[a_i] + wage[a_i,x_i] - agrid[a_i])/(1-β)
                U[a_i]     = log((1+r)*agrid[a_i] + b - agrid[a_i] + B)/(1-β)
            end
        else
            for a_i in 1:N_a
                wage[a_i,x_i] = xgrid[x_i]*2
                W[a_i,x_i] = log((1+r)*agrid[a_i] + wage[a_i,x_i] - agrid[a_i])/(1-β)
            end
        end
    end

    return W, U, wage
end

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Solve wage
@doc """
    SolveWage(para::ModelParams, ϵ=1e-6)

Applies steps #1-3 of the Bils, Chang, and Kim (2011) steady state equilibrium algorithm to obtain value functions, policy functions, and the wage mapping.

**Output:**
⋅ W             -- Approximated value of employment (matrix)
⋅ U             -- Approximated value of unemployment (vector)
⋅ emp_policy    -- Approximated asset policy function in case of employment (matrix)
⋅ unemp_policy  -- Approximated asset policy function in case of unemployment (vector)
⋅ wage          -- Approximated wage mapping (matrix)
"""
function SolveWage(para::ModelParams, ϵ=1e-5)

    W_old, U_old, wage_old = InitializeFunctions(para)

    """
    W_old       = ones(length(para.agrid),length(para.xgrid))
    U_old       = ones(length(para.agrid))
    c           = 0.9
    wage_old    = ones(length(para.agrid),length(para.xgrid))*(c*θ)
    """

    ζ_x         = 0.01
    difff       = 10
    counter     = 0

    while difff > ϵ
        W_new, U_new, emp_policy = SolveHHBellman(para, wage_old, W_old, U_old)
        wage_new, J              = UpdateWage(para, W_new, U_new, emp_policy, wage_old)
        difff                    = norm(wage_new - wage_old)
        wage_old                 = ζ_x*wage_new + (1-ζ_x)*wage_old
        W_old                    = W_new
        U_old                    = U_new
        counter = counter + 1
        println("Wage Solver Iteration: $(counter). Norm: $(difff).")
    end

    W, U, emp_policy, unemp_policy              = SolveHHBellman(para, wage_old, W_old, U_old)
    wage, J                                     = UpdateWage(para, W, U, emp_policy, wage_old)

    return W, U, J, emp_policy, unemp_policy, wage
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Reservation productivity, x(a;w)
@doc """
    ReservationProductivity(para::ModelParams, W, U)
"""
function ReservationProductivity(para::ModelParams, W, U)

        @unpack xgrid, agrid = para
        N_a = length(agrid)
        N_x = length(xgrid)
        x_comps = zeros(N_a,N_x)
        # If for any given (a,x), U(a) > W(a,x)
        # then store x_comps(a,x) = 1.0 (unemployment)
        # other store x_comps(a,x) = 0.0 (employment)
        for x_i in 1:N_x
            for a_i in 1:N_a
                x_comps[a_i,x_i] = -1*W[a_i,x_i] > -1*U[a_i]
            end
        end
        # x_grid = position in xgrid for which W surpasses U
        # indicating entry of xgrid that allows the switch
        # from unemployment to employment for a given `a`
        x_opt   = zeros(N_a)
        for a_i in 1:N_a
            x_index     = findfirst(x_comps[a_i,:] .== 0.0)
            if x_index > 1
                x_opt[a_i]  = (xgrid[x_index] - xgrid[x_index-1])/2
            else
                x_opt[a_i]  = 0.0
            end
        end

    return x_opt
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    InitializeTransitionMatrices(para::ModelParams, emp_policy, unemp_policy, x_star)
"""
function InitializeTransitionMatrices(para::ModelParams, emp_policy, unemp_policy, x_star)

    @unpack agrid, xgrid, π_x = para
    N_x  = length(xgrid)
    N_a  = length(agrid)
    H_ee = zeros(N_a*N_x,N_a*N_x) # trans. from employmed to employmed
    H_eu = zeros(N_a*N_x,N_a*N_x) # trans. from employed to unemployed
    H_ue = zeros(N_a*N_x,N_a*N_x) # trans. from unemployed to employed
    H_uu = zeros(N_a*N_x,N_a*N_x) # trans. from unemployed to unemployed

    ### Fill out the employment matrix
    Threads.@threads for a_i in 1:N_a
        for x_i in 1:N_x
            j       = a_i+N_a*(x_i-1) # index for (a,x)
            a_emp   = emp_policy[a_i,x_i] # index for a′ given employment
            a_unemp = unemp_policy[a_i] # index for a′ given unemployment
            x′_star = x_star[a_emp] # value of x′_star under employment
            for x_i′ in 1:N_x
                j′ = a_emp+N_a*(x_i′-1) # index for (a′,x′) given employment
                if xgrid[x_i′] > x′_star # if x′ is above reservation x
                    H_ee[j,j′] = π_x[x_i′] # prob. of (a,x) → (a′,x′) given employment
                else
                    H_eu[j,j′] = π_x[x_i′] # prob. of (a,x) → (a′,x′) given employment
                end
                j′ = a_unemp+N_a*(x_i′-1) # index for (a′,x′) given unemployment
                H_uu[j,j′] = π_x[x_i′]  # prob. of (a,x) → (a′,x′) given unemployment
                if x_i′ == 5 # if x′ = x̄ (x̄ is the middle (5th) entry of xgrid)
                    H_ue[j,j′] = 1 # prob. of (a,x) → (a′,x′) given unemployment
                end
            end
        end
    end

    return H_ee, H_eu, H_ue, H_uu
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    FindStationaryMeasures(para::ModelParams, emp_policy, unemp_policy)
"""
function FindStationaryMeasures(para::ModelParams,  emp_policy, unemp_policy, x_star)

    @unpack p_θ = para

    H_ee, H_eu, H_ue, H_uu = InitializeTransitionMatrices(para, emp_policy, unemp_policy, x_star)
    N_emp = size(H_ee)[1]
    π0_emp = ones(1,N_emp)/N_emp
    N_unemp = size(H_uu)[1]
    π0_unemp = ones(1,N_unemp)/N_unemp
    diff_emp = 1.
    diff_unemp = 1.
    while diff_emp > 1e-10 && diff_unemp > 1e-10
        π1_emp = π0_emp*H_ee + p_θ*π0_unemp*H_ue
        π1_unemp = π0_emp*H_eu + (1-p_θ)*π0_unemp*H_uu
        π1_emp = π1_emp/sum(π1_emp)
        π1_unemp = π1_unemp/sum(π1_unemp)
        diff_emp = norm(π1_emp-π0_emp,Inf)
        diff_unemp = norm(π1_unemp-π0_unemp,Inf)
        π0_emp = π1_emp
        π0_unemp = π1_unemp
    end

    π_emp   = π0_emp
    π_unemp = π0_unemp

    return π_emp, π_unemp
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    AdjustTheta(para::ModelParams, J)
"""
function AdjustTheta(para::ModelParams, J, π_unemp)

    @unpack κ, β, η, α, agrid = para

    N_a = length(agrid)

    # integral of J over ψ is the dot product of
    # J(:,x̄) and π_unemp(N_a*x̄-1:N_a*x̄)
    int_J = dot(J[:,5],π_unemp[(N_a*(4)+1):N_a*(5)])
    θ_new = ( κ/( β*η*int_J ) )^(1/(α-1))

    return θ_new
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    SolveEquilibrium(para::ModelParams)
"""
function SolveEquilibrium(para::ModelParams)

    para = para
    @unpack θ, agrid, xgrid = para

    N_a = length(agrid)
    N_x = length(xgrid)

    diff    = 1
    θ_old   = θ
    ζ_θ     = 0.05
    counterr = 0

    # Declare variables to be returned
    # at the end of the function
    W               = zeros(N_a,N_x)
    U               = zeros(N_a)
    J               = zeros(N_a,N_x)
    emp_policy      = zeros(N_a,N_x)
    unemp_policy    = zeros(N_a)
    wage            = zeros(N_a,N_x)
    x_star          = zeros(N_a)
    π_emp           = zeros(N_a*N_x)
    π_unemp         = zeros(N_a*N_x)

    while diff > 1e-4
    para = ModelParams(θ=θ_old)
    W, U, J, emp_policy, unemp_policy, wage = SolveWage(para)
    x_star = ReservationProductivity(para, W, U)
    π_emp, π_unemp = FindStationaryMeasures(para, emp_policy, unemp_policy, x_star)
    θ_new = AdjustTheta(para, J, π_unemp)
    diff = abs(θ_new - θ_old)
    θ_old = ζ_θ*θ_new + (1-ζ_θ)*θ_old
    counterr = counterr + 1
    println("Equilibrium Iteration: $(counterr). Norm: $(diff).")
    end

    θ = θ_old

    return para, θ, W, U, J, emp_policy, unemp_policy, wage, x_star, π_emp, π_unemp
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#end


"""
    modeconversion(previous_wavefields::Wavefields, wavefields::Wavefields,
        adjoint_wavefields::Wavefields; params=LMPParams())

Compute the mode conversion matrix `a` from the modes associated with `previous_wavefields`
to modes associated with `wavefields` and its `adjoint_wavefields`.

This is used in the approach known as full mode conversion [Pappert1972b].

# References

[Pappert1972b]: R. A. Pappert and R. R. Smith, “Orthogonality of VLF height gains in the
    earth ionosphere waveguide,” Radio Science, vol. 7, no. 2, pp. 275–278, 1972,
    doi: 10.1029/RS007i002p00275.
"""
function modeconversion(previous_wavefields::Wavefields{H}, wavefields::Wavefields{H},
                        adjoint_wavefields::Wavefields{H}; params=LMPParams()) where H

    @unpack wavefieldheights = params

    product = Vector{ComplexF64}(undef, numheights(wavefields))
    pproduct = similar(product)

    modes = eigenangles(wavefields)
    adjmodes = eigenangles(adjoint_wavefields)
    prevmodes = eigenangles(previous_wavefields)

    modes == adjmodes || @warn "Full mode conversion physics assumes adjoint wavefields are
        calculated for eigenangles of the original waveguide."

    nmodes = length(modes)
    nprevmodes = length(prevmodes)

    # `a` is total conversion
    a = Matrix{ComplexF64}(undef, nprevmodes, nmodes)
    for n in eachindex(modes)  # modes == adjmodes
        for m in eachindex(prevmodes)
            for i in eachindex(wavefieldheights)
                @inbounds f = wavefields[i,n][SVector(2,3,5,6)]  # Ey, Ez, Hy, Hz
                @inbounds fp = previous_wavefields[i,m][SVector(2,3,5,6)]  # Ey, Ez, Hy, Hz
                @inbounds g = SVector{4}(adjoint_wavefields[i,n][6],
                                         -adjoint_wavefields[i,n][5],
                                         -adjoint_wavefields[i,n][3],
                                         adjoint_wavefields[i,n][2])  # Hz, -Hy, -Ez, Ey

                gtranspose = transpose(g)
                product[i] = gtranspose*f
                pproduct[i] = gtranspose*fp
            end
            N, Nerr = romberg(heights(wavefields), product)  # normalization
            I, Ierr = romberg(heights(wavefields), pproduct)
            @inbounds a[m,n] = I/N
        end
    end

    return a
end


propositions = [T, F]
p = Proposition(rand(propositions))
q = Proposition(rand(propositions))
r = Proposition(rand(propositions))

## Axiom 1.1 [Commutativity] ##

@test (p * q ↔ q * p) == T
@test (p + q ↔ q + p) == T
@test ((p ↔ q) ↔ (q ↔ p)) == T

## Axiom 1.2 [Associativity] ##

@test (p * (q * r) ↔ (p * q) * r) == T
@test (p + (q + r) ↔ (p + q) + r) == T

## Axiom 1.3 [Distributivity] ##

@test (p + (q * r) ↔ (p + q) * (p + r)) == T
@test (p * (q + r) ↔ (p * q) + (p * r)) == T

## Axiom 1.4 [De Morgan]

@test (!(p * q) ↔ !p + !q) == T
@test (!(p + q) ↔ !p * !q) == T

## Axiom 1.5 [Negation] ##

@test (!!p ↔ p) == T

## Axiom 1.6 [Excluded Middle] ##

@test (p + !p ↔ T) == T

## Axiom 1.7 [Contradiction] ##

@test (p * !p ↔ F) == T

## Axiom 1.8 [Implication] ##

@test (p ⟹ q ↔ !p + q) == T

## Axiom 1.9 [Equality] ##

@test ((p ↔ q) ↔ (p ⟹ q) * (q ⟹ p)) == T

## Axiom 1.10 [or-simplification] ##

@test (p + p ↔ p) == T
@test (p + T ↔ T) == T
@test (p + F ↔ p) == T
@test (p + (p * q) ↔ p) == T

## Axiom 1.11 [and-simplification] ##

@test (p * p ↔ p) == T
@test (p * T ↔ p) == T
@test (p * F ↔ F) == T
@test (p * (p + q) ↔ p) == T

## Axiom 1.12 [Identity] ##

@test (p ↔ p) == T


## Create an extractor that extracts the modified original matrix, e.g.
## LD for BunchKaufman, UL for CholeskyDense, LU for LUDense and
## define size methods for Factorization types using it.

immutable BunchKaufman{T} <: Factorization{T}
    LD::Matrix{T}
    ipiv::Vector{BlasInt}
    uplo::Char
    symmetric::Bool
end

function bkfact!{T<:BlasReal}(A::StridedMatrix{T}, uplo::Symbol=:U, symmetric::Bool=issym(A))
    symmetric || error("The Bunch-Kaufman decomposition is only valid for symmetric matrices")
    LD, ipiv = LAPACK.sytrf!(string(uplo)[1] , A)
    BunchKaufman(LD, ipiv, string(uplo)[1], symmetric)
end
function bkfact!{T<:BlasComplex}(A::StridedMatrix{T}, uplo::Symbol=:U, symmetric::Bool=issym(A))
    LD, ipiv = (symmetric ? LAPACK.sytrf! : LAPACK.hetrf!)(string(uplo)[1] , A)
    BunchKaufman(LD, ipiv, string(uplo)[1], symmetric)
end
bkfact{T<:BlasFloat}(A::StridedMatrix{T}, uplo::Symbol=:U, symmetric::Bool=issym(A)) = bkfact!(copy(A), uplo, symmetric)
bkfact{T}(A::StridedMatrix{T}, uplo::Symbol=:U, symmetric::Bool=issym(A)) = bkfact!(convert(Matrix{promote_type(Float32,typeof(sqrt(one(T))))},A),uplo,symmetric)

convert{T}(::Type{BunchKaufman{T}},B::BunchKaufman) = BunchKaufman(convert(Matrix{T},B.LD),B.ipiv,B.uplo,B.symmetric)
convert{T}(::Type{Factorization{T}}, B::BunchKaufman) = convert(BunchKaufman{T}, B)

size(B::BunchKaufman) = size(B.LD)
size(B::BunchKaufman,d::Integer) = size(B.LD,d)
issym(B::BunchKaufman) = B.symmetric
ishermitian(B::BunchKaufman) = !B.symmetric

inv{T<:BlasReal}(B::BunchKaufman{T})=copytri!(LAPACK.sytri!(B.uplo, copy(B.LD), B.ipiv), B.uplo, true)

function inv{T<:BlasComplex}(B::BunchKaufman{T})
    if issym(B)
        copytri!(LAPACK.sytri!(B.uplo, copy(B.LD), B.ipiv), B.uplo)
    else
        copytri!(LAPACK.hetri!(B.uplo, copy(B.LD), B.ipiv), B.uplo, true)
    end
end

A_ldiv_B!{T<:BlasReal}(B::BunchKaufman{T}, R::StridedVecOrMat{T}) = LAPACK.sytrs!(B.uplo, B.LD, B.ipiv, R)
function A_ldiv_B!{T<:BlasComplex}(B::BunchKaufman{T}, R::StridedVecOrMat{T})
    (issym(B) ? LAPACK.sytrs! : LAPACK.hetrs!)(B.uplo, B.LD, B.ipiv, R)
end


using LidJul
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using Makie
using IterativeSolvers
using Preconditioners
using Random

function addmeasurements(msm,solvername,ti,ts,res,iterations=[])
    push!(msm,(solver=solvername,init_time=ti,solver_time=ts,residual=res,iterations=iterations))
end

function resnorm(splu,pref,b)
    ax=similar(b)
    LidJul.mul!(ax,splu,pref)
    norm(ax-b)
end

function test_solver!(msm,Solver,solver_args,splu,b)
    p=zeros(size(b))
    init_time=@elapsed solver=Solver(solver_args...)
    init_time=@elapsed solver=Solver(solver_args...)
    Random.seed!(1234)
    prand=rand(length(p))
    iterations=begin copyto!(p,prand) ; solve!(p,b,solver)  end
    solve_time=@elapsed  begin copyto!(p,prand) ; solve!(p,b,solver) end
    sname=string(Solver)
    length(solver_args)>1 && (sname=sname*"_"*string(solver_args[2]))

    addmeasurements(msm,sname,init_time,solve_time,resnorm(splu,p,b),iterations)
    p
end

function precoconstructor(A,preco)
    if preco==CholeskyPreconditioner
        return preco(A,4)
    elseif preco===nothing
        return IterativeSolvers.Identity()
    else
        return preco(A)
    end
end

include("measureplots.jl")

function testpoisson_GMG(n,bc)
    msm=Vector{NamedTuple{(:solver,:init_time,:solver_time,:residual,:iterations),Tuple{String,Float64,Float64,Float64,Any}}}()

    p=ones(n,n)
    pref=ones(n,n)
    b=zeros(n,n)
    #LP: Create a Poisson's RHS compatible with a possible full Neumann boundary conditions
    #See bellow
    for i=1:n
        for j=1:n
            b[i,j]=sin(3π*(i-1)/(n-1))*sin(6π*(j-1)/(n-1))
        end
    end

    L=1.0
    dx=L/n
    dy=L/n
    @show dx,dy

    l2D=Laplacian2D(n,n,L,L,bc...)
    if bc==(neumann,neumann,neumann,neumann)
        #LP: for all neumann conditions the laplacian equation has no
        #solution if the rhs b does not belong to the range of A.
        #This impose the integral of b to be zero on the domain and
        # the sum of b on the mesh to be also zero...
        #See Mc Cormick and other for more details.
        @show sum(b)
        @assert(sum(b)<1.e-8)

        #Note that if this compatibility condition is ensured, and
        # x is a solution, then x+Cte is also a solution.
    end

    @time s2D=sparse_corr(l2D)
    splu,tlu=PoissonSparseLU(s2D)
    ts=@elapsed solve!(pref,b,splu)

    @time s2D_nocorr=SparseArrays.sparse(l2D)
    splu_nocorr,tlu_nocorr=PoissonSparseLU(s2D_nocorr)
    test_solver!(msm,PoissonGMG,(l2D,GSSmoother),splu_nocorr,b)

    test_solver!(msm,PoissonGMG_new,(l2D,GSSmoother),splu_nocorr,b)

    sname=string(n)*tostring(bc)

    @save "msm"*sname*".jld2" msm

    # measureplot(msm,sname)

    return p,pref,p,splu,b,msm

end

function tostring(bcs)
     r=string()
     for bc in bcs
         bc==dirichlet && (r*="D")
         bc==neumann && (r*="N")
    end
    r
end

function dump_measurments(msm,sname)
    header=keys(msm[1])[1:end]
    nrows=length(msm)
    ncols=length(header)
    table=Matrix{Any}(undef,nrows,ncols)
    iterations=Vector{Union{Nothing,Vector{Float64}}}(undef,nrows)
    for i=1:nrows
        for j=1:ncols
            table[i,j]=msm[i][header[j]]
        end
        table[i]=replace(table[i],"Poisson"=>"")
        table[i]=replace(table[i],"Preconditioner"=>"")
        table[i]=replace(table[i],"SparseCG"=>"PCG")
    end
     #Remove some measurements :
     removed=["GMG_CGSmoother"]
     table=table[filter(i->!(table[i,1] in removed),1:size(table,1)),:]

     #add Total Time and niter
     header=[header[1:3]...,"Total Time",header[4],"niters",header[5]]
     table=hcat(table[:,1:3],table[:,2].+table[:,3],table[:,4],niterations.(residuals(table[:,5])),table[:,5])
     #Sort
     sp=sortperm(table[:, 4])
     table=table[sp,:]
     # iterations=iterations[sp,:]
     nr,nc=size(table)
     println("all methods")
     pretty_table(table[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)
end

function go()
    #Choose a power of two
    n=512
    #Choose boundary conditions
    # bc=(neumann,neumann,neumann,neumann)
    bc=(dirichlet,neumann,dirichlet,neumann)
    # bc=(dirichlet,dirichlet,dirichlet,dirichlet)

    sname=string(n)*tostring(bc)
    p,pref,pbj,s,b,msm=testpoisson_GMG(n,bc)
    @save "msm"*sname*".jld2" msm

    dump_measurments(msm,sname)
    @show msm
    msm

   #*** Makie Plots
    x = range(0,1,length=n)
    y = range(0,1,length=n)

    p=pref

    pmax=maximum(p)-minimum(p)

    p./=pmax

    s=Makie.surface(x,y,p)
    xm, ym, zm = minimum(scene_limits(s))
    Makie.contour!(s,x,y,p, levels = 15, linewidth = 2, transformation = (:xy, zm))
    Makie.wireframe!(s,x,y,p, overdraw = true, transparency = true, color = (:black, 0.1))
    display(AbstractPlotting.PlotDisplay(), s)
    resize!(s,(1600,800))
    Makie.save("makie"*string(n)*tostring(bc)*".png", s)
    nothing
end

msm=go()


#=

  Euler #31 in Julia.

  Problem 31
  """
  In England the currency is made up of pound, £, and pence, p, and
  there are eight coins in general circulation:

     1p, 2p, 5p, 10p, 20p, 50p, £1 (100p) and £2 (200p).

  It is possible to make £2 in the following way:

     1×£1 + 1×50p + 2×20p + 1×5p + 1×2p + 3×1p

  How many different ways can £2 be made using any number of coins?
  """

  This Julia program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

include("Euler.jl")

# DP approach
function coins(c,money,m)
    len = length(c)
    if m == len
        return 1
    end

    sum1 = 0
    for i in m:len
        if money - c[i] == 0
            sum1 += 1
        end
        if  money - c[i] > 0
            sum1 += coins(c, money-c[i], i)
        end
    end
    return sum1;
end

# 0.00058022s
function euler31a()
    c = [200,100,50,20,10,5,2,1]
    return coins(c, 200, 1)
end

run_euler(euler31a)


# # Stabilization of nonlinear systems

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Systems and Control/stabilization_of_nonlinear_systems.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Systems and Control/stabilization_of_nonlinear_systems.ipynb)
# **Adapted from**: Examples 1, 2 and 3 of [PPA04]
#
# [PPA04] Prajna, Stephen, Pablo A. Parrilo, and Anders Rantzer.
# *Nonlinear control synthesis by convex optimization*.
# IEEE Transactions on Automatic Control 49.2 (2004): 310-314.

using Test #src
using DynamicPolynomials
@polyvar x[1:2]

using SumOfSquares
using CSDP
using LinearAlgebra # for ⋅
using MultivariatePolynomials
divergence(f) = sum(differentiate.(f, x))
function controller(f, g, b, α, degs)
    solver = optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true)
    model = SOSModel(solver)
    a = 1
    monos = monomials(x, degs)
    N = length(monos) + 1
    @variable(model, c[1:N] in MOI.NormOneCone(N))
    c_poly = polynomial(c[2:end], monos)
    fagc = f * a + g * c_poly
    @constraint(model, b * divergence(fagc) - α * differentiate(b, x) ⋅ fagc in SOSCone())
    @objective(model, Min, c[1])
    optimize!(model)
    if termination_status(model) != MOI.OPTIMAL
        @warn("Termination status $(termination_status(model)): $(raw_status(model))")
    end
    u = value(c_poly) / value(a)
    return MultivariatePolynomials.mapcoefficientsnz(coef -> abs(coef) < 1e-6 ? 0.0 : coef, u)
end

import DifferentialEquations, Plots
function phase_plot(f, quiver_scaling, Δt, X0, solver = DifferentialEquations.Tsit5())
    ∇(vx, vy) = [fi(x[1] => vx, x[2] => vy) for fi in f]
    ∇pt(v, p, t) = ∇(v[1], v[2])
    function traj(v0)
        tspan = (0.0, Δt)
        prob = DifferentialEquations.ODEProblem(∇pt, v0, tspan)
        return DifferentialEquations.solve(prob, solver, reltol=1e-8, abstol=1e-8)
    end
    ticks = -5:0.5:5
    X = repeat(ticks, 1, length(ticks))
    Y = X'
    Plots.quiver(X, Y, quiver = (x, y) -> ∇(x, y) / quiver_scaling, linewidth=0.5)
    for x0 in X0
        Plots.plot!(traj(x0), vars=(1, 2), label = nothing)
    end
    Plots.plot!(xlims = (-5, 5), ylims = (-5, 5))
end

# ## Example 1

g = [0, 1]
f = [x[2] - x[1]^3 + x[1]^2, 0]
b = 3x[1]^2 + 2x[1]*x[2] + 2x[2]^2
u = controller(f, g, b, 4, 0:3)
@test monomials(u) == [x[2]^3, x[1], x[2]] #src

# We find the controller above which gives the following phase plot.

phase_plot(f + g * u, 200, 10.0, [[x1, x2] for x1 in -5:5:5, x2 in -5:5:5 if x1 != 0 || x2 != 0])

# ## Example 2

g = [0, 1]
f = [2x[1]^3 + x[1]^2*x[2] - 6x[1]*x[2]^2 + 5x[2]^3, 0]
b = x[1]^2 + x[2]^2
u = controller(f, g, b, 2.5, 0:3)
@test monomials(u) == [x[1]^3, x[1]^2*x[2], x[1]*x[2]^2, x[2]^3] #src

# We find the controller above which gives the following phase plot.

phase_plot(f + g * u, 2000, 5.0, [[-1.0, -5.0], [1.0, 5.0]])

# ## Example 3

g = [0, x[2]]
f = [-6x[1]*x[2]^2 - x[1]^2*x[2] + 2x[2]^3, 0]
b = x[1]^2 + x[2]^2
u = controller(f, g, b, 3, 0:2)
@test monomials(u) == [x[1]^2, x[2]^2] #src

# We find the controller above which gives the following phase plot.

X0 = [Float64[x1, x2] for x1 in -5:5:5, x2 in -5:5:5 if x2 != 0]
ε = 1e-4 # We separate the starting point slightly from the hyperplane `x2 = 0` which is invariant.
push!(X0, [-4,  ε])
push!(X0, [-3, -ε])
push!(X0, [ 3,  ε])
push!(X0, [ 4, -ε])
phase_plot(f + g * u, 2000, 10.0, X0)


# auxiliary/axpby.jl
#
# Simple wrapper for the operation of computing α * x + β * y, and two
# special singleton types (immutables) to remove any overhead of multiplication
# by one or addition by zero.

immutable Zero <: Integer
end
immutable One <: Integer
end

const _zero = Zero()
const _one = One()

axpby(α::One,  x, β::One,  y) = x+y
axpby(α::Zero, x, β::One,  y) = y
axpby(α::One,  x, β::Zero, y) = x
axpby(α::Zero, x, β::Zero, y) = zero(y)

axpby(α::One,  x, β,       y) = x+β*y
axpby(α::Zero, x, β,       y) = β*y
axpby(α,       x, β::Zero, y) = α*x
axpby(α,       x, β::One,  y) = α*x+y

axpby(α,       x, β,       y) = α*x+β*y