# Signal to Revise.jl that this file should be tracked as a package
#__revise_mode__ = :eval

# Types
abstract type AbstractGrid end

struct Grid{N, T} <: AbstractGrid
    r₀::SVector{N, T}
    l::SVector{N, T}
    Δl::SVector{N, T}
    size::SVector{N, Int}
    periodicity::SVector{N, Bool}

    function Grid(lower, upper, size, periodicity)
        return Grid(lower, upper, SVector(size), periodicity)
    end

    function Grid(lower, upper, size::SVector{N, <: Integer}, periodicity) where {N}
        l = SVector(upper .- lower)
        Δ = l ./ size
        T = eltype(Δ)
        return new{N, T}(SVector(lower), l, Δ, size, SVector(periodicity))
    end
end

function Base.show(io::IO, g::Grid{N, T}) where {N, T}
    context = IOContext(io, :compact => true)
    join(io, g.size, '×')
    print(io, "-element Grid{", N, ", ", T, "}:\n ")
    itr = Iterators.Stateful(zip(g.r₀, g.l))
    for (lo, l) in itr
        print(context, '[', lo, " .. ", lo + l, ']')
        if Base.peek(itr) !== nothing
            print(io, " × ")
        end
    end
    return nothing
end

function circ_index(r, r₀, l, Δl, m)
    n = floor(Int, mod(r - r₀, l) / Δl) + 1
    return ifelse(n > m, 1, n)
end

function Base.getindex(g::Grid{N}, r::Union{Real, NTuple{N, Real}, SVector{N}}) where {N}
    return CartesianIndex(Tuple(circ_index.(r, g.r₀, g.l, g.Δl, g.size)))
end

Base.size(g::Grid) = Tuple(g.size)
Base.ndims(::Grid{N}) where {N} = N


module Visualize

using GeometryBasics
using FileIO
using Plots

export Mesh3D, create_ply_file, show_closed_cross_section, show_multi_branch_cross_section

function Mesh3D(CrossSectionNodes, z, u, v, ϕ)

    #define planar cross-section geometry
    x=view(CrossSectionNodes,:,1)
    y=view(CrossSectionNodes,:,2)

    # write deformed cross-section in 3D coordinates, section by section along the z axis
    for i in 1:length(z)
        CrossSectionShape=hcat(x.+u[i]+x.*cos.(ϕ[i]).-y.*sin.(ϕ[i]),y.+v[i]+x.*sin.(ϕ[i]).+y.*cos.(ϕ[i]),ones(length(x))*z[i])
        if i==1
            global coordinates=CrossSectionShape   #start with first cross-section
        else
            coordinates=vcat(coordinates, CrossSectionShape)  #add following cross-sections along z
        end
    end

    # apply triangular mesh to 3D coordinates
    NumCrossSectionNodes=length(x)
    NumZNodes=length(z)

    #uppoly are triangles pointing towards +z, downpoly are triangles pointing towards -z
    uppoly=[[i i+1 i+1+NumCrossSectionNodes].+NumCrossSectionNodes*(j-1) for i=1:NumCrossSectionNodes-1 for j=1:NumZNodes-1]
    downpoly=[[i i+NumCrossSectionNodes i+NumCrossSectionNodes+1].+NumCrossSectionNodes*(j-1) for i=1:NumCrossSectionNodes-1 for j=1:NumZNodes-1]

    connectivity=vcat(uppoly,downpoly)  #combine all the triangles
    connectivity=permutedims(reshape(hcat(connectivity...), (length(connectivity[1]), length(connectivity))))  #convert into format to be read by Makie

    return coordinates, connectivity   #these are ready for Makie scenes, e.g., poly()

end

"""
    create_ply_file(coordinates, connectivity, filename)

Accepts triangular surface mesh `coordinates` and `connectivity` along with a `filename' and saves a .ply file.

https://en.wikipedia.org/wiki/PLY_(file_format)

The coordinates and connectivity are converted to GeometryBasics.jl primitives before saving with FileIO.jl.

"""

function create_ply_file(coordinates, connectivity, filename)

    points = [Point{3, Float64}(coordinates[i,1], coordinates[i,2], coordinates[i,3]) for i = 1:size(coordinates)[1]]

    facets = [TriangleFace{Cint}(connectivity[i, 1:3]) for i in eachindex(connectivity[:,1])]

    mesh = GeometryBasics.Mesh(points, facets)

    # using FileIO, MeshIO
    FileIO.save(filename, mesh)

end

function show_closed_cross_section(xcoords, ycoords, markershape, xlims, ylims)

    num_nodes = length(xcoords)

    plt = []

    for i = 1:num_nodes

        if i == 1

            plt = plot([xcoords[i], xcoords[i+1]], [ycoords[i], ycoords[i+1]], size = (600, 600), legend = false, linecolor = :black, markercolor = :black, markershape = markershape, seriestype = :line, xlims = xlims, ylims = ylims)

        elseif i<num_nodes

            plot!(plt, [xcoords[i], xcoords[i+1]], [ycoords[i], ycoords[i+1]], size = (600, 600), legend = false, linecolor = :black, markercolor = :black, markershape = markershape, seriestype = :line, xlims = xlims, ylims = ylims)

        elseif i==num_nodes

            plot!(plt, [xcoords[i], xcoords[1]], [ycoords[i], ycoords[1]], size = (600, 600), legend = false, linecolor = :black, markercolor = :black, markershape = markershape, seriestype = :line, xlims = xlims, ylims = ylims)

        end

    end

    return plt

end

function show_multi_branch_cross_section(xcoords, ycoords, element_connectivity, markershape, xlims, ylims)

    num_elem = size(element_connectivity)[1]

    plt = []

    for i = 1:num_elem

        node_i = Int(element_connectivity[i, 1])
        node_j = Int(element_connectivity[i, 2])

        if i == 1

            plt = plot([xcoords[node_i], xcoords[node_j]], [ycoords[node_i], ycoords[node_j]], size = (600, 600), legend = false, linecolor = :black, markercolor = :black, markershape = markershape, seriestype = :line, xlims = xlims, ylims = ylims)

        else

            plot!(plt, [xcoords[node_i], xcoords[node_j]], [ycoords[node_i], ycoords[node_j]], size = (600, 600), legend = false, linecolor = :black, markercolor = :black, markershape = markershape, seriestype = :line, xlims = xlims, ylims = ylims)

        end

    end

    return plt

end

end #module


include("bool_cube.jl")
include("bool_fun.jl")
include("ANF.jl")
include("bool_gen.jl")
include("boolean_utils.jl")


function normalize_factor{R <: Real}(arr::Array{R})
    arr ./ sum(arr)
end

function normpdf(mu, sigma)
    x -> (1/(sigma*sqrt(2*pi)))*e^((-(x-mu)^2)/(2*sigma^2))
end

function group_combinations(groups::Array{Array,1})
    if length(groups) > 0
        elem_type=eltype(groups[1])
        combinations=prod(map(length, groups))
        res=Array(elem_type, combinations, length(groups))

        groups=groups[sortperm(map(length, groups))]

        prev_group_elem_size=combinations

        for i=1:length(groups)

            g=groups[i]
            prev_steps=convert(Int,combinations/prev_group_elem_size)
            group_elem_size=convert(Int,prev_group_elem_size/length(g))

            step_elems=reduce(vcat,map(x -> fill(x,group_elem_size), g))
            steps=reduce(vcat,fill(step_elems,prev_steps))

            res[:,i]=steps
            prev_group_elem_size=group_elem_size
        end

        sets=Array(Set{elem_type},combinations)
        sets[:] = [ Set(res[i,:]) for i=1:combinations ]
    else
        []
    end
end

function single_to_multi_index(i::Integer, dims::(Integer,Integer))
    dim1, dim2 = dims
    convert(Integer,ceil(i/dim1)), i%dim1 > 0 ? i%dim1 : dim1
end

order_perm(A1::Array, A2::Array) = map(x -> findfirst(A1,x), A2)


module SigmaDC

using Mea.Periodize

function dfdw(beta_::Float64, ww::Float64)

    #if abs(beta_*ww) > 20.0
    #    return 0.0
    #end

    return (-beta_*exp(beta_*ww)/(1.0 + exp(beta_*ww))^2.0   )
end

function calc_sigmadc(modelvec::Periodize.ModelVector, beta_::Float64)

  cutoff = 20.0
  cutoffidx=0
  for (ii, ww) in enumerate(modelvec.wvec_)
    if abs(ww*beta_) < cutoff
      cutoffidx = ii
      #println("ii = ", ii)
      break
    end
  end

  #println("cutoffidx = ", cutoffidx)

  modelvec.wvec_ = modelvec.wvec_[cutoffidx:end-cutoffidx]
  modelvec.sEvec_c_ = modelvec.sEvec_c_[cutoffidx:end-cutoffidx,:,:]

  dos2 = Periodize.calcdos2(modelvec, maxevals=100000)

  integrand = zeros(Float64, size(dos2)[1])
  for ii in 1:size(dos2)[1]
      integrand[ii] = dfdw(beta_, dos2[ii, 1])*dos2[ii, 2]
  end

  sigmadc = 0.0

      for ii in 1:size(dos2)[1]-1
                              # y                                 x
          sigmadc += 0.5 * (integrand[ii] + integrand[ii+1]) * (dos2[ii+1, 1] - dos2[ii, 1])
      end
  sigmadc*= -2.0/(2.0*pi)
  println(sigmadc)
  return sigmadc
end

end


# This code demonstrates the use of SolveDSGE for solving linear rational expectations models.
# The demonstration model is a simple forward-looking new Keynesian model with monetary policy
# conducted using a Taylor rule.

# The model is as follows:

# pi_{t} = beta*E_{t}pi_{t+1} + kappa*y_{t} + u_{t}
#  y_{t} = E_{t}y_{t+1} - mu*(R_{t}-E_{t}pi_{t+1} + v_{t}
#  R_{t} = phi_pi*E_{t}pi_{t+1} + phi_y*y_{t}

#-------------------------------------------------------------------------------------------#

# Bring the SolveDSGE module into global scope

using SolveDSGE

# Assign values to model parameters

const beta = 0.99
const mu = 0.5
const phi_pi = 1.5
const rho_u = 0.8
const rho_v = 0.8
const kappa = 0.1
const phi_y = 0.5
const sigma = [1.0 0.0; 0.0 1.0]

const cutoff = 1.0
const tol    = 1e-8

# Solve the model using Blanchard and Kahn (1980)

# Construct the model matrices

b = [1.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0; 0.0 0.0 beta 0.0; 0.0 0.0 mu*(1.0-phi_pi) 1.0]
a = [rho_u 0.0 0.0 0.0; 0.0 rho_v 0.0 0.0; -1.0 0.0 1.0 -kappa; 0.0 -1.0 0.0 (1.0+mu*phi_y)]
c = [1.0 0.0; 0.0 1.0; 0.0 0.0; 0.0 0.0]

c = b\c
a = b\a

nx = 2
ny = 2

# Put the model into Blanchard_Kahn_Form type

m_bk = Blanchard_Kahn_Form(nx,ny,a,c,sigma)

# Do some basic checking to determine whether the model matrices are conformable

assessment = check_model_form(m_bk)
println(assessment)

# Solve the model by solving an eigenvalue problem

soln_bk_1 = solve_re(m_bk,cutoff)
responses_bk_1 = impulses(soln_bk_1,5,1)

# Solve the model by solving an algebraic Riccati equation

soln_bk_2 = solve_re(m_bk,cutoff,tol)
responses_bk_2 = impulses(soln_bk_2,5,1)

# Solve the model using Klein (2000)

# Construct the model matrices

b = [1.0 0.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0 0.0; 0.0 0.0 beta 0.0 0.0; 0.0 0.0 mu 1.0 0.0; 0.0 0.0 phi_pi 0.0 0.0]
a = [rho_u 0.0 0.0 0.0 0.0; 0.0 rho_v 0.0 0.0 0.0; -1.0 0.0 1.0 -kappa 0.0; 0.0 -1.0 0.0 1.0 mu; 0.0 0.0 0.0 -phi_y 1.0]
c = [1.0 0.0; 0.0 1.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]

nx = 2
ny = 3

# Put the model into Klein_Form type

m_k = Klein_Form(nx,ny,a,b,c,sigma)

# Do some basic checking to determine whether the model matrices are conformable

assessment = check_model_form(m_k)
println(assessment)

# Solve the model by solving a generalized eigenvalue problem

soln_k = solve_re(m_k,cutoff)
responses_k = impulses(soln_k,5,1)

# Solve the model using Sims (2001)

# Construct the model matrices

gamma0 = [1.0 0.0 0.0 0.0 0.0 0.0 0.0;
          0.0 1.0 0.0 0.0 0.0 0.0 0.0;
          0.0 0.0 1.0 0.0 0.0 0.0 0.0;
          0.0 0.0 0.0 1.0 0.0 0.0 0.0;
          0.0 0.0 0.0 -phi_y 1.0 -phi_pi 0.0;
          -1.0 0.0 1.0 -kappa 0.0 -beta 0.0;
          0.0 -1.0 0.0 1.0 mu -mu -1.0]

gamma1 = [rho_v 0.0 0.0 0.0 0.0 0.0 0.0;
          0.0 rho_u 0.0 0.0 0.0 0.0 0.0;
          0.0 0.0 0.0 0.0 0.0 1.0 0.0;
          0.0 0.0 0.0 0.0 0.0 0.0 1.0;
          0.0 0.0 0.0 0.0 0.0 0.0 0.0;
          0.0 0.0 0.0 0.0 0.0 0.0 0.0;
          0.0 0.0 0.0 0.0 0.0 0.0 0.0]

c = zeros(7)

psi = [1.0 0.0;
       0.0 1.0;
       0.0 0.0;
       0.0 0.0;
       0.0 0.0;
       0.0 0.0;
       0.0 0.0]

pi = [0.0 0.0;
      0.0 0.0
      1.0 0.0
      0.0 1.0;
      0.0 0.0;
      0.0 0.0;
      0.0 0.0]

# Put the model into Sims_Form type

m_sims = Sims_Form(gamma0,gamma1,c,psi,pi,sigma)

# Do some basic checking to determine whether the model matrices are conformable

assessment = check_model_form(m_sims)
println(assessment)

# Solve the model by solving a generalized eigenvalue problem

