export sample_points

import Distributions

"""
    sample_points(m::TriMesh, num_samples::Int=5000; returns_normals::Bool=false, eps::Number = 1e-6)

Uniformly samples `num_samples` points from the surface of TriMesh `m`.

`returns_normals` is optional keyword argument, to returns normals
from respective faces of samples. `eps` is optional keyword argument for
a small number to prevent division by zero for small surface areas.

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> points = sample_points(m, 5000)
julia> points, normals = sample_points(m, 5000; returns_normals=true)
```
"""
function sample_points(
    m::TriMesh{T,R,S},
    num_samples::Int = 5000;
    eps::Number = EPS,
)::S{T,3} where {T,R,S}
    verts_padded = get_verts_padded(m)
    faces_padded = get_faces_padded(m)
    faces_areas_padded = compute_faces_areas_padded(m)

    #(Fi,B)
    #TODO: condition for probvec fails in Float32
    faces_areas_prob = Zygote.ignore() do
        faces_areas_padded = Float64.(faces_areas_padded)
        faces_areas_prob =
            faces_areas_padded ./ max.(sum(faces_areas_padded; dims = 2), eps)
        faces_areas_prob[:, end, :] +=
            reshape(max.(1 .- sum(faces_areas_prob; dims = 2), 0.0), 1, :)
        return faces_areas_prob
    end
    samples = @ignore similar(verts_padded, 3, num_samples, m.N)
    samples = Zygote.bufferfrom(samples)

    for (i, _len) in enumerate(m._faces_len)
        # Distributions package expects Array!
        probvec = cpu(faces_areas_prob[1, 1:_len, i])
        dist = Distributions.Categorical(probvec)
        sample_faces_idx = @ignore rand(dist, num_samples)
        sample_faces = faces_padded[:, sample_faces_idx, i]
        samples[:, :, i] = _sample_points(
            S,
            verts_padded[:, 1:m._verts_len[i], i],
            sample_faces,
            num_samples,
        )
    end

    return copy(samples)
end

function _sample_points(
    ::Type{S},
    verts::AbstractArray{T,2},
    sample_faces::AbstractArray{R,2},
    num_samples::Int,
) where {S,T,R}

    v1 = verts[:, sample_faces[1, :]]
    v2 = verts[:, sample_faces[2, :]]
    v3 = verts[:, sample_faces[3, :]]
    (w1, w2, w3) = S.(_rand_barycentric_coords(num_samples))
    samples = (w1 .* v1) + (w2 .* v2) + (w3 .* v3)
    return samples
end

function _rand_barycentric_coords(num_samples::Int)
    u = sqrt.(rand(Float32, 1, num_samples))
    v = rand(Float32, 1, num_samples)
    w1 = 1.0f0 .- u
    w2 = u .* (1.0f0 .- v)
    w3 = u .* v
    return (w1, w2, w3)
end

"""
    normalize!(m::TriMesh)

Normalize each mesh in TriMesh `m` with mean centered at origin and unit standard deviation
and overwrite the `m` with normalized TriMesh.

See also: [`normalize`](@ref)

### Examples:

```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> normalize!(m)
```
"""
function normalize!(m::TriMesh{T,R,S}) where {T,R,S}
    verts_padded = get_verts_padded(m)
    _len = S(reshape(m._verts_len, 1, 1, :)) # move to `S` storage type
    _centroid = sum(verts_padded; dims = 2) ./ _len
    _correction = ((_centroid .^ 2) .* (m.V .- _len))
    _std =
        sqrt.(
            (sum((verts_padded .- _centroid) .^ 2; dims = 2) - _correction) ./ (_len .- 1),
        )
    verts_padded = (verts_padded .- _centroid) ./ max.(_std, EPS)
    m._verts_list = _padded_to_list(verts_padded, m._verts_len)
    return m
end

"""
    normalize(m::TriMesh)

Normalize each mesh in TriMesh `m` with mean centered at origin and unit standard deviation

See also: [`normalize!`](@ref)

### Examples:

```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> m = normalize(m)
```
"""
function normalize(m::TriMesh)
    m = deepcopy(m)
    normalize!(m)
    return m
end

"""
    scale!(m::TriMesh, factor::Number)
    scale!(m::TriMesh, factor::AbstractArray{<:Number})

Scale the TriMesh `m` by scaling factor `factor`
and overwrite `m` with scaled TriMesh. If `factor`
is array of size `(3, )`, then TriMesh will be scale
by scaling factor of respective dimension.

Scaling factor `factor` (each element in case of array)
should be strictly greater than `0.0`.

See also: [`scale`](@ref)

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> scale!(m, 1.0)
julia> scale!(m, [1.0, 1.0, 1.0])
```
"""
function scale!(m::TriMesh, factor::Float32)
    (factor > 0.0) || error("factor must be greater than 0.0")
    verts_packed = get_verts_packed(m)
    lmul!(factor, verts_packed)
    m._verts_packed = verts_packed
    return m
end

function scale!(m::TriMesh, factor::AbstractArray{Float32})
    (size(factor) == (3,)) ||
        error("factor must be (3, ), but instead got $(size(factor)) array")
    (factor .> 0.0) || error("factor must be greater than 0.0")
    verts_packed = get_verts_packed(m) .* reshape(factor, :, 1)
    m._verts_packed = verts_packed
    return m
end

scale!(m::TriMesh, factor) = scale!(m, Float32.(factor))

"""
    scale(m::TriMesh, factor::Number)
    scale(m::TriMesh, factor::AbstractArray{<:Number})

Scale the TriMesh `m` by scaling factor `factor`.
If `factor` is array of size `(3, )`, then TriMesh
will be scaleby scaling factor of respective dimension.

Scaling factor `factor` (each element in case of array)
should be strictly greater than `0.0`.

See also: [`scale!`](@ref)

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> m = scale(m, 1.0)
julia> m = scale!(m, [1.0, 1.0, 1.0])
```
"""
function scale(m::TriMesh, factor::Union{Float32,AbstractArray{Float32}})
    m = deepcopy(m)
    scale!(m, factor)
    return m
end

scale(m::TriMesh, factor) = scale(m, Float32.(factor))

"""
    rotate!(m::TriMesh, rotmat::AbstractArray{<:Number,2})
    rotate!(m::TriMesh, rotmat::AbstractArray{<:Number,3})

Rotate the TriMesh `m` by rotation matrix `rotmat`
and overwrite `m` with rotated TriMesh.

Rotation matrix `rotmat` should be of size `(3,3)` or `(3,3,B)`
where B is the batch size of TriMesh.

See also: [`rotate`](@ref)

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> rotmat = rand(3,3)
julia> rotate!(m, rotmat)
```
"""
function rotate!(m::TriMesh, rotmat::AbstractArray{Float32,2})
    size(rotmat) == (3, 3) ||
        error("rotmat must be (3, 3) array, but instead got $(size(rotmat)) array")
    verts_packed = transpose(rotmat) * get_verts_packed(m)
    m._verts_packed = verts_packed
    return m
end

function rotate!(m::TriMesh, rotmat::AbstractArray{Float32,3})
    size(rotmat) == (3, 3, m.N) ||
        error("rotmat must be (3, 3, $(m.N)) array, but instead got $(size(rotmat)) array")
    verts_padded = Flux.batched_mul(Flux.batched_transpose(rotmat), get_verts_padded(m))
    m._verts_padded = verts_padded
    return m
end

rotate!(m::TriMesh, rotmat::AbstractArray{<:Number}) = rotate!(m, Float32.(rotmat))

"""
    rotate(m::TriMesh, rotmat::AbstractArray{<:Number,2})

Rotate the TriMesh `m` by rotation matrix `rotmat`.

Rotation matrix `rotmat` should be of size `(3,3)`

See also: [`rotate!`](@ref)

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> rotmat = rand(3,3)
julia> m = rotate(m, rotmat)
```
"""
function rotate(m::TriMesh, rotmat::AbstractArray{<:Number})
    m = deepcopy(m)
    rotate!(m, rotmat)
    return m
end

"""
    realign!(src::TriMesh, tgt::TriMesh)
    realign!(src::TriMesh, tgt_min::AbstractArray{<:Number,2}, tgt_max::AbstractArray{<:Number,2})

Re-Align the TriMesh `src` with the axis aligned bounding box of mesh at `index` in TriMesh `tgt`
and overwrite `src` with re-aligned TriMesh.

See also: [`realign`](@ref)

### Examples:
```jldoctest
julia> src = load_trimesh("teapot.obj")
julia> tgt = scale(src, 2.0)
julia> realign!(src, tgt)
```
"""
function realign!(
    src::TriMesh,
    tgt_min::AbstractArray{Float32,2},
    tgt_max::AbstractArray{Float32,2},
)
    verts_padded = get_verts_padded(src)
    src_min = minimum(verts_padded, dims = 2)
    src_max = maximum(verts_padded, dims = 2)
    verts_padded =
        ((verts_padded .- src_min) ./ (src_max - src_min .+ EPS)) .* (tgt_max - tgt_min) .+
        tgt_min
    src._verts_packed = _padded_to_packed(verts_padded, src._verts_len)
    return src
end

realign!(
    src::TriMesh,
    tgt_min::AbstractArray{<:Number,2},
    tgt_max::AbstractArray{<:Number,2},
) = realign!(src, Float32.(tgt_min), Float32.(tgt_max))

realign!(src::TriMesh, tgt::AbstractArray{<:Number,2}) = realign!(src, Float32.(tgt))

function realign!(src::TriMesh, tgt::AbstractArray{Float32,2})
    tgt_min = minimum(tgt, dims = 2)
    tgt_max = maximum(tgt, dims = 2)
    realign!(src, tgt_min, tgt_max)
    return src
end

realign!(src::TriMesh, tgt::TriMesh, index::Integer = 1) =
    realign!(src, Float32.(get_verts_list(tgt)[index]))

"""
    realign(src::TriMesh, tgt::TriMesh, index::Integer=1)
    realign(src::TriMesh, tgt_min::AbstractArray{<:Number,2}, tgt_max::AbstractArray{<:Number,2})

Re-Align the TriMesh `src` with the axis aligned bounding box of mesh at `index` in TriMesh `tgt`.

See also: [`realign`](@ref)

### Examples:
```jldoctest
julia> src = load_trimesh("teapot.obj")
julia> tgt = scale(src, 2.0)
julia> src = realign(src, tgt)
```
"""
function realign(
    src::TriMesh,
    tgt_min::AbstractArray{<:Number,2},
    tgt_max::AbstractArray{<:Number,2},
)
    src = deepcopy(src)
    realign!(src, tgt_min, tgt_max)
    return src
end

function realign(src::TriMesh, tgt::TriMesh, index::Integer = 1)
    src = deepcopy(src)
    realign!(src, tgt, index)
    return src
end

"""
    translate!(m::TriMesh, vector::Number)
    translate!(m::TriMesh, vector::AbstractArray{<:Number})

Translate the TriMesh `m` by translating vector `vector`
and overwrite `m` with translated TriMesh. If `vector`
is a number, then TriMesh will be translated by same number
in all dimension.

See also: [`translate`](@ref)

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> translate!(m, 0.0)
julia> translate!(m, [0.0, 0.0, 0.0])
```
"""
translate!(m::TriMesh, vector::Float32) = translate!(m, fill(vector, (3,)))

function translate!(m::TriMesh, vector::AbstractArray{Float32})
    (size(vector) == (3,)) ||
        error("vector must be (3, ), but instead got $(size(vector)) array")
    verts_packed = get_verts_packed(m)
    verts_packed = verts_packed .+ reshape(vector, :, 1)
    m._verts_packed = verts_packed
    return m
end

translate!(m::TriMesh, vector) = translate!(m, Float32.(vector))

"""
    translate(m::TriMesh, vector::Number)
    translate(m::TriMesh, vector::AbstractArray{<:Number})

Translate the TriMesh `m` by translating vector `vector`.
If `vector` is a number, then TriMesh will be translated
by same number in all dimension.

See also: [`translate!`](@ref)

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> m = translate(m, 0.0)
julia> m = translate(m, [0.0, 0.0, 0.0])
```
"""
function translate(m::TriMesh, args...)
    m = deepcopy(m)
    translate!(m, args...)
end

"""
    offset!(m::TriMesh, offset_verts::AbstractArray{<:Number,2})

Add offset to the vertices of the TriMesh `m` by offset vertices `offset_verts_packed`
and overwrite `m` with updated TriMesh.

See also: [`offset`](@ref)

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> offset_verts = ones(get_verts_packed(m))
julia> offset!(m, offset_verts)
```
"""

function offset!(m::TriMesh, offset_verts_packed::AbstractArray{Float32,2})
    verts_packed = get_verts_packed(m)
    (size(offset_verts_packed) == size(verts_packed)) ||
        error("mesh and offset_verts size mismatch")
    verts_packed += offset_verts_packed
    m._verts_packed = verts_packed
    return m
end

offset!(m::TriMesh, offset_verts_packed::AbstractArray{<:Number,2}) =
    offset!(m, Float32.(offset_verts_packed))

"""
    offset(m::TriMesh, offset_verts_packed::AbstractArray{<:Number,2})

Add offset to the vertices of the TriMesh `m` by offset vertices `offset_verts_packed`.

See also: [`offset!`](@ref)

### Examples:
```jldoctest
julia> m = load_trimesh("teapot.obj")
julia> offset_verts = ones(get_verts_packed(m))
julia> m = offset(m, offset_verts)
```
"""
function offset(m::TriMesh, offset_verts_packed::AbstractArray{<:Number,2})
    m = deepcopy(m)
    offset!(m, offset_verts_packed)
end


"
 Test for projected newton QP solver

 author: Atsushi Sakai
"

using ArgMin
using Test

function test_solve_qp_with_projected_newton1()
    println("test_solve_qp_with_projected_newton1")
    n 		= 5
    g 		= randn(n)
    H 		= randn(n,n)
    H 		= H*H'
    lower 	= -ones(n)
    upper 	=  ones(n)

    xstar, status = solve_qp_with_projected_newton(
                                 H, g, lower, upper)

    @test status["optimal"] == true
end

