text:
push!(LOAD_PATH, "../src")

using GrammaticalEvolution
import GrammaticalEvolution.evaluate!

include("ExamplePopulation.jl")

# ground truth we're matching against
gt(x, y) = 2*x + 5*y

# action for creating
convert_number(lst) = float(join(lst))

function create_grammar()
  @grammar grammar begin
    start = ex
    ex = number | sum | product | (ex) | value
    sum = Expr(:call, :+, ex, ex)
    product = Expr(:call, :*, ex, ex)
    value = :x | :y | number
    number[convert_number] = digit + '.' + digit
    digit = 0:9
  end

  return grammar
end

function GrammaticalEvolution.evaluate!(grammar::Grammar, ind::ExampleIndividual)
  fitness::Array{Float64, 1} = []

  try
    ind.code = transform(grammar, ind)
    @eval fn(x, y) = $(ind.code)
  catch e
    # println("exception = $e")
    ind.fitness = Inf
    return
  end

  for x=0:10
    for y=0:10
      value = fn(x, y)
      diff = (value - gt(x, y)).^2
      if !isnan(diff) && diff > 0
        insert!(fitness, length(fitness)+1, sqrt(diff))
      elseif diff == 0
        insert!(fitness, length(fitness)+1, 0)
      end
    end
  end

  ind.fitness = mean(fitness)
end

function main()
  # our grammar
  grammar = create_grammar()

  # create population
  pop = ExamplePopulation(500, 100)

  fitness = Inf
  generation = 1
  while fitness > 1.0
    # generate a new population (based off of fitness)
    pop = generate(grammar, pop, 0.1, 0.2, 0.2)

    # population is sorted, so first entry it the best
    fitness = pop[1].fitness
    println("generation: $generation, max fitness=$fitness, code=$(pop[1].code)")
    generation += 1
  end
end

main()

max_stars_repo_name: abeschneider/GrammaticalEvolution

max_issues_repo_name: abeschneider/GrammaticalEvolution

max_forks_repo_name: abeschneider/GrammaticalEvolution


text:
using GNSSTools
using Statistics
using Base.Threads
using ProgressMeter
using JLD
using FFTW
using PyPlot
pygui(true)
include("ch4_constellation_hist_structs.jl")

# CN0s = [30, 35, 40, 45]  # dB⋅Hz
# h_parms = [1.9e-6, 2e-6, 9.7e-8, 9.1e-11, 9.9e-15]./16.3676e6^2  # TCXO
# h_parms = [2.2e-9, 1.6e-10, 1.6e-10, 7.1e-13, 7.2e-16]./10e6^2   # OCXO

# f_s = 5e6  # Hz
# sig_freq = L1_freq  # Hz
# f_d = 0.  # Hz
# f_if = 0.  # Hz
# fd_rates = 0:-10:-350  # Hz/s
# integration_times = range(1e-3, 20e-3, step=1e-3)
# t_length = maximum(integration_times)  # seconds
# iterations = 1
# prn = 1

# sigtype = define_l1ca_code_type(;sig_freq=sig_freq)
# signal = definesignal(sigtype, f_s, t_length; prn=prn, f_d=f_d, f_if=f_if,
#                       include_phase_noise=false)

# results = Array{Float64}(undef, length(CN0s), length(fd_rates), length(integration_times))

# p = Progress(iterations*length(CN0s)*length(fd_rates)*length(integration_times), 1, "Processing...")
# for i in 1:length(CN0s)
#     for j in 1:length(fd_rates)
#         for k in 1:length(integration_times)
#             SNR = 0.
#             for iteration in 1:iterations
#                 CN0 = CN0s[i]
#                 fd_rate = fd_rates[j]
#                 T = integration_times[k]
#                 ϕ₀ = rand(0:0.0001:2π)
#                 code_start_idx = rand(1:signal.sample_num)
#                 definesignal!(signal; fd_rate=fd_rate, CN0=CN0,
#                               code_start_idx=code_start_idx, phi=ϕ₀,
#                               new_thermal_noise=true, new_phase_noise=false)
#                 generatesignal!(signal)
#                 # println([i,j,k])
#                 replica = definesignal(sigtype, f_s, T;
#                                        skip_noise_generation=true,
#                                        allocate_noise_vectors=false)
#                 fd_course, n0_est, SNR_est = courseacquisition(signal,
#                                                                replica,
#                                                                prn)
#                 SNR += 10^(SNR_est/10)
#                 next!(p)
#             end
#             results[i,j,k] = 10*log10(SNR/iterations)
#         end
#     end
# end

# save("ch4_acquisition_performance_l1ca.jld", "results", results)

# results = load("ch4_acquisition_performance.jld", "results")

# max_val = maximum(results)
# min_val = minimum(results)
# heatmap_extent = [1e3*integration_times[1], 1e3*integration_times[end], fd_rates[1], fd_rates[end]]
# fig = figure(figsize=(7.5,2))
# for i in 1:length(CN0s)
#     ax = fig.add_subplot(1, length(CN0s), i)
#     global h = ax.imshow(results[i,:,:], aspect="auto", vmin=min_val, vmax=max_val, extent=heatmap_extent)
#     xlabel("Integration Time (ms)")
#     ylabel("Doppler Rate (Hz/s)")
#     title("C/N₀: $(CN0s[i])dB⋅Hz")
# end
# cbar_ax = fig.add_axes([0.92, 0.15, 0.04, 0.7])
# cbar = fig.colorbar(h, cax=cbar_ax)
# cbar.set_label("SNR (dB)")
# subplots_adjust(hspace=0.3, wspace=0.4, top=0.91, left=0.05, right=0.9, bottom=0.12)
# savefig("figures/ch4_acquisition_performance.pdf", dpi=300)

################################################################################
CN0s = [45]  # dB⋅Hz
h_parms = [1.9e-6, 2e-6, 9.7e-8, 9.1e-11, 9.9e-15]./16.3676e6^2  # TCXO
# h_parms = [2.2e-9, 1.6e-10, 1.6e-10, 7.1e-13, 7.2e-16]./10e6^2   # OCXO

directory = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/"

integration_times = [1e-3, 2e-3, 5e-3, 10e-3]  # seconds
sigtype = define_l1ca_code_type(maximum(integration_times))
f_s = 2.5e6  # Hz
sig_freq = L1_freq  # Hz
file_name = "ch4_acquisition_performance_nonheatmap_l1ca"

# sigtype = define_l5_code_type(;channel="Q")
# integration_times = [1e-3, 20e-3, 40e-3, 100e-3]  # seconds
# f_s = 25e6  # Hz
# sig_freq = L5_freq  # Hz
# file_name = "ch4_acquisition_performance_nonheatmap_l5"

f_d = 0.  # Hz
f_if = 0.  # Hz
# fd_rate_max = 10e3 * (sig_freq/L1_freq)
fd_rate_max = 10e3 * (sig_freq/L1_freq)
fd_rates = 0:-10:-fd_rate_max  # Hz/s
iterations = 100
prn = 1

results = Array{Float64}(undef, length(CN0s), length(fd_rates), length(integration_times))

p = Progress(iterations*length(CN0s)*length(fd_rates)*length(integration_times), 1, "Processing...")
@threads for k in 1:length(integration_times)
    T = integration_times[k]
    signal = definesignal(sigtype, f_s, T; prn=prn, f_if=f_if,
                      include_phase_noise=true, include_thermal_noise=true,
                      nADC=8,
                      receiver_h_parms=h_parms)
    replica = definereplica(sigtype, f_s, T; f_d=0)
    generatereplica!(replica)
    for i in 1:length(CN0s)
        for j in 1:length(fd_rates)
                SNR = 0.
                CN0 = CN0s[i]
                fd_rate = fd_rates[j]
                for iteration in 1:iterations
                    ϕ₀ = rand(0:0.0001:2π)
                    # f_d = rand(-500:0.1:500)
                    f_d = 0
                    code_start_idx = rand(1:signal.sample_num)
                    definesignal!(signal; fd_rate=fd_rate, CN0=CN0,
                                code_start_idx=code_start_idx, phi=ϕ₀,
                                new_thermal_noise=true, new_phase_noise=true,
                                f_d=f_d)
                    generatesignal!(signal)
                    corr_result = fft_correlate(signal.data[1:replica.sample_num],
                                                replica.data)
                    corr_result = abs2.(corr_result)
                    PS = maximum(corr_result)
                    PN = (sum(corr_result) - PS) / (length(corr_result) - 1)
                    SNR += PS/(2*PN)
                    next!(p)
                end
                results[i,j,k] = 10log10(SNR/iterations)
        end
    end
end

save(string(directory, file_name, ".jld"), "results", results,
     "integration_times", integration_times, "fd_rates", fd_rates,
     "sig_freq", sig_freq, "f_d", f_d, "f_if", f_if, "f_s", f_s)

doppler_file_name = string(directory, "ch1_doppler_curves.jld")
doppler_raw_gps, doppler_rate_raw_gps, sig_freq = load(doppler_file_name,
                                                       "doppler_raw_gps",
                                                       "doppler_rate_raw_gps",
                                                       "freq")
doppler_raw_iridium, doppler_rate_raw_iridium = load(doppler_file_name,
                                                     "doppler_raw_iridium",
                                                     "doppler_rate_raw_iridium")
doppler_raw_starlink, doppler_rate_raw_starlink = load(doppler_file_name,
                                                       "doppler_raw_starlink",
                                                       "doppler_rate_raw_starlink")
doppler_raw_oneweb, doppler_rate_raw_oneweb = load(doppler_file_name,
                                                   "doppler_raw_oneweb",
                                                   "doppler_rate_raw_oneweb")

const_names = ["GPS", "Iridium", "Starlink", "OneWeb"]
const_max_fd_rates = [minimum(doppler_rate_raw_gps),
                      minimum(doppler_rate_raw_iridium),
                      minimum(doppler_rate_raw_starlink),
                      minimum(doppler_rate_raw_oneweb)]

file_names = ["ch4_acquisition_performance_nonheatmap_l1ca",
              "ch4_acquisition_performance_nonheatmap_l5"]
fig = figure(figsize=(7.5, 3.5))
for j in 1:length(file_names)
    file_name = file_names[j]
    results, fd_rates, integration_times = load(string(directory, file_name, ".jld"),
                                                "results",
                                                "fd_rates",
                                                "integration_times")
    ax = fig.add_subplot(1,2,j)
    for i in 1:length(integration_times)
        plot(abs.(fd_rates), results[1,:,i],
            label="T = $(Int(integration_times[i]*1e3))ms")
    end
    for k in 1:length(const_names)
        axvline(x=abs(const_max_fd_rates[k]), c="grey", linestyle=":")
    end
    xlabel("Doppler Rate (Hz/s)")
    ylabel("SNR (dB)")
    xscale("log")
    legend()
    letter = ('a' + (j-1))
    title("($(letter))")
end
subplots_adjust(wspace=0.3, bottom=0.15, left=0.08, right=0.92, top=0.93)
savefig(string(directory, "figures/ch4_acquisition_high_doppler.svg"), dpi=300)

ext: jl


text:
using Glob

@test string.(glob("*tests.jl", cwd())) == glob("*tests.jl", pwd())
@test string.(readdir(glob"*tests.jl", cwd())) ==  readdir(glob"*tests.jl", pwd())