soln_sims = solve_re(m_sims,cutoff)

# Solve the model in structural form

# Construct the model matrices

a0 = [1.0 0.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0 0.0; -1.0 0.0 1.0 -kappa 0.0; 0.0 -1.0 0.0 1.0 mu; 0.0 0.0 0.0 -phi_y 1.0]
a1 = [rho_u 0.0 0.0 0.0 0.0; 0.0 rho_v 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0]
a2 = [0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 beta 0.0 0.0; 0.0 0.0 mu 1.0 0.0; 0.0 0.0 phi_pi 0.0 0.0]
a4 = [1.0 0.0; 0.0 1.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]

# Put the model in Structural_Form type

m_bp = Binder_Pesaran_Form(a0,a1,a2,a4,sigma)

# Do some basic checking to determine whether the model matrices are conformable

assessment = check_model_form(m_bp)
println(assessment)

# Solve the by by solving a generalized eigenvalue problem

soln_bp_1 = solve_re(m_bp,cutoff)
responses_bp_1 = impulses(soln_bp_1,5,1)

# Solve the model using Binder and Pesaran's (1995) brute force method

soln_bp_2 = solve_re(m_bp,cutoff,tol)
responses_bp_2 = impulses(soln_bp_2,5,1)


using Omega
using Test
using Statistics: mean

function testrid()
  θ = betarv(2.0, 2.0)
  x = bernoulli(θ)
  ridxθ = Omega.rid(x, θ)
  mean1, mean2 = rand((θ, lmean(ridxθ)))
  @test mean1 == mean2
  ω1 = Omega.rand(defΩ())
  ω2 = Omega.rand(defΩ())
  x_ = ridxθ(ω1)
  θ_ = θ(ω1)
  nsamples = 10000
  θ_approx = mean([x_(Omega.rand(defΩ())) for i = 1:nsamples])
  @test isapprox(θ_approx, θ_; atol = 0.01)
end

testrid()

function testrid2()
  t1 = normal(0.0, 1.0)
  t2 = normal(0.0, 1.0)
  x = normal(t1 + t2, 1.0)
  xrid = rid(x, t1, t2)
  samples = rand(t1, lmean(xrid) ==ₛ 0.0, 1000; alg = SSMH)
end

testrid2()


dt = 0.01
s = [10.0, 28.0, 8/3]
d = 3
function flow(u::Array{Float64,1})
    sigma, rho, beta = s
    x, y, z = u[1], u[2], u[3]
    return [sigma*(y - x), x*(rho - z) - y, x*y - beta*z]
end
function next(u, sn)
    k1 = dt*flow(u)
    k2 = dt*flow(u .+ k1)
    u_next = @. u + (k1 + k2)/2
	noise = sn*randn(d)
	u_next .+= noise
    return u_next
end


"""
    smoothspec(z, nas)

Smoothen field `z(ns,ns)` with a spectral method at scale `ns/nas`
Takes into account missing values.
"""
function smoothspec(zi, nas)

    iinan = findall(isnan.(zi))
    iinotnan = findall(.~isnan.(zi))
    zi[iinan] = 0.

    nss = size(zi)
    ns = nss[1]
    sdim = div(ns, nas) # the smoothing radius is one large scale pixel wide (setting the diameter to 1 pixel is wrong)

    kx = [collect(0:ns/2)' collect(-ns/2+1:-1)']'
    kx = kx*ones(1,ns)
    kx = kx.^2+permutedims(kx, [2 1]).^2

    zif = fft(zi)
    zif[kx[:].>(nas/2).^2] = 0.0
    zif[kx[:].==(nas/2).^2] = real(zif[kx[:].==(nas/2).^2]*2.0)
    zif = real(ifft(zif))

    return zif
end


# an idea: create one method named pivots with kwargs

function floorpivots{T,N}(ohlc::TimeArray{T,N})

    p  = lag(typical(ohlc))
    s1 = 2.*p .- lag(ohlc["High"])
    r1 = 2.*p .- lag(ohlc["Low"])
    s2 = p .- (r1 .- s1)
    r2 = (p .- s1) .+ r1
    s3 = p .- (r2 .- s1)
    r3 = (p .- s1) .+ r2

    TimeArray(s1.timestamp, [s3.values s2.values s1.values p.values r1.values r2.values r3.values],
              ["s3", "s2", "s1", "pivot", "r1", "r2", "r3"], ohlc.meta)

end

function woodiespivots{T,N}(ohlc::TimeArray{T,N})

    rng = lag(ohlc["High"]) .- lag(ohlc["Low"])

    p  = (lag(ohlc["High"]) .+ lag(ohlc["Low"]) .+ 2.*ohlc["Open"]) ./ 4
    s1 = 2.*p .- lag(ohlc["High"])
    r1 = 2.*p .- lag(ohlc["Low"])
    s2 = p .- rng
    r2 = p .+ rng
    s3 = s1 .- rng
    r3 = r1 .+ rng
    # s4 = s3 - range # can't get an answer that matches online calculators
    # r4 = r3 + range

    TimeArray(s1.timestamp, [s3.values s2.values s1.values p.values r1.values r2.values r3.values],
              ["s3", "s2", "s1", "pivot", "r1", "r2", "r3"], ohlc.meta)
end
#
# function camarillapivots{T,V}(hi::Array{SeriesPair{T,V},1},
#                               lo::Array{SeriesPair{T,V},1},
#                               cl::Array{SeriesPair{T,V},1})
#   #code here
#   ##  R4 = (H - L) * 1.1/2 + C
#   ##  R3 = (H - L) * 1.1/4 + C
#   ##  R2 = (H - L) * 1.1/6 + C
#   ##  R1 = (H - L) * 1.1/12 + C
#   ##  S1 = C - (H - L) * 1.1/12
#   ##  S2 = C - (H - L) * 1.1/6
#   ##  S3 = C - (H - L) * 1.1/4
#   ##  S4 = C - (H - L) * 1.1/2
# end
#
# function demarkpivots{T,V}(op::Array{SeriesPair{T,V},1},
#                            hi::Array{SeriesPair{T,V},1},
#                            lo::Array{SeriesPair{T,V},1},
#                            cl::Array{SeriesPair{T,V},1})
#   #code here
#   ## If Close < Open then X = (H + (L * 2) + C)
#   ## If Close > Open then X = ((H * 2) + L + C)
#   ## If Close = Open then X = (H + L + (C * 2))
#   ##
#   ## R1 = X / 2 - L
#   ## PP = X / 4
#   ## S1 = X / 2 - H
# end

function market_profile(x)
  #code here
end


