######################### Gradient Computation #######################
#TODO: LSQ Gradient
# Non-functional
function leastSqGrad(mesh::Mesh, matrix::AbstractArray{Float64, 2}, stencil=zeros(2,2))
    # Stencil should be a list of lists, with each sublist containing the cells contained in the stencil of the main cell

    #cells, cVols, cCenters, faces, fAVecs, fCenters, boundaryFaces = mesh
    #bdryFaceIndices = Array(nFaces-nBdryFaces:nFaces)

    nCells, nFaces, nBoundaries, nBdryFaces = unstructuredMeshInfo(mesh)
    nVars = size(matrix, 2) # Returns the length of the second dimension of "matrix"

    grad = zeros(nCells, nVars, 3)

    L11 = zeros(nCells,nVars)
    L12 = zeros(nCells,nVars)
    L13 = zeros(nCells,nVars)
    L22 = zeros(nCells,nVars)
    L23 = zeros(nCells,nVars)
    L33 = zeros(nCells,nVars)
    L1f = zeros(nCells,nVars)
    L2f = zeros(nCells,nVars)
    L3f = zeros(nCells,nVars)

    @fastmath for f in 1:nFaces-nBdryFaces

        c1 = mesh.faces[f][1]
        c2 = mesh.faces[f][2]

        dx = mesh.cCenter[c2][1] - mesh.cCenter[c1][1]
        dy = mesh.cCenter[c2][2] - mesh.cCenter[c1][2]
        dz = mesh.cCenter[c2][3] - mesh.cCenter[c1][3]

        weight = 1 / sqrt(dx*dx + dy*dy + dz*dz)

        wdx = weight * dx
        wdy = weight * dy
        wdz = weight * dz

        for v in 1:nVars
            dv = matrix[c2,v] - matrix[c1,v]
            wdv = weight * dv

            L11[c1,v] += wdx^2
            L12[c1,v] += wdx * wdy
            L13[c1,v] += wdx * wdz
            L22[c1,v] += wdy^2
            L23[c1,v] += wdy * wdz
            L33[c1,v] += wdz^2

            L1f[c1,v] += wdx * wdv
            L2f[c1,v] += wdy * wdv
            L3f[c1,v] += wdz * wdv

            L11[c2,v] += wdx^2
            L12[c2,v] += wdx * wdy
            L13[c2,v] += wdx * wdz
            L22[c2,v] += wdy^2
            L23[c2,v] += wdy * wdz
            L33[c2,v] += wdz^2

            L1f[c2,v] += wdx * wdv
            L2f[c2,v] += wdy * wdv
            L3f[c2,v] += wdz * wdv

        end

    end

    # Deal with boundary faces, and add them to the matrix vectors

    return grad
end

#=
    Takes the gradient of (scalar) data provided in matrix form (passed into arugment 'matrix'):
    Cell      x1      x2      x3
    Cell 1    x1_1    x2_1    x3_1
    Cell 2    x1_2    x2_2    x3_2
    ...
    (where x1, x2, and x3 are arbitrary scalars)
    and output a THREE-DIMENSIONAL gradient arrays of the following form
    Cell      x1          x2          x3
    Cell 1    grad(x1)_1  grad(x2)_1  grad(x3)_1
    Cell 2    grad(x1)_2  grad(x2)_2  grad(x3)_2
    ...
    Where each grad(xa)_b is made up of THREE elements for the (x,y,z) directions

    Ex. Gradient @ cell 1 of P would be: greenGaussGrad(mesh, P)[1, 1, :]
=#
function greenGaussGrad(mesh::Mesh, matrix::AbstractArray{Float64, 2}, valuesAtFaces=false)
    nCells, nFaces, nBoundaries, nBdryFaces = unstructuredMeshInfo(mesh)
    nVars = size(matrix, 2)

    #Interpolate values to faces if necessary
    if valuesAtFaces
        faceVals = matrix
    else
        faceVals = linInterp_3D(mesh, matrix)
    end

    # Create matrix to hold gradients
    grad = zeros(nCells, nVars, 3)

    # Integrate fluxes from each face
    @fastmath for f in eachindex(mesh.faces)
        ownerCell = mesh.faces[f][1]
        neighbourCell = mesh.faces[f][2]

        @inbounds for v in 1:nVars
            for d in 1:3
                # Every face has an owner
                grad[ownerCell, v, d] += mesh.fAVecs[f][d] * faceVals[f, v]

                # Boundary faces don't - could split into two loops
                if neighbourCell > -1
                    grad[neighbourCell, v, d] -= mesh.fAVecs[f][d] * faceVals[f, v]
                end
            end
        end
    end

    # Divide integral by cell volume to obtain gradients
    @simd for c in 1:nCells
        for v in 1:nVars
            for d in 1:3
                grad[c,v,d] /= mesh.cVols[c]
            end
        end
    end

    return grad
end

####################### Face value interpolation ######################
#=
    Interpolates to all INTERIOR faces
    Arbitrary value matrix interpolation
    Example Input Matrix:
    Cell      x1       x2       x3
    Cell 1    x1_c1    x2_c1    x3_c1
    Cell 2    x1_c2    x2_c2    x3_c2
    ...
    Outputs a matrix of the following form
    Cell      x1       x2       x3
    Face 1    x1_f1    x2_f1    x3_f1
    Face 2    x1_f2    x2_f2    x3_f2
    ...
=#
function linInterp_3D(mesh::Mesh, matrix::AbstractArray{Float64, 2}, faceVals=nothing)
    nCells, nFaces, nBoundaries, nBdryFaces = unstructuredMeshInfo(mesh)
    nVars = size(matrix, 2)

    # Make a new matrix if one is not passed in
    if faceVals == nothing
        faceVals = zeros(nFaces, nVars)
    end

    # Boundary face fluxes must be set separately
    for f in 1:nFaces-nBdryFaces
        # Find value at face using linear interpolation
        c1 = mesh.faces[f][1]
        c2 = mesh.faces[f][2]

        #TODO: Precompute these distances?
        c1Dist = 0
        c2Dist = 0
        for i in 1:3
            c1Dist += (mesh.cCenters[c1][i] - mesh.fCenters[f][i])^2
            c2Dist += (mesh.cCenters[c2][i] - mesh.fCenters[f][i])^2
        end
        totalDist = c1Dist + c2Dist

        for v in 1:nVars
            faceVals[f, v] = matrix[c1, v]*(c2Dist/totalDist) + matrix[c2, v]*(c1Dist/totalDist)
        end
    end

    return faceVals
end

# Similar to linInterp_3D. Instead of linearly interpolating, selects the maximum value of the two adjacent cell centers as the face value
function maxInterp(mesh::Mesh, vars...)
    nCells, nFaces, nBoundaries, nBdryFaces = unstructuredMeshInfo(mesh)

    nVars = size(vars, 1)
    faceVals = zeros(nFaces, nVars)

    # Boundary face fluxes must be set separately
    @inbounds @fastmath for f in 1:nFaces-nBdryFaces
        c1 = mesh.faces[f][1]
        c2 = mesh.faces[f][2]

        for v in 1:nVars
            faceVals[f, v] = max(vars[v][c1], vars[v][c2])
        end
    end

    return faceVals
end

# Similar to linInterp_3D. Instead of linearly interpolating, calculates the change (delta) of each variable across each face (required for JST method)
function faceDeltas(mesh::Mesh, sln::SolutionState)
    nCells, nFaces, nBoundaries, nBdryFaces = unstructuredMeshInfo(mesh)
    nVars = size(sln.cellState, 2)

    faceDeltas = zeros(nFaces, nVars)

    # Boundary face fluxes must be set separately (faceDelta is zero at all possible boundary conditions right now)
    @inbounds @fastmath for f in 1:nFaces-nBdryFaces
        ownerCell = mesh.faces[f][1]
        neighbourCell = mesh.faces[f][2]

        for v in 1:nVars
            faceDeltas[f, v] = sln.cellState[neighbourCell, v] - sln.cellState[ownerCell, v]
        end
    end

    return faceDeltas
end

######################### TimeStepping #################################
# Calculate sln.cellPrimitives and sln.cellFluxes from sln.cellState
function decodeSolution_3D(sln::SolutionState, fluid::Fluid)
    nCells = size(sln.cellState, 1)
    for c in 1:nCells
        # Updates cell primitives
        @views decodePrimitives3D!(sln.cellPrimitives[c,:], sln.cellState[c,:], fluid)
        # Updates mass, xMom, eV2 x,y,z-direction fluxes
        @views calculateFluxes3D!(sln.cellFluxes[c, :], sln.cellPrimitives[c,:], sln.cellState[c,:])
    end
end

#=
    Calculates sln.fluxResiduals from sln.faceFluxes
=#
function integrateFluxes_unstructured3D(mesh::Mesh, sln::SolutionState)
    nCells, nFaces, nBoundaries, nBdryFaces = unstructuredMeshInfo(mesh)

    # Recomputing flux balances, so wipe existing values
    fill!(sln.fluxResiduals, 0)
    nVars = size(sln.fluxResiduals, 2)

    #### Flux Integration ####
    @inbounds @fastmath for f in eachindex(mesh.faces)
        ownerCell = mesh.faces[f][1]
        neighbourCell = mesh.faces[f][2]

        for v in 1:nVars
            i1 = (v-1)*3 + 1
            i2 = i1+2
            @views flow = dot(sln.faceFluxes[f, i1:i2], mesh.fAVecs[f])

            # Subtract from owner cell
            sln.fluxResiduals[ownerCell, v] -= flow
            # Add to neighbour cell
            if neighbourCell > -1
                sln.fluxResiduals[neighbourCell, v] += flow
            end
        end
    end

    # Divide by cell volume
    for c in 1:nCells
        for v in 1:nVars
            sln.fluxResiduals[c,v] /= mesh.cVols[c]
        end
    end

    return sln.fluxResiduals
end


function cumulative(weights::Array{Int, 2})
	ret = copy(weights)
	dim = size(weights)
	for i in 2 : dim[1]
		for j in 1 : dim[2]
			min = ret[i - 1, j]
			if j > 1 && ret[i - 1, j - 1] < min
				min = ret[i - 1, j - 1]
			end
			if j < dim[2] && ret[i - 1, j + 1] < min
				min = ret[i - 1, j + 1]
			end
			ret[i, j] += min
		end
	end
	return ret
end

function back_track(weights::Array{Int, 2})
	dim = size(weights)

	buffer::Array{Int, 1} = ones(Int, dim[1])

	buffer[1] = 1
	for i in 2 : dim[2]
		if weights[dim[1], i] < weights[dim[1], buffer[1]]
			buffer[1] = i
		end
	end

	for i in dim[1] - 1 : -1 : 1
		k = dim[1] - i + 1
		j = buffer[k - 1]
		buffer[k] = j
		if j > 1 && weights[i, j - 1] < weights[i, buffer[k]]
			buffer[k] = j - 1
		end
		if j < dim[2] && weights[i, j + 1] < weights[i, buffer[k]]
			buffer[k] = j + 1
		end
	end

	ret::Array{Tuple{Int, Int}, 1} = Array{Tuple{Int, Int}, 1}(undef, dim[1])

	for i in 1 : dim[1]
		ret[i] = dim[1] - i + 1, buffer[i]
	end

	return ret
end


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 98ed9a78-f50c-11ea-1d20-7bcaee689a99
using Pkg, DrWatson

# ╔═╡ 98edd6be-f50c-11ea-26c3-ff6ac7641060
begin
	@quickactivate "StatisticalRethinkingTuring"
	using StatisticalRethinking
end

# ╔═╡ 6e665478-f50c-11ea-1861-3f509f92ac5c
md"## Fig3.6t.jl"

# ╔═╡ 98ee52ec-f50c-11ea-22bc-c1075a60bfed
md"### snippet 3.23"

# ╔═╡ 98f67062-f50c-11ea-23f3-539ae4f7e67b
begin
	N = 10000
	p =
	for j in 1:9

	  prob = j * 0.1
	  d = rand(Binomial(9, prob), N);
	  h = fit(Histogram, d, -0.5:1:9.5)

	  p[j] = plot(xlim=(0,9), xticks=0:9)
	  for (i, w) in enumerate(h.weights)
		plot!(p[j], [i-1, i-1], [0.0, w], color=:blue,
		  leg=false, title="prob=$(round(prob, digits=1))")
	  end
	end
end

# ╔═╡ 98f93450-f50c-11ea-3767-e526c9109464
plot(p..., layout=(3,3))

# ╔═╡ 99019636-f50c-11ea-3ac3-4107ebb330c5
md"## End of Fig3.6t.jl"

# ╔═╡ Cell order:
# ╟─6e665478-f50c-11ea-1861-3f509f92ac5c
# ╠═98ed9a78-f50c-11ea-1d20-7bcaee689a99
# ╠═98edd6be-f50c-11ea-26c3-ff6ac7641060
# ╟─98ee52ec-f50c-11ea-22bc-c1075a60bfed
# ╠═98f67062-f50c-11ea-23f3-539ae4f7e67b
# ╠═98f93450-f50c-11ea-3767-e526c9109464
# ╟─99019636-f50c-11ea-3ac3-4107ebb330c5


function c_R_to_G(Ns::Array{Int,1}, fR)
  out = zeros(Complex128,Ns[1],Ns[2],Ns[3])
  ccall( (:fftw_fw_fft3d, "fft3d"), Void,
    (Ptr{Complex128}, Ptr{Complex128}, Int64, Int64, Int64), fR, out,
    Ns[3],Ns[2],Ns[1] )
  return out
end

function c_G_to_R(Ns::Array{Int,1}, fG)
  out = zeros(Complex128,Ns[1],Ns[2],Ns[3])
  ccall( (:fftw_inv_fft3d, "fft3d"), Void,
    (Ptr{Complex128}, Ptr{Complex128}, Int64, Int64, Int64), fG, out,
    Ns[3],Ns[2],Ns[1] )
  return out
end

# generic, one-column version
function R_to_G( Ns::Array{Int,1}, fR )
  out = reshape( fft( reshape(fR,Ns[1],Ns[2],Ns[3]) ),
                   size(fR) )
end

function G_to_R( Ns::Array{Int,1}, fG )
  out = reshape( ifft( reshape(fG,Ns[1],Ns[2],Ns[3]) ),
                 size(fG) )
  return out
end

function test_main()
  Ns = [150,150,150]
  in1 = rand(Complex128,Ns[1],Ns[2],Ns[3])

  @time out1c = c_R_to_G(Ns,in1)
  @time in2c = c_G_to_R(Ns,out1c)

  @time out1 = R_to_G(Ns,in1)
  @time in2 = G_to_R(Ns,out1)

  println("diff out1 = ", sum( abs(out1c - out1) ))
  println("diff out1 = ", sum( abs(in2c - in2) ))
end

test_main()


# Betweenness centrality measures
# TODO - weighted, separate unweighted, edge betweenness

