#############################################################################
# fujishige_wolfe.jl
# Implements the the fujishige_wolfe algorithm, to find the point with
# minimum norm in the base polyhedron.
#############################################################################

export fujishige_wolfe

function fujishige_wolfe(p::AssocPoly, w::AbstractArray, Tol::Float64 = 1e-3)
  # step 1 initialization
  n = length(w)
  @assert length(p.V) == n
  V = sort(p.V)         # the sorted base set
  x = zeros(n)          # the candidate
  y = zeros(n + 1)
  for i = 1:n
    y[i + 1] = evaluate(p.F, V[1:i])[1]
    x[i] = y[i + 1] - y[i]
  end
  S = zeros(n, 1)       # the set of selected extreme points that constitute the affine set
  S[:, 1] = x
  # major cycle
  major_cycle(p, w, x, S, Tol)
end

function fujishige_wolfe(p::PosPoly, w::AbstractArray, Tol::Float64 = 1e-3)
  w₁ = copy(w)
  if any(w₁ .< 0)
    negind = find(w₁ .< 0)
    w₁[negind] = zeros(length(negind))
  end
  p₁ = SubmodPoly(p.F)
  fujishige_wolfe(p₁, w₁, Tol)
end

function fujishige_wolfe(p::SymPoly, w::AbstractArray, Tol::Float64 = 1e-3)
  w₁ = sign.(w) .* w
  p₁ = SubmodPoly(p.F)
  x = fujishige_wolfe(p₁, w₁, Tol)
  return sign.(w) .* x
end

# minimum-norm-point with warmstart
function fujishige_wolfe(p::AssocPoly, w::AbstractArray, x₀::AbstractArray, Tol::Float64 = 1e-3)
  n = length(w)
  @assert length(p.V) == n
  S = zeros(n, 1)
  S[:, 1] = x₀
  major_cycle(p, w, x₀, S, Tol)
end

function fujishige_wolfe(p::PosPoly, w::AbstractArray, x₀::AbstractArray, Tol::Float64 = 1e-3)
  w₁ = copy(w)
  if any(w₁ .< 0)
    negind = find(w₁ .< 0)
    w₁[negind] = zeros(length(negind))
  end
  p₁ = SubmodPoly(p.F)
  fujishige_wolfe(p₁, w₁, x₀, Tol)
end

function fujishige_wolfe(p::SymPoly, w::AbstractArray, x₀::AbstractArray, Tol::Float64 = 1e-3)
  w₁ = sign.(w) .* w
  p₁ = SubmodPoly(p.F)
  x = fujishige_wolfe(p₁, w₁, x₀, Tol)
  return sign.(w) .* x
end

function major_cycle(p::AssocPoly, w::AbstractArray, x::AbstractArray, S::AbstractMatrix, Tol::Float64 = 1e-3, nTol::Float64 = 1e-3)
  n = length(w)
  S1 = minimum_check(p, w, x, S, Tol)
	if S1 == S
    return x
  else
    S = S1
    colnum = size(S, 2)
    # step 3 find the point closes to w in the affine hull of points in S
    if colnum > size(S, 1)              # the affine hull is in face the whole space
      y0 = copy(w)
      push!(y0, 1)
      S1 = [S; ones(1, size(S, 2))]
      μ = S1\y0
      if all(μ .>= -nTol)                # μ is in the affine hull of points in set
        return w
      else                              # delete redundant points and compute the new x
        redind0 = find(μ .< -nTol)
        remind0 = deleteat!(collect(1:colnum), redind0)
        S = S[:, remind0]
        colnum = length(remind0)
        y = aff_proj(w, S)
        major_cycle(p, w, y, S, Tol)
      end
    else
      y = aff_proj(w, S)                 # the point found)
      y0 = copy(y)
      push!(y0, 1)
      S₁ = [S; ones(1, size(S, 2))]
      μ = S₁\y0                         # the coefficients of y in terms of points in S
      while any(μ .< -nTol)
        # step 4 find the closest point to y on the segment between x and y while staying in the convex hull of points in S, and set it to be x
        x₀ = copy(x)
        push!(x₀, 1)
        λ = S₁\x₀                           # the coefficients of x in terms of points in S
        ratios = λ ./ μ
        negratios = ratios[find(μ .< -nTol)]
        α = maximum(negratios)
        redind = find((ratios .<= α + nTol) .& (ratios .<= α + nTol))    # the index of the redundant points
        remind = deleteat!(collect(1:colnum), redind)
        S = S[:, remind]                           # the new set of points
        x = y                                      # the new candidate
        y = aff_proj(w, S)                          # the new optimal solution given S
        y0 = copy(y)
        push!(y0, 1)
        S₁ = [S; ones(1, size(S, 2))]
        μ = S₁\y0
      end
      major_cycle(p, w, y, S, Tol)
    end
  end
end

function minimum_check(p::AssocPoly, w::AbstractArray, x::AbstractArray, S::AbstractMatrix, Tol::Float64 = 1e-3)
  a = x - w
  z = fenchel(p, a)
  minprod = dot(z - w, a)
  curprod = dot(a, a)
  if curprod - minprod <= Tol
    return S
  else
    return [S z]
  end
end

function minimum_check(p::SubmodPoly, w::AbstractArray, x::AbstractArray, S::AbstractMatrix, Tol::Float64 = 1e-3)
  a = x - w
  pos = false
  if any(a .> 0)
    pos = true
    posind = find(a .> 0)
  end
  z = greedy(p.F, -a)
  if pos                                                 # x - w has positive elements
    s = z[posind] - w[posind]
    if any(s .> 0)                                       # some elements of z should be adjusted
      posind1 = find(s .> 0)
      z[posind[posind1]] = w[posind[posind1]]
    end
  end
  minprod = dot(z - w, a)
  curprod = dot(a, a)
  if curprod - minprod <= Tol
    return S
  else
    S1 = [S z]
    return [S z]
  end
end

function fujishige_wolfe(f::SubmodFunc, w::AbstractArray, Tol::Float64 = 1e-3)
  p = BasePoly(f)
  return fujishige_wolfe(p, w, Tol)
end


# Data for Robust Optimization
# Version: 1.0
# Date: April 7th, 2019
# 1. Data Import -------------------------------------------------------------------------------------------------------
file = openxl("data.xlsx")
vec_timeProdDev = readxl(file, "timeProdDev!B2:B11")
vec_costInvest = readxl(file, "costInvest!B2:B5")
vec_costProd = readxl(file, "costProd!B2:B11")
vec_demand = readxl(file, "demand!B2:B11")
mat_prodOnMachine = readxl(file, "prodOnMachine!B2:E11")
vec_timeAvai = readxl(file, "timeAvai!B2:B5")
vec_timeProd = readxl(file, "timeProd!B2:B11")
# 2. Data Process ------------------------------------------------------------------------------------------------------
# Function to calculate parameters
function paras_budgetMILP(num_x, num_y)
    vec_c = hcat(zeros(num_x))
    for i = 1: 4
        vec_c[(4 * i - 3): (4 * i), 1] = repeat([vec_costProd[i]], num_y)
    end
    vec_f = hcat(vec_costInvest)
    vec_b = hcat(repeat([0], 50))
    vec_b[1: 10, 1] = vec_demand
    # mat_a
    mat_a_1 = zeros(10, num_x)
    for i = 1: 10
        mat_a_1[i, (4 * i - 3): (4 * i)] = repeat([1], 4)
    end
    mat_a_2 = zeros(num_x, num_x)
    for i = 1: num_x
        mat_a_2[i, i] = -1
    end
    mat_a = vcat(mat_a_1, mat_a_2)
    # mat_b
    mat_b_1 = zeros(10, 4)
    mat_b_2 = zeros(num_x, 4)
    epsilon = 0.000001  # epsilon must be large enough.
    # ??? It seems that the epsilon can make a difference in the final result, thought it's large enough.
    for i = 1: 10
        for j = 1: 4
            mat_b_2[(4 * i - 4 + j), j] = mat_prodOnMachine[i, j] * 1 / epsilon
        end
    end
    mat_b = vcat(mat_b_1, mat_b_2)
    # mat_a_bar
    mat_a_bar = zeros(4, num_x)
    for i = 1: 10
        for j = 1: 4
            mat_a_bar[j, (4 * i - 4 + j)] = vec_timeProd[i]
        end
    end
    mat_a_hat = zeros(4, num_x)
    for i = 1: 10
        for j = 1: 4
            mat_a_hat[j, (4 * i - 4 + j)] = vec_timeProdDev[i]
        end
    end
    mat_b_bar = zeros(4, 4)
    for i = 1: 4
        mat_b_bar[i, i] = - vec_timeAvai[i]
    end
    vec_b_bar = hcat(zeros(4))
    #
    return (vec_c, vec_f, vec_b, mat_a, mat_b, mat_a_bar, mat_a_hat, mat_b_bar, vec_b_bar)
end

function paras_stochasticLP(num_x, num_y, penalty)
    (vec_c, vec_f, vec_b, mat_a, mat_b, mat_a_bar, mat_a_hat, mat_b_bar, vec_b_bar) = paras_budgetMILP(num_x, num_y)
    #
    vec_g = hcat(penalty * ones(10, 1))
    mat_d_1 = zeros(10, 10)
    mat_d_1[diagind(mat_d_1)] .= 1
    mat_d_2 = zeros(num_x, 10)
    mat_d = vcat(mat_d_1, mat_d_2)
    mat_d_bar = zeros(4, 10)
    return (vec_c, vec_f, vec_b, mat_a, mat_b, mat_a_bar, mat_a_hat, mat_b_bar, vec_b_bar, vec_g, mat_d, mat_d_bar)
end


using ITensors
using Test
using Suppressor

include("util.jl")