# ---
# title: 1248. Count Number of Nice Subarrays
# id: problem1248
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Two Pointers
# link: <https://leetcode.com/problems/count-number-of-nice-subarrays/description/>
# hidden: true
# ---
#
# Given an array of integers `nums` and an integer `k`. A continuous subarray is
# called **nice** if there are `k` odd numbers on it.
#
# Return _the number of **nice** sub-arrays_.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,1,2,1,1], k = 3
#     Output: 2
#     Explanation: The only sub-arrays with 3 odd numbers are [1,1,2,1] and [1,2,1,1].
#
#
# **Example 2:**
#
#
#
#     Input: nums = [2,4,6], k = 1
#     Output: 0
#     Explanation: There is no odd numbers in the array.
#
#
# **Example 3:**
#
#
#
#     Input: nums = [2,2,2,1,2,2,1,2,2,2], k = 2
#     Output: 16
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 50000`
#   * `1 <= nums[i] <= 10^5`
#   * `1 <= k <= nums.length`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# ## Non-zero inflated, continuous. e.g. LogNormal
# function expert_tn_pos(d::e, tl::Real, yl::Real, yu::Real, tu::Real) where {e<:NonZIContinuousExpert}

#     expert_tn = (tl == tu) ? logpdf.(d, tl) : logcdf.(d, tu) + log1mexp.(logcdf.(d, tl) - logcdf.(d, tu))
#     expert_tn = (tu == 0.) ? -Inf : expert_tn

#     return expert_tn
# end

# ## Zero inflated, continuous. e.g. ZILogNormal
# function expert_tn_pos(d::e, tl::Real, yl::Real, yu::Real, tu::Real) where {e<:ZIContinuousExpert}

#     expert_tn = (tl == tu) ? logpdf.(d, tl) : logcdf.(d, tu) + log1mexp.(logcdf.(d, tl) - logcdf.(d, tu))
#     expert_tn = (tu == 0.) ? -Inf : expert_tn

#     return expert_tn
# end

# ## Non-zero inflated, discrete. e.g. Poisson
# function expert_tn_pos(d::e, tl::Real, yl::Real, yu::Real, tu::Real) where {e<:NonZIDiscreteExpert}

#     expert_tn = (tl == tu) ? logpdf.(d, tl) : logcdf.(d, tu) + log1mexp.(logcdf.(d, ceil.(tl) .- 1) - logcdf.(d, tu))

#     return expert_tn
# end

# ## Zero inflated, discrete. e.g. ZIPoisson
# function expert_tn_pos(d::e, tl::Real, yl::Real, yu::Real, tu::Real) where {e<:ZIDiscreteExpert}

#     expert_tn = (tl == tu) ? logpdf.(d, tl) : logcdf.(d, tu) + log1mexp.(logcdf.(d, ceil.(tl) .- 1) - logcdf.(d, tu))

#     return expert_tn
# end


module layers
using LinearAlgebra
export halfspace,layer_plain,layer_patterned, modes_freespace,Layer,Halfspace

"""
    Layer(thi,V,W,X,q)
    Computed Eigenmodes and propagation properties of a layer
    thi: layer thickness
    V: coordinate transform between eigenmode amplitude and magnetic field
    W: coordinate transform between eigenmode amplitude and electric field
    X: transmission of amplitude vector through the layer
    q: effective propagation constant of the eigenmodes
"""
struct Layer
    thi::Float64
    V::AbstractArray{Complex{Float64},2}
    W::AbstractArray{Complex{Float64},2}
    X::AbstractArray{Complex{Float64},2}
    q::AbstractArray{Complex{Float64},2}
end
"""
Halfspace(V,W,Kz)
Computed Eigenmodes of an unpatterned, infinitely thin halfspace
V: coordinate transform between eigenmode amplitude and magnetic field
W: coordinate transform between eigenmode amplitude and electric field
Kz: z-axis component of the propagation vector
"""
struct Halfspace
    V::AbstractArray{Complex{Float64},2}
    W::AbstractArray{Complex{Float64},2}
    Kz::AbstractArray{Complex{Float64},2}
end
"""
modes_freespace(Kx,Ky)
Computes the eigenmodes of propagation through free space, for normalization
Kx: x-axis component of the propagation vector
Ky: y-axis component of the propagation vector
returns
V0: coordinate transform between free space eigenmode amplitude and magnetic field
Kz0: z-axis component of the propagation vector in free space
"""
function modes_freespace(Kx,Ky)
    #just because |k|=1
    Kz0=sqrt.(Complex.(I-Kx*Kx-Ky*Ky))
    #P0 is identity
    Q0=[Kx*Ky I-Kx*Kx;Ky*Ky-I -Ky*Kx]
    #propagation
    q0=1im*Kz0
    q0=[q0 q0*0;0*q0 q0]
    #Free space, so W is identity
    #W0=I+0*Q0
    V0=Q0/Diagonal(q0)
    return V0,Kz0
end
"""
modes_patterned(Kx,Ky,k0,thi,epsilon)
Computes the eigenmodes of propagation through a patterned layer
Kx: x-axis component of the propagation vector
Ky: y-axis component of the propagation vector
k0: is 2pi/lambda, to normalize the thickness
thi: layer thickness
epsilon: Permittivity of layer in reciprocal space
returns
V: coordinate transform between eigenmode amplitude and magnetic field in the layer
W: coordinate transform between eigenmode amplitude and electric field in the layer
X: transmission of amplitude vector through the layer
q: effective propagation constant of the eigenmodes
"""
function modes_patterned(Kx,Ky,k0,thi,epsilon)
    eta=I/epsilon
    P=[Kx*eta*Ky I-Kx*eta*Kx;Ky*eta*Ky-I -Ky*eta*Kx]
    Q=[Kx*Ky epsilon-Kx*Kx;Ky*Ky-epsilon -Ky*Kx]
    #eigenmodes
    ev=eigen(Matrix(P*Q))
    q=Diagonal(sqrt.(Complex.(ev.values)))
    q[real.(q).>0].*=-1
    #W is transform between amplitude vector and E-Field
    W=ev.vectors
    #V is transform between amplitude vector and H-Field
    V=Q*W/Diagonal(q)

    #X the factor applied to the amplitudes when propagatin through the layer
    X=exp(q*k0*thi)
    return V,W,X,q
end
"""
layer_patterned(Kx,Ky,k0,thi,epsilon)
Computes the eigenmodes of propagation through a patterned layer in a Layer object
Kx: x-axis component of the propagation vector
Ky: y-axis component of the propagation vector
k0: is 2pi/lambda, to normalize the thickness
thi: layer thickness
epsilon: Permittivity of layer in reciprocal space
returns
patterned layer object
"""
function layer_patterned(Kx,Ky,k0,thi,epsilon)
    V,W,X,q=modes_patterned(Kx,Ky,k0,thi,epsilon)
    return Layer(thi,V,W,X,q)
end
"""
modes_plain(Kx,Kyk0,thi,epsilon)
Computes the eigenmodes of propagation through a plain layer
Kx: x-axis component of the propagation vector
Ky: y-axis component of the propagation vector
k0: is 2pi/lambda, to normalize the thickness
thi: layer thickness
epsilon: Bulk permittivity of the layer
returns
V: coordinate transform between eigenmode amplitude and magnetic field in the layer
W: coordinate transform between eigenmode amplitude and electric field in the layer
X: transmission of amplitude vector through the layer
q: effective propagation constant of the eigenmodes
"""
function modes_plain(Kx,Ky,k0,thi,epsilon)
    #Kz=P*Q
    Kz=sqrt.(Complex.(epsilon*I-Kx*Kx-Ky*Ky))
    Q=[Kx*Ky epsilon*I-Kx*Kx;Ky*Ky-epsilon*I -Ky*Kx]
    q=[1im*Kz zeros(size(Kz));zeros(size(Kz)) 1im*Kz]
    q[real.(q).>0].*=-1
    #W is identity
    W=I
    V=Q/Diagonal(q)
    X=exp(Matrix(q*k0*thi))
    return V,W,X,q
end
"""
layer_plain(Kx,Kyk0,thi,epsilon)
Computes the eigenmodes of propagation through a plain layer, in a Layer object
Kx: x-axis component of the propagation vector
Ky: y-axis component of the propagation vector
k0: is 2pi/lambda, to normalize the thickness
thi: layer thickness
epsilon: Bulk permittivity of the layer
returns
layer object
"""
function layer_plain(Kx,Ky,k0,thi,epsilon)
    V,W,X,q=modes_plain(Kx,Ky,k0,thi,epsilon)
    return Layer(thi,V,V*0+W,X,q)
end
"""
modes_halfspace(Kx,Ky,epsilon)
Computes the eigenmodes of an infinitely thin halfspace
Kx: x-axis component of the propagation vector
Ky: y-axis component of the propagation vector
epsilon: bulk permittivity of the halfspace
returns
V: coordinate transform between eigenmode amplitude and magnetic field in the halfspace
W: coordinate transform between eigenmode amplitude and electric field in the halfspace
Kz: z-axis component of the propagation vector in the halfspace
"""
function modes_halfspace(Kx,Ky,epsilon)
    Kz=sqrt.(Complex.(epsilon*I-Kx*Kx-Ky*Ky))
    Q=[Kx*Ky epsilon*I-Kx*Kx;Ky*Ky-epsilon*I -Ky*Kx]
    q0=[1im*Kz zeros(size(Kz));zeros(size(Kz)) 1im*Kz]
    W=I
    V=Q/Diagonal(q0)
    return V,W+0*V,Kz
end
"""
halfspace(Kx,Ky,epsilon)
Computes the eigenmodes of an infinitely thin halfspace in a Halfspace object
Kx: x-axis component of the propagation vector
Ky: y-axis component of the propagation vector
epsilon: bulk permittivity of the halfspace
returns
halfspace object
"""
function halfspace(Kx,Ky,epsilon)
    V,W,Kz=modes_halfspace(Kx,Ky,epsilon)
    return Halfspace(V,W,Kz)
end

end


struct MaxUCB <: Function
    c::Float64
end

function (mu::MaxUCB)(n::StateNode)
    bestucb = -Inf
    best = first(children(n))
    tn = totaln(n)
    for sanode in children(n)
        if c == 0 || tn == 0 || (tn == 1 && n(sanode) == 0)
            ucb = q(sanode)
        else
            ucb = q(sanode) + c*sqrt(log(tn)/n(sanode))
        end
        @assert !isnan(ucb)
        @assert !isequal(ucb, -Inf)
        if ucb > bestucb
            bestucb = ucb
            best = sanode
        end
    end
    return best
end

maxtried(n::StateNode) = reduce(children(n)) do an1, an2
    if n(an1) > n(an2)
        return an1
    else
        return an2
    end
end


using LinearAlgebra, Plots, GLM, DataFrames, Polynomials
Data = DataFrame(j0_NMC =[2.8e-1;3.48e-1;4.53e-1;5.44e-1;6.02e-1;], j0_Graph = [2.95e-2;3.54e-2;5.76e-2;8.32e-2;1.30e-1], R_NMC = [1.8;1.48;1.17;1.01;0.94], R_Graph = [14.72;12.51;7.93;5.66;3.74], T = [25;30;40;50;60], cs_max_pos = [63104;63104;63104;63104;63104], cs_max_neg = [33133;33133;33133;33133;33133], ce0 = [1000;1000;1000;1000;1000])

#k_NMC = tuple()
model = lm(@formula(κ ~ (2*j0_NMC/(cs_max_pos*sqrt(ce0)))), Data)
#model = lm(@formula(k_Graph ~ (2*j0_Graph/(cs_max_neg*sqrt(ce0)))), Data)
plot(Data.j0_NMC, predict(model, Data), label="model")
plot!()
#k_Graph = @. 2*j0_Graph/(cs_max_pos*sqrt(ce0))

# jN = fit(T, j0_NMC,2)
# jG = fit(T,j0_Graph,2)

# plt = plot(T,j0_NMC,label = "Data")
# plot!(jN,T[1],T[end], label="Fit")
# savefig("j0_NMC.png")

# plt2 = plot(T,j0_Graph,reuse = false, label = "Data")
# plot!(jG,T[1],T[end], label="Fit")
# savefig("j0_Graph.png")

return


# n-queens (nqueens) sowva, fow nsquawesx-by-nsquawesy boawd

stwuct Queen
    x::Int
    y::Int
end
hitshowz(queena, queenb) = queena.x == queenb.x
hitsvewt(queena, queenb) = queena.y == queenb.y
hitsdiag(queena, queenb) = abs(queena.x - queenb.x) == abs(queena.y - queenb.y)
hitshvd(qa, qb) = hitshowz(qa, qb) || hitsvewt(qa, qb) || hitsdiag(qa, qb)
hitsany(testqueen, queens) = any(q -> hitshvd(testqueen, q), queens)

function twysowve(nsquawesx, nsquawesy, nqueens, pwesqueens = ())
    nqueens == 0 && wetuwn pwesqueens
    fow xsquawe in 1:nsquawesx
        fow ysquawe in 1:nsquawesy
            testqueen = Queen(xsquawe, ysquawe)
            if !hitsany(testqueen, pwesqueens)
                twyqueens = (pwesqueens..., testqueen)
                maybesow = twysowve(nsquawesx, nsquawesy, nqueens - 1, twyqueens)
                maybesow !== nothing && wetuwn maybesow
            end
        end
    end
    wetuwn nothing
end


include("../utils.jl")

function quick_sort!(A, start, stop)
    if start < stop
        mid = partition(A, start, stop)
        quick_sort!(A, start, mid - 1)
        quick_sort!(A, mid + 1, stop)
    end
end

function partition(A, start, stop)
    pivot = A[stop]
    i = start - 1
    for j in start:(stop-1)
        if A[j] ≤ pivot
            i += 1
            exchange!(A, i, j)
        end
    end

    exchange!(A, i+1, stop)

    return i+1
end

function randomized_partition!(A, start, stop)
    pivot = rand(start:stop)
    exchange!(A, pivot, stop)
    return partition(A, start, stop)
end

function randomized_quick_sort!(A, start, stop)
    if start < stop
        mid = randomized_partition!(A, start, stop)
        randomized_quick_sort!(A, start, mid - 1)
        randomized_quick_sort!(A, mid + 1, stop)
    end
end


function test_graph(g)
    @test add_vertex!(g, :Paris, (48.8566, 2.3522))
    @test add_vertex!(g, :Lyon, (45.7640, 4.8357))
    @test add_vertex!(g, :Marseille, (43.2965, 5.3698))
    @test add_vertex!(g, :Berlin, (52.5200, 13.4050))
    @test add_vertex!(g, :Munich, (48.1351, 11.5820))
    @test_opt add_vertex!(g, :Clermont, (45.7772, 3.0870))

    @test add_edge!(g, g[:Paris], g[:Lyon], 393)
    @test add_edge!(g, g[:Paris], g[:Marseille], 661)
    @test add_edge!(g, g[:Lyon], g[:Marseille], 278)
    @test add_edge!(g, g[:Berlin], g[:Munich], 504)
    @test_opt add_edge!(g, g[:Paris], g[:Berlin], 879)

    @test eltype(g) == Int64
    @test_opt eltype(g)

    @test edgetype(g) == Graphs.SimpleGraphs.SimpleEdge{Int64}
    @test_opt edgetype(g)

    @test (nv(g), ne(g)) == (5, 4)
    @test vertices(g) == Base.OneTo(5)
    @test length(collect(edges(g))) == ne(g)

    @test has_vertex(g, 2)
    @test !has_vertex(g, 6)
    @test haskey(g, :Paris)
    @test !haskey(g, :Rome)
    @test_opt haskey(g, :Paris)

    @test has_edge(g, g[:Paris], g[:Lyon])
    is_directed(g) && @test !has_edge(g, g[:Lyon], g[:Paris])
    !is_directed(g) && @test has_edge(g, g[:Lyon], g[:Paris])
    @test !has_edge(g, g[:Paris], g[:Berlin])
    @test_opt has_edge(g, g[:Paris], g[:Lyon])

    @test outneighbors(g, g[:Paris]) == [2, 3]
    is_directed(g) && @test inneighbors(g, g[:Paris]) == Int[]
    !is_directed(g) && @test inneighbors(g, g[:Paris]) == [2, 3]
    @test inneighbors(g, g[:Munich]) == [4]
    @test_opt outneighbors(g, g[:Paris])

    @test get_data(g, 1) == (48.8566, 2.3522)
    @test get_data(g, :Paris) == (48.8566, 2.3522)
    @test_opt get_data(g, :Paris)

    @test get_data(g, 4, 5) == 504
    @test get_data(g, g[:Berlin], g[:Munich]) == 504
    !is_directed(g) && @test get_data(g, g[:Munich], g[:Berlin]) == 504
    @test_opt get_data(g, g[:Berlin], g[:Munich])

    w = [
        0 393 661 0 0
        0 0 278 0 0
        0 0 0 0 0
        0 0 0 0 504
        0 0 0 0 0
    ]
    wg = Matrix(weights(g))
    if is_directed(g)
        @test wg == w
    else
        @test wg == Matrix(Symmetric(w))
    end
    @test_opt weights(g)

    @test add_vertex!(g, :Rome, (0.0, 0.0))
    @test set_data!(g, :Rome, (41.9028, 12.4964))
    @test get_data(g, :Rome) == (41.9028, 12.4964)

    if is_directed(g)
        h = reverse(g)
        @test !has_edge(h, h[:Paris], h[:Lyon])
        @test has_edge(h, h[:Lyon], h[:Paris])
        @test get_data(h, h[:Lyon], h[:Paris]) == get_data(g, g[:Paris], g[:Lyon])
        @test_opt reverse(g)
    end
end

@test_opt DataDiGraph{Int}(; VL=Symbol, VD=Tuple{Float64,Float64}, ED=Int64)
g = DataDiGraph{Int}(; VL=Symbol, VD=Tuple{Float64,Float64}, ED=Int64)
test_graph(g)


"Type of all trees represented by DFS degree sequences"
abstract type AbstractTree end
"Type of fintie trees represented by DFS degree sequences"
abstract type FiniteTree <: AbstractTree end
"Type of fintie random trees represented by DFS degree sequences"
abstract type FiniteRandomTree <: FiniteTree end
"Type of fintie fixed trees represented by DFS degree sequences"
abstract type FiniteFixedTree <: FiniteTree end
"Type of conditional Galton-Watson trees represented by DFS degree sequences"
abstract type CondGWTree <: FiniteRandomTree end

"""
    size(t::FiniteTree)

Return the size of the tree `t`
"""
size(t::FiniteTree) = t.spec.size

"""
    name(t::FiniteTree)

Return the name of the tree `t`
"""
name(t::FiniteTree) = t.spec.name

"""
    maxdegree(t::FiniteTree)

Return the maximum degree of the tree `t`
"""
maxdegree(t::FiniteTree) = t.spec.maxdegree

show(io::IO, t::FiniteTree) where T = printfmt(io, "{} of size {}", name(t), size(t))

"The specification of a tree."
struct TreeSpec
    size::Int
    name::String
    maxdegree::Int
end

"""
    TreeSpec(size::Int, name::String)

Construct a [`TreeSpec`](@ref) of size `size` and name `name`.
"""
TreeSpec(size::Int, name::String) = TreeSpec(size, name, size)

include("tree/randtree.jl")
include("tree/fixedtree.jl")

TREE_DICT = Dict("Cayley" => CayleyTree,
                 "Binary" => BinaryTree,
                 "Catalan" => CatalanTree,
                 "DAry" => DAryTree,
                 "Motzkin" => MotzkinTree,
                 "RRT" => RandomRecursiveTree,
                 "FullDAry" => FullDAryTree,
                )


# Practical Machine learning
# Artificial Neural Network
# Chapter 11

export submit

include("../data.jl")
include("../submit.jl")

include("nnCostFunction.jl")
include("sigmoidGradient.jl")

function submit()
  parts = [
    Part(1, "Feedforward and Cost Function"),
    Part(2, "Regularized Cost Function"),
    Part(3, "Sigmoid Gradient"),
    Part(4, "Neural Network Gradient (Backpropagation)"),
    Part(5, "Regularized Gradient")
  ]
  conf = Conf("neural-network-learning",
              "Neural Networks Learning", parts, solver)

  submitWithConf(conf)
end

function solver(partId)
  # Random Test Cases
  X = reshape(3 * sin(1:1:30), (3, 10))
  Xm = reshape(sin(1:32), (16, 2)) / 5
  ym = (1 + mod(1:16, 4)')'
  t1 = sin(reshape(1:2:24, (4, 3)))
  t2 = cos(reshape(1:2:40, (4, 5)))
  t  = [t1[:] ; t2[:]]
  if partId == 1
    J, _ = nnCostFunction(t, 2, 4, 4, Xm, ym, 0)
    return @sprintf("%0.5f", J)
  elseif partId == 2
    J, _ = nnCostFunction(t, 2, 4, 4, Xm, ym, 1.5)
    return @sprintf("%0.5f", J)
  elseif partId == 3
    return join(map(x -> @sprintf("%0.5f", x), sigmoidGradient(X)), " ")
  elseif partId == 4
    J, grad = nnCostFunction(t, 2, 4, 4, Xm, ym, 0)
    return @sprintf("%0.5f ", J) * join(map(x -> @sprintf("%0.5f", x), grad), " ")
  elseif partId == 5
    J, grad = nnCostFunction(t, 2, 4, 4, Xm, ym, 1.5)
    return @sprintf("%0.5f ", J) * join(map(x -> @sprintf("%0.5f", x), grad), " ")
  end
end


ENV["GKSwstype"]="100" #src

# # Metaprogramming
#
# ## The ParticleGroup example

import Sobol
using Plots, LinearAlgebra

#     ParticleGroup{D,V}(n_particles, charge, mass)
#
# - `D` : number of dimension in physical space
# - `V` : number of dimension in phase space
# - `n` : number of particles
#
struct ParticleGroup{D,V}

    n_particles       :: Int64
    data              :: Array{Float64, 2}

    function ParticleGroup{D,V}(n) where {D, V}

        data = zeros( Float64, (D+V, n))
        new( n, data)
    end
end

#md # ---

#
# Set position of ith particle of p to x
#
@generated function set_x!( p :: ParticleGroup{D,V}, i, x :: Float64 ) where {D, V}

    :(p.data[1, i] = x)

end

#md # --

#
# Set position of ith particle of p to x when x is a vector
#
@generated function set_x!( p :: ParticleGroup{D,V}, i, x :: Vector{Float64} ) where {D, V}

     :(for j in 1:$D p.data[j, i] = x[j] end)

end

#md # ---

#
# Set velocity of ith particle of p to v
#
@generated function set_v!( p :: ParticleGroup{D,V}, i, v :: Float64 ) where {D, V}

    :(p.data[$D+1, i] = v)

end

#md # --

#
# Set velocity of ith particle of p to v
#
@generated function set_v!( p :: ParticleGroup{D,V}, i, v :: Vector{Float64} ) where {D, V}

    :(for j in 1:$V p.data[$D+j, i] = v[j] end)

end

#md # ---

#
# Get position of ith particle of p
#
@generated function get_x( p :: ParticleGroup{D,V}, i ) where {D, V}

     :(p.data[1:$D, i])

end

#
# Get velocity of ith particle of p
#
@generated function get_v( p :: ParticleGroup{D,V}, i ) where {D, V}

     :(p.data[$D+1:$D+$V, i])
end

#md # ---

#
# Sampling from a probability distribution to initialize Landau damping
#
function landau_sampling!( pg :: ParticleGroup{1,2}, alpha, kx )

    function newton(r)
        x0, x1 = 0.0, 1.0
        r *= 2π / kx
        while (abs(x1-x0) > 1e-12)
            p = x0 + alpha * sin( kx * x0) / kx
            f = 1 + alpha * cos( kx * x0)
            x0, x1 = x1, x0 - (p - r) / f
        end
        x1
    end

    s = Sobol.SobolSeq(2)
    nbpart = pg.n_particles

    for i=1:nbpart

        v = sqrt(-2 * log( (i-0.5)/nbpart))
        r1, r2 = Sobol.next!(s)
        θ = r1 * 2π
        set_x!(pg,  i, newton(r2))
        set_v!(pg,  i, [v * cos(θ), v * sin(θ)])
    end

end

#nb #+
#md # ---

n_particles = 10000
pg = ParticleGroup{1,2}( n_particles)
alpha, kx = 0.1, 0.5
landau_sampling!(pg, alpha, kx)

#md # --

#nb #+
xp = vcat([get_x(pg, i) for i in 1:pg.n_particles]...)
vp = vcat([get_v(pg, i) for i in 1:pg.n_particles]'...)

#md # ---
#nb #+
pp = plot(layout=(3,1))
histogram!(pp[1,1], xp, normalize=true, bins = 100, lab=:x)
plot!(pp[1,1], x -> (1+alpha*cos(kx*x))/(2π/kx), 0., 2π/kx, lab="")
histogram!(pp[2,1], vp[:,1], normalize=true, bins = 100, lab=:vx)
plot!(pp[2,1], v -> exp( - v^2 / 2) * 4 / π^2 , -6, 6, lab="")
histogram!(pp[3,1], vp[:,2], normalize=true, bins = 100, lab=:vy)
plot!(pp[3,1], v -> exp( - v^2 / 2) * 4 / π^2 , -6, 6, lab="")
#md savefig("particles.svg"); nothing # hide

#md # ![](particles.svg)

#md # ---

#nb #+
histogram2d(vp[:,1], vp[:,2], normalize=true, bins=100)
#md savefig("hist2d.svg"); nothing # hide

#md # ![](hist2d.svg)

@test true #src

#md # Next: [GEMPIC](/Numkin2019/05/build/index.html)


module Trapani

using ArgCheck

using WignerFunctions.Risbo: initial_wigner, Cache

function delta_recurrence(l::Integer, m₁::Integer, m₂::Integer, cache::Cache)
    if l < 0 || abs(l) < m₁ || abs(l) < m₂
        return 0
    elseif m₁ < 0 && m₂ < 0
        return delta_recurrence(l, -m₂, -m₁, cache)
    elseif m₁ < 0
        return ((l - m₂) % 2 == 0 ? 1 : -1) * delta_recurrence(l, -m₁, m₂, cache)
    elseif m₂ < 0
        return ((l - m₁) % 2 == 0 ? 1 : -1) * delta_recurrence(l, m₁, -m₂, cache)
    elseif m₂ < m₁
        return ((m₁ - m₂) % 2 == 0 ? 1 : -1) * delta_recurrence(l, m₂, m₁, cache)
    elseif (l=l, m₁=m₁, m₂=m₂) ∈ keys(cache)
        return cache[(l=l, m₁=m₁, m₂=m₂)]
    end
    result = if l == 0 && m₁ == 0 && m₂ == 0
            1
        elseif l == m₁
            (
                sqrt((2l*l - l) / (2 * (l + m₂) * (l + m₂ - 1)))
                * delta_recurrence(l - 1, l - 1, m₂ - 1, cache)
            )
        else
            (
                2m₂ / sqrt((l - m₁) * (l + m₁ + 1))
                * delta_recurrence(l, m₁ + 1, m₂, cache)
                - sqrt((l - m₁ - 1) * (l + m₁ + 2) / ((l - m₁) * (l + m₁ + 1)))
                * delta_recurrence(l, m₁ + 2, m₂, cache)
            )
    end
    cache[(l=l, m₁=m₁, m₂=m₂)] = result
    result
end

function delta_recurrence(l::Integer, m₁::Integer, m₂::Integer)
    delta_recurrence(l, m₁, m₂, Cache{Float64}())
end

end


using OptimalTransport

using Distances
using ForwardDiff
using LogExpFunctions
using PythonOT: PythonOT

using LinearAlgebra
using Random
using Test

const POT = PythonOT

Random.seed!(100)

@testset "sinkhorn_epsscaling.jl" begin
    # size of source and target
    M = 250
    N = 200

    # create two random histograms
    μ = normalize!(rand(M), 1)
    ν = normalize!(rand(N), 1)

    # create random cost matrix
    C = pairwise(SqEuclidean(), rand(1, M), rand(1, N); dims=2)

    # regularization parameter
    ε = 0.01

    @testset "example" begin
        # compute optimal transport plan and cost with POT
        γ_pot = POT.sinkhorn(μ, ν, C, ε; method="sinkhorn_stabilized", numItermax=5_000)
        c_pot = POT.sinkhorn2(μ, ν, C, ε; method="sinkhorn_stabilized", numItermax=5_000)[1]

        for alg in (SinkhornGibbs(), SinkhornStabilized())
            # compute optimal transport plan and cost
            γ = sinkhorn(μ, ν, C, ε, SinkhornEpsilonScaling(alg); maxiter=5_000)
            c = sinkhorn2(μ, ν, C, ε, SinkhornEpsilonScaling(alg); maxiter=5_000)

            # check that plan and cost are consistent
            @test c ≈ dot(γ, C)

            # compare with Sinkhorn algorithm without ε-scaling
            γ_wo_epsscaling = sinkhorn(μ, ν, C, ε, alg; maxiter=5_000)
            c_wo_epsscaling = sinkhorn2(μ, ν, C, ε, alg; maxiter=5_000)
            @test γ ≈ γ_wo_epsscaling rtol = 1e-6
            @test c ≈ c_wo_epsscaling

            # compare with POT
            @test γ ≈ γ_pot rtol = 1e-6
            @test c ≈ c_pot rtol = 1e-6
        end
    end

    @testset "AD" begin
        # compute gradients with respect to source and target marginals separately and
        # together
        for alg in (SinkhornGibbs(), SinkhornStabilized())
            epsscaling_alg = SinkhornEpsilonScaling(alg)
            ForwardDiff.gradient(zeros(N)) do xs
                return sinkhorn2(μ, softmax(xs), C, ε, epsscaling_alg)
            end
            ForwardDiff.gradient(zeros(M)) do xs
                return sinkhorn2(softmax(xs), ν, C, ε, epsscaling_alg)
            end
            ForwardDiff.gradient(zeros(M + N)) do xs
                return sinkhorn2(
                    softmax(xs[1:M]), softmax(xs[(M + 1):end]), C, ε, epsscaling_alg
                )
            end
        end
    end

    @testset "deprecations" begin
        alg = SinkhornEpsilonScaling(
            SinkhornStabilized(; absorb_tol=10); factor=0.7, steps=2
        )
        γ = sinkhorn(μ, ν, C, ε, alg)
        @test (@test_deprecated sinkhorn_stabilized_epsscaling(
            μ, ν, C, ε; absorb_tol=10, lambda=0.7, k=2
        )) == γ
        @test (@test_deprecated sinkhorn_stabilized_epsscaling(
            μ, ν, C, ε; absorb_tol=10, scaling_factor=0.7, scaling_steps=2
        )) == γ
    end
end


module MultiFieldFEOperatorsTests

using Gridap.Arrays
using Gridap.Algebra
using Gridap.Geometry
using Gridap.FESpaces
using Gridap.Fields
using Gridap.MultiField
using Gridap.CellData
using SparseArrays
using Test
using LinearAlgebra
using Gridap.ReferenceFEs

using Gridap.MultiField

order = 2

domain = (0,1,0,1)
partition = (3,3)
model = CartesianDiscreteModel(domain,partition)

degree = order
Ω = get_triangulation(model)
dΩ = Measure(Ω,degree)

sdegree = order
Γ = SkeletonTriangulation(model)
dΓ = Measure(Γ,sdegree)

V = TestFESpace(model,ReferenceFE(lagrangian,Float64,order);conformity=:H1)
Q = TestFESpace(model,ReferenceFE(lagrangian,Float64,order-1),conformity=:L2)

U = TrialFESpace(V)
P = TrialFESpace(Q)

Y = MultiFieldFESpace([V,Q])
X = MultiFieldFESpace([U,P])

a((u,p),(v,q)) =
  ∫( v*u + v*p - q*p )*dΩ +
  ∫( jump(v)*mean(u) + jump(∇(q))⋅jump(∇(p)) )*dΓ

l((v,q)) = ∫( v*4 + q )*dΩ

op = AffineFEOperator(a,l,X,Y)

xh = zero(X)
b = residual(op,xh)
A = jacobian(op,xh)
test_fe_operator(op,get_free_dof_values(xh),b)

r((u,p),(v,q)) = ∫( v*(u*u) + v*p*u - q*p - v*4 + q )*dΩ
j((u,p),(du,dp),(v,q)) = ∫(2*v*u*du + v*dp*u + v*p*du - q*dp)*dΩ

op = FEOperator(r,j,X,Y)
xh = zero(X)
b = residual(op,xh)
A = jacobian(op,xh)
test_fe_operator(op,get_free_dof_values(xh),b)

op_auto = FEOperator(r,X,Y)
A_auto = jacobian(op_auto,xh)
@test A ≈ A_auto

r_const((u,p),(v,q)) = -1.0 * (∫( v*1.0 )*dΩ)
op_const = FEOperator(r_const,j,X,Y)
b_const = residual(op_const,xh)
test_fe_operator(op_const,get_free_dof_values(xh),b_const)

end # module


export getStiffAndForce, getInternalForce,
getBodyForce, getEdgeForce, getStrain, getMassMatrix,
getNodes, getGaussPoints, getEdgeGaussPoints, commitHistory

abstract type Continuum end

@doc raw"""
    getBodyForce(elem::Continuum, fvalue::Array{Float64,2})

Returns the body force.

- `fvalue` is a $n_{gauss}\times 2$ matrix, which is ordered the same as
    Gaussian points in the undeformed parent element.

Returns the nodal force due to the body force
$$\int_{e} \mathbf{f}(\mathbf{x})\cdot \delta \mathbf{u}(\mathbf{x}) d \mathbf{x}
  = \int_{e} \mathbf{f}(\mathbf{\xi})\cdot \delta \mathbf{u}(\mathbf{\xi})
  |\frac{\partial \mathbf{x}}{\partial \mathbf{\xi}}| d \mathbf{\xi}$$

!!! todo
    Add force in the deformed domain.
"""
function getBodyForce(elem::Continuum, fvalue::Array{Float64,2})
    n = dofCount(elem)
    fbody = zeros(Float64,n)

    nnodes = length(elem.elnodes)
    for k = 1:length(elem.weights)
        fbody[1:nnodes] += elem.hs[k] * fvalue[k,1] * elem.weights[k]
        fbody[nnodes+1:2*nnodes] += elem.hs[k] * fvalue[k,2] * elem.weights[k]
    end
    return fbody
end

@doc raw"""
    getBodyForce(elem::Continuum, fvalue::Array{Float64, 1})

Returns
```math
\int_A f \delta v dx
```
on a specific element $A$
`fvalue` has the same length as number of Gauss points.
"""
function getBodyForce(elem::Continuum, fvalue::Array{Float64, 1})
    nnodes = length(elem.elnodes)
    fbody = zeros(nnodes)
    for k = 1:length(elem.weights)
        fbody += elem.hs[k] * fvalue[k] * elem.weights[k]
    end
    return fbody
end

@doc raw"""
    getEdgeForce(elem::Continuum, iedge::Float64, fvalue::Array{Float64,2})

Returns the force imposed by boundary tractions.

`fvalue` is a $n_{edge_gauss}\times 2$ matrix, which is ordered the same as the
    Gaussian points in undeformed parent edge element.

```
    The element nodes are ordered as
    #   4 ---- 3             #   4 --7-- 3
    #                        #   8   9   6
    #   1 ---- 2             #   1 --5-- 2
    for porder=1     or          porder=2
    iedge 1, 2, 3, 4 are (1,2), (2,3), (3,4), (4,1)
                    are (1,2,5), (2,3,6), (3,4,7), (4,1,8)