@doc_str """
    betweenness_centrality(g[, vs])
    betweenness_centrality(g, k)
    parallel_betweenness_centrality(g[, vs])
    parallel_betweenness_centrality(g, k)
Calculate the [betweenness centrality](https://en.wikipedia.org/wiki/Centrality#Betweenness_centrality)
of a graph `g` across all vertices, a specified subset of vertices `vs`, or a random subset of `k`
vertices. Return a vector representing the centrality calculated for each node in `g`.
### Optional Arguments
- `normalize=true`: If true, normalize the betweenness values by the
total number of possible distinct paths between all pairsin the graphs.
For an undirected graph, this number is ``\\frac{(|V|-1)(|V|-2)}{2}``
and for a directed graph, ``{(|V|-1)(|V|-2)}``.
- `endpoints=false`: If true, include endpoints in the shortest path count.
Betweenness centrality is defined as:
``
bc(v) = \\frac{1}{\\mathcal{N}} \sum_{s \\neq t \\neq v}
\\frac{\\sigma_{st}(v)}{\\sigma_{st}}
``.
### References
- Brandes 2001 & Brandes 2008
"""
function betweenness_centrality(
    g::AbstractGraph,
    vs::AbstractVector = vertices(g);
    normalize=true,
    endpoints=false)

    n_v = nv(g)
    k = length(vs)
    isdir = is_directed(g)

    betweenness = zeros(n_v)
    for s in vs
        if degree(g, s) > 0  # this might be 1?
            state = dijkstra_shortest_paths(g, s; allpaths=true, trackvertices=true)
            if endpoints
                _accumulate_endpoints!(betweenness, state, g, s)
            else
                _accumulate_basic!(betweenness, state, g, s)
            end
        end
    end

    _rescale!(betweenness,
    n_v,
    normalize,
    isdir,
    k)

    return betweenness
end

betweenness_centrality(g::AbstractGraph, k::Integer; normalize=true, endpoints=false) =
    betweenness_centrality(g, sample(vertices(g), k); normalize=normalize, endpoints=endpoints)

function parallel_betweenness_centrality(
    g::AbstractGraph,
    vs::AbstractVector = vertices(g);
    normalize=true,
    endpoints=false)::Vector{Float64}

    n_v = nv(g)
    k = length(vs)
    isdir = is_directed(g)

    # Parallel reduction

    betweenness = @parallel (+) for s in vs
        temp_betweenness = zeros(n_v)
        if degree(g, s) > 0  # this might be 1?
            state = dijkstra_shortest_paths(g, s; allpaths=true, trackvertices=true)
            if endpoints
                _accumulate_endpoints!(temp_betweenness, state, g, s)
            else
                _accumulate_basic!(temp_betweenness, state, g, s)
            end
        end
        temp_betweenness
    end

    _rescale!(betweenness,
    n_v,
    normalize,
    isdir,
    k)

    return betweenness
end

parallel_betweenness_centrality(g::AbstractGraph, k::Integer; normalize=true, endpoints=false) =
    parallel_betweenness_centrality(g, sample(vertices(g), k); normalize=normalize, endpoints=endpoints)

function _accumulate_basic!(
    betweenness::Vector{Float64},
    state::DijkstraState,
    g::AbstractGraph,
    si::Integer
    )

    n_v = length(state.parents) # this is the ttl number of vertices
    δ = zeros(n_v)
    σ = state.pathcounts
    P = state.predecessors

    # make sure the source index has no parents.
    P[si] = []
    # we need to order the source vertices by decreasing distance for this to work.
    S = reverse(state.closest_vertices) #Replaced sortperm with this
    for w in S
        coeff = (1.0 + δ[w]) / σ[w]
        for v in P[w]
            if v > 0
                δ[v] += (σ[v] * coeff)
            end
        end
        if w != si
            betweenness[w] += δ[w]
        end
    end
    return nothing
end

function _accumulate_endpoints!(
    betweenness::Vector{Float64},
    state::DijkstraState,
    g::AbstractGraph,
    si::Integer
    )

    n_v = nv(g) # this is the ttl number of vertices
    δ = zeros(n_v)
    σ = state.pathcounts
    P = state.predecessors
    v1 = [1:n_v;]
    v2 = state.dists
    S = reverse(state.closest_vertices)
    s = vertices(g)[si]
    betweenness[s] += length(S) - 1    # 289

    for w in S
        coeff = (1.0 + δ[w]) / σ[w]
        for v in P[w]
            δ[v] += σ[v] * coeff
        end
        if w != si
            betweenness[w] += (δ[w] + 1)
        end
    end
    return nothing
end

function _rescale!(betweenness::Vector{Float64}, n::Integer, normalize::Bool, directed::Bool, k::Int)
    if normalize
        if n <= 2
            do_scale = false
        else
            do_scale = true
            scale = 1.0 / ((n - 1) * (n - 2))
        end
    else
        if !directed
            do_scale = true
            scale = 1.0 / 2.0
        else
            do_scale = false
        end
    end
    if do_scale
        if k > 0
            scale = scale * n / k
        end
        for v = 1:length(betweenness)
            betweenness[v] *= scale
        end
    end
    return nothing
end


#Demo 5: Logy plots
#-------------------------------------------------------------------------------
using InspectDR
using Colors

#==Input
===============================================================================#

#Constants
#-------------------------------------------------------------------------------
black = RGB24(0, 0, 0)
white = RGB24(1, 1, 1)
red = RGB24(1, 0, 0)
green = RGB24(0, 1, 0)
blue = RGB24(0, 0, 1)

#Input data
#-------------------------------------------------------------------------------
x = collect(1.0:1.0:100.0)
b = 2.0 #Base
k = 0.2 #Exponential coefficient

#==Generate plot
===============================================================================#
mplot = InspectDR.Multiplot(title="Log-Y Tests")
mplot.ncolumns = 1

plot_liny = InspectDR.Plot2D()
	plot_liny.axes = InspectDR.axes(:lin, :lin)
plot_logy = InspectDR.Plot2D()
	plot_logy.axes = InspectDR.axes(:lin, :log10)
	plot_logy.layout.grid = grid(vmajor=true, hmajor=true, hminor=true)

plotlist = [plot_liny, plot_logy]

for plot in plotlist
	wfrm = add(plot, x, b.^(k*x))
	wfrm.line = line(color=blue, width=3)

	add(mplot, plot)
	a = plot.annotation
	a.title = "y = $b^($(k)x)"
	a.xlabel = "x"
	a.ylabel = "y"
end

gplot = display(InspectDR.GtkDisplay(), mplot)

:DONE


using SparseArrays

function sub2ind(array_shape, rows, cols)
    out = (cols.-1).*array_shape[1] .+ rows
    return convert.(Int64, out)
end