@testset "SVD Algorithms" begin
  @testset "Matrix With Zero Sing Val" begin
    M = [
      1.0 2.0 5.0 4.0
      1.0 1.0 1.0 1.0
      0.0 0.5 0.5 1.0
      0.0 1.0 1.0 2.0
    ]
    U, S, V = NDTensors.svd_recursive(M)
    @test norm(U * LinearAlgebra.Diagonal(S) * V' - M) < 1E-13
  end

  @testset "Real Matrix" begin
    M = rand(10, 20)
    U, S, V = NDTensors.svd_recursive(M)
    @test norm(U * LinearAlgebra.Diagonal(S) * V' - M) < 1E-12

    M = rand(20, 10)
    U, S, V = NDTensors.svd_recursive(M)
    @test norm(U * LinearAlgebra.Diagonal(S) * V' - M) < 1E-12
  end

  @testset "Cplx Matrix" begin
    M = rand(ComplexF64, 10, 15)
    U, S, V = NDTensors.svd_recursive(M)
    @test norm(U * LinearAlgebra.Diagonal(S) * V' - M) < 1E-13

    M = rand(ComplexF64, 15, 10)
    U, S, V = NDTensors.svd_recursive(M)
    @test norm(U * LinearAlgebra.Diagonal(S) * V' - M) < 1E-13
  end

  @testset "Regression Test 1" begin
    # Implementation of the SVD was giving low
    # accuracy for this case
    M = rand(2, 2, 2, 2)

    M[:, :, 1, 1] = [
      7.713134067177845 -0.16367628720441685
      -1.5253996568409225 1.3577749944302373
    ]

    M[:, :, 2, 1] = [
      0.0 -2.1219889218225276
      -8.320068013774126 0.43565608213298096
    ]

    M[:, :, 1, 2] = [
      0.0 -8.662721825820825
      0.0 -0.46817091771736885
    ]

    M[:, :, 2, 2] = [
      0.0 0.0
      0.0 -8.159570989998151
    ]

    t1 = Index(2, "t1")
    t2 = Index(2, "t2")
    u1 = Index(2, "u1")
    u2 = Index(2, "u2")

    T = itensor(M, t1, t2, u1, u2)

    U, S, V = svd(T, (u1, t1))
    @test norm(U * S * V - T) / norm(T) < 1E-10
  end

  # TODO: remove this test, it takes a long time
  @testset "Ill-conditioned matrix" begin
    d = 5000
    i = Index(d, "i")
    T = itensor(make_illconditioned_matrix(dim(i)), i', i)

    @suppress begin
      F = svd(T, i'; alg="divide_and_conquer")
    end
    # Depending on the LAPACK implementation,
    # this sometimes works so don't test it
    #@test isnothing(F)

    F = svd(T, i'; alg="qr_iteration")
    @test !isnothing(F)
    @test F.U * F.S * F.V ≈ T
  end
end

nothing


using NodesAndModes
using StartUpDG
using UnPack
using LinearAlgebra
using SparseArrays
using FluxDiffUtils
using StaticArrays
using Plots

N = 7
K = 16
T = 1.5
CFL = .025

interval = 50

rd = RefElemData(Line(),N; quad_rule_vol=gauss_lobatto_quad(0,0,N))
VX,EToV = uniform_mesh(Line(),K)
md_BCs = MeshData(VX,EToV,rd)
md = make_periodic(md_BCs,rd)

@unpack M,Dr,Vf,LIFT = rd

Q = M*Dr
Q1 = diagm(1=>ones(N))-diagm(-1=>ones(N))
Q1[1,1] = -1
Q1[N+1,N+1] = 1
@. Q1 = .5*Q1
E = M*LIFT

fS(uL,uR) = @. (uL*uL + uL*uR + uR*uR)/6.0
dS(uL,uR) = @. .5*max(abs(uL),abs(uR))*(uL-uR)

function rhs(u,Qskew,E,rd,md)
    @unpack Vf = rd
    @unpack rxJ,nxJ,mapP = md
    uf = Vf*u
    return rxJ.*(Qskew*u) + .5*E*(uf[mapP].*nxJ)
end
function rhsf(u,E,rd,md)
    @unpack Vf = rd
    @unpack nxJ,mapP = md
    uf = Vf*u
    return .5*E*(uf[mapP].*nxJ)
end
function build_rhs_matrix(rhs::F,Np,K) where {F}
    A = zeros(Np*K,Np*K)
    u = zeros(Float64,Np,K)
    for i = 1:Np*K
        u[i] = 1.0
        A[:,i] .= vec(rhs(u))
        u[i] = 0.0
    end
    return droptol!(sparse(A),1e-12)
end
Qg = build_rhs_matrix(u->rhs(u,.5*(Q-Q'),E,rd,md),N+1,K)
Bg = build_rhs_matrix(u->rhsf(u,E,rd,md),N+1,K)
QgTr = -Qg # assume skew-symmetric
Q1g = build_rhs_matrix(u->rhs(u,.5*(Q1-Q1'),E,rd,md),N+1,K)
Q1gTr = -Q1g # assume skew-symmetric

# matrix used in convex limiting
ΔQ = .5*(Q1-Q1') - .5*(Q-Q')

function hadsum(ATr::SparseMatrixCSC,F::Fxn,u) where {Fxn}
    rhs = similar(u)
    rows = rowvals(ATr)
    vals = nonzeros(ATr)
    for i = 1:size(ATr,2)
        ui = u[i]
        val_i = zero(ui)
        for row_id in nzrange(ATr,i)
            j = rows[row_id]
            val_i += vals[row_id]*F(ui,u[j])
        end
        rhs[i] = val_i
    end
    return rhs
end

@unpack x,J = md
MJ = M*J
invm = 1 ./ vec(M*J)
x = vec(x)
u = @. (x > -1/3)*1.0

u1 = similar(u)
uN = similar(u)
rhsu = similar(u)
l = reshape(similar(u),N+1,K)
L = zeros(N+1,N+1)

dt = CFL*(x[2]-x[1]) / maximum(abs.(u))
Nsteps = ceil(Int,T/dt)
dt = T/Nsteps
ubounds = extrema(u)

init_mass = sum(vec(MJ).*u)
@gif for i = 1:Nsteps
    global u
    r1 = hadsum(Q1gTr,fS,u) + .5*hadsum(abs.(Q1gTr),dS,u)
    rN = hadsum(QgTr,fS,u) + .5*hadsum(abs.(Bg),dS,u)

    # low order update
    @. u1 = u - dt*invm*r1
    @. uN = u - dt*invm*rN

    u = reshape(u,N+1,K)

    # local limiting
    for e = 1:K
        u1e = @view reshape(u1,N+1,K)[:,e]
        uNe = @view reshape(uN,N+1,K)[:,e]
        uNmin = minimum(uNe)
        θ = max(0.0,min(1.0, minimum(@. u1e / (u1e-uNmin + 1e-7))))
        u[:,e] = @. u1e + θ*(uNe-u1e)
        # l[:,e] .= 0.0
        # if θ > 1e-10 # for plotting
            l[:,e] .= 1-θ
        # end
    end

    # # node-wise convex limiting
    # for e = 1:K
    #     uL,uR = NodesAndModes.meshgrid(@view u[:,e])
    #     A = ΔQ .* fS(uL,uR) - .5*abs.(.5*(Q1-Q1')) .* dS(uL,uR)
    #     λij = 1/(N+1)
    #     for i = 1:N+1
    #         for j = 1:N+1
    #             λij = MJ[i,e]/(N+1)
    #             L[i,j] = 1.0
    #             if u[i,e] + dt*A[i,j]/λij > 0 # need u[i,e] + lij * dt * Aij / λij > 0
    #                 L[i,j] = 1.0
    #             elseif abs(A[i,j]) > 1e-14
    #                 L[i,j] = λij * (1e-16 - u[i,e]) / A[i,j]
    #             end
    #         end
    #     end
    #     # fill!(L,minimum(L))
    #     u[:,e] = reshape(u1,N+1,K)[:,e] + dt*sum(min.(L,L').*A,dims=2) ./ @view MJ[:,e]
    #     l[:,e] = 1 .- sum(min.(L,L'),dims=2)/(N+1)
    # end

    u = vec(u)

    if i%interval==0
        println("$i / $Nsteps: $ubounds, $(extrema(u))")
        plot(x,u,mark=:dot,ms=2,ylims=ubounds .+ (-1,1))
        plot!(reshape(x,N+1,K),l,leg=false)
    end
end every interval

Δmass = sum(vec(MJ).*u)-init_mass
println("Δmass = $Δmass")


abstract type Shape{dim} end

getdim(::Shape{dim}) where {dim} = dim

#########################
# Triangles   #
#########################
struct Simplex{dim} <: Shape{dim} end
const Triangle = Simplex{2}
@inline get_num_faces(::Type{Triangle}) = 3
@inline get_num_vertices(::Type{Triangle}) = 3

function reference_coordinates(::Type{Triangle})
    [Tensors.Vec{2, Float64}((0.0, 0.0)),
    Tensors.Vec{2, Float64}((1.0, 0.0)),
    Tensors.Vec{2, Float64}((0.0, 1.0))]
end
function reference_edges(::Type{Triangle})
    [[Tensors.Vec{2, Float64}((1.0, 0.0)),Tensors.Vec{2, Float64}((0.0, 1.0))],
     [Tensors.Vec{2, Float64}((0.0, 1.0)),Tensors.Vec{2, Float64}((0.0, 0.0))],
     [Tensors.Vec{2, Float64}((0.0, 0.0)),Tensors.Vec{2, Float64}((1.0, 0.0))]]
end
@inline reference_edge_nodes(::Type{Triangle}) = ((1,2),(2,3),(3,1))

function gettopology(::Type{Triangle})
    return Dict(0=>3,1=>3,2=>1)
end

##########################
# Tetrahedron
######################

const Tetrahedron = Simplex{3}

get_num_faces(::Type{Tetrahedron}) = 3
get_num_vertices(::Type{Tetrahedron}) = 3

function reference_coordinates(::Type{Tetrahedron})
    [Tensors.Vec{3, Float64}((0.0, 0.0,0.0)),
     Tensors.Vec{3, Float64}((1.0, 0.0,0.0)),
     Tensors.Vec{3, Float64}((0.0,1.0, 0.0)),
     Tensors.Vec{3, Float64}((0.0,0.0, 1.0))]
end
function reference_edges(::Type{Tetrahedron})
    coords = reference_coordinates(Tetrahedron)
    [[coords[i[1]], coords[i[2]]] for i in reference_edge_nodes(Tetrahedron)]
end
function reference_faces(::Type{Tetrahedron})
    coords = reference_coordinates(Tetrahedron)
    [[coords[i[1]], coords[i[2]], coords[i[3]]] for i in reference_face_nodes(Tetrahedron)]
end
reference_edge_nodes(::Type{Tetrahedron}) = ((1,2),(2,3),(3,4),(4,1),(1,3),(2,4))
reference_face_nodes(::Type{Tetrahedron}) = ((1,2,3),(2,3,4),(1,3,4),(1,2,4))

function gettopology(::Type{Tetrahedron})
    return Dict(0=>4,1=>6,2=>4,3=>1)
end

"""
get points for a nodal basis of order `order` on a `dim`
    dimensional simplex
"""
function get_nodal_points(::Type{Simplex{dim}}, order) where {dim}
    points = Vector{Tensors.Vec{dim,Float64}}()
    vertices = reference_coordinates(Simplex{dim})
    topology = Dict{Int, Int}()
    append!(points, vertices)
    push!(topology, 0=>length(points))
    [append!(points, _interior_points(verts, order)) for verts in reference_edges(Simplex{dim})]
    push!(topology, 1=>length(points)-topology[0])
    if dim == 2
        append!(points, _interior_points(vertices, order))
        push!(topology, 2=>length(points)-topology[0]-topology[1])
    else
        [append!(points, _interior_points(verts, order)) for verts in reference_faces(Simplex{dim})]
        push!(topology, 2=>length(points)-topology[0]-topology[1])
        append!(points, _interior_points(vertices, order))
        push!(topology, 3=>length(points)-topology[0]-topology[1]-topology[2])
    end
    points, topology
end

function _interior_points(verts, order)
    n = length(verts)
    ls = [(verts[i] - verts[1])/order for i in 2:n]
    m = length(ls)
    grid_indices =  []
    if m == 1
        grid_indices = [[i] for i in 1:order-1]
    elseif m == 2 && order > 2
        grid_indices = [[i,j] for i in 1:order-1 for j in 1:order-i-1]
    elseif m == 3 && order > 3
        grid_indices = [[i,j,k] for i in 1:order-3 for j in 1:order-i-2 for k in 1:order-1-j-i]
    end
    pts = Vector{typeof(verts[1])}()
    for indices in grid_indices
        res = verts[1]
        for (i,ii) in enumerate(indices)
            res += (ii) * ls[m - i+1]
        end
        push!(pts,res)
    end
    pts
end

#########################
# Rectangles  #
#########################
struct HyperCube{dim} <: Shape{dim} end
const Rectangle = HyperCube{2}
@inline get_num_faces(::Type{Rectangle}) = 4
@inline get_num_vertices(::Type{Rectangle}) = 4

function reference_coordinates(::Type{Rectangle})
    [Tensors.Vec{2, Float64}((0.0, 0.0)),
    Tensors.Vec{2, Float64}((1.0, 0.0)),
    Tensors.Vec{2, Float64}((0.0, 1.0)),
    Tensors.Vec{2, Float64}((1.0, 1.0))]
end
function reference_edges(::Type{Rectangle})
    [[Tensors.Vec{2, Float64}((1.0, 0.0)),Tensors.Vec{2, Float64}((1.0, 1.0))],
     [Tensors.Vec{2, Float64}((1.0, 1.0)),Tensors.Vec{2, Float64}((0.0, 1.0))],
     [Tensors.Vec{2, Float64}((0.0, 1.0)),Tensors.Vec{2, Float64}((0.0, 0.0))],
     [Tensors.Vec{2, Float64}((0.0, 0.0)),Tensors.Vec{2, Float64}((1.0, 0.0))]]
end
@inline reference_edge_nodes(::Type{Rectangle}) = ((1,2),(2,3),(3,4),(4,1))

function gettopology(::Type{Rectangle})
    return Dict(0=>4,1=>4,2=>1)
end

##########################
# Hexahedron
######################

const Hexahedron = HyperCube{3}

get_num_faces(::Type{Hexahedron}) = 6
get_num_vertices(::Type{Hexahedron}) = 8

function reference_coordinates(::Type{Hexahedron})
    [Tensors.Vec{3}([0.0, 0.0, 0.0]),
    Tensors.Vec{3}([1.0, 0.0, 0.0]),
    Tensors.Vec{3}([1.0, 1.0, 0.0]),
    Tensors.Vec{3}([0.0, 1.0, 0.0]),
    Tensors.Vec{3}([0.0, 0.0, 1.0]),
    Tensors.Vec{3}([1.0, 0.0, 1.0]),
    Tensors.Vec{3}([1.0, 1.0, 1.0]),
    Tensors.Vec{3}([0.0, 1.0, 1.0])]
end

function reference_edges(::Type{Hexahedron} )
    coords = reference_coordinates(Hexahedron)
    [[coords[1],coords[2]],[coords[2],coords[3]],
    [coords[3],coords[4]],[coords[4],coords[1]],
    [coords[5],coords[6]],[coords[6],coords[7]],
    [coords[7],coords[8]],[coords[8],coords[5]],
    [coords[1],coords[4]],[coords[2],coords[3]],
    [coords[6],coords[7]],[coords[5],coords[8]]]
end

function reference_faces(::Type{Hexahedron})
    coords = reference_coordinates(Hexahedron)
    [[coords[1],coords[2],coords[3],coords[4]],
    [coords[5],coords[6],coords[7],coords[8]],
    [coords[2],coords[3],coords[7],coords[6]],
    [coords[1],coords[5],coords[8],coords[4]],
    [coords[1],coords[2],coords[5],coords[6]],
    [coords[3],coords[4],coords[8],coords[7]]]
end

reference_edge_nodes(::Type{Hexahedron}) = ((1,2),(2,3),(3,4),(4,1),(5,6),(6,7),(7,8),(8,5),(1,4),(2,3),(6,7),(5,8))
reference_face_nodes(::Type{Hexahedron}) = ((1,2,3,4),(6,5,8,7),(2,6,7,3),(1,4,8,5),(3,7,8,4),(1,5,6,2))

function gettopology(::Type{Hexahedron})
    return Dict(0=>8,1=>12,2=>6,3=>1)
end

"""
get points for a nodal basis of order `order` on a `dim`
    dimensional hypercube
"""
function get_nodal_points(::Type{HyperCube{dim}}, order) where {dim}
    points = Vector{Tensors.Vec{dim,Float64}}()
    vertices = reference_coordinates(HyperCube{dim})
    topology = Dict{Int, Int}()
    append!(points, vertices)
    push!(topology, 0=>length(points))
    [append!(points, _interior_points(verts, order)) for verts in reference_edges(HyperCube{dim})]
    push!(topology, 1=>length(points)-topology[0])
    if dim == 2
        append!(points, _cube_interior_points(vertices, order, dim))
        push!(topology, 2=>length(points)-topology[0]-topology[1])
    else
        [append!(points, _cube_interior_points(verts, order, dim)) for verts in reference_faces(HyperCube{dim})]
        push!(topology, 2=>length(points)-topology[0]-topology[1])
        append!(points, _cube_interior_points(vertices, order, dim))
        push!(topology, 3=>length(points)-topology[0]-topology[1]-topology[2])
    end
    points, topology
end

function _cube_interior_points(verts, order, dim)
    n = length(verts)
    ls = [(verts[i] - verts[1])/order for i in 2:n]
    m = length(ls)
    lx = ls[1]

    if dim == 2
        for point in ls
            if !any([x ≈ 0 for x in point])
                lx = point
            end
        end
        return [eltype(verts)((verts[1][1]+i*lx[1],verts[1][2]+j*lx[2])) for i in 1:order-1 for j in 1:order-1]
    elseif dim == 3
        if m == 3
            for point in ls
                if sum(point .≈ 0.0) == 1
                    lx = point
                end
            end
            if lx[1] ≈ 0.0
                return [eltype(verts)((verts[1][1],verts[1][2]+j*lx[2], verts[1][3]+k*lx[3])) for j in 1:order-1 for k in 1:order-1]
            elseif lx[2] ≈ 0.0
                return [eltype(verts)((verts[1][1]+i*lx[1],verts[1][2], verts[1][3]+k*lx[3])) for i in 1:order-1 for k in 1:order-1]
            elseif lx[3] ≈ 0.0
                return [eltype(verts)((verts[1][1]+i*lx[1],verts[1][2]+j*lx[2], verts[1][3])) for i in 1:order-1 for j in 1:order-1]
            end
        elseif m == 7 && order > 1
            for point in ls
                if !any([x ≈ 0 for x in point])
                    lx = point
                end
            end
            return [eltype(verts)((verts[1][1]+i*lx[1],verts[1][2]+j*lx[2], verts[1][3]+k*lx[3])) for i in 1:order-1 for j in 1:order-1 for k in 1:order-1]
        end
    else
        error("interior nodes for hypercube of dimension $dim not supported")
    end
    return []
end

############
# Segment
###########
const Segment = Simplex{1}
@inline get_num_faces(::Type{Segment}) = 1
@inline get_num_vertices(::Type{Segment}) = 2

function reference_coordinates(::Type{Segment})
    return [Tensors.Vec{1, Float64}((0.0,)),Tensors.Vec{1, Float64}((1.0,))]
end

function get_nodal_points(::Type{Segment}, order)
    points = Vector{Tensors.Vec{1,Float64}}()
    vertices = reference_coordinates(Segment)
    topology = Dict{Int, Int}()
    append!(points, vertices)
    push!(topology, 0=>length(points))
    append!(points, _interior_points(vertices, order))
    push!(topology, 1=>length(points)-topology[0])
    points, topology
end

### Utils
function map_shape_symbols(symbol)
    if symbol == :Triangle
        Triangle()
    elseif symbol == :Segment
        Segment()
    elseif symbol == :Hexahedron
        Hexahedron()
    elseif symbol == :Quad
        Rectangle()
    elseif symbol == :Tetrahedron
        Tetrahedron()
    else
        error("Shape not available")
    end
end


# Test on the idea of 0th order stream tracing.
# Assume the value is constant in each cell.

using FieldTracer, Meshes

"Check if tracing on mixed 2D unstructured grid passes."
function test_trace_unstructured2D()

   # Define coordinates
   coords = zeros(2,9)
   coords[:,1] = [1, 1]
   coords[:,2] = [3, 1]
   coords[:,3] = [4, 1]
   coords[:,4] = [1, 2]
   coords[:,5] = [2, 2]
   coords[:,6] = [1, 3]
   coords[:,7] = [2, 3]
   coords[:,8] = [3, 3]
   coords[:,9] = [4, 3]

   points = Point2[coords[:,i] for i in 1:size(coords,2)]
   tris = connect.([(2,3,9), (5,8,7), (5,2,8), (2,9,8)], Triangle)
   quads = connect.([(1,2,5,4),(4,5,7,6)], Quadrangle)
   mesh = SimpleMesh(points, [tris; quads])

   vx = fill(0.5, 6)
   vy = fill(0.3, 6)
   start = [1.5, 1.5]

   xs, ys = trace(mesh, vx, vy, [start[1]], [start[2]])

   if length(xs[1]) == 4 && xs[1][4] ≈ 4.000005 && ys[1][4] ≈ 3.000003
      return true
   else
      return false
   end
end

## Visualization
#=
using PyPlot
fig, ax = plt.subplots()
ax.set_aspect("equal")
plot(coords[1,:], coords[2,:], marker="o", ls="", color="crimson")
plot(xs, ys, ".-")
=#


import StatsBase.sample
function softmax_greedy(X, Y, stepsize, W, n_epoch, topk, exact=1)
    # SGD
    N, D = size(X)
    K = size(Y, 1)

    #n_epoch = 50
    #W = zeros(D, K)
    #tic()
    #XW = X*W
    #@show(toc())
    XT = X'
    for epoch = 1:n_epoch
        #grad = zeros(D, K)
        #XW = X*W

        perm = randperm(N)
        tic()
        prec = 0.0
        prec_rand = 0.0
        for i = 1:1000
            grad = zeros(D, K)
            sample_idx = perm[i]
            class = find(Y[:,sample_idx])[1]
            Wx = W'*XT[:,sample_idx]
            #get topk
            if exact == 1
                #exact topk
                top_idx = sortperm(Wx, rev=true)[1:topk]
            elseif exact == 2
                # greedy mips topk
                maxWx = zeros(K)
                for j in 1:K
                    maxWx[j] = maximum( W[:,j] .* XT[:,sample_idx] )
                end
                top_idx = sortperm(maxWx, rev=true)[1:topk]
            else
                # random topk
                top_idx = sample([1:K;], topk)
            end

            # check greedy mips
            maxWx = zeros(K)
            for j in 1:K
                maxWx[j] = maximum( W[:,j] .* XT[:,sample_idx] )
            end
            top_idx_gmips = sortperm(maxWx, rev=true)[1:topk]
            top_idx_true = sortperm(Wx, rev=true)[1:topk]
            top_idx_rand = sample([1:K;], topk)
            #println( sortperm(Wx, rev=true) )
            #if i == 500
                #@show(top_idx_gmips)
                #@show(top_idx_true)
                #@show( intersect(top_idx_gmips, top_idx_true) )
            #end
            prec += length(intersect(top_idx_gmips, top_idx_true[1:10]) )/10
            prec_rand +=length(intersect(top_idx_rand, top_idx_true[1:10]) )/10

            for kk in 1:topk
                if class == top_idx[kk]
                    deleteat!(top_idx, kk)
                    break
                end
            end

            noise_sample = []
            while length(noise_sample) < 10
                noise = rand(1:K)
                if noise in top_idx || noise == class
                    continue
                else
                    append!(noise_sample, noise)
                end
            end

            tmp = exp.(Wx[top_idx])
            deno_estimate = sum(tmp) + exp(Wx[class]) + (K - length(top_idx) - 1)/length(noise_sample) * sum( exp.(Wx[noise_sample]) )

            # for j = 1:K
            #     grad[:,j] = tmp[j]/deno_estimate * XT[:,sample_idx]
            # end
            # grad[:, class] -= XT[:,sample_idx]

            #update
            for ii in top_idx
                W[:,ii] -= stepsize* exp(Wx[ii])/deno_estimate * XT[:,sample_idx]
            end

            W[:,class] -= stepsize*(exp(Wx[class])/deno_estimate - 1) * XT[:,sample_idx]

        end
        @show(toq())
        @printf "%2.1e epoch complete \n" epoch

        # calculate loss
        prec = prec / 1000
        prec_rand = prec_rand / 1000
        @printf "prec@10 for gmips: %2.6e, prec@10 for random: %2.6e \n" prec prec_rand
        loss = 0
        risk = 0.0
        XW = X*W
        for i = 1:N
            Wxi = XW[i,:]
            class = find(Y[:,i])[1]
            deno = sum( exp.(Wxi) )
            loss += -log( exp(Wxi[class]) / deno )

            ii = indmax(Wxi)
            if ii != class
                risk += 1
            end
        end
        loss = loss / N
        loss = loss / log(2)
        risk = risk / N
        @printf "loss: %3.3e, risk %3.3e \n" loss risk

    end
    return W
end

# softmax
function softmax(X, Y, stepsize, W, n_epoch, topk)
    # SGD
    N, D = size(X)
    K = size(Y, 1)

    #n_epoch = 50
    #W = zeros(D, K)
    #tic()
    #XW = X*W
    #@show(toc())
    XT = X'
    for epoch = 1:n_epoch
        #grad = zeros(D, K)
        #XW = X*W

        perm = randperm(N)
        tic()
        for i = 1:1000
            #println(i)
            grad = zeros(D, K)
            sample_idx = perm[i]
            class = find(Y[:,sample_idx])[1]
            Wx = W'*XT[:,sample_idx]
            tmp = exp.(Wx)
            deno = sum(tmp)
            for j = 1:K
                grad[:,j] = tmp[j]/deno * XT[:,sample_idx]
            end
            grad[:, class] -= XT[:,sample_idx]

            #update
            W -= stepsize*grad
        end
        @show(toq())
        @printf "%2.1e epoch complete \n" epoch

        # show topk statistics
        # randomly choose 3 samples
        for i = 1:5
            sample_idx = rand(1:N)
            class = find(Y[:,sample_idx])[1]
            Wx = W'*XT[:,sample_idx]
            tmp = exp.(Wx)
            grad = zeros(D, K)
            deno = sum(tmp)
            for j = 1:K
                grad[:,j] = tmp[j]/deno * XT[:,sample_idx]
            end
            grad[:, class] -= XT[:,sample_idx]

            norm_grad = zeros(K)
            for j = 1:K
                norm_grad[j] = norm( grad[:,j], 1 )
            end

            sorted_tmp = sort(norm_grad, rev=true)
            total = sum(sorted_tmp)
            total_topk = sum(sorted_tmp[1:topk])
            true_tmp = norm_grad[class]
            # @show(norm_grad[class])
            # @show(sum(norm_grad))
            # @printf "    top %i propotion: %1.4e\n" topk total_topk/total
            # @printf "    true class propotion: %1.4e\n" true_tmp/total
        end

        # calculate loss
        loss = 0
        risk = 0.0
        XW = X*W
        for i = 1:N
            Wxi = XW[i,:]
            class = find(Y[:,i])[1]
            deno = sum( exp.(Wxi) )
            loss += -log( exp(Wxi[class]) / deno )

            ii = indmax(Wxi)
            if ii != class
                risk += 1
            end
        end
        loss = loss / N
        loss = loss / log(2)
        risk = risk / N
        @printf "loss: %3.3e, risk %3.3e \n" loss risk

    end
    return W
end

# pairwise rank loss
function train_pairwise(X, Y, stepsize, n_epoch, W)
    N, D = size(X)
    K = size(Y, 1)

    #W = zeros(D, K)
    XT = X'
    for epoch in 1:n_epoch
        perm = randperm(N)
        tic()
        for i = 1:N
            rand_class = rand(1:K)
            sample_idx = perm[i]
            class = find(Y[:,sample_idx])[1]
            if rand_class == class
                continue
            end
            tmp_class = dot(W[:,class], XT[:, sample_idx])
            tmp_other = dot(W[:,rand_class], XT[:, sample_idx])
            grad = exp(tmp_other - tmp_class) / (1 + exp(tmp_other - tmp_class)) * XT[:, sample_idx]

            # update
            W[:, class] += stepsize * grad
            W[:, rand_class] -= stepsize * grad

        end

        @show(toq())
        @printf "%2.1e epoch complete \n" epoch

        # calculate loss
        loss = 0
        risk = 0.0
        XW = X*W
        for i = 1:N
            Wxi = XW[i,:]
            class = find(Y[:,i])[1]
            #deno = sum( exp.(Wxi) )
            #loss += -log( exp(Wxi[class]) / deno )
            tmp_class = dot(W[:,class], XT[:, i])
            #for j = 1:K
            #    if j == class
            #        continue
            #    end
            #    tmp_other = XW[i,j]
            #    loss += log( 1 + exp(tmp_other - tmp_class ) )
            #end

            ii = indmax(Wxi)
            if ii != class
                risk += 1
            end
        end
        loss = loss / N
        loss = loss / log(2)
        risk = risk / N
        @printf "loss: %3.3e, risk %3.3e \n" loss risk

    end

    return W
end

import StatsBase.sample
# negative sampling
function train_neg_sampling(X, Y, stepsize, n_epoch, W, k)
    N, D = size(X)
    K = size(Y, 1)

    #W = zeros(D, K)
    XT = X'
    # set # neg samples to be 5
    #k = 5
    for epoch in 1:n_epoch
        perm = randperm(N)
        tic()
        for i = 1:N
            sample_idx = perm[i]
            # choose negative samples
            class = find(Y[:,sample_idx])[1]
            neg_samples = sample([1:K;], k)
            grad = 1 / (1 + exp( dot(W[:,class], XT[:,sample_idx]) )) * XT[:, sample_idx]
            # update
            W[:,class] += stepsize*grad
            for ss in neg_samples
                if ss == class
                    continue
                end
                grad = 1 / (1 + exp(-dot(W[:,ss], XT[:,sample_idx]))) * XT[:, sample_idx]
                # update
                W[:,ss] -= stepsize*grad
            end

        end

        @show(toq())
        @printf "%2.1e epoch complete \n" epoch

        # calculate loss
        loss = 0
        risk = 0.0
        XW = X*W
        for i = 1:N
            Wxi = XW[i,:]
            class = find(Y[:,i])[1]
            #deno = sum( exp.(Wxi) )
            #loss += -log( exp(Wxi[class]) / deno )
            tmp_class = dot(W[:,class], XT[:, i])
            #for j = 1:K
            #    if j == class
            #        continue
            #    end
            #    tmp_other = XW[i,j]
            #    loss += log( 1 + exp(tmp_other - tmp_class ) )
            #end

            ii = indmax(Wxi)
            if ii != class
                risk += 1
            end
        end
        loss = loss / N
        loss = loss / log(2)
        risk = risk / N
        @printf "loss: %3.3e, risk %3.3e \n" loss risk

    end

    return W
end


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

"""
    project_from_slave_to_master(Val{:Seg2}, xs, ns, xm1, xm2)

Find the projection of a slave node `xs`, having normal vector `ns`, onto master
elements with nodes (`xm1`, `xm2`). Returns master element dimensionless parameter
xi, that is,

    xm = 1/2*(1-xi)*xm1 + 1/2*(1+xi)*xm2

# Example

```jldoctest
xm1 = [7.0, 2.0]
xm2 = [4.0, -2.0]
xs = [0.0, 0.0]
ns = [3.0/5.0, -4.0/5.0]
xi2 = project_from_slave_to_master(Val{:Seg2}, xs, ns, xm1, xm2)
round(xi2, digits=6)

# output

1.833333
```

"""
function project_from_slave_to_master(::Type{Val{:Seg2}}, xs, ns, xm1, xm2)
    nom = ns[1]*(xm1[2] + xm2[2] - 2*xs[2]) - ns[2]*(xm1[1] + xm2[1] - 2*xs[1])
    denom = ns[1]*(xm1[2] - xm2[2]) + ns[2]*(xm2[1] - xm1[1])
    return nom/denom
end

"""
    project_from_master_to_slave(Val{:Seg2}, xm, xs1, xs2, ns1, ns2)

Find the projection of a master node `xm`, to the slave surface with nodes
(`xs1`, `xs2`), in direction of slave surface normal defined by (`ns1`, `ns2`).
Returns slave element dimensionless parameter, that is, to find coordinates
in slave side:

    xs = 1/2*(1-xi)*xs1 + 1/2*(1+xi)*xs2

# Example

```jldoctest
xm = [4.0, -2.0]
xs1 = [0.0, 0.0]
xs2 = [4.0, 3.0]
ns1 = [3.0/5.0, -4.0/5.0]
ns2 = sqrt(2)/2*[1, -1]
xi1 = project_from_master_to_slave(Val{:Seg2}, xm, xs1, xs2, ns1, ns2)
round(xi1, digits=6)

# output

-0.281575
```

"""
function project_from_master_to_slave(::Type{Val{:Seg2}}, xm, xs1, xs2, ns1, ns2)

    # Special case when normal is constant. Then we have
    # unique solution and linear equation to solve.
    if isapprox(ns1, ns2)
        nom = -2*ns1[1]*xm[2] + ns1[1]*xs1[2] + ns1[1]*xs2[2] + 2*ns1[2]*xm[1] - ns1[2]*xs1[1] - ns1[2]*xs2[1]
        denom = ns1[1]*xs1[2] - ns1[1]*xs2[2] - ns1[2]*xs1[1] + ns1[2]*xs2[1]
        return nom/denom
    end

    nom_b = 2*ns1[1]*xm[2] - 2*ns1[1]*xs1[2] - 2*ns1[2]*xm[1] + 2*ns1[2]*xs1[1] - 2*ns2[1]*xm[2] + 2*ns2[1]*xs2[2] + 2*ns2[2]*xm[1] - 2*ns2[2]*xs2[1]
    denom_b = ns1[1]*xs1[2] - ns1[1]*xs2[2] - ns1[2]*xs1[1] + ns1[2]*xs2[1] - ns2[1]*xs1[2] + ns2[1]*xs2[2] + ns2[2]*xs1[1] - ns2[2]*xs2[1]
    nom_c = -2*ns1[1]*xm[2] + ns1[1]*xs1[2] + ns1[1]*xs2[2] + 2*ns1[2]*xm[1] - ns1[2]*xs1[1] - ns1[2]*xs2[1] - 2*ns2[1]*xm[2] + ns2[1]*xs1[2] + ns2[1]*xs2[2] + 2*ns2[2]*xm[1] - ns2[2]*xs1[1] - ns2[2]*xs2[1]
    denom_c = ns1[1]*xs1[2] - ns1[1]*xs2[2] - ns1[2]*xs1[1] + ns1[2]*xs2[1] - ns2[1]*xs1[2] + ns2[1]*xs2[2] + ns2[2]*xs1[1] - ns2[2]*xs2[1]

    a = 1.0
    b = nom_b / denom_b
    c = nom_c / denom_c
    d = b^2 - 4*a*c
    if d < 0
        @warn("Mortar2D.calculate_projections(): negative discriminant $d")
        @warn("xm = $xm, xs1 = $xs1, xs2 = $xs2, ns1 = $ns1, ns2 = $ns2")
    end
    sols = [-b + sqrt(d), -b - sqrt(d)]/(2.0*a)
    return sols[argmin(abs.(sols))]
end


#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/build_literate/Analyze_NMR_PreSat_Data.ipynb)

# # Continuous Wave Saturation Experiments
# The following code analyzes data from a steady-state experiment similar to the original work of [Henkelman et al](https://doi.org/10.1002/mrm.1910290607). In this experiment, the magnetization of the coupled spin system is saturated with off-resonant continuous waves of the exponentially spaced frequencies:
Δ = exp.(range(log(0.01e3), log(100e3), length=20)) * 2π # rad/s
# and the amplitudes:
ω1_dB = -60:5:-5 # dB
ω1 = @. 10^(ω1_dB / 20) * π / 2 / 11.4e-6 # rad/s
# The waves were applied for 7 seconds to ensure a steady state. Thereafter, the magnetization was excited with a π/2-pulse and an FID was acquired. The repetition times was 30s to ensure full recovery to thermal equilibrium.

# We fit the data with Henkelman's closed form solution to this steady-state problem while assuming a Lorentzian lineshape for the free spin pool, and different lineshapes for the semi-solid spin pool:
g_Lorentzian(Δ, T2) = T2 / π / (1 + (T2 * Δ)^2)
g_Gaussian(Δ, T2) = T2 / sqrt(2π) * exp(-(T2 * Δ)^2 / 2)
g_superLorentzian(Δ, T2) = T2 * sqrt(2 / π) * quadgk(ct -> exp(- 2 *  (T2 * Δ / (3 * ct^2 - 1))^2) / abs(3 * ct^2 - 1), 0, sqrt(1 / 3), 1)[1];

# For this data analysis we need the following packages:
using MRIgeneralizedBloch
using QuadGK
using LsqFit
using LinearAlgebra
using Statistics
using Printf
using Plots
plotlyjs(bg=RGBA(31 / 255, 36 / 255, 36 / 255, 1.0), ticks=:native); # hide #!nb

# The raw data is stored in a separate [github repository](https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData) and the following functions return the URL to the individual files:
MnCl2_data(ω1_dB) = string("https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData/blob/main/20210419_1mM_MnCl2/ja_PreSat_v2%20(", ω1_dB, ")/1/data.2d?raw=true")
BSA_data(ω1_dB)   = string("https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData/blob/main/20210416_15%25BSA_2ndBatch/ja_PreSat_v2%20(", ω1_dB, ")/1/data.2d?raw=true");

# which can be loaded with functions implemented in this file:
include(string(pathof(MRIgeneralizedBloch), "/../../docs/src/load_NMR_data.jl"));

# We store the off-resonance frequencies and wave amplitudes in a single matrix for convenience:
x = zeros(Float64, length(ω1) * length(Δ), 2)
x[:,1] = repeat(Δ, length(ω1))
x[:,2] = vec(repeat(ω1, 1, length(Δ))');

# ## MnCl``_2`` Sample
# We load the first data point of each FID:
M = zeros(Float64, length(Δ), length(ω1))
for i = 1:length(ω1_dB)
    M[:,i] = load_first_datapoint(MnCl2_data(ω1_dB[i]); set_phase=:abs)
end
M ./= maximum(M);
# In contrast to the inversion-recovery experiment, the phase of the signal was not stable. Therefore, we took the absolute value of the signal by setting the flag `set_phase=:abs`.

# The MnCl``_2``-data can be described with a single compartment model:
function single_compartment_model(x, p)
    (m0, R1, T2) = p

    Δ  = @view x[:,1]
    ω1 = @view x[:,2]

    Rrf = @. π * ω1^2 * g_Lorentzian(Δ, T2)
    m = @. m0 * R1 / (R1 + Rrf)
    return m
end;
# (cf. Eqs. (14) and (15) in the [paper](https://arxiv.org/pdf/2107.11000.pdf)).

# As this model is merely a function of the relaxation times T₁ and T₂, we forgo a fitting routine and use the estimates from the [Inversion Recovery Experiments](@ref) instead:
R1 = 1.479 # 1/s
T2 = 0.075 # s
nothing # hide #md

# Visually, this model describes the data well:
p = plot(xlabel="Δ [rad/s]", ylabel="M / max(M)", xaxis=:log, legend=:none)
[scatter!(p, Δ, M[:,i], color=i) for i=1:length(ω1)]
[plot!(p, Δ, reshape(single_compartment_model(x, [1,R1,T2]), length(Δ), length(ω1))[:,i], color=i) for i=1:length(ω1)]
display(p) #!md
#md Main.HTMLPlot(p) #hide

# ## Bovine Serum Albumin Sample
# We acquired the same data for the BSA sample, which we load:
for i = 1:length(ω1_dB)
    M[:,i] = load_first_datapoint(BSA_data(ω1_dB[i]); set_phase=:abs)
end
M ./= maximum(M);

# We model the steady-state magnetization as described by [Henkelman et al.](https://doi.org/10.1002/mrm.1910290607):
function Henkelman_model(x, p; lineshape=:superLorentzian)
    (m0, m0s, R1f, R1s, T2f, T2s, Rx) = p

    m0s /= 1 - m0s # switch from m0s + m0f = 1 to m0f = 1 normalization

    Δ  = @view x[:,1]
    ω1 = @view x[:,2]

    Rrf_f = @. π * ω1^2 * g_Lorentzian(Δ, T2f)

    if lineshape == :Lorentzian
        Rrf_s = @. π * ω1^2 * g_Lorentzian(Δ, T2s)
    elseif lineshape == :Gaussian
        Rrf_s = @. π * ω1^2 * g_Gaussian(Δ, T2s)
    elseif lineshape == :superLorentzian
        Rrf_s = @. π * ω1^2 * g_superLorentzian(Δ, T2s)
    end

    m = @. m0 * (R1s * Rx * m0s + Rrf_s * R1f + R1f * R1s + R1f * Rx) / ((R1f + Rrf_f + Rx * m0s) * (R1s + Rrf_s + Rx) - Rx^2 * m0s)
    return m
end;

# Here, we use a fitting routine to demonstrate the best possible fit with each of the three lineshapes. We define an initialization for the fitting routine `p0 = [m0, m0s, R1f, R1s, T2f, T2s, Rx]` and set some reasonable bounds:
p0   = [  1,0.01,   1,   5,0.052,  1e-5, 40]
pmin = [  0,   0,   0,   0,0.052,  1e-6,  1]
pmax = [Inf,   1, Inf, Inf,0.052, 10e-3,100];
# Note that we fixed T₂ᶠ = 52ms to the value estimated with the [Inversion Recovery Experiments](@ref) as T₂ᶠ is poorly defined by this saturation experiment.

# ### Super-Lorentzian Lineshape
# Fitting the model with a super-Lorentzian lineshape to the data achieves good concordance:
fit = curve_fit((x, p) -> Henkelman_model(x, p; lineshape=:superLorentzian), x, vec(M), p0, lower=pmin, upper=pmax)
fit_std = stderror(fit)

p = plot(xlabel="Δ [rad/s]", ylabel="M / max(M)", xaxis=:log, legend=:none)
[scatter!(p, Δ, M[:,i], color=i) for i=1:length(ω1)]
[plot!(p, Δ, reshape(Henkelman_model(x, fit.param), length(Δ), length(ω1))[:,i], color=i) for i=1:length(ω1)]
display(p) #!md
#md Main.HTMLPlot(p) #hide

# ### Lorentzian Lineshape
# The Lorentzian lineshape, on the other hand, does not fit the data well:
fit = curve_fit((x, p) -> Henkelman_model(x, p; lineshape=:Lorentzian), x, vec(M), p0, lower=pmin, upper=pmax)
fit_std = stderror(fit)

p = plot(xlabel="Δ [rad/s]", ylabel="M / max(M)", xaxis=:log, legend=:none)
[scatter!(p, Δ, M[:,i], color=i) for i=1:length(ω1)]
[plot!(p, Δ, reshape(Henkelman_model(x, fit.param), length(Δ), length(ω1))[:,i], color=i) for i=1:length(ω1)]
display(p) #!md
#md Main.HTMLPlot(p) #hide

# ### Gaussian Lineshape
# And the Gaussian lineshape does not not fit the data well either:
fit = curve_fit((x, p) -> Henkelman_model(x, p; lineshape=:Lorentzian), x, vec(M), p0, lower=pmin, upper=pmax)
fit_std = stderror(fit)

p = plot(xlabel="Δ [rad/s]", ylabel="M / max(M)", xaxis=:log, legend=:none)
[scatter!(p, Δ, M[:,i], color=i) for i=1:length(ω1)]
[plot!(p, Δ, reshape(Henkelman_model(x, fit.param), length(Δ), length(ω1))[:,i], color=i) for i=1:length(ω1)]
display(p) #!md
#md Main.HTMLPlot(p) #hide


# mpirun -n 16 julia MPI_forward.jl
using Revise
using ADSeismic
using ADCME
using JLD2
ADCME.options.customop.verbose = false
mpi_init()
M = N = Int(sqrt(mpi_size()))
n = 100÷M

param = MPIElasticPropagatorParams(NX=100, NY=100, n=n, NSTEP=1000, DELTAT=1e-4,
    DELTAX=1.0, DELTAY=1.0,
    USE_PML_XMIN=true, USE_PML_XMAX = true, USE_PML_YMIN = true, USE_PML_YMAX = true)
compute_PML_Params!(param)
vp = 3300.
vs = 3300. / 1.732
rho = 2800.
λ, μ, ρ = compute_lame_parameters(param, vp, vs, rho)
source = Ricker(param, 50.0, 200.0)

# close("all")
# plot(run(sess, source))
# savefig("test.png")
srci = [div(param.NX,5)]
srcj = [div(param.NY,2)]
srctype = [0]
srcv = reshape(source, :, 1)
src = MPIElasticSource(param, srci, srcj, srctype, srcv)

propagator = MPIElasticPropagatorSolver(param, src, ρ, λ, μ)

sess = Session(); init(sess)
@info "starting..."
Vx, Vy, Sxx, Syy, Sxy = run(sess, [propagator.vx, propagator.vy, propagator.sigmaxx, propagator.sigmayy, propagator.sigmaxy])

if mpi_rank()==0 && mpi_size()==1
    # close("all")
    # pcolormesh(rcvv, cmap="gray")
    # colorbar()
    # xlabel("Location")
    # ylabel("Time")
    # savefig("receiver.png")

    close("all")
    p = visualize_wavefield(Vx, param)
    saveanim(p, "elastic_wavefied.gif")

end

if mpi_size()>1
    mpi_finalize()
end


import KernelDensity.UnivariateKDE
o1 = UncertainValue(Normal, 0, 0.5)
o2 = UncertainValue(Normal, 2.0, 0.1)
o3 = UncertainValue(Uniform, 0, 4)
o4 = UncertainValue(Uniform, rand(100))
o5 = UncertainValue(Beta, 4, 5)
o6 = UncertainValue(Gamma, 4, 5)
o7 = UncertainValue(Frechet, 1, 2)
o8 = UncertainValue(BetaPrime, 1, 2)
o9 = UncertainValue(BetaBinomial, 10, 3, 2)
o10 = UncertainValue(Binomial, 10, 0.3)
o11 = UncertainValue(rand(100), rand(100))
o12 = UncertainValue(2)
o13 = UncertainValue(2.3)
o14 = UncertainValue([2, 3, 4], [0.3, 0.4, 0.1])

uvals = [o1, o2, o3, o4, o5, o6, o7, o8, o9, o11, o12, o13, o14]
D = UncertainDataset(uvals)
UV = UncertainValueDataset(uvals)
UIDX = UncertainIndexDataset(uvals)
UIV = UncertainIndexValueDataset(UIDX, UV)

n = length(D)
@test resample(D) isa Vector
@test resample(UV) isa Vector
@test resample(UIDX) isa Vector
@test resample(D, 10) isa Vector
@test resample(UV, 10) isa Vector
#@test resample(UIV, 10) isa Vector

######################################################################
# Resampling datasets consisting of uncertain values furnished by
# theoretical distributions
######################################################################
measurements = [UncertainValue(Normal, 0, 0.1) for i = 1:5]
d = UncertainDataset(measurements)

@test resample(d) isa Vector{T} where T <: Real
@test resample(d, 10) isa Vector{Vector{T}} where T <: Real

##########################################################################
# Resampling datasets consisting of uncertain values furnished by
# theoretical distributions with parameters estimated from empirical data.
##########################################################################
measurements = [UncertainValue(Normal, rand(Normal(), 1000)) for i = 1:5]
d = UncertainDataset(measurements)

@test resample(d) isa Vector{T} where T <: Real
@test resample(d, 10) isa Vector{Vector{T}} where T <: Real

##########################################################################
# Resampling datasets consisting of uncertain values furnished by
# kernel density estimations to the distributions.
##########################################################################
measurements = [UncertainValue(UnivariateKDE, rand(Normal(), 1000)) for i = 1:5]
d = UncertainDataset(measurements)
@test resample(d) isa Vector{T} where T <: Real
@test resample(d, 10) isa Vector{Vector{T}} where T <: Real

##########################################################################
# Resampling datasets consisting of a mixture of different types of
# uncertain values
##########################################################################
measurements = [UncertainValue(rand(100));
                UncertainValue(Normal, rand(Normal(), 100));
                UncertainValue(Normal, 0, 2)]
d = UncertainDataset(measurements)
@test resample(d) isa Vector{T} where T <: Real
@test resample(d, 10) isa Vector{Vector{T}} where T <: Real

iv = UncertainIndexValueDataset(measurements, measurements)

@test resample(iv) isa Tuple{Vector{T}, Vector{T}} where T <: Real
@test resample(iv, 5) isa  Vector{Tuple{Vector{T}, Vector{T}}} where T <: Real
@test length(resample(iv, 5)) == 5


import Base: *, +

using FFTW

abstract type AbstractOperator end

struct PointWiseSpectralOperator{O} <: AbstractOperator
    operation::O
end

*(a::AbstractField, b::AbstractField) = a.data .* b.data
+(a::AbstractField, b::AbstractField) = a.data .+ b.data

function (operator::PointWiseSpectralOperator)(a::ScalarField)
    â = fft(a.data)
    data = operator.operation .* â
    return ifft(data)
end

function (operator::PointWiseSpectralOperator)(a::AbstractArray)
    â = fft(a)
    data = operator.operation .* â
    return ifft(data)
end


using InteractiveUtils
using BenchmarkTools
using PyCall

versioninfo()

a = rand(10^7) # 1D vector of random numbers, uniform on [0,1)

function mysum(arr)
    s = 0.0
    for a in arr
        s+=a
    end
    s
end

function mysum_simd(arr)
    s = 0.0
    @simd for a in arr
        s += a
    end
    s
end

j_bench=@benchmark sum($a)
println("Julia embedded sum: Fastest time was $(minimum(j_bench.times) / 1e6) msec")

npsum=pyimport("numpy")["sum"]
np_bench=@benchmark npsum($a)
println("np.sum: Fastest time was $(minimum(np_bench.times) / 1e6) msec")

j_bench_hand = @benchmark mysum($a)
println("mysum: Fastest time was $(minimum(j_bench_hand.times)/1e6) msec")

elapsed=@elapsed mysum(a)
@show elapsed

elapsed=@elapsed mysum_simd(a)
@show elapsed

# このコードは動かない・・・
#j_bench_hand_simd @benchmark $(mysum_simd(a))
#println("mysum_simd: Fastest time was $(minimum(j_bench_hand_simd.times)/1e6) msec")

#= output
Julia Version 0.7.0
Commit a4cb80f3ed (2018-08-08 06:46 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin14.5.0)
  CPU: Intel(R) Core(TM) M-5Y51 CPU @ 1.10GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.0 (ORCJIT, broadwell)
Julia embedded sum: Fastest time was 5.339243 msec
np.sum: Fastest time was 5.530426 msec
mysum: Fastest time was 11.793079 msec
elapsed = 0.025699647
elapsed = 0.026391872
=#


# This file is a part of NeighborhoodApproximationIndex.jl

using SimilaritySearch
using Test

function test_knr_vectors(perf, knr, lower)
    @time p = probe(perf, knr)
    @test p.macrorecall >= lower
    p
end

@testset "indexing vectors with Knr" begin
    # NOTE: The following algorithms are complex enough to say we are testing
    # it doesn't have syntax errors and achieves some expected quality
    # a more grained test functions are required
    k = 10
    n = 10000 # number of items in the dataset
    m = 100  # number of queries
    dim = 8  # vector's dimension
    db = [rand(Float32, dim) for i in 1:n]
    queries = [rand(Float32, dim) for i in 1:m]

    dist = L2Distance()
    seq = ExhaustiveSearch(dist, db)
    perf = Performance(seq, queries, k)

    knr = Knr(dist, db; numrefs=64, kbuild=2)
    p = test_knr_vectors(perf, knr, 0.7)
    @info "-- $(string(knr)) $(p)"
    optimize!(perf, knr; recall=0.95)
    p = test_knr_vectors(perf, knr, 0.95)
    @info "-- Optimized $(string(knr)) $(knr.opts) -- $(p)"
end


cd(@__DIR__)
using Pkg
Pkg.activate(@__DIR__); Pkg.instantiate()

using OrdinaryDiffEq, Test

include("prinz_neuron_setup.jl")

byhand_prob = ODEProblem{true}(prinz_neuron!, u0, (0.,2000.), p)
byhand_sol = solve(byhand_prob, Rosenbrock23(), reltol=1e-8, abstol=1e-8, saveat=0.025);
current_mtk_sol = solve(sim, Rosenbrock23(), reltol=1e-8, abstol=1e-8, saveat=0.025);

tsteps = 0.0:0.025:2000.0
byhand_out = Array(byhand_sol(tsteps, idxs=2))
current_mtk_out = current_mtk_sol[neuron.sys.Vₘ]

@test isapprox(byhand_out, current_mtk_out, rtol=0.001)


### Routines for constructing smooth functions that generate field data ###

import Base: *, -, +, show

export AbstractSpatialField, Gaussian, radius, center, strength,
       EmptySpatialField,
       SpatialGaussian,
       GeneratedField, datatype, grid,
       PulseField

abstract type AbstractSpatialField end
abstract type AbstractGeneratedField end

## Empty spatial field

"""
    EmptySpatialField()

Create a blank spatial field. This is primarily useful for initializing a
sum of spatial fields.

# Example
```jldoctest
julia> g = EmptySpatialField()
EmptySpatialField()

julia> g(2,3)
0.0
```
"""
struct EmptySpatialField <: AbstractSpatialField end
(g::EmptySpatialField)(x,y) = Float64(0)

## Gaussian

"""
    Gaussian(σ,x0,A)

Construct a 1-d Gaussian function centered at `x0` with standard deviation `σ`
and amplitude `A`. The resulting function can be evaluated at any real-valued
number.

# Example

```jldoctest
julia> g = Gaussian(0.2,0,1)
Gaussian(0.2, 0, 1, 2.8209479177387813)

julia> g(0.2)
1.0377687435514866
```
"""
struct Gaussian
  σ :: Real
  x0 :: Real
  A :: Real
  fact :: Float64
end

Gaussian(σ,x0,A) = Gaussian(σ,x0,A,A/sqrt(π*σ^2))

radius(g::Gaussian) = g.σ
center(g::Gaussian) = g.x0
strength(g::Gaussian) = g.A

@inline gaussian(r;tol=6.0) = abs(r) < tol ? exp(-r^2) : 0.0

(g::Gaussian)(x) = g.fact*gaussian((x-center(g))/radius(g))

## Spatial Gaussian field ##

struct SpatialGaussian <: AbstractSpatialField
  gx :: Gaussian
  gy :: Gaussian
  A :: Real
end

SpatialGaussian(σx::Real,σy::Real,x0::Real,y0::Real,A::Real) =
          SpatialGaussian(Gaussian(σx,x0,A),Gaussian(σy,y0,1),A)

SpatialGaussian(σ::Real,x0::Real,y0::Real,A::Real) = SpatialGaussian(σ,σ,x0,y0,A)

(g::SpatialGaussian)(x,y) = g.gx(x)*g.gy(y)

## Scaling spatial fields

struct ScaledField{N <: Real, P <: AbstractSpatialField} <: AbstractSpatialField
    s::N
    p::P
end
function show(io::IO, p::ScaledField)
    print(io, "$(p.s) × ($(p.p))")
end
s::Number * p::AbstractSpatialField = ScaledField(s, p)
-(p::AbstractSpatialField) = ScaledField(-1, p)
(p::ScaledField)(x,y) = p.s*p.p(x,y)

struct AddedFields{T <: Tuple} <: AbstractSpatialField
    ps::T
end
function show(io::IO, Σp::AddedFields)
    println(io, "AddedFields:")
    for p in Σp.ps
        println(io, "  $p")
    end
end

## Adding spatial fields together

"""
    p₁::AbstractSpatialField + p₂::AbstractSpatialField

Add the fields so that `(p₁ + p₂)(x,y) = p₁(x,y) + p₂(x,y)`.

"""
+(p::AbstractSpatialField, Σp::AddedFields) = AddedFields((p, Σp.ps...))
+(Σp::AddedFields, p::AbstractSpatialField) = AddedFields((Σp.ps..., p))
function +(Σp₁::AddedFields, Σp₂::AddedFields)
    AddedFields((Σp₁..., Σp₂...))
end

-(p₁::AbstractSpatialField, p₂::AbstractSpatialField) = p₁ + (-p₂)

+(p::AbstractSpatialField...) = AddedFields(p)

function (Σp::AddedFields)(x,y)
    f = 0.0
    for p in Σp.ps
        f += p(x,y)
    end
    f
end

## For generating an actual instance of a spatial field
"""
    GeneratedField(d::GridData,field::AbstractSpatialField...,grid::PhysicalGrid)

Create an instance of a spatial field function `field` on scalar grid data `d`,
based on a grid `grid`. After creating the instance `g = GeneratedField(d,field,grid)`,
then the resulting grid data can be accessed by typing `g()`. For vector grid data,
a separate `field` must be supplied for each component.
"""
struct GeneratedField{T <: GridData}
    fielddata :: T
    fieldfcns :: Vector{AbstractSpatialField}
    grid :: PhysicalGrid
end

function GeneratedField(d::ScalarGridData,field::AbstractSpatialField,g::PhysicalGrid)
    xg, yg = coordinates(d,g)
    tmp = _generatedfield(xg,yg,field,d)
    GeneratedField{typeof(tmp)}(tmp,AbstractSpatialField[field],g)
end

function GeneratedField(d::CollectedGridData,fields::Vector{AbstractSpatialField},g::PhysicalGrid)
  @assert length(fields) == _numberofcomponents(typeof(d))
  tmp = similar(d)
  cnt = 0
  for (i,fname) in enumerate(propertynames(d))
    if typeof(getfield(d,fname)) <: GridData
        cnt += 1
        x, y = coordinates(getfield(d,fname),g)
        getfield(tmp,fname) .= _generatedfield(x,y,fields[cnt],getfield(d,fname))
    end
  end
  GeneratedField{typeof(tmp)}(tmp,fields,g)
end

GeneratedField(d::VectorGridData,
      fieldu::AbstractSpatialField,fieldv::AbstractSpatialField,
      g::PhysicalGrid) = GeneratedField(d,AbstractSpatialField[fieldu,fieldv],g)

_generatedfield(xg,yg,field::AbstractSpatialField,d::ScalarGridData) =
          typeof(d)(field.(xg*ones(1,length(yg)),ones(length(xg))*yg'))

datatype(f::GeneratedField{T}) where {T} = T
grid(f::GeneratedField{T}) where {T} = f.grid

(f::GeneratedField{T})() where {T} = f.fielddata

## For generating a transient form of a spatial field

"""
    PulseField(g::GeneratedField,t0,σ)

Create a transient form of a generated spatial field, useful for
introducing a short-lived forcing field onto the grid. The supplied field `g`
is modulated by a Gaussian pulse centered at time `t0` with width `σ`. This pulse's
maximum is unity, so that the amplitude of the overall field is set by `g`. The
resulting object can be evaluated with a single argument (time) and returns a
`ScalarGridData` type object of the same type contained in `g`.
"""
struct PulseField
  gfield :: GeneratedField
  timemod :: Gaussian
end

PulseField(gfield::GeneratedField{T},t0::Real,σt::Real) where {T <: GridData} =
            PulseField(gfield,Gaussian(σt,t0,sqrt(π*σt^2)))

(f::PulseField)(t) = f.timemod(t)*f.gfield()

datatype(f::PulseField) = datatype(f.gfield)

## Allow all spatial fields to accept a vector argument

#for T in InteractiveUtils.subtypes(AbstractSpatialField)
#    @eval (f::$T)(x::Vector{R}) where {R <: Float64} = f(x...)
#end


function correlate_eigenmodes(U, M, V; tracked_modes=1:size(U,1), rtol=1e-3, atol=1e-3)

    # get number of (not necessarily unique) eigenvalues
    nev = size(U,1)

    # construct the correlation matrix
    C = U*M*V

    # get the row permutation that puts maximum values on the diagonals
    rval = Vector{real(eltype(C))}(undef, nev) # magnitude of values in chosen row
    cval = Vector{real(eltype(C))}(undef, nev) # magnitude of values in chosen column
    perm = zeros(Int, nev) # permutation
    corruption = Vector{real(eltype(C))}(undef, nev) # corruption index for each correlation

    # tracked modes
    for i in tracked_modes

        # get correlation matrix entry magnitudes for this row
        for j = 1:nev
            rval[j] = abs(C[i,j])
        end

        # rank modes correlations from best to worst for this row
        ranked_modes = sortperm(rval, rev=true)

        # choose the best unassigned mode correlation for this row
        j1 = ranked_modes[findfirst((x) -> !(x in perm), ranked_modes)]

        # get correlation matrix entry magnitudes for the chosen column
        for k = 1:nev
            cval[k] = abs(C[k,j1])
        end

        # find approximate multiplicity of left and/or right eigenvector
        multiplicity = max(
            count(x -> isapprox(x, rval[j1]; rtol, atol), rval),
            count(x -> isapprox(x, cval[i]; rtol, atol), cval)
        )

        if multiplicity < nev
            # determine n+1 best unassigned fit, where n is the approximate multiplicity
            if j1 in ranked_modes[1:multiplicity]
                j2 = ranked_modes[multiplicity+1]
            else
                j2 = ranked_modes[1]
            end

            # assign best eigenmode fit, create corruption index
            perm[i] = j1
            corruption[i] = rval[j2]/rval[j1]
        else
            # assign best eigenmode fit, create corruption index
            perm[i] = j1
            corruption[i] = 0.0
        end
    end

    # remaining modes
    for i = 1:nev

        if !(i in tracked_modes)

            # get correlation matrix entry magnitudes for this row
            for j = 1:nev
                rval[j] = abs(C[i,j])
            end

            # rank modes correlations from best to worst for this row
            ranked_modes = sortperm(rval, rev=true)

            # choose the best unassigned mode correlation for this row
            j1 = ranked_modes[findfirst((x) -> !(x in perm), ranked_modes)]

            # get correlation matrix entry magnitudes for the chosen column
            for k = 1:nev
                cval[k] = abs(C[k,j1])
            end

            # find approximate multiplicity of left and/or right eigenvector
            multiplicity = max(
                count(x -> isapprox(x, rval[j1]; rtol, atol), rval),
                count(x -> isapprox(x, cval[i]; rtol, atol), cval)
            )

            # determine n+1 best unassigned fit, where n is the approximate multiplicity
            if j1 in ranked_modes[1:multiplicity]
                j2 = ranked_modes[multiplicity+1]
            else
                j2 = ranked_modes[1]
            end

            # assign best eigenmode fit, create corruption index
            perm[i] = j1
            corruption[i] = rval[j2]/rval[j1]
        end

    end

    return perm, corruption
end


const _hermite = OrthogonalPolynomialsQuasi.Hermite()
# perhaps call fastGaussQuadrature directly
const DEFAULT_HERMITE_INTEGRATOR = sqrt(2π)*expectation(Normal(); n=100)
const DEFAULT_SQUARED_HERMITE_INTEGRATOR = sqrt(π)*expectation(Normal(0, 1/sqrt(2)); n=101)

# physicist's hermite polynomials

function hermite_fun_noexp(xs, n)
	xs = big.(xs)
	n = big.(n )
	Float64.(_hermite[xs, n+1]./sqrt(2^n*factorial(n)*sqrt(pi)))
end

function hermite_fun(xs, n)
	hermite_fun_noexp(xs, n).*exp.(-xs.^2/2)
end

hermite_cf_coefficient(n) = im^n*sqrt(2*π)

#fourier_transform(t, order) = hermite_fun(t, order)*hermite_cf_coefficient(order)

# Hermite = Mixture of Hermite Basis Functions
struct HermiteBasisFunction
	q::Integer
end

(herm::HermiteBasisFunction)(x) = MCEB.hermite_fun(x, herm.q)

function expectation(herm::HermiteBasisFunction,
	                 base_integrator::Expectation = DEFAULT_HERMITE_INTEGRATOR)
	hermite_integrator = deepcopy(base_integrator)
	wts = hermite_fun_noexp.(nodes(hermite_integrator), herm.q)
	Expectations.weights(hermite_integrator) .*= wts
	hermite_integrator
end

# wow this should really be dispatching on integrator
function (target::MarginalDensityTarget{<:DiscretizedStandardNormalSample})(integrator::Expectation)
    Z_disc = location(target)
    grid = Z_disc.samples.mhist.grid
    i = Z_disc.bin
    if i == 1
        base_integrand = μ -> cdf(Normal(μ), first(grid))
    elseif i == lastindex(Z_disc.samples.mhist)
        base_integrand = μ -> ccdf(Normal(μ), last(grid))
    else
        base_integrand = μ -> (cdf(Normal(μ), grid[i]) -  cdf(Normal(μ), grid[i-1]))
    end
    # refactor into riesz representer.. -> Oh already thought about it...
    integrator(base_integrand)
end

# Here ``g^*(t) = \\int \\exp(2 i \\pi \\mu) g(\\mu)d\\mu$ is the Fourier transform of ``g``.
"""
	HermitePriorClass(qmax, sobolev_order, sobolev_bound, solver)

Class of densities ``g \\in \\mathcal{G}`` that can be written as ``g(\\mu) = \\sum_{j=0}^{qmax} h_j(\\mu)``, where
``h_j`` is the ``j``-th Hermite function and further satisfy the Sobolev constraint
(with `sobolev_order` corresponding to ``b`` and `sobolev_bound` to ``C``):

```math
\\int_{-\\infty}^{\\infty} |g^*(t)|^2(t^2+1)^b dt \\leq 2\\pi C
```
The `solver` object is a `MathOptInterface` compatible optimizer such as `Mosek.Optimizer` that will be
used to solve the convex modulus of continuity (or worst-case-bias estimation) problem.
"""
Base.@kwdef mutable struct HermitePriorClass{T<:Real} <: ConvexPriorClass
	qmax::Integer
	sobolev_order::T  #not necessarily but should suffice for now
	sobolev_bound::Float64
	integrators = expectation.(HermiteBasisFunction.(0:qmax), Ref(DEFAULT_HERMITE_INTEGRATOR))
	squared_integrator = DEFAULT_SQUARED_HERMITE_INTEGRATOR
	solver = nothing
	solver_params = []
end

function add_prior_variables!(model::JuMP.Model, hermite_class::HermitePriorClass; var_name = "πs") # adds constraints
    qmax = hermite_class.qmax
	nparams = qmax + 1
    tmp_vars = @variable(model, [i=1:nparams])
    model[Symbol(var_name)] = tmp_vars
	cf_coefs = hermite_cf_coefficient.(0:qmax) .* hermite_fun_noexp.(0.0, 0:qmax)
	con = @constraint(model, dot(tmp_vars, real.(cf_coefs)) == 1.0)
	con2 = @constraint(model, dot(tmp_vars, imag.(cf_coefs)) == 0.0)
	#tmp_grid = -7:0.02:4
	tmp_grid = sort([nodes(hermite_class.squared_integrator); nodes(hermite_class.integrators[1])])
	# rewrite below
	A_mat_prior = [hermite_fun(x, qidx)  for x in tmp_grid, qidx=0:qmax]
	nonneg_constraint = @constraint(model, A_mat_prior*tmp_vars .>= 0)

	integrator2 = hermite_class.squared_integrator
	b = hermite_class.sobolev_order
	sobolev_bound = sqrt(hermite_class.sobolev_bound*2π)

	ft_mat = [hermite_fun_noexp(x, qidx)*hermite_cf_coefficient(qidx)  for x in integrator2.nodes, qidx=0:qmax]
	real_ft = (real(ft_mat) * tmp_vars) .* sqrt.(integrator2.weights .* (integrator2.nodes.^2 .+ 1).^b)
	imag_ft = imag(ft_mat) * tmp_vars .* sqrt.(integrator2.weights .* (integrator2.nodes.^2 .+ 1).^b)

	@constraint(model, [sobolev_bound; real_ft; imag_ft] in SecondOrderCone())

	#@constraint(model, sum( real_ft.^2 .* integrator2.weights) + sum( imag_ft.^2 .* integrator2.weights) <= sobolev_bound)

    tmp_vars
end

function marginalize(hermite_class::HermitePriorClass,
                     Zs_discr::DiscretizedStandardNormalSamples,
                     param_vec) #-> A*param_vec
	qdtrs =  hermite_class.integrators
	Zs_discr_length = length(Zs_discr.mhist.grid) + 1
	A_mat = [ MarginalDensityTarget(Zs_discr(i))(qdtr) for i=1:Zs_discr_length, qdtr in qdtrs]
    A_mat*param_vec
end

function linear_functional(hermite_class::HermitePriorClass,
                           target::PriorDensityTarget, #let's hardcode to this for now.
                           param_vec)
	x = location(target)
	qmax = hermite_class.qmax
    v = hermite_fun.(x, 0:qmax)
    dot(v, param_vec)
end

struct HermiteQuasiDistribution<: Distribution{Univariate, Continuous}
	q::Integer
	coefs::Vector{Float64}
end

function pdf(herm_dbn::HermiteQuasiDistribution, x)
	 dot(hermite_fun.(x, 0:herm_dbn.q), herm_dbn.coefs)
end

function cf(herm_dbn::HermiteQuasiDistribution, t)
	 cf_coefs = hermite_cf_coefficient.(0:herm_dbn.q)
	 dot(hermite_fun.(t, 0:herm_dbn.q), herm_dbn.coefs .* cf_coefs)
end

function marginalize(hermite_dbn::HermiteQuasiDistribution,
                     Zs_discr::DiscretizedStandardNormalSamples)
	#temporary hack
   hm_prior_class = HermitePriorClass(qmax=hermite_dbn.q,
                    sobolev_order =2 , sobolev_bound = 10.0)
   ws = marginalize(hm_prior_class, Zs_discr, hermite_dbn.coefs)
   mhist = Zs_discr.mhist
   mhist = @set mhist.hist.weights = ws
end

function (hermite_class::HermitePriorClass)(param_vec)
    q = hermite_class.qmax
  	HermiteQuasiDistribution(q, param_vec)
end

# temporary hacks
struct GaussianHermite
	f
end

function marginalize(hermite_dbn::HermiteQuasiDistribution, Z::StandardNormalSample)
  αs = hermite_dbn.coefs
  qmax = hermite_dbn.q
  GaussianHermite(x -> sum([αs[i+1]*_tmp_conv(i,x) for i in 0:qmax]))
end

function _tmp_conv(j, x)
	j = big(j)
	big(1/sqrt(2*sqrt(pi)*2^j*factorial(j))*exp(-x^2/4)*x^j)
end
pdf(g::GaussianHermite, x) = Float64(g.f(x))
broadcastable(g::GaussianHermite) = Ref(g)


module PolynomialProbabilities
using Distributions

import Base.getindex, Base.append!, Base.isequal, Base.==, Base.length

export PolynomialProbability, PolynomialProbabilityPiece, append!, ==, getindex, length

type PolynomialProbabilityPiece
    t0::Float64
    t1::Float64
    polys::Array{Float64, 2}
    exp::Vector{Float64}

    # Setting all attributes
    PolynomialProbabilityPiece(t0::Float64, t1::Float64, polys::Array{Float64, 2}, exp::Vector{Float64}) = new(t0, t1, polys, exp)
    # Uniform distribution
    PolynomialProbabilityPiece(t0::Float64, t1::Float64, constant::Float64) = new(t0, t1, eye(1) * constant, [0])
end

function ==(pp1::PolynomialProbabilityPiece, pp2::PolynomialProbabilityPiece)
    return pp1.t0 == pp2.t0 && pp1.t1 == pp2.t1 && pp1.polys == pp2.polys && pp1.exp == pp2.exp
end

type PolynomialProbability <: ContinuousUnivariateDistribution
    pieces::Vector{PolynomialProbabilityPiece}

    # Empty distribution
    PolynomialProbability() = new([])
    # Full domain simple distribution
    PolynomialProbability(polys::Array{Float64, 2}, exp::Vector{Float64}) = new([PolynomialProbabilityPiece(-Inf, Inf, polys, exp)])
end

length(pp::PolynomialProbability) = length(pp.pieces)
getindex(pp::PolynomialProbability, ii::Int64) = pp.pieces[ii]

function ==(pp1::PolynomialProbability, pp2::PolynomialProbability)
    if length(pp1) != length(pp2)
        return false
    end

    for ii = 1:length(pp1)
        if pp1[ii] != pp2[ii]
            return false
        end
    end

    return true
end

function append!(pp::PolynomialProbability, t0::Float64, t1::Float64, polys::Array{Float64, 2}, exp::Vector{Float64})
    piece = PolynomialProbabilityPiece(t0, t1, polys, exp)
    push!(pp.pieces, piece)

    pp
end

function append!(pp::PolynomialProbability, t0::Float64, t1::Float64, constant::Float64)
    piece = PolynomialProbabilityPiece(t0, t1, constant)
    push!(pp.pieces, piece)

    pp
end

include("ppsimple.jl")

end # module


# Define some new primitives: conv4 and pool

"""

`conv4(w,x;kwargs...)` executes convolutions or cross-correlations
using filters specified with `w` over tensor `x`.  Currently 4 or 5
dimensional KnetArrays with Float32 or Float64 entries are supported.

If `w` has dimensions (W1,W2,...,I,O) and `x` has dimensions
(X1,X2,...,I,N), the result y will have dimensions (Y1,Y2,...,O,N)
where

    Yi=1+floor(Xi-Wi+2*padding[i]/stride[i])

Here I is the number of input channels, O is the number of output
channels, N is the number of instances, and Wi,Xi,Yi are spatial
dimensions.  Padding and stride are keyword arguments that can be
specified as a single number (in which case they apply to all
dimensions), or an array/tuple with entries for each spatial
dimension.

Here is a description of all available keyword arguments:

* padding: the number of extra zeros implicitly concatenated at the start and at the end of each dimension. Default=0.
* stride: the number of elements to slide to reach the next filtering window. Default=1.
* upscale: upscale factor for each dimension. Default=1.
* mode: 0 for convolution and 1 for cross-correlation.  Default=0.
* alpha: can be used to scale the result. Default=1.
* algo: specifies which convolution algorithm shoud be used to compute the results. Default=0. See the CUDNN User Guide for details.
* workSpace: data pointer to GPU memory to a workspace needed to able to execute the specified algorithm. Default=C_NULL.
* workSpaceSizeInBytes: the size in bytes of the provided workSpace. Default=0.
* handle: handle to a previously created cuDNN context. Default=Knet allocated context.

"""
function conv4{T}(w::KnetArray{T},x::KnetArray{T};
                  handle=cudnnhandle, alpha=one(T), beta=zero(T),
                  algo=0, workSpace=C_NULL, workSpaceSizeInBytes=0, o...)
    y = similar(x, cdims(w,x;o...))
    @cuda(cudnn, cudnnConvolutionForward,
          (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,   UInt32,Cptr,     Csize_t,             Ptr{T},Cptr,Ptr{T}),
          handle,Ref(alpha),TD(x),x,FD(w),w,CD(w,x;o...),algo,workSpace,workSpaceSizeInBytes,Ref(beta),TD(y),y)
    return y
end

function conv4x{T}(w::KnetArray{T},x::KnetArray{T},dy::KnetArray{T};
                   handle=cudnnhandle, alpha=one(T), beta=zero(T),
                   algo=0, workSpace=C_NULL, workSpaceSizeInBytes=0, o...)
    dx = similar(x)
    @cuda(cudnn,cudnnConvolutionBackwardData,
          (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,     UInt32,Cptr,     Csize_t,             Ptr{T},Cptr,Ptr{T}),
          handle,Ref(alpha),FD(w),w,TD(dy),dy,CD(w,x;o...),algo,workSpace,workSpaceSizeInBytes,Ref(beta),TD(dx),dx)
    return dx
end

function conv4w{T}(w::KnetArray{T},x::KnetArray{T},dy::KnetArray{T};
                   handle=cudnnhandle, alpha=one(T), beta=zero(T),
                   algo=0, workSpace=C_NULL, workSpaceSizeInBytes=0, o...)
    dw = similar(w)
    @cuda(cudnn,cudnnConvolutionBackwardFilter,
          (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,     UInt32,Cptr,     Csize_t,             Ptr{T},Cptr,Ptr{T}),
          handle,Ref(alpha),TD(x),x,TD(dy),dy,CD(w,x;o...),algo,workSpace,workSpaceSizeInBytes,Ref(beta),FD(dw),dw)
    return dw
end

@primitive conv4(w,x; o...),dy  conv4w(w,x,dy;o...)  conv4x(w,x,dy;o...)
@zerograd  conv4x(w,x,dy;o...)
@zerograd  conv4w(w,x,dy;o...)

"""

`pool(x;kwargs...)` computes pooling of input values (i.e., the
maximum or average of several adjacent values) to produce an output
with smaller height and/or width.  Currently 4 or 5 dimensional
KnetArrays with Float32 or Float64 entries are supported.

If `x` has dimensions (X1,X2,...,I,N), the result y will have
dimensions (Y1,Y2,...,I,N) where

   Yi=1+ceil(Xi-window[i]+2*padding[i])/stride[i]

Here I is the number of input channels, N is the number of instances,
and Xi,Yi are spatial dimensions.  Window, padding and stride are
keyword arguments that can be specified as a single number (in which
case they apply to all dimensions), or an array/tuple with entries for
each spatial dimension.

Here is a description of all available keyword arguments:

* window: the pooling window size for each dimension. Default=2.
* padding: the number of extra zeros implicitly concatenated at the start and at the end of each dimension. Default=0.
* stride: the number of elements to slide to reach the next pooling window. Default=same as window.
* mode: 0 for max, 1 for average including padded values, 2 for average excluding padded values.  Default=0.
* maxpoolingNanOpt: Nan numbers are not propagated if 0, they are propagated if 1. Default=0.
* alpha: can be used to scale the result. Default=1.
* handle: Handle to a previously created cuDNN context. Default=Knet allocated context.

"""
function pool{T}(x::KnetArray{T}; handle=cudnnhandle, alpha=one(T), beta=zero(T), o...)
    y = similar(x, pdims(x; o...))
    @cuda(cudnn, cudnnPoolingForward,
          (Cptr, Cptr,      Ptr{T},    Cptr,Ptr{T},Ptr{T},   Cptr,Ptr{T}),
          handle,PD(x;o...),Ref(alpha),TD(x),x,    Ref(beta),TD(y),y)
    return y
end

function poolx{T}(x::KnetArray{T},y::KnetArray{T},dy::KnetArray{T};
                  handle=cudnnhandle, alpha=one(T), beta=zero(T), o...)
    dx = similar(x)
    @cuda(cudnn,cudnnPoolingBackward,
          (Cptr,Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Ptr{T},Cptr,Ptr{T}),
          handle,PD(x;o...),Ref(alpha),TD(y),y,TD(dy),dy,TD(x),x,Ref(beta),TD(dx),dx)
    return dx
end

@primitive pool(x;o...),dy,y  poolx(x,y,dy;o...)
@zerograd  poolx(x,y,dy;o...)

# cudnn descriptors

type TD; ptr
    function TD(a::KnetArray)
        d = Cptr[0]
        @cuda(cudnn,cudnnCreateTensorDescriptor,(Ptr{Cptr},),d)
        n = ndims(a)
        sz = [Cint(size(a,n-i+1)) for i=1:n]
        st = [Cint(stride(a,n-i+1)) for i=1:n]
        @cuda(cudnn,cudnnSetTensorNdDescriptor,
              (Cptr,UInt32,Cint,Ptr{Cint},Ptr{Cint}),
              d[1], DT(a), n, sz, st)
        td = new(d[1])
        finalizer(td, x->@cuda(cudnn,cudnnDestroyTensorDescriptor,(Cptr,),x.ptr))
        return td
    end
end

type FD; ptr
    function FD(a::KnetArray)
        d = Cptr[0]
        @cuda(cudnn,cudnnCreateFilterDescriptor,(Ptr{Cptr},),d)
        n = ndims(a)
        sz = [Cint(size(a,n-i+1)) for i=1:n]
        @cuda(cudnn,cudnnSetFilterNdDescriptor,
              (Cptr,UInt32,UInt32,Cint,Ptr{Cint}),
              d[1], DT(a), 0,     n,   sz)
        fd = new(d[1])
        finalizer(fd, x->@cuda(cudnn,cudnnDestroyFilterDescriptor,(Cptr,),x.ptr))
        return fd
    end
end

type CD; ptr
    function CD(w::KnetArray,x::KnetArray; padding=0, stride=1, upscale=1, mode=0)
        d = Cptr[0]
        @cuda(cudnn,cudnnCreateConvolutionDescriptor,(Ptr{Cptr},),d)
        nd = ndims(x)-2
        @cuda(cudnn,cudnnSetConvolutionNdDescriptor,
              (Cptr,Cint,Ptr{Cint},Ptr{Cint},Ptr{Cint},UInt32,UInt32),
              d[1],nd,cdsize(padding,nd),cdsize(stride,nd),cdsize(upscale,nd),mode,DT(x))
        cd = new(d[1])
        finalizer(cd, x->@cuda(cudnn,cudnnDestroyConvolutionDescriptor,(Cptr,),x.ptr))
        return cd
    end
end

type PD; ptr
    function PD(x::KnetArray; window=2, padding=0, stride=window, mode=0, maxpoolingNanOpt=0)
        d = Cptr[0]
        @cuda(cudnn,cudnnCreatePoolingDescriptor,(Ptr{Cptr},),d)
        nd = ndims(x)-2
        @cuda(cudnn,cudnnSetPoolingNdDescriptor,
              (Cptr,UInt32,UInt32,Cint,Ptr{Cint},Ptr{Cint},Ptr{Cint}),
              d[1],mode,maxpoolingNanOpt,nd,cdsize(window,nd),cdsize(padding,nd),cdsize(stride,nd))
        pd = new(d[1])
        finalizer(pd, x->@cuda(cudnn,cudnnDestroyPoolingDescriptor,(Cptr,),x.ptr))
        return pd
    end
end

import Base: unsafe_convert
unsafe_convert(::Type{Cptr}, td::TD)=td.ptr
unsafe_convert(::Type{Cptr}, fd::FD)=fd.ptr
unsafe_convert(::Type{Cptr}, cd::CD)=cd.ptr
unsafe_convert(::Type{Cptr}, pd::PD)=pd.ptr

function cdsize(w, nd)
    if isa(w,Integer)
        fill(Cint(w),nd)
    elseif length(w)==nd
        [ Cint(w[nd-i+1]) for i=1:nd ]
    else
        throw(DimensionMismatch())
    end
end

DT(::KnetArray{Float32})=UInt32(0)
DT(::KnetArray{Float64})=UInt32(1)
DT(::KnetArray{Float16})=UInt32(2)

function cdims{T,N}(w::KnetArray{T,N},x::KnetArray{T,N}; padding=0, stride=1, o...)
    ntuple(N) do i
        if i < N-1
            pi = (if isa(padding,Number); padding; else padding[i]; end)
            si = (if isa(stride,Number); stride; else stride[i]; end)
            1 + div(size(x,i) - size(w,i) + 2*pi, si)
        elseif i == N-1
            size(w,N)
        else # i == N
            size(x,N)
        end
    end
end

function pdims{T,N}(x::KnetArray{T,N}; window=2, padding=0, stride=window, o...)
    ntuple(N) do i
        if i < N-1
            wi = (if isa(window,Number); window; else window[i]; end)
            pi = (if isa(padding,Number); padding; else padding[i]; end)
            si = (if isa(stride,Number); stride; else stride[i]; end)
            1 + ceil(Int, (size(x,i) + 2*pi - wi) / si)
        else
            size(x,i)
        end
    end
end

"""

mat(x) reshapes x into a two-dimensional matrix.  For 1-D inputs mat
returns `reshape(x, (length(x),1))`.  For inputs with more than two
dimensions of size (X1,X2,...,XD), mat returns

    reshape(x, (X1*X2*...*X[D-1],XD))

"""
function mat(x)
    if ndims(x) > 2
        xn = size(x,ndims(x))
        reshape(x, (div(length(x),xn),xn))
    elseif ndims(x)==2
        x
    elseif ndims(x)==1
        reshape(x, (length(x),1))
    else
        throw(MethodError(mat,x))
    end
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

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

                     @variable(m, -2.3024137037198598 <= q <= 2.756640266423038)

                     add_NL_constraint(m, :(swish(0.8745011162287222 + -0.4747792679508138*swish(0.20979439694353186 + 0.2523963619940135*$(x[1]) + 0.6039440929888729*$(x[2])) + 0.06822380173369114*swish(-0.192203947089018 + 0.35618645260012016*$(x[1]) + -0.9380253880003773*$(x[2]))) + swish(-0.5284568968024108 + 0.8782860084347535*swish(0.20979439694353186 + 0.2523963619940135*$(x[1]) + 0.6039440929888729*$(x[2])) + 0.990986450920027*swish(-0.192203947089018 + 0.35618645260012016*$(x[1]) + -0.9380253880003773*$(x[2]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using OrdinaryDiffEq, DiffEqCallbacks, Test

function some_dynamics(u,p,t)
    du = zeros(length(u))
    for i in 1:length(u)
        for j in 1:length(u)
            du[i] += p[i,j]*u[j]*u[i]
        end
    end
    return du
end

p = rand(4,4)
startp = copy(p)

u0 = zeros(length(p[1,:])) .+ 0.1
tspan = (0.,1.)
prob = ODEProblem(some_dynamics,u0,tspan,p)
cb = PresetTimeCallback(0.5,integrator -> integrator.p .= rand(4,4))
sol = solve(prob,Tsit5(),callback = cb)
@test 0.5 ∈ sol.t
@test p != startp

p = rand(4,4)
startp = copy(p)

prob = ODEProblem(some_dynamics,u0,tspan,p)
cb = PresetTimeCallback([0.3,0.6],integrator -> integrator.p .= rand(4,4))
sol = solve(prob,Tsit5(),callback = cb)
@test 0.3 ∈ sol.t
@test 0.6 ∈ sol.t
@test p != startp


# Copyright (c) 2019 Arpit Bhatia and contributors                               #src
#                                                                                #src
# Permission is hereby granted, free of charge, to any person obtaining a copy   #src
# of this software and associated documentation files (the "Software"), to deal  #src
# in the Software without restriction, including without limitation the rights   #src
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell      #src
# copies of the Software, and to permit persons to whom the Software is          #src
# furnished to do so, subject to the following conditions:                       #src
#                                                                                #src
# The above copyright notice and this permission notice shall be included in all #src
# copies or substantial portions of the Software.                                #src
#                                                                                #src
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR     #src
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,       #src
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE    #src
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER         #src
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,  #src
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE  #src
# SOFTWARE.                                                                      #src

# # Portfolio Optimization

# **Originally Contributed by**: Arpit Bhatia

# Optimization models play an increasingly important role in financial
# decisions. Many computational finance problems can be solved efficiently using
# modern optimization techniques.
#
# This tutorial solves the famous Markowitz Portfolio Optimization problem with
# data from [lecture notes from a course taught at Georgia Tech by Shabir Ahmed](https://www2.isye.gatech.edu/~sahmed/isye6669/).

# This tutorial uses the following packages
using JuMP
import Ipopt
import Statistics

# Suppose we are considering investing 1000 dollars in three non-dividend paying
# stocks, IBM (IBM), Walmart (WMT), and Southern Electric (SEHI), for a
# one-month period.

# This means we will use the money to buy shares of the three stocks at the
# current market prices, hold these for one month, and sell the shares off at
# the prevailing market prices at the end of the month.
#
# As a rational investor, we hope to make some profit out of this endeavor,
# i.e., the return on our investment should be positive.
#
# Suppose we bought a stock at $p$ dollars per share in the beginning of the
# month, and sold it off at $s$ dollars per share at the end of the month. Then
# the one-month return on a share of the stock is $ \frac{s-p}{p} $.

# Since the stock prices are quite uncertain, so is the end-of-month return on
# our investment. Our goal is to invest in such a way that the expected
# end-of-month return is at least \$50 or 5%. Furthermore, we want to make
# sure that the “risk” of not achieving our desired return is minimum.

# Note that we are solving the problem under the following assumptions:
# 1. We can trade any continuum of shares.
# 2. No short-selling is allowed.
# 3. There are no transaction costs.

# We model this problem by taking decision variables $x_{i}, i=1,2,3,$ denoting
# the dollars invested in each of the 3 stocks.

# Let us denote by $\tilde{r}_{i}$ the random variable corresponding to the
# monthly return (increase in the stock price) per dollar for stock $i$.

# Then, the return (or profit) on $x_{i}$ dollars invested in stock $i$ is
# $\tilde{r}_{i} x_{i},$ and the total (random) return on our investment is
# $\sum_{i=1}^{3} \tilde{r}_{i} x_{i}.$ The expected return on our investment is
# then $\mathbb{E}\left[\sum_{i=1}^{3} \tilde{r}_{i} x_{i}\right]=\sum_{i=1}^{3} \overline{r}_{i} x_{i},$
# where $\overline{r}_{i}$ is the expected value of the $\tilde{r}_{i}.$

# Now we need to quantify the notion of “risk” in our investment.

# Markowitz, in his Nobel prize winning work, showed that a rational investor’s
# notion of minimizing risk can be closely approximated by minimizing the
# variance of the return of the investment portfolio. This variance is given by:

# ```math
# \operatorname{Var}\left[\sum_{i=1}^{3} \tilde{r}_{i} x_{i}\right] = \sum_{i=1}^{3} \sum_{j=1}^{3} x_{i} x_{j} \sigma_{i j}
# ```

# where $\sigma_{i j}$ is the covariance of the return of stock $i$ with stock $j$.

# Note that the right hand side of the equation is the most reduced form of the
# expression and we have not shown the intermediate steps involved in getting to
# this form. We can also write this equation as:

# ```math
# \operatorname{Var}\left[\sum_{i=1}^{3} \tilde{r}_{i} x_{i}\right] =x^{T} Q x
# ```

# Where $Q$ is the covariance matrix for the random vector $\tilde{r}$.

# Finally, we can write the model as:

# ```math
# \begin{aligned}
# \min x^{T} Q x \\
# \text { s.t. } \sum_{i=1}^{3} x_{i} \leq 1000.00 \\
# \overline{r}^{T} x \geq 50.00 \\
# x \geq 0
# \end{aligned}
# ```

# After that long discussion, lets now use JuMP to solve the portfolio
# optimization problem for the data given below.

# | Month        |  IBM     |  WMT    |  SEHI  |
# |--------------|----------|---------|--------|
# | November-00  |  93.043  |  51.826 |  1.063 |
# | December-00  |  84.585  |  52.823 |  0.938 |
# | January-01   |  111.453 |  56.477 |  1.000 |
# | February-01  |  99.525  |  49.805 |  0.938 |
# | March-01     |  95.819  |  50.287 |  1.438 |
# | April-01     |  114.708 |  51.521 |  1.700 |
# | May-01       |  111.515 |  51.531 |  2.540 |
# | June-01      |  113.211 |  48.664 |  2.390 |
# | July-01      |  104.942 |  55.744 |  3.120 |
# | August-01    |  99.827  |  47.916 |  2.980 |
# | September-01 |  91.607  |  49.438 |  1.900 |
# | October-01   |  107.937 |  51.336 |  1.750 |
# | November-01  |  115.590 |  55.081 |  1.800 |

stock_data = [
    93.043 51.826 1.063;
    84.585 52.823 0.938;
    111.453 56.477 1.000;
    99.525 49.805 0.938;
    95.819 50.287 1.438;
    114.708 51.521 1.700;
    111.515 51.531 2.540;
    113.211 48.664 2.390;
    104.942 55.744 3.120;
    99.827 47.916 2.980;
    91.607 49.438 1.900;
    107.937 51.336 1.750;
    115.590 55.081 1.800;
]

# Calculating stock returns

stock_returns = Array{Float64}(undef, 12, 3)
for i in 1:12
    stock_returns[i, :] = (stock_data[i + 1, :] .- stock_data[i, :]) ./ stock_data[i, :]
end
stock_returns

# Calculating the expected value of monthly return:

r = Statistics.mean(stock_returns, dims = 1)

# Calculating the covariance matrix Q

Q = Statistics.cov(stock_returns)

# JuMP Model

portfolio = Model(Ipopt.Optimizer)
set_silent(portfolio)
@variable(portfolio, x[1:3] >= 0)
@objective(portfolio, Min, x' * Q * x)
@constraint(portfolio, sum(x) <= 1000)
@constraint(portfolio, sum(r[i] * x[i] for i = 1:3) >= 50)
optimize!(portfolio)

objective_value(portfolio)

#-

value.(x)


"""
    _bisection(f, left::T, right::T; max_iters=500, tol=1e-14) -> Union{Nothing, T}

Finds the root `x, f(x) ≈ 0` in interval `[left, right]`, returns nothing if convergence fails.
Errors if the interval is non-bracketing (sign of `f(left)` identical to sign of `f(right)`).
"""
function _bisection(f, left, right; max_iters=500, tol=1e-14)
    for _ in 1:max_iters
        f_left, f_right = f(left), f(right)
        signbit(f_left) == signbit(f_right) && error("Interval became non-bracketing.")

        # Terminate if interval length ~ eps()
        mid = (left + right) / 2
        if left == mid || right == mid
            return mid
        end
        # Terminate if within tol of 0; otherwise, bisect
        f_mid = f(mid)
        if abs(f_mid) < tol
            return mid
        end
        if signbit(f_mid) == signbit(f_left)
            left = mid
        elseif signbit(f_mid) == signbit(f_right)
            right = mid
        end
    end
    return nothing
end


function parse_line_from_stdin()
    map(x->parse(Int64,x), convert(Array{String,1},split(readline(STDIN), " ")))
end

N,num_query = parse_line_from_stdin()

Arrs = Array{Array{Int64,1},1}(N)
for i in 1:N
    Arrs[i] = Array{Int64,1}(0)
end

lastAns = 0
for i in 1:num_query
    query_type,x,y = parse_line_from_stdin()

    if query_type == 0
        ind = mod(x$lastAns,N)+1
        push!(Arrs[ind],y)
    elseif query_type == 1
        ind = mod(x$lastAns,N)+1
        seq = Arrs[ind]
        size = length(seq)
        val = seq[y%size+1]
        println(val)
        lastAns = val
    else
        error("unkown query type")
    end
end


using LockPatternComplexity
using Plots
using Random

function plot_line_types(grid)
    (n, _) = grid.n
    cmap = color_map(grid)
    t = LineTypes(grid)
    lines = (NTuple{2,Vector{Int}}[], NTuple{2,Vector{Int}}[])
    colors = (RGB[], RGB[])
    lines2 = (NTuple{2,Vector{Int}}[], NTuple{2,Vector{Int}}[])
    colors2 = (RGB[], RGB[])
    for (x, y, k) in zip(t.x, t.y, t.k)
        if y ≥ 0
            push!(lines[1], ([1, 1 + x], [1, 1 + y]))
            push!(colors[1], cmap[(x, y)])
            if k > 1
                push!(lines2[1], ([1, 1 + k * x], [1, 1 + k * y]))
                push!(colors2[1], cmap[(x, y)])
            end
        else
            push!(lines[2], ([1, 1 + x], [n, n + y]))
            push!(colors[2], cmap[(x, y)])
            if k > 1
                push!(lines2[2], ([1, 1 + k * x], [n, n + k * y]))
                push!(colors2[2], cmap[(x, y)])
            end
        end
    end
    plt = (
        canvas(; size = (200, 200), backgroud = :lightgray),
        canvas(; size = (200, 200), backgroud = :lightgray)
    )
    plot_lines!(plt[1], lines[1], colors[1]; linewidth = 3)
    plot_lines!(plt[2], lines[2], colors[2]; linewidth = 3)
    plot_lines!(plt[1], lines2[1], colors2[1]; linewidth = 3, linestyle = :dash)
    plot_lines!(plt[2], lines2[2], colors2[2]; linewidth = 3, linestyle = :dash)
    plot_grid!(plt[1], grid; marker = :circle, color = :black, markersize = 5)
    plot_grid!(plt[2], grid; marker = :circle, color = :black, markersize = 5)
    return plt
end

grid = Grid(5)
pattern = (shuffle(1:length(grid)), shuffle(1:length(grid)))

begin
    plt = canvas(; size = (200, 200), backgroud = :lightgray)
    plot_grid!(plt, grid; marker = :circle, color = :black, markersize = 5)
    savefig(plt, "docs/src/plots/grid.svg")
end

begin
    plt = plot_line_types(grid)
    savefig(plt[1], "docs/src/plots/line_types_1.svg")
    savefig(plt[2], "docs/src/plots/line_types_2.svg")
end

for i = 1:2
    plt = canvas(size = (200, 200), backgroud = :lightgray)
    plot_lock_pattern!(plt, grid, pattern[i])
    savefig(plt, "docs/src/plots/lock_pattern_$i.svg")
end

# TODO: symmetries


@testset "Convex Hull" begin
    A = zeros(50, 30)
    A = convert(Array{Bool}, A)
    A[25, 1] = 1
    A[1, 10] = 1
    A[10, 10] = 1
    A[10, 30] = 1
    A[40, 30] = 1
    A[40, 10] = 1
    A[50, 10] = 1
    B = @inferred convexhull(A)
    C = CartesianIndex{}[]
    push!(C, CartesianIndex{}(25, 1))
    push!(C, CartesianIndex{}(1, 10))
    push!(C, CartesianIndex{}(10, 30))
    push!(C, CartesianIndex{}(40, 30))
    push!(C, CartesianIndex{}(50, 10))
    @test typeof(B) == Array{CartesianIndex{2},1}
    @test sort(B) == sort(C)

    A = [
        0.0,
        0.0,
        1.0,
        0.0,
        0.0,
        0.0,
        1.0,
        1.0,
        0.0,
        0.0,
        1.0,
        0.0,
        0.0,
        1.0,
        1.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        1.0,
        0.0,
        0.0,
    ]
    A = reshape(A, 5, 5)
    A = convert(Array{Bool}, A)
    B = B = @inferred convexhull(A)
    C = CartesianIndex{}[]
    push!(C, CartesianIndex{}(1, 3))
    push!(C, CartesianIndex{}(3, 1))
    push!(C, CartesianIndex{}(3, 5))
    push!(C, CartesianIndex{}(5, 3))
    @test typeof(B) == Array{CartesianIndex{2},1}
    @test sort(B) == sort(C)
end