text:
# ---
# title: 905. Sort Array By Parity
# id: problem905
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/sort-array-by-parity/description/>
# hidden: true
# ---
#
# Given an array `A` of non-negative integers, return an array consisting of all
# the even elements of `A`, followed by all the odd elements of `A`.
#
# You may return any answer array that satisfies this condition.
#
#
#
# **Example 1:**
#
#
#
#     Input: [3,1,2,4]
#     Output: [2,4,3,1]
#     The outputs [4,2,3,1], [2,4,1,3], and [4,2,1,3] would also be accepted.
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 5000`
#   2. `0 <= A[i] <= 5000`
#
#
## @lc code=start
using LeetCode

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


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module Query

import ..PkgError
using ..Types
import Pkg3.equalto

# If there are explicitly required packages, dicards all versions outside
# the allowed range.
# It also propagates requirements: when all allowed versions of a required package
# require some other package, this creates a new implicit requirement.
# The propagation is tracked so that in case a contradiction is detected the error
# message allows to determine the cause.
# This is a pre-pruning step, so it also creates some structures which are later used by pruning
function filter_versions(reqs::Requires, deps::Dict{String,Dict{VersionNumber,Available}},
                         bktrc::ResolveBacktrace, uuid_to_name::Dict{String, String})
    allowed = Dict{String,Dict{VersionNumber,Bool}}()
    staged = copy(reqs)
    while !isempty(staged)
        staged_next = Requires()
        for (p,vs) in staged
            # Parse requirements and store allowed versions.
            depsp = deps[p]
            if !haskey(allowed, p)
                allowedp = Dict{VersionNumber,Bool}(vn=>true for vn in keys(depsp))
                allowed[p] = allowedp
                seen = false
            else
                allowedp = allowed[p]
                oldallowedp = copy(allowedp)
                seen = true
            end
            for vn in keys(depsp)
                allowedp[vn] &= vn ∈ vs
            end
            @assert !isempty(allowedp)
            if !any(values(allowedp))
                name = haskey(uuid_to_name, p) ? uuid_to_name[p] : "UNKNOWN"
                uuid_short = p[1:8]
                err_msg = "Unsatisfiable requirements detected for package $name [$uuid_short]:\n"
                err_msg *= string(bktrc[p])
                err_msg *= """The intersection of the requirements is $(bktrc[p].versionreq).
                              None of the available versions can satisfy this requirement."""
                throw(PkgError(err_msg))
            end

            # If we've seen this package already and nothing has changed since
            # the last time, we stop here.
            seen && allowedp == oldallowedp && continue

            # Propagate requirements:
            # if all allowed versions of a required package require some other package,
            # then compute the union of the allowed versions for that other package, and
            # treat that as a new requirement.
            # Start by filtering out the non-allowed versions
            fdepsp = Dict{VersionNumber,Available}(vn=>depsp[vn] for vn in keys(depsp) if allowedp[vn])
            # Collect all required packages
            isreq = Dict{String,Bool}(rp=>true for a in values(fdepsp) for rp in keys(a.requires))
            # Compute whether a required package appears in all requirements
            for a in values(fdepsp), rp in keys(isreq)
                haskey(a.requires, rp) || (isreq[rp] = false)
            end
            staged_new = Set{String}()
            for a in values(fdepsp)
                for (rp,rvs) in a.requires
                    # Skip packages that may not be required
                    isreq[rp] || continue
                    # Compute the union of the version sets
                    snvs = get!(staged_next, rp, copy(rvs))
                    union!(snvs, rvs)
                    push!(staged_new, rp)
                end
            end
            for rp in staged_new
                srvs = staged_next[rp]
                isreq[rp] || continue
                bktrcp = get!(bktrc, rp) do; ResolveBacktraceItem(); end
                push!(bktrcp, p=>bktrc[p], srvs)
                if isa(bktrcp.versionreq, VersionSet) && isempty(bktrcp.versionreq)
                    name = haskey(uuid_to_name, rp) ? uuid_to_name[rp] : "UNKNOWN"
                    uuid_short = rp[1:8]
                    err_msg = "Unsatisfiable requirements detected for package $name [$uuid_short]:\n"
                    err_msg *= string(bktrcp)
                    err_msg *= "The intersection of the requirements is empty."
                    throw(PkgError(err_msg))
                end
            end
        end
        staged = staged_next
    end

    filtered_deps = Dict{String,Dict{VersionNumber,Available}}()
    for (p,depsp) in deps
        filtered_deps[p] = Dict{VersionNumber,Available}()
        allowedp = get(allowed, p, Dict{VersionNumber,Bool}())
        fdepsp = filtered_deps[p]
        for (vn,a) in depsp
            get(allowedp, vn, true) || continue
            fdepsp[vn] = a
        end
    end

    return filtered_deps, allowed
end

# Reduce the number of versions by creating equivalence classes, and retaining
# only the highest version for each equivalence class.
# Two versions are equivalent if:
#   1) They appear together as dependecies of another package (i.e. for each
#      dependency relation, they are both required or both not required)
#   2) They have the same dependencies
# Preliminarily calls filter_versions.
function prune_versions(reqs::Requires, deps::Dict{String,Dict{VersionNumber,Available}}, bktrc::ResolveBacktrace, uuid_to_name::Dict{String, String})
    filtered_deps, allowed = filter_versions(reqs, deps, bktrc, uuid_to_name)

    # To each version in each package, we associate a BitVector.
    # It is going to hold a pattern such that all versions with
    # the same pattern are equivalent.
    vmask = Dict{String,Dict{VersionNumber, BitVector}}()

    # For each package, we examine the dependencies of its versions
    # and put together those which are equal.
    # While we're at it, we also collect all dependencies into alldeps
    alldeps = Dict{String,Set{VersionSet}}()
    for (p,fdepsp) in filtered_deps
        # Extract unique dependencies lists (aka classes), thereby
        # assigning an index to each class.
        uniqdepssets = unique(a.requires for a in values(fdepsp))

        # Store all dependencies seen so far for later use
        for r in uniqdepssets, (rp,rvs) in r
            get!(alldeps, rp) do; Set{VersionSet}() end
            push!(alldeps[rp], rvs)
        end

        # If the package has just one version, it's uninteresting
        length(deps[p]) == 1 && continue

        # Grow the pattern by the number of classes
        luds = length(uniqdepssets)
        @assert !haskey(vmask, p)
        vmask[p] = Dict{VersionNumber,BitVector}()
        vmaskp = vmask[p]
        for vn in keys(fdepsp)
            vmaskp[vn] = falses(luds)
        end
        for (vn,a) in fdepsp
            vmind = findfirst(equalto(a.requires), uniqdepssets)
            @assert vmind > 0
            vm = vmaskp[vn]
            vm[vmind] = true
        end
    end

    # Produce dependency patterns.
    for (p,vss) in alldeps, vs in vss
        # packages with just one version, or dependencies
        # which do not distiguish between versions, are not
        # interesting
        (length(deps[p]) == 1 || vs == VersionSet()) && continue

        # Store the dependency info in the patterns
        @assert haskey(vmask, p)
        for (vn,vm) in vmask[p]
            push!(vm, vn in vs)
        end
    end

    # At this point, the vmask patterns are computed. We divide them into
    # classes so that we can keep just one version for each class.
    pruned_vers = Dict{String,Vector{VersionNumber}}()
    eq_classes = Dict{String,Dict{VersionNumber,Vector{VersionNumber}}}()
    for (p, vmaskp) in vmask
        vmask0_uniq = unique(values(vmaskp))
        nc = length(vmask0_uniq)
        classes = [VersionNumber[] for c0 = 1:nc]
        for (vn,vm) in vmaskp
            c0 = findfirst(equalto(vm), vmask0_uniq)
            push!(classes[c0], vn)
        end
        map(sort!, classes)

        # For each nonempty class, we store only the highest version)
        pruned_vers[p] = VersionNumber[]
        prunedp = pruned_vers[p]
        eq_classes[p] = Dict{VersionNumber,Vector{VersionNumber}}()
        eqclassp = eq_classes[p]
        for cl in classes
            if !isempty(cl)
                vtop = maximum(cl)
                push!(prunedp, vtop)
                @assert !haskey(eqclassp, vtop)
                eqclassp[vtop] = cl
            end
        end
        sort!(prunedp)
    end
    # Put non-allowed versions into eq_classes
    for (p, allowedp) in allowed
        haskey(eq_classes, p) || continue
        eqclassp = eq_classes[p]
        for (vn, a) in allowedp
            a && continue
            eqclassp[vn] = [vn]
        end
    end
    # Put all remaining packages into eq_classes
    for (p, depsp) in deps
        haskey(eq_classes, p) && continue
        eq_classes[p] = Dict{VersionNumber,Vector{VersionNumber}}()
        eqclassp = eq_classes[p]
        for vn in keys(depsp)
            eqclassp[vn] = [vn]
        end
    end

    # Recompute deps. We could simplify them, but it's not worth it
    new_deps = Dict{String,Dict{VersionNumber,Available}}()

    for (p,depsp) in deps
        @assert !haskey(new_deps, p)
        if !haskey(pruned_vers, p)
            new_deps[p] = depsp
            continue
        end
        new_deps[p] = Dict{VersionNumber,Available}()
        pruned_versp = pruned_vers[p]
        for (vn,a) in depsp
            vn ∈ pruned_versp || continue
            new_deps[p][vn] = a
        end
    end

    #println("pruning stats:")
    #numvers = 0
    #numdeps = 0
    #for (p,d) in deps, (vn,a) in d
    #    numvers += 1
    #    for r in a.requires
    #        numdeps += 1
    #    end
    #end
    #numnewvers = 0
    #numnewdeps = 0
    #for (p,d) in new_deps, (vn,a) in d
    #    numnewvers += 1
    #    for r in a.requires
    #        numnewdeps += 1
    #    end
    #end
    #println("  before: vers=$numvers deps=$numdeps")
    #println("  after: vers=$numnewvers deps=$numnewdeps")
    #println()

    return new_deps, eq_classes
end
prune_versions(env, deps::Dict{String,Dict{VersionNumber,Available}}) =
    prune_versions(env, Dict{String,VersionSet}(), deps, ResolveBacktrace())
prune_versions(env, deps::Dict{String,Dict{VersionNumber,Available}}, bktrc::ResolveBacktrace) =
    prune_versions(env, Dict{String,VersionSet}(), deps, bktrc)

# Build a graph restricted to a subset of the packages
function subdeps(deps::Dict{String,Dict{VersionNumber,Available}}, pkgs::Set{String})
    sub_deps = Dict{String,Dict{VersionNumber,Available}}()
    for p in pkgs
        haskey(sub_deps, p) || (sub_deps[p] = Dict{VersionNumber,Available}())
        sub_depsp = sub_deps[p]
        for (vn,a) in deps[p]
            sub_depsp[vn] = a
        end
    end

    return sub_deps
end

# Build a subgraph incuding only the (direct and indirect) dependencies
# of a given package set
function dependencies_subset(deps::Dict{String,Dict{VersionNumber,Available}}, pkgs::Set{String})
    staged::Set{String} = filter(p->p in keys(deps), pkgs)
    allpkgs = copy(staged)
    while !isempty(staged)
        staged_next = Set{String}()
        for p in staged, a in values(get(deps, p, Dict{VersionNumber,Available}())), rp in keys(a.requires)
            rp ∉ allpkgs && rp ≠ "julia" && push!(staged_next, rp)
        end
        union!(allpkgs, staged_next)
        staged = staged_next
    end

    return subdeps(deps, allpkgs)
end

function prune_dependencies(reqs::Requires, deps::Dict{String,Dict{VersionNumber,Available}}, uuid_to_name::Dict{String, String})
    bktrc = ResolveBacktrace()
    for (p,vs) in reqs
        bktrc[p] = ResolveBacktraceItem(:required, vs)
    end
    return prune_dependencies(reqs, deps, bktrc, uuid_to_name)
end

function prune_dependencies(reqs::Requires, deps::Dict{String,Dict{VersionNumber,Available}},
                            bktrc::ResolveBacktrace, uuid_to_name::Dict{String, String})
    deps = dependencies_subset(deps, Set{String}(keys(reqs)))
    deps, _ = prune_versions(reqs, deps, bktrc, uuid_to_name)

    return deps
end

end # module


text:
using Test
using LinearAlgebra
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

function nonpoly_displacement(alpha, x)
    u1 = alpha * x[2] * sin(pi * x[1])
    u2 = alpha * (x[1]^3 + cos(pi * x[2]))
    return [u1, u2]
end

function nonpoly_body_force(lambda, mu, alpha, x)
    b1 = alpha * (lambda + 2mu) * pi^2 * x[2] * sin(pi * x[1])
    b2 =
        -alpha * (6mu * x[1] + (lambda + mu) * pi * cos(pi * x[1])) +
        alpha * (lambda + 2mu) * pi^2 * cos(pi * x[2])
    return [b1, b2]
end

function onboundary(x, L, W)
    return x[2] ≈ 0.0 || x[1] ≈ L || x[2] ≈ W || x[1] ≈ 0.0
end

function solve_plane_interface(x0, normal, nelmts, polyorder, numqp, theta)
    L = 1.0
    W = 1.0
    lambda, mu = 1.0, 2.0
    dx = 1.0 / nelmts
    penalty = theta / dx * (lambda + mu)
    alpha = 0.1
    stiffness = CutCell.HookeStiffness(lambda, mu, lambda, mu)

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCell.Mesh([0.0, 0.0], [L, W], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs =
        CutCell.levelset_coefficients(x -> plane_distance_function(x, normal, x0), mesh)

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition =
        CutCell.coherent_interface_condition(basis, interfacequads, stiffness, cutmesh, penalty)

    displacementbc = CutCell.DisplacementCondition(
        x -> nonpoly_displacement(alpha, x),
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onboundary(x, L, W),
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh)
    CutCell.assemble_body_force_linear_form!(
        sysrhs,
        x -> nonpoly_body_force(lambda, mu, alpha, x),
        basis,
        cellquads,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(sysmatrix, sysrhs, displacementbc, cutmesh)

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    sol = matrix \ rhs
    disp = reshape(sol, 2, :)

    err =
        mesh_L2_error(disp, x -> nonpoly_displacement(alpha, x), basis, cellquads, cutmesh)
    return err
end

function solve_curved_interface(xc, radius, nelmts, polyorder, numqp, penaltyfactor)
    L = 1.0
    W = 1.0
    lambda, mu = 1.0, 2.0
    dx = 1.0 / nelmts
    penalty = penaltyfactor / dx * (lambda + mu)
    alpha = 0.1
    stiffness = CutCell.HookeStiffness(lambda, mu, lambda, mu)

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCell.Mesh([0.0, 0.0], [L, W], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs =
        CutCell.levelset_coefficients(x -> circle_distance_function(x, xc, radius), mesh)

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition =
        CutCell.coherent_interface_condition(basis, interfacequads, stiffness, cutmesh, penalty)

    displacementbc = CutCell.DisplacementCondition(
        x -> nonpoly_displacement(alpha, x),
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onboundary(x, L, W),
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh)
    CutCell.assemble_body_force_linear_form!(
        sysrhs,
        x -> nonpoly_body_force(lambda, mu, alpha, x),
        basis,
        cellquads,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(sysmatrix, sysrhs, displacementbc, cutmesh)

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    sol = matrix \ rhs
    disp = reshape(sol, 2, :)

    err =
        mesh_L2_error(disp, x -> nonpoly_displacement(alpha, x), basis, cellquads, cutmesh)
    return err
end

function mean(v)
    return sum(v)/length(v)
end

x0 = [0.5, 0.0]
normal = [1.0, 0.0]
powers = [3,4,5]
nelmts = [2^p + 1 for p in powers]
polyorder = 1
numqp = 3
theta = 1e2

err = [solve_plane_interface(x0, normal, ne, polyorder, numqp, theta) for ne in nelmts]
u1err = [er[1] for er in err]
u2err = [er[2] for er in err]
dx = 1.0 ./ nelmts

u1rate = diff(log.(u1err)) ./ diff(log.(dx))
u2rate = diff(log.(u2err)) ./ diff(log.(dx))

# println("Plane interface convergence of linear elements : ", u1rate[end], "    ", u2rate[end])
@test allapprox(u1rate,repeat([2.0],length(u1rate)),0.1)
@test allapprox(u2rate,repeat([2.0],length(u2rate)),0.1)

polyorder = 2
numqp = 4
theta = 1e2
powers = [2,3,4]
nelmts = [2^p + 1 for p in powers]

err = [solve_plane_interface(x0, normal, ne, polyorder, numqp, theta) for ne in nelmts]
u1err = [er[1] for er in err]
u2err = [er[2] for er in err]
dx = 1.0 ./ nelmts

u1rate = diff(log.(u1err)) ./ diff(log.(dx))
u2rate = diff(log.(u2err)) ./ diff(log.(dx))

# println("Plane interface convergence of quadratic elements : ", u1rate[end], "    ", u2rate[end])
@test allapprox(u1rate,repeat([3.0],length(u1rate)),0.05)
@test allapprox(u2rate,repeat([3.0],length(u2rate)),0.05)

xc = [1.0, 0.5]
radius = 0.45
polyorder = 2
penaltyfactor = 1e2
powers = [3,4,5]
nelmts = [2^p + 1 for p in powers]
numqp = required_quadrature_order(polyorder) + 2
nelmts = [2^p + 1 for p in powers]

err = [solve_curved_interface(xc, radius, ne, polyorder, numqp, penaltyfactor) for ne in nelmts]
u1err = [er[1] for er in err]
u2err = [er[2] for er in err]
dx = 1.0 ./ nelmts

u1rate = diff(log.(u1err)) ./ diff(log.(dx))
u2rate = diff(log.(u2err)) ./ diff(log.(dx))

# println("Curved interface convergence of quadratic elements : ", u1rate[end], "    ", u2rate[end])
@test allapprox(u1rate,repeat([3.0],length(u1rate)),0.05)
@test allapprox(u2rate,repeat([3.0],length(u2rate)),0.05)

xc = [0.5, 0.5]
radius = 0.25
polyorder = 2
penaltyfactor = 1e2
powers = [3,4,5]
nelmts = [2^p + 1 for p in powers]
numqp = required_quadrature_order(polyorder) + 2
nelmts = [2^p + 1 for p in powers]

err = [solve_curved_interface(xc, radius, ne, polyorder, numqp, penaltyfactor) for ne in nelmts]
u1err = [er[1] for er in err]
u2err = [er[2] for er in err]
dx = 1.0 ./ nelmts

u1rate = diff(log.(u1err)) ./ diff(log.(dx))
u2rate = diff(log.(u2err)) ./ diff(log.(dx))

# println("Curved interface convergence of quadratic elements : ", u1rate[end], "    ", u2rate[end])
@test allapprox(u1rate,repeat([3.0],length(u1rate)),0.05)
@test allapprox(u2rate,repeat([3.0],length(u2rate)),0.05)


text:
using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack
using ToeplitzMatrices

using KernelAbstractions
using CUDAapi
if CUDAapi.has_cuda_gpu()
    using CuArrays
    CuArrays.allowscalar(false)
end

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D
using Basis2DTri
using UniformTriMesh

push!(LOAD_PATH, "./examples/EntropyStableEuler")
using EntropyStableEuler

using SetupDG

S_N(x) = @. sin(pi*x/h)/(2*pi/h)/tan(x/2)
"""
Vandermonde matrix of sinc basis functions determined by h,
evaluated at r
"""
function vandermonde_Sinc(h,r)
    N = convert(Int, 2*pi/h)
    V = zeros(length(r),N)
    for n = 1:N
        V[:,n] = S_N(r.-n*h)
    end
    V[1,1] = 1
    V[end,end] = 1
    return V
end

"Constants"
const sp_tol = 1e-12
const has_gpu = CUDAapi.has_cuda_gpu()

"Program parameters"
compute_L2_err = false

"Approximation Parameters"
N_P   = 2;    # The order of approximation in polynomial dimension
Np_P  = Int((N_P+1)*(N_P+2)/2)
Np_F  = 8;    # The order of approximation in Fourier dimension
K1D   = 20;   # Number of elements in polynomial (x,y) dimension
CFL   = 1.0;
T     = 1.0;  # End time

"Time integration Parameters"
rk4a,rk4b,rk4c = rk45_coeffs()
CN = (N_P+1)*(N_P+2)*3/2  # estimated trace constant for CFL
dt = CFL * 2 / CN / K1D
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

"Initialize Reference Element in Fourier dimension"
h = 2*pi/Np_F
column = [0; .5*(-1).^(1:Np_F-1).*cot.((1:Np_F-1)*h/2)]
column2 = [-pi^2/3/h^2-1/6; -((-1).^(1:Np_F-1)./(2*(sin.((1:Np_F-1)*h/2)).^2))]
Dt = Array{Float64,2}(Toeplitz(column,column[[1;Np_F:-1:2]]))
D2t = Array{Float64,2}(Toeplitz(column2,column2[[1;Np_F:-1:2]]))
t = LinRange(h,2*pi,Np_F)

"Initialize Reference Element in polynomial dimension"
rd = init_reference_tri(N_P);
@unpack fv,Nfaces,r,s,VDM,V1,Dr,Ds,rf,sf,wf,nrJ,nsJ,rq,sq,wq,Vq,M,Pq,Vf,LIFT = rd
Nq_P = length(rq)
Nfp_P = length(rf)
Nh_P = Nq_P+Nfp_P # Number of hybridized points
Lq = LIFT

"Mesh related variables"
# First initialize 2D triangular mesh
VX,VY,EToV = uniform_tri_mesh(K1D,K1D)
@. VX = 1+VX
@. VY = 1+VY
md = init_mesh((VX,VY),EToV,rd)
VX = repeat(VX,2)
VY = repeat(VY,2)
VZ = [2/Np_F*ones((K1D+1)*(K1D+1),1); 2*ones((K1D+1)*(K1D+1),1)]
EToV = [EToV EToV.+(K1D+1)*(K1D+1)]

# Make domain periodic
@unpack Nfaces,Vf = rd
@unpack xf,yf,K,mapM,mapP,mapB = md
LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
mapPB = build_periodic_boundary_maps(xf,yf,LX,LY,Nfaces*K,mapM,mapP,mapB)
mapP[mapB] = mapPB
@pack! md = mapP

# Initialize 3D mesh
@unpack x,y,xf,yf,xq,yq,rxJ,sxJ,ryJ,syJ,J,nxJ,nyJ,sJ,mapM,mapP,mapB = md
x,y,xf,yf,xq,yq,rxJ,sxJ,ryJ,syJ,J,nxJ,nyJ,sJ = (x->reshape(repeat(x,inner=(1,Np_F)),size(x,1),Np_F*K)).((x,y,xf,yf,xq,yq,rxJ,sxJ,ryJ,syJ,J,nxJ,nyJ,sJ))
z,zq,zf = (x->reshape(repeat(collect(2/Np_F:(2/Np_F):2),inner=(1,x),outer=(K,1))',x,Np_F*K)).((Np_P,Nq_P,Nfp_P))
mapM = reshape(1:Nfp_P*Np_P*K,Nfp_P,Np_P*K)
mapP_2D = (x->mod1(x,Nfp_P)+div(x-1,Nfp_P)*Nfp_P*Np_F).(mapP)
mapP = reshape(repeat(mapP_2D,inner=(1,Np_F)),Nfp_P,Np_F,K)
for j = 1:Np_F
    mapP[:,j,:] = mapP[:,j,:].+(j-1)*Nfp_P
end
mapP = reshape(mapP,Nfp_P,Np_F*K)

# scale by Fourier dimension
M = h*M
wq = h*wq
wf = h*wf

# Hybridized operators
Vh = [Vq;Vf]
rxJ,sxJ,ryJ,syJ = (x->mapslices((y->Vh*y),x,dims=(1,2))).((rxJ,sxJ,ryJ,syJ))
Ef = Vf*Pq
Br,Bs = (x->diagm(wf.*x)).((nrJ,nsJ))
Qr,Qs = (x->Pq'*M*x*Pq).((Dr,Ds))
Qrh,Qsh = (x->1/2*[x[1]-x[1]' Ef'*x[2];
                   -x[2]*Ef   x[2]]).(((Qr,Br),(Qs,Bs)))
Qrh_skew,Qsh_skew = (x->1/2*(x-x')).((Qrh,Qsh))
Qt = Dt
Qth = Qt # Not the SBP operator, weighted when flux differencing
Ph = [Vq;Vf]*Pq # TODO: refactor
LIFTq = Vq*Lq

# TODO: assume mesh uniform affine, so Jacobian are constants
# TODO: fix other Jacobian parts
JP = 1/K1D^2
JF = 1/pi
J = JF*JP
wq = J*wq
wf = JF*wf
Lq = 1/JP*Lq
Qrh = JF*Qrh
Qsh = JF*Qsh
Qth = JP*Qth
Qrh_skew = 1/2*(Qrh-Qrh')
Qsh_skew = 1/2*(Qsh-Qsh')
VPh = Vq*[Pq Lq]*diagm(1 ./ [wq;wf])

function flux_differencing_xy!(∇fh,Qh,Qlog,ops_flux,geo_flux,param_flux)
    K,Nq_P,Nh_P,Np_F,Nd = param_flux
    rxJ,sxJ,ryJ,syJ = geo_flux
    Qrh_skew,Qsh_skew,Qth,wq = ops_flux
    if isa(∇fh[1],Array)
        kernel! = update∇fh_xy_kernel!(CPU(),4)
    else
        kernel! = update∇fh_xy_kernel!(CUDA(),256)
    end
    kernel!(∇fh,Qh,Qlog,Qrh_skew,Qsh_skew,Nh_P,Np_F,Nd,rxJ,sxJ,ryJ,syJ,ndrange=(K,Np_F))
end

#TODO: clean up function arguments
@kernel function update∇fh_xy_kernel!(∇fh,Qh,Qlog,Qrh_skew,Qsh_skew,Nh_P,Np_F,Nd,rxJ_mat,sxJ_mat,ryJ_mat,syJ_mat)
    k,nf = @index(Global, NTuple)
    j = nf+Np_F*(k-1)
    # TODO: find a fix for not passing matrices?
    rxJ,sxJ,ryJ,syJ = rxJ_mat[1,j],sxJ_mat[1,j],ryJ_mat[1,j],syJ_mat[1,j]
    for col_idx = 1:Nh_P
        for row_idx = col_idx:Nh_P
            Fxj_tmp, Fyj_tmp,_ = euler_fluxes(Qh[1][row_idx,j],Qh[2][row_idx,j],Qh[3][row_idx,j],Qh[4][row_idx,j],Qh[5][row_idx,j],
                                               Qh[1][col_idx,j],Qh[2][col_idx,j],Qh[3][col_idx,j],Qh[4][col_idx,j],Qh[5][col_idx,j],
                                               Qlog[1][row_idx,j],Qlog[2][row_idx,j],
                                               Qlog[1][col_idx,j],Qlog[2][col_idx,j])
            var_Qrh = Qrh_skew[row_idx,col_idx]
            var_Qsh = Qsh_skew[row_idx,col_idx]
            for d = 1:Nd
                update_val = 2*((rxJ*var_Qrh+sxJ*var_Qsh)*Fxj_tmp[d]
                               +(ryJ*var_Qrh+syJ*var_Qsh)*Fyj_tmp[d])
                ∇fh[d][row_idx,j] += update_val
                ∇fh[d][col_idx,j] -= update_val
            end
        end
    end
end

function flux_differencing_z!(∇fh,Qh,Qlog,ops_flux,geo_flux,param_flux)
    K,Nq_P,Nh_P,Np_F,Nd = param_flux
    rxJ,sxJ,ryJ,syJ = geo_flux
    Qrh_skew,Qsh_skew,Qth,wq = ops_flux
    if isa(∇fh[1],Array)
        kernel! = update∇fh_z_kernel!(CPU(),4)
    else
        kernel! = update∇fh_z_kernel!(CUDA(),256)
    end
    kernel!(∇fh,Qth,Np_F,Nd,wq,J,Qlog,Qh,ndrange=(K,Nq_P))
end

# TODO: clean up
@kernel function update∇fh_z_kernel!(∇fh,Qth,Np_F,Nd,wq,J,Qlog,Qh)
    k,nh = @index(Global,NTuple)
    j_idx = (k-1)*Np_F+1:k*Np_F
    wqn = 2/J*wq[nh]
    for col_idx = 1:Np_F
        for row_idx = col_idx:Np_F
            _,_,f_tmp = euler_fluxes((x->x[nh,j_idx[row_idx]]).(Qh),(x->x[nh,j_idx[col_idx]]).(Qh),(x->x[nh,j_idx[row_idx]]).(Qlog),(x->x[nh,j_idx[col_idx]]).(Qlog))
            var_Qth = wqn*Qth[row_idx,col_idx]
            for i = 1:Nd
                ∇fh[i][nh,j_idx[row_idx]] += var_Qth*f_tmp[i]
                ∇fh[i][nh,j_idx[col_idx]] -= var_Qth*f_tmp[i]
            end
        end
    end
end

# Wrapper
function update_rhs_flux!(rhsQ,Nh_P,Nq_P,K,Np_F,Nfp_P,mapP,Nd,Qh,nxJ,nyJ,Qlog)
    if isa(rhsQ[1],Array)
        kernel! = update_rhs_flux_kernel!(CPU(),8)
    else
        kernel! = update_rhs_flux_kernel!(CUDA(),256)
    end
    kernel!(rhsQ,Nh_P,Nq_P,K,Np_F,Nfp_P,mapP,Nd,Qh,nxJ,nyJ,Qlog,ndrange=size(rhsQ[1]))
end

@kernel function update_rhs_flux_kernel!(rhsQ,Nh_P,Nq_P,K,Np_F,Nfp_P,mapP,Nd,Qh,nxJ,nyJ,Qlog)
    row_idx, col_idx = @index(Global,NTuple)

    tmp_idx = mapP[row_idx+(col_idx-1)*Nfp_P]
    r_idx = Nq_P+Nh_P*div(tmp_idx-1,Nfp_P)+mod1(tmp_idx,Nfp_P)
    tmp_flux_x, tmp_flux_y,_ = euler_fluxes((x->x[Nq_P+row_idx,col_idx]).(Qh),(x->x[r_idx]).(Qh),
                                            (x->x[Nq_P+row_idx,col_idx]).(Qlog),(x->x[r_idx]).(Qlog))
    tmp_nxJ = nxJ[row_idx,col_idx]
    tmp_nyJ = nyJ[row_idx,col_idx]
    for d = 1:Nd
        # normal_flux(fx,fy,u) = fx.*nxJ + fy.*nyJ - LFc.*(u[mapP]-u)
        # TODO: Lax Friedrichs
        rhsQ[d][row_idx,col_idx] = tmp_flux_x[d]*tmp_nxJ + tmp_flux_y[d]*tmp_nyJ
    end
end

# TODO: refactor
ops = (Vq,Vf,wq,wf,Pq,Lq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,VPh)
mesh = (rxJ,sxJ,ryJ,syJ,sJ,nxJ,nyJ,JP,JF,J,h,mapM,mapP)
param = (K,Np_P,Nfp_P,Np_F,Nq_P,Nh_P)
function rhs(Q,ops,mesh,param,compute_rhstest,has_gpu)

    Vq,Vf,wq,wf,Pq,Lq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq = ops
    rxJ,sxJ,ryJ,syJ,sJ,nxJ,nyJ,JP,JF,J,h,mapM,mapP = mesh
    K,Np_P,Nfp_P,Np_F,Nq_P,Nh_P = param
    Nd = length(Q) # number of components

    # TODO: cleaner syntax
    if has_gpu
        Vq,Vf,wq,wf,Pq,Lq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,rxJ,sxJ,ryJ,syJ,sJ,nxJ,nyJ,mapM,mapP = (x->CuArray(x)).((Vq,Vf,wq,wf,Pq,Lq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,rxJ,sxJ,ryJ,syJ,sJ,nxJ,nyJ,mapM,mapP))
        rhsQ = (CuArray(zeros(Nfp_P,K*Np_F)),CuArray(zeros(Nfp_P,K*Np_F)),CuArray(zeros(Nfp_P,K*Np_F)),CuArray(zeros(Nfp_P,K*Np_F)),CuArray(zeros(Nfp_P,K*Np_F)))
        # TODO: why adding this if statement reduces allocation even when has_gpu = false?
        Qh = (CuArray(zeros(Nh_P,K*Np_F)),CuArray(zeros(Nh_P,K*Np_F)),CuArray(zeros(Nh_P,K*Np_F)),CuArray(zeros(Nh_P,K*Np_F)),CuArray(zeros(Nh_P,K*Np_F)))
        ∇fh = (CuArray(zeros(Nh_P,K*Np_F)),CuArray(zeros(Nh_P,K*Np_F)),CuArray(zeros(Nh_P,K*Np_F)),CuArray(zeros(Nh_P,K*Np_F)),CuArray(zeros(Nh_P,K*Np_F)))
        VU = (CuArray(zeros(Nq_P,K*Np_F)),CuArray(zeros(Nq_P,K*Np_F)),CuArray(zeros(Nq_P,K*Np_F)),CuArray(zeros(Nq_P,K*Np_F)),CuArray(zeros(Nq_P,K*Np_F)))
        tmp = CuArray(zeros(Nh_P,K*Np_F))
        tmp2 = CuArray(zeros(Nh_P,K*Np_F))
    else
        rhsQ = (zeros(Nfp_P,K*Np_F),zeros(Nfp_P,K*Np_F),zeros(Nfp_P,K*Np_F),zeros(Nfp_P,K*Np_F),zeros(Nfp_P,K*Np_F))
        Qh = (zeros(Nh_P,K*Np_F),zeros(Nh_P,K*Np_F),zeros(Nh_P,K*Np_F),zeros(Nh_P,K*Np_F),zeros(Nh_P,K*Np_F))
        # TODO: cleaner syntax, avoid storage?
        ∇fh = (zeros(Nh_P,K*Np_F),zeros(Nh_P,K*Np_F),zeros(Nh_P,K*Np_F),zeros(Nh_P,K*Np_F),zeros(Nh_P,K*Np_F))
        VU = (zeros(Nq_P,K*Np_F),zeros(Nq_P,K*Np_F),zeros(Nq_P,K*Np_F),zeros(Nq_P,K*Np_F),zeros(Nq_P,K*Np_F))
        tmp = zeros(Nh_P,K*Np_F)
        tmp2 = zeros(Nh_P,K*Np_F)
    end
    vector_norm(U) = sum((x->x.^2).(U))
    # VU = v_ufun(Q...)
    @. VU[4] = Q[2]^2+Q[3]^2+Q[4]^2 # rhoUnorm
    @. VU[5] = Q[5]-.5*VU[4]/Q[1] # rhoe
    @. VU[1] = log(0.4*VU[5]/(Q[1]^1.4)) # sU #TODO: hardcoded gamma
    @. VU[1] = (-Q[5]+VU[5]*(2.4-VU[1]))/VU[5]
    @. VU[2] = Q[2]/VU[5]
    @. VU[3] = Q[3]/VU[5]
    @. VU[4] = Q[4]/VU[5]
    @. VU[5] = -Q[1]/VU[5]

    Qh = (x->Ph*x).(VU)
    # (ρ,ρu,ρv,ρw,E) = u_vfun(Qh...)
    @. tmp = Qh[2]^2+Qh[3]^2+Qh[4]^2 #vUnorm
    @. tmp2 = (0.4/((-Qh[5])^1.4))^(1/0.4)*exp(-(1.4 - Qh[1] + tmp/(2*Qh[5]))/0.4) # rhoeV
    @. Qh[1] = tmp2.*(-Qh[5])
    @. Qh[2] = tmp2.*Qh[2]
    @. Qh[3] = tmp2.*Qh[3]
    @. Qh[4] = tmp2.*Qh[4]
    @. Qh[5] = tmp2.*(1-tmp/(2*Qh[5]))

    # TODO: Lax Friedrichs Dissipation flux
    # (ρM,ρuM,ρvM,ρwM,EM) = Uf
    # ρuM_n = @. (ρuM*nxJ+ρvM*nyJ)/sJ # TODO: 3D lax-friedrichs?
    # lam = abs.(wavespeed(ρM,ρuM_n,EM))
    # LFc = .5*max.(lam,lam[mapP]).*sJ

    # (Qh[2][Nq_P+1:end,:].*nxJ+Qh[3][Nq_P+1:end,:].*nyJ)./sJ

    # β = betafun(ρ,ρu,ρv,ρw,E)
    # Qh2 = (ρ,ρu./ρ,ρv./ρ,ρw./ρ,β)
    @. tmp = Qh[1]/(2*0.4*(Qh[5]-.5*(Qh[2]^2+Qh[3]^2+Qh[4]^2)/Qh[1])) #beta
    @. Qh[2] = Qh[2]./Qh[1]
    @. Qh[3] = Qh[3]./Qh[1]
    @. Qh[4] = Qh[4]./Qh[1]
    @. Qh[5] = tmp

    # TODO: implement Lax Friedrichs
    # TODO: storing flux, so avoid calculate flux on face quad point again?
    Qlog = (log.(Qh[1]),log.(Qh[5]))
#    update_rhs_flux!(rhsQ,Nh_P,Nq_P,K,Np_F,Nfp_P,mapP,Nd,Qh,nxJ,nyJ,Qlog)
    ev = update_rhs_flux!(rhsQ,Nh_P,Nq_P,K,Np_F,Nfp_P,mapP,Nd,Qh,nxJ,nyJ,Qlog)
    wait(ev)
    rhsQ = (x->Vq*Lq*x).(rhsQ) # TODO: put it into update_rhs_flux!

    # Flux differencing
    # TODO: fix Jacobian
    param_flux = (K,Nq_P,Nh_P,Np_F,Nd)
    geo_flux = (rxJ,sxJ,ryJ,syJ)
    ops_flux = (Qrh_skew,Qsh_skew,Qth,wq)
    # flux_differencing_xy!(∇fh,Qh,Qlog,ops_flux,geo_flux,param_flux)
    ev2 = flux_differencing_xy!(∇fh,Qh,Qlog,ops_flux,geo_flux,param_flux)
    wait(ev2)
    ev3 = flux_differencing_z!(∇fh,Qh,Qlog,ops_flux,geo_flux,param_flux)
    wait(ev3)
    ∇f = (x->VPh*x).(∇fh)
    rhsQ = @. -(∇f+rhsQ)

    rhstest = 0
    if compute_rhstest
        for fld in eachindex(rhsQ)
            rhstest += sum(wq.*VU[fld].*rhsQ[fld])
        end
    end
    # rhstest = 0
    # rhsQ = [zeros(size(Q[1])), zeros(size(Q[1])),zeros(size(Q[1])),zeros(size(Q[1])),zeros(size(Q[1]))]
    return rhsQ,rhstest
end

xq,yq,zq = (x->reshape(x,Nq_P,Np_F*K)).((xq,yq,zq))
# All directions
println(" ")
println("======= All directions =======")
ρ_exact(x,y,z,t) = @. 1+0.2*sin(pi*(x+y+z-3/2*t))
ρ = @. 1+0.2*sin(pi*(xq+yq+zq))
u = ones(size(xq))
v = -1/2*ones(size(xq))
w = ones(size(xq))
p = ones(size(xq))
Q_exact(x,y,z,t) = (ρ_exact(x,y,z,t),u,v,w,p)

Q = primitive_to_conservative(ρ,u,v,w,p)
Q = collect(Q)
resQ = [zeros(size(Q[1])) for _ in eachindex(Q)]
if has_gpu
    Q = [CuArray(Q[1]),CuArray(Q[2]),CuArray(Q[3]),CuArray(Q[4]),CuArray(Q[5])]
    resQ = [CuArray(resQ[1]),CuArray(resQ[2]),CuArray(resQ[3]),CuArray(resQ[4]),CuArray(resQ[5])]
end
# rhs(Q,ops,mesh,param,false,has_gpu)
# @btime rhs(Q,ops,mesh,param,false,has_gpu)
@time begin
for i = 1:Nsteps
    rhstest = 0
    for INTRK = 1:5
        compute_rhstest = INTRK==5
        rhsQ,rhstest = rhs(Q,ops,mesh,param,compute_rhstest,has_gpu)
        @. resQ = rk4a[INTRK]*resQ + dt*rhsQ
        @. Q += rk4b[INTRK]*resQ
    end

    if i%10==0 || i==Nsteps
        println("Time step: $i out of $Nsteps with rhstest = $rhstest")
    end
end

end # time

rq2,sq2,wq2 = quad_nodes_2D(N_P+2)
Vq2 = vandermonde_2D(N_P,rq2,sq2)/VDM
xq2,yq2,zq2 = (x->Vq2*x).((x,y,z))
ρ = Vq2*Pq*Q[1]
ρ_ex = ρ_exact(xq2,yq2,zq2,T)
Q = (x->Vq2*Pq*x).(Q)
p = pfun(Q[1],(Q[2],Q[3],Q[4]),Q[5])
Q = (Q[1],Q[2]./Q[1],Q[3]./Q[1],Q[4]./Q[1],p)
Q_ex = Q_exact(xq2,yq2,zq2,T)

L2_err = 0.0
for fld in eachindex(Q)
    global L2_err
    L2_err += sum(h*J*wq2.*(Q[fld]-Q_ex[fld]).^2)
end
println("L2err at final time T = $T is $L2_err\n")

max_stars_repo_name: yiminllin/ESDG-Fourier

max_issues_repo_name: yiminllin/ESDG-Fourier

max_forks_repo_name: yiminllin/ESDG-Fourier


text:
#==============================================================================
    Code for solving the Hamiltonian Jacobi Bellman for
	   a Ramsey Model with a diffusion process for capital

 	Based on  Matlab code from Ben Moll:
        http://www.princeton.edu/~moll/HACTproject.htm

        Updated to julia 1.0.0
==============================================================================#

using Distributions, Plots, LinearAlgebra, SparseArrays
using Random

Random.seed!(12364)

γ= 2.0 #gamma parameter for CRRA utility
ρ = 0.05 #the discount rate
α = 1/3 # the curvature of the production function (cobb-douglas)
δ = 0.05 # the depreciation rate
σ = 0.5 #variance term
n=0.51
T = 1000 # set the number of external time periods for misspecification loop

# Add in the misspecification
σ_g =0.02
Σ_g=[σ_g] #create cell with the agent's guess for σ

#=============================================================================
 	k our capital follows a process that depends on
            dk = (s⋅f(k) - (δ-σ^2)k)dt - σ⋅k⋅dW_[t]
        Where W is a Wiener process
=============================================================================#
#K_starting point, for mean of z process

k_st = ((α)/(ρ+δ+n-σ^2))^(1/(1-α))

# create the grid for k
H = 100 #number of points on grid
k_min = 0.1*k_st # min value
k_max = 10*k_st # max value
k = LinRange(k_min, k_max, H)
dk = (k_max-k_min)/(H-1)
k_min_1=k_min
k_max_1=k_max

#create matrices for k and z
kk = k*1

# use Ito's lemma to find the drift and variance of our optimization equation

# parameters for our loops
max_it = 100
ε = 0.1^(6)
Δ = 1000

# create a random distribution to pull from for the updating rule
Dist = Bernoulli(.45)

# set up all of these empty matrices for the misspecified version later on
Vaf_2, Vab_2, c_2 = [zeros(H,1) for i in 1:3]
Val_2 =[]
savings=[]

# set up all of these empty matrices for the steady state algorithm
Vaf_1, Vab_1, c_1 = [zeros(H,1) for i in 1:3]

y = pdf.(Normal(.99, σ), k) #check
plot(k,y, grid=false,
		xlabel="k", ylabel="Probability",
		legend=false, color="purple", title="PDF of k")
#png("PDF_k")

#============================================================================#
# Now it's time to solve the model, first put in a guess for the value function
v0 = (kk.^α).^(1-γ)/(1-γ)/ρ
v_1=v0

maxit= 30 #set number of iterations (only need 6 to converge)
dist = [] # set up empty array for the convergence criteria

for j = 1:maxit
    V=v_1

    #Now set up the forward difference

    Vaf_1[1:H-1,:] = (V[2:H, :] - V[1:H-1,:])/dk
    Vaf_1[H,:] .= (k_max.^α - (δ+n-σ^2).*k_max).^(-γ) # imposes a constraint

    #backward difference
    Vab_1[2:H,:] = (V[2:H, :] - V[1:H-1,:])/dk
    Vab_1[1,:] .= (k_min.^α - (δ+n-σ^2).*k_min).^(-γ)

    #I_concave = Vab .> Vaf # indicator for whether the value function is concave

    # Consumption and savings functions
    cf = Vaf_1.^(-1/γ)
    drf =(1 .-(cf/kk.^α))*kk.^α-(δ+n-σ^2).*kk

    # consumption and saving backwards difference
    cb = Vab_1.^(-1.0/γ)
    drb = (1 .-(cb/k.^α))*kk.^α-(δ+n-σ^2).*kk
    #println(sb)
    #consumption and derivative of the value function at the steady state

    c0 = kk.^α - (δ+n-σ^2).*kk
    Va0 = c0.^(-γ)

    # df chooses between the forward or backward difference

    If = drf.>0 # positive drift will ⇒ forward difference
    Ib = drb.<0 # negative drift ⇒ backward difference
    I0=(1.0.-If-Ib) # at steady state

    Va_upwind = Vaf_1.*If + Vab_1.*Ib + Va0.*I0 # need to include SS term

    global c_1 = Va_upwind.^(-1/γ)
    u = (c_1.^(1-γ))/(1-γ)

    # Now to constuct the A matrix
    X = -min.(drb, 0)/dk + (σ.*kk/(2*dk))
    Y = -max.(drf, 0)/dk + min.(drb, 0)/dk - (σ.*kk/(dk))
    Z = max.(drf, 0)/dk + (σ.*kk/(2*dk))

    updiag = 0
    updiag =[updiag; Z[1:H-1,1]; 0]
    updiag =(updiag[:])

    centerdiag=reshape(Y, H, 1)
    centerdiag = (centerdiag[:]) # for tuples

   lowdiag = X[2:H, 1]
   lowdiag=(lowdiag)

   # spdiags in Matlab allows for automatic trimming/adding of zeros
       # spdiagm does not do this
  A = sparse(Diagonal(centerdiag))+ [zeros(1, H); sparse(Diagonal(lowdiag)) zeros(H-1,1)] + sparse(Diagonal(updiag))[2:end, 1:H] # trim first element

  B = (1/Δ + ρ)*sparse(I,H,H) - A

  u_stacked= reshape(u, H, 1)
  V_stacked = reshape(V,H, 1)

  b = u_stacked + (V_stacked./Δ)

  V_stacked = B\b

  V = reshape(V_stacked, H, 1)

  V_change = V-v_1

  global v_1= V
  # need push function to add to an already existing array
  push!(dist, findmax(abs.(V_change))[1])
  if dist[j].< ε
      println("Value Function Converged Iteration=")
      println(j)
      break
  end

end

# calculate the savings for kk
ss_1 = kk.^α - (δ+n-σ^2).*kk

#==================== Now resolve the model with σ misspecified =======#
for t in 1:T
    # Now it's time to solve the model, first put in a guess for the value function
    global v0 = (kk.^α).^(1-γ)/(1-γ)/ρ
    v_2=v0
    global dist = [] # needs to be in the time loop to work properly
    indicator = rand(Dist)
    for j = 1:maxit
        V=v_2

        #Now set up the forward difference

        Vaf_2[1:H-1,:] = (V[2:H, :] - V[1:H-1,:])/dk
        Vaf_2[H,:] .= (k_max.^α - (δ+n-σ_g^2).*k_max).^(-γ) # imposes a constraint

        #backward difference
        Vab_2[2:H,:] = (V[2:H, :] - V[1:H-1,:])/dk
        Vab_2[1,:] .= (k_min.^α - (δ+n-σ_g^2).*k_min).^(-γ)

        #I_concave = Vab .> Vaf # indicator for whether the value function is concave

        # Consumption and savings functions
        cf = Vaf_2.^(-1/γ)
        drf =(1 .-(cf/kk.^α))*kk.^α-(δ+n-σ_g^2).*kk

        # consumption and saving backwards difference
        cb = Vab_2.^(-1.0/γ)
        drb = (1 .-(cb/k.^α))*kk.^α-(δ+n-σ_g^2).*kk
        #println(sb)
        #consumption and derivative of the value function at the steady state

        c0 = kk.^α - (δ+n-σ_g^2).*kk
        Va0 = c0.^(-γ)

        # df chooses between the forward or backward difference

        If = drf.>0 # positive drift will ⇒ forward difference
        Ib = drb.<0 # negative drift ⇒ backward difference
        I0=(1.0.-If-Ib) # at steady state

        Va_upwind = Vaf_2.*If + Vab_2.*Ib + Va0.*I0 # need to include SS term

        global c_2 = Va_upwind.^(-1/γ)
        u = (c_2.^(1-γ))/(1-γ)

        # Now to constuct the A matrix
        X = -min.(drb, 0)/dk + (σ_g.*kk/(2*dk))
        Y = -max.(drf, 0)/dk + min.(drb, 0)/dk - (σ_g.*kk/(dk))
        Z = max.(drf, 0)/dk + (σ_g.*kk/(2*dk))

        updiag = 0
        updiag =[updiag; Z[1:H-1,1]; 0]
        updiag =(updiag[:])

        centerdiag=reshape(Y, H, 1)
        centerdiag = (centerdiag[:]) # for tuples

       lowdiag = X[2:H, 1]
       lowdiag=(lowdiag)

       # spdiags in Matlab allows for automatic trimming/adding of zeros
           # spdiagm does not do this
      A = sparse(Diagonal(centerdiag))+ [zeros(1, H); sparse(Diagonal(lowdiag)) zeros(H-1,1)] + sparse(Diagonal(updiag))[2:end, 1:H] # trim first element

      B = (1/Δ + ρ)*sparse(I,H,H) - A

      u_stacked= reshape(u, H, 1)
      V_stacked = reshape(V,H, 1)

      b = u_stacked + (V_stacked./Δ)

      V_stacked = B\b

      V = reshape(V_stacked, H, 1)

      V_change = V-v_2

      v_2= V
      # need push function to add to an already existing array
      push!(dist, findmax(abs.(V_change))[1])
      if dist[j].< ε
        push!(Val_2, v_2)
        break
      end
    end
    push!(Σ_g, σ_g)
    # calculate the savings for kk
    ss_2 = kk.^α - (δ+n-σ_g^2).*kk
    push!(savings, ss_2)
    # add in updating
    if indicator ==1
		σ_noise = σ + rand(Normal(0,.1),1)[1] # generate a noisy observation
        global σ_g = σ_g + .01(σ_noise-σ_g) # use RLS to update using the noisy obs.
    end
    println(t)
end

# Plot the savings vs. k
plot(k[1:H-1], c_2[1:H-1], grid=false,
		xlabel="k", ylabel="c(k)",
        xlims=(k[1],k[end]),legend=:bottomright,
		label="Guess", title="Optimal Consumption Policies")
plot!(k[1:H-1],c_1[1:H-1], label="True Value", line=:dash)
#png("OptimalCons")

plot(k, savings[1], grid=false,
		xlabel="k", ylabel="s(k)",
        xlims=(k[1],k[end]),label="Period 1",
        title="Optimal Savings Policies",
		color=:hotpink, legend=:topleft, line=:dashdot)
plot!(k, zeros(H,1), color=:black, label="")
plot!(k,savings[100], label="Period 100", color=:blue,line=:dot)
plot!(k,savings[500], label="Period 500", color=:green, line=:dashdotdot)
plot!(k,savings[1000], label="Period 1,000", color=:orange)
plot!(k,ss_1, label="True Value", line=:dash, color=:black, legend=:bottomleft)
png("OptimalSavings_2")

plot(k, savings, grid=false,
		xlabel="k", ylabel="s(k)",
        xlims=(k[1],k[end]),label="", title="Optimal Savings Policies")
plot!(k, zeros(H,1), line=:dash, color=:black, label="", legend=:topleft)
plot!(k,ss_1, color=:black, label="True Value", legend=:bottomleft)
#png("OptimalSavings_All")

plot(k,Val_2[1], label="Period 1 ", color=:hotpink, line=:dashdot)
plot!(k,Val_2[100], label="Period 100", color=:blue, line=:dot)
plot!(k,Val_2[500], label="Period 500", color=:green, line=:dashdotdot)
plot!(k,Val_2[end], label="Period 1,000", color=:orange )
plot!(k,v_1, grid=false,
		xlabel="k", ylabel="V(k)",
		xlims=(k[1],k[end]), title="Value Functions",
        legend=:bottomright, color=:black, line=:dash,
        label="True Value")
png("Value_function_vs_k_2")

plot(k, Val_2, grid=false,
		xlabel="k", ylabel="V(k)",
		xlims=(k[1],k[end]), title="Value Functions",
		label="", legend=:bottomright)
plot!(k,v_1, label="True Value", color=:black, line=:dash)
#png("Value_function_vs_k_all")

plot([1:T], Σ_g[2:end], grid=false,
		xlabel="time", ylabel="\$\\sigma\$",
		title="\$ \\textrm{Estimate of } \\sigma \\textrm{ over time}\$",
		label="Misspecfication", legend=:bottomright)
plot!([1:T],ones(T).*σ, label="True Value", color=:black, line=:dash)
png("sigma")


text:
using Test
using LinearAlgebra
using MatrixProductStates

@testset "Random MPS and basis states" begin
    # Use left canonical form
    for N = 10:2:20
        mps_real = random_mps_obc(N, 5, 3, Float64)
        mps_complex = random_mps_obc(N, 5, 3, ComplexF64)
        gaugeMPS!(mps_real, :left, true)
        gaugeMPS!(mps_complex, :left, true)
        @test isapprox(calculate_overlap(mps_real, mps_real), 1.0)
        @test isapprox(calculate_overlap(mps_complex, mps_complex), 1.0 + 0.0im)
    end
    # Use right canonical form
    for N = 10:2:20
        mps_real = random_mps_obc(N, 5, 3, Float64)
        mps_complex = random_mps_obc(N, 5, 3, ComplexF64)
        gaugeMPS!(mps_real, :right, true)
        gaugeMPS!(mps_complex, :right, true)
        @test isapprox(calculate_overlap(mps_real, mps_real), 1.0)
        @test isapprox(calculate_overlap(mps_complex, mps_complex), 1.0 + 0.0im)
    end

    # Invalid configurations which should produce an error
    configuration = [0; 1; 2; 1; 2]
    @test_throws ArgumentError basis_state_obc(configuration)
    configuration = [1; 1; 3; 1; 2]
    @test_throws ArgumentError basis_state_obc(configuration)

    # Some valid configurations
    v = Vector{Vector{Float64}}(undef, 3)
    v[1] = [1.0; 0; 0]
    v[2] = [0; 1.0; 0]
    v[3] = [0; 0; 1.0]
    for N = 4:10
        configuration = rand([1; 2; 3], N)
        mps = basis_state_obc(configuration, 3)
        @test isapprox(calculate_overlap(mps, mps), 1.0)
    end
end

@testset "Identity operator" begin
    for N = 10:2:20
        for d = 2:5
            mps_complex = random_mps_obc(N, 5, d, ComplexF64)
            gaugeMPS!(mps_complex, :left, true)
            id_mpo = getIdentityMPO(N, d)
            mps_operator_applied = apply_operator(id_mpo, mps_complex)
            @test isapprox(calculate_overlap(mps_complex, mps_operator_applied), 1.0 + 0.0im)
        end
    end
end

@testset "Canonical form" begin
    mps = random_mps_obc(10, 5, 2, ComplexF64)
    mps_left_gauged = gaugeMPS(mps, :left, true)
    mps_right_gauged = gaugeMPS(mps, :right, true)
    # Check the left canonical gauge
    for i = 1:length(mps_left_gauged)
        Dr = size(mps_left_gauged[i], 2)
        res = contract_tensors(conj(mps_left_gauged[i]), [1; 3], mps_left_gauged[i], [1; 3])
        @test isapprox(res, Matrix((1.0 + 0.0im) * I, Dr, Dr))
    end
    # Check the right canonical gauge
    for i = 1:length(mps_right_gauged)
        Dl = size(mps_right_gauged[i], 1)
        res = contract_tensors(conj(mps_right_gauged[i]), [2; 3], mps_right_gauged[i], [2; 3])
        @test isapprox(res, Matrix((1.0 + 0.0im) * I, Dl, Dl))
    end
end

@testset "Inplace operator application" begin
    # Check that we get an error, if we try to overwrite a real MPS with a complex MPO
    mps = random_mps_obc(10, 9, 2, Float64)
    H = getHeisenbergMPO(10, 1, 1)
    @test_throws InexactError apply_operator!(H, mps)

    # Compute result of the contraction and store in a new MPS, compare it with result of inplace contraction
    for N = 10:2:20
        # Prepare a random MPS and gauge it
        mps = random_mps_obc(N, 15, 2, ComplexF64)
        gaugeMPS!(mps, :left, true)
        # Some MPOs for testing
        H = getIsingMPO(N, 1.0, 0.9)
        # Apply the Hamiltonian generating a copy
        mps_new = apply_operator(H, mps)
        gaugeMPS!(mps_new, :left, true)
        # Now in place
        apply_operator!(H, mps)
        gaugeMPS!(mps, :left, true)
        # Check the expectation values
        @test isapprox(calculate_overlap(mps, mps_new), 1.0 + 0.0im)
    end
end

@testset "Summing MPSs" begin
    for N = 10:2:20
        # Prepare a random MPS and gauge it
        mps = random_mps_obc(N, 15, 2, ComplexF64)
        gaugeMPS!(mps, :left, true)
        # Sum it with itself
        mps2 = sum_states(mps, mps)
        # Check the expectation values
        @test isapprox(calculate_overlap(mps2, mps), 2.0)
        @test isapprox(calculate_overlap(mps2, mps2), 4.0)
    end
end

@testset "Summing MPOs" begin
    for N = 10:2:20
        # Prepare a random MPS and gauge it
        mps = random_mps_obc(N, 15, 2, ComplexF64)
        gaugeMPS!(mps, :left, true)
        # Some MPOs for testing
        id_mpo = getIdentityMPO(N, 2)
        H = getIsingMPO(N, 1.0, 0.9)
        # Compute sums thereof
        O1 = sum_operators(id_mpo, H)
        O2 = sum_operators(H, H)
        # Check the expectation values
        @test isapprox(expectation_value(mps, O1), 1.0 + expectation_value(mps, H))
        @test isapprox(expectation_value(mps, O2), 2.0 * expectation_value(mps, H))
    end
end

@testset "Basis states and contracting virtual indices of an MPS" begin
    # Some experiments with product states
    v = Vector{Vector{Float64}}(undef, 3)
    v[1] = [1.0; 0; 0]
    v[2] = [0; 1.0; 0]
    v[3] = [0; 0; 1.0]
    for N = 4:10
        # A MPS representing a simple basis state
        configuration = rand([1; 2; 3], N)
        mps = basis_state_obc(configuration, 3)
        # Construct the corresponding state vector
        state_vector = 1
        for i = 1:N
            state_vector = kron(state_vector, v[configuration[i]])
        end
        # Compare to the result from the MPS
        state_vector_mps = contract_virtual_indices(mps)
        @test isapprox(state_vector' * state_vector_mps, 1.0)
    end

    # Now try with the GHZ state on 10 sites
    all_zeros = basis_state_obc(ones(Int64, 10))
    all_ones = basis_state_obc(2 * ones(Int64, 10))
    mps_ghz = sum_states(all_zeros, all_ones)
    gaugeMPS!(mps_ghz, :left, true)
    v1 = zeros(2^10)
    v1[1] = 1
    v2 = zeros(2^10)
    v2[end] = 1
    state_vector_ghz = 1 / sqrt(2) * (v1 + v2)
    state_vector_mps = contract_virtual_indices(mps_ghz)
    @test isapprox(state_vector_ghz' * state_vector_mps, 1.0)
end

@testset "Contracting virtual indices of an MPO" begin
    # Test with the identity
    id_mpo = getIdentityMPO(5, 2)
    id_mpo_matrix = contract_virtual_indices(id_mpo)
    @test isapprox(id_mpo_matrix, Matrix(1.0I, 2^5, 2^5))

    # A simple instance of the Ising Hamiltonian
    J = 0.9
    lambda = 1.1
    ising_mpo = getIsingMPO(4, J, lambda)
    Id, X, _, Z = getPauliMatrices()
    Hising = -J * (kron(X, X, Id, Id) + kron(Id, X, X, Id) + kron(Id, Id, X, X)) - lambda * (kron(Z, Id, Id, Id) + kron(Id, Z, Id, Id) + kron(Id, Id, Z, Id) + kron(Id, Id, Id, Z))
    Hising_mpo = contract_virtual_indices(ising_mpo)
    @test isapprox(Hising, Hising_mpo)
end

@testset "Test entropy computation" begin
    # A product state should always have vanishing entropy
    mps = random_mps_obc(10, 1, 2)
    gaugeMPS!(mps, :left, true)
    for i = 1:length(mps)
        @test abs(compute_entropy(mps, i)) < 1E-12
    end
    @test abs(compute_entropy(mps)) < 1E-12

    # Now prepare the GHZ state
    all_zeros = basis_state_obc(ones(Int64, 10))
    all_ones = basis_state_obc(2 * ones(Int64, 10))
    mps = sum_states(all_zeros, all_ones)
    gaugeMPS!(mps, :left, true)
    for i = 1:length(mps)-1
        @test isapprox(compute_entropy(mps, i), 1.0)
    end

    # Compare the entropy computation with the default argument to what one would expect
    mps = random_mps_obc(11, 4, 3)
    entropy1 = compute_entropy(mps)
    entropy2 = compute_entropy(mps, 6)
    @test isapprox(entropy1, entropy2)

    mps = random_mps_obc(8, 6, 3)
    entropy1 = compute_entropy(mps)
    entropy2 = compute_entropy(mps, 4)
    @test isapprox(entropy1, entropy2)
end

@testset "Compressing an MPS via SVD" begin
    mps = basis_state_obc([1; 2; 1; 2; 1; 2])
    @test_throws AssertionError svd_compress_mps(mps, 0, 0.0)

    # Generate an empty state corresponding to a 0
    mps = random_mps_obc(4, 1, 2)
    mps[1] = 0.0 * mps[1]
    # Build an equal weight superposition of all basis states which is a simple product state
    for i1 = 1:2
        for i2 = 1:2
            for i3 = 1:2
                for i4 = 1:2
                    tmp = basis_state_obc([i1; i2; i3; i4], 2)
                    mps = sum_states(mps, tmp)
                end
            end
        end
    end
    # Now we compress the MPS, which is exactly represented by a MPS with bond dimension one
    res = svd_compress_mps(mps, 1)
    res2 = svd_compress_mps(mps, 0, 1E-6)
    res3 = svd_compress_mps(mps, 1, 1E-6)
    # The norm will be 16, as we added all 16 configurations without normalizing the result
    @test isapprox(calculate_overlap(mps, res), 16.0)
    @test isapprox(calculate_overlap(mps, res2), 16.0)
    @test isapprox(calculate_overlap(mps, res3), 16.0)
end

@testset "Sampling from an MPS" begin
    # Some experiments with product states for which the outcome has to be deterministic
    for N = 4:6
        configuration = rand([1; 2; 3], N)
        mps = basis_state_obc(configuration, 3)
        for i = 1:10
            sample = sample_from_mps(mps)
            @test isequal(sample, configuration)
        end
    end
    # Now prepare the GHZ state
    all_zeros = basis_state_obc(ones(Int64, 4))
    all_ones = basis_state_obc(2 * ones(Int64, 4))
    mps_ghz = sum_states(all_zeros, all_ones)
    gaugeMPS!(mps_ghz, :left, true)
    prob = zeros(2)
    nsamples = 100000
    # Check the outcome for the first 10 samples
    for i = 1:10
        sample = sample_from_mps(mps_ghz)
        @test (sample == [1; 1; 1; 1] || sample == [2; 2; 2; 2])
        if sample == [1; 1; 1; 1]
            prob[1] += 1
        else
            prob[2] += 1
        end
    end
    # For efficiency do not check the other ones anymore
    for i = 11:nsamples
        sample = sample_from_mps(mps_ghz)
        if sample == [1; 1; 1; 1]
            prob[1] += 1
        else
            prob[2] += 1
        end
    end
    @test (sum(abs.(prob ./ nsamples - 0.5 * ones(2))) < 1E-2)
end

@testset "MPO decomposition" begin
    # Check if erroneous inputs are detected
    H = zeros(4, 5)
    @test_throws ArgumentError decompose_into_mpo(H, 2)
    H = zeros(5, 5)
    @test_throws ArgumentError decompose_into_mpo(H, [2; 3])
    # All local dimensions the same
    for i = 1:5
        A = rand(2, 2)
        B = rand(2, 2)
        C = rand(2, 2)
        H = kron(A, B, C)
        mpo = decompose_into_mpo(H, 2)
        H_mpo = contract_virtual_indices(mpo)
        @test isapprox(H, H_mpo)
    end
    # Different local dimensions
    for i = 1:10
        A = rand(2, 2)
        B = rand(4, 4)
        C = rand(3, 3)
        H = kron(A, B, C)
        mpo = decompose_into_mpo(H, [2; 4; 3])
        H_mpo = contract_virtual_indices(mpo)
        @test isapprox(H, H_mpo)
    end
end


text:
using ModelingToolkit, DiffEqOperators, LinearAlgebra, OrdinaryDiffEq
using ModelingToolkit: operation, istree, arguments

# Define some variables
@parameters t x
@variables u(..) v(..)
Dt = Differential(t)
Dx = Differential(x)
Dxx = Differential(x)^2
eqs  = [Dt(u(t,x)) ~ Dxx(u(t,x)),
        Dt(v(t,x)) ~ Dxx(v(t,x))]
bcs = [u(0,x) ~ - x * (x-1) * sin(x),
       v(0,x) ~ - x * (x-1) * sin(x),
       u(t,0) ~ 0.0, u(t,1) ~ 0.0,
       v(t,0) ~ 0.0, v(t,1) ~ 0.0]

domains = [t ∈ IntervalDomain(0.0,1.0),
           x ∈ IntervalDomain(0.0,1.0)]

pdesys = PDESystem(eqs,bcs,domains,[t,x],[u(t,x),v(t,x)])
discretization = MOLFiniteDifference([x=>0.1],t;grid_align=edge_align)
prob = discretize(pdesys,discretization) # This gives an ODEProblem since it's time-dependent
sol = solve(prob,Tsit5())

@parameters t x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dt = Differential(t)
t_min= 0.
t_max = 2.0
x_min = 0.
x_max = 2.
y_min = 0.
y_max = 2.

# 3D PDE
eq  = Dt(u(t,x,y)) ~ Dxx(u(t,x,y)) + Dyy(u(t,x,y))

analytic_sol_func(t,x,y) = exp(x+y)*cos(x+y+4t)
# Initial and boundary conditions
bcs = [u(t_min,x,y) ~ analytic_sol_func(t_min,x,y),
       u(t,x_min,y) ~ analytic_sol_func(t,x_min,y),
       u(t,x_max,y) ~ analytic_sol_func(t,x_max,y),
       u(t,x,y_min) ~ analytic_sol_func(t,x,y_min),
       u(t,x,y_max) ~ analytic_sol_func(t,x,y_max)]

# Space and time domains
domains = [t ∈ IntervalDomain(t_min,t_max),
           x ∈ IntervalDomain(x_min,x_max),
           y ∈ IntervalDomain(y_min,y_max)]
pdesys = PDESystem([eq],bcs,domains,[t,x,y],[u(t,x,y)])

# Method of lines discretization
dx = 0.1; dy = 0.2
discretization = MOLFiniteDifference([x=>dx,y=>dy],t)
prob = ModelingToolkit.discretize(pdesys,discretization)
sol = solve(prob,Tsit5())

# Diffusion in a sphere
@parameters t r
@variables u(..)
Dt = Differential(t)
Dr = Differential(r)
Drr = Dr^2
eq  = Dt(u(t,r)) ~ (1/r^2 * Dr(r^2 * Dr(u(t,r))))
bcs = [u(0,r) ~ - r * (r-1) * sin(r),
       Dr(u(t,0)) ~ 0.0, u(t,1) ~ sin(1)]

domains = [t ∈ IntervalDomain(0.0,1.0),
           r ∈ IntervalDomain(0.0,1.0)]

pdesys = PDESystem(eq,bcs,domains,[t,r],[u(t,r)])
discretization = MOLFiniteDifference([r=>0.1],t)
prob = discretize(pdesys,discretization) # This gives an ODEProblem since it's time-dependent
sol = solve(prob,Tsit5())


text:
using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges

include("utilities.jl")

struct DummyModelAttribute <: MOI.AbstractModelAttribute end
struct DummyEvaluator <: MOI.AbstractNLPEvaluator end
struct DummyVariableAttribute <: MOI.AbstractVariableAttribute end
struct DummyConstraintAttribute <: MOI.AbstractConstraintAttribute end

@testset "Substitution of variables" begin
    mock = MOIU.MockOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()))
    bridged = MOIB.Variable.Vectorize{Float64}(mock)
    x, cx = MOI.add_constrained_variable(bridged, MOI.GreaterThan(1.0))
    fx = MOI.SingleVariable(x)
    y = MOI.get(mock, MOI.ListOfVariableIndices())[1]
    fy = MOI.SingleVariable(y)

    c2fx = MOI.add_constraint(bridged, 2.0fx, MOI.GreaterThan(1.0))
    @test MOI.get(bridged, MOI.ConstraintFunction(), c2fx) ≈ 2.0fx
    @test MOI.get(bridged, MOI.ConstraintSet(), c2fx) == MOI.GreaterThan(1.0)
    # The constant was moved to the set
    @test MOI.get(mock, MOI.ConstraintFunction(), c2fx) ≈ 2.0fy
    @test MOI.get(mock, MOI.ConstraintSet(), c2fx) == MOI.GreaterThan(-1.0)

    MOI.set(bridged, MOI.ConstraintSet(), c2fx, MOI.GreaterThan(2.0))
    @test MOI.get(bridged, MOI.ConstraintFunction(), c2fx) ≈ 2.0fx
    @test MOI.get(bridged, MOI.ConstraintSet(), c2fx) == MOI.GreaterThan(2.0)
    @test MOI.get(mock, MOI.ConstraintFunction(), c2fx) ≈ 2.0fy
    @test MOI.get(mock, MOI.ConstraintSet(), c2fx) == MOI.GreaterThan(0.0)

    MOI.set(bridged, MOI.ConstraintFunction(), c2fx, 3.0fx)
    @test MOI.get(bridged, MOI.ConstraintFunction(), c2fx) ≈ 3.0fx
    @test MOI.get(bridged, MOI.ConstraintSet(), c2fx) == MOI.GreaterThan(2.0)
    @test MOI.get(mock, MOI.ConstraintFunction(), c2fx) ≈ 3.0fy
    @test MOI.get(mock, MOI.ConstraintSet(), c2fx) == MOI.GreaterThan(-1.0)

    MOI.set(bridged, MOI.ConstraintSet(), c2fx, MOI.GreaterThan(4.0))
    @test MOI.get(bridged, MOI.ConstraintFunction(), c2fx) ≈ 3.0fx
    @test MOI.get(bridged, MOI.ConstraintSet(), c2fx) == MOI.GreaterThan(4.0)
    @test MOI.get(mock, MOI.ConstraintFunction(), c2fx) ≈ 3.0fy
    @test MOI.get(mock, MOI.ConstraintSet(), c2fx) == MOI.GreaterThan(1.0)

    MOI.set(bridged, MOI.ObjectiveFunction{typeof(2.0fx)}(), 2.0fx)
    @test MOI.get(mock, MOI.ObjectiveFunction{typeof(2.0fy)}()) ≈ 2.0fy + 2.0

    MOI.set(bridged, DummyModelAttribute(), 2.0fx + 1.0)
    @test MOI.get(bridged, DummyModelAttribute()) ≈ 2.0fx + 1.0
    @test MOI.get(mock, DummyModelAttribute()) ≈ 2.0fy + 3.0

    z = MOI.add_variable(bridged)
    fz = MOI.SingleVariable(z)
    for (attr, index) in [
        (DummyVariableAttribute(), z),
        (DummyConstraintAttribute(), c2fx)]
        MOI.set(bridged, attr, index, 1.0fx)
        @test MOI.get(bridged, attr, index) ≈ 1.0fx
        @test MOI.get(mock, attr, index) ≈ 1.0fy + 1.0

        MOI.set(bridged, attr, [index], [3.0fx + 1.0fz])
        @test MOI.get(bridged, attr, [index])[1] ≈ 3.0fx + 1.0fz
        @test MOI.get(mock, attr, [index])[1] ≈ 3.0fy + 1.0fz + 3.0
    end

    @testset "RawSolver" begin
        MOI.get(bridged, MOI.RawSolver()) === mock
    end

    @testset "HeuristicCallback" begin
        attr = MOI.HeuristicCallback()
        f(callback_data) = nothing
        MOI.set(bridged, attr, f)
        @test MOI.get(bridged, attr) === f
    end

    @testset "CallbackVariablePrimal" begin
        attr = MOI.CallbackVariablePrimal(nothing)
        err = ErrorException("No mock callback primal is set for variable `$z`.")
        @test_throws err MOI.get(bridged, attr, z)
        MOI.set(mock, attr, y, 1.0)
        MOI.set(mock, attr, z, 2.0)
        @test MOI.get(bridged, attr, z) == 2.0
        err = ArgumentError("Variable bridge of type `$(typeof(MOIB.bridge(bridged, x)))` does not support accessing the attribute `$attr`.")
        @test_throws err MOI.get(bridged, attr, x)
    end

    @testset "LazyConstraint" begin
        sub = MOI.LazyConstraint(nothing)
        @test MOI.supports(bridged, sub)
        MOI.submit(bridged, sub, 2.0fx, MOI.GreaterThan(1.0))
        mock.submitted[sub][1][1] ≈ 2.0fy + 2.0
        mock.submitted[sub][1][2] == MOI.GreaterThan(1.0)
    end
    @testset "HeuristicSolution" begin
        sub = MOI.HeuristicSolution(nothing)
        @test MOI.supports(bridged, sub)
        @testset "Non-bridged variable" begin
            MOI.submit(bridged, sub, [z], [1.0])
            mock.submitted[sub][1][1] == [z]
            mock.submitted[sub][1][2] == [1.0]
        end
        @testset "Bridged variable" begin
            err = ErrorException("Cannot substitute `$x` as it is bridged into `$(1.0fy + 1.0)`.")
            @test_throws err MOI.submit(bridged, sub, [x], [1.0])

        end
    end

    nlp_data = MOI.NLPBlockData(
        [MOI.NLPBoundsPair(1.0, 2.0)],
        DummyEvaluator(), false)
    MOI.set(bridged, MOI.NLPBlock(), nlp_data)
    for nlp_data in [MOI.get(bridged, MOI.NLPBlock()), MOI.get(mock, MOI.NLPBlock())]
        @test nlp_data.constraint_bounds == [MOI.NLPBoundsPair(1.0, 2.0)]
        @test nlp_data.evaluator isa DummyEvaluator
        @test nlp_data.has_objective == false
    end
end

# Model not supporting Interval
MOIU.@model(NoIntervalModel,
            (),
            (MOI.EqualTo, MOI.GreaterThan, MOI.LessThan, MOI.Interval),
            (MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives,
             MOI.NormInfinityCone, MOI.NormOneCone, MOI.SecondOrderCone,
             MOI.RotatedSecondOrderCone, MOI.GeometricMeanCone,
             MOI.RelativeEntropyCone, MOI.NormSpectralCone, MOI.NormNuclearCone,
             MOI.PositiveSemidefiniteConeTriangle, MOI.ExponentialCone),
            (MOI.PowerCone, MOI.DualPowerCone),
            (),
            (MOI.ScalarAffineFunction, MOI.ScalarQuadraticFunction),
            (MOI.VectorOfVariables,),
            (MOI.VectorAffineFunction, MOI.VectorQuadraticFunction))

mock = MOIU.MockOptimizer(NoIntervalModel{Float64}())
bridged_mock = MOIB.Constraint.LessToGreater{Float64}(MOIB.Constraint.SplitInterval{Float64}(mock))

@testset "Unsupported constraint attribute" begin
    attr = MOIT.UnknownConstraintAttribute()
    err = ArgumentError(
        "Bridge of type `$MOI.Bridges.Constraint.SplitIntervalBridge{Float64,$MOI.SingleVariable,$MOI.Interval{Float64},$MOI.GreaterThan{Float64},$MOI.LessThan{Float64}}` " *
        "does not support accessing the attribute `$attr`.")
    x = MOI.add_variable(bridged_mock)
    ci = MOI.add_constraint(bridged_mock, MOI.SingleVariable(x),
                            MOI.Interval(0.0, 1.0))
    @test !MOI.supports(bridged_mock, attr, typeof(ci))
    @test_throws err MOI.get(bridged_mock, attr, ci)
end

@testset "Issue #453" begin
    MOI.empty!(bridged_mock)
    MOIU.loadfromstring!(bridged_mock, """
        variables: x
        maxobjective: 3.0x
        c: 2.0x in Interval(1.0, 4.0)
        d: x in LessThan(1.5)
    """)
    x = MOI.get(bridged_mock, MOI.VariableIndex, "x")
    @test isa(x, MOI.VariableIndex)
    c1 = MOI.get(bridged_mock, MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64}}, "c")
    @test isa(c1, MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64}})
    c2 = MOI.get(bridged_mock, MOI.ConstraintIndex, "c")
    @test c1 == c2
    d1 = MOI.get(bridged_mock, MOI.ConstraintIndex{MOI.SingleVariable, MOI.LessThan{Float64}}, "d")
    @test isa(d1, MOI.ConstraintIndex{MOI.SingleVariable, MOI.LessThan{Float64}})
    d2 = MOI.get(bridged_mock, MOI.ConstraintIndex, "d")
    @test d1 == d2
end

MOI.empty!(bridged_mock)

@testset "Name test" begin
    MOIT.nametest(bridged_mock)
    # Test that names are deleted in `MOI.empty!`
    MOIT.nametest(bridged_mock)
end

@testset "Copy test" begin
    MOIT.failcopytestc(bridged_mock)
    MOIT.failcopytestia(bridged_mock)
    MOIT.failcopytestva(bridged_mock)
    MOIT.failcopytestca(bridged_mock)
    MOIT.copytest(bridged_mock, MOIU.Model{Float64}())
end

@testset "Custom test" begin
    model = MOIB.Constraint.SplitInterval{Int}(NoIntervalModel{Int}())
    @test !MOIB.supports_bridging_constraint(model, MOI.VectorAffineFunction{Float64}, MOI.Interval{Float64})

    x, y = MOI.add_variables(model, 2)
    @test MOI.get(model, MOI.NumberOfVariables()) == 2

    f1 = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(3, x)], 7)
    c1 = MOI.add_constraint(model, f1, MOI.Interval(-1, 1))

    @test MOI.get(model, MOI.ListOfConstraints()) == [(MOI.ScalarAffineFunction{Int},MOI.Interval{Int})]
    test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.GreaterThan{Int}, 0)
    test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.Interval{Int}, 1)
    @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Int},MOI.Interval{Int}}())) == [c1]

    f2 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2, -1], [x, y]), 2)
    c2 = MOI.add_constraint(model, f1, MOI.GreaterThan(-2))

    @test MOI.get(model, MOI.ListOfConstraints()) == [(MOI.ScalarAffineFunction{Int},MOI.GreaterThan{Int}), (MOI.ScalarAffineFunction{Int},MOI.Interval{Int})]
    test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.GreaterThan{Int}, 1)
    test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.Interval{Int}, 1)
    @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Int},MOI.Interval{Int}}())) == [c1]
    @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Int},MOI.GreaterThan{Int}}())) == [c2]

    n = 4
    z = MOI.add_variables(model, n)
    scon_indices = MOI.ConstraintIndex{MOI.SingleVariable, MOI.Interval{Int}}[]
    for (i, v) in enumerate([x; y; z])
        f = MOI.SingleVariable(v)
        c = MOI.add_constraint(model, f, MOI.Interval(i, 2i))
        push!(scon_indices, c)

        @test Set(MOI.get(model, MOI.ListOfConstraints())) == Set([(MOI.ScalarAffineFunction{Int},MOI.GreaterThan{Int}), (MOI.ScalarAffineFunction{Int},MOI.Interval{Int}), (MOI.SingleVariable, MOI.Interval{Int})])
        test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.GreaterThan{Int}, 1)
        test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.Interval{Int}, 1)
        test_num_constraints(model, MOI.SingleVariable, MOI.Interval{Int}, i)
        @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Int},MOI.Interval{Int}}())) == [c1]
        @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Int},MOI.GreaterThan{Int}}())) == [c2]
        # The indices should be returned in order of creation
        @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.SingleVariable,MOI.Interval{Int}}())) == scon_indices
    end

    vcon_indices = MOI.ConstraintIndex{MOI.VectorOfVariables, MOI.Nonnegatives}[]
    for (i, v) in enumerate(z)
        f = MOI.VectorOfVariables([v])
        c = MOI.add_constraint(model, f, MOI.Nonnegatives(1))
        push!(vcon_indices, c)

        @test Set(MOI.get(model, MOI.ListOfConstraints())) == Set([(MOI.ScalarAffineFunction{Int},MOI.GreaterThan{Int}), (MOI.ScalarAffineFunction{Int},MOI.Interval{Int}), (MOI.SingleVariable, MOI.Interval{Int}), (MOI.VectorOfVariables, MOI.Nonnegatives)])
        test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.GreaterThan{Int}, 1)
        test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.Interval{Int}, 1)
        test_num_constraints(model, MOI.SingleVariable, MOI.Interval{Int}, n + 2)
        test_num_constraints(model, MOI.VectorOfVariables, MOI.Nonnegatives, i)
        @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Int},MOI.Interval{Int}}())) == [c1]
        @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Int},MOI.GreaterThan{Int}}())) == [c2]
        # The indices should be returned in order of creation
        @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.SingleVariable,MOI.Interval{Int}}())) == scon_indices
        @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.VectorOfVariables,MOI.Nonnegatives}())) == vcon_indices
    end

    @test MOI.is_valid(model, c2)
    MOI.delete(model, c2)

    @test Set(MOI.get(model, MOI.ListOfConstraints())) == Set([(MOI.ScalarAffineFunction{Int},MOI.Interval{Int}), (MOI.SingleVariable, MOI.Interval{Int}), (MOI.VectorOfVariables, MOI.Nonnegatives)])
    test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.GreaterThan{Int}, 0)
    test_num_constraints(model, MOI.ScalarAffineFunction{Int}, MOI.Interval{Int}, 1)
    test_num_constraints(model, MOI.SingleVariable, MOI.Interval{Int}, n + 2)
    test_num_constraints(model, MOI.VectorOfVariables, MOI.Nonnegatives, n)
    @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Int},MOI.Interval{Int}}())) == [c1]
    # The indices should be returned in order of creation
    @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.SingleVariable,MOI.Interval{Int}}())) == scon_indices
    @test (@inferred MOI.get(model, MOI.ListOfConstraintIndices{MOI.VectorOfVariables,MOI.Nonnegatives}())) == vcon_indices
end

@testset "Continuous Linear" begin
    exclude = ["partial_start"] # VariablePrimalStart not supported.
    MOIT.contlineartest(bridged_mock, MOIT.TestConfig(solve=false), exclude)
end

@testset "Show" begin
    @test sprint(show, bridged_mock) == raw"""
    MOIB.Constraint.SingleBridgeOptimizer{MOIB.Constraint.LessToGreaterBridge{Float64,F,G} where G<:MOI.AbstractScalarFunction where F<:MOI.AbstractScalarFunction,MOIB.Constraint.SingleBridgeOptimizer{MOIB.Constraint.SplitIntervalBridge{Float64,F,S,LS,US} where US<:MOI.AbstractSet where LS<:MOI.AbstractSet where S<:MOI.AbstractSet where F<:MOI.AbstractFunction,MOIU.MockOptimizer{NoIntervalModel{Float64}}}}
    with 1 constraint bridge
    with inner model MOIB.Constraint.SingleBridgeOptimizer{MOIB.Constraint.SplitIntervalBridge{Float64,F,S,LS,US} where US<:MOI.AbstractSet where LS<:MOI.AbstractSet where S<:MOI.AbstractSet where F<:MOI.AbstractFunction,MOIU.MockOptimizer{NoIntervalModel{Float64}}}
      with 0 constraint bridges
      with inner model MOIU.MockOptimizer{NoIntervalModel{Float64}}"""
end


text:
module Codim1LimitCycle  # TODO: merge it to Codim1?  Should I?  Maybe
                         # just rename Codim1 to Codim1FixedPoint?

using LinearAlgebra

using ..Continuations: AbstractContinuationProblem, AbstractProblemCache
import ..Continuations: get_prob_cache, get_u0, residual!, residual_jacobian!,
    residual, isindomain

# TimeKind trait:
using ..BifurcationsBase: timekind, Continuous, Discrete
import ..BifurcationsBase: TimeKind

# StateKind trait:
using ..BifurcationsBase: statekind, MutableState, ImmutableState
import ..BifurcationsBase: StateKind

# ContinuationKind trait:
using ..BifurcationsBase: contkind, LimitCycleCont

import ..BifurcationsBase: BifurcationProblem

using Setfield: Lens

include("problem.jl")
include("solver.jl")
include("factories.jl")
include("data.jl")

end  # module


text:
include("utils.jl")
mode = "training"

noise = 0.0
if length(ARGS)==1
    global noise = parse(Float64, ARGS[1])
end
@info noise

# Global parameters
K_CONST =  9.869232667160130e-16 * 86400 * 1e3
ALPHA = 1.0
GRAV_CONST = 9.8    # gravity constant
SRC_CONST = 86400.0 #
n = 15
m = 30
h = 30.0 # meter
NT  = 50
Δt = 1000/NT
# Δt = 20.0 # day

ρw = 501.9
ρo = 1053.0
μw = 0.1
μo = 1.0
K_init = 20.0 .* ones(n,m) # initial guess of permeability
g = GRAV_CONST
φ0 = 0.25 .* ones(n,m)
qw = zeros(NT+1, n, m)
qw[:,9,3] .= 0.005 * (1/h^2)/10.0 * SRC_CONST
qo = zeros(NT+1, n, m)
qo[:,9,28] .= -0.005 * (1/h^2)/10.0 * SRC_CONST
sw0 = zeros(n, m)
K = 20.0 .* ones(n,m) # millidarcy
K[8:10,:] .= 120.0

# E = 6e9
E = 6.e9
ν = 0.35
D = E/(1+ν)/(1-2ν)*[1-ν ν 0;ν 1-ν 0;0 0 (1-2ν)/2]

if mode!="data"
    A = Variable(diagm(0=>ones(3)))
    global D = 6.e9 * spd(A)
end

Z = zeros(n, m)
for j = 1:n
    for i = 1:m
        Z[j, i] = (j-0.5)*h
    end
end

bdnode = bcnode("left | right", m, n, h)

InteractionM = constant(compute_interaction_matrix(m, n, h)')
StiffM = compute_fem_stiffness_matrix(D, m, n, h)/E
StiffM, _ = fem_impose_Dirichlet_boundary_condition_experimental(StiffM, bdnode, m, n, h)
StiffM = StiffM*E

qo, qw = constant(qo), constant(qw)
function porosity(u)
    ε = compute_fvm_mechanics_term(u, m, n, h)/h^2
    ε = reshape(ε, (n, m))
    out = 1 - constant(1 .- φ0) .* exp(-ε)
end

################### solid equations
function solid(Ψ2)
    p = Ψ2 #+ ρw*g*Z
    # p = p - ave_normal(p)
    rhs = InteractionM*reshape(p, (-1,))
    mask = ones(2*(m+1)*(n+1))
    mask[[bdnode;bdnode .+ (m+1)*(n+1)]] .= 0.0
    rhs = rhs .* mask
    u = StiffM\rhs
end
###################

################### fluid  equations
function ave_normal(quantity)
    aa = sum(quantity)
    return aa/(m*n)
end

function fluid(i, u, uold, sw, p)
    # step 0: compute porosity and its rate
    φ = porosity(u)
    φold = porosity(uold)
    dotφ = (φ-φold)/Δt

    # step 1: update p
    λw = sw.*sw/μw
    λo = (1-sw).*(1-sw)/μo
    λ = λw + λo
    q = qw[i] + qo[i] + λw/(λo+1e-16).*qo[i]

    potential_c = (ρw - ρo)*g .* Z

    # Step 2: implicit potential
    Θ = upwlap_op(K * K_CONST, λo, potential_c, h, constant(0.0))

    load_normal = (Θ+q/ALPHA+dotφ) - ave_normal(Θ+q/ALPHA+dotφ)

    p = upwps_op(K * K_CONST, λ, load_normal, tf.zeros_like(p), h, constant(0.0), constant(0))

    # step 3: implicit transport
    sw1 = sat_op2(sw, dotφ, p, K * K_CONST , φ, qw[i], qo[i], μw, μo, sw, Δt, h)
    return sw1, p, φ
end
###################

function simulate()
    ta_u, ta_sw, ta_p, ta_φ, ta_σ = TensorArray(NT+1), TensorArray(NT+1), TensorArray(NT+1), TensorArray(NT+1), TensorArray(NT+1)
    ta_u = write(ta_u, 1, constant(zeros(2*(m+1)*(n+1))))
    ta_u = write(ta_u, 2, constant(zeros(2*(m+1)*(n+1))))
    ta_sw = write(ta_sw, 1, constant(sw0))
    ta_sw = write(ta_sw, 2, constant(sw0))
    ta_p = write(ta_p, 1, constant(zeros(n, m)))
    ta_p = write(ta_p, 2, constant(zeros(n, m)))
    ta_φ = write(ta_φ, 1, constant(φ0))
    ta_φ = write(ta_φ, 2, constant(φ0))
    ta_σ = write(ta_σ, 1, constant(zeros(4*m*n, 3)))
    ta_σ = write(ta_σ, 2, constant(zeros(4*m*n, 3)))
    i = constant(2, dtype=Int32)
    function condition(i, tas...)
        i <= NT
    end
    function body(i, tas...)
        ta_u, ta_sw, ta_p, ta_φ, ta_σ = tas
        u = read(ta_u, i)
        uold = read(ta_u, i-1)
        sw, p = read(ta_sw, i), read(ta_p, i)
        sw, p, φ = fluid(i, u, uold, sw, p)
        unew = solid(p)
        ta_sw = write(ta_sw, i+1, sw)
        ta_p = write(ta_p, i+1, p)
        ta_u = write(ta_u, i+1, unew)
        ta_φ = write(ta_φ, i+1, φ)

        ε = eval_strain_on_gauss_pts(unew, m, n, h)
        σ = ε * D
        ta_σ = write(ta_σ, i+1, σ)
        i+1, ta_u, ta_sw, ta_p, ta_φ, ta_σ
    end
    _, ta_u, ta_sw, ta_p, ta_φ, ta_σ = while_loop(condition, body, [i, ta_u, ta_sw, ta_p, ta_φ, ta_σ])
    out_u, out_sw, out_p, out_φ, out_σ = stack(ta_u), stack(ta_sw), stack(ta_p), stack(ta_φ), stack(ta_σ)
    set_shape(out_u, NT+1, 2*(m+1)*(n+1)), set_shape(out_sw, NT+1, n, m),
    set_shape(out_p, NT+1, n, m), set_shape(out_φ, NT+1, n, m), set_shape(out_σ, NT+1, 4*m*n, 3)
end

u, S2, Ψ2, φ, σ = simulate()
uobs = u[:, 1:m+1]

if mode!="data"
    uobs_ = readdlm("linear_obs.txt")
    uobs_ = uobs_ .* (1. .+ randn(size(uobs_)...)*noise)
    global loss = mean((uobs-uobs_)^2)
end

sess = Session(); init(sess)
if mode=="data"
    writedlm("linear_obs.txt", run(sess, uobs))
    o1, o2, o3, o4, o5 = run(sess, [u, S2, Ψ2, φ, σ])
    error("Stop")
end

@info run(sess, loss)
loss_ = BFGS!(sess, loss)
res = [noise;run(sess, D)[:]]'
open("data/linear.txt", "a") do io
    writedlm(io, res)
end
# writedlm("data/loss_linear$noise.txt", reshape(loss_, :, 1))

# anim = visualize_displacement(50*o1, m, n, h)
# saveanim(anim, "data/linear_disp.gif")
# visualize_obs(o1)
# savefig("data/linear_obs")
# anim = visualize_saturation(o2, m, n, h)
# saveanim(anim, "data/linear_sat.gif")
# anim = visualize_potential(o3, m, n, h)
# saveanim(anim, "data/linear_potential.gif")
# anim = visualize_potential(o4, m, n, h)
# saveanim(anim, "data/linear_pressure.gif")
# anim = visualize_von_mises_stress(o5, m, n, h)
# saveanim(anim, "data/linear_vm.gif")

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
using RDatasets

include("censored.jl")
include("cross_validate.jl")
#=
Initialize example fails because we're passing -1 as the third argument to
```
evaluate(::MultinomialOrdinalLoss, ::Array{Float64,1}, ::Int32)
```
src/evaluate_fit.jl:15
...
initialize.jl:11
=#
# include("initialize.jl")
include("precision_at_k.jl")
include("simple_glrms.jl")
include("fit_rdataset.jl")


text:
module CameraModel

import GeometryTypes

export AbstractCamParams, RgbdCamParams,

       KinectV1Cam, Aug_ICL_NUIM_CamParam, TumCamROS,

       applyT, applyT!

include("camera_params.jl")
include("common_cameras.jl")
include("rigid_transformation.jl")
end # module


text:
function create_request_func(url::String, headers::Dict{String,String})
    return function (method::Symbol, dest::String; body="")
        HTTP.request(method, url * dest; headers=headers, body=body)
    end
end

function get_gateway_url()
    return try
        resp = HTTP.get("$URL/v$API_VERSION/gateway")
        url = get(JSON.parse(String(resp.body)), "url", "")
        isempty(url) && error("failed to decode the message")
        "$url?v=$API_VERSION&encoding=json", nothing
    catch e
        err_msg = "on getting a geteway-url from Discord"
        e isa ErrorException && (err_msg *= ": $(e.msg)")
        nothing, ErrorException(err_msg)
    end
end

throw_on_err(err::Union{Nothing,Any}) = (isnothing(err) || throw(err))

function read_json(io)
    return try
        data = readavailable(io)
        JSON.parse(String(data); dicttype=Dict{Symbol,Any}), nothing
    catch e
        nothing, e
    end
end

function write_json(io, data)
    return try
        write(io, json(data)), nothing
    catch e
        nothing, e
    end
end


text:
using BinDeps

@BinDeps.setup

ENV["JULIA_ROOT"] = abspath(JULIA_HOME, "../../")

# include alias for WinRPM library
libsndfile = library_dependency("libsndfile", aliases=["libsndfile-1"])

provides(AptGet, "libsndfile1-dev", libsndfile)
provides(Pacman, "libsndfile", libsndfile)

@osx_only begin
    using Homebrew
    provides(Homebrew.HB, "libsndfile", libsndfile)
end

@windows_only begin
    using WinRPM
    provides(WinRPM.RPM, "libsndfile1", libsndfile, os = :Windows)
end

@BinDeps.install Dict(:libsndfile => :libsndfile)


text:
struct TimeSection
	name :: String
	start :: Float64
end
export TimeSection

function Base.push!(
	stack :: Vector{TimeSection}, name :: String
)
	push!(stack, TimeSection(name, time()))
end

function Base.append!(
	stack :: Vector{TimeSection}, names :: AbstractVector{String}
)
	now = time()
	return append!(stack, (TimeSection(name, now) for name in names))
end

function close_and_print!(
	stack :: Vector{TimeSection}, name :: String,
	io = stdout
) :: Float64
	now = time()
	top = last(stack)
	@assert top.name == name
	println(io, "$(name) = $(now - top.start)")
	pop!(stack)
	return now
end
export close_and_print!

function close_and_print!(
	stack :: Vector{TimeSection},
	names :: Vector{String},
	io = stdout
) :: Float64
	now = time()
	to_pop_idxs = length(stack):-1:(length(stack) - length(names) + 1)
	for (idx_names, idx_stack) in enumerate(to_pop_idxs)
		if names[idx_names] != stack[idx_stack].name
			error(
				"$(names[idx_names]) != $(stack[idx_stack])." *
				"The sections must be closed in reverse order of opening (FIFO)."
			)
		end
	end
	for e in (stack[i] for i in to_pop_idxs)
		println(io, "$(e.name) = $(now - e.start)")
	end
	deleteat!(stack, reverse(to_pop_idxs)) # idxs must be sorted
	return now
end
export close_and_print!


text:
@with_kw mutable struct StopIntersectionDriver <: DriverModel{LonAccelDirection}
    a::LonAccelDirection = LonAccelDirection(0., 1)
    navigator::RouteFollowingIDM = RouteFollowingIDM()
    intersection::Vector{Lane{Float64}} = Lane{Float64}[]
    intersection_entrances::Vector{Lane{Float64}} = Lane{Float64}[]
    intersection_exits::Vector{Lane{Float64}} = Lane{Float64}[]
    horizon::Float64 = 20.0 # distance from the intersection to start the logic
    stop_delta::Float64 = 0. # precision for stopping slightly before the line
    accel_tol::Float64 = 1e-2 # if |a| < accel_tol then a = 0.
    priorities::Dict{Tuple{LaneTag, LaneTag}, Bool} = Dict{Tuple{LaneTag, LaneTag}, Bool}()
    # states
    have_wait_list::Bool = false
    wait_list::Vector{Int64} = Int64[]
    priority::Bool = false
    stop::Bool = false # switch to true when stopped at the intersection, stays true until the end
    n_yield::Int64 = 0 # number of vehicles to yield to
end

get_name(::StopIntersectionDriver) = "StopIntersectionDriver"
Base.rand(model::StopIntersectionDriver) = model.a

function AutomotiveDrivingModels.set_desired_speed!(model::StopIntersectionDriver, v_des::Float64)
    model.navigator.v_des = v_des
    model
end

function AutomotiveDrivingModels.reset_hidden_state!(model::StopIntersectionDriver)
    model.have_wait_list = false
    model.wait_list = Int64[]
    model.priority = false
    model.stop = false
    model.n_yield = 0
    model
end

function AutomotiveDrivingModels.observe!(model::StopIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    AutomotiveDrivingModels.observe!(model.navigator, scene, roadway, egoid) # set the direction
    dir = model.navigator.dir
    a_lon =0.
    a_lon_idm = model.navigator.a
    if isempty(model.intersection) || model.priorities[(model.navigator.route[1].tag,model.navigator.route[end].tag)] # no intersection, go
        a_lon = model.navigator.a
    else
        if !model.priority && !model.stop # reach stop line
            a_lon = min(a_lon_idm, stop_at_end(model, ego, roadway))
        elseif !model.priority && model.stop # wait
            a_lon = -model.navigator.d_max # negative to make sure v stays 0
        else
            a_lon = a_lon_idm # just idm
        end
    end
    # if !model.have_wait_list
    #     grow_wait_list!(model, scene, roadway, egoid)
    #     model.have_wait_list = true
    # end
    # if !model.stop
    #     grow_wait_list!(model, scene, roadway, egoid)
    #     update_stop!(model, ego, roadway)
    # else
    #     ungrow_wait_list!(model, scene, roadway, egoid)
    # end
    if !model.stop
        update_stop!(model, ego, roadway)
    end
    if !model.priority
        grow_wait_list!(model, scene, roadway, egoid)
        ungrow_wait_list!(model, scene, roadway, egoid)
        update_priority!(model, scene, roadway, egoid)
    end

    # println("stop ", model.stop)
    # println("priority ", model.priority)
    # println("n_yield ", model.n_yield)
    # println("a_lon ", a_lon)
    # println("v ", ego.state.v)
    model.a = LonAccelDirection(a_lon + 2*model.accel_tol*randn(), dir) # add noise to break ties #XXX remove with priorities
    d_stop = get_dist_to_end(ego, roadway)
    # println("veh ", ego.id, " wait_list ", model.wait_list, " stop ", model.stop, " priority ", model.priority, " dstop ", d_stop, " a_lon ", a_lon, " v ", ego.state.v )
    # println("veh ", ego.id, " taking action ", model.a)
    model
end

"""
If other vehicles reaches the intersection before the ego vehicle they are added to the waitlist
"""
function grow_wait_list!(model::StopIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    ego_dist = get_dist_to_end(ego, roadway)
    ego_tts = tts(ego_dist, ego.state.v)
    n_yield = 0
    for veh in scene
        if veh.id == egoid
            continue
        end
        if is_at_intersection(model, veh, roadway)
            veh_dist = get_dist_to_end(veh, roadway)
            veh_tts = tts(veh_dist, veh.state.v)
            veh_lane = get_lane(roadway, veh)
            if veh_tts < ego_tts || veh_lane ∈ get_exit_lanes(model.intersection_entrances, roadway)
            # if veh_dist < ego_dist #|| is_exiting(model, veh, roadway)
                n_yield += 1
                if !(veh.id ∈ model.wait_list)
                    push!(model.wait_list, veh.id)
                end
            end
        end
    end
    model.n_yield = length(model.wait_list)
end

"""
Remove from the wait list vehicles that are now exiting the intersection
"""
function ungrow_wait_list!(model::StopIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    ego_dist = get_dist_to_end(ego, roadway)
    n = length(model.wait_list)
    to_remove = Int64[]
    for i=1:n
        vehid = model.wait_list[i]
        veh_ind = findfirst(vehid, scene)
        if veh_ind == nothing
            push!(to_remove, i)
            continue
        end
        veh = scene[veh_ind]
        veh_dist = get_dist_to_end(veh, roadway)
        lane = get_lane(roadway, veh)
        # if !(lane ∈ model.intersection_entrances) &&  veh.state.posF.s > veh.def.length # vehicle has exited for more than 1 car length
        if is_exiting(model, veh, roadway) || (ego_dist < veh_dist && is_at_intersection(model, veh, roadway))
            push!(to_remove, i)
            model.n_yield -= 1
        end
    end
    deleteat!(model.wait_list, to_remove)
end

"""
Robust function to compute the time to drive dist at speed v
"""
function tts(dist::Float64, v::Float64)
    if isapprox(dist, 0, atol=1e-1)
        return 0
    end
    return dist/v
end

function is_intersection_cleared(model::StopIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int, delta::Float64=2.0)
    for veh in scene
        if veh.id == egoid
            continue
        end
        if !(is_exiting(model, veh, roadway) || is_at_intersection(model, veh, roadway))
            return false
        end
    end
    return true
end

function is_exiting(model::StopIntersectionDriver, veh::Vehicle, roadway::Roadway)
    lane = get_lane(roadway, veh)
    return lane ∈ model.intersection_exits
end

"""
return true if veh is before the intersection
"""
function is_at_intersection(model::StopIntersectionDriver, veh::Vehicle, roadway::Roadway)
    lane = get_lane(roadway, veh)
    return lane ∈ model.intersection_entrances
end

"""
Check if the ego car has priority or not
"""
function update_priority!(model::StopIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    model.priority = isempty(model.wait_list) && model.stop ||
                       !(get_lane(roadway, ego) ∈ model.intersection_entrances)
end

## Fancier priority model
function check_priority(model::StopIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    lane = get_lane(roadway, ego)
    route = model.navigator.route
    next_lane = intersect(lane.exits, model.navigator.route)

end

max_stars_repo_name: MaximeBouton/AutomotivePOMDPs

max_issues_repo_name: MaximeBouton/AutomotivePOMDPs

max_forks_repo_name: MaximeBouton/AutomotivePOMDPs