function test_solve_qp_with_projected_newton2()
    println("test_solve_qp_with_projected_newton2")
    n 		= 3
    q = [3.0;4.0;-5.0]
    A = [1.0 0.0 -2.0;
         1.0 2.0 0.0;
         0.0 0.0 3.0]
    A = A*A'
    xmin = [-3.0;-5.0;-2.0]
    xmax = [1.0;1.0;1.0]

    xstar, status = solve_qp_with_projected_newton(
                                 A, q, xmin, xmax)
    @test xstar[1] ≈ 0.791667 atol=0.01
    @test xstar[2] ≈ -0.958333 atol=0.01
    @test xstar[3] ≈ 1.0 atol=0.01

end

test_solve_qp_with_projected_newton1()
test_solve_qp_with_projected_newton2()


function approxextrema(I)
    arr = sample(I)
    return (minimum(arr), maximum(arr))
end

"""
    estimatequantile(array, p)

Quickly estimates the quantile `p` of a possibly large array by using a subset of the data.
"""
function estimatequantile(array, p)
    try
        return quantile(sample(array; n=1e5), p)
    catch
        @warn "quantile could not be estimated! (maybe only NaNs)"
        return 0
    end
end

function sample(I; n=10000)
    n = min(n, length(I))
    len = ceil(Int, √n) # take len blocks of len elements
    startindices = round.(Int, range(firstindex(I) - 1, lastindex(I) - len; length=len))
    indices = vcat((i .+ (1:len) for i in startindices)...)
    ret = filter(isfinite, I[indices])
    if isempty(ret)
        ret = filter(isfinite, I)
    end
    return ret
end

function get_corner_indices(I; max_length=10)
    d = size(I)
    n = min.(max_length, ceil.(Int, d ./ 3)) # n voxels for each dim
    getrange(num, len) = [1:num, (len-num+1):len] # first and last voxels
    return collect(Iterators.product(getrange.(n, d)...))
end

# estimate noise parameters from corner without signal
"""
    estimatenoise(image::AbstractArray)

Estimates the noise from the corners of the image.
It assumes that at least one corner is without signal and only contains noise.
"""
function estimatenoise(image::AbstractArray)
    corners = get_corner_indices(image)
    (lowestmean, ind) = findmin(mean.(filter(isfinite, image[I...]) for I in corners))
    sigma = std(filter(isfinite, image[corners[ind]...]))
    if isnan(sigma) # no corner available
        # estimation that is correct if half the image is signal and half noise
        sigma = 2estimatesigma_from_quantile(image, 1/4)
        lowestmean = sigma / 2
    end
    return lowestmean, sigma
end

# sigma is only calculated for quantile (biased)
function estimatesigma_from_quantile(image, quantile)
    q = estimatequantile(image, quantile)
    samples = filter(x -> x < q, sample(image))
    return std(samples)
end

"""
    robustmask!(image)
    robustmask!(image; maskedvalue)

Creates a mask and applies it inplace.
It assumes that at least one corner is without signal and only contains noise.
"""
function robustmask!(image; maskedvalue=if eltype(image) <: AbstractFloat NaN else 0 end)
    image[.!robustmask(image)] .= maskedvalue
    image
end
"""
    robustmask(weight::AbstractArray)

Creates a mask from a weights images.
It assumes that at least one corner is without signal and only contains noise.
"""
function robustmask(weight::AbstractArray)
    μ, σ = estimatenoise(weight)
    m = mean(filter(isfinite, weight[weight .> 5σ]))
    mask = weight .> maximum((5σ, m/5))
    # remove small holes and minimally grow
    boxsizes=[[3,3] for i in 1:ndims(weight)]
    return gaussiansmooth3d(mask; nbox=2, boxsizes) .> 0.55
end

getcomplex(fnmag::AbstractString, fnphase::AbstractString) = getcomplex(niread(fnmag), niread(fnphase))

function getcomplex(mag, phase)
    minp, maxp = approxextrema(phase)
    mag .* exp.((2im * pi / (maxp - minp)) .* phase)
end

function readfromtextheader(filename, searchstring)
    for line in readlines(open(filename, "r"))
        if occursin(searchstring, line)
            # regex to search for "= " or ": " and return the following non-whitespace characters
            return match(r"(?<=(= |: ))(\S+)", line).match
        end
    end
end

# root sum of squares combination
"""
    RSS(mag; dims=ndims(mag))

Performs root-sum-of-squares combination along the last dimension of `mag`.
The dimension can be specificed via the `dims` keyword argument.
"""
RSS(mag; dims=ndims(mag)) = dropdims(.√sum(mag.^Float32(2); dims); dims)

function getscaledimage(array, div::Number, offset = 0, type::Symbol = :trans)
    array = reshape(array, size(array)[1:2]) # drops trailing singleton dimensions
    scaled = if offset != 0
        (array .- offset) .* (1 / div) .+ 0.5
    else
        array .* (1 / div)
    end
    scaled[isnan.(scaled) .| (scaled .< 0)] .= 0
    scaled[scaled .> 1] .= 1
    if type == :trans
        scaled = reverse(permutedims(scaled, [2 1]); dims = 1)
    else
    end
    scaled
end

function getscaledimage(array, type::Symbol = :trans)
    scaled = robustrescale(array, 0, 1, threshold=true)
    getscaledimage(scaled, 1, 0, type)
end

"""
    robustrescale(array, newmin, newmax; threshold=false, mask=trues(size(array)), datatype=Float64)

Rescales the image to the the new range, disregarding outliers.
Only values inside `mask` are used for estimating the rescaling option
"""
robustrescale(array, newmin, newmax; threshold=false, mask=trues(size(array)), datatype=Float64) =
    robustrescale!(datatype.(array), newmin, newmax; threshold, mask)

function robustrescale!(array, newmin, newmax; threshold=false, mask=trues(size(array)))
    mask[isnan.(array)] .= false
    q = [0.01, 0.99] # quantiles
    oldq = estimatequantile(array[mask], q)
    oldrange = (oldq[2] - oldq[1]) / (q[2] - q[1])
    oldmin = oldq[1] - q[1] * oldrange
    newrange = newmax - newmin

    array .= (array .- oldmin) .* (newrange / oldrange) .+ newmin

    if threshold
        array[array .< newmin] .= newmin
        array[array .> newmax] .= newmax
    end
    array
end

function rescale(array, newmin, newmax; datatype=eltype(array))
    rescale!(datatype.(array), newmin, newmax)
end

function rescale!(array, newmin, newmax)
    oldmin, oldmax = approxextrema(array)
    factor = (newmax - newmin) / (oldmax - oldmin)
    array .= (array .- oldmin) .* factor .+ newmin
end

"""
    to_dim(V::AbstractVector, dim::Int)

    to_dim(a::Real, dim::Int)

Converts a vector or number to a higher dimension.

# Examples
```julia-repl
julia> to_dim(5, 3)
1×1×1 Array{Int64, 3}:
[:, :, 1] =
 5
julia> to_dim([1,2], 2)
 1×2 Matrix{Int64}:
  1  2
```
"""
to_dim(a::Real, dim::Int) = to_dim([a], dim)
to_dim(V::AbstractArray, dim::Int) = reshape(V, ones(Int, dim-1)..., :)

"""
    getHIP(mag, phase; echoes=[1,2])

    getHIP(compl; echoes=[1,2])

Calculates the Hermitian Inner Product between the specified echoes.
"""
function getHIP(mag, phase; echoes=[1,2])
    e1, e2 = echoes
    compl = zeros(ComplexF64, size(mag)[1:3])
    for iCha in 1:size(mag, 5)
        compl .+= exp.(1.0im .* (phase[:,:,:,e2,iCha] .- phase[:,:,:,e1,iCha])) .* mag[:,:,:,e1,iCha] .* mag[:,:,:,e2,iCha]
    end
    compl
end

function getHIP(compl; echoes=[1,2])
    e1, e2 = echoes
    c = zeros(eltype(compl), size(compl)[1:3])
    for iCha in 1:size(compl, 5)
        c .+=  compl[:,:,:,e2,iCha] .* conj.(compl[:,:,:,e1,iCha])
    end
    return c
end


push!(LOAD_PATH, abspath(@__DIR__,".."))

using SSNVariability ; const S = SSNVariability
using Test
using LinearAlgebra,Calculus,Statistics,Distributions
using Plots,NamedColors ; theme(:dark)
using Random
Random.seed!(0)

using OrdinaryDiffEq, StochasticDiffEq
# import DiffEqNoiseProcess: CorrelatedWienerProcess,WienerProcess

##

function eiplot(t::V,e::V,i::V) where V<:Vector{<:Real}
  plt=plot()
  plot!(plt,t,hcat(e,i) ; color=[:red :blue] , linewidth=3,leg=false)
end

function eiplot(t::V,e::M,i::M) where {V<:Vector,M<:Matrix{<:Real}}
  mp(mat)=mapslices(norm,mat;dims=1)[:]
  return eiplot(t,mp(e),mp(i))
end

function plotvs(x::AbstractArray{<:Real},y::AbstractArray{<:Real})
  x,y=x[:],y[:]
  @info """
  The max differences between the two are $(extrema(x .-y ))
  """
  plt=plot()
  scatter!(plt,x,y;leg=false,ratio=1,color=:white)
  lm=xlims()
  plot!(plt,identity,range(lm...;length=3);linestyle=:dash,color=:yellow)
  return plt
end

##


@testset "scomponents" begin
    A = load_matrix_network("cores_example")
    sizes = [15;1;5]
    components = 3
    @testset "MatrixNetwork" begin
        cc = scomponents(MatrixNetwork(A))
        @test cc.sizes == sizes
        @test length(cc.map) == 21
        @test cc.number == components
    end
    @testset "SparseMatrixCSC" begin
        cc = scomponents(A)
        @test cc.sizes == sizes
        @test length(cc.map) == 21
        @test cc.number == components
    end
    @testset "CSR" begin
        cc = scomponents(sparse_to_csr(A)...)
        @test cc.sizes == sizes
        @test length(cc.map) == 21
        @test cc.number == components
    end
    @testset "triplet" begin
        cc = scomponents(findnz(A)[1], findnz(A)[2])
        @test cc.sizes == sizes
        @test length(cc.map) == 21
        @test cc.number == components
    end
    @testset "strong_components_map" begin
        sci = strong_components_map(MatrixNetwork(A))
        @test strong_components_map(A) == sci
        @test strong_components_map(sparse_to_csr(A)...) == sci
        @test strong_components_map(findnz(A)[1], findnz(A)[2]) == sci
    end
    @testset "empty" begin
        @test scomponents(empty_graph(0)).sizes == []
    end
end


# using Revise, Plots
using OrdinaryDiffEq, ForwardDiff, Test
	using BifurcationKit, LinearAlgebra, Parameters, Setfield
	const BK = BifurcationKit
	const FD = ForwardDiff

norminf = x -> norm(x, Inf)

function Fsl!(f, u, p, t)
	@unpack r, μ, ν, c3, c5 = p
	u1 = u[1]
	u2 = u[2]

	ua = u1^2 + u2^2

	f[1] = r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1
	f[2] = r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2
	return f
end

Fsl(x, p) = Fsl!(similar(x), x, p, 0.)

Fode(f, x, p, t) = Fsl!(f, x, p, t)
dFsl(x, dx, p) = FD.derivative(t -> Fsl(x .+ t .* dx, p), 0.)

par_sl = (r = 0.5, μ = 0., ν = 1.0, c3 = 1.0, c5 = 0.0,)
u0 = [.001, .001]

function FslMono!(f, x, p, t)
	u = x[1:2]
	du = x[3:4]
	Fsl!(f[1:2], u, p, t)
	f[3:4] .= dFsl(u, du, p)
end

####################################################################################################
par_hopf = (@set par_sl.r = 0.1)
####################################################################################################
# continuation
optconteq = ContinuationPar(ds = -0.01, detectBifurcation = 3, pMin = -0.5, nInversion = 4)
br, = continuation(Fsl, u0, par_hopf, (@lens _.r), optconteq)
####################################################################################################
prob = ODEProblem(Fode, u0, (0., 100.), par_hopf)
probMono = ODEProblem(FslMono!, vcat(u0, u0), (0., 100.), par_hopf)
####################################################################################################
sol = solve(probMono, KenCarp4(), abstol =1e-9, reltol=1e-6)
sol = solve(prob, KenCarp4(), abstol =1e-9, reltol=1e-6)
# plot(sol[1,:], sol[2,:])
####################################################################################################
section(x, T) = x[1] #* x[end]
# standard simple shooting
M = 1
dM = 1
_pb = ShootingProblem(Fsl, par_hopf, prob, KenCarp4(), 1, section; abstol =1e-10, reltol=1e-9)

initpo = [0.13, 0., 6.]
res = _pb(initpo, par_hopf)

# test of the differential of the shooting method

_dx = rand(3)
resAD = FD.derivative(z -> _pb(initpo .+ z .* _dx, par_hopf), 0.)
resFD = (_pb(initpo .+ 1e-8 .* _dx, par_hopf) - _pb(initpo, par_hopf)) .* 1e8
resAN = _pb(initpo, par_hopf, _dx; δ = 1e-8)
@test norm(resAN - resFD, Inf) < 4e-6
@test norm(resAN - resAD, Inf) < 4e-6
####################################################################################################
# test shooting interface M = 1
_pb = ShootingProblem(Fsl, par_hopf, prob, KenCarp4(), [initpo[1:end-1]]; abstol =1e-10, reltol=1e-9)
res = _pb(initpo, par_hopf)
res = _pb(initpo, par_hopf, initpo)