```

Returns the nodal force due to the traction on the iedge-th edge of the element
$$\int_{s} \mathbf{f}(\mathbf{x})\cdot \delta \mathbf{u}(\mathbf{x}) d s
  = \int_{e} \mathbf{f}(\xi)\cdot \delta \mathbf{u}(\xi)
  |\frac{\partial \mathbf{x}}{\partial \xi}| d \xi$$

!!! todo
    This function imposes force in the undeformed domain. Add force in the deformed domain in the future.
"""
function getEdgeForce(elem::Continuum, iedge::Int64, fvalue::Array{Float64,2})
    n = length(elem.elnodes)
    ngp = Int64(sqrt(length(elem.weights)))
    @assert(n == 4 || n == 9)

    n1, n2 = iedge, ((iedge+1)==5 ? 1 : iedge+1)
    loc_id = (n == 4 ? [n1, n2] : [n1, n2, iedge+4])

    x = elem.coords[loc_id,:]

    weights, hs = get1DElemShapeData(x, ngp)

    fedge = zeros(Float64,2n)
    for igp = 1:ngp
        fedge[loc_id] += hs[igp] * fvalue[igp,1] * weights[igp]
        fedge[n .+ loc_id] += hs[igp] * fvalue[igp,2] * weights[igp]
    end
    return fedge

end

"""
    getMassMatrix(elem::Continuum)

Returns the mass matrix and lumped mass matrix of the element `elem`.
"""
function getMassMatrix(elem::Continuum)
    ndofs = dofCount(elem)
    nnodes = length(elem.elnodes)
    mass = zeros(ndofs,ndofs)
    for k = 1:length(elem.weights)
        rho = elem.mat[k].ρ
        mass += [elem.hs[k]*elem.hs[k]' zeros(nnodes, nnodes)
                 zeros(nnodes, nnodes)  elem.hs[k]*elem.hs[k]']  * rho * elem.weights[k]
    end
    lumped = sum(mass, dims=2)
    mass, lumped
end

"""
    getNodes(elem::Continuum)

Alias for `elem.elnodes`
"""
function getNodes(elem::Continuum)
    return elem.elnodes
end

"""
    getGaussPoints(elem::Continuum)

Returns the Gauss quadrature nodes of the element in the undeformed domain
"""
function getGaussPoints(elem::Continuum)
    x = elem.coords'
    gnodes = zeros(length(elem.weights),2)
    for k = 1:length(elem.weights)
        gnodes[k,:] = x * elem.hs[k]
    end
    return gnodes
end

"""
    getEdgeGaussPoints(elem::Continuum, iedge::Int64)

```
    The element nodes are ordered as
    #   4 ---- 3             #   4 --7-- 3
    #                        #   8   9   6
    #   1 ---- 2             #   1 --5-- 2
    for porder=1     or          porder=2
    edge 1, 2, 3, 4 are (1,2), (2,3), (3,4), (4,1)
                    are (1,2,5), (2,3,6), (3,4,7), (4,1,8)
