text:
"""
    nh_pressure_buoy(param_set, a_up, b_up, ρ0, asp_ratio, bcs)

Returns the value of perturbation pressure gradient
for updraft i following [He2020](@cite), given:

 - `a_up`: updraft area
 - `b_up`: updraft buoyancy
 - `ρ0`: reference density
 - `asp_ratio`: the specific aspect ratio of the updraft
 - `bcs`: a `NamedTuple` of BCs with updraft area fraction and buoyancy entries
"""
function nh_pressure_buoy(param_set::APS, a_up, b_up, ρ0, asp_ratio::FT, bcs) where {FT}
    Ifb = CCO.InterpolateC2F(; bcs.b_up...)
    Ifa = CCO.InterpolateC2F(; bcs.a_up...)

    α_b::FT = CPEDMF.α_b(param_set)
    α₂ = FT(0)

    return @. Int(Ifa(a_up) > 0) * -α_b / (1 + α₂ * asp_ratio^2) * ρ0 * Ifa(a_up) * Ifb(b_up)
end

"""
    nh_pressure_adv(param_set, updraft_top, a_up, ρ0, w_up, bcs)

Returns the value of perturbation pressure gradient
for updraft i following [He2020](@cite), given:

 - `updraft_top`: the height of the updraft in the previous timestep
 - `a_up`: updraft area
 - `ρ0`: reference density
 - `w_up`: updraft vertical velocity
 - `bcs`: a `NamedTuple` of BCs with updraft velocity and area fraction entries
"""
function nh_pressure_adv(param_set::APS, updraft_top::FT, a_up, ρ0, w_up, bcs) where {FT}

    Ifa = CCO.InterpolateC2F(; bcs.a_up...)
    Ifc = CCO.InterpolateF2C()
    ∇ = CCO.DivergenceC2F(; bcs.w_up...)
    α_a::FT = CPEDMF.α_a(param_set)
    wvec = CC.Geometry.WVector

    return @. Int(Ifa(a_up) > 0) * ρ0 * Ifa(a_up) * α_a * w_up * ∇(wvec(Ifc(w_up)))
end;

"""
    nh_pressure_drag(param_set, updraft_top, a_up, ρ0, w_up, w_en, bcs)

Returns the value of perturbation pressure gradient
for updraft i following [He2020](@cite), given:

 - `updraft_top`: the height of the updraft in the previous timestep
 - `a_up`: updraft area
 - `ρ0`: reference density
 - `w_up`: updraft vertical velocity
 - `w_en`: environment vertical velocity
 - `bcs`: a `NamedTuple` of BCs with an updraft area fraction entry
"""
function nh_pressure_drag(param_set::APS, updraft_top::FT, a_up, ρ0, w_up, w_en, bcs) where {FT}

    Ifa = CCO.InterpolateC2F(; bcs.a_up...)
    α_d::FT = CPEDMF.α_d(param_set)
    H_up_min::FT = CPEDMF.H_up_min(param_set)

    # drag as w_dif and account for downdrafts
    return @. Int(Ifa(a_up) > 0) * -1 * ρ0 * Ifa(a_up) * α_d * (w_up - w_en) * abs(w_up - w_en) /
              max(updraft_top, H_up_min)
end;


text:
module LabVIEW0

"""
    lv_dir()

Locate the LabVIEW folder within the project
"""
function lv_dir()
    lv_path = @__DIR__
    println(lv_path)
    return lv_path
end
export lv_dir

using Reexport
@reexport using LVServer

end

ext: jl


text:
@testset "DescentGraph tests" begin
    @testset "Vector{Vector} constructor tests" begin
        data = [[0., 0.],
                [0., 1.]]
        graph = DescentGraph(data, 1)
        @test size(graph.indices) == (1, 2)
        @test size(graph.distances) == (1, 2)
        @test graph.indices == [2 1]
        @test graph.distances == [1. 1.]

        for np = 20:10:50, k = 1:2:10
            data = [rand(5) for _ in 1:np]
            graph = DescentGraph(data, k)
            @test size(graph.indices) == (k, np)
            @test size(graph.distances) == (k, np)
        end
    end
    @testset "Matrix constructor tests" begin
        data = [0. 0.; 0. 1.]
        @test DescentGraph(data, 1) isa DescentGraph
    end
    @testset "eltype stability tests" begin
        # Float32
        data = [rand(Float32, 10) for _ in 1:100]
        @inferred DescentGraph(data, 5)
        # Integer
        data = [rand([0, 1], 10) for _ in 1:100]
        @inferred DescentGraph(data, 5, Hamming())
    end
end

@testset "Distances tests" begin
    @testset "Euclidean" begin
        data = [[0., 0., 0.],
                [0., 0., 1.],
                [0., 1., 0.],
                [0., 1., 1.],
                [1., 0., 0.],
                [1., 0., 1.],
                [1., 1., 0.],
                [1., 1., 1.]]
        _3nn = brute_knn(data, Euclidean(), 3)
        true_3nn = transpose([2 3 5;
                    1 4 6;
                    1 4 7;
                    2 3 8;
                    1 6 7;
                    2 5 8;
                    3 5 8;
                    4 6 7])
        @testset "brute knn test" begin
            @test sort(getindex.(_3nn, 1), dims=1) == sort(true_3nn, dims=1)
        end
        @testset "basic_nndescent tests" begin
            graph = DescentGraph(data, 3)
            @test_skip sort(graph.indices, dims=1) == sort(true_3nn, dims=1)
        end
    end
    @testset "CosineDist" begin
        data = [rand(50) for _ in 1:10]
        true_3nn = brute_knn(data, CosineDist(), 3)
        desc_3nn = DescentGraph(data, 3, CosineDist()).indices
        @test_skip getindex.(true_3nn, 1) == desc_3nn
    end
    @testset "Hamming" begin
        data = [rand([0, 1], 50) for _ in 1:10]
        true_3nn = brute_knn(data, Hamming(), 3)
        desc_3nn = DescentGraph(data, 3, Hamming()).indices
        @test_skip getindex.(true_3nn, 1) == desc_3nn
    end
end

@testset "search tests" begin
    data = [[0., 0., 0.],
            [0., 0., 1.],
            [0., 1., 0.],
            [0., 1., 1.],
            [1., 0., 0.],
            [1., 0., 1.],
            [1., 1., 0.],
            [1., 1., 1.]]
    queries = [[0., 0., 0.4],
               [0., 1., 0.4]]

    graph = DescentGraph(data, 3)
    true_inds = [1 3;
                 2 4]
    true_dists = [.4 .4;
                  .6 .6]
    inds, dists = search(graph, queries, 2, 4)
    @test all(inds .== true_inds)
    @test all(dists .== true_dists)

    @testset "search matrix queries tests" begin
        queries = [0. 0.; 0. 1.; 0.4 0.4]
        inds, dists = search(graph, queries, 2, 4)
        @test all(inds .== true_inds)
        @test all(dists .== true_dists)
    end

end


text:
function test_print(name, value)
    println("=================================================================")
    println(name)
    println(value)
    println("-----------------------------------------------------------------")
end # test_print


text:
using Term
using Test
using Suppressor

include("__test_utils.jl")

using TimerOutputs: TimerOutputs, @timeit
const TIMEROUTPUT = TimerOutputs.TimerOutput()

tprint("\n[bold blue]Runing all tests measuring timing and allocations\n")

# ? 0  - misc
tprint("[bold green]Running: '00_misc.jl' ")
@time @timeit_include("00_misc.jl")

# ? 1  - text utils
tprint("[bold green]Running: '01_test_text_utils.jl' ")
@time @timeit_include("01_test_text_utils.jl")

 # ? 2 ansi
tprint("\n\n[bold green]Running: '02_test_ansi.jl")  # ansi & col' ")
@time @timeit_include("02_test_ansi.jl")

 # ? 3 measure
tprint("\n\n[bold green]Running: '03_test_measure.jl' ")
@time @timeit_include("03_test_measure.jl")

 # ? 4 markup
tprint("\n\n[bold green]Running: '04_test_markup_and_style.jl' ")
@time @timeit_include("04_test_markup_and_style.jl")

 # ? 5 macros
tprint("\n\n[bold green]Running: '05_test_macros.jl' ")
@time @timeit_include("05_test_macros.jl")

 # ? 6 renderables
tprint("\n\n[bold green]Running: '06_test_renderables.jl' ")
@time @timeit_include("06_test_renderables.jl")

# ? 7 panel
tprint("\n\n[bold green]Running: '07_test_panel.jl' ")
@time @timeit_include("07_test_panel.jl")

# ? 8 layout
tprint("\n\n[bold green]Running: '08_test_layout.jl' ")
@time @timeit_include("08_test_layout.jl")

#  ? 9 inspect
tprint("\n\n[bold green]Running: '09_test_inspect.jl' ")
@time @timeit_include("09_test_inspect.jl")

# ? 11 console
tprint("\n\n[bold green]Running: '11_test_console.jl' ")
@time @timeit_include("11_test_console.jl")

# # ? 12 logging
tprint("\n\n[bold green]Running: '12_test_logging.jl' ")
@time @timeit_include("12_test_logging.jl")

# ? 13 box
tprint("\n\n[bold green]Running: '13_test_box.jl' ")
@time @timeit_include("13_test_box.jl")

# ? 14 highlight
tprint("\n\n[bold green]Running: '14_test_highlight.jl' ")
@time @timeit_include("14_test_highlight.jl")

# ? 15 progress
tprint("\n\n[bold green]Running: '15_test_progress.jl' ")
@time @timeit_include("15_test_progress.jl")

# ? 16 Tree
tprint("\n\n[bold green]Running: '16_test_tree.jl' ")
@time @timeit_include("16_test_tree.jl")

# ? 999 EXAMPLES
tprint("\n\n[bold green]Running: '999_test_examples.jl' ")
@time @timeit_include("999_test_examples.jl")

# ? 99 ERRORS
tprint("\n\n[bold green]Running: '99_test_errors.jl' ")
@time @timeit_include("99_test_errors.jl")

show(TIMEROUTPUT; compact = true, sortby = :firstexec)
println("\n")


text:
using Revise
using VCFTools
using MendelImpute
using GeneticVariation
using Random
using JLSO
using LinearAlgebra

function filter_and_mask(data, maf)
    # filter chromosome data for unique snps
    #println("filtering for unique snps")
    #full_record_index = .!find_duplicate_marker(data)
    #@time VCFTools.filter(data, full_record_index, 1:nsamples(data),
    #    des = "full.uniqueSNPs.vcf.gz")

    # summarize data
    println("summarizing data")
    total_snps, samples, _, _, _, maf_by_record, _ = gtstats("full.uniqueSNPs.vcf.gz")

    # keep snps with at least 5 copies of minor alleles
    snps_tokeep = findall(x -> x ≥ 5 / 2samples, maf_by_record)

    # generate target panel with all snps
    println("generating complete target panel")
    n = 1000
    sample_idx = falses(samples)
    sample_idx[1:n] .= true
    shuffle!(sample_idx)
    @time VCFTools.filter("full.uniqueSNPs.vcf.gz", snps_tokeep,
        sample_idx, des = "target.full.vcf.gz", allow_multiallelic=false)

    #also generate reference panel without target samples
    println("generating reference panel without target samples")
    @time VCFTools.filter("full.uniqueSNPs.vcf.gz", snps_tokeep,
        .!sample_idx, des = "ref.excludeTarget.vcf.gz", allow_multiallelic=false)

    # generate target file with 1000 samples and typed snps with certain maf
    println("generating target file with typed snps only")
    my_maf = findall(x -> x > maf, maf_by_record)
    p = length(my_maf)
    record_idx = falses(total_snps)
    record_idx[my_maf] .= true
    @time VCFTools.filter("full.uniqueSNPs.vcf.gz", record_idx, sample_idx,
        des = "target.typedOnly.maf$maf.vcf.gz", allow_multiallelic=false)

    # unphase and mask 1% entries in target file
    println("unphasing and masking entries in target file with typed snps only")
    masks = falses(p, n)
    missingprop = 0.005
    for j in 1:n, i in 1:p
        rand() < missingprop && (masks[i, j] = true)
    end
    @time mask_gt("target.typedOnly.maf$maf.vcf.gz", masks,
        des="target.typedOnly.maf$maf.masked.vcf.gz", unphase=true)

    # compress reference file to jlso format
    max_d   = [1000]
    reffile = "ref.excludeTarget.vcf.gz"
    tgtfile = "target.typedOnly.maf$maf.masked.vcf.gz"
    H, H_sampleID, H_chr, H_pos, H_ids, H_ref, H_alt = convert_ht(Bool, reffile, trans=true, save_snp_info=true, msg="importing reference data...")
    X, X_sampleID, X_chr, X_pos, X_ids, X_ref, X_alt = VCFTools.convert_gt(UInt8, tgtfile, trans=true, save_snp_info=true, msg = "Importing genotype file...")
    for d in max_d
        outfile = "ref.excludeTarget.maxd$d.jlso"
        @time compress_haplotypes(H, X, outfile, X_pos, H_sampleID, H_chr, H_pos, H_ids, H_ref, H_alt, d, 0, 0.0)
    end
end
Random.seed!(2020)
data = "/mnt/AppRun/msprime/full3.vcf"
maf = 0.05
@time filter_and_mask(data, maf)

max_stars_repo_name: Hua-Zhou/MendelImpute


text:
include("../src/NMRSpectraSimulator.jl")
import .NMRSpectraSimulator

using LinearAlgebra
using FFTW

import PlotlyJS
using Plots; plotly()

import OffsetArrays
import Interpolations

import BSON
import JSON

import Random
Random.seed!(25)

# simulation compounds. # paths.
plot_title = "resonance groupings"

# name of molecules to simulate.
# Limit to only one compound in this tutorial since we just want to visualize the different resonance groups for one compound.
#molecule_names = ["L-Histidine";]

# molecule_names = ["D-(+)-Glucose";]
# root_folder = "/home/roy/MEGAsync/outputs/NMR/experiments/BMRB-500-100mM"
# project_name = "D-(+)-Glucose"

# molecule_names = ["L-Glutamine";]
# root_folder = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "bmse000038_glutamine"
#
molecule_names = ["L-Leucine";]
root_folder = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
project_name = "gissmo_900_leucine"
#
# molecule_names = ["L-Leucine";]
# root_folder = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "gissmo_900_leucine_1"
#
# molecule_names = ["L-Serine";]
# root_folder = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "bmse000867_serine"
#
# molecule_names = ["L-Isoleucine";]
# root_folder = "/home/roy/MEGAsync/outputs/NMR/experiments/BMRB-500-100mM"
# project_name = "L-Isoleucine"

# get mapping from molecule names to their spin system info json files.
H_params_path = "/home/roy/Documents/repo/NMRData/input/coupling_info"
dict_compound_to_filename = JSON.parsefile("/home/roy/Documents/repo/NMRData/input/compound_mapping/select_compounds.json")

# where the bson file is located.
project_path = joinpath(root_folder, project_name)
load_path = joinpath(project_path, "experiment.bson")

# where to save the resultant plot. Warning: This script will create the following path if it doesn't it exist.
save_folder = joinpath(project_path, "plots")
isdir(save_folder) || mkpath(save_folder)