# test the jacobian of the functional in the case M=1
_Jad = FD.jacobian( x -> _pb(x, par_hopf), initpo)
_Jana = _pb(Val(:JacobianMatrix), initpo, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-7

_pb2 = ShootingProblem(Fsl, par_hopf, prob, Rodas4(), probMono, Rodas4(autodiff=false), [initpo[1:end-1]]; abstol =1e-10, reltol=1e-9)
res = _pb2(initpo, par_hopf)
res = _pb2(initpo, par_hopf, initpo)
BK.isSimple(_pb2)

# we test this using Newton - Continuation
ls = GMRESIterativeSolvers(reltol = 1e-5, N = length(initpo))
optn = NewtonPar(verbose = false, tol = 1e-9,  maxIter = 20, linsolver = ls)
deflationOp = BK.DeflationOperator(2.0, (x,y) -> dot(x[1:end-1], y[1:end-1]), 1.0, [zeros(3)])
outpo, _, flag,_ = newton(_pb,
	initpo, par_hopf,
	optn,
	normN = norminf)
	@test flag

BK.getPeriod(_pb, outpo, par_hopf)
BK.getAmplitude(_pb, outpo, par_hopf)
BK.getMaximum(_pb, outpo, par_hopf)
BK.getTrajectory(_pb, outpo, par_hopf)

opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.01, ds= -0.01, pMax = 4.0, maxSteps = 30, detectBifurcation = 2, nev = 2, newtonOptions = @set optn.tol = 1e-7)
	br_pok2, upo , _= continuation(
		_pb,
		outpo, par_hopf, (@lens _.r),
		opts_po_cont;
		tangentAlgo = BorderedPred(),
		verbosity = 0,
		plot = false,
		# plotSolution = (x, p; kwargs...) -> plot!(x[1:end-1]; kwargs...),
		printSolution = (u, p) -> norm(u[1:2]),
		normC = norminf)
# plot(br_pok2)
####################################################################################################
# test automatic branch switching
JFsl(x, p) = FD.jacobian(t -> Fsl(t, p), x)
d1Fsl(x, p, dx) = FD.derivative(t -> Fsl(x .+ t .* dx, p), 0.)
d2Fsl(x, p, dx1, dx2) = FD.derivative(t -> d1Fsl(x .+ t .* dx2, p, dx1), 0.)
d3Fsl(x, p, dx1, dx2, dx3) = FD.derivative(t -> d2Fsl(x .+ t .* dx3, p, dx1, dx2), 0.)
jet = (Fsl, JFsl, d2Fsl, d3Fsl)

br_pok2, = continuation(jet..., br, 1, opts_po_cont, ShootingProblem(1, par_hopf, prob, KenCarp4());printSolution = (u, p) -> norm(u[1:2]), normC = norminf)