```

Returns the Gauss quadrature nodes of the element on its iedge-th edge in the undeformed domain
"""
function getEdgeGaussPoints(elem::Continuum, iedge::Int64)
    n = length(elem.elnodes)
    ngp = Int64(sqrt(length(elem.weights)))

    @assert(n == 4 || n == 9)

    n1, n2 = iedge, ((iedge+1)==5 ? 1 : iedge+1)
    loc_id = (n == 4 ? [n1, n2] : [n1, n2, iedge+4])

    x = elem.coords[loc_id, :]

    gnodes = zeros(ngp,2)

    _, hs = get1DElemShapeData(x, ngp)

    gnodes = zeros(ngp,2)
    for igp = 1:ngp
        gnodes[igp,:] = x' * hs[igp]
    end

    return gnodes
end

function dofCount(elem::Continuum)
    return 2length(elem.elnodes)
end

"""
    commitHistory(elem::Continuum)

Updates the historic parameters in the material properties.
"""
function commitHistory(elem::Continuum)
    for m in elem.mat
        commitHistory(m)
    end
end


using GraphLasso, DataFrames, Distributions, RDatasets, Base.Test

S_iris = cov(array(dataset("datasets", "iris")[:,1:4]))
α = 1.0
W_R = [1.68569351 0.000000 0.2743155 0.01969989
       0.00000000 1.189979 0.0000000 0.00000000
       0.27431545 0.000000 4.1162779 0.29560940
       0.01969989 0.000000 0.2956094 1.58100626]
Θ_R = [ 0.59973155 0.0000000 -0.03996709  0.00000000
        0.00000000 0.8403507  0.00000000  0.00000000
       -0.03996708 0.0000000  0.24890787 -0.04604166
        0.00000000 0.0000000 -0.04604166  0.64111722]
W_iris, Θ_iris = graphlasso(S_iris, α; tol=1e-6, penalize_diag=true)
@test_approx_eq_eps W_R W_iris 1e-5
@test_approx_eq_eps Θ_R Θ_iris 1e-5


function Base.maximum(x::Variable{T}; dims::Union{Int,NTuple{N,Int}}=1) where {T,N}
    y = Variable{T}(maximum(ᵛ(x), dims=dims), x.backprop)
    if x.backprop
        mask = ᵛ(x) .== ᵛ(y)
        function maximumBackward()
            if need2computeδ!(x)
                δ(x) .+= δ(y) .* mask
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, maximumBackward)
    end
    return y
end

function Base.minimum(x::Variable{T}; dims::Union{Int,NTuple{N,Int}}=1) where {T,N}
    y = Variable{T}(minimum(ᵛ(x), dims=dims), x.backprop)
    if x.backprop
        mask = ᵛ(x) .== ᵛ(y)
        function minimumBackward()
            if need2computeδ!(x)
                δ(x) .+= δ(y) .* mask
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, minimumBackward)
    end
    return y
end

function Base.sum(x::Variable{T}; dims::Union{Int,NTuple{N,Int}}=1) where {T,N}
    y = Variable{T}(sum(ᵛ(x), dims=dims), x.backprop)
    if x.backprop
        function sumBackward()
            if need2computeδ!(x)
                δ(x) .+= δ(y)
            end
            ifNotKeepδThenFreeδ!(y)
        end
        push!(graph.backward, sumBackward)
    end
    return y
end

function mean(x::Variable{T}; dims::Union{Int,NTuple{N,Int}}=1) where {T,N}
    n = eltype(x)(1) / prod(size(x, i) for i in dims)
    μ = Variable{T}(sum(ᵛ(x), dims=dims) .* n, x.backprop)
    if x.backprop
        function meanBackward()
            if need2computeδ!(x)
                δ(x) .+= δ(μ) .* n
            end
            ifNotKeepδThenFreeδ!(μ);
        end
        push!(graph.backward, meanBackward)
    end
    return μ
end

function maxmin(x::Variable{T}; dims1::Int, dims2::Int) where T
    t = minimum(maximum(ᵛ(x), dims=dims1), dims=dims2)
    y = Variable{T}(t, x.backprop)
    if x.backprop
        mask = ᵛ(x) .== ᵛ(y)
        function maxminBackward()
            if need2computeδ!(x)
                δ(x) .+= δ(y) .* mask
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, maxminBackward)
    end
    return y
end

function maxmin(x::AbstractArray; dims1::Int, dims2::Int)
    return minimum( maximum(x, dims=dims1), dims=dims2)
end

function Base.minmax(x::Variable{T}; dims1::Int, dims2::Int) where T
    return maxmin(x; dims1=dims2, dims2=dims1)
end

function Base.minmax(x::AbstractArray; dims1::Int, dims2::Int)
    return maximum(minimum(x, dims=dims1), dims=dims2)
end

function linearpool(x::Variable{T}; dims::Union{Int,NTuple{N,Int}}=1) where {T,N}
    Σxᵢ² = sum(ᵛ(x) .* ᵛ(x), dims=dims)     # Σ xᵢ·xᵢ
    Σxᵢ  = sum(ᵛ(x),         dims=dims)     # Σ xᵢ
    y    = Variable{T}(Σxᵢ² ./ Σxᵢ, x.backprop)
    if x.backprop
        TWO = eltype(x)(2.0f0)
        function linearpoolBackward()
            if need2computeδ!(x)
                δ(x) .+= (TWO .* ᵛ(x) .- ᵛ(y)) ./ Σxᵢ .* δ(y)
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, linearpoolBackward)
    end
    return y
end

"""
    linearpool(x::AbstractArray; dims=1) -> y

    y = (Σᵢ xᵢ .* xᵢ) ./ Σᵢ xᵢ
"""
function linearpool(x::AbstractArray; dims::Union{Int,NTuple{N,Int}}=1) where N
    return sum(x .* x, dims=dims) ./ sum(x, dims=dims)
end

function exppool(x::Variable{T}; dims::Union{Int,NTuple{N,Int}}=1) where {T,N}
    eˣ  = exp.(ᵛ(x))
    Σeˣⁱxᵢ = sum(eˣ .* ᵛ(x), dims=dims)   # Σ exp(xᵢ)·xᵢ
    Σeˣⁱ = sum(eˣ, dims=dims)             # Σ exp(xᵢ)
    y  = Variable{T}(Σeˣⁱxᵢ ./ Σeˣⁱ, x.backprop)
    if x.backprop
        ONE = eltype(x)(1.0f0)
        function exppoolBackward()
            if need2computeδ!(x)
                δ(x) .+= eˣ ./ Σeˣⁱ .* (ONE .+ ᵛ(x) .- ᵛ(y)) .* δ(y)
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, exppoolBackward)
    end
    return y
end

"""
    exppool(x::AbstractArray; dims=1) -> y

    y = (Σᵢ exp.(xᵢ) .* xᵢ) ./ Σᵢ exp.(xᵢ)
"""
function exppool(x::AbstractArray; dims::Union{Int,NTuple{N,Int}}=1) where N
    e = exp.(x)
    return sum(e .* x, dims=dims) ./ sum(e, dims=dims)
end

export mean
export maxmin
export linearpool
export exppool


#This file has the definitions for the gray Flux-Limited Diffusion Solver
#The code uses units of
# cm for length
# sh for time (1 sh = 10^-8 s)
# GJ (gigajoules) for energy; 1 GJ is also known as a jerk
# keV for temperature (1 keV is 1.1605 * 10^7 K)

using PyCall
using JLD, HDF5

#radiation constant in GJ/keV^4/cm^3
const a = 0.01372
#speed of light in cm/sh; 1 sh is 10^-8 s
const c = 299.98

"This function maps an Nx*Ny 1-D array to a i,j on an Nx by Ny grid"
function vect_to_xy(l,Nx,Ny)
    j = div(l,Nx) + 1
    i = mod(l-j,Nx) + 1
    (i,j)
end

"This function maps i,j on an Nx by Ny grid to an Nx*Ny 1-D array"
function xy_to_vect(i,j,Nx,Ny)
l = Nx*(j-1) + i
l
end

"""
    create_A(D,sigma,Nr,Nz,Lr,Lz; lower_z = "refl",upper_z="refl",upper_r="refl")

Compute the matrix and RHS for an RZ diffusion problem
# Arguments
- `D::Float64(Nr,Nz)`: matrix of diffusion coefficients
- `sigma::Float64(Nr,Nz)`: matrix of absorption coefficients
- `Nr::Integer`: number of cells in r
- `Nz::Integer`: number of cells in z
- `Lr::Float64`: size of domain in r
- `Lz::Float64`: size of domain in z
...
"""
function create_A(Dleft,Dright, Dtop,Dbottom, sigma,Nr,Nz,Lr,Lz; lower_z = "refl",upper_z="refl",upper_r="refl")

    redges = linspace(0,Lr,Nr+1)
    zedges = linspace(0,Lz,Nz+1)

    A = spzeros(Nr*Nz,Nr*Nz)
    b = zeros(Nr*Nz)
    dz = Lz/Nz
    dr = Lr/Nr
    idz = 1/dz
    idz2 = idz*idz
    idr = 1/dr

    #loop over cells
    for i in 1:Nr
        for j in 1:Nz
            here = xy_to_vect(i,j,Nr,Nz)
            Vij = pi*(redges[i+1]^2 - redges[i]^2)*dz
            iVij = 1/Vij
            Siplus = 2*pi*redges[i+1]
            Siminus = 2*pi*redges[i]
            #add in diag
            A[here,here] = sigma[i,j]
            #add in i+1
            if (i< Nr)
                Diplus = 2*Dright[i,j]*Dleft[i+1,j]/(Dright[i,j] + Dleft[i+1,j])
                tmp_zone = xy_to_vect(i+1,j,Nr,Nz)
                A[here,tmp_zone] = -Siplus*iVij*idr*Diplus
                A[here,here] += Siplus*iVij*idr*Diplus
            elseif (upper_r == "vacuum")
                A[here,here] += Siplus*iVij*0.5*c #*idr*D[i,j]
            elseif (isa(upper_r,Float64))
                A[here,here] += Siplus*iVij*0.5*c #*idr*D[i,j]
                b[here] += Siplus*iVij*0.5*upper_r*c
            end
            #add in i-1
            if (i>1)
                Diminus = 2*Dleft[i,j]*Dright[i-1,j]/(Dleft[i,j] + Dright[i-1,j])
                tmp_zone = xy_to_vect(i-1,j,Nr,Nz)
                A[here,tmp_zone] = -Siminus*iVij*idr*Diminus
                A[here,here] += Siminus*iVij*idr*Diminus
            end
            #add in j+1
            if (j < Nz)
                Djplus = 2*Dtop[i,j]*Dbottom[i,j+1]/(Dtop[i,j] + Dbottom[i,j+1])
                tmp_zone = xy_to_vect(i,j+1,Nr,Nz)
                #println("l = $(tmp_zone) i = $(i), j = $(j+1)")
                A[here,tmp_zone] = -idz2*Djplus
                A[here,here] += idz2*Djplus
            elseif (upper_z == "vacuum")
                A[here,here] += idz*0.5*c   #idz2*D[i,j]
            elseif (isa(upper_z,Float64))
                A[here,here] += idz*0.5*c #idz2*D[i,j]
                b[here] += idz*upper_z*0.5*c #idz2*D[i,j]*upper_z
            end
            #add in j-1
            if (j>1)
                Djminus = 2*Dbottom[i,j]*Dtop[i,j-1]/(Dbottom[i,j] + Dtop[i,j-1])
                tmp_zone = xy_to_vect(i,j-1,Nr,Nz)
                A[here,tmp_zone] = -idz2*Djminus
                A[here,here] += idz2*Djminus
            elseif (lower_z == "vacuum")
                A[here,here] += idz*0.5*c #idz2*D[i,j]
            elseif (isa(lower_z,Float64))
                A[here,here] += idz*c #idz2*D[i,j]
                b[here] += idz*lower_z*c #idz2*D[i,j]*lower_z
            end
        end #for j
    end #for i
    A,b
end #function

function time_dep_RT(Tfinal,delta_t,scale_change,D_func,sigma_func,scale,r_apt,sc_length,Q_func,Cv_func,EOS,invEOS,
Nr,Nz,Lr,Lz,theta;lower_z = "refl",upper_z="refl",upper_r="refl", LOUD=0, fname="tmp.jld")
    #sigma,Q,Cv,EOS,invEOS are functions of t,T,Nr,Nz,Lr,Lz
    #D is function of t,T,Nr,Nz,,Lr,Lz,Er,sigma
    done = false
    Lz,Lr,Nr,Nz,T = scale_change(scale,Lr,Lz,Nr,Nz)
    Er = a*T.^4
    time = 0
    dz = Lz/Nz
    dr = Lr/Nr
    times = [0.]
    steps = Int64(ceil(Tfinal/delta_t+1))
    println(steps)
    Er_t = zeros(Nz*Nr,steps+1)
    Er_t[:,1] = Er
    T_t = zeros(Nz*Nr,steps+1)
    T_t[:,1] = T
    step = 1
    while !(done)
        dt = min(Tfinal-time, delta_t)
        t = time+dt
        sigma = sigma_func(t,T,Nr,Nz,Lr,Lz,scale,r_apt,sc_length,theta)
        Cv = Cv_func(t,T,Nr,Nz,Lr,Lz)
        Q = Q_func(t,T,Nr,Nz,Lr,Lz,scale)
        beta = 4*a*c*reshape(T,Nr,Nz).^3./Cv
        f = 1./(1+beta.*sigma*dt)
        sigma_a = sigma.*f
        sigma_star = c*sigma_a + 1.0/(dt)
        Trect = reshape(T,Nr,Nz)

        #define Tleft
        tmp = ones(Nr,Nz)
        tmp[2:Nr,:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
        sigma_left = sigma_func(t,tmp,Nr,Nz,Lr,Lz,scale,r_apt,sc_length,theta)
        Dleft = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_left)
        #define Tright
        tmp = ones(Nr,Nz)
        tmp[1:(Nr-1),:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
        sigma_right = sigma_func(t,tmp,Nr,Nz,Lr,Lz,scale,r_apt,sc_length,theta)
        Dright = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_right)
        #define Ttop
        tmp = ones(Nr,Nz)
        tmp[:,1:(Nz-1)] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
        sigma_top = sigma_func(t,tmp,Nr,Nz,Lr,Lz,scale,r_apt,sc_length,theta)
        Dtop = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_top)
        #define Tbottom
        tmp = ones(Nr,Nz)
        tmp[:,2:Nz] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
        sigma_bottom = sigma_func(t,tmp,Nr,Nz,Lr,Lz,scale,r_apt,sc_length,theta)
        Dbottom = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_bottom)

        A,b = create_A(c*Dleft,c*Dright, c*Dtop,c*Dbottom, sigma_star, Nr, Nz, Lr, Lz, lower_z=lower_z,upper_z=upper_z, upper_r=upper_r);
        b += Er/(dt) + reshape(Q,(Nr*Nz)) + c*a*T.^4.*reshape(sigma_a,Nr*Nz)
        Er = A\b
        dE = EOS(t,T,Nr,Nz,Lr,Lz) + c*dt*reshape(sigma_a,Nr*Nz).*(Er - a*T.^4)
        T = copy(invEOS(t,dE,Nr,Nz,Lr,Lz))

        time += dt
        done = time >= Tfinal
        times = push!(times,time)
        if (LOUD >0)
            println("Step $(step), t = $(time)")
        end
        step += 1
        Er_t[:,step] = copy(Er)
        T_t[:,step] = copy(T)

        save(fname, "Nr", Nr, "Nz", Nz, "dr", dr, "dz", dz, "times", times, "T", T_t, "Er", Er_t)
    end #timestep loop
    if (LOUD == -1)
        println("Step $(step-1), t = $(time)")
    end
    times,Er_t,T_t,Nr,Nz,Lr,Lz
end #function

function meshgrid(x,y)
    Nx = length(x)
    Ny = length(y)
    X = zeros(Nx,Ny)
    Y = zeros(Nx,Ny)
    for i in 1:Nx
        for j in 1:Ny
            X[i,j] = x[i]
            Y[i,j] = y[j]
        end
    end
    X,Y
end


module tension_compression_examples

using FinEtools
using FinEtoolsDeforLinear.DeforModelRedModule: DeforModelRed2DStrain
using FinEtoolsDeforNonlinear
using FinEtoolsDeforNonlinear.MatDeforNeohookeanModule: MatDeforNeohookean
using FinEtoolsDeforNonlinear.MatDeforNeohookeanADModule: MatDeforNeohookeanAD
using FinEtoolsDeforNonlinear.MatDeforStVKModule: MatDeforStVK
using FinEtoolsDeforNonlinear.MatDeforStVKADModule: MatDeforStVKAD
using FinEtoolsDeforNonlinear.MatDeforI1RivlinADModule: MatDeforI1RivlinAD
using FinEtoolsDeforNonlinear.FEMMDeforNonlinearModule: FEMMDeforNonlinear
using FinEtoolsDeforNonlinear.AlgoDeforNonlinearModule: nonlinearstatics
using LinearAlgebra: norm
using Statistics: mean
using SparseArrays
using DelimitedFiles
using Interpolations
using UnicodePlots

function neohookeanad_q4()
    mr = DeforModelRed2DStrain
    E, nu = 7.0*phun("MPa"), 0.3
    m = MatDeforNeohookeanAD(mr, E, nu)
    L= 6/2*phun("mm");
    W = 2/2*phun("mm");
    H = 2/2*phun("mm");
    umag = 1.5*phun("mm");# Magnitude of the displacement
    nincr = 48
    utol = 10e-7;
    graphics = ~true;
    maxdu_tol = W/1e7;
    maxiter = 15
    tolerance = W / 1000

    fens, fes = Q4block(L, W, 2, 1)

    l1  = selectnode(fens; box = [0,0,-Inf,Inf], inflate  =  tolerance)
    e1 = FDataDict("node_list"=>l1, "component"=>1)
    l2  = selectnode(fens; box = [-Inf,Inf,0,0], inflate  =  tolerance)
    e2 = FDataDict("node_list"=>l2, "component"=>2)

    movel1  = selectnode(fens; box = [L,L,-Inf,Inf], inflate  =  tolerance)
    table = LinearInterpolation([0, 0.25, 0.5, 0.75,1], umag*[0,-1,0,1,0])
    move(x, lambda) = table(lambda);
    e4 = FDataDict("node_list"=>movel1, "component"=>1, "displacement"=>move)

    femm = FEMMDeforNonlinear(mr, IntegDomain(fes, GaussRule(2, 2), H), m)

    region1 = FDataDict("femm"=>femm)
    modeldata =  FDataDict("fens"=> fens, "regions"=>  [region1],  "essential_bcs"=>[e1, e2, e4])

    modeldata["load_multipliers"] = (1:nincr)./nincr*1.0;
    modeldata["maxdu_tol"] = maxdu_tol;
    modeldata["maxiter"] = maxiter;
    modeldata["line_search"]  = true;
    modeldata["iteration_observer"] = (lambda, iter, du, modeldata) -> begin
        @show lambda, iter, modeldata["maxdu"], modeldata["maxbal"]
        # print("-\n")
        # @show modeldata["un1"]
        # vtkexportmesh("junk$(iter).vtk", fens, fes; vectors = [("u", modeldata["un1"].values)])
        # print("-\n")
    end
    Ux = FFlt[]; Rx = FFlt[]
    modeldata["increment_observer"] = (lambda, modeldata) -> begin
        push!(Ux, mean(modeldata["un1"].values[movel1,1]));
        push!(Rx, sum(modeldata["reactions"].values[movel1,1]));
    end

    modeldata = nonlinearstatics(modeldata);
    # @show Ux, Rx
    @show (minimum(Ux), minimum(Rx), maximum(Ux), maximum(Rx))

    pl = lineplot((L .+ Ux) ./ L, Rx, canvas = DotCanvas)
    display(pl)
end # function neohookeanad_q4

function allrun()
    println("#####################################################")
    println("# neohookeanad_q4 ")
    neohookeanad_q4()
    return true
end # function allrun

end # module tension_compression_examples


#=
Author : Ng Yeesian
Modified : Joshua Ong / Guy Farmer
generates the stochastic deployment model
=#
struct StochasticDeployment_rem <: DeploymentModel
    m::JuMP.Model
    x::Vector{JuMP.VariableRef}
    y::Array{JuMP.VariableRef,3}
    slack::Vector{JuMP.VariableRef}
end
deployment(m::StochasticDeployment_rem) = [round(Int,x) for x in JuMP.value.(m.x)]

function StochasticDeployment_rem(p::DeploymentProblem; nperiods=params.nperiods, tol=params.δ,
    solver=Gurobi.Optimizer(OutputFlag=0), extra_amb = 1)

    old_x = [1, 1, 1, 2, 0, 0, 1, 2, 2, 0, 2, 0, 0, 1, 2, 2, 1, 2, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1, 2, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 1, 0, 2, 1]
    nperiods = min(length(p.train), nperiods)
    demand = p.demand[p.train,:]
    I = 1:p.nlocations
    J = 1:p.nregions
    T = 1:nperiods

    m = Model(GLPK.Optimizer) #using GLPK
    JuMP.@variable(m, x[1:p.nlocations] >= 0, Int)
    JuMP.@variable(m, y[1:p.nlocations,1:p.nregions,1:nperiods] >= 0, Int)
    JuMP.@variable(m, z[1:p.nregions,1:nperiods] >= 0, Int)
    JuMP.@variable(m, slack[1:p.nlocations] <= 0, Int) #slack variable

    JuMP.@objective(m, Min, sum(z[j,t] for j=J, t=T) + tol*sum(y[i,j,t] for i=I, j=J, t=T));

    JuMP.@constraint(m, sum(x[i] for i=I) <= p.nambulances - extra_amb);

    #ensure new solution is within slack of the old x.
    for i in I #collect(1:44)
        JuMP.@constraint(m, x[i] - slack[i] == old_x[i]);
    end
    JuMP.@constraint(m, sum(slack[i] for i in I) == extra_amb) #make sure slack is the +/- n ambulances of the old solution

    #Guarantees coverage for all regions
    for j in J
        JuMP.@constraint(m, sum(x[i] for i in filter(i->p.coverage[j,i], I)) >= 1);
    end

    # flow constraints at each station
    for i in I, t in T
        outflow = JuMP.@expression(m, sum(y[i,j,t] for j in filter(j->p.coverage[j,i], J)));
        JuMP.@constraint(m, x[i] >= outflow);
    end

    # shortfall from satisfying demand/calls
    for j in J, t in T
        inflow = JuMP.@expression(m, sum(y[i,j,t] for i in filter(i->p.coverage[j,i], I)));
        JuMP.@constraint(m, z[j,t] >= demand[t,j] - inflow);
    end

    # any location can only house so many hospitals
    for i in I
        JuMP.@constraint(m, x[i] <= 5);
    end

    StochasticDeployment_rem(m, x, y, slack)
end

#status = JuMP.optimize!(model.m)
optimize!(model::StochasticDeployment_rem) = JuMP.optimize!(model.m)
optimize!(model::StochasticDeployment_rem, p::DeploymentProblem) = JuMP.optimize!(model.m)


using Asap, JSON, Statistics

# Taken from NCCR compas.dev example files
test1 = JSON.parsefile("examples/compas_fea_truss_frame.json")

# nodes (in mm)
ns = test1["node_list"]

# elements
es = test1["element_list"]

# material properties
mp = test1["material_properties"]
E = Float64(mp["youngs_modulus"]) # kn/cm2
A = mp["cross_sec_area"] #cm^2
G = Float64(mp["shear_modulus"]) #kN/cm^2
Iz = mp["Iz"] #cm^4
Iy = mp["Iy"] # cm^4
J = mp["Jx"] #cm^4

# Create nodes
nodes = Vector{Node}()

for i = 1:length(ns)
    n = ns[i]
    position = [n["point"]["X"], n["point"]["Y"], n["point"]["Z"]] ./ 10

    if n["is_grounded"] > 0
        fixtype = :fixed
    else
        fixtype = :free
    end
    push!(nodes, Node(position, :frame, fixtype))
end

# Create elements
elements = Vector{Element}()
for e in es
    indices = e["end_node_ids"] .+ 1
    push!(elements, Element(nodes, indices, E, A, G, Iz, Iy, J))
end

# Create self-weight loads
lengths = [e.length for e in elements]
P = mean(lengths) * A * mp["density"]/100^3 / 2
Pvec = [0, 0, -P, 0, 0, 0]

loads = Vector{Load}()
for node in nodes
    push!(loads, Load(nodes, node.position, Pvec))
end

structure = Structure(nodes, elements, loads)

analyze(structure)

geo = Geometry(structure)
structurePlot(geo; scaleToForce = true)


function policy_opt(p ;sⁱ = false, sᶠ = false, sᴱ = false)
  poltype = [sⁱ, sᶠ, sᴱ]
  npol    = sum(poltype)
  polinit = ones(npol)*.05

# first solve without policy
  eqInit = [ 2.00, 0.55, 0.06,
             0.72, 0.83, 1.73]

  eq,res = solveBGP(p,eqInit)
  eqInit.= res.zero

  cevMax = [calCev(eq.g, eq.cactivtot,eq.g, eq.cactivtot,p)]
  @printf("OPTIMAL POLICY\n")
  @printf("----------------------------\n")
  @printf("Baseline Welfare is %3.4f \n", cevMax[1])

  # optimization
  # ------------------------------------------------------------------------------------------------
  opt = Opt(:LN_NELDERMEAD, npol)
  ftol_rel!(opt,1e-8)
  min_objective!(opt, (x, grad) -> policy_obj(x, poltype, p, eqInit, cevMax, eq.g, eq.cactivtot))
  minf,minpol,ret = NLopt.optimize(opt, polinit)
  nEvals = opt.numevals
  #-------------------------------------------------------------------------------------------------

  if ret == :FTOL_REACHED
  	@printf("Optimal policy found!\n\n")
  	optpolval = zeros(3)
    optpolval[poltype] .= minpol
	res = Dict(:sⁱ=> optpolval[1], :sᶠ=>optpolval[2], :sᴱ=>optpolval[3], :cev=>-minf, :ret=>ret, :opt=>opt, :nEvals=>nEvals)
  	return res
  else
	@printf("Optimal policy could not solved!\n\n")
	return ret
  end

end

function policy_obj(polguess, poltype, p, eqInit, cevMax, gBase, cactivtotBase)

  polval = zeros(3)
  polval[poltype] .= polguess

  pNew = Params(p, (sⁱ = polval[1], sᶠ = polval[2], sᴱ = polval[3]))

  eq, res = solveBGP(pNew, eqInit)

  if res.f_converged
  	eqInit .= res.zero
  	# welfare calculation
  	cev = calCev(gBase, cactivtotBase, eq.g, eq.cactivtot, pNew)
  else
    cev = -Inf
  end
  if cev>cevMax[1]
  	@printf("Welfare improved to %3.4f (", cev)
	poltype[1] ? @printf("sⁱ = %2.2f", pNew.sⁱ) : nothing
	poltype[1]&(poltype[2] | poltype[3]) ? @printf(",") : nothing
	poltype[2] ? @printf("sᶠ = %2.2f", pNew.sᶠ) : nothing
	poltype[2]&poltype[3] ? @printf(",") : nothing
	poltype[3] ? @printf("sᴱ = %2.2f", pNew.sᴱ) : nothing
	@printf(")\n")
	cevMax[1] = cev
  end
  return -cev
end

function calCev(gBase, cactivtotBase, gPol, cactivtotPol, p)
#	welf  = (1.0/(1.0-p.σ))*((eq.cactivtot^((1.0 - p.σ)/(p.ε-1.0))/(p.ρ-(1.0-p.σ)*eq.g))-(1.0/p.ρ))
	cev  = ((p.ρ-(1.0-p.σ)*gBase)/(p.ρ-(1.0-p.σ)*gPol))^(1.0/(1.0-p.σ))*((cactivtotPol/cactivtotBase)^(1.0/(p.ε-1.0)));
end


module BitsTallies

import Base: getindex, +, -

const BITS = 5
const mask = UInt8(2 ^ BITS - 1)
const offsets = [BITS * i for i in 0:25]
const charmasks = [UInt128(1) << x for x in offsets]

"""
The BitsTally uses a single unsigned integer to store counts
for every letter in 'a':'z'. This lets us compute those tallies
without allocating an array, and compare them using simple
integer equality.

The way we do this is to divide up the internal 128-bit integer into 26
blocks. Each block is 5 bits long, except the final block which is only 3 bits
long. We can use each of those blocks to store a 5-bit count of a given
letter. We store the counts for the letter 'a' in the low-order bits, which
leaves the awkward 3-bit block for 'z'. So this will overflow if a word
contains more than 7 'z's. or more than 31 of any other letter.
"""
struct BitsTally
    data::UInt128

    BitsTally(data=0) = new(data)
end

"""
Add a single count for char c to the tally t.
"""
+(t::BitsTally, c::Char) = BitsTally(t.data + charmasks[convert(Int, c - 'a' + 1)])

-(t::BitsTally, c::Char) = BitsTally(t.data - charmasks[convert(Int, c - 'a' + 1)])

function BitsTally(s::String)
    t = BitsTally()
    for c in s
        t += c
    end
    t
end

"""
Strip the appropriate bits out of the storage UInt128 to return the count of
the given letter as a UInt8.
"""
getindex(t::BitsTally, c::Char) = mod((t.data >> offsets[convert(Int, c - 'a' + 1)]) & mask, UInt8)

isanagram(t1::BitsTally, t2::BitsTally) = t1.data == t2.data

"""
Returns true iff t1 is a transaddition of t2. That is, the letters of t2
can be rearranged along with one extra letter to make the letters of t1.
"""
function istransaddition(t1::BitsTally, t2::BitsTally)
    d = t1.data - t2.data
    for m in charmasks
        if d == m
            return true
        end
    end
    false
end

end


using Test

using LinearAlgebra, Combinatorics, Statistics
using OrdinaryDiffEq, ForwardDiff, DiffResults
using DoubleFloats

include("../src/type.jl")
include("testfunctions.jl")

include("../src/metrics/Minkowski.jl")
include("../src/metrics/KerrSchild.jl")
include("../src/geosol.jl")
include("../src/broyden.jl")
include("../src/outlier.jl")
include("../src/squirrel.jl")

η = X -> ημν(typeof(X[1]),4)

@testset "All tests:" begin

    @time @testset "Minkowski tests:" begin
            include("minkowski_test.jl")
        end

    @time @testset "geosol tests:" begin
            include("geosol_test.jl")
        end

    @time @testset "broyden tests:" begin
            include("broyden_test.jl")
        end

    @time @testset "outlier detection tests:" begin
            include("outlier_test.jl")
        end

    @time @testset "squirrel tests:" begin
            include("squirrel_test.jl")
        end

    @time @testset "seval tests:" begin
            include("seval_test.jl")
    end

end

nothing


using ConvDiffMIPDECO
using jInv.Mesh
using jInv.ForwardShare
using jInv.InverseSolve
using jInv.LinearSolvers
using MUMPSjInv
using MAT
# using jInvVis
using LinearAlgebra

# read results from l-curve
dataset    = "Peaks2D"
reg        = "wTVReg"
m          = [256 128]
noiseLevel = 0.1

filename= joinpath(dirname(pathof(ConvDiffMIPDECO)),"..","examples","Peaks2D.mat")
data = matread(filename)
domain = data["domain"]
v      = data["v"]
sig    = data["sig"]
dtrue  = data["dtrue"]
rec    = data["rec"];

M = getRegularMesh(domain,m)

resFile = joinpath(dirname(pathof(ConvDiffMIPDECO)),"..","examples",
                    "$(dataset)-$(reg)-noise-$(noiseLevel)-$(m[1])x$(m[2])-penalty.mat")
res = matread(resFile)
dobs = res["dobs"]

srcRelaxed   = res["SourcesRelaxed"]
alphaRelaxed = res["alphaRelaxed"]
Ms           = res["Ms"]

println("solve using alpha = $(alphaRelaxed)")
noiseLevel = norm(dobs-dtrue)/norm(dtrue)
println("estimated noise level: $noiseLevel")
if noiseLevel > 0.5
    error("noise level too large, double check that data was loaded correctly. ")
end

# build inverse problem
x1,x2 = getNodalAxes(M)
P = interpmat(x1,x2, rec);
pFor = getConvDiffFEMParam(M,v=v,sig=sig,P=P,Ainv=getMUMPSsolver());

## configure misfit
Wt         = ones(size(dobs))
sigback    = 0.0
pMis       = getMisfitParam(pFor,Wt,dobs,SSDFun)

## Configure regularization
mref       = zeros(M.nc)
reg        = (m,mr,M,I=1.0) -> wTVReg(m,mr,M,eps=1e-8)

## Configure optimization
maxIter    = 50
minUpdate  = 1e-3
HesPrec    = getSSORRegularizationPreconditioner(1.0,1e-15,50)
cgit       = 5
pcgTol     = 1e-1
modFun     = identityMod
boundsLow  = 0*ones(M.nc)
boundsHigh = 1*ones(M.nc)
maxStep	   = 0.1*maximum(boundsHigh)

## store the configuration
pInv       = getInverseParam(M,modFun,reg,alphaRelaxed,mref,
                             boundsLow,boundsHigh,maxStep=maxStep,
                            pcgMaxIter=cgit,pcgTol=pcgTol,minUpdate=minUpdate,maxIter=maxIter,
                            HesPrec=HesPrec);

neighborhoods = (s->trues(size(s)), s->dilation(s,pInv.MInv,1), s->dilation(s,pInv.MInv,1,0.1,0.9))
neighborhoodsNames = ["all","dilation1", "dilation2"];

nrow = size(Ms,2)
ncol = length(neighborhoods)

results = zeros(m[1],m[2],nrow,length(neighborhoods))
init = zeros(m[1],m[2],nrow,length(neighborhoods))
His = zeros(pInv.maxIter,3,nrow,length(neighborhoods))
times = zeros(nrow,length(neighborhoods))

println("solve using alpha = $(pInv.alpha)")

for k1=1:nrow
    for k2=1:length(neighborhoods)
        println("\n --- mipdecoHeuristic starting with Ms[:,$k1] using $(neighborhoodsNames[k2])--")
        (t0,nsolve0) = (pFor.Ainv.solveTime, pFor.Ainv.nSolve)

        src0 = naiveRounding(Ms[:,k1])
        init[:,:,k1,k2] = src0

        times[k1,k2] = @elapsed begin
            mcTR,DcTR,flagTR,his = mipdecoHeuristic(src0,pInv,pMis,getNeighborhood=neighborhoods[k2],out=0)
            results[:,:,k1,k2] = mcTR
            His[:,:,k1,k2] = his;
		end
        nIter = findlast(His[:,1,k1,k2].>0)

        # p1 = viewImage2D(Ms[:,k1],M)
        # title!("relaxed solution")
        # plot!(rec[:,1],rec[:,2],seriestype=:scatter,legend=false)
        # p2 = viewImage2D(1.0.*src0,M)
        # title!("starting guess")
        # p3 = viewImage2D(1.0.*mcTR,M)
        # title!("MIPDECO result")
        # p4 = viewImage2D(1.0.*(mcTR-src0),M)
        # title!("update")
        # pt = plot(p1,p2,p3,p4)
        # display(pt)

        println("\t\t initial obj. :\t\t$(His[1,1,k1,k2]) ")
        println("\t\t final obj. :\t\t$(His[nIter,1,k1,k2]) ")
        println("\t\t relative :\t\t\t$(His[nIter,1,k1,k2]/His[1,1,k1,k2]) ")
        println("\t\t no. PDE solves:\t$(2*nIter) ")
        println("\t\t runtime: \t\t\t$(times[k1,k2])")
    end
end

res["SourcesMIPDECOPenalty"]=results
res["InitMIPDECOPenalty"]=init
res["HisMIPDECOPenalty"] = His
res["timesMIPDECOPenalty"] = times
matwrite(resFile,res)


# Code for time and frequency response of LTI systems
using LinearAlgebra
# include("Utilities.jl");

# Frequency response at user specified frequencies
function FrequencyResponse(sys::StateSpace,no::Integer,ni::Integer, Om::Array{Float64,1})::Vector{ComplexF64}
    CheckIODimensions(sys,no,ni);

    B = sys.B[:,ni];
    C = reshape(sys.C[no,:],1,length(sys.C[no,:])); # Must ensure row vector
    D = sys.D[no,ni];
    I = im*eye(size(sys.A)[1]);
    TF = Vector{ComplexF64}(undef,length(Om));
    for (i,om) in enumerate(Om)
        G = C*inv(om*I-sys.A)*B; # This is vector of one element
        TF[i] = G[1]+D; # that is why we need G[1], otherwise the addition is across different types.
    end
    return(TF);
end

# Frequency response with no user specified om. This makes an adhoc grid between [1E-3 100*wn.max], with 300 points.
function FrequencyResponse(sys::StateSpace,no::Integer,ni::Integer)::Tuple{Array{Complex{Float64},1},Array{Float64,1}}

    CheckIODimensions(sys,no,ni);

    l,ev = eigen(sys.A);
    lmax = maximum(abs.(l)); # Find highest natural frequency
    Om = logspace(1E-3,lmax*100,300); # Making adhoc decision about 100 times highest natural frequency.

    B = sys.B[:,ni];
    C = reshape(sys.C[no,:],1,length(sys.C[no,:])); # Must ensure row vector
    D = sys.D[no,ni];
    I = im*eye(size(sys.A)[1]);
    TF = Vector{ComplexF64}(undef,length(Om));
    for (i,om) in enumerate(Om)
        G = C*inv(om*I-sys.A)*B; # This is vector of one element ...
        TF[i] = G[1]+D; # ... that is why we need G[1], otherwise the addition is across different types.
    end
    return(TF,Om);
end

function StepResponse(sys::StateSpace,no::Integer,ni::Integer)::Tuple{Array{Float64,1}, Array{Float64,1}}
    l,ev = eigen(sys.A);
    lmin = minimum(abs.(l)); # Find the slowest mode
    lmax = maximum(abs.(l)); # Find the fastest mode
    tmax = 2.5*(2*pi)/lmin; # Quite adhoc -- 10 times the slowest time-constant.
    dt =  (2*pi)/lmax/50; # Quite adhoc -- Sampling is 50 times the fastest mode.
    t = Vector(0:dt:tmax);
    u = ones(length(t));
    y = ForcedResponse(sys,no,ni,t,u);
    return(t,y);
end

function StepResponse(sys::StateSpace,no::Integer,ni::Integer,t::Vector{Float64})::Vector{Float64}
    u = ones(length(t));
    y = ForcedResponse(sys,no,ni,t,u);
    return(y);
end

function ForcedResponse(sys::StateSpace,no::Integer,ni::Integer,t::Vector{Float64},u::Vector{Float64},x0::Vector{Float64}=Array{Float64}(undef,0))::Vector{Float64}
    CheckIODimensions(sys,no,ni);

    # Obtain a discrete-time system.
    dt = t[2]-t[1]; # Assumed to be uniform dt -- probably should put a check in here.
    ns = size(sys.A)[2];

    # We are only looking at SISO time response
    B = sys.B[:,ni];
    C = reshape(sys.C[no,:],1,length(sys.C[no,:])); # Must ensure row vector
    D = sys.D[no,ni];

    M = dt*[sys.A B; zeros(1,ns+1)]
    S = exp(M);
    F = S[1:ns,1:ns];
    G = S[1:ns,ns+1:ns+1];

    # Check initial condition
    if isempty(x0)
        x = zeros(ns);
    else
        x = x0;
    end

    # Time march
    y = similar(t);
    for i = 1:length(t)
        y[i] = (C*x)[1] + D*u[i];
        x = F*x + G*u[i];
    end
    return(y);
end

function ImpulseResponse(sys::StateSpace,no::Integer,ni::Integer)::Tuple{Array{Float64,1}, Array{Float64,1}}
    l,ev = eigen(sys.A);
    lmin = minimum(abs.(l)); # Find the slowest mode
    lmax = maximum(abs.(l)); # Find the fastest mode
    tmax = 2.5*(2*pi)/lmin; # Quite adhoc -- 2.5 times the slowest time-constant.
    dt =  (2*pi)/lmax/50; # Quite adhoc -- Sampling is 50 times the fastest mode.
    t = Vector(0:dt:tmax);
    y = ImpulseResponse(sys,no,ni,t);
    return(t,y);
end

function ImpulseResponse(sys::StateSpace,no::Integer,ni::Integer,t::Vector{Float64})::Vector{Float64}
    y = similar(t);
    B = sys.B[:,ni];
    C = reshape(sys.C[no,:],1,length(sys.C[no,:])); # Must ensure row vector
    for i=1:length(t)
        y[i] = (C*exp(sys.A*t[i])*B)[1];
    end
    return(y);
end

function InitialResponse(sys::StateSpace,x0::Vector{Float64})::Tuple{Vector{Float64}, Matrix{Float64}}
    l,ev = eigen(sys.A);
    lmin = minimum(abs.(l)); # Find the slowest mode
    lmax = maximum(abs.(l)); # Find the fastest mode
    tmax = 2.5*(2*pi)/lmin; # Quite adhoc -- 2.5 times the slowest time-constant.
    dt =  (2*pi)/lmax/50; # Quite adhoc -- Sampling is 50 times the fastest mode.
    t = Vector(0:dt:tmax);
    y = InitialResponse(sys,x0,t);
    return(t,y);
end

function InitialResponse(sys::StateSpace,x0::Vector{Float64},t::Vector{Float64})::Matrix{Float64}
    y = Matrix{Float64}(undef,length(t),length(x0));
    for i=1:length(t)
        y[i,:] = reshape(exp(sys.A*t[i])*x0,1,length(x0));
    end
    return(y);
end


struct FDHeatProblem{TBCL, TBCR,TTS}
    a # (a,b) define the domain
    b
    Δx # mesh spacing
    nx # number of points in the interior, excluding a and b
    x # mesh on which we solve the heat equation
    Δt # time step
    nΔt # number of time steps
    tmax # final time
    u0 # Initial condition
    A # the matrix (default to sparse)
    rhs # right hand side of our system
    left_bc::TBCL
    right_bc::TBCR
    time_step::TTS
end

struct DirichletBC
    u_bc
end

struct NeumannBC
    v_bc
end

struct PeriodicBC
end

struct BEuler
end


points(H::Hermite,n)=gausshermite(n)[1]

plan_transform(H::Hermite,v::Vector) = gausshermite(length(v))
plan_itransform(H::Hermite,cfs::Vector) = points(H,length(cfs))
function transform(H::Hermite,vals,plan::Tuple{Vector,Vector})
    x,w = plan
    V=hermitep(0:length(vals)-1,x)'
    nrm=(V.^2)*w

    V*(w.*vals)./nrm
end
itransform(H::Hermite,cfs,plan::Vector) = hermitep(0:length(cfs)-1,tocanonical(H,plan))*cfs


function poly_eval(p::MP.AbstractPolynomial{JuMP.AffExpr},
                   a::AbstractVector{Float64})
    vars = MP.variables(p)
    aff = zero(JuMP.AffExpr)
    for term in MP.terms(p)
        mono = MP.monomial(term)
        aff = MA.add_mul!!(aff, mono(vars => a), MP.coefficient(term))
    end
    return aff
end
function sublevel_eval(ell::Sets.Ellipsoid,
                       a::AbstractVector)
    return quad_form(ell.Q, a)
end
function sublevel_eval(set::Union{Sets.PolySet,
                                  Sets.ConvexPolySet},
                       a::AbstractVector)
    return poly_eval(MP.polynomial(set.p), a)
end
function sublevel_eval(set::Sets.Householder, a::AbstractVector, β)
    x = Sets.space_variables(set)
    z = Sets.perspective_variable(set)
    # Avoid large values, with high degree polynomials, it might cause issues
    if false
        scale_a = norm(a)
        scale_β = abs(β)
        scaling = max(scale_a, scale_β) / sqrt(scale_a * scale_β)
    else
        scaling = 1.0
    end
    return Sets.perspective_gauge0(set)(z => β / scaling, x => a / scaling)
end


using Random, Test, FFTW
using FourierTools
using ImageTransformations
using IndexFunArrays
using Zygote
using NDTools

Random.seed!(42)

include("fft_helpers.jl")
include("fftshift_alternatives.jl")
include("utils.jl")
include("fourier_shifting.jl")
include("fourier_shear.jl")
include("fourier_rotate.jl")
include("resampling_tests.jl")
include("convolutions.jl")
include("custom_fourier_types.jl")
include("damping.jl")
include("czt.jl")

return


include("./ols.jl")

data = readcsv("data/gasoline.csv")

# Get the columns that correspond with our X values
# and create a new matrix to hold them
# syntax: data[row, column]
# to get a whole column do data[:, column]
# to get a whole row do data[row, :]
x = [data[:,2] data[:,3] data[:,4]]
y = data[:,6]

# Create a new OLS object
reg = ols.tols(y, x, "Octane Rating", ["Error", "Component 1", "Component 2", "Component 3"])


"""
NueralNetwork module

MIT License
Copyright (c) 2019 amenoyoya
"""

# 数値微分関数
numeric_diff(f, x) = begin
    h = 1e-4 # 10^(-4)
    (f(x + h) - f(x - h)) / 2h # 中心差分から微分値を計算
end

# 数値微分による勾配計算
## 関数 f(Float64配列), Float64配列 x -> Float64配列 勾配
numeric_gradient(f, x::Array{Float64,1})::Array{Float64,1} = begin
    h = 1e-4 # 10^(-4)
    grad = Array{Float64, 1}(undef, length(x)) # xと同じ長さの配列 [undef, undef, ...] を生成
    # 各変数ごとの数値微分を配列にまとめる
    for i in 1:length(x)
        # 指定indexの変数に対する中心差分を求める
        org = x[i]
        x[i] = org + h
        f1 = f(x) # f([..., x[i] + h, ...])
        x[i] = org - h
        f2 = f(x) # f([..., x[i] - h, ...])
        grad[i] = (f1 - f2) / 2h # i番目の変数に対する数値微分
        x[i] = org # x[i]の値をもとに戻す
    end
    return grad
end

# 数値微分による勾配計算（バッチ対応版）
## 関数(Array{Float64,2})::Float64, 入力値::Array{Float64,2} -> 勾配::Array{Float64,2}
numeric_gradient(f, x::Array{Float64,2})::Array{Float64,2} = begin
    h = 1e-4 # 10^(-4)
    grad = Array{Float64, 2}(undef, size(x, 1), size(x, 2)) # xと同じ次元の行列を生成
    # 各変数ごとの数値微分を行列にまとめる
    for row in 1:size(x, 1), col in 1:size(x, 2)
        # 指定indexの変数に対する中心差分を求める
        org = x[row, col]
        x[row, col] = org + h
        f1 = f(x) # f([..., x[row, col] + h, ...]) -> Float64
        x[row, col] = org - h
        f2 = f(x) # f([..., x[row, col] - h, ...]) -> Float64
        grad[row, col] = (f1 - f2) / 2h # (row, col)番目の変数に対する数値微分
        x[row, col] = org # x[i]の値をもとに戻す
    end
    return grad
end

# 勾配降下法
## 損失関数 f, ベクトル x, 学習率 lr, 試行回数 ste_num -> 最適化ベクトル
gradient_descent(f, x::Array{Float64,1}, lr::Float64=0.01, step_num::Int=100)::Array{Float64,1} = begin
    step_num > 0 ? gradient_descent(f, x - lr * numeric_gradient(f, x), lr, step_num - 1) : x
end


include("../../utils.jl")

function pollardfactor(n::T=big(1208925819691594988651321)) where T<:Integer
    for c in T(1):(n - 3)
        G, r, q = ones(T,3)
        y = 2
        m::T = 1900
        ys::T = 0
        x::T = 0
        while G == 1
            x = y
            for i in 1:r
                y = (y^2 + c) % n
            end
            k = T(0)
            G = T(1)
            while k < r && G == 1
                for i in 1:min(r - k, m)
                    ys = y
                    y = (y^2 + c) % n
                    q = (q * abs(x - y)) % n
                end
                G = gcd(q, n)
                k += m
            end
            r *= 2
        end
        G == n && (G = T(1))
        while G == 1
            ys = (ys^2 + c) % n
            G = gcd(abs(x - ys), n)
        end
        if G != n
            return G
        end
    end
end

@gctime pollardfactor()


# estimate GARCH(1,1) with AR(1) in mean
using Statistics, DelimitedFiles, Econometrics
function main()
include("garch11.jl")
# weekly close price of NSYE, data provided with GRETL
data = readdlm("nysewk.txt")
# compute weekly percentage growth
y = 100.0 * log.(data[2:end] ./ data[1:end-1])
# Constrained maximization of logL
# note that the objective has a minus sign in front, as fmincon
# minimizes, but we want to maximize the logL
# get constrained estimates to use as input for MLE
thetastart = [mean(y); 0.0; var(y); 0.1; 0.1]
obj = theta -> -sum(garch11(theta, y))
thetahat, logL, junk  = fmincon(obj, thetastart, [], [], [-Inf, -1.0, 0.0, 0.0, 0.0], [Inf, 1.0, Inf, 1.0, 1.0])
# now do MLE to see formatted results.
# NOTE TO SELF: this won't work when the constraints are binding
# should add a method to mle for constrained problems.
model = theta -> garch11(theta, y)
thetahat, logL, junk, converged = mleresults(model, thetahat, "GARCH(1,1) example")
return
end
main()


@node Categorical Stochastic [ out, p ]

@average_energy Categorical (q_out::Categorical, q_p::Any) = -sum(probvec(q_out) .* mean(log, q_p))


immutable InverseGaussian <: ContinuousUnivariateDistribution
    μ::Float64
    λ::Float64

    function InverseGaussian(μ::Real, λ::Real)
        (μ > zero(μ) && λ > zero(λ)) || error("InverseGaussian's μ and λ must be positive")
        new(float64(μ), float64(λ))
    end

    InverseGaussian(μ::Real) = InverseGaussian(μ, 1.0)
    InverseGaussian() = new(1.0, 1.0)
end

@distr_support InverseGaussian 0.0 Inf

#### Parameters

shape(d::InverseGaussian) = d.λ
params(d::InverseGaussian) = (d.μ, d.λ)

#### Statistics

mean(d::InverseGaussian) = d.μ

var(d::InverseGaussian) = d.μ^3 / d.λ

skewness(d::InverseGaussian) = 3.0 * sqrt(d.μ / d.λ)

kurtosis(d::InverseGaussian) = 15.0 * d.μ / d.λ

function mode(d::InverseGaussian)
    μ, λ = params(d)
    r = μ / λ
    μ * (sqrt(1.0 + 2.25 * r^2) - 1.5 * r)
end

#### Evaluation

function pdf(d::InverseGaussian, x::Float64)
    if x > 0.0
        μ, λ = params(d)
        return sqrt(λ / (twoπ * x^3)) * exp(-λ * (x - μ)^2 / (2.0 * μ^2 * x))
    else
        return 0.0
    end
end

function logpdf(d::InverseGaussian, x::Float64)
    if x > 0.0
        μ, λ = params(d)
        return 0.5 * (log(λ) - (log2π + 3.0 * log(x)) - λ * (x - μ)^2 / (μ^2 * x))
    else
        return -Inf
    end
end

function cdf(d::InverseGaussian, x::Float64)
    if x > 0.0
        μ, λ = params(d)
        u = sqrt(λ / x)
        v = x / μ
        return Φ(u * (v - 1.0)) + exp(2.0 * λ / μ) * Φ(-u * (v + 1.0))
    else
        return 0.0
    end
end

function ccdf(d::InverseGaussian, x::Float64)
    if x > 0.0
        μ, λ = params(d)
        u = sqrt(λ / x)
        v = x / μ
        Φc(u * (v - 1.0)) - exp(2.0 * λ / μ) * Φ(-u * (v + 1.0))
    else
        return 1.0
    end
end

function logcdf(d::InverseGaussian, x::Float64)
    if x > 0.0
        μ, λ = params(d)
        u = sqrt(λ / x)
        v = x / μ
        a = logΦ(u * (v -1.0))
        b = 2.0 * λ / μ + logΦ(-u * (v + 1.0))
        a + log1pexp(b - a)
    else
        return -Inf
    end
end

function logccdf(d::InverseGaussian, x::Float64)
    if x > 0.0
        μ, λ = params(d)
        u = sqrt(λ / x)
        v = x / μ
        a = logΦc(u * (v - 1.0))
        b = 2.0 * λ / μ + logΦ(-u * (v + 1.0))
        a + log1mexp(b - a)
    else
        return 0.0
    end
end

# TODO: need a more accurate method
function quantile(d::InverseGaussian, p::Float64)
    if 0.0 < p < 1.0
        # Whitmore and Yalovsky (1978) approximation
        μ, λ = params(d)
        w = Φinv(p)
        φ = λ / μ
        x = μ * exp(w / sqrt(φ) - 0.5 / φ)
        # one multiplicitive Newton iteration
        xn = x * exp((p - cdf(d, x)) / (pdf(d, x) * exp(x)))
        # additive Newton iterations
        while !isapprox(xn,x)
            x = xn
            xn = x + (p - cdf(d, x)) / pdf(d, x)
        end
        return xn
    elseif p <= 0.0
        return 0.0
    else
        return Inf
    end
end

#### Sampling

# rand method from:
#   John R. Michael, William R. Schucany and Roy W. Haas (1976)
#   Generating Random Variates Using Transformations with Multiple Roots
#   The American Statistician , Vol. 30, No. 2, pp. 88-90
function rand(d::InverseGaussian)
    μ, λ = params(d)
    z = randn()
    v = z * z
    w = μ * v
    x1 = μ + μ / (2.0 * λ) * (w - sqrt(w * (4.0 * λ + w)))
    p1 = μ / (μ + x1)
    u = rand()
    u >= p1 ? μ^2 / x1 : x1
end