# spin-Hamiltonian-related.
tol_coherence = 1e-2 # resonances are pairs of eigenvalues of the Hamiltonian that have quantum numbers that differ by -1. This is the tolerance away from -1 that is allowed.
α_relative_threshold = 0.05 # resonances with relative amplitude less than this factor compared to the maximum resonance in the spin group will be removed. Set to 0.0 to see every single resonance component.
Δc_partition_radius = 0.17 # determines how many resonances get grouped together. Larger number means less groups and thus more resonances per group.
SH_config_path = "/home/roy/Documents/repo/NMRData/input/SH_configs/select_compounds_SH_configs.json"

# surrogate-related.
# default values. The alternative is to load from a config file.
Δr_default = 1.0 # the samples used to build the surrogate is taken every `Δr` radian on the frequency axis. Decrease for improved accuracy at the expense of computation resources.
Δκ_λ_default = 0.05 # the samples used to build thes urrogate for κ_λ are taken at this sampling spacing. Decrease for improved accuracy at the expense of computation resources.
Δcs_max_scalar_default = 0.2 # In units of ppm. interpolation border that is added to the lowest and highest resonance frequency component of the mixture being simulated.
κ_λ_lb_default = 0.5 # interpolation lower limit for κ_λ.
κ_λ_ub_default = 2.5 # interpolation upper limit for κ_λ.
λ0 = 3.4
surrogate_config_path = "/home/roy/Documents/repo/NMRData/input/surrogate_configs/select_compounds_surrogate_configs.json"

dummy_SSFID = NMRSpectraSimulator.SpinSysParamsType1(0.0) # level 2 model.

##### end user inputs.

### load block.
dic = BSON.load(load_path)
fs = dic[:fs]
SW = dic[:SW]
ν_0ppm = dic[:ν_0ppm]
λ_0ppm = dic[:λ_0ppm]

#ν_0ppm += 42e6

hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs
ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

# css_sys = [[1.1696, 1.1696, 1.1696, 3.64079, 3.64079]]
# J_vals_sys = [[-12.5, -12.5, 7.094311, 7.094311, -12.5, 7.094311, 7.094311, 7.094311, 7.094311, -12.4]]
# J_inds_sys = [[(1, 2), (1, 3), (1, 4), (1, 5), (2, 3), (2, 4), (2, 5), (3, 4), (3, 5), (4, 5)]]
# N_spins_sys = [5;]
#
# intermediates_sys = NMRSpectraSimulator.prepcouplingalgorithm(N_spins_sys)
#
# i = 1
# αs, Ωs, coherence_mat_sys, eig_vals_sys, Q_sys, coherence_state_pairs_sys,
#        H_sys, H1, H2, M_sys = NMRSpectraSimulator.evalSCalgorithm!(css_sys[i],
#        J_vals_sys[i], J_inds_sys[i],
#        intermediates_sys[i],
#        ppm2hzfunc;
#        tol_coherence = 0.1)

name = molecule_names[1]
base_path = H_params_path

config_path = SH_config_path

load_path = joinpath(base_path, dict_compound_to_filename[name]["file name"])
H_IDs, H_css, J_IDs, J_vals = NMRSpectraSimulator.loadcouplinginfojson(load_path)

db_dict = JSON.parsefile(config_path)
dict = db_dict[name]

tol_coherence = dict["coherence tolerance"]
α_relative_threshold = dict["relative amplitude threshold"]
Δc_partition_radius = dict["maximum Δc deviation"]
simple_coherence_atol = dict["simple coherence absolute tolerance"]

J_inds_sys, J_inds_sys_local, J_IDs_sys, J_vals_sys, H_inds_sys,
    cs_sys, H_inds_singlets, cs_singlets, H_inds, J_inds,
    g = NMRSpectraSimulator.setupcsJ(H_IDs, H_css, J_IDs, J_vals)

N_spins_singlet = length.(H_inds_singlets)

N_spins_sys = collect( length(cs_sys[m]) for m = 1:length(cs_sys) )
intermediates_sys = NMRSpectraSimulator.prepcouplingalgorithm(N_spins_sys)

# css_sys becomes cs_sys
# cs_singlets_compact becomes cs_singlets.

αs_inp, Ωs_inp, coherence_mat_sys, eig_vals_sys, Q_sys,
coherence_state_pairs_sys, H_sys, states_sys, ms_sys,
M_sys = NMRSpectraSimulator.setupcompoundspectrum!(cs_sys,
    J_vals_sys, J_inds_sys_local, intermediates_sys,
    ppm2hzfunc, cs_singlets, N_spins_singlet, fs, SW;
    tol_coherence = tol_coherence)
#
T = Float64
αs_spin_sys = copy(αs_inp)
Ωs_spin_sys = copy(Ωs_inp)
αs_singlets = Vector{T}(undef, 0)
Ωs_singlets = Vector{T}(undef, 0)

αs, Ωs, part_inds_compound,
Δc_m_compound, Δc_bar = NMRSpectraSimulator.partitionresonances(coherence_state_pairs_sys,
ms_sys, αs_spin_sys, Ωs_spin_sys;
α_relative_threshold = α_relative_threshold,
Δc_partition_radius = Δc_partition_radius,
simple_coherence_atol = simple_coherence_atol)

M2 = collect( sum(ms_sys[1][i]) for i = 1:length(ms_sys[1]) )

cp = coherence_state_pairs_sys[1]
as = αs[1]
Fs = Ωs[1]
Ps = hz2ppmfunc.( Fs ./ (2*π) )

cb = Δc_bar[1]
cbr = collect( round.(cb[i], digits = 2) for i = 1:length(cb))

inds = part_inds_compound[1]
energies = collect( sum(as[inds[k]]) for k = 1:length(inds) )

import Statistics
freqs = collect( collect(Ps[inds[k]]) for k = 1:length(inds) )
avg_freqs = collect( Statistics.mean(Ps[inds[k]]) for k = 1:length(inds) )
[avg_freqs energies 1:length(energies)]
println("[avg_freqs energies 1:length(energies)]")
display([avg_freqs energies 1:length(energies)])

#[freqs 1:length(freqs)]

#inds2 = [2; 3; 4; 13; 14; 11; 12]
#hcat(cbr[inds2]...) # 4 peaks.

inds1 = [2; 3; 8]
hcat(cbr[inds1]...) # 3 peaks.