function compute_warping_operator(flow)
    # ref: https://github.com/HendrikMuenster/flexBox/blob/79abc7285703911cca2653434cca5bcefc79c722/operators/generateMatrix/warpingOperator.m
    """
    u: gray image
    v: flow [H, W, 2]
    """
    H, W, dims = size(flow)
    shape = (H,W)
    XX, YY = repeat(collect(1:W)', H, 1), repeat(collect(1:H), 1, W)
    targetPoint2Mat = XX + flow[:,:,1]
    targetPoint1Mat = YY + flow[:,:,2]
    x1 = floor.(targetPoint1Mat).-1
    x2 = x1 .+ 1
    x3 = x1 .+ 2
    x4 = x1 .+ 3
    y1 = floor.(targetPoint2Mat).-1
    y2 = y1 .+ 1
    y3 = y1 .+ 2
    y4 = y1 .+ 3
    v2 = targetPoint1Mat .- x2
    v1 = targetPoint2Mat .- y2
    indicator = ones(H,W)
    indicator = (x1 .> 0).*indicator
    indicator = (x4 .<= W).*indicator
    indicator = (y1 .> 0).*indicator
    indicator = (y4 .<= H).*indicator
    indicator = indicator .> 0
    v1 = v1[indicator]
    v2 = v2[indicator]
    idxList = sub2ind(shape, YY, XX)
    idxList_ = collect(Iterators.flatten(idxList))
    #find the indices
    list_J  = vcat(
        idxList_[sub2ind(shape, x2[indicator], y2[indicator])],
        idxList_[sub2ind(shape, x3[indicator], y2[indicator])],
        idxList_[sub2ind(shape, x2[indicator], y3[indicator])],
        idxList_[sub2ind(shape, x3[indicator], y3[indicator])])
    # list_J  = vcat(
    #     idxList_[sub2ind(shape, x1[indicator], y1[indicator])],
    #     idxList_[sub2ind(shape, x2[indicator], y1[indicator])],
    #     idxList_[sub2ind(shape, x3[indicator], y1[indicator])],
    #     idxList_[sub2ind(shape, x4[indicator], y1[indicator])],
    #     idxList_[sub2ind(shape, x1[indicator], y2[indicator])],
    #     idxList_[sub2ind(shape, x2[indicator], y2[indicator])],
    #     idxList_[sub2ind(shape, x3[indicator], y2[indicator])],
    #     idxList_[sub2ind(shape, x4[indicator], y2[indicator])],
    #     idxList_[sub2ind(shape, x1[indicator], y3[indicator])],
    #     idxList_[sub2ind(shape, x2[indicator], y3[indicator])],
    #     idxList_[sub2ind(shape, x3[indicator], y3[indicator])],
    #     idxList_[sub2ind(shape, x4[indicator], y3[indicator])],
    #     idxList_[sub2ind(shape, x1[indicator], y4[indicator])],
    #     idxList_[sub2ind(shape, x2[indicator], y4[indicator])],
    #     idxList_[sub2ind(shape, x3[indicator], y4[indicator])],
    #     idxList_[sub2ind(shape, x4[indicator], y4[indicator])])
    list_J = collect(Iterators.flatten(list_J))
    list_val  = ( (1 .-v1) .* (1 .-v2))
    list_val = vcat(list_val,(v2 .* (1 .-v1)))
    list_val = vcat(list_val,(v1 .* (1 .-v2)))
    list_val = vcat(list_val,(v1 .* v2))
    # list_val =vcat(
    #     (v1.*v2.*(v1 .- 1).^2 .*(v2 .- 1).^2)/4,
    #     -1 .*(v1.*(v1 .- 1).^2 .*(3 .*v2.^3 .- 5 .*v2.^2 .+ 2))./4,
    #     -1 .*(v1.*v2.*(v1 .- 1).^2 .*(- 3 .*v2.^2 + 4 .*v2 .+ 1))./4,
    #     -1 .*(v1.*v2.^2 .*(v1 .- 1).^2 .*(v2 .- 1))./4,
    #     -1 .*(v2.*(v2 .- 1).^2 .*(3 .*v1.^3 .- 5 .*v1.^2 .+ 2))./4,
    #     ((3 .*v1.^3 .- 5 .*v1.^2 .+ 2).*(3 .*v2.^3 .- 5 .*v2.^2 .+ 2))./4,
    #     (v2.*(- 3 .*v2.^2 .+ 4 .*v2 .+ 1).*(3 .*v1.^3 .- 5 .*v1.^2 .+ 2))./4,
    #     (v2.^2 .*(v2 .- 1).*(3 .*v1.^3 .- 5 .*v1.^2 .+ 2))./4,
    #     -1 .*(v1.*v2.*(v2 .- 1).^2 .*(- 3 .*v1.^2 .+ 4 .*v1 .+ 1))./4,
    #     (v1.*(- 3 .*v1.^2 .+ 4 .*v1 .+ 1).*(3 .*v2.^3 .- 5 .*v2.^2 .+ 2))./4,
    #     (v1.*v2.*(- 3 .*v1.^2 .+ 4 .*v1 .+ 1).*(- 3 .*v2.^2 .+ 4 .*v2 .+ 1))./4,
    #     (v1.*v2.^2 .*(v2 .- 1).*(- 3 .*v1.^2 .+ 4 .*v1 .+ 1))./4,
    #     -1 .*(v1.^2 .*v2.*(v1 .- 1).*(v2 .- 1).^2)./4,
    #     (v1.^2 .*(v1 .- 1).*(3 .*v2.^3 .- 5 .*v2.^2 .+ 2))./4,
    #     (v1.^2 .*v2.*(v1 .- 1).*(- 3 .*v2.^2 .+ 4 .*v2 .+ 1))./4,
    #     (v1.^2 .*v2.^2 .*(v1 .- 1).*(v2 .- 1))./4
    # )
    list_val = collect(Iterators.flatten(list_val))
    list_I = idxList[indicator]
    #list_I =collect(Iterators.flatten(repeat(list_I,1,16)))
    list_I =collect(Iterators.flatten(repeat(list_I,1,4)))
    # list_I_py, list_J_py, list_val_py = py"py_warping_operator"(flow)
    # list_I_py = list_I_py.+1
    # list_J_py = list_J_py.+1
    # W_mat_py = sparse(list_I_py, list_J_py, list_val_py, H*W, H*W)
    W_mat = dropzeros!(sparse(list_I, list_J, list_val, H*W, H*W))
    return W_mat
end

img2 = zeros(128,128)
H, W = size(img2)
img2[30:60,40:50] .= 255.0

flow = zeros(H, W, 2)
flow[30:60,40:50, 1] .= -1.0

Wop = compute_warping_operator(flow)
    img1_ = Wop * vec(img2)
img1 = reshape(img1_, H, W)


## plotting code using Plots

"""

Plotting of implicit functions.

An implicit function or equation is defined in terms of a logical
condition and a function of two variables. These produce `Predicate`
objects.

Predicates may be plotted over a specified region, the default begin [-5,5] x [-5,5].

The algorithm, breaks the region into blocks. The ultimate resolution
is given by `W=2^n` and `H=2^m`. The algorithm used, due to Tupper, colors
region if it knows the predicate is true or false, and otherwise
resolves the region into 4 equal-sized subregions and test each subregion
again to determine if it is true, false, or still maybe. This repeats until `W` and `H`
can't be subdivided.

The `plot`ting of a predicate simply plots each block that is knows satisfies the predicate "black" and each ambiguous block "red." By taking `m` and `n` larger the graphs look less "blocky" but take more time to render.

For text-based plots, the `asciigraph` function is available.

Examples:
```
using Plots, ImplicitEquations
pyplot()

a,b = -1,2
f(x,y) = y^4 - x^4 + a*y^2 + b*x^2
plot(f == 0)

## trident of Newton
c,d,e,h = 1,1,1,1
f(x,y) = x*y
g(x,y) =c*x^3 + d*x^2 + e*x + h
plot(eq(f,g), title="Trident of Newton") ## aka f ⩵ g (using Unicode\Equal[tab])

## inequality
f(x,y)= (y-5)*cos(4*sqrt((x-4)^2 + y^2))
g(x,y) = x*sin(2*sqrt(x^2 + y^2))
r = f < g
plot(r, (-10, 10), (-10, 10), N=9, M=9)  # (xmin, xmax), (ymin, ymax),
```
"""
plot_implicit = nothing

## Helpers to convert
linterp(A,B,a,b,W) = (a + A/W*(b-a),a + B/W*(b-a))

function xyrange(u, L, R, B, T, W, H; offset=0)
    tuple(linterp(u.x.val.lo, u.x.val.hi + offset, L, R, W)...,
          linterp(u.y.val.lo, u.y.val.hi + offset, B, T, H)...)
end

function get_xs_ys(rs, L, R, B, T, W, H)
    xs = Float64[]
    ys = Float64[]
    for u in rs
        x0,x1,y0,y1 = ImplicitEquations.xyrange(u, L,R,B,T,W,H)
        append!(xs, [x0,x1,x1,x0,x0]); push!(xs, NaN)
        append!(ys, [y0,y0,y1,y1,y0]); push!(ys, NaN)
    end
    if length(xs) > 0
        pop!(xs)
        pop!(ys)
    else
        xs = Float64[NaN]
        ys = Float64[NaN]
    end

    xs, ys
end

## A plot recipe
## x, y describe region to plot over
## N, M give no. of pixels 2^N by 2^M
## red and black are used for colors.
@recipe function f(p::Predicate, x=(-5,5), y=(-5,5);
                   N=8,
                   M=8,              # oddly m as keyword fails. 9/8 too slow
                   red=nothing,      # or :red ...
                   black=:black
                   )

    L, R = extrema(x)
    B, T = extrema(y)

    W = 2^N
    H = 2^M

    r, b, w = ImplicitEquations.GRAPH(p, L, R, B, T, W, H)

    ## add red as a series
    if length(r) > 0 && red != nothing
        @series begin
            xs, ys = get_xs_ys(r, L, R, B, T, W, H)

            seriestype := :shape
            fillcolor := red
            linewidth := 0
            x := xs
            y := ys

            ()
        end
    end

    ## do black
    seriestype --> :shape
    xlims --> [L, R]
    ylims --> [B, T]
    legend --> false

    fillcolor --> black
    linewidth --> 0

    xs, ys = get_xs_ys(b, L, R, B, T, W, H)
    x --> xs
    y --> ys

    ()

end


#---------------------------------------------------------------------#
#This function computes the Derivative Mapping for General 2D Quad Grids.
#Written by F.X. Giraldo on 5/2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
function map_deriv(ψ,dψ,f,Np,Nq,DFloat)

    #Allocate arrays
    f_ξ=zeros(DFloat,Nq,Nq,Nq)
    f_η=zeros(DFloat,Nq,Nq,Nq)
    f_ζ=zeros(DFloat,Nq,Nq,Nq)

    #Loop through Integration Points
    for n=1:Nq, m=1:Nq, l=1:Nq
        sum_ξ=0
        sum_η=0
        sum_ζ=0
        for k=1:Np, j=1:Np, i=1:Np
            sum_ξ=sum_ξ + dψ[i,l]*ψ[j,m]*ψ[k,n]*f[i,j,k]
            sum_η=sum_η + ψ[i,l]*dψ[j,m]*ψ[k,n]*f[i,j,k]
            sum_ζ=sum_ζ + ψ[i,l]*ψ[j,m]*dψ[k,n]*f[i,j,k]
        end #i,j,k
        f_ξ[l,m,n]=sum_ξ
        f_η[l,m,n]=sum_η
        f_ζ[l,m,n]=sum_ζ
    end #l,m,n
    return (f_ξ,f_η,f_ζ)
end


module TensorLabXD

# Exports
# Types:
export Sector, AbstractIrrep, Irrep
export FusionStyle, UniqueFusion, MultipleFusion, MultiplicityFreeFusion,
        SimpleFusion, GenericFusion
export BraidingStyle, SymmetricBraiding, Bosonic, Fermionic, Anyonic
export Z2Irrep, Z3Irrep, Z4Irrep, ZNIrrep, U1Irrep, SU2Irrep, CU1Irrep
export Fermion, FermionParity, FermionNumber, FermionSpin
export FibonacciAnyon
export IsingAnyon

export VectorSpace, Field, ElementarySpace, InnerProductSpace, EuclideanSpace # abstract vector spaces
export ComplexSpace, CartesianSpace, GeneralSpace, GradedSpace # concrete spaces
export ZNSpace, Z2Space, Z3Space, Z4Space, U1Space, CU1Space, SU2Space
export Vect, Rep # space constructors
export CompositeSpace, ProductSpace # composite spaces
export FusionTree
export IndexSpace, TensorSpace, AbstractTensorMap, AbstractTensor, TensorMap, Tensor # tensors and tensor properties
export TensorMapSpace, EuclideanTensorSpace, EuclideanTensorMapSpace
export AbstractEuclideanTensorMap, EuclideanTensorMap
export TruncationScheme
export SpaceMismatch, SectorMismatch, IndexError # error types

# general vector space methods
export space, field, dual, dim, dims, fuse, flip, isdual, insertunit

# partial order for vector spaces
export infimum, supremum, isisomorphic, ismonomorphic, isepimorphic

# methods for sectors and properties thereof
export sectortype, sectors, hassector, Nsymbol, Fsymbol, Rsymbol, Bsymbol,
        frobeniusschur, twist
export fusiontrees, braid, permute, transpose
export Trivial, ZNSpace, SU2Irrep, U1Irrep, CU1Irrep # Fermion
# other fusion tree manipulations, should not be exported:
# export insertat, split, merge, repartition, artin_braid,
#        bendleft, bendright, foldleft, foldright, cycleclockwise, cycleanticlockwise

# some unicode
export ⊕, ⊗, ×, ⊠, ℂ, ℝ, ℤ, ←, →, ≾, ≿, ≅, ≺, ≻
export ℤ₂, ℤ₃, ℤ₄, U₁, SU, SU₂, CU₁
export fℤ₂, fU₁, fSU₂
export ℤ₂Space, ℤ₃Space, ℤ₄Space, U₁Space, CU₁Space, SU₂Space

# tensor maps
export domain, codomain, numind, numout, numin, spacetype, storagetype, eltype
export domainind, codomainind, allind
export tensormaptype
export blocksectors, blockdim, block, blocks

# random methods for constructor
export randuniform, randnormal, randisometry, randhaar

# special purpose constructors
export zero, one, one!, id, isomorphism, unitary, isometry

# tensor algebra and factorizations
export dot, norm, normalize, normalize!, tr
export mul!, lmul!, rmul!, adjoint!, pinv, axpy!, axpby!
export leftorth, rightorth, leftnull, rightnull,
        leftorth!, rightorth!, leftnull!, rightnull!,
        tsvd!, tsvd, eigen, eigen!, eig, eig!, eigh, eigh!, exp, exp!,
        isposdef, isposdef!, ishermitian, sylvester
export braid!, permute!, transpose!, twist!
export catdomain, catcodomain

export OrthogonalFactorizationAlgorithm, QR, QRpos, QL, QLpos, LQ, LQpos, RQ, RQpos,
        SVD, SDD, Polar

# tensor operations
export @tensor, @tensoropt, @ncon, ncon, @planar, @planar2
export scalar, add!, contract!

# truncation schemes
export notrunc, truncerr, truncdim, truncspace, truncbelow

# Imports
using TupleLabXD
using TupleLabXD: StaticLength

using StridedTensorXD

using TensorContractionsXD: TensorContractionsXD, @tensor, @tensoropt, @ncon, ncon
const TO = TensorContractionsXD

using LRUCache

using HalfIntegers
using WignerSymbols

using Base: @boundscheck, @propagate_inbounds, OneTo, tail, front,
            tuple_type_head, tuple_type_tail, tuple_type_cons,
            SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype
using Base.Iterators: product, filter
import Base: min, max

import LinearAlgebra
using LinearAlgebra: norm, dot, normalize, normalize!, tr,
                        axpy!, axpby!, lmul!, rmul!, mul!,
                        adjoint, adjoint!, transpose, transpose!,
                        pinv, sylvester,
                        eigen, eigen!, svd, svd!,
                        isposdef, isposdef!, ishermitian,
                        Diagonal, Hermitian
import Base.Meta

const IndexTuple{N} = NTuple{N, Int}

# Auxiliary files
include("auxiliary/auxiliary.jl")
include("auxiliary/dicts.jl")
include("auxiliary/linalg.jl")
include("auxiliary/random.jl")

# experiment with different dictionaries
const SectorDict{K, V} = SortedVectorDict{K, V}
const FusionTreeDict{K, V} = Dict{K, V}
fusiontreedict(I) = FusionStyle(I) isa UniqueFusion ? SingletonDict : FusionTreeDict

# Exception types:
abstract type TensorException <: Exception end

# Exception type for all errors related to sector mismatch
struct SectorMismatch{S<:Union{Nothing, String}} <: TensorException
    message::S
end
SectorMismatch()=SectorMismatch{Nothing}(nothing)
Base.show(io::IO, ::SectorMismatch{Nothing}) = print(io, "SectorMismatch()")
Base.show(io::IO, e::SectorMismatch) = print(io, "SectorMismatch(", e.message, ")")

# Exception type for all errors related to vector space mismatch
struct SpaceMismatch{S<:Union{Nothing, String}} <: TensorException
    message::S
end
SpaceMismatch()=SpaceMismatch{Nothing}(nothing)
Base.show(io::IO, ::SpaceMismatch{Nothing}) = print(io, "SpaceMismatch()")
Base.show(io::IO, e::SpaceMismatch) = print(io, "SpaceMismatch(", e.message, ")")

# Exception type for all errors related to invalid tensor index specification.
struct IndexError{S<:Union{Nothing, String}} <: TensorException
    message::S
end
IndexError()=IndexError{Nothing}(nothing)
Base.show(io::IO, ::IndexError{Nothing}) = print(io, "IndexError()")
Base.show(io::IO, e::IndexError) = print(io, "IndexError(", e.message, ")")

# typerepr
type_repr(T::Type) = repr(T)

# Definitions and methods for superselection sectors (quantum numbers)
include("sectors/sectors.jl")

# Constructing and manipulating fusion trees and iterators thereof
include("fusiontrees/fusiontrees.jl")

# Definitions and methods for vector spaces
include("spaces/vectorspaces.jl")

# Definitions and methods for tensors
include("tensors/abstracttensor.jl")
include("tensors/tensortreeiterator.jl")
include("tensors/tensor.jl")
include("tensors/adjoint.jl")
include("tensors/linalg.jl")
include("tensors/tensoroperations.jl")
include("tensors/indexmanipulations.jl")
include("tensors/truncation.jl")
include("tensors/factorizations.jl")
include("tensors/planar.jl")
include("tensors/braidingtensor.jl")

# deprecations: to be removed in version 1.0 or sooner
include("auxiliary/deprecate.jl")

end


## pav.jl  Pool Adjacent Violators algorithm for isotonic regression
## (c) 2015 David A. van Leeuwen
##
## Licensed under the MIT software license, see LICENSE.md
## This code is largely based on code from the Bosaris toolkit

## pav(y) returns the isotonic regression of the predictor y.
function pav{T<:Real}(y::Vector{T})
    yy = similar(y, Float64)
    n = length(y)
    i = similar(y, Int)         # index
    l = similar(y, Int)         # length
    ci = 1
    i[ci] = l[ci] = 1
    yy[ci] = y[1]
    for j=2:n
        ci += 1
        i[ci] = j
        l[ci] = 1
        yy[ci] = y[j]
        while ci ≥ 2 && yy[max(ci-1,1)] ≥ yy[ci]
            nl = l[ci-1] + l[ci]
            yy[ci-1] += l[ci] * (yy[ci] - yy[ci-1]) / nl
            ci -= 1
            l[ci] = nl
        end
    end
    while n > 0
        for j=i[ci]:n
            yy[j] = yy[ci]
        end
        n = i[ci] - 1
        ci -= 1
    end
    yy
end

## return the optimal log-likelihood ratios for target and non-target scores
function optllr{T<:Real}(tar::Vector{T}, non::Vector{T}; laplace=true)
    ntar = length(tar)
    nnon = length(non)
    ntar > 0 && nnon > 0 || error("Lenghts must be nonzero")
    o = sortperm([tar; non])
    pideal = zeros(ntar + nnon + 4laplace)
    ## set pideal=1 for target scores
    for i=1:ntar+nnon
        if o[i] ≤ ntar         # tar sorted before non, a target score
            pideal[i+2laplace] = 1.0
        end
    end
    if laplace
        pideal[1] = pideal[end-1] = 1.0
    end
    ## optimal posterior, given monotonicity constraints
    popt = pav(pideal)          # optimal posterior
    postlo = logit(popt)[1+2laplace:end-2laplace] # posterior log odds
    priorlo = log(length(tar) / length(non))
    llrs = Array(T, ntar+nnon)
    llrs[o] = postlo - priorlo
    llrs[1:ntar], llrs[ntar+1:end]
end


#
# Löwner-John ellipsoid example
#
# Compute the inner and outer ellipsoids of a polytope.
#
# Adapter from example in the MOSEK distro.
#
# The inner ellipsoidal approximation to a polytope
#
#    S = { x \in R^n | Ax < b }.
#
# maximizes the volume of the inscribed ellipsoid,
#
#    { x | x = C*u + d, || u ||_2 <= 1 }.
#
# The volume is proportional to det(C)^(1/n), so the
# problem can be solved as
#
#   maximize         t
#   subject to       t       <= det(C)^(1/n)
#               || C*ai ||_2 <= bi - ai^T * d,  i=1,...,m
#                 C is PSD
#
# which is equivalent to a mixed conic quadratic and semidefinite
# programming problem.
#
#
# The outer ellipsoidal approximation to a polytope given
# as the convex hull of a set of points
#
#     S = conv{ x1, x2, ... , xm }
#
# minimizes the volume of the enclosing ellipsoid,
#
#   { x | || P*(x-c) ||_2 <= 1 }
#
# The volume is proportional to det(P)^{-1/n}, so the problem can
# be solved as
#
#   minimize         t
#   subject to       t       >= det(P)^(-1/n)
#               || P*xi + c ||_2 <= 1,  i=1,...,m
#                 P is PSD.
#
# References:
# [1] "Lectures on Modern Optimization", Ben-Tal and Nemirovski, 2000.
#

using JuMP

# Models the convex set
#
#   S = { (x, t) \in R^n x R | x >= 0, t <= (x1 * x2 * ... * xn)^(1/n) }
#
function geometric_mean(m, x, t)
    n = length(x)
    if n == 1
        @constraint(m,t-x[1] <= 0)
    else
        t2 = @variable(m)
        @constraint(m,1.0*[t2; x[n]; t] in MOI.PowerCone(1-1.0/n))
        geometric_mean(m,x[1:n-1], t2)
    end
end

# Purpose: Models the hypograph of the n-th power of the
# determinant of a positive definite matrix. See [1,2] for more details.
#
#   The convex set (a hypograph)
#
#   C = { (X, t) \in S^n_+ x R |  t <= det(X)^{1/n} },
#
#   can be modeled as the intersection of a semidefinite cone
#
#   [ X, Z; Z^T Diag(Z) ] >= 0
#
#   and a number of rotated quadratic cones and affine hyperplanes,
#
#   t <= (Z11*Z22*...*Znn)^{1/n}  (see geometric_mean).
function det_root(m,X,t)
    n = floor(Int,sqrt(length(X)))

    @variable(m,Y[1:2*n,1:2*n],PSD)
    @constraint(m,reshape(Y[n+1:2*n,1:n] - Y[n+1:2*n,n+1:2*n],n*n) in MOI.Zeros(n*n))
    @constraint(m,reshape(X - Y[1:n,1:n], n*n) in MOI.Zeros(n*n))

    geometric_mean(m,[Y[i+n,i+n] for i in 1:n],t)
end

#  The inner ellipsoidal approximation to a polytope
#
#     S = { x \in R^n | Ax < b }.
#
#  maximizes the volume of the inscribed ellipsoid,
#
#     { x | x = C*u + d, || u ||_2 <= 1 }.
#
#  The volume is proportional to det(C)^(1/n), so the
#  problem can be solved as
#
#    maximize         t
#    subject to       t       <= det(C)^(1/n)
#                || C*ai ||_2 <= bi - ai^T * d,  i=1,...,m
#                C is PSD
#
#  which is equivalent to a mixed conic quadratic and semidefinite
#  programming problem.
function lownerjohn_inner(solver :: MOI.AbstractOptimizer,
                          A      :: Array{Float64,2},
                          b      :: Vector{Float64})
    m,n = size(A)

    model = direct_model(solver)

    @variable(model, t >= 0)
    @variable(model, C[1:n,1:n], PSD)
    @variable(model, d[1:n])

    # [ b-A*d ; A*C ] in QCone

    for i in 1:m
        @constraint(model, [ b[i] - (A[i,:]' * d) ; (A[i,:]'*C)'] in MOI.SecondOrderCone(n+1))
    end

    det_root(model,C,t)

    @objective(model, Max, t)

    optimize(model)

    reshape([JuMP.resultvalue(item) for item in C],(n,n)),[JuMP.resultvalue(item) for item in d]
end

# The outer ellipsoidal approximation to a polytope given
# as the convex hull of a set of points
#
#   S = conv{ x1, x2, ... , xm }
#
# minimizes the volume of the enclosing ellipsoid,
#
#   { x | || P*x-c ||_2 <= 1 }
#
# The volume is proportional to det(P)^{-1/n}, so the problem can
# be solved as
#
#   maximize         t
#   subject to       t       <= det(P)^(1/n)
#               || P*xi - c ||_2 <= 1,  i=1,...,m
#               P is PSD.
function lownerjohn_outer(solver :: MOI.AbstractOptimizer,
                          x      :: Array{Float64,2})
    m,n = size(x)
    model = direct_model(solver)

    @variable(model,t >= 0)
    @variable(model,P[1:n,1:n],PSD)
    @variable(model,c[1:n])

    # (1, Px-c) in cone
    for i in 1:m
        @constraint(model,[1 ; (x[i,:]' * P)' - c] in MOI.SecondOrderCone(n+1))
    end

    # t <= det(P)^{1/n}
    det_root(model,P,t)

    @objective(model,Max,t)

    optimize(model)

    [ JuMP.resultvalue(item) for item in P ], [ JuMP.resultvalue(item) for item in c ]
end

# Probably do something else here:
using MosekTools
solver = Mosek.Optimizer

# Test

# Vertices of a pentagon in 2D
p = [0.0  0.0;
     1.0  3.0;
     5.5  4.5;
     7.0  4.0;
     7.0  1.0;
     3.0 -2.0]

n = size(p,1)
# The hyperplane representation of the same polytope
A = [ (-p[:,2] + [p[n,2];p[1:n-1,2]])  (p[:,1] - [p[n,1];p[1:n-1,1]]) ]
b = A[:,1] .* p[:,1] + A[:,2] .* p[:,2]

Po, co = lownerjohn_outer(solver(), p)
Ci, di = lownerjohn_inner(solver(), A, b)
Poinv = Po ^ -1

println("Solution:")
println("  Outer ellipsoid: Po = $Po, co = $co")
println("  Inner ellipsoid: Ci = $Ci, di = $di")

if Pkg.installed("Plots") != nothing && Pkg.installed("GR") != nothing
    println("Plotting solution...")
    import Plots
    Plots.gr()
    Plots.plot([p[:,1] ; p[1,1]], [p[:,2] ; p[1,2]])

    Plots.plot!(t -> Ci[1,:]' * [ cos(t),sin(t) ] + di[1],
                t -> Ci[2,:]' * [ cos(t),sin(t) ] + di[2],
                0,2*pi)

    Plots.plot!(t -> Poinv[1,:]' * ( [cos(t),sin(t)] + c ),
                t -> Poinv[2,:]' * ( [cos(t),sin(t)] + c ),
                0,2*pi)

    println("Plot done. Press enter to continue.")
    readline()
else
    println("Install Plots package to plot solution")
end


module GreekModule
   σ(x) = 1 ./ (1 + exp.(-x))
   logσ(x) = - log1p.(exp.(-x))
end


export 𝒯hcubature, generate_𝒯integrand

function generate_𝒯integrand(func, t; args=())
    function integrand(x)
        y = 1.0
        sv = similar(x)
        sv[1] = x[1]*t
        for i in 2:length(x)
            sv[i] = x[i]*sv[i-1]
            y = y*sv[i-1]
        end
        func(sv, args...)*y*t
    end
end

function 𝒯hcubature(f, t, ndim, args=(); rtol=1e-6, atol=1e-8)
    integrand = generate_𝒯integrand(f, t, args=args)
    a = zeros(ndim)
    b = ones(ndim)
    hcubature(integrand, a, b, rtol=rtol, atol=atol)
end


module NnetTest

using Learn
# using StatPlots; glvisualize(leg=false, show=true)

function doit()

    nin = 10
    nout = 1
    τ = 1000
    w1 = randn(nout, nin)
    w2 = randn(nout, nin)
    # f(x) = tanh(w1*x + w2*sin(20x) + 0.05randn(nout))[1] > 0.0 ? 1.0 : 0.0
    f(x) = tanh(10w1*x)[1] > 0.0 ? 1.0 : 0.0
    inputs = randn(nin,τ)
    targets = zeros(nout,τ)
    for (i,x) in enumerate(DataSubset(inputs, 1:nobs(inputs)))
        targets[i] = f(x)
    end
    # scatter(inputs', targets', layout=nin, m=(3,0.2,stroke(0)), smooth=true)

    # build our model
    nh = 10
    chain = Chain(
        Affine(nin,nh),
        Activation(:relu,nh),
        Affine(nh,nout),
        Activation(:logistic,nout)
    )
    obj = objective(chain, CrossentropyLoss(), L1Penalty(1e-6))

    learner = make_learner(
        GradientLearner(FixedLR(1e-4), SGD()),
        maxiter = 10000,
        oniter = (model,i) -> begin
            @show i
        end
    )
    learn!(obj, learner, MiniBatches((inputs, targets), 5))

    # eval
    for (i,(input,target)) in enumerate(eachobs((inputs,targets)))
        @show i,transform!(chain,input),target
    end

    # store for inspection
    obj,learner
end
obj,learner = doit()
end #module


type SDF_PID_VD <: SDFDynamicSystem
    omg::Float64
    zta::Float64
    mu::Float64
    xi::Float64
    kpa::Float64
    alpha::Float64
    u0::Float64
    v0::Float64
    ud0::Float64
    vd0::Float64

    SDF_PID_VD(omg::Float64, zta::Float64, mu::Float64, xi::Float64, kpa::Float64, alpha::Float64) = new(omg, zta, mu, xi, kpa, alpha, 0.0, 0.0, 0.0, 0.0)
    SDF_PID_VD(omg::Float64, zta::Float64, mu::Float64, xi::Float64, kpa::Float64) = new(omg, zta, mu, xi, kpa, 1.0, 0.0, 0.0, 0.0, 0.0)
end

function systemEquation(s::SDF_PID_VD, p::Function)
    function dydt(y::AbstractArray{Float64,1}, t::Float64)
        u, v, ud, vd = y
        omg2 = s.omg*s.omg
        dvdt = p(t)-2.0*s.zta*s.omg*v-omg2*u-s.kpa*omg2*(u-ud)

        if s.alpha == 1.0
            dvddt = (s.kpa*omg2*(u-ud) - 2.0*s.xi*s.omg*vd)/s.mu
        else
            dvddt = (s.kpa*omg2*(u-ud) - 2.0*s.xi*s.omg*sign(vd)*abs(vd)^s.alpha)/s.mu
        end
        [v; dvdt; vd; dvddt]
    end
    dydt
end

function response(s::SDF_PID_VD, p::Function, t::AbstractArray{Float64,1})
    y0 = [s.u0; s.v0; s.ud0; s.vd0]

    f = systemEquation(s, p)

    y = ode4(f, y0, t)

    u = y[:,1]
    v = y[:,2]
    ud = y[:,3]
    vd = y[:,4]
    fdz = 2.0*s.zta*s.omg*v
    fdx = s.kpa*s.omg*s.omg*(u-ud)
    fdd = 2.0*s.xi*s.omg*sign(vd).*abs(vd).^s.alpha
    a = map(p,t)-fdz-fdx-s.omg*s.omg*u

    [u v a fdz fdx ud fdd]
end


# This is a minimal working example for running LMlovo on several
# workers in Julia and saving them into a shared matrix

using Random
using Revise
using Distributed
using SharedArrays

@everywhere using RAFF
@everywhere using Base.CoreLogging

@everywhere CoreLogging._min_enabled_level[] = CoreLogging.Warn + 1

@everywhere gmodel!(x, t_, g) = begin
    g[1] = exp(t_ * x[2])
    g[2] = t_ * x[1] * exp(t_ * x[2]);
end

@everywhere model(x, t) = x[1] * exp(t * x[2])

n = 2

np = 1000

p = 700

# Set the seed for generating the same data
Random.seed!(123456789)

data, xSol = RAFF.generateNoisyData(model, n, np, p)

result = SharedArray{Float64, 2}(n, np)

f = @sync @distributed for i = 1:np

    # Starting point
    x = zeros(Float64, 2)

    # Call function and store results
    s, x, iter, p, f = LMlovo(model, gmodel!, x, data, 2, i, MAXITER=1000)

    result[:, i] .= x

    println("Finished. p = $(p) and f = $(f).")

end


using AirSeaFluxes
using Test

@testset "AirSeaFluxes.jl" begin
    mld=10 #mixed layer depth (m)
    tim=86400*30 #relaxation time scale (s)
    pisvel=mld/tim #piston velocity (m/s)
    Co=0.0 #ocean value (e.g. concentration of some compound)
    Ca=1.0 #atmospeheric value (e.g. equivalent compound concentration)
    flx=simpleflux(Ca,Co,pisvel)
    @test isapprox(flx,3.858024691358025e-6; atol=1e-2)

    atemp=300.;aqh=0.001;speed=1.; sst=10.
    all=bulkformulae(atemp,aqh,speed,sst)
    @test isapprox(all["hl"],-3.0606099804357885; rtol=1e-2)
    @test isapprox(all["hs"],2.0282408526727473; rtol=1e-2)
    @test isapprox(all["evap"],1.2244888899523058e-9; rtol=1e-2)
    @test isapprox(all["tau"],0.00915548218468587; rtol=1e-2)
end


part1(input) = flash(input, part = 1)
part2(input) = flash(input, part = 2)

function flash(input; part)
    octopi = reduce(hcat, (parse.(Int, collect(line))
                           for line in eachline(input)))
    num_flashes = 0
    i = 0
    while true
        i += 1
        octopi .+= 1
        while true
            flashes = octopi .> 9
            (n = count(flashes)) == 0 && break
            num_flashes += n
            for Δi in -1:1, Δj in -1:1
                Δi == Δj == 0 && continue
                octopi[max(1 + Δi, 1):min(end + Δi, end),
                       max(1 + Δj, 1):min(end + Δj, end)] .+=
                    flashes[max(1 - Δi, 1):min(end - Δi, end),
                            max(1 - Δj, 1):min(end - Δj, end)]
            end
            octopi[flashes] .= -10
        end
        part == 1 && i == 100 && return num_flashes
        octopi .= max.(0, octopi)
        part == 2 && all(==(0), octopi) && return i
    end
end


using LinearAlgebra, Test, Random
using Knet.KnetArrays: A_mul_Bt, At_mul_B, At_mul_Bt, KnetArray
using Knet.Ops20: mat
using CUDA: CUDA, functional
using AutoGrad: gradcheck
include("combinatorics.jl")

#Random.seed!(42)
nsample(a,n)=collect(a)[randperm(length(a))[1:n]]

@testset "linalg" begin
    for t in (Float32,Float64)
        a = rand(t,3,5)
        b = rand(t,5,2)
        mmul(w)=w[1]*w[2]
        @test gradcheck(mmul, (a,b))
        if CUDA.functional()
            c = a * b
            ka = KnetArray(a)
            kb = KnetArray(b)
            kc = ka * kb
            @test isapprox(c, Array(kc))
            @test gradcheck(mmul, (ka,kb))
            # Issue #456: 2nd derivative for MLP
            mmulABt(w)=A_mul_Bt(w[1],w[2])
            mmulAtB(w)=At_mul_B(w[1],w[2])
            mmulAtBt(w)=At_mul_Bt(w[1],w[2])
            @test gradcheck(mmulABt, (ka,kb'))
            @test gradcheck(mmulAtB, (ka',kb))
            @test gradcheck(mmulAtBt, (ka',kb'))
        end

        if CUDA.functional()
            # cannot gradcheck axpy!, scal! overwriting
            d = rand(t,3,5)
            kd = KnetArray(d)
            r = rand()
            @test isapprox(axpy!(r,a,d), axpy!(r,ka,kd))
            @test isapprox(lmul!(r,d), lmul!(r,kd))
            @test isapprox(rmul!(d,r), rmul!(kd,r))
        end

        @test gradcheck(transpose, a)
        at = a'
        bt = b'
        s = a[1]
        mmul1(w)=w[1]'*w[2]
        mmul2(w)=w[1]*w[2]'
        mmul3(w)=w[1]'*w[2]'
        @test gradcheck(mmul1, Any[at,b])
        @test gradcheck(mmul2, Any[a,bt])
        @test gradcheck(mmul3, Any[at,bt])
        @test gradcheck(mmul1, Any[a,s])
        @test gradcheck(mmul1, Any[s,b])
        @test gradcheck(mmul2, Any[a,s])
        @test gradcheck(mmul2, Any[s,b])
        @test gradcheck(mmul3, Any[a,s])
        @test gradcheck(mmul3, Any[s,b])
        if CUDA.functional()
            kat = ka'
            kbt = kb'
            @test isapprox(at, Array(kat))
            @test gradcheck(transpose, ka)
            @test isapprox(kat'*kb, at'*b)
            @test isapprox(ka*kbt', a*bt')
            @test isapprox(kat'*kbt', at'*bt')
            @test gradcheck(mmul1, Any[kat,kb])
            @test gradcheck(mmul2, Any[ka,kbt])
            @test gradcheck(mmul3, Any[kat,kbt])
            @test gradcheck(mmul1, Any[ka,s])
            @test gradcheck(mmul1, Any[s,kb])
            @test gradcheck(mmul2, Any[ka,s])
            @test gradcheck(mmul2, Any[s,kb])
            @test gradcheck(mmul3, Any[ka,s])
            @test gradcheck(mmul3, Any[s,kb])
        end

        @test gradcheck(mat, a)
        if CUDA.functional()
            @test isapprox(mat(a), Array(mat(ka)))
            @test gradcheck(mat, ka)
        end

        for p in collect(permutas(1:2))
            p2(x) = permutedims(x,p)
            @test gradcheck(p2, a)
            if CUDA.functional()
                @test isapprox(p2(a), Array(p2(ka)))
                @test gradcheck(p2, ka)
            end
        end

        a3 = rand(2,3,4)
        if CUDA.functional(); k3 = KnetArray(a3); end
        for p in collect(permutas(1:3))
            p3(x) = permutedims(x,p)
            @test gradcheck(p3, a3)
            if CUDA.functional()
                @test isapprox(p3(a3), Array(p3(k3)))
                @test gradcheck(p3, k3)
            end
        end

        a4 = rand(2,3,4,5)
        if CUDA.functional(); k4 = KnetArray(a4); end
        for p in nsample(permutas(1:4),6)
            p4(x) = permutedims(x,p)
            @test gradcheck(p4, a4)
            if CUDA.functional()
                @test isapprox(p4(a4), Array(p4(k4)))
                @test gradcheck(p4, k4)
            end
        end

        a5 = rand(2,3,4,5,6)
        if CUDA.functional(); k5 = KnetArray(a5); end
        for p in nsample(permutas(1:5),6)
            p5(x) = permutedims(x,p)
            @test gradcheck(p5, a5)
            if CUDA.functional()
                @test isapprox(p5(a5), Array(p5(k5)))
                @test gradcheck(p5, k5)
            end
        end

        # transpose and matmul should work with vectors
        if CUDA.functional()
            a6 = rand(t,3)
            b6 = rand(t,3,3)
            c6 = KnetArray(a6)
            d6 = KnetArray(b6)
            @test c6' ≈ a6'
            @test c6' * c6 ≈ a6' * a6 # (a6'*a6)::Number, (Array(a6')*a6)::Vector (not Matrix), (c6'*c6)::Number.
            @test c6 * c6' ≈ a6 * a6'
            @test b6 * a6 ≈ d6 * c6 # Matrix * Vector => Vector
        end
    end
end

nothing


using LinearAlgebra, ForwardDiff, StaticArrays
using IterativeSolvers
using Plots

"""
    particle dynamics
    - 3D particle subject to contact forces

    - configuration: q = (x, y, z) ∈ R³
    - impact force (magnitude): n ∈ R₊
    - friction force: β ∈ R⁴₊
        - contact force: λ = (β, n) ∈ R⁴₊ × R₊
        - friction coefficient: μ ∈ R₊

    Discrete Mechanics and Variational Integrators
        pg. 363
"""
struct Particle
      m # mass
      g # gravity
      μ # friction coefficient

      nq # configuration dimension
end

# mass matrix
function mass_matrix(model)
    m = model.m

    Diagonal(@SVector [m, m, m])
end

# gravity
function gravity(model, q)
    m = model.m
    g = model.g

    @SVector [0.0, 0.0, m * g]
end

# signed distance function
function signed_distance(model, q)
    q[3]
end

# contact force Jacobian
function jacobian(model, q)
    Diagonal(@SVector ones(model.nq))
end

# dynamics
function dynamics(model, q1, q2, q3, λ, h)
    nq = model.nq
    SVector{nq}(mass_matrix(model) * (2.0 * q2 - q1 - q3) / h
        - h * gravity(model, q2)
        + h * jacobian(model, q3)' * λ)
end

"""
    step particle system
        solves 1-step feasibility problem
"""
function _step(q1, q2, h;
    tol = 1.0e-8, z0_scale = 0.001, step_type = :gmres, max_iter = 100)
    # 1-step optimization problem:
    #     find z
    #     s.t. r(z) = 0
    #
    # z = (q, n, b, ψ, η, s)
    #     ψ, η are dual variables for friction subproblem
    #     s are slack variables for convenience

    # initialize
    z = 1.0e-1 * ones(15)
    z[1:3] = copy(q2)

    # if z[3] < 0.1
    #     z[3] = 0.1
    # end

    ρ = 1.0 # barrier parameter
    flag = false

    for k = 1:5
        function r(z)
            # system variables
            q3 = view(z, 1:3)
            n = z[4]
            b = view(z, 5:8)

            λ = [(b[1] - b[3]); (b[2] - b[4]); n] # contact forces
            ϕ = signed_distance(model, q3)        # signed-distance function
            vT = (view(q3, 1:2) - view(q2, 1:2)) ./ h

            # friction subproblem variables
            ψ = z[9]
            η = view(z, 10:13)

            # slack
            s1 = z[14]
            s2 = z[15]

            # action optimality conditions
            [dynamics(model, q1, q2, q3, λ, h);
             s1 - ϕ;
             n * s1 - ρ;

             # maximum dissipation optimality conditions
             [vT; -vT] + ψ * ones(4) - η;
             s2 - (model.μ * n - sum(b));
             ψ * s2 - ρ;
             b .* η .- ρ]
        end

        # Jacobian
        function R(z)
            # differentiate r
            _R = ForwardDiff.jacobian(r, z)

            return _R
        end

        function check_variables(z)
            # q3 = view(z, 1:3)
            n = z[4]
            b = view(z, 5:8)

            # friction subproblem variables
            ψ = z[9]
            η = view(z, 10:13)
            s1 = z[14]
            s2 = z[15]

            if n <= 0.0
                return true
            end

            if any(b .<= 0.0)
                return true
            end

            if any(η .<= 0.0)
                return true
            end

            if ψ <= 0.0
                return true
            end

            if s1 <= 0.0
                return true
            end

            if s2 <= 0.0
                return true
            end

            return false
        end

        extra_iters = 0

        for i = 1:max_iter
            # compute residual, residual Jacobian
            res = r(z)

            if norm(res) < tol
                println("   iter ($i) - norm: $(norm(res))")
                # return z, true
                flag = true
                continue
            end

            jac = R(z)

            # compute step
            # if step_type == :gmres
            #     Δ = gmres(jac, res, abstol = 1.0e-12, maxiter = i + extra_iters)
            # else
            # Δ = (jac' * jac + 0.0e-6 * I) \ (jac' * res) # damped least-squares direction
            Δ = jac \ res
            # end

            # line search the step direction
            α = 1.0

            iter = 0
            while check_variables(z - α * Δ) # backtrack inequalities
                α = 0.5 * α
                # println("   α = $α")
                iter += 1
                if iter > 50
                    @error "backtracking line search fail"
                    flag = false
                    @show n = (z - α * Δ)[4]
                    @show b = view(z - α * Δ, 5:8)
                    @show ψ = (z - α * Δ)[9]
                    @show η = view(z - α * Δ, 10:13)
                    @show s = (z - α * Δ)[14:15]
                    return z, false
                end
            end

            while norm(r(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(res)^2.0
                α = 0.5 * α
                # println("   α = $α")
                iter += 1
                if iter > 50
                    @error "line search fail"
                    flag = false
                    return z, false
                end
            end

            # update
            z .-= α * Δ
        end

        ρ = 0.1 * ρ
        println("ρ: $ρ")
    end

    return z, flag
end

"""
    simulate
    - solves 1-step feasibility problem for T time steps
    - initial configurations: q1, q2 (note this can encode initial velocity)
    - time step: h
"""
function simulate(q1, q2, T, h;
    z0_scale = 0.0, step_type = :gmres)
    println("simulation")

    # initialize histories
    q = [q1, q2]
    n = [0.0]
    b = [zeros(4)]

    # step
    for t = 1:T
        println("   t = $t")
        z_sol, status = _step(q[end-1], q[end], h,
            z0_scale = z0_scale, step_type = step_type)

        if !status
            @error "failed step (t = $t)"
            return q, n, b
        else
            push!(q, view(z_sol, 1:3))
            push!(n, z_sol[4])
            push!(b, view(z_sol, 5:8))
        end
    end

    return q, n, b
end

# simulation setup
model = Particle(1.0, 9.81, 1.0, 3)
h = 0.01

# initial conditions
# v1 = [1.0; 1.0; 0.0]
# q1 = [0.0; 0.0; 1.0]

v1 = [1.0; 10.0; 0.0]
q1 = [0.0; 0.0; 1.0]

v2 = v1 - gravity(model, q1) * h
q2 = q1 + 0.5 * (v1 + v2) * h

q_sol, y_sol, b_sol = simulate(q1, q2, 1000, h)

plot(hcat(q_sol...)[3:3, :]', xlabel = "", label = "z")
plot!(h .* hcat(y_sol...)', xlabel = "", label = "n", linetype = :steppost)

plot(hcat(q_sol...)[1:2, :]', xlabel = "", label = ["x" "y"])
plot!(h * hcat(b_sol...)', label = ["b1" "b2" "b3" "b4"], linetype = :steppost)


using DelimitedFiles
movements = readdlm("subdirections_t.txt",' ')

@show movements

function movesub(movements)

    depth = 0
    position = 0
    aim = 0

    for move in eachindex(movements)
        if move > length(movements)/2
            break
        end

        direction = movements[move,1]
        distance = movements[move,2]
        @show direction, distance

        if direction == "forward"
            position += distance
            depth += (aim*distance)
        elseif direction == "down"
            aim += distance
        elseif direction == "up"
            aim -= distance
        end
#        @show depth,position
    end

    @show depth*position
end

movesub(movements)


# returns vector of function value and gradient, using diff results
function value_gradient(f, x::AbstractVector, diffresults::Val{true})#  = Val(true))
    r = DiffResults.GradientResult(x) # this could take pre-allocated temporary storage
    ForwardDiff.gradient!(r, f, x)
    return vcat(r.value, r.derivs[1]) # d+1 # ... to avoid this
end
function value_gradient(f, x::AbstractVector, diffresults::Val{false} = Val(false))
    ∇ = ForwardDiff.gradient(f, x)
    return vcat(f(x), ∇)
end

function value_gradient_hessian(f, x::AbstractVector)
    d = length(x)
    dd = d^2 + d + 1
    fx = similar(x, dd)
    fx[1] = f(x)
    grad = @view fx[2:d+1]
    hess = @view fx[d+2:end]
    results = DiffResults.DiffResult(grad, hess)
    g(x) = ForwardDiff.gradient(f, x)
    ForwardDiff.jacobian!(results, g, x)
    # IDEA: try to minimize evaluations to a single pass
    # hessian = reshape(hessian, d, d) # needs to be d + 1  by d
    # result = DiffResults.DiffResult(valgrad, hessian)
    # println(result)
    # ForwardDiff.jacobian!(result, x->value_gradient(f, x), x)
    return fx

end

function value_derivative(f, x::Real, diffresults::Val{true} = Val(true))
    r = DiffResults.DiffResult(zero(x), zero(x)) # this could take pre-allocated temporary storage
    r = ForwardDiff.derivative!(r, f, x)
    return vcat(r.value, r.derivs[1])
end
function value_derivative(f, x::Real, diffresults::Val{false})
    ∇ = ForwardDiff.derivative(f, x)
    return vcat(f(x), ∇)
end

################################################################################
issquare(A::AbstractMatrix) = size(A, 1) == size(A, 2)
issquare(A::Real) = true
# converts y in various forms to a vector of numbers
_vec(Y::AbstractMatrix) = vec(Y)
_vec(y::AbstractVector{<:Number}) = y
# this needs some thinking if we want to make static arrays efficient
function _vec(y::VecOfVec{<:Number}) # different behavior to vec!
	d = length(y[1])
	u = zeros(length(y)*d)
	for i in eachindex(y)
		u[d*(i-1)+1:d*i] .= y[i]
	end
	return u
end

########################### Laplace Approximation ##############################
# IDEA: expand this into variational approximations (mean field etc.)
# IDEA: connect with GP classification
# ForwardDiff.hessian(f, x::Real) = FD.derivative(y->FD.derivative(f, y), x)
# approximates a pdf f with a normal distribution centered at x using laplace's method
# TODO: height of Gaussian?
function laplace(f, x::AbstractVector)
    H = ForwardDiff.hessian(y->log(f(y)), x)
    Gaussian(x, inverse(H))
end


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ b88059f0-f793-11ea-0031-23a60382d51a
using Pkg

# ╔═╡ b8809a8c-f793-11ea-382c-5763b1d11ca5
begin
	using Optim
	using StatisticalRethinking
end

# ╔═╡ fdf51ec0-f794-11ea-327a-33958a9cd52a
md"## Intro-logpdf-2s.jl"

# ╔═╡ 3330c924-f793-11ea-28b7-6d61030db6f9
md"##### This script shows clip-04-26-29s.jl using Optim and a loglik function."

# ╔═╡ b8812010-f793-11ea-19d0-61be51252ae5
md"## snippet 4.26"

# ╔═╡ 5d5852d4-34cf-11eb-2618-4b83ea072b1e
begin
	df = DataFrame(CSV.read(sr_datadir("Howell1.csv"), DataFrame; delim=';'))
	df2 = filter(row -> row[:age] >= 18, df)
end;

# ╔═╡ b88ee9c0-f793-11ea-08c9-e9b05a6e3990
md"## snippet 4.27"

# ╔═╡ 4433d922-34cf-11eb-379a-cfe1dbd939cb
# Our first model:

stan4_1 = "
  μ ~ Normal(178,20) # prior
  σ ~ Uniform(0, 50) # prior
  height ~ Normal(μ, σ) # likelihood
";

# ╔═╡ b898fb18-f793-11ea-202f-672f3b026e96
md"## snippet 4.28"

# ╔═╡ b8998394-f793-11ea-021c-bf822c791396
function loglik(x)
  ll = 0.0
  ll += log(pdf(Normal(178, 20), x[1]))
  ll += log(pdf(Uniform(0, 50), x[2]))
  ll += sum(logpdf.(Normal(x[1], x[2]), df2.height))
  -ll
end

# ╔═╡ b8a3c502-f793-11ea-34f8-eddc53070e3b
md"## snippet 4.29"

# ╔═╡ 9e16251c-34cf-11eb-2af5-21ac06afd361
# Start values

begin
	lower = [0.0, 0.0]
	upper = [250.0, 50.0]
	x0 = [170.0, 10.0]
end

# ╔═╡ b8bcaa18-f793-11ea-3e04-0fa9333a1435
res = optimize(loglik, lower, upper, x0)

# ╔═╡ b8bd8988-f793-11ea-0f93-4bf863950fe3
Optim.minimizer(res)

# ╔═╡ b8cab216-f793-11ea-1880-9dc5d6a84e79
# Our second model:

stan4_2 = "
  μ ~ Normal(178, 0.1) # prior
  σ ~ Uniform(0, 50) # prior
  height ~ Normal(μ, σ) # likelihood
";

# ╔═╡ 9c55acca-f794-11ea-26ae-c505f54777f2
function loglik2(x)
  ll = 0.0
  ll += log(pdf(Normal(178, 0.1), x[1]))
  ll += log(pdf(Uniform(0, 50), x[2]))
  ll += sum(logpdf.(Normal(x[1], x[2]), df2.height))
  -ll
end

# ╔═╡ b8dd1d66-f793-11ea-1210-99abfb4919c9
x1 = [178.0, 40.0] # Initial values can't be outside the the box.

# ╔═╡ b8ecbb40-f793-11ea-219a-4d4d89afbede
res2 = optimize(loglik2, lower, upper, x1)

# ╔═╡ d9e99bd8-34ce-11eb-1779-231c74d97cd1
Optim.minimizer(res2)

# ╔═╡ b8f4b8ec-f793-11ea-1f5e-6d1225f2da1a
md"##### Notice the increase of σ."

# ╔═╡ b8fc6e14-f793-11ea-275c-35a406bdaefc
md"## End of intro-logpdfs-2.jl"

# ╔═╡ Cell order:
# ╟─fdf51ec0-f794-11ea-327a-33958a9cd52a
# ╟─3330c924-f793-11ea-28b7-6d61030db6f9
# ╠═b88059f0-f793-11ea-0031-23a60382d51a
# ╠═b8809a8c-f793-11ea-382c-5763b1d11ca5
# ╟─b8812010-f793-11ea-19d0-61be51252ae5
# ╠═5d5852d4-34cf-11eb-2618-4b83ea072b1e
# ╠═b88ee9c0-f793-11ea-08c9-e9b05a6e3990
# ╠═4433d922-34cf-11eb-379a-cfe1dbd939cb
# ╟─b898fb18-f793-11ea-202f-672f3b026e96
# ╠═b8998394-f793-11ea-021c-bf822c791396
# ╟─b8a3c502-f793-11ea-34f8-eddc53070e3b
# ╠═9e16251c-34cf-11eb-2af5-21ac06afd361
# ╠═b8bcaa18-f793-11ea-3e04-0fa9333a1435
# ╠═b8bd8988-f793-11ea-0f93-4bf863950fe3
# ╠═b8cab216-f793-11ea-1880-9dc5d6a84e79
# ╠═9c55acca-f794-11ea-26ae-c505f54777f2
# ╠═b8dd1d66-f793-11ea-1210-99abfb4919c9
# ╠═b8ecbb40-f793-11ea-219a-4d4d89afbede
# ╠═d9e99bd8-34ce-11eb-1779-231c74d97cd1
# ╟─b8f4b8ec-f793-11ea-1f5e-6d1225f2da1a
# ╟─b8fc6e14-f793-11ea-275c-35a406bdaefc


# fitting

"""
    struct FittedCone{A<:AbstractArray, R<:Real} <: FittedShape

Cone primitive, defined by its apex,
its axis (that points from the apex towards the opening),
and its opening angle in radians.
Also stored, if the normals point outwards of the shape.
"""
struct FittedCone{A<:AbstractArray, R<:Real} <: FittedShape
    apex::A
    # points from apex towards the opening
    # unit length
    axis::A
    # opening angle in radians
    opang::R
    outwards::Bool
end

Base.show(io::IO, x::FittedCone) =
    print(io, """cone, ω: $(x.opang)""")

Base.show(io::IO, ::MIME"text/plain", x::FittedCone{A, R}) where {A, R} =
    print(io, """FittedCone{$A, $R}\napex: $(x.apex), axis: $(x.axis), ω: $(x.opang), $(x.outwards ? "outwards" : "inwards" )""")

strt(x::FittedCone) = "cone"

function defaultshapeparameters(::Type{FittedCone})
    return (cone=(ϵ=0.3, α=deg2rad(5), minconeopang=deg2rad(2)),)
end

function setconeOuterity(fc, b)
    FittedCone(fc.apex, fc.axis, fc.opang, b)
end

## fitting

function fit3pointcone(p, n)
    # rank of the coefficient matrix
    ftype = eltype(eltype(p))
    r = Array{ftype,2}(undef, (3,3))
    for i in 1:3; for j in 1:3; r[i,j] = n[i][j]; end; end
    rank(r) == 3 || return nothing
    ds = [dot(p[i], n[i]) for i in 1:3]
    # rank of the augmented matrix
    rv = hcat(r, -1 .* ds)
    rank(rv) == 3 || return nothing
    # apex
    ap = SVector{3, ftype}(r\ds)
    # axis
    axis3p = [ap+((p[i]-ap)/norm(p[i]-ap)) for i in 1:3]
    ax = normalize(cross(axis3p[2]-axis3p[1], axis3p[3]-axis3p[1]))
    midp = sum(axis3p)/3
    dirv = normalize(midp-ap)
    ax = dot(ax, dirv) < 0 ? -1*ax : ax
    # opening angle
    angles = [acos(clamp(dot(normalize(p[i]-ap), ax), -1, 1)) for i in 1:3]
    opangle = 2*sum(angles)/3
    return FittedCone(ap, ax, opangle, true)
end

"""
    project2cone(cone, p)

Return the distance from the cone's surface and normal of the surface at that point.
"""
function project2cone(cone, p)
    # vector to the point from the apex
    to_point = cone.apex-p
    to_pointn = normalize(to_point)
    # rotation axis defined by the above vector and the axis of the cone
    rot_ax = normalize(cross(cone.axis, to_pointn))
    # normal: from the rotation axis and the axis fo the cone
    # need to be rotated to get the normal of the cone
    comp_n = normalize(cross(cone.axis, rot_ax))
    # rotation matrix from the rotation axis and opening angle
    rM = rodriguesrad(rot_ax, -cone.opang/2)
    # normal of the cone at the point
    current_normal = normalize(rM*comp_n)
    # distance from the surface of the cone to the point
    dist = dot(-current_normal, -to_point)
    # return the distance and the normal
    return (dist, current_normal)
end

function validatecone(cone, ps, ns, params)
    #@unpack α_cone, ϵ_cone, minconeopang = params
    @extract params : params_cone=cone
    @extract params_cone : α_cone=α ϵ_cone=ϵ minconeopang
    calcs = [project2cone(cone, ps[i]) for i in eachindex(ps)]
    for i in eachindex(calcs)
        if calcs[i][1] > ϵ_cone
            return nothing
        end
    end

    #cone with small opening angle is filtered
    cone.opang < minconeopang && return nothing

    lp = size(ps, 1)
    norm_ok = falses(lp)
    invnorm_ok = falses(lp)

    thr = cos(α_cone)
    for i in eachindex(calcs)
        dotp = dot(calcs[i][2], ns[i])
        norm_ok[i] = dotp > thr
        invnorm_ok[i] = dotp < -thr
    end

    norm_ok == trues(lp) && return setconeOuterity(cone, true)
    invnorm_ok == trues(lp) && return setconeOuterity(cone, false)
    return nothing
end

"""
    fit(::Type{FittedCone}, p, n, pc, params)

Fit a cone to 3 points. Normals are expected to be normalized.
Return `nothing` if points do not fit to a cone.
"""
function fit(::Type{FittedCone}, p, n, pc, params)
    fcone = fit3pointcone(p, n)
    fcone === nothing && return nothing
    valid_cone = validatecone(fcone, p, n, params)
    return valid_cone
end

## scoring

function compatiblesCone(cone, points, normals, params)
    #@unpack α_cone, ϵ_cone = params
    @extract params : params_cone=cone
    @extract params_cone : α_cone=α ϵ_cone=ϵ
    calcs = (project2cone(cone, points[i]) for i in eachindex(points))

    # eps check
    # c1 = [abs(calcs[i][1]) < ϵ_cone for i in eachindex(calcs)]
    #if cone.outwards
    #    c2=[isparallel(calcs[i][2], normals[i], α_cone) && c1[i] for i in eachindex(calcs)]
    #else
    #    c2=[isparallel(-calcs[i][2], normals[i], α_cone) && c1[i] for i in eachindex(calcs)]
    #end

    zcn = zip(calcs, normals)
    if cone.outwards
        c2=[isparallel(c[2], n, α_cone) && (abs(c[1]) < ϵ_cone) for (c,n) in zcn]
    else
        c2=[isparallel(-c[2], n, α_cone) && (abs(c[1]) < ϵ_cone) for (c,n) in zcn]
    end
    return c2
end

function scorecandidate(pc, candidate::FittedCone, subsetID, params)
    ps = @view pc.vertices[pc.subsets[subsetID]]
    ns = @view pc.normals[pc.subsets[subsetID]]
    ens = @view pc.isenabled[pc.subsets[subsetID]]

    cp = compatiblesCone(candidate, ps, ns, params)
    # verti: összes pont indexe, ami enabled és kompatibilis
    # lenne, ha működne, de inkább a boolean indexelést machináljuk
    inder = cp.&ens
    inpoints = (pc.subsets[subsetID])[inder]
    score = estimatescore(length(pc.subsets[subsetID]), pc.size, length(inpoints))
    return (score, inpoints)
end

## refit

"""
    refit(s::T, pc, params) where {T<:FittedCone}

Refit cone.
"""
function refit(s::T, pc, params) where {T<:FittedCone}
    # TODO: use octree for that
    p = @view pc.vertices[pc.isenabled]
    n = @view pc.normals[pc.isenabled]
    cp = compatiblesCone(s, p, n, params)
    return ExtractedShape(s, ((1:pc.size)[pc.isenabled])[cp])
end


# Rational number reconstruction implementation borrowed from CLUE
# and modified a bit to suit the 'Modern Computer Algebra' definitions
# Returns a rational r // h of QQ field in a canonical form such that
#   r // h ≡ a (mod m)
#
# let n = max( λ(a), λ(m) ) , where λ(x) is a number of bits for x
# O(n^2)
function rational_reconstruction(a::I, m::I) where {I<:Union{Int, BigInt}}
    bnd = sqrt(float(m) / 2)

    steps = 0

    # @warn "" bnd
    U = (I(1), I(0), m)
    V = (I(0), I(1), a)
    while abs(V[3]) >= bnd
        q = div(U[3], V[3])
        T = U .- q .* V
        U = V
        V = T

        steps += 1
    end

    t = abs(V[2])
    r = V[3] * sign(V[2])

    # @warn "" r t

    # changed from `<= bnd` to `<= m / bnd`
    # we can speed up this !
    if t <= bnd && gcd(r, t) == 1
        # return (r, t)
        return (true, steps)
    end

    return (false, steps)
end

function maximal_quotient_rational_reconstruction(a::I, m::I) where {I<:Union{Int, BigInt}}

    steps = 0

    T = 2^20 * log2(m)

    n, d = 0, 0

    t0, r0 = 0, m
    t1, r1 = 1, a

    U = (I(1), I(0), m)
    V = (I(0), I(1), a)
    while r1 != 0 && r0 > T
        q = div(r0, r1)

        if q > T
            n, d, T = r1, t1, q
        end

        steps += 1

        r0, r1 = r1, r0 - q*r1
        t0, t1 = t1, t0 - q*t1
    end

    if d == 0 || gcd(n, d) != 1
        return (false, steps)
    end

    if d < 0
        n, d = -n, -d
    end

    return (true, steps)
end


facts("Sparse (semi-)nonnegative Tucker decomposition") do

context("nonnegative decomposition") do
    # An example of nonnegative and semi-nonnegative Tucker decomposition
    tucker_orig = rand_tucker((4, 5, 6), (40, 50, 60), factors_nonneg=true, core_nonneg=true)
    tnsr_orig = compose(tucker_orig)

    tnsr_max = maximum(tnsr_orig)
    map!(x -> x / tnsr_max, tucker_orig.core)
    map!(x -> x / tnsr_max, tnsr_orig)

    tnsr = add_noise(tnsr_orig, 0.6, true)

    # Solve the problem
    @time tucker_spnn = spnntucker(tnsr, size(tucker_orig.core), tol=1E-4, ini_decomp=:hosvd,
                             core_nonneg=true,
                             max_iter=1000, verbose=true, lambdas=fill(0.1, 4))

    # Reporting
    @fact rel_residue(tucker_spnn) --> less_than(0.05)
    info("Relative error of decomposition : $(rel_residue(tucker_spnn))")
end

context("semi-nonnegative decomposition") do
    tucker_orig = rand_tucker((4, 5, 6), (40, 50, 60), factors_nonneg=true, core_nonneg=false)
    tnsr_orig = compose(tucker_orig)

    tnsr_max = maximum(tnsr_orig)
    map!(x -> x / tnsr_max, tucker_orig.core)
    map!(x -> x / tnsr_max, tnsr_orig)

    tnsr = add_noise(tnsr_orig, 0.6, false)

    # Solve the problem
    @time tucker_spnn = spnntucker(tnsr, size(tucker_orig.core), tol=1E-4, ini_decomp=:hosvd,
                             core_nonneg=false,
                             max_iter=1000, verbose=true, lambdas=fill(0.1, 4))

    # Reporting
    @fact rel_residue(tucker_spnn) --> less_than(0.05)
    info("Relative error of decomposition : $(rel_residue(tucker_spnn))")
end

end


using Pumas, Plots, CSV, Random
Random.seed!(0)

single_dose_regimen = DosageRegimen(100, time=0)
first(single_dose_regimen.data)

s1 = Subject(id=1, evs=single_dose_regimen,cvs=(Wt=70,))

choose_covariates() = (Wt = rand(55:80),)

pop = Population(map(i -> Subject(id = i,evs = single_dose_regimen, cvs =  choose_covariates()),1:24))

pop[1].covariates

mymodel = @model begin
  @param   begin
    tvcl ∈ RealDomain(lower=0, init = 1.0)
    tvv ∈ RealDomain(lower=0, init = 20)
    tvka ∈ RealDomain(lower = 0, init= 1)
    Ω ∈ PDiagDomain(init=[0.09,0.09, 0.09])
    σ_prop ∈ RealDomain(lower=0,init=0.04)
  end

  @random begin
    η ~ MvNormal(Ω)
  end

  @pre begin
    CL = tvcl * (Wt/70)^0.75 * exp(η[1])
    V  = tvv * (Wt/70) * exp(η[2])
    Ka = tvka * exp(η[3])
  end
  @covariates Wt

  @dynamics Depots1Central1
    #@dynamics begin
    #    Depot' =  -Ka*Depot
    #    Central' =  Ka*Depot - (CL/V)*Central
    #end

  @derived begin
      cp = @. 1000*(Central / V)
      dv ~ @. Normal(cp, sqrt(cp^2*σ_prop))
    end
end

param = init_param(mymodel)

obs = simobs(mymodel, pop, param, obstimes=0:1:72)
plot(obs)

simdf = DataFrame(obs)
first(simdf, 6)

simdf[!, :route] .= "ev"

timeu = u"hr"
concu = u"mg/L"
amtu  = u"mg"

ncadf = read_nca(simdf, id=:id, time=:time, conc=:cp, amt=:amt,
    route=:route,timeu=timeu, concu=concu, amtu=amtu, lloq=0.4concu)

plot(ncadf)

auc = NCA.auc(ncadf)

cmax = NCA.cmax(ncadf)

report = NCAReport(ncadf)
report = NCA.to_dataframe(report)
first(report,6)

simdf.cmt = ifelse.(ismissing.(simdf.cmt), 2, simdf.cmt)
est_df = simdf[.!((simdf.dv .== 0.0) .& (simdf.cmt .==2)),:]
first(est_df,6)

data = read_pumas(est_df ,cvs = [:Wt], dvs=[:dv])

res = fit(mymodel,data,param,Pumas.FOCEI())

infer(res)

preds = DataFrame(predict(res))
first(preds, 6)

resids = DataFrame(wresiduals(res))
first(resids, 6)

ebes = DataFrame(empirical_bayes(res))
first(ebes, 6)

resout = DataFrame(inspect(res))
first(resout, 6)


using jlmie
using Plots

# settings
# structure
nmat = 4  # refractive index of a sphere
radius = 100*1e-9

# environment
nenv = 1.00

# wavelenth
lbd0 = (400:1000)*1e-9; lbdp = lbd0.*1e9;

# other settings
nmax = -1  # -1: namx large enough (determined from x)

# main
m,x = jlmie_mx(nmat,radius,lbd0,nenv)
Qsca = jlmie_Qsca(m,x,nmax)
println("calculation completed")

# show results
# validation can be done by comparing to ITMO Mie Calculator
# https://physics.itmo.ru/en/mie
plt = plot(lbdp,Qsca,
    xlabel       = "Wavelength (nm)",
    ylabel       = "Scattering efficiency",
    legend       = false,
    # label        ="LABEL",
    # xlims        =(-3,3),
    # ylims        =(-3,3),
    # aspect_ratio =0.5,
    # title        ="TITLE",
    # linecolor    =:blue,
    # linewidth    =5,
    # linestyle    =:dot,
    # size         =(400,300),
) # https://qiita.com/I_ppp/items/dca3552affa6a672e4bd
display(plt)

# save
if true
    runningfilename = splitext(splitpath(@__FILE__)[end])[1]
    outfilename = ".\\result\\" * runningfilename
    savefig(plt,outfilename)
    println("Graph has been saved in result directory")
end


# Testing whether SpectralGaussianSimulation does what I think it does
using GeoStats
using SpectralGaussianSimulation
using DirectGaussianSimulation
using Distances
include("CovFuncs.jl")
include("Utils.jl")

n = 100
N = n^2
NSamples = 2000

l = 0.5
# covfunc!(r) = exponential!(r, l)
covfunc!(r) = matern32!(r, l)

# simProblem = SimulationProblem(RegularGrid((-2.0, -2.0.0, 0.0), (1.0, 1.0), dims=(n, n)), :z => Float64, NSamples)

# simSolver = SpecGaussSim(:z => (variogram=MaternVariogram(range=l, order=1.5),))

# simSolution = solve(simProblem, simSolver)
# x = coordinates(domain(simSolution))
# Y = reduce(hcat, simSolution.realizations[:z])

x, Y = specSampleGrid2D(n, NSamples, 3, MaternVariogram(range=l, order=1.5))

Θ = covfunc!(Distances.pairwise(Euclidean(), x[:, inds], dims=2))
YExact = rand(MvNormal(Θ), NSamples)

ΘEmp = cov(YExact, dims=2)
ΘEmpApp = cov(Y[inds, :], dims=2)

@show norm(Θ - ΘEmp) / norm(Θ)
@show norm(Θ - ΘEmpApp) / norm(Θ)


using FFTW

function compressaudiofft(channels, threshold)
    # Split channels
    left = channels[:,1]
    right = channels[:,2]

    # Perfom transforms
    fleft = fft(left)
    fright = fft(right)

    # Compress
    tleft = threshold * maximum(abs.(fleft))
    ileft = abs.(fleft) .> tleft
    tright = threshold * maximum(abs.(fright))
    iright = abs.(fright) .> tright
    return (fleft[ileft], ileft, fright[iright], iright, length(left))
end

function decompressaudiofft(compression)
    # Unpack compression
    cleft, ileft, cright, iright, len = compression

    # Decompress left channel
    fleft = Array{Complex{Float64},1}(undef, length(ileft))
    j = 1
    for i = 1:length(fleft)
        if ileft[i]
            fleft[i] = cleft[j]
            j += 1
        end
    end

    # Decompress right channel
    fright = Array{Complex{Float64},1}(undef, length(iright))
    j = 1
    for i = 1:length(fright)
        if iright[i]
            fright[i] = cright[j]
            j += 1
        end
    end

    # Perform inverse transforms
    left = ifft(fleft)
    right = ifft(fright)

    return real.([left[1:len] right[1:len]])
end


"""
Computes line profiles and associated quantities.
"""

const invSqrtPi = 1. / sqrt(π)

"""
Compute scaled complex complementary error function using
[Humlicek (1982, JQSRT 27, 437)](https://doi.org/10.1016/0022-4073(82)90078-4)
W4 rational approximations.
Here, z is defined as z = v + i * a, and returns w(z) = H(a,v) + i * L(a, v).
"""
function humlicek(z::Complex)
    s = abs(real(z)) + imag(z)
    if s > 15.0
        # region I
        w = im * invSqrtPi * z / (z * z - 0.5)
    elseif s > 5.5
        # region II
        zz = z * z
        w = im * (z * (zz * invSqrtPi - 1.4104739589)) / (0.75 + zz * (zz - 3.0))
    else
        x, y = real(z), imag(z)
        t = y - im * x
        if y >= 0.195 * abs(x) - 0.176
            # region III
            w = ((16.4955 + t * (20.20933 + t * (11.96482 + t * (3.778987 + 0.5642236 * t))))
               / (16.4955 + t * (38.82363 + t * (39.27121 + t * (21.69274 + t * (6.699398 + t))))))
        else
            # region IV
            u = t * t
            nom = t * (36183.31 - u * (3321.99 - u * (1540.787 -  u *
                   (219.031 - u * (35.7668 - u * (1.320522 - u * .56419))))))
            den = 32066.6 - u * (24322.8 - u * (9022.23 - u * (2186.18 -
                    u * (364.219 - u * (61.5704 - u * (1.84144 - u))))))
            w = exp(u) - nom / den
        end
    end
    return w
end

"""
    voigt_profile(a::AbstractFloat, v::AbstractFloat, ΔD::Unitful.Quantity)

Compute the normalised Voigt profile, given a damping constant `a`, dimensionless
wavelength or frequency `v`, and Doppler width `ΔD` (wavelength or frequency).
In the case of wavelength, v = (λ - λ0) / ΔD.
Uses Humlicek's W4 approximation. Returns in inverse units of ΔD.
"""
function voigt_profile(a::AbstractFloat, v::AbstractFloat, ΔD::Unitful.Quantity)
    z = v + a * im
    profile = real(humlicek(z))
    return profile * invSqrtPi / ΔD
end

"""
    dispersion_profile(a::AbstractFloat, v::AbstractFloat, ΔD::Unitful.Quantity)

Compute the normalised dispersion (or Faraday) profile, given a damping constant `a`,
dimensionless wavelength or frequency `v`, and Doppler width `ΔD` (wavelength or frequency).
In the case of wavelength, v = (λ - λ0) / ΔD.
Uses Humlicek's W4 approximation. Returns in inverse units of ΔD.
"""
function dispersion_profile(a::AbstractFloat, v::AbstractFloat, ΔλD::Unitful.Length)
    z = v + a * im
    profile = imag(humlicek(z))
    return profile * invSqrtPi / ΔλD
end

"""
    doppler_width(
        λ0::Unitful.Length,
        mass::Unitful.Mass,
        temperature::Unitful.Temperature
    )
    doppler_width(
        ν0::Unitful.Frequency,
        mass::Unitful.Mass,
        temperature::Unitful.Temperature
    )

Compute Doppler width in wavelength or frequency units, given a rest wavelength/frequency,
mass of atom/ion, and temperature. Returns in typical units for UV/optical/infrared lines.
"""
function doppler_width(
    λ0::Unitful.Length,
    mass::Unitful.Mass,
    temperature::Unitful.Temperature
)
    return (λ0 / c_0 * sqrt(2 * k_B * temperature / mass)) |> u"nm"
end

function doppler_width(
    ν0::Unitful.Frequency,
    mass::Unitful.Mass,
    temperature::Unitful.Temperature
)
    return (ν0 / c_0 * sqrt(2 * k_B * temperature / mass)) |> u"THz"
end


# Timescale for diffusion across one cell
min_Δxyz(grid) = min(grid.Δx, grid.Δy, grid.Δz)
min_Δxy(grid) = min(grid.Δx, grid.Δy)
min_Δz(grid) = grid.Δz

cell_diffusion_timescale(model) = cell_diffusion_timescale(model.closure, model.diffusivities, model.grid)

"Returns the time-scale for diffusion on a regular grid across a single grid cell."
function cell_diffusion_timescale(closure::ConstantIsotropicDiffusivity, diffusivities, grid)
    Δ = min_Δxyz(grid)
    max_κ = maximum(closure.κ)
    return min(Δ^2 / closure.ν, Δ^2 / max_κ)
end

function cell_diffusion_timescale(closure::ConstantAnisotropicDiffusivity, diffusivies, grid)
    Δh = min_Δxy(grid)
    Δz = min_Δz(grid)
    max_κh = maximum(closure.κh)
    max_κv = maximum(closure.κv)
    return min(Δz^2 / closure.νv, Δh^2 / closure.νh,
               Δz^2 / max_κv, Δh^2 / max_κh)
end

function cell_diffusion_timescale(closure::AnisotropicBiharmonicDiffusivity, diffusivies, grid)
    Δh = min_Δxy(grid)
    Δz = min_Δz(grid)
    max_κh = maximum(closure.κh)
    max_κv = maximum(closure.κv)
    return min(Δz^4 / closure.νv, Δh^4 / closure.νh,
               Δz^4 / max_κv, Δh^4 / max_κh)
end

function cell_diffusion_timescale(closure::Union{AbstractSmagorinsky, AbstractLeith}, diffusivies, grid)
    Δ = min_Δxyz(grid)
    min_Pr = minimum(closure.Pr)
    max_κ = maximum(closure.κ)
    max_νκ = maximum(diffusivities.νₑ.data.parent) * max(1, 1/min_Pr)
    return min(Δ^2 / max_νκ, Δ^2 / max_κ)
end

function cell_diffusion_timescale(closure::AbstractAnisotropicMinimumDissipation, diffusivies, grid)
    Δ = min_Δxyz(grid)
    max_ν = maximum(diffusivities.νₑ.data.parent)
    max_κ = max(Tuple(maximum(κₑ.data.parent) for κₑ in diffusivities.κₑ)...)
    return min(Δ^2 / max_ν, Δ^2 / max_κ)
end


function ManifoldTests.test_forwarddiff(M::Manifold, pts, tv)
    return for (p, X) in zip(pts, tv)
        exp_f(t) = distance(M, p, exp(M, p, t[1] * X))
        d12 = norm(M, p, X)
        for t in 0.1:0.1:0.9
            Test.@test d12 ≈ ForwardDiff.derivative(exp_f, t)
        end

        retract_f(t) = distance(M, p, retract(M, p, t[1] * X))
        for t in 0.1:0.1:0.9
            Test.@test ForwardDiff.derivative(retract_f, t) ≥ 0
        end
    end
end


#=

julia julia_callback.jl
  2.457 ms (100002 allocations: 10.68 MiB)

=#
using BenchmarkTools

function rober(t, u)
    k1 = 0.04
    k2 = 3e7
    k3 = 1e4
    y1, y2, y3 = u
    dy1 = -k1 * y1 + k3 * y2 * y3
    dy2 = k1 * y1 - k2 * y2 * y2 - k3 * y2 * y3
    dy3 = k2 * y2 * y2
    return [dy1, dy2, dy3]
end

rober(0, [1.0, 0.0, 0.0])

function call_function(func)
    u = [1.0, 0.0, 0.0]
    dt = 0.1
    times = dt * (0:1e5)
    for time in times
        u = func(time, u)
    end
end

@btime call_function(rober)


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using JuliaFEM.Testing
importall Base
import JuliaFEM: get_basis, get_dbasis, get_integration_points

type MyQuad4 <: AbstractElement
end

function get_basis(element::Element{MyQuad4}, ip, time)
    1/4*[(1-ip[1])*(1-ip[2]) (1+ip[1])*(1-ip[2]) (1+ip[1])*(1+ip[2]) (1-ip[1])*(1+ip[2])]
end

function get_dbasis(element::Element{MyQuad4}, ip, time)
    1/4*[-(1-ip[2])    (1-ip[2])   (1+ip[2])  -(1+ip[2])
         -(1-ip[1])   -(1+ip[1])   (1+ip[1])   (1-ip[1])]
end

function get_integration_points(element::MyQuad4)
    [
        (1.0, 1.0/sqrt(3.0)*[-1, -1]),
        (1.0, 1.0/sqrt(3.0)*[ 1, -1]),
        (1.0, 1.0/sqrt(3.0)*[ 1,  1]),
        (1.0, 1.0/sqrt(3.0)*[-1,  1])
    ]
end

function length(element::Element{MyQuad4})
    return 4
end

function size(element::Element{MyQuad4})
    return (2, 4)
end

@testset "test new element" begin
    el = Element(MyQuad4)
    el["geometry"] = Vector{Float64}[[0.0,0.0], [1.0,0.0], [1.0,1.0], [0.0,1.0]]
    el["displacement"] = Vector{Float64}[[0.0,0.0], [0.0,0.0], [1.0,0.0], [0.0,0.0]]
    @test isapprox(el("geometry", [0.0, 0.0], 0.0), [0.5, 0.5])
    @test isapprox(el("displacement", [0.0, 0.0], 0.0), [0.25, 0.0])
    el["temperature thermal conductivity"] = 6.0
    dim = length(el)
    K = zeros(dim, dim)
    A = 0.0
    time = 0.0
    for ip in get_integration_points(el)
        dN = el(ip, time, Val{:Grad})
        detJ = el(ip, time, Val{:detJ})
        w = ip.weight*detJ
        c = el("temperature thermal conductivity", ip, time)
        K += w*c*dN'*dN
        A += w
    end
    @test isapprox(A, 1.0)
    K_expected = [
        4.0  -1.0  -2.0  -1.0
       -1.0   4.0  -1.0  -2.0
       -2.0  -1.0   4.0  -1.0
       -1.0  -2.0  -1.0   4.0]
    @test isapprox(K, K_expected)
end


using Distributions, SpecialFunctions, Roots, LinearAlgebra
"""
TSmileOriginal: SMiLe rule (Faraji, et al., 2018)
    for the Categorical task (section 2.3.2)
    - No knowledge of environment's prior (i.e. stochasticity) for the calculation
        of Scc (it is there for the calculation of the alphas though)
    - Optimization to find γ
"""
struct TSmileOriginal
    ns::Int
    m::Float64
    stochasticity::Float64 # This is used for the alphas
    Ps1::Array{Float64, 1}
    alphas::Array{Float64,1}
end
function TSmileOriginal(;ns = 10, m = .1, stochasticity = .01)
    Ps1 = Array{Float64, 1}(undef, ns)
    alphas = Array{Float64, 1}(undef, ns)
    alphas = stochasticity .* ones(ns)
    TSmileOriginal(ns, m, stochasticity, Ps1, alphas)
end
export TSmileOriginal
function update!(learnerT::TSmileOriginal, s1)
    betas = ones(learnerT.ns)
    betas[s1] += 1.
    Scc = KL(learnerT.alphas, betas)
    Bmax = KL(betas, learnerT.alphas)
    B = learnerT.m * Scc/(1. + learnerT.m * Scc) * Bmax
    γ0 = find_γ0(betas, learnerT.alphas, B)
    for s in 1:learnerT.ns
        learnerT.alphas[s] = (1. - γ0) * learnerT.alphas[s] + γ0 * betas[s]
    end
    computePs1!(learnerT)
end
export update!
function find_γ0(betas, alphas, B)
    f = γ -> KL(γ .* betas .+ (1. - γ) .* alphas, alphas) - B
    if abs(f(0.)) < 5*eps()
        γ0 = 0.
    elseif abs(f(1.)) < 5*eps()
        γ0 = 1.
    else
        γ0 = find_zero(f, (0., 1.))
    end
    γ0
end
function KL(α1, α2)
    max(lbeta(α2) - lbeta(α1) + dot(α1 .- α2, digamma.(α1) .- digamma(sum(α1))), 0.)
end
function lbeta(α::Array{Float64,1})
    α0 = 0.
    lmnB = 0.
    for i in 1:length(α)
        αi = α[i]
        α0 += αi
        lmnB += lgamma(αi)
    end
    lmnB -= lgamma(α0)
    lmnB
end
function computePs1!(learnerT::Union{TVarSmile, TSmileOriginal})
    for s in 1:learnerT.ns
        learnerT.Ps1[s] = learnerT.alphas[s] / sum(learnerT.alphas)
    end
end

# function computePs1mode!(learnerT::TSmile)
#     for s in 1:learnerT.ns
#         learnerT.Ps1[s] = (learnerT.alphas[s] - 1. + eps())/ sum(learnerT.alphas .-1. .+ eps())
#     end
# end


struct GaussianRewards <: RewardDist
    means::Vector{Function}
    sd::Vector{Float64}
end

function GaussianRewards(means::Vector{<:Function}, sd::Float64)
    sd1 = fill(sd, length(means))
    return GaussianRewards(means, sd1)
end

function sample(x::Vector{Float64}, a::Int64, rdist::GaussianRewards)
    mu = rdist.means[a](x)
    sd = rdist.sd[a]
    return rand(Normal(mu, sd))
end

function expected_regret(x::Vector{Float64}, a::Int64, rdist::GaussianRewards)
    mu = [rdist.means[a1](x) for a1 in 1:length(rdist.means)]
    mu_max = findmax(mu)[1]
    return mu_max - mu[a]
end


using Revise
using AdFem
using ADCME
using Random
using ProgressMeter
using PyPlot
using DelimitedFiles
using Distributions

np = pyimport("numpy")
n = 10
m = 15
h = 0.1

xy = Dict(
    1=>(1.95,2.05,0.32,0.38),
    2=>(1.8,2.2,0.2,0.45),
    3=>(1.7,2.2,0.18,0.5),
    4=>(1.5,2.3,0.15,0.5),
    5=>(1.2,3.0,0.0,0.5)
)
for (ii,σ0) in enumerate([0.01 0.05 0.1 0.2 0.5])
xmin,xmax, ymin,ymax = xy[ii]

Random.seed!(233)
obs_id = rand(interior_node("all", m, n, h), 10)
obs_id = [obs_id; obs_id .+ (m+1)*(n+1)]
function solve_elasticity(c)
    c = constant(c)
    E, ν = c[1], c[2]
    H = 1e6*E/(1+ν)/(1-2ν)*tensor([
             1-ν ν 0.0
             ν 1-ν 0.0
             0.0 0.0 (1-2ν)/2
         ])
    K = compute_fem_stiffness_matrix(H, m, n, h)
    bdedge = bcedge("right", m, n, h)
    bdnode = bcnode("left", m, n, h)
    K, _ = fem_impose_Dirichlet_boundary_condition(K, bdnode, m, n, h)
    F = compute_fem_traction_term([zeros(4*m*n) 1e6*ones(4*m*n)], bdedge, m, n, h)
    sol = K\F
    sol = sol[obs_id]
    return sol
end

sol = solve_elasticity([2.0;0.35])
sess = Session(); init(sess)
sol_ = run(sess, sol)
# # visualize_displacement(sol_*1e3, m, n, h)

c = Variable(0.1*rand(2))
sol = solve_elasticity(c)

loss = sum((sol-sol_)^2)
sess = Session(); init(sess)
BFGS!(sess, loss)

y = sol_
hs = run(sess, sol)
H = run(sess, gradients(sol, c))
R = 0.1^2*diagm(0=>ones(length(sol_)))
s = run(sess, c)
Q = diagm(0=>ones(2))

μ, V = uqnlin(y, hs, H, R, s, Q)
d = Distributions.MvNormal(μ, (V+V')/2)

x0 = LinRange(xmin, xmax, 100)
y0 = LinRange(ymin, ymax, 100)
X, Y = np.meshgrid(x0, y0)
Z = hcat([[pdf(d, [X[i, j];Y[i, j]]) for i = 1:100] for j = 1:100]...)
mcmc = readdlm("simdata/sigma$σ0-100000.txt")

close("all")
contour(X, Y, Z, levels=4)
scatter(mcmc[:,1], mcmc[:,2], s=2, c="red", label="RW-MCMC", alpha=0.5)
legend()
xlabel("E")
ylabel("\$\\nu\$")
xlim(xmin, xmax)
ylim(ymin, ymax)
savefig("sigma$σ0.png")
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -5.25849806212995 <= q <= 4.375154350776637)

                     add_NL_constraint(m, :((0.4440484214047524 + 0.27117166180904073*$(x[1]) + 0.1549579763005995*$(x[2]) + -0.19609804892644567*$(x[3]))*(1 + erf((0.4440484214047524 + 0.27117166180904073*$(x[1]) + 0.1549579763005995*$(x[2]) + -0.19609804892644567*$(x[3]))/sqrt(2)))/2 + (-0.1025904891815026 + 0.10197317322059929*$(x[1]) + 0.049423996875141984*$(x[2]) + -0.5007584267765015*$(x[3]))*(1 + erf((-0.1025904891815026 + 0.10197317322059929*$(x[1]) + 0.049423996875141984*$(x[2]) + -0.5007584267765015*$(x[3]))/sqrt(2)))/2 + (0.34943625748753826 + 0.7627977483870185*$(x[1]) + -0.7008024309541692*$(x[2]) + -0.24608961291909015*$(x[3]))*(1 + erf((0.34943625748753826 + 0.7627977483870185*$(x[1]) + -0.7008024309541692*$(x[2]) + -0.24608961291909015*$(x[3]))/sqrt(2)))/2 + (-0.3923242036950221 + -0.06896448884223272*$(x[1]) + -0.2210802938542722*$(x[2]) + -0.35764467125421584*$(x[3]))*(1 + erf((-0.3923242036950221 + -0.06896448884223272*$(x[1]) + -0.2210802938542722*$(x[2]) + -0.35764467125421584*$(x[3]))/sqrt(2)))/2 + (-0.7402418416924226 + -0.21635543628426968*$(x[1]) + -0.713523096277699*$(x[2]) + 0.2551851437719974*$(x[3]))*(1 + erf((-0.7402418416924226 + -0.21635543628426968*$(x[1]) + -0.713523096277699*$(x[2]) + 0.2551851437719974*$(x[3]))/sqrt(2)))/2 - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


module RiskPerfTests

using Test
using Dates
using Statistics
using RiskPerf

# load test data produced using R PerformanceAnalytics package
asset_returns = parse.(Float64, readlines("data/recon_asset_rets.csv")[2:end])
market_returns = parse.(Float64, readlines("data/recon_market_rets.csv")[2:end])
rf = parse.(Float64, readlines("data/recon_rf_rets.csv")[2:end])

@test std(rf) ≈ 1.03135983021606e-07

@test skewness(asset_returns) ≈ -0.0904045109144944
@test skewness(asset_returns; method=:moment) ≈ -0.0904045109144944
@test skewness(asset_returns; method=:fisher_pearson) ≈ -0.09054037810903758                # bug in R's PerformanceAnalytics, corrected here
@test skewness(asset_returns; method=:sample) ≈ -0.0906763586376902

@test kurtosis(asset_returns) ≈ 0.0619856901585289
@test kurtosis(asset_returns; method=:excess) ≈ 0.0619856901585289
@test kurtosis(asset_returns; method=:moment) ≈ 3.06198569015853
@test kurtosis(asset_returns; method=:cornish_fisher) ≈ -0.158076535757453

@test downside_deviation(asset_returns, 0.001; method=:full) ≈ 0.00102022195386797
@test downside_deviation(asset_returns, rf; method=:full) ≈ 0.000474556578712896            # bug in R's PerformanceAnalytics, corrected here
@test downside_deviation(asset_returns, 0.001; method=:partial) ≈ 0.00123267724486424
@test downside_deviation(asset_returns, rf; method=:partial) ≈ 0.00085648312494845          # bug in R's PerformanceAnalytics, corrected here

@test lower_partial_moment(asset_returns, 0.001, 2, :full) ≈ 1.04085283515417e-06
@test lower_partial_moment(asset_returns, rf, 2, :full) ≈ 2.25203946399689e-07
@test lower_partial_moment(asset_returns, 0.001, 2, :partial) ≈ 1.51949319000609e-06
@test lower_partial_moment(asset_returns, rf, 2, :partial) ≈ 7.33563343321463e-07

@test higher_partial_moment(asset_returns, 0.001, 2, :full) ≈ 2.15200955473486e-07
@test higher_partial_moment(asset_returns, rf, 2, :full) ≈ 1.0554005124107e-06
@test higher_partial_moment(asset_returns, 0.001, 2, :partial) ≈ 6.83177636423766e-07
@test higher_partial_moment(asset_returns, rf, 2, :partial) ≈ 1.5229444623531e-06

@test upside_potential_ratio(asset_returns, 0.001; method=:full) ≈ 0.199853514985918        # bug in R's PerformanceAnalytics, corrected here
@test upside_potential_ratio(asset_returns, rf; method=:full) ≈ 1.49034594998309            # bug in R's PerformanceAnalytics, corrected here
@test upside_potential_ratio(asset_returns, 0.001; method=:partial) ≈ 0.525105446510415     # bug in R's PerformanceAnalytics, corrected here
@test upside_potential_ratio(asset_returns, rf; method=:partial) ≈ 1.19157955999556         # bug in R's PerformanceAnalytics, corrected here

@test volatility(asset_returns) ≈ 0.0010120151595378
@test tracking_error(asset_returns, market_returns) ≈ 0.00111905105002574
@test information_ratio(asset_returns, market_returns) ≈ 0.34034622188359
@test omega_ratio(asset_returns, 0.001) ≈ 0.297008426326568

@test sharpe_ratio(asset_returns) ≈ 0.511256629034303
@test sharpe_ratio(asset_returns; risk_free=rf) ≈ 0.501375462696111
@test adjusted_sharpe_ratio(asset_returns) ≈ 0.5069731154546914  # slightly different to R's PerformanceAnalytics due to their annualization (0.506476646827294)

@test sortino_ratio(asset_returns) ≈ 1.10000435875069
@test sortino_ratio(asset_returns; MAR=rf) ≈ 1.0692077438794

α1, β1 = capm(asset_returns, market_returns)
@test α1 ≈ 0.000512345109212185
@test β1 ≈ 0.0370188032088717

α2, β2 = capm(asset_returns, market_returns; risk_free=rf)
@test α2 ≈ 0.000502715256287843
@test β2 ≈ 0.037019963055605

γ = jensen_alpha(asset_returns, market_returns; risk_free=rf)
@test γ ≈ α2

γ = modified_jensen(asset_returns, market_returns; risk_free=rf)
@test γ ≈ α2/β2

@test treynor_ratio(asset_returns, market_returns) ≈ 0.013976666292467
@test treynor_ratio(asset_returns, market_returns; risk_free=rf) ≈ 0.0137061068404259

@test value_at_risk(asset_returns, 0.05, :historical) ≈ -0.00115939244815142
@test value_at_risk(asset_returns, 0.05, :gaussian) ≈ -0.00114638483011466
@test value_at_risk(asset_returns, 0.05, :cornish_fisher) ≈ -0.00117095813580772

@test expected_shortfall(asset_returns, 0.05, :historical) ≈ -0.00164648648602694
@test expected_shortfall(asset_returns, 0.05, :gaussian) ≈ -0.00156905316259027
@test expected_shortfall(asset_returns, 0.05, :cornish_fisher) ≈ -0.0016287316713825

@test length(drawdowns(asset_returns)) == length(asset_returns)
@test mean(drawdowns(asset_returns)) ≈ -0.000418043468523338
@test std(drawdowns(asset_returns)) ≈ 0.000757130672676854
@test mean(drawdowns(asset_returns; geometric=true)) ≈ -0.000517982256217951
@test std(drawdowns(asset_returns; geometric=true)) ≈ 0.000916577011725993

pnl = [-50, 100, -20, 30, -120, -10, 200, 10, 18, -12, -20, -30, 0];
@test length(drawdowns_pnl(pnl)) == length(pnl)
@test mean(drawdowns_pnl(pnl)) ≈ -37.53846153846154
@test std(drawdowns_pnl(pnl)) ≈ 45.40303841634277

# https://breakingdownfinance.com/finance-topics/modern-portfolio-theory/marginal-contribution-to-risk-mctr/
w = [0.042, 0.25, 0.32, 0.22, 0.168]
covmat = [
    0.001369 0.001184 0.000161 0.001924 0.002886;
    0.001184 0.006400 0.000580 0.008840 0.006240;
    0.000161 0.000580 0.000841 0.001885 0.001131;
    0.001924 0.008840 0.001885 0.016900 0.003380;
    0.002886 0.006240 0.001131 0.003380 0.067600
]
ctb = risk_contribution(w, covmat)

@test all(isapprox.(ctb, [0.009944 0.217649 0.059174 0.284805 0.428427]'; rtol=0.0001))

end