# test matrix-free computation of floquet coefficients
eil = EigKrylovKit(dim = 2, x₀=rand(2))
opts_po_contMF = @set opts_po_cont.newtonOptions.eigsolver = eil
opts_po_contMF = @set opts_po_cont.detectBifurcation=0
br_pok2, = continuation(jet...,br,1, opts_po_contMF, ShootingProblem(1, par_hopf, prob, Rodas4());printSolution = (u, p) -> norm(u[1:2]), normC = norminf, plot=false)
####################################################################################################
# test shooting interface M > 1
_pb = ShootingProblem(Fsl, par_hopf, prob, KenCarp4(), [initpo[1:end-1],initpo[1:end-1],initpo[1:end-1]]; abstol =1e-10, reltol=1e-9)
initpo = [0.13, 0, 0, 0.13, 0, 0.13 , 6.3]
res = _pb(initpo, par_hopf)
res = _pb(initpo, par_hopf, initpo)
# test the jacobian of the functional in the case M=1
_Jad = FD.jacobian( x -> _pb(x, par_hopf), initpo)
_Jana = _pb(Val(:JacobianMatrix), initpo, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-7
####################################################################################################
# Single Poincaré Shooting with hyperplane parametrization
normals = [[-1., 0.]]
centers = [zeros(2)]

probPsh = PoincareShootingProblem(2, par_hopf,prob, Rodas4(), probMono, Rodas4(); abstol=1e-10, reltol=1e-9)
probPsh = PoincareShootingProblem(2, par_hopf,prob, Rodas4(); rtol = abstol=1e-10, reltol=1e-9)

probPsh = PoincareShootingProblem(Fsl, par_hopf,
		prob, Rodas4(),
		probMono, Rodas4(),
		normals, centers; rtol = abstol =1e-10, reltol=1e-9)

initpo_bar = BK.R(probPsh, [0, 0.4], 1)

BK.E(probPsh, [1.0,], 1)
initpo_bar = [0.4]

probPsh(initpo_bar, par_hopf)

# test of the analytical formula for jacobian of the functional
_Jad = BifurcationKit.finiteDifferences( x-> probPsh(x, par_hopf), initpo_bar)
_Jana = probPsh(Val(:JacobianMatrix), initpo_bar, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-4

ls = GMRESIterativeSolvers(reltol = 1e-7, N = length(initpo_bar), maxiter = 500, verbose = false)
	eil = EigKrylovKit(dim = 1, x₀=rand(1))
	optn = NewtonPar(verbose = false, tol = 1e-8,  maxIter = 140, linsolver = ls, eigsolver = eil)
	deflationOp = BK.DeflationOperator(2.0, dot, 1.0, [zero(initpo_bar)])
	outpo, _, flag, = newton(probPsh,
			initpo_bar, par_hopf,
			optn; normN = norminf)
	@test flag

BK.getPeriod(probPsh, outpo, par_hopf)
BK.getAmplitude(probPsh, outpo, par_hopf)
BK.getMaximum(probPsh, outpo, par_hopf)
BK.getTrajectory(probPsh, outpo, par_hopf)

probPsh = PoincareShootingProblem(Fsl, par_hopf,
		prob, Rodas4(),
		# probMono, Rodas4(autodiff=false),
		normals, centers; rtol = 1e-8)

opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.015, ds= -0.01, pMax = 4.0, maxSteps = 50, newtonOptions = setproperties(optn;tol = 1e-9, eigsolver = eil), detectBifurcation = 1)
	br_pok2, = continuation(
		probPsh, outpo, par_hopf, (@lens _.r),
		opts_po_cont; verbosity = 0,
		tangentAlgo = BorderedPred(),
		plot = false,
		# plotSolution = (x, p;kwargs...) -> plot!(x; kwargs...),
		printSolution = (u, p) -> norm(u), normC = norminf)
# plot(br_pok2)
####################################################################################################
# normals = [[-1., 0.], [1, -1]]
# centers = [zeros(2), zeros(2)]
# initpo_bar = [1.04, -1.04/√2]

normals = [[-1., 0.], [1, 0]]
centers = [zeros(2), zeros(2)]
initpo_bar = [0.2, -0.2]

probPsh = PoincareShootingProblem(Fsl, par_hopf, prob, KenCarp4(), normals, centers; abstol=1e-10, reltol=1e-9)
# version with analytical jacobian
probPsh2 = PoincareShootingProblem(Fsl, par_hopf, prob, KenCarp4(), normals, centers; abstol=1e-10, reltol=1e-9, δ = 0)

# test of the analytical formula for jacobian of the functional
_Jad = BifurcationKit.finiteDifferences( x-> probPsh(x, par_hopf), initpo_bar)
# _Jphifd = BifurcationKit.finiteDifferences(x->probPsh.flow(x, par_hopf, Inf64), [0,0.4]; δ=1e-8)
_Jana = probPsh(Val(:JacobianMatrix), initpo_bar, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-5

ls = GMRESIterativeSolvers(reltol = 1e-5, N = length(initpo_bar), maxiter = 500, verbose = false)
	eil = EigKrylovKit(dim = 1, x₀=rand(1))
	optn = NewtonPar(verbose = false, tol = 1e-9,  maxIter = 140, linsolver = ls, eigsolver = eil)
	deflationOp = DeflationOperator(2.0, dot, 1.0, [zero(initpo_bar)])
	outpo, = newton(probPsh2, initpo_bar, par_hopf, optn; normN = norminf)
	outpo, = newton(probPsh, initpo_bar, par_hopf, optn; normN = norminf)

getPeriod(probPsh, outpo, par_hopf)

opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.025, ds= -0.01, pMax = 4.0, maxSteps = 50, newtonOptions = (@set optn.tol = 1e-9), detectBifurcation = 1, nev = 2)
	br_pok2, = continuation(probPsh, outpo, par_hopf, (@lens _.r),
		opts_po_cont; verbosity = 0,
		tangentAlgo = BorderedPred(),
		plot = false,
		# plotSolution = (x, p;kwargs...) -> plot!(x; kwargs...),
		printSolution = (u, p) -> norm(u), normC = norminf)

####################################################################################################
normals = [[-1., 0.], [1, 0], [0, 1]]
centers = [zeros(2), zeros(2), zeros(2)]
initpo = [[0., 0.4], [0, -.3], [0.3, 0]]

probPsh = PoincareShootingProblem(Fsl, par_hopf, prob, KenCarp4(), normals, centers; abstol=1e-10, reltol=1e-9)

initpo_bar = reduce(vcat, [BK.R(probPsh, initpo[ii], ii) for ii in eachindex(centers)])
# same with projection function
initpo_bar = reduce(vcat, BK.projection(probPsh, initpo))

# test of the other projection function
BK.projection(probPsh, reduce(hcat, initpo)')

probPsh(initpo_bar, par_hopf; verbose = true)

# test of the analytical formula for jacobian of the functional
_Jad = BifurcationKit.finiteDifferences( x-> probPsh(x, par_hopf), initpo_bar)
_Jana = probPsh(Val(:JacobianMatrix), initpo_bar, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-5

ls = GMRESIterativeSolvers(reltol = 1e-7, N = length(initpo_bar), maxiter = 10, verbose = false)
	optn = NewtonPar(verbose = false, tol = 1e-9,  maxIter = 50, linsolver = ls, eigsolver = eil)
	outpo, = newton(probPsh, initpo_bar, par_hopf, optn; normN = norminf)

for ii=1:length(normals)
	@show BK.E(probPsh	, [outpo[ii]], ii)
end

getPeriod(probPsh, outpo, par_hopf)

opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.025, ds= -0.005, pMax = 4.0, maxSteps = 10, newtonOptions = setproperties(optn; tol = 1e-8), detectBifurcation = 1)
	br_hpsh, upo , _= continuation(probPsh, outpo, par_hopf, (@lens _.r),
		opts_po_cont; normC = norminf)
# plot(br_hpsh)
####################################################################################################
# test automatic branch switching with most possible options
# calls with analytical jacobians
br_psh, = continuation(jet..., br,1, (@set opts_po_cont.ds = 0.005), PoincareShootingProblem(1, par_hopf, prob, KenCarp4(), probMono, KenCarp4(); abstol=1e-10, reltol=1e-9); normC = norminf)

opts_po_cont = @set opts_po_cont.detectBifurcation = 0
for M in [1,2], linearPO in [:autodiffMF, :MatrixFree, :autodiffDense, :FiniteDifferencesDense]
	@show M, linearPO
	br_psh, = continuation(jet..., br, 1, (@set opts_po_cont.ds = 0.005), PoincareShootingProblem(M, par_hopf, prob, Rodas4P(); abstol=1e-10, reltol=1e-9, parallel = true); normC = norminf, updateSectionEveryStep = 2, linearPO = linearPO == :autodiffMF ? :FiniteDifferencesDense : linearPO, verbosity = 0)

	br_ssh, = continuation(jet..., br, 1, (@set opts_po_cont.ds = 0.005), ShootingProblem(M, par_hopf, prob, Rodas4P(); abstol=1e-10, reltol=1e-9, parallel = true); normC = norminf, updateSectionEveryStep = 2, linearPO = linearPO, verbosity = 0)
end

br_psh, = continuation(jet..., br,1, (@set opts_po_cont.ds = 0.005), PoincareShootingProblem(2, par_hopf, prob, KenCarp4(), probMono, KenCarp4(); abstol=1e-10, reltol=1e-9); normC = norminf)

# plot(br_hpsh, vars = (:step, :p))
# plot(br_hpsh)


using OrdinaryDiffEq, Test

jac_called = Ref(false)
tgrad_called = Ref(false)

function Lotka(du,u,p,t)
  du[1] = u[1] - u[1] * u[2] # REPL[7], line 3:
  du[2] = -3 * u[2] + 1 * u[1] * u[2]
  nothing
end

function Lotka_jac(J,u,p,t)
  jac_called.x = true
  J[1,1] = 1.0 - u[2]
  J[1,2] = -u[1]
  J[2,1] = 1 * u[2]
  J[2,2] = -3 + u[1]
  nothing
end

function Lotka_tgrad(grad,u,p,t)
  tgrad_called.x = true
  grad[1] = 1 * 0
  grad[2] = 1 * 0
  nothing
end

Lotka_f = ODEFunction(Lotka,jac=Lotka_jac,tgrad=Lotka_tgrad)
prob = ODEProblem(Lotka_f,ones(2),(0.0,10.0))

good_sol = solve(prob,Rosenbrock23())

@test jac_called[]
@test tgrad_called[]

prob2 = ODEProblem(Lotka,ones(2),(0.0,10.0))

sol = solve(prob2,Rosenbrock23(autodiff=true))
@test ≈(good_sol[end],sol[end],rtol=1e-2)

sol = solve(prob2,Rosenbrock23(autodiff=true,chunk_size=1))
@test ≈(good_sol[end],sol[end],rtol=1e-2)

sol = solve(prob2,Rosenbrock23(autodiff=false))
@test ≈(good_sol[end],sol[end],rtol=1e-2)


if basename(pwd()) == "aoc"
    cd("2016/2")
end

offset(c::Char) =
    if c == 'U'
        CartesianIndex(-1, 0)
    elseif c == 'D'
        CartesianIndex(1, 0)
    elseif c == 'L'
        CartesianIndex(0, -1)
    elseif c == 'R'
        CartesianIndex(0, 1)
    end

function part1(filename::AbstractString)
    keypad = [1 2 3; 4 5 6; 7 8 9]
    position = CartesianIndex(2, 2)
    result = Int[]
    for line in eachline(filename)
        for c in line
            newposition = position + offset(c)
            if checkbounds(Bool, keypad, newposition)
                position = newposition
            end
        end
        push!(result, keypad[position])
    end
    foldl((a, b) -> 10a + b, result, init = 0)
end

@assert part1("example.txt") == 1985
part1("input.txt")

function part2(filename::AbstractString)
    keypad = [
        '0' '0' '1' '0' '0'
        '0' '2' '3' '4' '0'
        '5' '6' '7' '8' '9'
        '0' 'A' 'B' 'C' '0'
        '0' '0' 'D' '0' '0'
    ]
    position = CartesianIndex(3, 1)
    result = ""
    for line in eachline(filename)
        for c in line
            newposition = position + offset(c)
            if checkbounds(Bool, keypad, newposition) && keypad[newposition] != '0'
                position = newposition
            end
        end
        result *= keypad[position]
    end
    result
end

@assert part2("example.txt") == "5DB3"
part2("input.txt")


# This is a Julia version of Solution of the optimal control problem
# based on code written by Andrea Walther. See:
# Walther, Andrea, and Narayanan, Sri Hari Krishna. Extending the Binomial Checkpointing
# Technique for Resilience. United States: N. p., 2016. https://www.osti.gov/biblio/1364654.

using Checkpointing
using Zygote

include("optcontrolfunc.jl")

function header()
        println("**************************************************************************")
        println("*              Solution of the optimal control problem                   *")
        println("*                                                                        *")
        println("*                     J(y) = y_2(1) -> min                               *")
        println("*           s.t.   dy_1/dt = 0.5*y_1(t) + u(t),            y_1(0)=1      *")
        println("*                  dy_2/dt = y_1(t)^2 + 0.5*u(t)^2         y_2(0)=0      *")
        println("*                                                                        *")
        println("*                  the adjoints equations fulfill                        *")
        println("*                                                                        *")
        println("*         dl_1/dt = -0.5*l_1(t) - 2*y_1(t)*l_2(t)          l_1(1)=0      *")
        println("*         dl_2/dt = 0                                      l_2(1)=1      *")
        println("*                                                                        *")
        println("*   with Revolve for Online and (Multi-Stage) Offline Checkpointing      *")
        println("*                                                                        *")
        println("**************************************************************************")

        println("**************************************************************************")
        println("*        The solution of the optimal control problem above is            *")
        println("*                                                                        *")
        println("*        y_1*(t) = (2*e^(3t)+e^3)/(e^(3t/2)*(2+e^3))                     *")
        println("*        y_2*(t) = (2*e^(3t)-e^(6-3t)-2+e^6)/((2+e^3)^2)                 *")
        println("*          u*(t) = (2*e^(3t)-e^3)/(e^(3t/2)*(2+e^3))                     *")
        println("*        l_1*(t) = (2*e^(3-t)-2*e^(2t))/(e^(t/2)*(2+e^3))                *")
        println("*        l_2*(t) = 1                                                     *")
        println("*                                                                        *")
        println("**************************************************************************")

        return
end

function muoptcontrol(scheme, steps)
    println( "\n STEPS    -> number of time steps to perform")
    println("SNAPS    -> number of checkpoints")
    println("INFO = 1 -> calculate only approximate solution")
    println("INFO = 2 -> calculate approximate solution + takeshots")
    println("INFO = 3 -> calculate approximate solution + all information ")
    println(" ENTER:   STEPS, SNAPS, INFO \n")

    # F   : output
    # F_H : input
    # L   : seed the output adjoint
    # L_H : set input adjoint to 0
    F = [1.0, 0.0]
    F_H = [0.0, 0.0]
    t = 0.0
    h = 1.0/steps
    model = Model(F, F_H, t, h)
    # Just make sure it's all zero.
    shadowmodel = Model([0.0,0.0], [0.0,0.0], 0.0, 0.0)

    function foo(model::Model)
        @checkpoint_struct scheme model shadowmodel for i in 1:steps
            model.F_H .= model.F
            advance(model)
            model.t += h
        end
        return model.F[2]
    end
    g = Zygote.gradient(foo,model)

    F = model.F
    L = [g[1].F[1], g[1].F[2]]

    F_opt = Array{Float64, 1}(undef, 2)
    L_opt = Array{Float64, 1}(undef, 2)
    opt_sol(F_opt,1.0)
    opt_lambda(L_opt,0.0)
    println("\n\n")
    println("y_1*(1)  = " , F_opt[1] , " y_2*(1)  = " , F_opt[2])
    println("y_1 (1)  = " , F[1] , "  y_2 (1)  = " , F[2] , " \n\n")
    println("l_1*(0)  = " , L_opt[1] , "  l_2*(0)  = " , L_opt[2])
    println("l_1 (0)  = " , L[1]     , "  sl_2 (0)  = " , L[2] , " ")
    return F, L, F_opt, L_opt
end

# main(10,3,3)


export SuperpositionMeasure

"""
    struct SuperpositionMeasure{X,NT} <: AbstractMeasure
        components :: NT
    end

Superposition of measures is analogous to mixture distributions, but (because
measures need not be normalized) requires no scaling.

The superposition of two measures μ and ν can be more concisely written as μ + ν.

Superposition measures satisfy

    basemeasure(μ + ν) == basemeasure(μ) + basemeasure(ν)
"""
struct SuperpositionMeasure{NT} <: AbstractMeasure
    components :: NT
end

# SuperpositionMeasure(ms :: AbstractMeasure...) = SuperpositionMeasure{X,length(ms)}(ms)

# SuperpositionMeasure(m::NTuple{N, Measure{X}}) where {N,X} = SuperpositionMeasure(m...)

# Base.length(::SuperpositionMeasure{X,N}) where {X,N} = N

# function Base.:+(μ::SuperpositionMeasure{N1}, ν::SuperpositionMeasure{N2}) where {N1,N2}
#     components = (μ.components..., ν.components...)
#     SuperpositionMeasure{X, N1+N2}(components)
# end

# function Base.:+(μ::AbstractMeasure, ν::SuperpositionMeasure{X,N}) where {X,N}
#     components = (μ, ν.components...)
#     SuperpositionMeasure{X,N+1}(components)
# end

# function Base.:+(μ::SuperpositionMeasure{X,N}, ν::AbstractMeasure) where {X,N}
#     components = (μ.components..., ν)
#     SuperpositionMeasure{X,N+1}(components)
# end

function Base.:+(μ::AbstractMeasure, ν::AbstractMeasure)
    components = (μ, ν)
    SuperpositionMeasure(components)
end

logdensity(μ::SuperpositionMeasure, x) = logsumexp((logdensity(m,x) for m in μ.components))

basemeasure(μ::SuperpositionMeasure) = SuperpositionMeasure(map(basemeasure, μ.components))

# TODO: Fix `rand` method (this one is wrong)
# function Base.rand(μ::SuperpositionMeasure{X,N}) where {X,N}
#     return rand(rand(μ.components))
# end


# Arithmetic

@simple_unary(Base.abs)
@simple_unary(Base.exp)
@simple_unary(Base.log)
@simple_unary(Base.log10)
@simple_unary(Base.log2)
@simple_unary(Base.sqrt)
@simple_unary(Base.cbrt)
@simple_unary(Base.sign)
@simple_unary_self_inverse(Base.:-)
@simple_unary_self_inverse(Base.inv)

@simple_binary(Base.:+)
@simple_binary(Base.:-)
@simple_binary(Base.:*)
@simple_binary(Base.:/)
@simple_binary(Base.:^)
@simple_binary(Base.min)
@simple_binary(Base.max)

# Logical

@simple_unary_self_inverse(Base.:!)

@simple_binary(Base.:>)
@simple_binary(Base.:<)
@simple_binary(Base.:>=)
@simple_binary(Base.:<=)

# Linear algebra

@simple_binary(LinearAlgebra.dot)


using JuMP, Ipopt
import Base:show,gamma
export DemonCooling, t, gamma, spec, steps, tol, hamiltonian, timer, probability, cooling!
const MAX_COOLING_FAILURE_COUNT = 200
cooling_module_id = 1

type DemonCooling{N}<:AQCprocessor{N}
    t::Real
    gamma::Real
    H::AbstractMatrix
    k::Int # k steps
    tol::Int # tolerance for error correction, at least $tol zeros
    optimize::Int
    failure_count::Int
end

t(dc::DemonCooling) = dc.t
gamma(dc::DemonCooling) = dc.gamma
spec(dc::DemonCooling) = dc.spectrum
steps(dc::DemonCooling) = dc.k
tol(dc::DemonCooling) = dc.tol
timer(dc::DemonCooling) = dc.t

function DemonCooling(H::AbstractSparseMatrix; k::Int=5, tol::Int=2, optimize=:off)
    if optimize == :static
        t,gamma_ = static_opt_parameters(H)
        return DemonCooling{Int(log2(size(H)[1]))}(t,gamma_,H,k,tol,0,0)
    elseif optimize == :off
        maxEigen = eigs(H;nev=1,which=:LR)[1][1]|>real
        t = π/(2*maxEigen)
        gamma_ = 0
        return DemonCooling{Int(log2(size(H)[1]))}(t,gamma_,H,k,tol,1,0)
    elseif optimize == :runtime
        return DemonCooling{Int(log2(size(H)[1]))}(0.0,0.0,H,k,tol,2,0)
    end
end

function static_opt_parameters(H::AbstractSparseMatrix)
    eigens = eigfact(full(H))
    spec = real(eigens[:values])
    val,idx = findmin(spec)
    m = Model()
    @variable(m, t)
    @variable(m, gamma_)
    @NLobjective(m, Min, exp(sum(1-2*sin(e*t-gamma_) for e in spec)/(1-2*sin(val*t-gamma_))))
    solve(m)
    return getvalue(t),getvalue(gamma_)
end

function multi_cooling_obj(t,gamma_,e,k,j,p)
    return binomial(k,j)*p*(0.5-sin(e*t-gamma_)^j*(0.5+sin(e*t-gamma_))^(k-j))
end

function ∇multi_cooling_obj(g,t,gamma_,e,k,j,p)
    phi = e*t-gamma_
    g[1] = e*(k-j)*cos(phi)*(0.5-sin(phi))^j*(0.5+sin(phi))^(k-j-1)-e*j*cos(phi)*(0.5-sin(phi))^(j-1)*(0.5+sin(phi))^(k-j)

    g[2] = (j-k)*cos(phi)*(0.5-sin(phi))^j*(0.5+sin(phi))^(k-j-1)+j*cos(phi)*(0.5-sin(phi))^(j-1)*(0.5+sin(phi))^(-j+k)
end

function cooling_obj_register(H::AbstractSparseMatrix,k::Int, tol::Int, state::AbstractVector)
    eigens = eigfact(full(H))
    spec = real(eigens[:values])
    val,idx = findmin(spec)
    vecs = eigens[:vectors]
    p = abs2(state.'*eigens[:vectors])
    ground = vecs[:,idx]
    obj(x,y) = sum(sum(
            multi_cooling_obj(x,y,spec[i],k,j,p[i])*multi_cooling_obj(x,y,spec[idx],k,j,p[idx])
            for i in eachindex(spec)) for j =tol:k)
    ∇obj(g,x,y) = sum(sum(
                        ∇multi_cooling_obj(g,x,y,spec[i],k,j,p[i])*multi_cooling_obj(x,y,spec[idx],k,j,p[idx])+∇multi_cooling_obj(g,x,y,spec[idx],k,j,p[idx])*multi_cooling_obj(x,y,spec[i],k,j,p[i])
                        for i in eachindex(spec)) for j =tol:k)
    cooling_obj_name = Symbol("cool_obj$(cooling_module_id)")
    global cooling_module_id += 1
    JuMP.register(cooling_obj_name,2,obj,∇obj)
    return cooling_obj_name
end

function runtime_opt_parameters(dc::DemonCooling,aqc::AQC)
    cooling_obj_name = cooling_obj_register(dc.H,dc.k,dc.tol,coeff(aqc.state))
    m = Model()
    spec = real(eigvals(full(dc.H)))
    minEigen, maxEigen = minimum(spec), maximum(spec)
    @variable(m, t_)
    @variable(m, gamma_)
    @constraint(m, -pi<=minEigen*t_-gamma_<=pi)
    @constraint(m, -pi<=maxEigen*t_-gamma_<=pi)
    @NLobjective(m, Max, $(cooling_obj_name)(t_,gamma_))
    solve(m)
    return getvalue(t_),getvalue(gamma_)
end

DemonCooling(H::AbstractMatrix;k::Int=5, tol::Int=2) = DemonCooling(sparse(H);k=k,tol=tol)

function show{N}(io::IO,dc::DemonCooling{N})
    println(io,"----------------------------")
    println(io,"$N qubits demon-like cooling")
    println(io,"----------------------------")
    println(io,"t = $(dc.t)")
    println(io,"gamma = $(dc.gamma)")
    println(io,"for $(dc.k) steps")
    println(io,"tolerence: larger than $(dc.tol) zeros")
end

function run!{N}(dc::DemonCooling{N}, aqc::AQC)
    if dc.optimize == 2
        dc.t,dc.gamma = runtime_opt_parameters(dc,aqc)
    end
    ret=cooling!(aqc.state, dc)
    run_t = 0.0
    if ret==false
        dc.failure_count += 1
        aqc.state = BellState(N)
        run_t = run!(aqc)
        # aqc.process_list = append!(aqc.process_list[1:aqc.status],aqc.process_list)
        # aqc.state = BellState(N)
        return run_t+t(dc)
    elseif ret
        return t(dc)
    else
    end
end

# return a list of probability for each
# reference: README
function probability{V<:AbstractVector,T,N}(
    state::ComputState{V,T,N},
    dc::DemonCooling{N}
    )
    ret = zeros(Float64,dc.k+1)
    eigens = eigfact(full(dc.H))
    spec = real(eigens[:values])
    A = spec
    map!(x->0.5*sin(x*dc.t-dc.gamma), A)
    p = coeff(state).'*eigens[:vectors]
    for j = 0:dc.k
        for i = 1:2^N
            ret[j+1] += binomial(dc.k,j)*abs2(p[i])*(0.5-A[i])^j*(0.5+A[i])^(dc.k-j)
        end
    end
    return ret
end

# return a cooled state
function cooling!{V<:AbstractVector,T,N}(
    state::ComputState{V,T,N},
    dc::DemonCooling{N}
    )
    p = probability(state, dc)
    k = steps(dc)
    eigens = eigfact(full(dc.H))
    spec = real(eigens[:values])
    A = spec
    map!(x->0.5*sin(x*dc.t-dc.gamma), A)

    dice = rand()
    if dice<sum(p[1:tol(dc)])
        return false
    end
    for j=tol(dc)+1:k
        if dice<sum(p[1:j+1])
            state.coeffs = 0.5^k*binomial(k,j)*(I-im*exp(im*dc.gamma)*expm(-im*full(dc.H)*dc.t))^j*(I+im*exp(im*dc.gamma)*expm(-im*full(dc.H)*dc.t))^(k-j)*state.coeffs
            normalize!(state.coeffs)
            return true
        end
    end
    return false
end


"""
Lennard Jones equation of state

"""
function LJSinglephase(T,ro)
	x=[0.8623085097507421; 2.976218765822098; -8.402230115796038; 0.1054136629203555; -0.8564583828174598;
1.582759470107601; 0.7639421948305453; 1.753173414312048; 2.798291772190376e03; -4.8394220260857657e-02;
0.9963265197721935; -3.698000291272493e01; 2.084012299434647e01; 8.305402124717285e01; -9.574799715203068e02;
-1.477746229234994e02; 6.398607852471505e01; 1.603993673294834e01; 6.805916615864377e01; -2.791293578795945e03;
-6.245128304568454; -8.116836104958410e03; 1.488735559561229e01; -1.059346754655084e04; -1.131607632802822e02;
-8.867771540418822e03; -3.986982844450543e01; -4.689270299917261e03; 2.593535277438717e02; -2.6945235894349033e03;
-7.218487631550215e02; 1.721802063863269e02;]

	a=Array{Float64}(undef, 8)
	b=Array{Float64}(undef, 6)
	G=Array{Float64}(undef, 6)

	a[1]=x[1]*T+x[2]*sqrt(T)+x[3]+x[4]/T+x[5]/(T^2);
	a[2]=x[6]*T+x[7]+x[8]/T+x[9]/(T^2);
	a[3]=x[10]*T+x[11]+x[12]/T;
	a[4]=x[13];
	a[5]=x[14]/T+x[15]/T^2;
	a[6]=x[16]/T;
	a[7]=x[17]/T+x[18]/T^2;
	a[8]=x[19]/T^2;

	b[1]=x[20]/T^2+x[21]/T^3;
	b[2]=x[22]/T^2+x[23]/T^4;
	b[3]=x[24]/T^2+x[25]/T^3;
	b[4]=x[26]/T^2+x[27]/T^4;
	b[5]=x[28]/T^2+x[29]/T^3;
	b[6]=x[30]/T^2+x[31]/T^3+x[32]/T^4;

	gm=3.0
	F=exp(-gm*ro*ro)
	G[1]=(1.0-F)/(2*gm)
	G[2]=-(F*ro^2-2*G[1])/(2*gm)
	G[3]=-(F*ro^4-4*G[2])/(2*gm)
	G[4]=-(F*ro^6-6*G[3])/(2*gm)
	G[5]=-(F*ro^8-8*G[4])/(2*gm)
	G[6]=-(F*ro^10-10*G[5])/(2*gm)

	pres=ro*T+a[1]*ro^2+a[2]*ro^3+a[3]*ro^4+a[4]*ro^5+a[5]*ro^6+a[6]*ro^7+a[7]*ro^8+a[8]*ro^9+F*(b[1]*ro^3+b[2]*ro^5+b[3]*ro^7+b[4]*ro^9+b[5]*ro^11+b[6]*ro^13)
	Ar=a[1]*ro^1+a[2]*ro^2/2.0+a[3]*ro^3/3.0+a[4]*ro^4/4.0+a[5]*ro^5/5.0+a[6]*ro^6/6.0+a[7]*ro^7/7.0+a[8]*ro^8/8.0+b[1]*G[1]+b[2]*G[2]+b[3]*G[3]+b[4]*G[4]+b[5]*G[5]+b[6]*G[6]
	Ge=Ar+pres/ro-T
	dmudn=a[1]+a[2]*ro+a[3]*ro^2+a[4]*ro^3+a[5]*ro^4+a[6]*ro^5+a[7]*ro^6+a[8]*ro^7+F*(b[1]*ro+b[2]*ro^3+b[3]*ro^5+b[4]*ro^7+b[5]*ro^9+b[6]*ro^11)
	dmudn+=a[1]+a[2]*2.0*ro+a[3]*3.0*ro^2+a[4]*4.0*ro^3+a[5]*5.0*ro^4+a[6]*6.0*ro^5+a[7]*7.0*ro^6+a[8]*8.0*ro^7-F*2*gm*(b[1]*ro^3+b[2]*ro^5+b[3]*ro^7+b[4]*ro^9+b[5]*ro^11+b[6]*ro^13)+2*F*(b[1]*ro+b[2]*2.0*ro^3+b[3]*3.0*ro^5+b[4]*4.0*ro^7+b[5]*5.0*ro^9+b[6]*6.0*ro^11)

	return [pres; Ar; Ge; dmudn]
end


using FountainCodes
using Random, StatsBase, Distributions
using LinearAlgebra, SparseArrays

function r10_example(nsrc=10, nenc=20)
    code = R10(nsrc)
    println("Example for $code w. $nsrc source symbols and $nenc coded symbols")
    src = rand(GF256, nsrc)
    println("source symbols:")
    println(src)
    println()

    # compute the intermediate symbols
    A = constraint_matrix(code, 0:nsrc-1)
    pre = zeros(GF256, size(A, 2))
    pre[end-nsrc+1:end] .= src
    int = decode(A, pre)
    println("intermediate symbols:")
    println(int)
    println()

    # randomly selected nenc unique encodes symbol identifiers, corresponding to the nenc encoded
    # symbols received by the decoder (the upper limit is chosen arbitrarily)
    Xs = sample(0:1000, nenc, replace=false)

    # compute the encoded symbols
    G = generator_matrix(code, Xs)
    enc = G'*int
    println("encoded symbols:")
    println(enc)
    println()

    # decode the intermediate symbols from the encoded symbols
    A = constraint_matrix(code, Xs)
    pre = zeros(GF256, size(A, 2))
    pre[end-nenc+1:end] .= enc
    dec_int = decode(A, pre)
    println("decoded intermediate symbols:")
    println(dec_int)
    println()

    # generate the source symbols from the decoded intermediate symbols
    G = generator_matrix(code, 0:nsrc-1)
    dec_src = G'*dec_int
    println("decoded source symbols:")
    println(dec_src)
end

function rq_example(nsrc=10, nenc=20)
    code = RQ(nsrc)
    nsrc = code.Kp # w. zero-padding up to the next supported number of source symbols
    println("Example for $code w. $nsrc source symbols and $nenc coded symbols")
    src = rand(GF256, nsrc)
    println("source symbols:")
    println(src)
    println()

    # compute the intermediate symbols
    A = constraint_matrix(code, 0:nsrc-1)
    pre = zeros(GF256, size(A, 2))
    pre[end-nsrc+1:end] .= src
    int = decode(A, pre)
    println("intermediate symbols:")
    println(int)
    println()

    # randomly selected nenc unique encodes symbol identifiers, corresponding to the nenc encoded
    # symbols received by the decoder (the upper limit is chosen arbitrarily)
    Xs = sample(0:1000, nenc, replace=false)

    # compute the encoded symbols
    G = generator_matrix(code, Xs)
    enc = G'*int
    println("encoded symbols:")
    println(enc)
    println()

    # decode the intermediate symbols from the encoded symbols
    A = constraint_matrix(code, Xs)
    pre = zeros(GF256, size(A, 2))
    pre[end-nenc+1:end] .= enc
    dec_int = decode(A, pre)
    println("decoded intermediate symbols:")
    println(dec_int)
    println()

    # generate the source symbols from the decoded intermediate symbols
    G = generator_matrix(code, 0:nsrc-1)
    dec_src = G'*dec_int
    println("decoded source symbols:")
    println(dec_src)
end

function lt_example(nsrc=10, nenc=20)
    dd = Soliton(nsrc, nsrc, 1e-6)
    code = LT{GF256}(nsrc, dd) # w. generator matrix coefficients of type GF256
    println("Example for $code w. $nsrc source symbols and $nenc coded symbols")
    src = rand(GF256, nsrc)
    println("source symbols:")
    println(src)
    println()

    # randomly selected nenc unique encodes symbol identifiers, corresponding to the nenc encoded
    # symbols received by the decoder (the upper limit is chosen arbitrarily)
    Xs = sample(0:1000, nenc, replace=false)

    # compute the encoded symbols
    G = generator_matrix(code, Xs)
    enc = G'*src
    println("encoded symbols:")
    println(enc)
    println()

    # decode the source symbols from the encoded symbols
    dec = decode(G, enc)
    println("decoded intermediate symbols:")
    println(dec)
end

using DelimitedFiles

function ldpc_example(nerasures=20)
    # read the parity check matrix
    # (transposed, since the decoder expects columns to correspond to constraints and rows to symbols)
    # despite it being a binary matrix, we're using symbols of type GF256 since there are no
    # optimizations specific to the binary field
    H = SparseMatrixCSC{GF256,Int}(readdlm("./examples/H_612_1224.txt")')
    # H = SparseMatrixCSC{GF256,Int}(readdlm("./examples/H_2400_4800.txt")')
    nenc, nconstraints = size(H)

    # since we do not consider encoding, we use the all-zeros vector, which is a valid codeword in
    # any linear code
    enc = zeros(GF256, nenc)

    # randomly select nerasures symbols to erase
    Xs = collect(1:nenc)
    shuffle!(Xs)
    Xs_rec = sort!(Xs[1:nenc-nerasures])
    Xs_era = sort!(Xs[end-nerasures+1:end])
    H_rec = H[Xs_rec, :] # constraints corresponding to received symbols
    H_era = H[Xs_era, :] # constraints corresponding to erased symbols
    enc_rec = enc[Xs_rec]
    enc_era = enc[Xs_era]
    println("erased encoded symbols:")
    println(enc_era)
    println()

    # note that H'*enc = H_rec'*enc_rec + H_era'*enc_era = 0, i.e., we can solve for enc_era from
    # the system of equations H_era'*enc_era = - H_rec'*enc_rec = H_rec'*enc_rec,
    # where the second equality is due to the fact that -GF256(x) = GF256(x)
    # finally, assuming that the code is systematic, the source symbols can be recovered from the
    # decoded enc vector by selecting the subset of its indices corresponding to the systematic entries
    rhs = H_rec'*enc_rec
    enc_era_dec = decode(H_era, rhs)
    println("decoded erased symbols:")
    println(enc_era_dec)
end


function simple_toy_MetNet()
    net = Dict()
    net[:S] =
    # rxns: gt    ferm  resp  ldh   lt   biom    atpm  # mets
        [   1.0  -1.0   0.0   0.0   0.0   0.0    0.0;  #  G
            0.0   2.0  18.0   0.0   0.0  -55.0  -5.0;  #  E
            0.0   2.0  -1.0  -1.0   0.0   0.0    0.0;  #  P
            0.0   0.0   0.0   1.0   1.0   0.0    0.0;  #  L
        ]

    net[:mets] = ["G", "E", "P", "L"]
    net[:b] =    [0.0, 0.0, 0.0, 0.0] # demand

    net[:metNames] = ["Glucose", "Energy", "Intermediate Product" , "Lactate"];

    net[:rxns] = ["gt"  ,"ferm" ,"resp" , "ldh" ,  "lt" , "biom", "atpm"];
    net[:lb] =   [0.0   , 0.0   , 0.0   ,  0.0  , -100.0,   0.0,     0.5];
    net[:ub] =   [10.0 , 100.0 , 100.0 , 100.0 ,    0.0, 100.0,    100.0];
    net[:c] =    [ 0.0 ,   0.0 ,   0.0 ,   0.0 ,    0.0,   0.0,      0.0];
    net[:rxnNames] = ["Glucose transport", "Fermentation", "Respiration",
        "Lactate DH", "Lactate transport", "Biomass production rate", "atp demand"];

    return MetNet(;net...)
end
const TOY_MODEL_BIOMASS_IDER = "biom"
const TOY_MODEL_COST_IDER = "tot_cost"

function  toy_model(;resp_cost::Real = -0.1, E_demand::Real = 5.0)
    @assert resp_cost <= 0.0

    model = simple_toy_MetNet()

    # Atpm
    lb!(model, "atpm", E_demand)

    # cost
    # Add a new metabolite simulating the cost penalazing
    # reaction fluxes.
    # A new balance equations is then added:
    #     Σ(rᵢ*costᵢ) + tot_cost = 0
    # Because the cost coefficients (costᵢ) < 0, the system must allocate
    # the fluxes (rᵢ) so that Σ(rᵢ*costᵢ) = tot_cost, and tot_cost
    # are usually bounded [0.0, 1.0]

    # We will add 1 met and 1 rxn
    M, N = size(model)
    model = expanded_model(model, M + 1, N + 1)
    # Add cost met
    set_met!(model, M + 1, Met("cost"; rxns = ["resp"], S = [resp_cost]))
    set_rxn!(model, N + 1, Rxn("tot_cost"; mets = ["cost"], S = [1.0], lb = 0.0, ub = 1.0))
    return model
end


using JuMP, Mosek, MosekTools, LinearAlgebra, Suppressor, StatsBase#, SCS#, Ipopt

# Solve the perspective relaxation (naive but cheaper formulation)
function solvepersprelax(X, y, m, n, p, lambda, gamma)
    mod=Model(Mosek.Optimizer)

    @variable(mod, beta[1:p, 1:n])
    @variable(mod, W[1:n, 1:n], Symmetric)
    @variable(mod, theta[1:p, 1:p], Symmetric)
    @variable(mod, t) # Epigraph variable for rotated SOC constraint

    @constraint(mod, Symmetric(Matrix(1.0I, n, n)-W) in PSDCone());

    @constraint(mod, Symmetric([theta beta; beta' W]) in PSDCone());

    @constraint(mod, [t; 1.0; vec(y-X*beta)] in RotatedSecondOrderCone());

    @objective(mod, Min, (0.5/m)*t+0.5*1.0/(2.0*gamma)*sum(theta[i,i] for i=1:p)+lambda*sum(W[i,i] for i=1:n))

    @suppress optimize!(mod)

    ofv_dual=objective_value(mod)
    u,v,=svd(value.(W))
    # Use a cutoff of 10^-4 for singular values, to avoid e.g. Julia counting singular values of 10^-12 as part of the rank
    ofv_primal=(0.5/m)*LinearAlgebra.dot(y.-X*value.(beta), y.-X*value.(beta))+1.0/(2.0*gamma)*norm(value.(beta))^2+lambda*sum(abs.(v).>1e-4)
    # Todo: round greedily and resolve second stage problem??

    return value.(beta), ofv_dual, ofv_primal, value.(W)
end

# Solve the perspective relaxation (expensive DCL version)
function solvepersprelax_DCL(X, y, m, n, p, lambda, gamma)
    mod=Model(Mosek.Optimizer)

    @variable(mod, beta[1:p, 1:n])
    @variable(mod, B[1:p, 1:p], Symmetric)
    @variable(mod, W[1:n, 1:n], Symmetric)

    @constraint(mod, Symmetric(Matrix(1.0I, n, n)-W) in PSDCone());

    @constraint(mod, Symmetric([B beta; beta' W]) in PSDCone());

    @objective(mod, Min, (0.5/m)*sum(y[i,j]^2 for i=1:m for j=1:n)-(1.0/m)*LinearAlgebra.dot(y, X*beta)+0.5*LinearAlgebra.dot(B, 1.0/gamma*Matrix(1.0I, p, p)+(1.0/m)X'*X)+lambda*sum(W[i,i] for i=1:n))

    @suppress optimize!(mod)

    @show ofv_dual=objective_value(mod)

    u,v,=svd(value.(W))
    # Use a cutoff of 10^-4 for singular values, to avoid e.g. Julia counting singular values of 10^-12 as part of the rank
    ofv_primal=(0.5/m)*LinearAlgebra.dot(y.-X*value.(beta), y.-X*value.(beta))+1.0/(2.0*gamma)*norm(value.(beta))^2+lambda*sum(abs.(v).>1e-4)

    return value.(beta), ofv_dual, ofv_primal, value.(W)
end

# Solve the nuclear norm "relaxation" (not actually a relaxation without big-M constraints and projection matrices)
function solveNNrelax(X, y, m, n, p, lambda, gamma)
    mod=Model(Mosek.Optimizer)

    @variable(mod, beta[1:p, 1:n])
    @variable(mod, t) # Epigraph variable for rotated SOC constraint
    @variable(mod, U[1:p, 1:p])
    @variable(mod, V[1:n, 1:n])

    @constraint(mod, [t; 1.0; vec(y-X*beta)] in RotatedSecondOrderCone());

    @constraint(mod, Symmetric([U beta; beta' V]) in PSDCone());

    @objective(mod, Min, (0.5/m)*t+1.0/(2.0*gamma)*sum(beta[i,j]^2 for i=1:p for j=1:n)+0.5*lambda*sum(U[i,i] for i=1:p)+0.5*lambda*sum(V[i,i] for i=1:n))

    @suppress optimize!(mod)

    # Todo: round greedily and resolve second stage problem??

    return value.(beta), objective_value(mod)
end


using ScRNAseq
using DelimitedFiles
cd(dirname(@__FILE__))

X=readdlm("X.txt",',',Int16)
genelist=vec(readdlm("genelist.txt",String))

X,genelist=ScRNAseq.QualityControl.selectg(X,genelist)
X1=ScRNAseq.Transformation.pearsonresiduals(X)
X2=ScRNAseq.Normalization.norm_libsize(X)
Y=ScRNAseq.Embedding.umap(X1)

# Y2=ScRNAseq.Embedding.tsne(X)
using Plots
scatter3d(Y[:,1],Y[:,2],Y[:,3])


#This example fits a GP regression model to the Mauna Loa CO2 data set. This data set is regularly updated and can be found at co2now.org/images/stories/data/co2-mlo-monthly-noaa-esrl.xls
#
#This example follows from Chapter 5 of Gaussian Processes for Machine Learning, Rasmussen and Williams (2006)
############################################################################################

using Gadfly, GaussianProcesses

data = readcsv("CO2_data.csv")

year = data[:,1]; co2 = data[:,2];
x = year[year.<2004]; y = co2[year.<2004];
xpred = year[year.>=2004]; ypred = co2[year.>=2004];

mConst = MeanConst(mean(y))       #Fit the constant mean function

#Kernel is represented as a sum of kernels
kernel = SE(4.0,4.0) + Periodic(0.0,1.0,0.0)*SE(4.0,0.0) + RQ(0.0,0.0,-1.0) + SE(-2.0,-2.0)

gp = GP(x,y,mConst,kernel,-2.0)   #Fit the GP

plot(gp,clim=(2004.0,2024.0))  #Gadfly can take a while to load


# Tested in Julia 1.5.1
using DelimitedFiles, Combinatorics
data = readdlm("./data/adventInput01.txt", Int)
function ReportRepairPerm(data; nums = 2, sumto = 2020)
    for p in permutations(data, nums)
        if sum(p) == sumto
            print(p, " Sum: $(sum(p)),  Product : $(prod(p))")
            break
        end
    end
end
ReportRepairPerm(data)
ReportRepairPerm(data, nums = 3)


function eval_TIFₜ(xₜ::Vector{Float64}, coreₜ::MaximinObjectiveCore) # x = I
    return minimum(LinearAlgebra.BLAS.gemv('N', coreₜ.IIF, xₜ))::Float64
end

function eval_TIFₜ(xₜ::Vector{Float64}, coreₜ::CCMaximinObjectiveCore) # x = I
    IIF = coreₜ.IIF
    α = coreₜ.alpha
    K, I, R = size(IIF)
    alphaR = Int(ceil(R * (α)))
    αQle = Inf
    if K > 1
        for k = 1:K
            αQle = min(αQle, sort(LinearAlgebra.BLAS.gemv('T', IIF[k, :, :], xₜ))[alphaR])
        end
    else
        αQle = sort(_gemvblasT(IIF[1, :, :], xₜ, R))[alphaR]
    end
    return αQle::Float64
end

function eval_TIFₜ(xₜ::Vector{Float64}, coreₜ::MinimaxObjectiveCore) # x = I
    IIF = coreₜ.IIF
    targets = coreₜ.targets
    return -maximum(abs.(LinearAlgebra.BLAS.gemv('N', IIF, xₜ) - targets))::Float64
end


function outer_dGx_pos(loc_globaldata, loc_ghost_holder, dist_length, idx, gamma, phi_i, phi_k, G_i, G_k, result, qtilde_i, qtilde_k, ∑_Δx_Δf, ∑_Δy_Δf, power, limiter_flag, vl_const, Gxp)

    ∑_Δx_sqr = zero(Float64)
    ∑_Δy_sqr = zero(Float64)
    ∑_Δx_Δy = zero(Float64)

    fill!(∑_Δx_Δf, zero(Float64))
    fill!(∑_Δy_Δf, zero(Float64))

    x_i = loc_globaldata[idx].x
    y_i = loc_globaldata[idx].y

    nx = loc_globaldata[idx].nx
    ny = loc_globaldata[idx].ny

    tx = ny
    ty = -nx

    for itm in loc_globaldata[idx].xpos_conn

        if itm <= dist_length
            globaldata_itm = loc_globaldata[itm]
        else
            globaldata_itm = loc_ghost_holder[1][itm]
        end
        x_k = globaldata_itm.x
        y_k = globaldata_itm.y

        delx = x_k - x_i
        dely = y_k - y_i

        dels = delx*tx + dely*ty
        deln = delx*nx + dely*ny

        dist = hypot(dels, deln)
        weights = dist ^ power

        dels_weights = dels*weights
        deln_weights = deln*weights

        ∑_Δx_sqr = ∑_Δx_sqr + dels*dels_weights
        ∑_Δy_sqr = ∑_Δy_sqr + deln*deln_weights

        ∑_Δx_Δy = ∑_Δx_Δy + dels*deln_weights

        @. qtilde_i = loc_globaldata[idx].q - 0.5*(delx*loc_globaldata[idx].dq1 + dely*loc_globaldata[idx].dq2)
        @. qtilde_k = globaldata_itm.q - 0.5*(delx*globaldata_itm.dq1 + dely*globaldata_itm.dq2)

        if limiter_flag == 1
            venkat_limiter(qtilde_i, loc_globaldata[idx], vl_const, phi_i)
            venkat_limiter(qtilde_k, globaldata_itm, vl_const, phi_k)
            @. qtilde_i = loc_globaldata[idx].q - 0.5 * phi_i * (delx*loc_globaldata[idx].dq1 + dely*loc_globaldata[idx].dq2)
            @. qtilde_k = globaldata_itm.q - 0.5 * phi_k * (delx*globaldata_itm.dq1 + dely*globaldata_itm.dq2)
        end

        qtilde_to_primitive(result, qtilde_i, gamma)
        flux_quad_GxIII(G_i, nx, ny, result[1], result[2], result[3], result[4])

        qtilde_to_primitive(result, qtilde_k, gamma)
        flux_quad_GxIII(G_k, nx, ny, result[1], result[2], result[3], result[4])

        for i in 1:4
            ∑_Δx_Δf[i] += (G_k[i] - G_i[i]) * dels_weights
            ∑_Δy_Δf[i] += (G_k[i] - G_i[i]) * deln_weights
        end
    end

    det = ∑_Δx_sqr*∑_Δy_sqr - ∑_Δx_Δy*∑_Δx_Δy
    one_by_det = one(Float64) / det

    for iter in 1:4
        Gxp[iter] = (∑_Δx_Δf[iter] *∑_Δy_sqr - ∑_Δy_Δf[iter] *∑_Δx_Δy)*one_by_det
    end
    return nothing
end

function outer_dGx_neg(loc_globaldata, loc_ghost_holder, dist_length, idx, gamma, phi_i, phi_k, G_i, G_k, result, qtilde_i, qtilde_k, ∑_Δx_Δf, ∑_Δy_Δf, power, limiter_flag, vl_const, Gxn)

    ∑_Δx_sqr = zero(Float64)
    ∑_Δy_sqr = zero(Float64)
    ∑_Δx_Δy = zero(Float64)

    fill!(∑_Δx_Δf, zero(Float64))
    fill!(∑_Δy_Δf, zero(Float64))

    x_i = loc_globaldata[idx].x
    y_i = loc_globaldata[idx].y

    nx = loc_globaldata[idx].nx
    ny = loc_globaldata[idx].ny

    tx = ny
    ty = -nx

    for itm in loc_globaldata[idx].xneg_conn

        if itm <= dist_length
            globaldata_itm = loc_globaldata[itm]
        else
            globaldata_itm = loc_ghost_holder[1][itm]
        end
        x_k = globaldata_itm.x
        y_k = globaldata_itm.y

        delx = x_k - x_i
        dely = y_k - y_i

        dels = delx*tx + dely*ty
        deln = delx*nx + dely*ny

        dist = hypot(dels, deln)
        weights = dist ^ power

        dels_weights = dels*weights
        deln_weights = deln*weights

        ∑_Δx_sqr = ∑_Δx_sqr + dels*dels_weights
        ∑_Δy_sqr = ∑_Δy_sqr + deln*deln_weights

        ∑_Δx_Δy = ∑_Δx_Δy + dels*deln_weights

        @. qtilde_i = loc_globaldata[idx].q - 0.5*(delx*loc_globaldata[idx].dq1 + dely*loc_globaldata[idx].dq2)
        @. qtilde_k = globaldata_itm.q - 0.5*(delx*globaldata_itm.dq1 + dely*globaldata_itm.dq2)

        if limiter_flag == 1
            venkat_limiter(qtilde_i, loc_globaldata[idx], vl_const, phi_i)
            venkat_limiter(qtilde_k, globaldata_itm, vl_const, phi_k)
            @. qtilde_i = loc_globaldata[idx].q - 0.5 * phi_i * (delx*loc_globaldata[idx].dq1 + dely*loc_globaldata[idx].dq2)
            @. qtilde_k = globaldata_itm.q - 0.5 * phi_k * (delx*globaldata_itm.dq1 + dely*globaldata_itm.dq2)
        end

        qtilde_to_primitive(result, qtilde_i, gamma)
        flux_quad_GxIV(G_i, nx, ny, result[1], result[2], result[3], result[4])

        qtilde_to_primitive(result, qtilde_k, gamma)
        flux_quad_GxIV(G_k, nx, ny, result[1], result[2], result[3], result[4])

        for i in 1:4
            ∑_Δx_Δf[i] += (G_k[i] - G_i[i]) * dels_weights
            ∑_Δy_Δf[i] += (G_k[i] - G_i[i]) * deln_weights
        end
    end
    det = ∑_Δx_sqr*∑_Δy_sqr - ∑_Δx_Δy*∑_Δx_Δy
    one_by_det = one(Float64) / det
    for iter in 1:4
        Gxn[iter] = (∑_Δx_Δf[iter] * ∑_Δy_sqr - ∑_Δy_Δf[iter] * ∑_Δx_Δy)*one_by_det
    end
end

function outer_dGy_pos(loc_globaldata, loc_ghost_holder, dist_length, idx, gamma, phi_i, phi_k, G_i, G_k, result, qtilde_i, qtilde_k, ∑_Δx_Δf, ∑_Δy_Δf, power, limiter_flag, vl_const, Gyp)

    ∑_Δx_sqr = zero(Float64)
    ∑_Δy_sqr = zero(Float64)
    ∑_Δx_Δy = zero(Float64)

    fill!(∑_Δx_Δf, zero(Float64))
    fill!(∑_Δy_Δf, zero(Float64))

    x_i = loc_globaldata[idx].x
    y_i = loc_globaldata[idx].y

    nx = loc_globaldata[idx].nx
    ny = loc_globaldata[idx].ny

    tx = ny
    ty = -nx

    for itm in loc_globaldata[idx].ypos_conn
        if itm <= dist_length
            globaldata_itm = loc_globaldata[itm]
        else
            globaldata_itm = loc_ghost_holder[1][itm]
        end
        x_k = globaldata_itm.x
        y_k = globaldata_itm.y

        delx = x_k - x_i
        dely = y_k - y_i

        dels = delx*tx + dely*ty
        deln = delx*nx + dely*ny

        dist = hypot(dels, deln)
        weights = dist ^ power

        dels_weights = dels*weights
        deln_weights = deln*weights

        ∑_Δx_sqr = ∑_Δx_sqr + dels*dels_weights
        ∑_Δy_sqr = ∑_Δy_sqr + deln*deln_weights

        ∑_Δx_Δy = ∑_Δx_Δy + dels*deln_weights

        @. qtilde_i = loc_globaldata[idx].q - 0.5*(delx*loc_globaldata[idx].dq1 + dely*loc_globaldata[idx].dq2)
        @. qtilde_k = globaldata_itm.q - 0.5*(delx*globaldata_itm.dq1 + dely*globaldata_itm.dq2)

        if limiter_flag == 1
            venkat_limiter(qtilde_i, loc_globaldata[idx], vl_const, phi_i)
            venkat_limiter(qtilde_k, globaldata_itm, vl_const, phi_k)
            @. qtilde_i = loc_globaldata[idx].q - 0.5 * phi_i * (delx*loc_globaldata[idx].dq1 + dely*loc_globaldata[idx].dq2)
            @. qtilde_k = globaldata_itm.q - 0.5 * phi_k * (delx*globaldata_itm.dq1 + dely*globaldata_itm.dq2)
        end

        qtilde_to_primitive(result, qtilde_i, gamma)
        flux_Gyp(G_i, nx, ny, result[1], result[2], result[3], result[4])

        qtilde_to_primitive(result, qtilde_k, gamma)
        flux_Gyp(G_k, nx, ny, result[1], result[2], result[3], result[4])

        for i in 1:4
            ∑_Δx_Δf[i] += (G_k[i] - G_i[i]) * dels_weights
            ∑_Δy_Δf[i] += (G_k[i] - G_i[i]) * deln_weights
        end

    end
    det = ∑_Δx_sqr*∑_Δy_sqr - ∑_Δx_Δy*∑_Δx_Δy
    one_by_det = one(Float64) / det
    for iter in 1:4
        Gyp[iter] = (∑_Δy_Δf[iter] * ∑_Δx_sqr - ∑_Δx_Δf[iter] * ∑_Δx_Δy)*one_by_det
    end
end


import Dates

s = Dates.Date(1901, 1, 1)
c = s
e = Dates.Date(2000, 12, 31)
d = Dates.Day(1)

counter = 0
while c != e
  global c, counter
  if Dates.day(c) == 1 && Dates.dayofweek(c) == 7
    counter += 1
  end
  c += d
end
println(counter)


type DirichiletLift

    N::Int
    nd::Int
    nv::Int
    idx::Array{Int,1}
    vidx::Array{Int,1}
    A::Array{Float64,2}
    function DirichiletLift(n, is_dirichilet)
        nd = 0
        for i = 1:n
            nd = is_dirichilet[i]?nd+1:nd
        end
        nv = n - nd
        idx = zeros(Int, nd)
        vidx = zeros(Int, n-nd)
        icnt = 1
        ivar = 1
        for i = 1:n
            if is_dirichilet[i]
                idx[icnt] = i
                icnt += 1
            else
                vidx[ivar] = i
                ivar += 1
            end
        end

        A = zeros(nv, nd)

        new(n, nd, nv, idx, vidx, A)
    end

end

function set_dirichilet_matrix(dlift::DirichiletLift, M)
    A = dlift.A
    N = dlift.N
    nd = dlift.nd
    nv = dlift.nv
    idx = dlift.idx
    vidx = dlift.vidx

    for i = 1:nv
        iv = vidx[i]
        for k = 1:nd
            A[i,k] = M[iv, idx[k]]
        end
    end

end

function lift(dlift::DirichiletLift, f)

    idx = dlift.idx
    nv = dlift.nv
    nd = dlift.nd
    vidx = dlift.vidx
    A = dlift.A

    for i = 1:nv
        iv = vidx[i]
        for k = 1:nd
            f[iv] -= A[i,k] * f[idx[k]]
        end
    end
end

function lift(dlift::DirichiletLift, f, u)

    idx = dlift.idx
    nv = dlift.nv
    nd = dlift.nd
    vidx = dlift.vidx
    A = dlift.A

    for i = 1:nv
        iv = vidx[i]
        for k = 1:nd
            f[iv] -= A[i,k] * u[idx[k]]
        end
    end
end


# # Landcover data

# In this example, we will look at landcover data in Corsica

using SimpleSDMLayers
using Plots

# To avoid loading the (rather large) dataset of land cover at once, we will
# rely on a bounding box:

boundaries = (left=8.25, right=10.0, bottom=41.2, top=43.2)

# The numbers corresponding to the layers can be found in the documentation for
# `SimpleSDMPredictor` - urban area is 9.

urban = SimpleSDMPredictor(EarthEnv, LandCover, 9; boundaries...)

# This dataset is returning data as `UInt8` (as it represents a proportion of
# the pixel occupied by the type), but this is not something that can be plotted
# efficiently (because we rely on `NaN` to indicate no data, and `UInt8` has no
# `NaN`). So in the next step, we will manipulate this object a little bit to
# have something more workable.

urban = convert(Float16, urban)

# Why `Float16`? It is the smallest floating point type with a `NaN` value
# (`NaN16`). We will replace the values of 0 by `nothing`, to only see the
# pixels with *at least* some urban cover:

replace!(urban, zero(eltype(urban)) => nothing)

# With this done, we can plot the results:

plot(urban; c=:berlin, clim=(0, 100))

# Unsuprisingly to anyone who had the chance to visit Corsica, it is not a very
# densely urbanized island. This is a good time to question whether we can look
# at (i) which landcover type dominates within each pixel, and (ii) how
# heterogeneous the land use within each pixel is.

# First, we will download all values for the landcover layers, including open
# water (12).

landcover =
    convert.(
        Float16, SimpleSDMPredictor(EarthEnv, LandCover, 1:12; full=false, boundaries...)
    )

plot(
    plot.(landcover, leg=false, c=:oleron, clim=(0, 100))...;
    leg=false,
    grid=:none,
    frame=:none,
)

# To perform the actual analysis, we will define a `shannon` function, which
# will return the entropy of the land use categories:

function shannon(x)
    v = filter(n -> n > zero(eltype(x)), x)
    length(v) == 0 && return NaN
    v = v ./ sum(v)
    return -sum(v .* log2.(v))
end

entropy = mosaic(shannon, landcover)

# We can also get the index of the most common layer within the pixel:

consensus = mosaic(x -> last(findmax(x)), landcover)

# We may not be *that* interested in fully open water, so let's define a mask to
# remove it:

openwater = broadcast(!isequal(12), consensus)

# All that's left to do is to plot after applying this mask, and we now get a
# map of most common land cover type (left), and land cover heterogeneity
# (right)

p1 = plot(mask(openwater, consensus); c=:terrain, frame=:none)
p2 = plot(mask(openwater, entropy); c=:lapaz, frame=:none)
plot(p1, p2)


##############################################################################
#
# TF
#
##############################################################################

tf(dtm::Matrix{T}) where {T <: Real} = tf!(dtm, Array{Float64}(undef, size(dtm)...))

tf(dtm::SparseMatrixCSC{T}) where {T <: Real} =  tf!(dtm, similar(dtm, Float64))

tf!(dtm::AbstractMatrix{T}) where {T <: Real} = tf!(dtm, dtm)

tf!(dtm::SparseMatrixCSC{T}) where {T <: Real} = tf!(dtm, dtm)

tf(dtm::DocumentTermMatrix) = tf(dtm.dtm)

# The second Matrix will be overwritten with the result
# Will work correctly if dtm and tfidf are the same matrix
function tf!(dtm::AbstractMatrix{T1}, tf::AbstractMatrix{T2}) where {T1 <: Real, T2 <: AbstractFloat}
    n, p = size(dtm)

    # TF tells us what proportion of a document is defined by a term
    for i in 1:n
        words_in_document = 0
        for j in 1:p
            words_in_document += dtm[i, j]
        end
        tf[i, :] = dtm[i, :] ./ max(words_in_document, one(T1))
    end

    return tf
end

# assumes second matrix has same nonzeros as first one
function tf!(dtm::SparseMatrixCSC{T}, tf::SparseMatrixCSC{F}) where {T <: Real, F <: AbstractFloat}
    rows = rowvals(dtm)
    dtmvals = nonzeros(dtm)
    tfvals = nonzeros(tf)
    @assert size(dtmvals) == size(tfvals)

    # TF tells us what proportion of a document is defined by a term
    words_in_documents = sum(dtm,dims=2)

    n, p = size(dtm)
    for i = 1:p
       for j in nzrange(dtm, i)
          row = rows[j]
          tfvals[j] = dtmvals[j] / max(words_in_documents[row], one(T))
       end
    end
    tf
end

##############################################################################
#
# TF-IDF
#
##############################################################################

tf_idf(dtm::Matrix{T}) where {T <: Real} = tf_idf!(dtm, Array{Float64}(undef, size(dtm)...))

tf_idf(dtm::SparseMatrixCSC{T}) where {T <: Real} =  tf_idf!(dtm, similar(dtm, Float64))

tf_idf!(dtm::AbstractMatrix{T}) where {T <: Real} = tf_idf!(dtm, dtm)

tf_idf!(dtm::SparseMatrixCSC{T}) where {T <: Real} = tf_idf!(dtm, dtm)

tf_idf(dtm::DocumentTermMatrix) = tf_idf(dtm.dtm)

# This does not make sense, since DocumentTermMatrix is based on an array of integers
#tf_idf!(dtm::DocumentTermMatrix) = tf_idf!(dtm.dtm)

# The second Matrix will be overwritten with the result
# Will work correctly if dtm and tfidf are the same matrix
function tf_idf!(dtm::AbstractMatrix{T1}, tfidf::AbstractMatrix{T2}) where {T1 <: Real, T2 <: AbstractFloat}
    n, p = size(dtm)

    # TF tells us what proportion of a document is defined by a term
    tf!(dtm, tfidf)

    # IDF tells us how rare a term is in the corpus
    documents_containing_term = vec(sum(dtm .> 0, dims=1))
    idf = log.(n ./ documents_containing_term)

    # TF-IDF is the product of TF and IDF
    for i in 1:n
        for j in 1:p
           tfidf[i, j] = tfidf[i, j] * idf[j]
        end
    end

    return tfidf
end

# sparse version
function tf_idf!(dtm::SparseMatrixCSC{T}, tfidf::SparseMatrixCSC{F}) where {T <: Real, F <: AbstractFloat}
    rows = rowvals(dtm)
    dtmvals = nonzeros(dtm)
    tfidfvals = nonzeros(tfidf)
    @assert size(dtmvals) == size(tfidfvals)

    n, p = size(dtm)

    # TF tells us what proportion of a document is defined by a term
    words_in_documents = F.(sum(dtm, dims=2))
    oneval = one(F)

    # IDF tells us how rare a term is in the corpus
    documents_containing_term = vec(sum(dtm .> 0, dims=1))
    idf = log.(n ./ documents_containing_term)

    for i = 1:p
       for j in nzrange(dtm, i)
          row = rows[j]
          tfidfvals[j] = dtmvals[j] / max(words_in_documents[row], oneval) * idf[i]
       end
    end

    tfidf
end


# This file includes the testset for modulation pulses

@testset "PulseTestset" begin
    # Fields
    for type in [RectangularPulse, RaisedCosinePulse]
        @test :amplitude in fieldnames(type)
        @test :duration in fieldnames(type)
    end

    # Defaults
    pulse = RectangularPulse()
    @test pulse.amplitude ≈ 1.
    @test pulse.duration  ≈ 1.
    pulse = RaisedCosinePulse()
    @test pulse.amplitude ≈ 1.
    @test pulse.duration  ≈ 1.

    # Energy
    pulse = RectangularPulse(2., 4.)
    @test energy(pulse) ≈ 16.
    pulse = RaisedCosinePulse(2., 4.)
    @test energy(pulse) ≈ 4. * 3 /

    # Waveforms
    pulse = RectangularPulse(4., 1.)
    @test pulse(-1) ≈ 0.
    @test pulse(2.) ≈ 0.
    @test pulse(1/2) ≈ pulse.amplitude

    pulse = RaisedCosinePulse(4., 1.)
    @test pulse(-1.) ≈ 0.
    @test pulse(1/2) ≈ pulse.amplitude / 2 * (1 - cos(2π / pulse.duration * 1/2))
    @test pulse(2.) ≈ 0.
end # Test set.


######################################################################
# range.jl
#
# range only sensing
######################################################################

type RangeOnly <: Sensor
	noise_sigma::Float64		# noise std deviation

	# for discretized measurements
	num_bins::Int
	bin_range::UnitRange{Int64}

	function RangeOnly(noise_sigma::Real, num_bins::Int)
		new( float(noise_sigma), num_bins, 0:(num_bins-1) )
	end
	RangeOnly(noise_sigma::Real) = RangeOnly(noise_sigma, 1000)
end

function observe(m::SearchDomain, ro::RangeOnly, p::Pose)
	dx = (p[1] - m.theta[1])
	dy = (p[2] - m.theta[2])
	return sqrt(dx*dx + dy*dy) + ro.noise_sigma*randn()
end

# just returns a density. used in particle filtering
function O(ro::RangeOnly, theta::LocTuple, p::Pose, o::Float64)
	dx = (p[1] - theta[1])
	dy = (p[2] - theta[2])
	o_diff = sqrt(dx*dx + dy*dy) - o

	d = Normal(0, ro.noise_sigma)
	p = pdf(d, o_diff)
	return p
end

# just returns a density. used in particle filtering
# TODO: make this be correct
function O(ro::RangeOnly, theta::LocTuple, p::Pose, o::Int)
	dx = (p[1] - theta[1])
	dy = (p[2] - theta[2])
	o_diff = sqrt(dx*dx + dy*dy) - o

	d = Normal(0, ro.noise_sigma)
	p = pdf(d, o_diff)
	return p
end

# here, num_bins isn't important; we just bin to nearest integer
obs2bin(o::Float64, s::RangeOnly) = round(Int, o, RoundDown)


module QuadOsc

export quadosc

include("SeriesAccelerations.jl")
using .SeriesAccelerations
using QuadGK

@doc raw"""
    quadosc(fn, a, Inf, fnzeros; ...)

Integrate the function `fn(x)` from `a` to `Inf`. The function `fnzeros(n)`
takes an integer `n` and is such that `fn(fnzeros(n)) == 0`. The algorithm
works by integrating between successive zeros, and accelerating the alternating
series.

The argument `pren` is the number of intervals to integrate before applying the
series acceleration.

`atol` and `rtol` specify the absolute and relative tolerances for determining
convergence.

`order` is passed on to `quadgk()` of the
[QuadGK](https://github.com/JuliaMath/QuadGK.jl) package.

`nconvergences` is the number of iterations before convergence is declared.

See `?QuadOsc.accel_cohen_villegas_zagier` for details on the series
acceleration.
"""
function quadosc(f::Function, a::Number, b::Number, zerosf::Function; pren=2,
		 atol=zero(Float64), rtol=sqrt(eps(Float64)), order=7,
                 #nconvergences=ceil(Int,-1.31*log10(rtol)),
                 #accelerator=accel_cohen_villegas_zagier)
                 nconvergences=5,
                 accelerator=accel_wynn_eps)
    @assert b == Inf
    T = Float64

    i1 = findfirst(n -> zerosf(n) - a >= 0, 1:typemax(Int))
    z1 = zerosf(i1)
    Ipre, Epre = quadgk(f, a, z1; atol=atol, rtol=rtol, order=order)

    z0 = z1
    for i=1:pren
        i1 += 1
        z1 = zerosf(i1)
        I, E = quadgk(f, z0, z1; atol=atol, rtol=rtol, order=order)
        Ipre += I
        Epre += E
        z0 = z1
    end

    I = T(0)
    oldI = T(0)
    ak = T[]
    ek = T[]
    while nconvergences > 0
        i1 += 1
        z1 = zerosf(i1)
        I, E = quadgk(f, z0, z1; atol=atol, rtol=rtol)
        push!(ak, I)
        push!(ek, E)
        z0 = z1
        I = accelerator(ak)

        adiff = abs(I - oldI)
        rdiff = adiff * 2 / abs(I + oldI)
        if adiff <= atol || rdiff <= rtol
            nconvergences -= 1
        end
        oldI = I
    end
    I = Ipre + I
    E = Epre + abs(accelerator(ek))
    return I, E
end

end # module

# vim: set sw=4 et sts=4 :


mutable struct AMGPreconditioner{Tv, Ti} <: AbstractPreconditioner{Tv,Ti}
    A::ExtendableSparseMatrix{Tv,Ti}
    fact::AlgebraicMultigrid.Preconditioner
    max_levels::Int
    max_coarse::Int
    function AMGPreconditioner{Tv,Ti}(;max_levels=10, max_coarse=10) where {Tv,Ti}
        precon=new()
        precon.max_levels=max_levels
        precon.max_coarse=max_coarse
        precon
    end
end

"""
```
AMGPreconditioner(;max_levels=10, max_coarse=10, valuetype=Float64,indextype=Int64)
AMGPreconditioner(matrix;max_levels=10, max_coarse=10)
```

Create the  [`AMGPreconditioner`](@ref) wrapping the Ruge-Stüben AMG preconditioner from [AlgebraicMultigrid.jl](https://github.com/JuliaLinearAlgebra/AlgebraicMultigrid.jl)
"""
AMGPreconditioner(;valuetype::Type=Float64, indextype::Type=Int64,kwargs...)=AMGPreconditioner{valuetype,indextype}(;kwargs...)

@eval begin
    @makefrommatrix AMGPreconditioner
end

function update!(precon::AMGPreconditioner{Tv,Ti}) where {Tv,Ti}
    @inbounds flush!(precon.A)
    precon.fact=AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben(precon.A.cscmatrix))
end


"""
    xT,RT,ll = smooth(kf::KalmanFilter, u::Vector, y::Vector)
Returns smoothed estimates of state `x` and covariance `R` given all input output data `u,y`
"""
function smooth(kf::KalmanFilter, u::AbstractVector, y::AbstractVector)
    reset!(kf)
    T            = length(y)
    x,xt,R,Rt,ll = forward_trajectory(kf, u, y)
    xT           = similar(xt)
    RT           = similar(Rt)
    xT[end]      = xt[end]      |> copy
    RT[end]      = Rt[end]      |> copy
    for t = T-1:-1:1
        C     = Rt[t]*kf.A/R[t+1]
        xT[t] = xt[t] .+ C*(xT[t+1] .- x[t+1])
        RT[t] = Rt[t] .+ symmetrize(C*(RT[t+1] .- R[t+1])*C')
    end
    xT,RT,ll
end

function smooth(pf::AbstractParticleFilter, M, u, y)
    xf,wf,wef,ll = forward_trajectory(pf, u, y)
    smooth(pf::AbstractParticleFilter, xf, wf, wef, ll, M, u, y)
end

"""
    xb,ll = smooth(pf, M, u, y)

Perform particle smoothing using forward-filtering, backward simulation. Return smoothed particles and loglikelihood.
See also `smoothed_trajs`, `smoothed_mean`, `smoothed_cov`
"""
function smooth(pf::AbstractParticleFilter, xf, wf, wef, ll, M, u, y)
    T = length(y)
    N = num_particles(pf)
    f = dynamics(pf)
    df = dynamics_density(pf)
    @assert M <= N "Must extend cache size of bins and j to allow this"
    xb = Array{particletype(pf)}(undef,M,T)
    j = resample(ResampleSystematic, wef[:,T], M)
    # @show Set(j)
    for i = 1:M
        xb[i,T] = xf[j[i], T]
    end
    wb = Vector{Float64}(undef,N)
    @inbounds for t = T-1:-1:1
        # tset = Set{Int}()
        for m = 1:M
            for n = 1:N
                wb[n] = wf[n,t] + logpdf(df, xb[m,t+1], f(xf[n,t],u[t],t), t)
            end
            i = draw_one_categorical(pf,wb)
            # push!(tset, i)
            xb[m,t] = xf[i, t]
        end
        # @show tset
    end
    return xb,ll
end

function sse(f::AbstractFilter, u, y, λ=1)
    reset!(f)
    ll = sum(zip(u, y)) do (u,y)
        ll, e = f(u,y)
        dot(e, λ, e)
    end
end

"""
    ll = loglik(filter,u,y)
Calculate loglikelihood for entire sequences `u,y`
"""
function loglik(f::AbstractFilter,u,y)
    reset!(f)
    ll = sum(x->f(x[1],x[2])[1], zip(u, y))
end

function loglik(pf::AuxiliaryParticleFilter,u,y)
    reset!(pf)
    ll = sum(t->pf(u[t],y[t],y[t+1],t)[1], 1:length(u)-1)
    ll + pf.pf(u[end],y[end], length(u))[1]
end

"""
    ll(θ) = log_likelihood_fun(filter_from_parameters(θ::Vector)::Function, priors::Vector{Distribution}, u, y)

returns function θ -> p(y|θ)p(θ)
"""
function log_likelihood_fun(filter_from_parameters,priors::Vector{<:Distribution},u,y)
    n = numargs(filter_from_parameters)
    pf = nothing
    function (θ)
        pf === nothing && (pf = filter_from_parameters(θ))
        length(θ) == length(priors) || throw(ArgumentError("Input must have same length as priors"))
        ll = sum(i->logpdf(priors[i], θ[i]), eachindex(priors))
        isfinite(ll) || return -Inf
        pf = filter_from_parameters(θ,pf)
        ll += loglik(pf,u,y)
    end
end

function naive_sampler(θ₀)
    !any(iszero.(θ₀)) || throw(ArgumentError("Naive sampler does not work if initial parameter vector contains zeros (it was going to return θ -> θ .+ rand(MvNormal(0.1abs.(θ₀))), but that is not a good ideas if θ₀ is zero."))
    θ -> θ .+ rand(MvNormal(0.1abs.(θ₀)))
end

"""
    metropolis(ll::Function(θ), R::Int, θ₀::Vector, draw::Function(θ) = naive_sampler(θ₀))

Performs MCMC sampling using the marginal Metropolis (-Hastings) algorithm
`draw = θ -> θ'` samples a new parameter vector given an old parameter vector. The distribution must be symmetric, e.g., a Gaussian. `R` is the number of iterations.
See `log_likelihood_fun`

# Example:
```julia
filter_from_parameters(θ) = ParticleFilter(N, dynamics, measurement, MvNormal(n,exp(θ[1])), MvNormal(p,exp(θ[2])), d0)
priors = [Normal(0,0.1),Normal(0,0.1)]
ll     = log_likelihood_fun(filter_from_parameters,priors,u,y,1)
θ₀ = log.([1.,1.]) # Initial point
draw = θ -> θ .+ rand(MvNormal(0.1ones(2))) # Function that proposes new parameters (has to be symmetric)
burnin = 200 # If using threaded call, provide number of burnin iterations
# @time theta, lls = metropolis(ll, 2000, θ₀, draw) # Run single threaded
# thetam = reduce(hcat, theta)'
@time thetalls = LowLevelParticleFilters.metropolis_threaded(burnin, ll, 5000, θ₀, draw) # run on all threads, will provide (2000-burnin)*nthreads() samples
histogram(exp.(thetalls[:,1:2]), layout=3)
plot!(thetalls[:,3], subplot=3) # if threaded call, log likelihoods are in the last column
```
"""
function metropolis(ll, R, θ₀, draw = naive_sampler(θ₀))
    params    = Vector{typeof(θ₀)}(undef,R)
    lls       = Vector{Float64}(undef,R)
    params[1] = θ₀
    lls[1]    = ll(θ₀)
    for i = 2:R
        θ = draw(params[i-1])
        lli = ll(θ)
        if rand() < exp(lli-lls[i-1])
            params[i] = θ
            lls[i] = lli
        else
            params[i] = params[i-1]
            lls[i] = lls[i-1]
        end
    end
    params, lls
end

function metropolis_threaded(burnin, args...)
    res = []
    mtx = ReentrantLock()
    Threads.@threads for i = 1:Threads.nthreads()
        p,l = metropolis(args...)
        resi = [reduce(hcat,p)' l]
        resi = resi[burnin+1:end,:]
        lock(mtx)
        push!(res, resi)
        unlock(mtx)
    end
    reduce(vcat,res)
end

"""
    smoothed_mean(xb)
Helper function to calculate the mean of smoothed particle trajectories
"""
function smoothed_mean(xb)
    M,T = size(xb)
    n = length(xb[1])
    xbm = vec(mean(xb,dims=1))
    reduce(hcat, xbm)
end

"""
    smoothed_cov(xb)
Helper function to calculate the covariance of smoothed particle trajectories
"""
function smoothed_cov(xb)
    M,T = size(xb)
    n = length(xb[1])
    xbc = [cov(copy(reshape(reinterpret(Float64, xb[:,t]),n,M)),dims=2) for t = 1:T]
end

"""
    smoothed_trajs(xb)
Helper function to get particle trajectories as a 3-dimensions array (N,M,T) instead of matrix of vectors.
"""
function smoothed_trajs(xb)
    M,T = size(xb)
    n = length(xb[1])
    copy(reshape(reinterpret(Float64, xb), n,M,T))
end


# # 2D decaying turbulence
#
# A simulation of decaying two-dimensional turbulence.

using Pkg; Pkg.instantiate()

using FourierFlows, Printf, Random, JLD2

using Random: seed!
using FFTW: rfft, irfft
using LinearAlgebra: ldiv!
using FourierFlows: parsevalsum

import GeophysicalFlows.TwoDNavierStokes
import GeophysicalFlows.TwoDNavierStokes: energy, enstrophy
import GeophysicalFlows: peakedisotropicspectrum

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)

# ## Numerical, domain, and simulation parameters
#
# First, we pick some numerical and physical parameters for our model.

α = 1                      # aspect ratio parameter α = Ly / Lx
ny, Ly  = 2*2*256, 2π            # grid resolution and domain length
nx, Lx  = Int(2*2*256 / α), Ly / α

## Then we pick the time-stepper parameters
	 dt =0.5*1e-3  # timestep
nsubs = 1000    # number of steps between each plot

 ν = 0.2*1e-4

initial_data_wavenumber = 2.0
initial_data_bandwidth = 0.5

k₀ = initial_data_wavenumber

 tfinal = 600.0 #/ (ν * k₀^2)
 nsteps = Int(round(tfinal / dt))

 grid = TwoDGrid(dev, nx, Lx, ny, Ly)
 x, y = gridpoints(grid)

 # No Forcing
 calcF!(args...) = nothing

# Initial condition

# seed!(1234)

E0 = 0.1

function random_initial_condition(grid, initial_data_wavenumber, initial_data_bandwidth; energy=0.1)
	K = @. sqrt(grid.Krsq)
	forcing_spectrum = @. exp(-(K - initial_data_wavenumber)^2 / (2*initial_data_bandwidth^2))

	energy_before_normalization = parsevalsum(forcing_spectrum .* grid.invKrsq / 2, grid) / (grid.Lx * grid.Ly)

	forcing_spectrum *= energy / energy_before_normalization

	ξ = exp.(2π * im * rand(eltype(grid), size(K)))

	qh = @. sqrt(forcing_spectrum) * ξ
	q  = irfft(qh, grid.nx)

	return q
end

ζ₀ = random_initial_condition(grid, initial_data_wavenumber, initial_data_bandwidth; energy=E0)

# confirm that energy is E0
# prob = TwoDNavierStokes.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, ν=ν, dt=dt, calcF=calcF!, stepper="ETDRK4")
# TwoDNavierStokes.set_ζ!(prob, ζ₀)
# dx, dy = Lx/nx, Ly/ny
# println(sum(@. (prob.vars.u^2 + prob.vars.v^2)/2 * dx * dy / (Lx * Ly)))

# plot initial condition
# using Plots
# p =heatmap(grid.x, grid.y, ζ₀', aspectratio=1)
# display(p)

include("simulation.jl")


@testset "fastmod" begin
    a = convert(Vector{UInt32},sample(1:typemax(UInt32),10_000))
    a2 = copy(a)
    b = convert(Vector{UInt32},sample(1:typemax(UInt32),1_000))

    @testset for m in b
            modm = Mod(m)
            for i in eachindex(a)
                c1 = mod(a[i],m)
                c2 = fastmod(a[i],modm)
                @test c1==c2
            end
    end
end

@testset "fastmod!" begin
    a = convert(Vector{UInt32},sample(1:typemax(UInt32),10_000))
    a2 = copy(a)
    b = convert(Vector{UInt32},sample(1:typemax(UInt32),1_000))

    @testset for m in b
        modm = Mod(m)
        fastmod!(a,modm)
        inplace_mod!(a2,m)
        @test all(a.==a2)
    end
end

@testset "divtest" begin
    a = convert(Vector{UInt32},sample(1:typemax(UInt32),10_000))
    a2 = copy(a)
    b = convert(Vector{UInt32},sample(1:typemax(UInt32),1_000))
    c1 = similar(a)
    c2 = similar(a)

    @testset for m in b
        modm = Mod(m)
        for i in eachindex(a)
            c1[i] = (mod(a[i],m) == 0)
            c2[i] = (divtest(a[i],modm))
        end
        @test all(c1.==c2)
    end
end