[hcat(cbr[collect(1:length(inds))]...); collect(1:length(inds))'] # all groups. last row is index.

#5, 6, 7 are the weak peaks for 1.

#
# [cp[inds[1]] as[inds[1]]]
#

#
#
# k_fin = findfirst(xx->isapprox(xx,0,atol = 1e-5), energies)
# if typeof(k_fin) == Nothing
#     k_fin = length(cbr)
# end
# [hcat(cbr[1:k_fin]...); collect(1:k_fin)']
#
# [energies[1:k_fin] 1:length(energies[1:k_fin])]
# #[freqs[1:k_fin] 1:length(freqs[1:k_fin])]
#
# flags = trues(length(cbr))
# flags[1:k_fin] .= false
# cb2_not = cbr[flags]
#
# tmp1 = collect( sum(cbr[k][1:3]) for k = 1:length(cbr) )
# tmp2 = collect( sum(cbr[k][4:5]) for k = 1:length(cbr) )
# [tmp1 tmp2 flags 1:length(cbr)]
#
# # cp_groups = collect( cp[inds[k]] for k = 1:length(inds) )
# #
# # # check.
# # c1 = cp_groups[1]
# # m = ms_sys[1]
# # Δc_m_compound[1][inds[1]] - collect( m[c1[i][1]] - m[c1[i][2]] for i = 1:length(c1) )

@assert 1==21

include("plot.jl")


text:
using Test, SpikeSynchrony

@testset "SPIKE distance" begin include("SPIKEdistance_tests.jl") end
@testset "vanRossum" begin include("vanRossum_tests.jl") end


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "basiclu"
version = v"2.1.0"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/ERGO-Code/basiclu.git", "fe63cd34ab0259329e5a979df7f19ff0de59cfae")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/basiclu/
ARGS=(CC99="cc" CFLAGS="-std=c99" CPPFLAGS="-DBASICLU_NOTIMER")
if [[ ${target} == *mingw32* ]]; then
    ARGS+=(UNAME="Windows")
else
    ARGS+=(UNAME="$(uname)")
fi
make "${ARGS[@]}"
cp lib/* "${libdir}/."
mkdir -p ${includedir}
cp include/* "${includedir}/."
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms(;experimental=true)

# The products that we will ensure are always built
products = [
    LibraryProduct("libbasiclu", :libbasiclu)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="CompilerSupportLibraries_jll", uuid="e66e0078-7015-5450-92f7-15fbd957f2ae"))
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6")

max_stars_repo_name: jameskermode/Yggdrasil

max_issues_repo_name: jameskermode/Yggdrasil

max_forks_repo_name: jameskermode/Yggdrasil


text:
using DSGE, JLD2
import Test: @test, @testset
import DSGE: @test_matrix_approx_eq
import DataStructures: OrderedDict

### Model
m = OneAssetHANK()
endo = m.endogenous_states
### Model indices
@testset "Model indices" begin
    # Endogenous states
    @test n_states(m) == length(endo[:value_function]) + length(endo[:inflation]) + length(endo[:distribution]) + length(endo[:monetary_policy]) + 5
    @test get_setting(m, :n_jump_vars) == length(endo[:value_function]) + length(endo[:inflation])
    @test get_setting(m, :n_state_vars) == length(endo[:distribution]) - get_setting(m, :n_state_vars_unreduce) + length(endo[:monetary_policy])

    # Exogenous shocks
    @test n_shocks_exogenous(m) == 1

    # Expectation shocks
    @test n_shocks_expectational(m) == get_setting(m, :n_jump_vars)
end

### Steady state
#=
# Load test values. Only using params for now, hard-coded other values,
# but we list the other sets of parameters here if user wants to
# edit this test file.
params = load("../../../test_outputs/one_asset_hank/saved_params.jld2", "params")
# grids = load("../../../test_outputs/one_asset_hank/saved_params.jld2", "grids")
# init_params = load("../../../test_outputs/one_asset_hank/saved_params.jld2", "init_params")
# approx_params = load("../../../test_outputs/one_asset_hank/saved_params.jld2", "approx_params")
# red_params = load("../../../test_outputs/one_asset_hank/saved_params.jld2", "red_params")
vars_ss_mat = load("../../../test_outputs/one_asset_hank/vars_ss.jld2")

# Update parameters
m[:coefrra] = params[:coefrra]
m[:frisch] = params[:frisch]
m[:meanlabeff] = params[:meanlabeff]
m[:maxhours] = params[:maxhours]
m[:ceselast] = params[:ceselast]
m[:labtax] = params[:labtax]
m[:govbondtarget] = params[:govbondtarget]
m[:labdisutil] = params[:labdisutil]
m[:lumptransferpc] = params[:lumptransferpc]
m[:priceadjust] = params[:priceadjust]
m[:taylor_inflation] = params[:taylor_inflation]
m[:taylor_outputgap] = params[:taylor_outputgap]
m[:govbcrule_fixnomB] = params[:govbcrule_fixnomB]
m[:σ_MP] = params[:ssigma_MP]
m[:θ_MP] = params[:ttheta_MP]
=#
# Update settings
update!(m.settings[:r0], Setting(:r0, .005))
update!(m.settings[:rmax], Setting(:rmax, .08 ))
update!(m.settings[:rmin], Setting(:rmin, .001))
update!(m.settings[:ρ0], Setting(:ρ0, .02))
update!(m.settings[:ρmax], Setting(:ρmax, .05))
update!(m.settings[:ρmin], Setting(:ρmin, .005))
update!(m.settings[:iterate_r], Setting(:iterate_r, false))
update!(m.settings[:iterate_ρ], Setting(:iterate_ρ, true))
update!(m.settings[:I], Setting(:I, 100))
update!(m.settings[:amin], Setting(:amin, 0.))
update!(m.settings[:amax], Setting(:amax, 40.))
update!(m.settings[:agridparam], Setting(:agridparam, 1))
#update!(m.settings[:a], Setting(:a, construct_asset_grid(get_setting(m, :I), get_setting(m, :agridparam), get_setting(m, :amin), get_setting(m, :amax))))
update!(m.settings[:J], Setting(:J, 2))
update!(m.settings[:ygrid_combined], Setting(:ygrid_combined, [0.2, 1.]))
update!(m.settings[:ymarkov_combined], Setting(:ymarkov_combined, [-.5 .5; .0376 -.0376]))
#update!(m.settings[:g_z] , Setting(:g_z, compute_stationary_income_distribution(get_setting(m, :ymarkov_combined), get_setting(m, :J))))
#update!(m.settings[:zz], Setting(:zz, construct_labor_income_grid(get_setting(m, :ygrid_combined), get_setting(m, :g_z), m[:meanlabeff].value, get_setting(m, :I))))
update!(m.settings[:z], Setting(:z, get_setting(m, :zz)[1, :]))
update!(m.settings[:n_jump_vars], Setting(:n_jump_vars, 201))
update!(m.settings[:n_state_vars], Setting(:n_state_vars, 200))
update!(m.settings[:n_state_vars_unreduce], Setting(:n_state_vars_unreduce, 0))
update!(m.settings[:maxit_HJB], Setting(:maxit_HJB, 500))
update!(m.settings[:tol_HJB], Setting(:tol_HJB, 1e-8))
update!(m.settings[:Δ_HJB], Setting(:Δ_HJB, 1e6))
update!(m.settings[:Ir], Setting(:Ir, 100))
update!(m.settings[:crit_S], Setting(:crit_S, 1e-5))
update!(m.settings[:maxit_kfe], Setting(:maxit_kfe, 1000))
update!(m.settings[:tol_kfe], Setting(:tol_kfe, 1e-12))
update!(m.settings[:Δ_kfe], Setting(:Δ_kfe, 1e6))
update!(m.settings[:niter_hours], Setting(:niter_hours, 10))
update!(m.settings[:reduce_state_vars], Setting(:reduce_state_vars, true))
update!(m.settings[:reduce_v], Setting(:reduce_v, true))
update!(m.settings[:krylov_dim], Setting(:krylov_dim, 20))
update!(m.settings[:n_knots], Setting(:n_knots, 12))
update!(m.settings[:c_power], Setting(:c_power, 1))
knots = collect(range(get_setting(m, :amin), stop=get_setting(m, :amax), length=(get_setting(m, :n_knots) - 1)))
knots = (get_setting(m, :amax) .- get_setting(m, :amin))/(2^get_setting(m, :c_power) - 1) * ((knots .- get_setting(m, :amin)) / (get_setting(m, :amax) .- get_setting(m, :amin)) .+ 1).^get_setting(m, :c_power) .+ get_setting(m, :amin) .- (get_setting(m, :amax) - get_setting(m, :amin))/(2^get_setting(m, :c_power) - 1)
update!(m.settings[:knots_dict], Setting(:knots_dict, Dict(1 => knots)))
update!(m.settings[:spline_grid], Setting(:spline_grid, get_setting(m, :a)))
update!(m.settings[:n_prior], Setting(:n_prior, 1))
update!(m.settings[:n_post], Setting(:n_post, 2))

# compare steady state values
steadystate!(m)
#=
@testset "Steady State" begin
    @test @test_matrix_approx_eq vec(m[:V_ss].value) vec(vars_ss_mat["V_SS"])
    @test @test_matrix_approx_eq vec(m[:g_ss].value) vec(vars_ss_mat["g_SS"])
    @test @test_matrix_approx_eq vec(m[:u_ss].value) vec(vars_ss_mat["u_SS"])
    @test @test_matrix_approx_eq vec(m[:c_ss].value) vec(vars_ss_mat["c_SS"])
    @test @test_matrix_approx_eq vec(m[:h_ss].value) vec(vars_ss_mat["h_SS"])
    @test @test_matrix_approx_eq vec(m[:s_ss].value) vec(vars_ss_mat["s_SS"])
    @test m[:B_ss].value ≈ vars_ss_mat["B_SS"] # approx to allow machine error
    @test m[:N_ss].value ≈ vars_ss_mat["N_SS"]
    @test m[:Y_ss].value ≈ vars_ss_mat["Y_SS"]
    @test m[:G_ss].value ≈ vars_ss_mat["G_SS"]
    @test m[:T_ss].value ≈ vars_ss_mat["T_SS"]
end
=#
### Equilibrium conditions
Γ0, Γ1, Ψ, Π, C = eqcond(m)

# Transition and measurement equations
TTT, RRR, CCC = solve(m)
# measurement here

# Matrices have expected dimensions
@testset "Canonical and State Transition Matrices" begin
    n_vars = n_states(m)
    @test size(Γ0) == (n_vars, n_vars)
    @test size(Γ1) == (n_vars, n_vars)
    @test size(C) == (n_vars,)
    @test size(Ψ) == (n_vars, 1)
    @test size(Π) == (n_vars, n_shocks_expectational(m))
    if get_setting(m, :reduce_v)
        @test size(TTT) == (get_setting(m, :n_state_vars_red) + get_setting(m, :n_splined), get_setting(m, :n_state_vars_red) + get_setting(m, :n_splined))
        @test size(RRR) == (get_setting(m, :n_state_vars_red) + get_setting(m, :n_splined), 1)
        @test size(vec(CCC)) == (get_setting(m, :n_state_vars_red) + get_setting(m, :n_splined),)
    elseif get_setting(m, :reduce_state_vars)
        @test size(TTT) == (get_setting(m, :n_jump_vars) + get_setting(m, :n_state_vars_red), get_setting(m, :n_state_vars_red))
        @test size(RRR) == (get_setting(m, :n_jump_vars) + get_setting(m, :n_state_vars_red), 1)
        @test size(vec(CCC)) == (get_setting(m, :n_jump_vars) + get_setting(m, :n_state_vars_red),)
    else
        @test size(TTT) == (n_vars, n_vars)
        @test size(RRR) == (n_vars, 1)
        @test size(vec(CCC)) == (n_vars,)
    end
end


text:
using LinearAlgebra

# order 2
Qp = [
  -1//4 5//4 -1//2 0 0 0
  -1//4 -5//4 2 -1//2 0 0
  0 0 -3//2 2 -1//2 0
  0 0 0 -3//2 2 -1//2
  0 0 0 0 -5//4 5//4
  0 0 0 0 -1//4 -1//4
]
Qm = Matrix(-Qp')
B = zero(Qp); B[1,1] = -1; B[end,end] = 1
M = Diagonal([1//4, 5//4, 1, 1, 5//4, 1//4])
Dp = M \ (Qp + B / 2)
Dm = M \ (Qm + B / 2)

# order 3
Qp = [
  -1//12 3//4 -1//6 0 0 0
  -5//12 -5//12 1 -1//6 0 0
  0 -1//3 -1//2 1 -1//6 0
  0 0 -1//3 -1//2 1 -1//6
  0 0 0 -1//3 -5//12 3//4
  0 0 0 0 -5//12 -1//12
]
Qm = Matrix(-Qp')
B = zero(Qp); B[1,1] = -1; B[end,end] = 1
M = Diagonal([5//12, 13//12, 1, 1, 13//12, 5//12])
Dp = M \ (Qp + B / 2)
Dm = M \ (Qm + B / 2)

# order 4
Qp = [
    -1//48  205//288 -29//144   1//96   0      0      0     0     0 0 0
  -169//288 -11//48   33//32  -43//144  1//12  0      0     0     0 0 0
    11//144 -13//32  -29//48  389//288 -1//2   1//12  0     0     0 0 0
     1//32  -11//144 -65//288 -13//16  3//2   -1//2   1//12 0     0 0 0
     0        0        0       -1//4   -5//6   3//2  -1//2  1//12 0 0 0
     0        0        0        0      -1//4  -5//6   3//2 -1//2  1//12 0 0
  0 0 0 0 0 -1//4 -5//6 3//2 -1//2 1//12 0
  0 0 0 0 0 0 -1//4 -13//16 389//288 -43//144 1//96
  0 0 0 0 0 0 0 -65//288 -29//48 33//32 -29//144
  0 0 0 0 0 0 0 -11//144 -13//32 -11//48 205//288
  0 0 0 0 0 0 0 1//32 11//144 -169//288 -1//48
]
Qm = Matrix(-Qp')
B = zero(Qp); B[1,1] = -1; B[end,end] = 1
m = [49//144, 61//48, 41//48, 149//144]
M = Diagonal(vcat(m, ones(Int, size(Qp,1)-2*length(m)), m[end:-1:1]))
Dp = M \ (Qp + B / 2)
Dm = M \ (Qm + B / 2)

# order 5
Qp = [
    -1//120 941//1440 -47//360 -7//480    0      0      0     0     0 0 0
  -869//1440 -11//120 25//32  -43//360    1//30  0      0     0     0 0 0
    29//360 -17//32  -29//120 1309//1440 -1//4 1//30  0     0     0 0 0
     1//32  -11//360 -661//1440 -13//40   1     -1//4  1//30 0     0 0 0
     0        0        1//20   -1//2     -1//3   1     -1//4  1//30 0 0 0
     0        0        0        1//20    -1//2  -1//3   1     -1//4  1//30 0 0
  0 0 0 0 1//20   -1//2     -1//3   1     -1//4  1//30 0
  0 0 0 0 0 1//20 -1//2 -13//40 1309//1440 -43//360 -7//480
  0 0 0 0 0 0 1//20 -661//1440 -29//120 25//32 -47//360
  0 0 0 0 0 0 0 -11//360 -17//32 -11//120 941//1440
  0 0 0 0 0 0 0 1//32 29//360 -869//1440 -1//120
]
Qm = Matrix(-Qp')
B = zero(Qp); B[1,1] = -1; B[end,end] = 1
# m = [251//720, 299//240, 41//48, 149//144]
m = [251//720, 299//240, 211//240, 739//720]
M = Diagonal(vcat(m, ones(Int, size(Qp,1)-2*length(m)), m[end:-1:1]))
Dp = M \ (Qp + B / 2)
Dm = M \ (Qm + B / 2)

# order 6
Qptmp = [
  -265//128688 1146190567//1737288000 -1596619//18384000 -55265831//579096000 26269819//3474576000 2464501//144774000 0 0 0 0 0 0
  -1116490567//1737288000 -8839//214480 190538869//347457600 102705469//694915200 413741//9651600 -191689861//3474576000 0 0 0 0 0 0
  1096619//18384000 -135385429//347457600 -61067//321720 45137333//57909600 -253641811//694915200 70665929//579096000 -1//60 0 0 0 0 0
  66965831//579096000 -208765789//694915200 -17623253//57909600 -18269//45960 410905829//347457600 -477953317//1158192000 2//15 -1//60 0 0 0 0
  -49219819//3474576000 293299//9651600 26422771//694915200 -141938309//347457600 -346583//643440 2217185207//1737288000 -1//2 2//15 -1//60 0 0 0
  -2374501//144774000 142906261//3474576000 -3137129//579096000 -29884283//1158192000 -630168407//1737288000 -3559//6128 4//3 -1//2 2//15 -1//60 0 0
  0 0 0 0 1//30 -2//5 -7//12 4//3 -1//2 2//15 -1//60 0
  0 0 0 0 0 1//30 -2//5 -7//12 4//3 -1//2 2//15 -1//60
]
bnd = size(Qptmp, 1)-2
Qp = zeros(eltype(Qptmp), 2*bnd+2, 2*bnd+2)
Qp[1:size(Qptmp,1), 1:size(Qptmp,2)] = Qptmp
for i in 1:bnd
  Qp[end+1-i, end+1-size(Qptmp,1):end] = Qptmp[end:-1:1, i]
end
Qm = Matrix(-Qp')
B = zero(Qp); B[1,1] = -1; B[end,end] = 1
m = [13613//43200, 12049//8640, 535//864, 1079//864, 7841//8640, 43837//43200]
M = Diagonal(vcat(m, ones(Int, size(Qp,1)-2*length(m)), m[end:-1:1]))
Dp = M \ (Qp + B / 2)
Dm = M \ (Qm + B / 2)

# order 7
Qptmp = Rational{Int128}[
  -265//300272 1587945773//2432203200 -1926361//25737600 -84398989//810734400 48781961//4864406400 3429119//202683600 0 0 0 0 0 0 0
  -1570125773//2432203200 -26517//1501360 240029831//486440640 202934303//972881280 118207//13512240 -231357719//4864406400 0 0 0 0 0 0 0
  1626361//25737600 -206937767//486440640 -61067//750680 49602727//81073440 -43783933//194576256 51815011//810734400 -1//140 0 0 0 0 0 0
  91418989//810734400 -53314099//194576256 -33094279//81073440 -18269//107240 440626231//486440640 -365711063//1621468800 1//15 -1//140 0 0 0 0 0
  -62551961//4864406400 799//35280 82588241//972881280 -279245719//486440640 -346583//1501360 2312302333//2432203200 -3//10 1//15 -1//140 0 0 0 0
  -3375119//202683600 202087559//4864406400 -11297731//810734400 61008503//1621468800 -1360092253//2432203200 -10677//42896 1 -3//10 1//15 -1//140 0 0 0
  0 0 0 -1//105 1//10 -3//5 -1//4 1 -3//10 1//15 -1//140 0 0
  0 0 0 0 -1//105 1//10 -3//5 -1//4 1 -3//10 1//15 -1//140 0
  0 0 0 0 0 -1//105 1//10 -3//5 -1//4 1 -3//10 1//15 -1//140
]
bnd = size(Qptmp, 1)-3
Qp = zeros(eltype(Qptmp), 2*bnd+3, 2*bnd+3)
Qp[1:size(Qptmp,1), 1:size(Qptmp,2)] = Qptmp
for i in 1:bnd
  Qp[end+1-i, end+1-size(Qptmp,1):end] = Qptmp[end:-1:1, i]
end
Qm = Matrix(-Qp')
B = zero(Qp); B[1,1] = -1; B[end,end] = 1
m = [19087//60480, 84199//60480, 18869//30240, 37621//30240, 55031//60480, 61343//60480]
M = Diagonal(vcat(m, ones(Int, size(Qp,1)-2*length(m)), m[end:-1:1]))
Dp = M \ (Qp + B / 2)
Dm = M \ (Qm + B / 2)


text:
include("Graph.jl")
include("Centrality.jl")

function doExp(ags1)
    G0 = readDiGraph("data/"*ags1*".txt")
    G = LSCC(G0)
    lg = open("log.txt", "a")
    println(lg, ags1, " ", G0.n, " ", G0.m, " ", G.n, " ", G.m)
    println(lg, "PageRank : ", getPc(PageRank(G)))
    println(lg, "Eigenvector : ", getPc(EigenvectorCentrality(G)))
    println(lg)
end

doExp(ARGS[1])


text:
function get_key(obj::SatelliteSystem)
    return jcall(obj, "getKey", jchar, ())
end

function parse_satellite_system(::Type{SatelliteSystem}, arg0::JString)
    return jcall(SatelliteSystem, "parseSatelliteSystem", SatelliteSystem, (JString,), arg0)
end

function value_of(::Type{SatelliteSystem}, arg0::JString)
    return jcall(SatelliteSystem, "valueOf", SatelliteSystem, (JString,), arg0)
end

function values(::Type{SatelliteSystem})
    return jcall(SatelliteSystem, "values", Vector{SatelliteSystem}, ())
end


text:
# idea: use a type parameterized on functions as the stack element
using DataStructures
import Base: convert, promote_rule, +, -, *, /, >, <, log, exp, <=, >=

type RDiff <: Real
    val::Float64
    adj::Float64
    vsym::Symbol
    asym::Symbol
    RDiff(val, adj, vsym, asym) = new(val, adj, vsym, asym)
end
typealias Converts Union{Float64, Int, Irrational, Rational}
RDiff(x::Converts) = RDiff(x, 0., gensym(), gensym())

convert(::Type{RDiff}, x::Converts) = RDiff(x)
promote_rule(::Type{Bool}, ::Type{RDiff}) = RDiff
promote_rule(::Type{Float64}, ::Type{RDiff}) = RDiff
promote_rule{T<:Integer}(::Type{T}, ::Type{RDiff}) = RDiff

value(x::RDiff) = x.val
adjoint(x::RDiff) = x.adj
value(x::Array{RDiff}) = map(value, x)
adjoint(x::Array{RDiff}) = map(adjoint, x)
function incadjoint!{N}(x::Array{RDiff, N}, y::Array{Float64, N})
    for i in eachindex(x)
        x[i].adj += y[i]
    end
end
>(x::RDiff, y::RDiff) = x.val > y.val
<(x::RDiff, y::RDiff) = x.val < y.val
>=(x::RDiff, y::RDiff) = x.val >= y.val
<=(x::RDiff, y::RDiff) = x.val <= y.val

type RNode{F<:Function, T1<:Tuple, T2}
    op::F
    inputs::T1
    outputs::T2
end

const global varstack = Stack(RNode)
const global exqueue = Queue(Expr)
const global exstack = Stack(Expr)

function grad(f, N)
    y = Array{RDiff}(N)

    function ∇f(x, storage)
        copy!(y, x)
        f(y)
        for node in varstack
            backprop(node)
        end
        clear!(varstack)
        for i in eachindex(storage)
            storage[i] = y[i].val
        end
    end

    return ∇f
end

function clear!(s::Stack)
    while !isempty(s)
        pop!(s)
    end
    nothing
end

backprop(n::RNode) = backprop(n.op, n.inputs..., n.outputs)

####################################
# HERE BE MONSTERS...
####################################

macro diff_rule(sig, code)
    fn = _fn_of_vals(sig)
    op = sig.args[1]
    args_typed = sig.args[2:end]
    argsyms = Symbol[s.args[1] for s in sig.args[2:end]]
    code_ex = Expr(:quote, code)
    ex = quote
        function $op($(args_typed...))
            res = map(RDiff, $fn)
            # n = RNode($op, tuple($(argsyms...)), res)
            # push!(varstack, n)
            ex_f = forward_codegen($op, tuple($(argsyms...)), res)
            vdict = Dict(zip($argsyms, [$(argsyms...)]))
            push!(vdict, :res => res)
            ex_r = reverse_codegen(vdict, $code_ex)
            enqueue!(exqueue, ex_f)
            push!(exstack, ex_r)
            res
        end

        function backprop(::typeof($op), $(args_typed...), res)
            $code
            nothing
        end
    end
    esc(ex)
end

_fn_of_vals(x) = x
function _fn_of_vals(ex::Expr)
    if ex.head == :(::)
        out = :(value($(ex.args[1])))
    else
        out = copy(ex)
        for i in eachindex(out.args)
            out.args[i] = _fn_of_vals(out.args[i])
        end
    end
    out
end

function forward_codegen(f::Function, inputs, output)
    ex = :($(output.vsym) = $f($([a.vsym for a in inputs]...)))
end
function reverse_codegen(args, ex)
    rdict = Dict{Expr, Symbol}()
    for (k, v) in args
        push!(rdict, Expr(:., k, QuoteNode(:adj)) => v.asym)
        push!(rdict, Expr(:., k, QuoteNode(:val)) => v.vsym)
    end
    _replacer(rdict, ex)
end

_replacer(d::Dict, x) = x
_replacer(d::Dict, ex::Expr) = haskey(d, ex) ? d[ex] : Expr(ex.head, map(x -> _replacer(d, x), ex.args)...)

####################################
# </MONSTERS>
####################################

@diff_rule log(x::RDiff) begin
    x.adj += res.adj / x.val
end

@diff_rule exp(x::RDiff) begin
    x.adj += res.adj * x.val
end

@diff_rule +(x::RDiff, y::RDiff) begin
    x.adj += res.adj
    y.adj += res.adj
end

@diff_rule -(x::RDiff, y::RDiff) begin
    x.adj += res.adj
    y.adj += -res.adj
end

@diff_rule *(x::RDiff, y::RDiff) begin
    x.adj += res.adj * y.val
    y.adj += res.adj * x.val
end

@diff_rule /(x::RDiff, y::RDiff) begin
    x.adj += res.adj / y.val
    y.adj += -res.adj * x.val / y.val^2
end

@diff_rule -(x::RDiff) begin
    x.adj += -res.adj
end

@diff_rule *(x::Array{RDiff}, y::Array{RDiff}) begin
    incadjoint!(x, adjoint(res) * value(y)')
    incadjoint!(y, value(x)' * adjoint(res))
end

@diff_rule +(x::Array{RDiff}, y::Array{RDiff}) begin
    incadjoint!(x, adjoint(res))
    incadjoint!(y, adjoint(res))
end

#= TODO
@diff_rule map(f::Function, x::Array{RDiff}) begin
    incadjoint!(x, adjoint(res) .* grad_expr(f)(x))
end
=#


text:
using Core.Compiler: PhiNode, GotoNode, GotoIfNot, ReturnNode, SSAValue, BasicBlock, Argument, PiNode
using IRTools

struct FutureLoopMask end

is_header(block_cfg, block_id) = any(map(pred -> pred > block_id, block_cfg.preds))

function get_block_for_statement(ir::IR)
  block_for_statement = Dict{SSAValue, Int64}()
  for b in blocks(ir)
    for (ssavalue, stmt) in b
      if ssavalue != nothing
        block_for_statement[ssavalue] = b.id
      end
    end
  end
  block_for_statement
end

function add_cond_for_block(block_to_conds::Dict{Int64, Vector{Tuple{Union{SSAValue, Argument}, Bool}}}, block_id, cond)
  if haskey(block_to_conds, block_id)
    if !(cond in block_to_conds[block_id])
      push!(block_to_conds[block_id], cond)
    end
  else
    block_to_conds[block_id] = [cond]
  end
end

_select(conds::NTuple{0}, first_vals::NTuple{0}, second_vals::NTuple{0}) = ()
function _select(conds::NTuple{N}, first_vals, second_vals) where {N}
  cond = first(conds).value
  v1 = first(first_vals)
  v2 = first(second_vals)
  other_conds = Base.tail(conds)
  other_v1s = Base.tail(first_vals)
  other_v2s = Base.tail(second_vals)
  (cond ? v1 : v2, _select(other_conds, other_v1s, other_v2s)...)
end

function select(conds::Mask{N}, first_vals::AbstractVec{T,N}, second_vals::AbstractVec{S,N}) where {T,S,N}
  res = _select(conds.data, first_vals.data, second_vals.data)
  vect(res...)
end

function select(conds::Mask{N}, first_val::T, second_val::S) where {T <: ScalarTypes, S <: ScalarTypes, N}
  if any(conds)
    first_val
  else
    second_val
  end
end

function fix_phinode(expr, old_to_new_ssavalue, old_to_new_block=Dict())
  if expr isa PhiNode
    new_values = map(x->if x isa SSAValue old_to_new_ssavalue[x] else x end,expr.values)
    new_edges = map(x->if haskey(old_to_new_block, x) old_to_new_block[x] else x end, expr.edges)
    new_expr = PhiNode(new_edges, new_values)
    return new_expr
  end
  return expr
end

function get_phinodes_values(ir, block_to_cond)
  conds = block_to_cond.vals
  phinodes_values = Dict{SSAValue, SSAValue}()
  for b in blocks(ir)
    for (ssavalue, stmt) in b
      if stmt.expr isa PhiNode && !(ssavalue in conds)
        in_loop_values = map(x -> x[1] >= b.id && x[2] isa SSAValue && x[2], zip(stmt.expr.edges, stmt.expr.values))
        in_loop_values = filter(x -> x != false, in_loop_values)
        for value in in_loop_values
          phinodes_values[value] = ssavalue
        end
      end
    end
  end
  phinodes_values
end

function pass_loops_result_value(ir, block_to_cond)
  block_for_statement = get_block_for_statement(ir)
  phinodes_values = get_phinodes_values(ir, block_to_cond)
  new_ir = IR(ir.lines, ir.args)
  old_to_new_ssavalue = Dict{SSAValue, SSAValue}()
  for b in blocks(ir)
    for (ssavalue, stmt) in b
      if stmt.expr isa Expr
        new_args = map(x->if x isa SSAValue old_to_new_ssavalue[x] else x end,stmt.expr.args)
        new_stmt = Expr(:call, new_args[1], new_args[2:length(new_args)]...)
        push!(new_ir, new_stmt)
        new_ssavalue = SSAValue(length(new_ir.defs))
        if haskey(phinodes_values, ssavalue)
          cond = old_to_new_ssavalue[block_to_cond[b.id]]
          header_value = old_to_new_ssavalue[phinodes_values[ssavalue]]
          push!(new_ir, xcall(Hydra, :select, cond, new_ssavalue, header_value))
          new_ssavalue = SSAValue(length(new_ir.defs))
        end
        old_to_new_ssavalue[ssavalue] = new_ssavalue
      else
        fix_stmt(ssavalue, stmt, new_ir, old_to_new_ssavalue)
      end
    end
    if (b.id != length(blocks(ir)))
      block!(new_ir)
    end
  end
  new_ir = map(x->fix_phinode(x, old_to_new_ssavalue), new_ir)
  new_ir, old_to_new_ssavalue
end

and_mask(masks...) = (&)(masks...)
or_mask(masks...) = (|)(masks...)
not_mask(mask) = ~mask

function fix_stmt(ssavalue, stmt, ir, old_to_new_ssavalue)
  replacement(x::SSAValue) = old_to_new_ssavalue[x]
  replacement(x) = x
  if stmt.expr isa GotoIfNot
    new_stmt = GotoIfNot(old_to_new_ssavalue[stmt.expr.cond], stmt.expr.dest)
    push!(ir, new_stmt)
  elseif stmt.expr isa GotoNode
    push!(ir, stmt)
  elseif stmt.expr isa ReturnNode
    push!(ir, ReturnNode(replacement(stmt.expr.val)))
  elseif stmt.expr isa PhiNode
    push!(ir, stmt)
    new_ssavalue = SSAValue(length(ir.defs))
    old_to_new_ssavalue[ssavalue] = new_ssavalue
  elseif stmt.expr isa PiNode
    push!(ir, stmt)
    new_ssavalue = SSAValue(length(ir.defs))
    old_to_new_ssavalue[ssavalue] = new_ssavalue
  elseif stmt.expr isa Expr
    new_args = map(x->if x isa SSAValue old_to_new_ssavalue[x] else x end,stmt.expr.args)
    new_stmt = Expr(:call, new_args[1], new_args[2:length(new_args)]...)
    push!(ir, new_stmt)
    new_ssavalue = SSAValue(length(ir.defs))
    old_to_new_ssavalue[ssavalue] = new_ssavalue
  elseif stmt.expr isa GlobalRef
    push!(ir, stmt.expr)
    new_ssavalue = SSAValue(length(ir.defs))
    old_to_new_ssavalue[ssavalue] = new_ssavalue
  end
end

function pass_create_masks(ir)
  cfg = CFG(ir)
  forward_ir = IR(ir.lines, ir.args)
  old_to_new_ssavalue = Dict{SSAValue, SSAValue}()
  block_to_block_mask = Dict{Tuple{Int64, Int64}, Union{SSAValue, Argument}}()
  block_to_block_mask[(0,1)] = Argument(0)
  cfg.blocks[1] = Hydra.BasicBlock(cfg.blocks[1].stmts, [0], cfg.blocks[1].succs)
  block_to_mask = Dict{Int64, Union{SSAValue, Argument}}()
  for (b, b_cfg) in zip(blocks(ir), cfg.blocks)
    forward_preds = filter(k-> k == 0 || k < b.id, b_cfg.preds)
    backward_preds = filter(k-> k != 0 && k > b.id, b_cfg.preds)
    succs = b_cfg.succs
    if length(backward_preds) == 0
      incoming_masks = map(x->block_to_block_mask[(x,b.id)], forward_preds)
      if length(incoming_masks) == 0
        continue
      end
      cond = incoming_masks[1]
      for other_cond in incoming_masks[2:end]
        push!(forward_ir, xcall(Hydra, :or_mask, cond, other_cond))
        cond = SSAValue(length(forward_ir.defs))
      end
      block_to_mask[b.id] = cond
    else
      push!(forward_ir, FutureLoopMask())
      block_to_mask[b.id] = SSAValue(length(forward_ir.defs))
    end
    for (ssavalue,stmt) in b
      if !iscontrol(stmt.expr)
        fix_stmt(ssavalue, stmt, forward_ir, old_to_new_ssavalue)
      end
    end
    if length(succs) == 1
      block_to_block_mask[(b.id, succs[1])] = block_to_mask[b.id]
    elseif length(succs) > 1
      cond = old_to_new_ssavalue[b.bb.gotos[1].expr.cond]
      else_dest = b.bb.gotos[1].expr.dest
      if_dest = b.id + 1
      push!(forward_ir, xcall(Hydra, :and_mask, block_to_mask[b.id], cond))
      if_cond = SSAValue(length(forward_ir.defs))
      block_to_block_mask[(b.id, if_dest)] = if_cond
      push!(forward_ir, xcall(Hydra, :not_mask, cond))
      not_cond = SSAValue(length(forward_ir.defs))
      push!(forward_ir, xcall(Hydra, :and_mask, block_to_mask[b.id], not_cond))
      else_cond = SSAValue(length(forward_ir.defs))
      block_to_block_mask[(b.id, else_dest)] = else_cond
    end
    for goto in b.bb.gotos
      fix_stmt(nothing, goto, forward_ir, old_to_new_ssavalue)
    end
    if (b.id != length(blocks(ir)))
      block!(forward_ir)
    end
  end
  forward_ir = map(x->fix_phinode(x, old_to_new_ssavalue), forward_ir)
  backward_ir = IR(ir.lines, ir.args)
  for (b, b_cfg) in zip(blocks(forward_ir), cfg.blocks)
    for (ssavalue, stmt) in b
      if stmt.expr isa FutureLoopMask
        preds = b_cfg.preds
        conds = map(k->block_to_block_mask[(k,b.id)], preds)
        loop_mask = PhiNode(preds, conds)
        push!(backward_ir, loop_mask)
      else
        push!(backward_ir, stmt.expr)
      end
    end
    if (b.id != length(blocks(ir)))
      block!(backward_ir)
    end
  end
  backward_ir, block_to_mask, block_to_block_mask
end

function pass_from_phi_to_select(ir, block_to_block_mask)
  new_ir = IR(ir.lines, ir.args)
  old_to_new_ssavalue = Dict{SSAValue, SSAValue}()
  for (b, b_cfg) in zip(blocks(ir), CFG(ir).blocks)
    for (ssavalue, stmt) in b
      if stmt.expr isa PhiNode
        if is_header(b_cfg, b.id)
          push!(new_ir, stmt)
          new_ssavalue = SSAValue(length(new_ir.defs))
          old_to_new_ssavalue[ssavalue] = new_ssavalue
          continue
        end
        conds = map(x -> old_to_new_ssavalue[block_to_block_mask[(x,b.id)]], stmt.expr.edges)
        values = map(x -> old_to_new_ssavalue[x], stmt.expr.values)
        second_cond = conds[2]
        first_val, second_val = values[1], values[2]
        push!(new_ir, xcall(Hydra, :select, second_cond, second_val, first_val))
        new_ssavalue = SSAValue(length(new_ir.defs))
        for (condition, value) in zip(conds[3:end], values[3:end])
          push!(new_ir, xcall(Hydra, :select, condition, value, new_ssavalue))
          new_ssavalue = SSAValue(length(new_ir.defs))
        end
        old_to_new_ssavalue[ssavalue] = new_ssavalue
      else
        fix_stmt(ssavalue, stmt, new_ir, old_to_new_ssavalue)
      end
    end
    if (b.id != length(blocks(ir)))
      block!(new_ir)
    end
  end
  map(x->fix_phinode(x, old_to_new_ssavalue), new_ir)
end

function pass_call(ir::IR, block_to_cond)
  new_ir = IR(ir.lines, ir.args)
  for b in blocks(ir)
    for (ssavalue, stmt) in b
      if stmt.expr isa Expr
        is_and_mask(x) = x == GlobalRef(Hydra, :and_mask)
        is_not_mask(x) = x == GlobalRef(Hydra, :not_mask)
        is_or_mask(x) = x == GlobalRef(Hydra, :or_mask)
        new_stmt = nothing
        funct = stmt.expr.args[1]
        if is_and_mask(funct) || is_not_mask(funct) || is_or_mask(funct)
          new_stmt = stmt
        else
          new_stmt = xcall(Hydra, :spmd, block_to_cond[b.id], stmt.expr.args...)
        end
        # new_stmt = xcall(Hydra, :spmd, stmt.expr.args[1], block_to_cond[b.id], stmt.expr.args[2:length(stmt.expr.args)]...)

        push!(new_ir, new_stmt)
      else
        push!(new_ir, stmt)
      end
    end
    if (b.id != length(blocks(ir)))
      block!(new_ir)
    end
  end
  new_ir
end

function get_all_loop_headers(ir)
  cfg = CFG(ir).blocks
  loop_footers = filter(x->any(x[1] .> x[2].succs), collect(enumerate(cfg)))
  loop_headers = []
  for (footer_id, footer) in loop_footers
    possible_header_indices = footer.succs[1]:footer_id-1
    header_id = filter(b_id->any(cfg[b_id].succs .== footer_id+1),possible_header_indices)[1]
    push!(loop_headers, header_id)
  end
  loop_headers
end

function for_loop_aggregator(masks)
  return any(masks)
end

function pass_delete_gotos(ir, block_to_block_mask)
  loop_headers = get_all_loop_headers(ir)
  new_ir = IR(ir.lines, ir.args)
  old_to_new_ssavalue = Dict{SSAValue, SSAValue}()
  cfg = CFG(ir).blocks
  for b in blocks(ir)
    for (ssavalue, stmt) in b
      if stmt.expr isa GotoIfNot
        if b.id in loop_headers
          mask = block_to_block_mask[(b.id, b.id+1)]
          push!(new_ir, xcall(Hydra, :for_loop_aggregator, mask))
          new_ssavalue = SSAValue(length(new_ir.defs))
          new_stmt = GotoIfNot(new_ssavalue, stmt.expr.dest)
          push!(new_ir, new_stmt)
        else
          continue
        end
      elseif stmt.expr isa GotoNode
        if stmt.expr.label < b.id
          push!(new_ir, stmt)
        end
      else
        fix_stmt(ssavalue, stmt, new_ir, old_to_new_ssavalue)
      end
    end
    if (b.id != length(blocks(ir)))
      block!(new_ir)
    end
  end
  map(x->fix_phinode(x, old_to_new_ssavalue), new_ir)
end

unwraptype(::Type{SVec{T,N}}) where {T,N} = T
unwraptype(::Type{VecArray{C,A,N}}) where {C,A,N} = C
unwraptype(x) = x

function pass(ir)
  ir, block_to_cond, block_to_block_mask = pass_create_masks(ir)
  ir = pass_call(ir, block_to_cond)
  ir = pass_from_phi_to_select(ir, block_to_block_mask)
  ir, old_to_new_ssavalue = pass_loops_result_value(ir, block_to_cond)
  block_to_block_mask = Dict(k=>v isa SSAValue ? old_to_new_ssavalue[v] : v for (k,v) in block_to_block_mask)
  ir = pass_delete_gotos(ir, block_to_block_mask)
  ir
end

@generated function spmd(mask, f, args...)
  f <: Core.IntrinsicFunction && return :(error("Can't SPMD intrinsic $f"))
  T = Tuple{f,unwraptype.(args)...}
  m = meta(T)
  m == nothing && return :(error($"Can't SPMD $T"))
  ir = IR(m)
  ir = pass(ir)
  argnames!(m, :f, :args)
  ir = spliceargs!(m, ir, (Symbol("#self#"), typeof(spmd)), (:mask, mask))
  ir = varargs!(m, ir, 3)
  update!(m, ir)
  return m.code
end

using InteractiveUtils
using InteractiveUtils: typesof

macro code_spmd(ex)
  @capture(ex, f_(xs__)) || error("@code_spmd f(xs...)")
  :(IRTools.code_ir($(esc(f)), typesof($(xs...))) |> pass)
end


text:
module PinGo

using Luxor
using ExprTools
using Luxor: preview
export Bingo, save, preview, generate

include("bingo.jl")
include("spells.jl")

end


text:
function search(model, hint, fun = (x) -> false; maxprint = 50,
    fields = [:rxns, :mets, :rxnNames, :metNames, :genes])

hint = string(hint)
hint == "" && (println("0 found!!!"), return)

fun_ = [] # custom filter
eqs_ = [] # equals
stw_ = [] # starts with
ctn_ = [] # contains
edw_ = [] # ends with

up = uppercase
up_hint = up(hint)
for field in fields
    dat = getfield(model, field)
    push!(fun_, map(idx -> (dat[idx], idx), findall(fun, dat))...)
    push!(eqs_, map(idx -> (dat[idx], idx), findall(x-> up(x) == up_hint, dat))...)
    push!(stw_, map(idx -> (dat[idx], idx), findall(x-> startswith(up(x), up_hint), dat))...)
    push!(ctn_, map(idx -> (dat[idx], idx), findall(x-> occursin(up_hint, up(x)), dat))...)
    push!(edw_, map(idx -> (dat[idx], idx), findall(x-> endswith(up(x), up_hint), dat))...)
end

# print
all_res = [sort!(fun_); sort!(eqs_); sort!(stw_); sort!(ctn_); sort!(edw_)] |> unique
c = 0
println("$(length(all_res)) found!!!")
for (res, idx) in all_res

    print(idx, ": ")
    _print_bold(res, hint)

    println()
    c == maxprint && (println("..."), return)
    c += 1
end
end

function _print_bold(text, in_bold)
    up_in_bold = uppercase(in_bold)
    up_text = uppercase(text)

    # I know maybe this is not the best way :)
    # This is only printing in bold the matches
    ci = 1 # current index
    while true
        hr_ = findnext(up_in_bold, up_text, ci)

        # No more hint
        if isnothing(hr_)
            print(text[ci:end])
            break;
        end

        # no-hint before hint
        print(text[ci:(first(hr_) - 1)])

        # hint
        printstyled(text[hr_], bold = true)

        ci = first(hr_) + length(in_bold)
        ci > length(text) && break
    end
end


text:
@doc raw"""
    Symplectic{n, 𝔽} <: AbstractEmbeddedManifold{𝔽, DefaultIsometricEmbeddingType}

The symplectic manifold consists of all ``2n \times 2n`` matrices which preserve
the canonical symplectic form over ``𝔽^{2n × 2n} \times 𝔽^{2n × 2n}``,
````math
    \omega\colon 𝔽^{2n × 2n} \times 𝔽^{2n × 2n} \rightarrow 𝔽,
    \quad \omega(x, y) = p^T Q_{2n} q, \; x, y \in 𝔽^{2n × 2n},
````
where
````math
Q_{2n} =
\begin{bmatrix}
  0_n & I_n \\
 -I_n & 0_n
\end{bmatrix}.
````
That is, the symplectic manifold consists of
````math
\operatorname{Sp}(2n, ℝ) = \bigl\{ p ∈ ℝ^{2n × 2n} \, \big| \, p^TQ_{2n}p = Q_{2n} \bigr\},
````
with ``0_n`` and ``I_n`` denoting the ``n × n`` zero-matrix
and indentity matrix in ``ℝ^{n \times n}`` respectively.

The tangent space at a point ``p`` is given by [^Bendokat2021]
````math
\begin{align*}
    T_p\operatorname{Sp}(2n)
        &= \{X \in \mathbb{R}^{2n \times 2n} \;|\; p^{T}Q_{2n}X + X^{T}Q_{2n}p = 0 \}, \\
        &= \{X = pQS \;|\; S ∈ R^{2n × 2n}, S^T = S \}.
\end{align*}
````

# Constructor
    Symplectic(2n, field=ℝ) -> Symplectic{div(2n, 2), field}()

Generate the (real-valued) symplectic manifold of ``2n \times 2n`` symplectic matrices.
The constructor for the [`Symplectic`](@ref) manifold accepts the even column/row embedding
dimension ``2n`` for the real symplectic manifold, ``ℝ^{2n × 2n}``.
"""
struct Symplectic{n,𝔽} <: AbstractDecoratorManifold{𝔽} end

function active_traits(f, ::Symplectic, args...)
    return merge_traits(IsEmbeddedManifold(), IsDefaultMetric(RealSymplecticMetric()))
end

function Symplectic(n::Int, field::AbstractNumbers=ℝ)
    n % 2 == 0 || throw(ArgumentError("The dimension of the symplectic manifold
                        embedding space must be even. Was odd, n % 2 == $(n % 2)."))
    return Symplectic{div(n, 2),field}()
end

@doc raw"""
    RealSymplecticMetric <: RiemannianMetric

The canonical Riemannian metric on the symplectic manifold,
defined pointwise for ``p \in \operatorname{Sp}(2n)`` by [^FioriSimone2011]
````math
\begin{align*}
    & g_p \colon T_p\operatorname{Sp}(2n) \times T_p\operatorname{Sp}(2n) \rightarrow ℝ, \\
    & g_p(Z_1, Z_2) = \operatorname{tr}((p^{-1}Z_1)^T (p^{-1}Z_2)).
\end{align*}
````
This metric is also the default metric for the [`Symplectic`](@ref) manifold.
"""
struct RealSymplecticMetric <: RiemannianMetric end

@doc raw"""
    ExtendedSymplecticMetric <: AbstractMetric

The extension of the [`RealSymplecticMetric`](@ref) at a point `p \in \operatorname{Sp}(2n)`
as an inner product over the embedding space ``ℝ^{2n \times 2n}``, i.e.
````math
    \langle x, y \rangle_{p} = \langle p^{-1}x, p^{-1}\rangle_{\operatorname{Fr}}
    = \operatorname{tr}(x^T(pp^T)^{-1}y), \;\forall\; x, y \in ℝ^{2n \times 2n}.
````
"""
struct ExtendedSymplecticMetric <: AbstractMetric end

@doc raw"""
    SymplecticMatrix{T}

A lightweight structure to represent the action of the matrix
representation of the canonical symplectic form,
````math
Q_{2n}(λ) = λ
\begin{bmatrix}
0_n & I_n \\
 -I_n & 0_n
\end{bmatrix} \quad \in ℝ^{2n \times 2n},
````
such that the canonical symplectic form is represented by
````math
\omega_{2n}(x, y) = x^TQ_{2n}(1)y, \quad x, y \in ℝ^{2n}.
````

The entire matrix is however not instantiated in memory, instead a scalar
``λ`` of type `T` is stored, which is used to keep track of scaling and transpose operations
applied  to each `SymplecticMatrix`.
For example, given `Q = SymplecticMatrix(1.0)` represented as `1.0*[0 I; -I 0]`,
the adjoint `Q'` returns `SymplecticMatrix(-1.0) = (-1.0)*[0 I; -I 0]`.
"""
struct SymplecticMatrix{T}
    λ::T
end
SymplecticMatrix() = SymplecticMatrix(1)
SymplecticMatrix(λ::T) where {T<:Number} = SymplecticMatrix{T}(λ)

function SymplecticMatrix(arrays::Vararg{AbstractArray})
    TS = Base.promote_type(map(eltype, arrays)...)
    return SymplecticMatrix(one(TS))
end

@doc raw"""
    change_representer(::Symplectic, ::EuclideanMetric, p, X)
    change_representer!(::Symplectic, Y, ::EuclideanMetric, p, X)

Compute the representation of a tangent vector ``ξ ∈ T_p\operatorname{Sp}(2n, ℝ)`` s.t.
````math
    g_p(c_p(ξ), η) = ⟨ξ, η⟩^{\text{Euc}} \;∀\; η ∈ T_p\operatorname{Sp}(2n, ℝ).
````
with the conversion function
````math
    c_p : T_p\operatorname{Sp}(2n, ℝ) \rightarrow T_p\operatorname{Sp}(2n, ℝ), \quad
    c_p(ξ) = \frac{1}{2} pp^T ξ + \frac{1}{2} pQ ξ^T pQ.
````

Each of the terms ``c_p^1(ξ) = p p^T ξ`` and ``c_p^2(ξ) = pQ ξ^T pQ`` from the
above definition of ``c_p(η)`` are themselves metric compatible in the sense that
````math
    c_p^i : T_p\operatorname{Sp}(2n, ℝ) \rightarrow \mathbb{R}^{2n \times 2n}\quad
    g_p^i(c_p(ξ), η) = ⟨ξ, η⟩^{\text{Euc}} \;∀\; η ∈ T_p\operatorname{Sp}(2n, ℝ),
````
for ``i \in {1, 2}``. However the range of each function alone is not confined to
``T_p\operatorname{Sp}(2n, ℝ)``, but the convex combination
````math
    c_p(ξ) = \frac{1}{2}c_p^1(ξ) + \frac{1}{2}c_p^2(ξ)
````
does have the correct range ``T_p\operatorname{Sp}(2n, ℝ)``.
"""
change_representer(::Symplectic, ::EuclideanMetric, p, X)

function change_representer!(::Symplectic, Y, ::EuclideanMetric, p, X)
    Q = SymplecticMatrix(p, X)
    pT_X = p' * X
    Y .= (1 / 2) .* p * (pT_X .+ Q * pT_X' * Q)
    return Y
end

@doc raw"""
    change_representer(MetMan::MetricManifold{𝔽, Euclidean{Tuple{m, n}, 𝔽}, ExtendedSymplecticMetric},
                       EucMet::EuclideanMetric, p, X)
    change_representer!(MetMan::MetricManifold{𝔽, Euclidean{Tuple{m, n}, 𝔽}, ExtendedSymplecticMetric},
                        Y, EucMet::EuclideanMetric, p, X)

Change the representation of a matrix ``ξ ∈ \mathbb{R}^{2n \times 2n}``
into the inner product space ``(ℝ^{2n \times 2n}, g_p)`` where the inner product
is given by
``g_p(ξ, η) = \langle p^{-1}ξ, p^{-1}η \rangle = \operatorname{tr}(ξ^T(pp^T)^{-1}η)``,
as the extension of the [`RealSymplecticMetric`](@ref) onto the entire embedding space.

By changing the representation we mean to apply a mapping
````math
    c_p : \mathbb{R}^{2n \times 2n} \rightarrow \mathbb{R}^{2n \times 2n},
````
defined by requiring that it satisfy the metric compatibility condition
````math
    g_p(c_p(ξ), η) = ⟨p^{-1}c_p(ξ), p^{-1}η⟩ = ⟨ξ, η⟩^{\text{Euc}}
        \;∀\; η ∈ T_p\operatorname{Sp}(2n, ℝ).
````
In this case, we compute the mapping
````math
    c_p(ξ) = pp^T ξ.
````
"""
function change_representer(
    ::MetricManifold{𝔽,Euclidean{Tuple{m,n},𝔽},ExtendedSymplecticMetric},
    ::EuclideanMetric,
    p,
    X,
) where {𝔽,m,n}
    return p * p' * X
end

function change_representer!(
    ::MetricManifold{𝔽,Euclidean{Tuple{m,n},𝔽},ExtendedSymplecticMetric},
    Y,
    ::EuclideanMetric,
    p,
    X,
) where {𝔽,m,n}
    Y .= p * p' * X
    return Y
end

@doc raw"""
    check_point(M::Symplectic, p; kwargs...)

Check whether `p` is a valid point on the [`Symplectic`](@ref) `M`=$\operatorname{Sp}(2n)$,
i.e. that it has the right [`AbstractNumbers`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#number-system) type and $p^{+}p$ is (approximately)
the identity, where $A^{+} = Q_{2n}^TA^TQ_{2n}$ is the symplectic inverse, with
````math
Q_{2n} =
\begin{bmatrix}
0_n & I_n \\
 -I_n & 0_n
\end{bmatrix}.
````
The tolerance can be set with `kwargs...` (e.g. `atol = 1.0e-14`).
"""
function check_point(M::Symplectic{n,ℝ}, p; kwargs...) where {n,ℝ}
    # Perform check that the matrix lives on the real symplectic manifold:
    expected_zero = norm(inv(M, p) * p - LinearAlgebra.I)
    if !isapprox(expected_zero, zero(eltype(p)); kwargs...)
        return DomainError(
            expected_zero,
            (
                "The point p does not lie on $(M) because its symplectic" *
                " inverse composed with itself is not the identity."
            ),
        )
    end
    return nothing
end

@doc raw"""
    check_vector(M::Symplectic, p, X; kwargs...)

Checks whether `X` is a valid tangent vector at `p` on the [`Symplectic`](@ref)
`M`=``\operatorname{Sp}(2n)``, i.e. the [`AbstractNumbers`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#number-system) fits and
it (approximately) holds that ``p^{T}Q_{2n}X + X^{T}Q_{2n}p = 0``,
where
````math
Q_{2n} =
\begin{bmatrix}
0_n & I_n \\
 -I_n & 0_n
\end{bmatrix}.
````
The tolerance can be set with `kwargs...` (e.g. `atol = 1.0e-14`).
"""
check_vector(::Symplectic, ::Any...)

function check_vector(M::Symplectic{n}, p, X; kwargs...) where {n}
    Q = SymplecticMatrix(p, X)
    tangent_requirement_norm = norm(X' * Q * p + p' * Q * X, 2)
    if !isapprox(tangent_requirement_norm, 0.0; kwargs...)
        return DomainError(
            tangent_requirement_norm,
            (
                "The matrix X is not in the tangent space at point p of the" *
                " manifold $(M), as X'Qp + p'QX is not the zero matrix."
            ),
        )
    end
    return nothing
end

ManifoldsBase.default_inverse_retraction_method(::Symplectic) = CayleyInverseRetraction()

ManifoldsBase.default_retraction_method(::Symplectic) = CayleyRetraction()

@doc raw"""
    distance(M::Symplectic, p, q)

Compute an approximate geodesic distance between two Symplectic matrices
``p, q \in \operatorname{Sp}(2n)``, as done in [^WangRealSymplecticGroup].
````math
    \operatorname{dist}(p, q)
        ≈ ||\operatorname{Log}(p^+q)||_{\operatorname{Fr}},
````
where the ``\operatorname{Log}(\cdot)`` operator is the matrix logarithm.

This approximation is justified by first recalling the Baker-Campbell-Hausdorf formula,
````math
\operatorname{Log}(\operatorname{Exp}(A)\operatorname{Exp}(B))
 = A + B + \frac{1}{2}[A, B] + \frac{1}{12}[A, [A, B]] + \frac{1}{12}[B, [B, A]]
    + \ldots \;.
````
Then we write the expression for the exponential map from ``p`` to ``q`` as
````math
    q =
    \operatorname{exp}_p(X)
    =
    p \operatorname{Exp}((p^{+}X)^T)
    \operatorname{Exp}([p^{+}X - (p^{+}X)^T]),
    X \in T_p\operatorname{Sp},
````
and with the geodesic distance between ``p`` and ``q`` given by
``\operatorname{dist}(p, q) = ||X||_p = ||p^+X||_{\operatorname{Fr}}``
we see that
````math
    \begin{align*}
    ||\operatorname{Log}(p^+q)||_{\operatorname{Fr}}
    &= ||\operatorname{Log}\left(
        \operatorname{Exp}((p^{+}X)^T)
        \operatorname{Exp}(p^{+}X - (p^{+}X)^T)
    \right)||_{\operatorname{Fr}} \\
    &= ||p^{+}X + \frac{1}{2}[(p^{+}X)^T, p^{+}X - (p^{+}X)^T]
            + \ldots ||_{\operatorname{Fr}} \\
    &≈ ||p^{+}X||_{\operatorname{Fr}} = \operatorname{dist}(p, q).
    \end{align*}
````
"""
function distance(M::Symplectic{n}, p, q) where {n}
    return norm(log(symplectic_inverse_times(M, p, q)))
end

embed(::Symplectic, p) = p
embed(::Symplectic, p, X) = X

@doc raw"""
    exp(M::Symplectic, p, X)
    exp!(M::Symplectic, q, p, X)

The Exponential mapping on the Symplectic manifold with the
[`RealSymplecticMetric`](@ref) Riemannian metric.

For the point ``p \in \operatorname{Sp}(2n)`` the exponential mapping along the tangent
vector ``X \in T_p\operatorname{Sp}(2n)`` is computed as [^WangRealSymplecticGroup]
````math
    \operatorname{exp}_p(X) = p \operatorname{Exp}((p^{-1}X)^T)
                                \operatorname{Exp}(p^{-1}X - (p^{-1}X)^T),
````
where ``\operatorname{Exp}(\cdot)`` denotes the matrix exponential.

[^WangRealSymplecticGroup]:
    > Wang, Jing and Sun, Huafei and Fiori, Simone:
    > A Riemannian-steepest-descent approach for optimization on the real symplectic group,
    > Mathematical Methods in the Applied Sciences, 41(11) pp. 4273-4286, 2018
    > doi [10.1002/mma.4890](https://doi.org/10.1002/mma.4890)
"""
exp(::Symplectic, ::Any...)

function exp!(M::Symplectic, q, p, X)
    p_star_X = symplectic_inverse_times(M, p, X)
    q .= p * exp(Array(p_star_X')) * exp(p_star_X - p_star_X')
    return q
end

get_embedding(::Symplectic{n,ℝ}) where {n} = Euclidean(2n, 2n; field=ℝ)

@doc raw"""
    gradient(M::Symplectic, f, p, backend::RiemannianProjectionBackend;
             extended_metric=true)
    gradient!(M::Symplectic, f, p, backend::RiemannianProjectionBackend;
             extended_metric=true)

Compute the manifold gradient ``\text{grad}f(p)`` of a scalar function
``f \colon \operatorname{Sp}(2n) \rightarrow ℝ`` at
``p \in \operatorname{Sp}(2n)``.

The element ``\text{grad}f(p)`` is found as the Riesz representer of the differential
``\text{D}f(p) \colon T_p\operatorname{Sp}(2n) \rightarrow ℝ`` w.r.t.
the Riemannian metric inner product at ``p`` [^FioriSimone2011].
That is, ``\text{grad}f(p) \in T_p\operatorname{Sp}(2n)`` solves the relation
````math
    g_p(\text{grad}f(p), X) = \text{D}f(p) \quad\forall\; X \in T_p\operatorname{Sp}(2n).
````

The default behaviour is to first change the representation of the Euclidean gradient from
the Euclidean metric to the [`RealSymplecticMetric`](@ref) at ``p``, and then we projecting
the result onto the correct tangent tangent space ``T_p\operatorname{Sp}(2n, ℝ)``
w.r.t the Riemannian metric ``g_p`` extended to the entire embedding space.

# Arguments:
- `extended_metric = true`: If `true`, compute the gradient ``\text{grad}f(p)`` by
    first changing the representer of the Euclidean gradient of a smooth extension
    of ``f``, ``∇f(p)``, w.r.t. the [`RealSymplecticMetric`](@ref) at ``p``
    extended to the entire embedding space, before projecting onto the correct
    tangent vector space w.r.t. the same extended metric ``g_p``.
    If `false`, compute the gradient by first projecting ``∇f(p)`` onto the
    tangent vector space, before changing the representer in the tangent
    vector space to comply with the [`RealSymplecticMetric`](@ref).

[^FioriSimone2011]:
    > Simone Fiori:
    > Solving minimal-distance problems over the manifold of real-symplectic matrices,
    > SIAM Journal on Matrix Analysis and Applications 32(3), pp. 938-968, 2011.
    > doi [10.1137/100817115](https://doi.org/10.1137/100817115).
"""
function gradient(
    M::Symplectic,
    f,
    p,
    backend::RiemannianProjectionBackend;
    extended_metric=true,
)
    Y = allocate_result(M, gradient, p)
    return gradient!(M, f, Y, p, backend, extended_metric=extended_metric)
end

function gradient!(
    M::Symplectic,
    f,
    X,
    p,
    backend::RiemannianProjectionBackend;
    extended_metric=true,
)
    _gradient!(f, X, p, backend.diff_backend)
    if extended_metric
        MetricM = MetricManifold(get_embedding(M), ExtendedSymplecticMetric())
        change_representer!(MetricM, X, EuclideanMetric(), p, X)
        return project!(MetricM, X, p, X)
    else
        project!(M, X, p, X)
        return change_representer!(M, X, EuclideanMetric(), p, X)
    end
end

@doc raw"""
    inner(::Symplectic{n, ℝ}, p, X, Y)

Compute the canonical Riemannian inner product [`RealSymplecticMetric`](@ref)
````math
    g_p(X, Y) = \operatorname{tr}((p^{-1}X)^T (p^{-1}Y))
````
between the two tangent vectors ``X, Y \in T_p\operatorname{Sp}(2n)``.
"""
function inner(M::Symplectic{n,ℝ}, p, X, Y) where {n}
    p_star = inv(M, p)
    return dot((p_star * X), (p_star * Y))
end

@doc raw"""
    inv(::Symplectic, A)
    inv!(::Symplectic, A)

Compute the symplectic inverse ``A^+`` of matrix ``A ∈ ℝ^{2n × 2n}``. Given a matrix
````math
A ∈ ℝ^{2n × 2n},\quad
A =
\begin{bmatrix}
A_{1,1} & A_{1,2} \\
A_{2,1} & A_{2, 2}
\end{bmatrix}
````
the symplectic inverse is defined as:
````math
A^{+} := Q_{2n}^T A^T Q_{2n},
````
where
````math
Q_{2n} =
\begin{bmatrix}
0_n & I_n \\
 -I_n & 0_n
\end{bmatrix}.
````
The symplectic inverse of A can be expressed explicitly as:
````math
A^{+} =
\begin{bmatrix}
  A_{2, 2}^T & -A_{1, 2}^T \\[1.2mm]
 -A_{2, 1}^T &  A_{1, 1}^T
\end{bmatrix}.
````
"""
function Base.inv(::Symplectic{n,ℝ}, A) where {n}
    Ai = similar(A)
    checkbounds(A, 1:(2n), 1:(2n))
    @inbounds for i in 1:n, j in 1:n
        Ai[i, j] = A[j + n, i + n]
    end
    @inbounds for i in 1:n, j in 1:n
        Ai[i + n, j] = -A[j + n, i]
    end
    @inbounds for i in 1:n, j in 1:n
        Ai[i, j + n] = -A[j, i + n]
    end
    @inbounds for i in 1:n, j in 1:n
        Ai[i + n, j + n] = A[j, i]
    end
    return Ai
end

function inv!(::Symplectic{n,ℝ}, A) where {n}
    checkbounds(A, 1:(2n), 1:(2n))
    @inbounds for i in 1:n, j in 1:n
        A[i, j], A[j + n, i + n] = A[j + n, i + n], A[i, j]
    end
    @inbounds for i in 1:n, j in i:n
        if i == j
            A[i, j + n] = -A[i, j + n]
        else
            A[i, j + n], A[j, i + n] = -A[j, i + n], -A[i, j + n]
        end
    end
    @inbounds for i in 1:n, j in i:n
        if i == j
            A[i + n, j] = -A[i + n, j]
        else
            A[i + n, j], A[j + n, i] = -A[j + n, i], -A[i + n, j]
        end
    end
    return A
end

@doc raw"""
    inverse_retract(M::Symplectic, p, q, ::CayleyInverseRetraction)

Compute the Cayley Inverse Retraction ``X = \mathcal{L}_p^{\operatorname{Sp}}(q)``
such that the Cayley Retraction from ``p`` along ``X`` lands at ``q``, i.e.
``\mathcal{R}_p(X) = q`` [^Bendokat2021].

First, recall the definition the standard symplectic matrix
````math
Q =
\begin{bmatrix}
 0    & I \\
-I  & 0
\end{bmatrix}
````
as well as the symplectic inverse of a matrix ``A``, ``A^{+} = Q^T A^T Q``.

For ``p, q ∈ \operatorname{Sp}(2n, ℝ)`` then, we can then define the
inverse cayley retraction as long as the following matrices exist.
````math
    U = (I + p^+ q)^{-1}, \quad V = (I + q^+ p)^{-1}.
````

If that is the case, the inverse cayley retration at ``p`` applied to ``q`` is
````math
\mathcal{L}_p^{\operatorname{Sp}}(q) = 2p\bigl(V - U\bigr) + 2\bigl((p + q)U - p\bigr)
                                        ∈ T_p\operatorname{Sp}(2n).
````

[^Bendokat2021]:
    > Bendokat, Thomas and Zimmermann, Ralf:
	> The real symplectic Stiefel and Grassmann manifolds: metrics, geodesics and applications
	> arXiv preprint arXiv:2108.12447, 2021 (https://arxiv.org/abs/2108.12447)
"""
inverse_retract(::Symplectic, p, q, ::CayleyInverseRetraction)

function inverse_retract_caley!(M::Symplectic, X, p, q)
    U_inv = lu(add_scaled_I!(symplectic_inverse_times(M, p, q), 1))
    V_inv = lu(add_scaled_I!(symplectic_inverse_times(M, q, p), 1))

    X .= 2 .* ((p / V_inv .- p / U_inv) .+ ((p + q) / U_inv) .- p)
    return X
end

@doc raw"""
    manifold_dimension(::Symplectic{n})

Returns the dimension of the symplectic manifold
embedded in ``ℝ^{2n \times 2n}``, i.e.
````math
    \operatorname{dim}(\operatorname{Sp}(2n)) = (2n + 1)n.
````
"""
manifold_dimension(::Symplectic{n}) where {n} = (2n + 1) * n

@doc raw"""
    project(::Symplectic, p, A)
    project!(::Symplectic, Y, p, A)

Given a point ``p \in \operatorname{Sp}(2n)``,
project an element ``A \in \mathbb{R}^{2n \times 2n}`` onto
the tangent space ``T_p\operatorname{Sp}(2n)`` relative to
the euclidean metric of the embedding ``\mathbb{R}^{2n \times 2n}``.

That is, we find the element ``X \in T_p\operatorname{SpSt}(2n, 2k)``
which solves the constrained optimization problem
````math
    \operatorname{min}_{X \in \mathbb{R}^{2n \times 2n}} \frac{1}{2}||X - A||^2, \quad
    \text{s.t.}\;
    h(X) \colon= X^T Q p + p^T Q X = 0,
````
where ``h\colon\mathbb{R}^{2n \times 2n} \rightarrow \operatorname{skew}(2n)`` defines
the restriction of ``X`` onto the tangent space ``T_p\operatorname{SpSt}(2n, 2k)``.
"""
project(::Symplectic, p, A)

function project!(::Symplectic, Y, p, A)
    Q = SymplecticMatrix(Y, p, A)
    Q_p = Q * p

    function h(X)
        XT_Q_p = X' * Q_p
        return XT_Q_p .- XT_Q_p'
    end

    # Solve for Λ (Lagrange mutliplier):
    pT_p = p' * p  # (2k × 2k)
    Λ = sylvester(pT_p, pT_p, h(A) ./ 2)

    Y[:, :] = A .- Q_p * (Λ .- Λ')
    return Y
end

@doc raw"""
    project!(::MetricManifold{𝔽,Euclidean,ExtendedSymplecticMetric}, Y, p, X) where {𝔽}

Compute the projection of ``X ∈ R^{2n × 2n}`` onto ``T_p\operatorname{Sp}(2n, ℝ)`` w.r.t.
the Riemannian metric ``g`` [`RealSymplecticMetric`](@ref).
The closed form projection mapping is given by [^Gao2021Riemannian]
````math
    \operatorname{P}^{T_p\operatorname{Sp}(2n)}_{g_p}(X) = pQ\operatorname{sym}(p^TQ^TX),
````
where ``\operatorname{sym}(A) = \frac{1}{2}(A + A^T)``.
This function is not exported.
"""
function project!(
    ::MetricManifold{𝔽,Euclidean{Tuple{m,n},𝔽},ExtendedSymplecticMetric},
    Y,
    p,
    X,
) where {m,n,𝔽}
    Q = SymplecticMatrix(p, X)

    pT_QT_X = p' * Q' * X
    symmetrized_pT_QT_X = (1 / 2) .* (pT_QT_X + pT_QT_X')

    Y .= p * Q * (symmetrized_pT_QT_X)
    return Y
end

@doc raw"""
    project_normal!(::MetricManifold{𝔽,Euclidean,ExtendedSymplecticMetric}, Y, p, X)

Project onto the normal of the tangent space ``(T_p\operatorname{Sp}(2n))^{\perp_g}`` at
a point ``p ∈ \operatorname{Sp}(2n)``, relative to the riemannian metric
``g`` [`RealSymplecticMetric`](@ref).
That is,
````math
(T_p\operatorname{Sp}(2n))^{\perp_g} = \{Y \in \mathbb{R}^{2n \times 2n} :
                        g_p(Y, X) = 0 \;\forall\; X \in T_p\operatorname{Sp}(2n)\}.
````
The closed form projection operator onto the normal space is given by [^Gao2021Riemannian]
````math
\operatorname{P}^{(T_p\operatorname{Sp}(2n))\perp}_{g_p}(X) = pQ\operatorname{skew}(p^TQ^TX),
````
where ``\operatorname{skew}(A) = \frac{1}{2}(A - A^T)``.
This function is not exported.

[^Gao2021Riemannian]:
    > Gao, Bin and Son, Nguyen Thanh and Absil, P-A and Stykel, Tatjana:
    > Riemannian optimization on the symplectic Stiefel manifold,
    > SIAM Journal on Optimization 31(2), pp. 1546-1575, 2021.
    > doi [10.1137/20M1348522](https://doi.org/10.1137/20M1348522)
"""
function project_normal!(
    ::MetricManifold{𝔽,Euclidean{Tuple{m,n},𝔽},ExtendedSymplecticMetric},
    Y,
    p,
    X,
) where {m,n,𝔽}
    Q = SymplecticMatrix(p, X)

    pT_QT_X = p' * Q' * X
    skew_pT_QT_X = (1 / 2) .* (pT_QT_X .- pT_QT_X')

    Y .= p * Q * skew_pT_QT_X
    return Y
end

@doc raw"""
    rand(::SymplecticStiefel; vector_at=nothing,
        hamiltonian_norm = (vector_at === nothing ? 1/2 : 1.0))

Generate a random point on ``\operatorname{Sp}(2n)`` or a random
tangent vector ``X \in T_p\operatorname{Sp}(2n)`` if `vector_at` is set to
a point ``p \in \operatorname{Sp}(2n)``.

A random point on ``\operatorname{Sp}(2n)`` is constructed by generating a
random Hamiltonian matrix ``Ω \in \mathfrak{sp}(2n,F)`` with norm `hamiltonian_norm`,
and then transforming it to a symplectic matrix by applying the Cayley transform
````math
    \operatorname{cay}\colon \mathfrak{sp}(2n,F) \rightarrow \operatorname{Sp}(2n),
    \; \Omega \mapsto (I - \Omega)^{-1}(I + \Omega).
````
To generate a random tangent vector in ``T_p\operatorname{Sp}(2n)``, this code employs the
second tangent vector space parametrization of [Symplectic](@ref).
It first generates a random symmetric matrix ``S`` by `S = randn(2n, 2n)`
and then symmetrizes it as `S = S + S'`.
Then ``S`` is normalized to have Frobenius norm of `hamiltonian_norm`
and `X = pQS` is returned, where `Q` is the [`SymplecticMatrix`](@ref).
"""
function Base.rand(
    M::Symplectic;
    vector_at=nothing,
    hamiltonian_norm=(vector_at === nothing ? 1 / 2 : 1.0),
)
    if vector_at === nothing
        Ω = rand_hamiltonian(M; frobenius_norm=hamiltonian_norm)
        return (I - Ω) \ (I + Ω)
    else
        random_vector(M, vector_at; symmetric_norm=hamiltonian_norm)
    end
end

function random_vector(::Symplectic{n}, p::AbstractMatrix; symmetric_norm=1.0) where {n}
    # Generate random symmetric matrix:
    S = randn(2n, 2n)
    S .= (S + S')
    S *= symmetric_norm / norm(S)
    Q = SymplecticMatrix(p)
    lmul!(Q, S)
    return p * S
end

function rand_hamiltonian(::Symplectic{n}; frobenius_norm=1.0) where {n}
    A = randn(n, n)
    B = randn(n, n)
    C = randn(n, n)
    B = (1 / 2) .* (B .+ B')
    C = (1 / 2) .* (C .+ C')
    Ω = [A B; C -A']
    return frobenius_norm * Ω / norm(Ω, 2)
end

@doc raw"""
    retract(::Symplectic, p, X, ::CayleyRetraction)
    retract!(::Symplectic, q, p, X, ::CayleyRetraction)

Compute the Cayley retraction on ``p ∈ \operatorname{Sp}(2n, ℝ)`` in
the direction of tangent vector ``X ∈ T_p\operatorname{Sp}(2n, ℝ)``,
as defined in by Birtea et al in proposition 2 [^birtea2020optimization].

Using the symplectic inverse of a matrix ``A \in ℝ^{2n \times 2n}``,
``
A^{+} := Q_{2n}^T A^T Q_{2n}
``
where
````math
Q_{2n} =
\begin{bmatrix}
0_n & I_n \\
 -I_n & 0_n
\end{bmatrix},
````
the retraction
``\mathcal{R}\colon T\operatorname{Sp}(2n) \rightarrow \operatorname{Sp}(2n)``
is defined pointwise as
````math
\begin{align*}
\mathcal{R}_p(X) &= p \operatorname{cay}\left(\frac{1}{2}p^{+}X\right), \\
                 &= p \operatorname{exp}_{1/1}(p^{+}X), \\
                 &= p (2I - p^{+}X)^{-1}(2I + p^{+}X).
\end{align*}
````
Here
``\operatorname{exp}_{1/1}(z) = (2 - z)^{-1}(2 + z)``
denotes the Padé (1, 1) approximation to ``\operatorname{exp}(z)``.

[^birtea2020optimization]:
    > Birtea, Petre and Ca{\c{s}}u, Ioan and Com{\u{a}}nescu, Dan:
    > Optimization on the real symplectic group,
    > Monatshefte f{\"u}r Mathematik, Springer
    > doi [10.1007/s00605-020-01369-9](https://doi.org/10.1007/s00605-020-01369-9)
"""
retract(M::Symplectic, p, X)

function retract_caley!(M::Symplectic, q, p, X)
    p_star_X = symplectic_inverse_times(M, p, X)

    divisor = lu(2 * I - p_star_X)
    q .= p * (divisor \ add_scaled_I!(p_star_X, 2.0))
    return q
end

Base.show(io::IO, ::Symplectic{n,𝔽}) where {n,𝔽} = print(io, "Symplectic{$(2n), $(𝔽)}()")

@doc raw"""
    symplectic_inverse_times(::Symplectic, p, q)
    symplectic_inverse_times!(::Symplectic, A, p, q)

Directly compute the symplectic inverse of ``p \in \operatorname{Sp}(2n)``,
multiplied with ``q \in \operatorname{Sp}(2n)``.
That is, this function efficiently computes
``p^+q = (Q_{2n}p^TQ_{2n})q \in ℝ^{2n \times 2n}``,
where ``Q_{2n}`` is the [`SymplecticMatrix`](@ref)
of size ``2n \times 2n``.
"""
function symplectic_inverse_times(M::Symplectic{n}, p, q) where {n}
    A = similar(p)
    return symplectic_inverse_times!(M, A, p, q)
end

function symplectic_inverse_times!(::Symplectic{n}, A, p, q) where {n}
    # we write p = [p1 p2; p3 p4] (and q, too), then
    p1 = @view(p[1:n, 1:n])
    p2 = @view(p[1:n, (n + 1):(2n)])
    p3 = @view(p[(n + 1):(2n), 1:n])
    p4 = @view(p[(n + 1):(2n), (n + 1):(2n)])
    q1 = @view(q[1:n, 1:n])
    q2 = @view(q[1:n, (n + 1):(2n)])
    q3 = @view(q[(n + 1):(2n), 1:n])
    q4 = @view(q[(n + 1):(2n), (n + 1):(2n)])
    A1 = @view(A[1:n, 1:n])
    A2 = @view(A[1:n, (n + 1):(2n)])
    A3 = @view(A[(n + 1):(2n), 1:n])
    A4 = @view(A[(n + 1):(2n), (n + 1):(2n)])
    mul!(A1, p4', q1) # A1 = p4'q1
    mul!(A1, p2', q3, -1, 1) # A1 -= p2'p3
    mul!(A2, p4', q2) # A2 = p4'q2
    mul!(A2, p2', q4, -1, 1) #A2 -= p2'q4
    mul!(A3, p1', q3) #A3 = p1'q3
    mul!(A3, p3', q1, -1, 1) # A3 -= p3'q1
    mul!(A4, p1', q4) # A4 = p1'q4
    mul!(A4, p3', q2, -1, 1) #A4 -= p3'q2
    return A
end

ndims(Q::SymplecticMatrix) = 2
copy(Q::SymplecticMatrix) = SymplecticMatrix(copy(Q.λ))
Base.eltype(::SymplecticMatrix{T}) where {T} = T
function Base.convert(::Type{SymplecticMatrix{T}}, Q::SymplecticMatrix) where {T}
    return SymplecticMatrix(convert(T, Q.λ))
end

function Base.show(io::IO, Q::SymplecticMatrix)
    s = "$(Q.λ)"
    if occursin(r"\w+\s*[\+\-]\s*\w+", s)
        s = "($s)"
    end
    return print(io, typeof(Q), "(): $(s)*[0 I; -I 0]")
end

(Base.:-)(Q::SymplecticMatrix) = SymplecticMatrix(-Q.λ)

function (Base.:^)(Q::SymplecticMatrix, n::Integer)
    return ifelse(
        n % 2 == 0,
        UniformScaling((-1)^(div(n, 2)) * (Q.λ)^n),
        SymplecticMatrix((-1)^(div(n - 1, 2)) * (Q.λ)^n),
    )
end

(Base.:*)(x::Number, Q::SymplecticMatrix) = SymplecticMatrix(x * Q.λ)
(Base.:*)(Q::SymplecticMatrix, x::Number) = SymplecticMatrix(x * Q.λ)
function (Base.:*)(Q1::SymplecticMatrix, Q2::SymplecticMatrix)
    return LinearAlgebra.UniformScaling(-Q1.λ * Q2.λ)
end

Base.transpose(Q::SymplecticMatrix) = -Q
Base.adjoint(Q::SymplecticMatrix) = SymplecticMatrix(-conj(Q.λ))
Base.inv(Q::SymplecticMatrix) = SymplecticMatrix(-(1 / Q.λ))

(Base.:+)(Q1::SymplecticMatrix, Q2::SymplecticMatrix) = SymplecticMatrix(Q1.λ + Q2.λ)
(Base.:-)(Q1::SymplecticMatrix, Q2::SymplecticMatrix) = SymplecticMatrix(Q1.λ - Q2.λ)

(Base.:+)(Q::SymplecticMatrix, p::AbstractMatrix) = p + Q
function (Base.:+)(p::AbstractMatrix, Q::SymplecticMatrix)
    # When we are adding, the Matrices must match in size:
    two_n, two_k = size(p)
    if (two_n % 2 != 0) || (two_n != two_k)
        throw(
            ArgumentError(
                "'p' must be square with even row and dimension, " *
                "was: ($(two_n), $(two_k)) != (2n, 2n).",
            ),
        )
    end
    n = div(two_n, 2)

    # Allocate new memory:
    TS = typeof(one(eltype(p)) + one(eltype(Q)))
    out = copyto!(similar(p, TS), p)

    add_scaled_I!(view(out, 1:n, (n + 1):(2n)), Q.λ)
    add_scaled_I!(view(out, (n + 1):(2n), 1:n), -Q.λ)
    return out
end

# Binary minus:
(Base.:-)(Q::SymplecticMatrix, p::AbstractMatrix) = Q + (-p)
(Base.:-)(p::AbstractMatrix, Q::SymplecticMatrix) = p + (-Q)

function (Base.:*)(Q::SymplecticMatrix, p::AbstractVecOrMat)
    two_n = size(p)[1]
    if two_n % 2 != 0
        throw(ArgumentError("'p' must have even row dimension, was: $(two_n) != 2n."))
    end
    n = div(two_n, 2)

    # Allocate new memory:
    TS = typeof(one(eltype(p)) + one(eltype(Q)))
    Qp = similar(p, TS)

    # Perform left mulitply by λ*Q:
    mul!((@inbounds view(Qp, 1:n, :)), Q.λ, @inbounds view(p, (n + 1):lastindex(p, 1), :))
    mul!((@inbounds view(Qp, (n + 1):lastindex(Qp, 1), :)), -Q.λ, @inbounds view(p, 1:n, :))

    return Qp
end

function (Base.:*)(p::AbstractMatrix, Q::SymplecticMatrix)
    two_k = size(p)[2]
    if two_k % 2 != 0
        throw(ArgumentError("'p' must have even column dimension, was: $(two_k) != 2k."))
    end
    k = div(two_k, 2)

    # Allocate new memory:
    TS = typeof(one(eltype(p)) + one(eltype(Q)))
    pQ = similar(p, TS)

    # Perform right mulitply by λ*Q:
    mul!((@inbounds view(pQ, :, 1:k)), -Q.λ, @inbounds view(p, :, (k + 1):lastindex(p, 2)))
    mul!((@inbounds view(pQ, :, (k + 1):lastindex(pQ, 2))), Q.λ, @inbounds view(p, :, 1:k))
    return pQ
end

function LinearAlgebra.lmul!(Q::SymplecticMatrix, p::AbstractVecOrMat)
    # Perform left multiplication by a symplectic matrix,
    # overwriting the matrix p in place:
    two_n = size(p)[1]
    if two_n % 2 != 0
        throw(ArgumentError("'p' must have even row dimension, was: $(two_n) != 2n."))
    end
    n = div(two_n, 2)

    half_row_p = copy(@inbounds view(p, 1:n, :))

    mul!((@inbounds view(p, 1:n, :)), Q.λ, @inbounds view(p, (n + 1):lastindex(p, 1), :))

    mul!(
        (@inbounds view(p, (n + 1):lastindex(p, 1), :)),
        -Q.λ,
        @inbounds view(half_row_p, :, :)
    )
    return p
end

function LinearAlgebra.rmul!(p::AbstractMatrix, Q::SymplecticMatrix)
    # Perform right multiplication by a symplectic matrix,
    # overwriting the matrix p in place:
    two_k = size(p)[2]
    if two_k % 2 != 0
        throw(ArgumentError("'p' must have even column dimension, was: $(two_k) != 2k."))
    end
    k = div(two_k, 2)

    half_col_p = copy(@inbounds view(p, :, 1:k))

    mul!((@inbounds view(p, :, 1:k)), -Q.λ, @inbounds view(p, :, (k + 1):lastindex(p, 2)))

    mul!(
        (@inbounds view(p, :, (k + 1):lastindex(p, 2))),
        Q.λ,
        @inbounds view(half_col_p, :, :)
    )

    return p
end

function LinearAlgebra.mul!(A::AbstractVecOrMat, Q::SymplecticMatrix, p::AbstractVecOrMat)
    size_p = size(p)
    two_n = size_p[1]
    if two_n % 2 != 0
        throw(ArgumentError("'p' must have even row dimension, was: $(two_n) != 2n."))
    end
    n = div(two_n, 2)
    k = length(size_p) == 1 ? 0 : div(size_p[2], 2)

    # k == 0 means we're multiplying with a vector:
    @boundscheck k == 0 ? checkbounds(A, 1:(2n), 1) : checkbounds(A, 1:(2n), 1:(2k))

    # Perform left multiply by λ*Q:
    mul!((@inbounds view(A, 1:n, :)), Q.λ, @inbounds view(p, (n + 1):lastindex(p, 1), :))
    mul!((@inbounds view(A, (n + 1):lastindex(A, 1), :)), -Q.λ, @inbounds view(p, 1:n, :))
    return A
end

function LinearAlgebra.mul!(A::AbstractVecOrMat, p::AbstractMatrix, Q::SymplecticMatrix)
    two_n, two_k = size(p)
    if two_k % 2 != 0
        throw(ArgumentError("'p' must have even col dimension, was: $(two_k) != 2k."))
    end
    n = div(two_n, 2)
    k = div(two_k, 2)

    # n == 0 means we're multiplying with a vector:
    @boundscheck n == 0 ? checkbounds(A, 1, 1:(2k)) : checkbounds(A, 1:(2n), 1:(2k))

    # Perform right multiply by λ*Q:
    mul!((@inbounds view(A, :, 1:k)), -Q.λ, @inbounds view(p, :, (k + 1):lastindex(p, 2)))
    mul!((@inbounds view(A, :, (k + 1):lastindex(A, 2))), Q.λ, @inbounds view(p, :, 1:k))
    return A
end

function add_scaled_I!(A::AbstractMatrix, λ::Number)
    LinearAlgebra.checksquare(A)
    @inbounds for i in axes(A, 1)
        A[i, i] += λ
    end
    return A
end


text:
using SparseArrays: AbstractSparseMatrix

#  ======================================================
#  AdjacencyMatrix
#  ======================================================

"""
    AdjacencyMatrix(g::AbstractGraph)

A matrix view of a graph.

Entry `(i, j)` is `true` if `i -> j` is and edge of `g` and `false` otherwise.

As this is as a view, the entries and the size of this matrix can change
when the underlying graph changes. The view itself is immutable. Convert
to a `Matrix` or `SparseMatrixCSC` to get a mutable matrix that does
not change when the graph does.

### See also
[`adjacency_matrix`](@ref), [`ValMatrix`](@ref), [`weigths`](@ref)

### Examples

```jldoctest
julia> g = complete_bipartite_graph(2, 2)
{4, 4} undirected simple Int64 graph

julia> AdjacencyMatrix(g)
4×4 AdjacencyMatrix{SimpleGraph{Int64}}:
 0  0  1  1
 0  0  1  1
 1  1  0  0
 1  1  0  0

julia> AdjacencyMatrix(g) |> Matrix
4×4 Matrix{Bool}:
 0  0  1  1
 0  0  1  1
 1  1  0  0
 1  1  0  0
```
"""
struct AdjacencyMatrix{G <: AbstractGraph} <: AbstractSparseMatrix{Bool, Int}
    graph::G
end

Base.getindex(matrix::AdjacencyMatrix, s, d) = has_edge(matrix.graph, s, d)

function Base.size(matrix::AdjacencyMatrix)
    nvg = Int(nv(matrix.graph))
    return (nvg, nvg)
end

"""
    adjacency_matrix(g::AbstractValGraph)

Create an `AdjacencyMatrix` view from a graph `g`.

### See also
[`AdjacencyMatrix`](@ref), [`ValMatrix`](@ref), [`weights`](@ref)

### Examples

```jldoctest
julia> gv = ValGraph(star_graph(4), edgeval_types=(Float64,), edgeval_init=(s, d) -> (1.0, ))
{4, 3} undirected ValGraph with
              eltype: Int64
  vertex value types: ()
    edge value types: (Float64,)
   graph value types: ()

julia> adjacency_matrix(gv)
4×4 AdjacencyMatrix{ValGraph{Int64, Tuple{}, Tuple{Float64}, Tuple{}, Tuple{}, Tuple{Vector{Vector{Float64}}}}}:
 0  1  1  1
 1  0  0  0
 1  0  0  0
 1  0  0  0
```
"""
Graphs.adjacency_matrix(g::AbstractValGraph) = AdjacencyMatrix(g)

## ---------------------------------------------------------------
##       LinearAlgebra
## ---------------------------------------------------------------

function LinearAlgebra.ishermitian(matrix::AdjacencyMatrix)

    if !is_directed(matrix.graph)
        return true
    end
    invoke(LinearAlgebra.ishermitian, Tuple{AbstractSparseMatrix{Bool, Int}}, matrix)
end

function LinearAlgebra.issymmetric(matrix::AdjacencyMatrix)

    if !is_directed(matrix.graph)
        return true
    end
    invoke(LinearAlgebra.issymmetric, Tuple{AbstractSparseMatrix{Bool, Int}}, matrix)
end

function LinearAlgebra.adjoint(matrix::AdjacencyMatrix)

    if !is_directed(matrix.graph)
        return matrix
    end
    invoke(LinearAlgebra.adjoint, Tuple{AbstractSparseMatrix{Bool, Int}}, matrix)
end

function LinearAlgebra.transpose(matrix::AdjacencyMatrix)

    if !is_directed(matrix.graph)
        return matrix
    end
    invoke(LinearAlgebra.transpose, Tuple{AbstractSparseMatrix{Bool, Int}}, matrix)
end

# TODO consider implementing matrix x matrix multiplication
# TODO consider implementing the 5-argument version of mul! instead
function LinearAlgebra.mul!(y::AbstractVector, matrix::AdjacencyMatrix, b::AbstractVector)

    Base.require_one_based_indexing(y, matrix, b)
    n = length(y)
    size(matrix, 1) == n || throw(DimensionMismatch())
    length(b) == n || throw(DimensionMismatch())
    g = matrix.graph

    for i in 1:n
        s = false * zero(eltype(b))
        @simd for j in outneighbors(g, i)
            @inbounds s += true * b[j]
        end
        @inbounds y[i] = s
    end
    return y
end

##  ------------------------------------------------------
##  AdjacencyMatrix -> SparseMatrixCSC
##  ------------------------------------------------------

function SparseMatrixCSC(matrix::AdjacencyMatrix{<: Union{ValGraph, ValDiGraph}})

    g = matrix.graph
    n = Int(nv(g))
    nnz = is_directed(g) ? ne(g) : (2 * ne(g) - num_self_loops(g))

    colptr = Vector{Int}(undef, n + 1)
    rowval = Vector{Int}(undef, nnz)
    nzval = fill(true, nnz)

    nnz_idx = 1
    @inbounds colptr[1] = 1
    for v in vertices(g)
        for u in inneighbors(g, v)
            @inbounds rowval[nnz_idx] = u
            nnz_idx += 1
        end
        @inbounds colptr[v + 1] = nnz_idx
    end

    return SparseMatrixCSC(n, n, colptr, rowval, nzval)
end

# TODO Create directly without transpose
function SparseMatrixCSC(matrix::AdjacencyMatrix{<: ValOutDiGraph})

    g = matrix.graph
    n = Int(nv(g))
    nnz = is_directed(g) ? ne(g) : (2 * ne(g) - num_self_loops(g))

    colptr = Vector{Int}(undef, n + 1)
    rowval = Vector{Int}(undef, nnz)
    nzval = fill(true, nnz)

    nnz_idx = 1
    @inbounds colptr[1] = 1
    for u in vertices(g)
        for v in outneighbors(g, u)
            @inbounds rowval[nnz_idx] = v
            nnz_idx += 1
        end
        @inbounds colptr[u + 1] = nnz_idx
    end

    return SparseMatrixCSC(transpose(SparseMatrixCSC(n, n, colptr, rowval, nzval)))
end

#  ======================================================
#  ValMatrix
#  ======================================================

"""
    ValMatrix{Tv, :< AbstractGraph, key}

A matrix view of the edge values for a specific key of a graph.

As this is as a view, the entries and the size of this matrix can change
when the underlying graph changes. The view itself is immutable. Convert
to a `Matrix` or `SparseMatrixCSC` to get a mutable matrix that does
not change when the graph does.
"""
struct ValMatrix{Tv, G <: AbstractValGraph, key} <: AbstractSparseMatrix{Tv, Int}

    graph::G
    zero_value::Tv
end

"""
    ValMatrix(g::AbstractGraph, key, zero_value)

Construct a new `ValMatrix` view for a graph `g` where the values are
the edge values specified by `key`. Entries that are not in the graph
are represented by `zero_value` in the matrix.

### See also
[`AdjacencyMatrix`](@ref), [`adjacency_matrix`](@ref), [`weights`](@ref)

### Examples
```jldoctest
julia> gv = ValDiGraph(path_digraph(3),  edgeval_types=(a=Float64, b=String), edgeval_init=(s, d) -> (rand(MersenneTwister(0)), "\$s-\$d"))
{3, 2} directed ValDiGraph with
              eltype: Int64
  vertex value types: ()
    edge value types: (a = Float64, b = String)
   graph value types: ()

julia> ValMatrix(gv, 1, 0.0)
3×3 ValMatrix{Float64, ValDiGraph{Int64, Tuple{}, NamedTuple{(:a, :b), Tuple{Float64, String}}, Tuple{}, Tuple{}, NamedTuple{(:a, :b), Tuple{Vector{Vector{Float64}}, Vector{Vector{String}}}}}, 1}:
 0.0  0.823648  0.0
 0.0  0.0       0.823648
 0.0  0.0       0.0

julia> ValMatrix(gv, :b, nothing)
3×3 ValMatrix{Union{Nothing, String}, ValDiGraph{Int64, Tuple{}, NamedTuple{(:a, :b), Tuple{Float64, String}}, Tuple{}, Tuple{}, NamedTuple{(:a, :b), Tuple{Vector{Vector{Float64}}, Vector{Vector{String}}}}}, :b}:
 nothing  "1-2"    nothing
 nothing  nothing  "2-3"
 nothing  nothing  nothing
```
"""
function ValMatrix(g::AbstractValGraph, key::Union{Integer, Symbol}, zero_value)

    T = edgevals_type(g, key)
    Z = typeof(zero_value)
    Tv = Union{T, Z}

    return ValMatrix{Tv, typeof(g), key}(g, zero_value)
end

function Base.size(matrix::ValMatrix)

    nvg = Int(nv(matrix.graph))
    return (nvg, nvg)
end

function Base.getindex(matrix::ValMatrix{Tv, G, key}, s, d) where {Tv, G, key}

    return get_edgeval_or(matrix.graph, s, d, key, matrix.zero_value)
end

LinearAlgebra.ishermitian(::ValMatrix{ <: Real, <: ValGraph}) = true
LinearAlgebra.issymmetric(::ValMatrix{ <: Any, <: ValGraph}) = true

function LinearAlgebra.ishermitian(matrix::ValMatrix{Tv}) where {Tv <: Real}

    if !is_directed(matrix.graph)
        return true
    end
    invoke(LinearAlgebra.ishermitian, Tuple{AbstractSparseMatrix{Tv, Int}}, matrix)
end

function LinearAlgebra.issymmetric(matrix::ValMatrix{Tv}) where {Tv}

    if !is_directed(matrix.graph)
        return true
    end
    invoke(LinearAlgebra.issymmetric, Tuple{AbstractSparseMatrix{Tv, Int}}, matrix)
end

LinearAlgebra.adjoint(matrix::ValMatrix{ <: Real, <: ValGraph}) = matrix
LinearAlgebra.transpose(matrix::ValMatrix{ <: Any, <: ValGraph}) = matrix

function LinearAlgebra.adjoint(matrix::ValMatrix{Tv}) where {Tv <: Real}

    if !is_directed(matrix.graph)
        return matrix
    end
    invoke(LinearAlgebra.adjoint, Tuple{AbstractSparseMatrix{Tv, Int}}, matrix)
end

function LinearAlgebra.transpose(matrix::ValMatrix{Tv}) where {Tv}

    if !is_directed(matrix.graph)
        return matrix
    end
    invoke(LinearAlgebra.transpose, Tuple{AbstractSparseMatrix{Tv, Int}}, matrix)
end

##  ------------------------------------------------------
##  SparseMatrixCSC
##  ------------------------------------------------------

function SparseMatrixCSC(matrix::ValMatrix{Tv, <: Union{ValGraph, ValDiGraph}, key}) where {Tv, key}

    g = matrix.graph
    n = Int(nv(g))
    nnz = is_directed(g) ? ne(g) : (2 * ne(g) - num_self_loops(g))

    colptr = Vector{Int}(undef, n + 1)
    rowval = Vector{Int}(undef, nnz)
    nzval = Vector{Tv}(undef, nnz)

    nnz_idx = 1
    @inbounds colptr[1] = 1
    for v in vertices(g)
        for (u, val) in zip(inneighbors(g, v), inedgevals(g, v, key))
            @inbounds rowval[nnz_idx] = u
            @inbounds nzval[nnz_idx] = val
            nnz_idx += 1
        end
        @inbounds colptr[v + 1] = nnz_idx
    end

    return SparseMatrixCSC(n, n, colptr, rowval, nzval)
end

# TODO instead of using transpose, a better way should be found that
# does not allocate a temporary SparseMatrixCSC
function SparseMatrixCSC(matrix::ValMatrix{Tv, <: ValOutDiGraph, key}) where {Tv, key}

    g = matrix.graph
    n = Int(nv(g))
    nnz = ne(g)

    colptr = Vector{Int}(undef, n + 1)
    rowval = Vector{Int}(undef, nnz)
    nzval = Vector{Tv}(undef, nnz)

    nnz_idx = 1
    @inbounds colptr[1] = 1
    for u in vertices(g)
        for (v, val) in zip(outneighbors(g, u), outedgevals(g, u, key))
            @inbounds rowval[nnz_idx] = v
            @inbounds nzval[nnz_idx] = val
            nnz_idx += 1
        end
        @inbounds colptr[u + 1] = nnz_idx
    end

    return SparseMatrixCSC(transpose(SparseMatrixCSC(n, n, colptr, rowval, nzval)))
end

##  ------------------------------------------------------
##  weights
##  ------------------------------------------------------

"""
    weights(g::AbstractValGraph[, key]; zerovalue)

Return a matrix where entry (i,j) is the value of the edge `i -- j` for the specific `key` in `g`.

If `g` has no edge values, `key` should be omitted and `DefaultDistance(g)` is returned.
Otherwise the result is a `ValMatrix`.
If `g` has a single value per edge, `key` can be omitted.
If the optional argument `zerovalue` is specified, then this value will be used
if entry (i, j) is not an edge in `g`. `zerovalue` cannot be used for graphs
without any edge values.
"""
function weights end

function Graphs.weights(g::AbstractValGraph, key; zerovalue=zero(edgevals_type(g, key)))

    return ValMatrix(g, key, zerovalue)
end

function Graphs.weights(g::AbstractValGraph; kwargs...)

    g isa ZeroEdgeValGraph && return Graphs.DefaultDistance(nv(g))
    g isa OneEdgeValGraph && return weights(g, 1; kwargs...)

    throw(ArgumentError("g has more than one edge value key. Must specify key."))
end