using TSVD
using LinearAlgebra
using ToeplitzMatrices
using ARMA

make_mpr(data, basis) = pinv(basis)*data
make_time_domain(mpr, basis) = basis*mpr
function make_std_residuals(data, basis)
    mpr = make_mpr(data,basis) # model pulse reduced
    td = make_time_domain(mpr,basis)
    std_residuals = std(data-td, dims=1)[1,:]
end

function getall(ljh, maxrecords=typemax(Int))
    records = collect(ljh[1:min(Int(maxrecords),length(ljh))])
    pulses = Array{Float32}(undef, ljh.record_nsamples,length(records))
    for (i,record) in enumerate(records)
        pulses[:,i]=record.data
    end
    pulses
end

function evenly_distributed_inds(inds,n_wanted)
    frac_keep = n_wanted/length(inds)
    keep_each_n = max(1,floor(Int,1/frac_keep))
    inds[1:keep_each_n:length(inds)]
end

function choose_new_train_inds(residuals, train_inds, frac_keep)
    n_keep = round(Int, frac_keep*length(train_inds))
    sort(train_inds, by = i -> residuals[i])[1:n_keep]
end

function train_loop(data, n_pulses_for_train, n_basis, n_presamples, n_loop, frac_keep_per_loop, make_basis)
    last_train_inds = train_inds = evenly_distributed_inds(1:size(data,2),n_pulses_for_train)
    # Use only a 3-Dimensional basis the first time through. This helps reject pileup before it
    # can sneak into the SVD.
    nb = min(n_basis, 3)
    for i=1:n_loop
        basis, singular_values = make_basis(data[:,train_inds],nb)
        residuals = make_std_residuals(data,basis)
        last_train_inds = train_inds
        train_inds = choose_new_train_inds(residuals, train_inds, frac_keep_per_loop)
        if i==n_loop # variables created in loop are not availble outside the loop
            return basis, residuals, last_train_inds, train_inds, singular_values
        end
        nb = n_basis
    end
end

function TSVD_tsvd(data_train::Matrix{<:AbstractFloat},n_basis)
    U, S, V = TSVD.tsvd(data_train, n_basis)
    U, S
end

"""    TSVD_tsvd_mass3(data_train::Matrix{<:AbstractFloat}, n_basis, n_presamples, noise_model::ARMA.ARMAModel, noise_solver::ARMA.ARMASolver)

This function is called as the first step towards creating a basis with `n_basis` elements (minimum 3).
Elements will be 1) constant, 2) derivative-like, 3) pulse-like with first `n_presamples` equal to zero.
Elements 4 and higher are based on SVD of the noise whitened residuals after projecting into a basis of
just the first 3 elements. This ensures that adding additional components will not change the projector
for the first 3 components, and therefore the energy resolution obtained with only the first 3 projectors
will be independing of number of components.

The basis is paired with projectors, which allow you to project into the basis while minimizing
the noise whitened residuals (aka the mahalonobis distance).

This function returns `U,S` much like a truncated singular value decomposition, but such that the
basis caluclated from `U,S` will meet the description above."""
function TSVD_tsvd_mass3(data_train::Matrix{<:AbstractFloat}, n_basis, n_presamples,
        noise_model::ARMA.ARMAModel, noise_solver::ARMA.ARMASolver)
    @assert n_basis>=3 "mean, derivative and average pulse are 3 components, must request at least 3 components"
    # Average pulse is the pretrigger-mean-subtracted average pulse, rescaled to have a maximum value of 1.0
    average_pulse = mean(data_train, dims=2)[:]
    if n_presamples > 0
        average_pulse .-= mean(average_pulse[1:n_presamples])
    end
    average_pulse[1:n_presamples] .= 0.0
    average_pulse /= maximum(abs.(average_pulse))
    # Calculate the derivative like component; assume clean baseline, so low derivative at start
    derivative_like = [0.0;diff(average_pulse)]
    constant_component = ones(size(data_train,1))
    mass3_basis = hcat(constant_component,derivative_like,average_pulse)
    if n_basis == 3
        return mass3_basis, [NaN,NaN,NaN]
    end
    projectors3, _ = computeprojectors(mass3_basis, noise_model)
    mpr = projectors3 * data_train # model pulse reduced
    td = mass3_basis * mpr
    data_residual = data_train .- td
    # Whiten residual before taking the TSVD
    white_residual = ARMA.whiten(noise_solver, data_residual)
    Uwhite, S, _ = TSVD.tsvd(white_residual, n_basis-3)
    # Unwhiten and renormalize U before combining
    Uunnorm = ARMA.unwhiten(noise_solver, Uwhite)
    U = Uunnorm ./ sqrt.(sum(Uunnorm.^2, dims=1))
    U_combined = hcat(mass3_basis,U)
    S_combined = vcat([NaN,NaN,NaN],S) # first 3 elemnts are not from svd, have no meaningful singular value
    U_combined, S_combined
end

"""
    TSVD_tsvd_mass3(data_train::Matrix{<:AbstractFloat}, n_basis, n_presamples, autocorr::Vector)

Create mass3+SVD as the other method does, but use the noise autocorrelation (`autocorr`)
instead of using a noise model. This should mirror the MASS behavior.
"""
function TSVD_tsvd_mass3(data_train::Matrix{<:AbstractFloat}, n_basis, n_presamples, autocorr::Vector)
    @assert n_basis>=3 "mean, derivative and average pulse are 3 components, must request at least 3 components"
    # Average pulse is the pretrigger-mean-subtracted average pulse, rescaled to have a maximum value of 1.0
    average_pulse = mean(data_train, dims=2)[:]
    if n_presamples > 0
        average_pulse .-= mean(average_pulse[1:n_presamples])
    end
    average_pulse[1:n_presamples] .= 0.0
    average_pulse /= maximum(abs.(average_pulse))
    # Calculate the derivative like component; assume clean baseline, so low derivative at start
    derivative_like = [0.0;diff(average_pulse)]
    constant_component = ones(size(data_train,1))
    mass3_basis = hcat(constant_component,derivative_like,average_pulse)
    if n_basis == 3
        return mass3_basis, [NaN,NaN,NaN]
    end
    projectors3, _ = computeprojectors(mass3_basis, autocorr)
    mpr = projectors3 * data_train # model pulse reduced
    td = mass3_basis * mpr
    data_residual = data_train .- td
    # Whiten residual before taking the TSVD
    R = SymmetricToeplitz(autocorr)
    Winv = cholesky(R).L  # The inverse of Winv will whiten a data record.
    white_residual = Winv \ data_residual
    Uwhite, S, _ = TSVD.tsvd(white_residual, n_basis-3)
    # Unwhiten and renormalize U before combining
    Uunnorm = Winv * Uwhite
    U = Uunnorm ./ sqrt.(sum(Uunnorm.^2, dims=1))
    U_combined = hcat(mass3_basis, U)
    S_combined = vcat([NaN,NaN,NaN], S) # first 3 elemnts are not from svd, have no meaningful singular value
    U_combined, S_combined
end

function full_svd(data_train, n_basis)
    U,S,V = svd(data_train)
    U[:,1:n_basis], S[1:n_basis]
end
tsvd_dict = Dict(
    "TSVD" => TSVD_tsvd,
    "full" => full_svd,
    "noisemass3" => TSVD_tsvd_mass3,
    "TSVDmass3" => TSVD_tsvd_mass3)

struct SVDBasis <: AbstractBasisAnalyzer
    basis::Array{Float32,2} # unwhitened time domain representation, tall: basis*mpr = reconsituted pulse
    projectors::Array{Float32,2} # whitened time domain representation, wide: projectors*data = mpr
    projector_covariance::Array{Float32,2}
    noise_result::NoiseAnalysis.NoiseResult
end
(a::SVDBasis)(record) = record2dataproduct(a,record)

struct SVDBasisWithCreationInfo
    svdbasis::SVDBasis
    singular_values::Vector{Float32}
    example_pulses::Array{UInt16,2}
    std_residuals_of_example_pulses::Vector{Float32}
    percentiles_of_sample_pulses::Vector{Float32}
    n_loop::Int
    noise_model_file::String
    pulse_file::String
    std_residuals::Vector{Float32}
    tsvd_method::String
    channel_number::Int
    noise_std_dev::Float64
end

function create_basis_one_channel(ljh, noise_result, frac_keep, n_loop, n_pulses_for_train, n_basis,
    tsvd_method_string)
    data = getall(ljh, ceil(Int,n_pulses_for_train/frac_keep))
    create_basis_one_channel(data, noise_result, frac_keep, n_loop,
        n_pulses_for_train, n_basis, tsvd_method_string,
        LJH.first_rising_sample(ljh[1]),
        LJH.filename(ljh), LJH.channel(ljh))
end

function create_basis_one_channel(data::Matrix{<:AbstractFloat}, noise_result, frac_keep, n_loop, n_pulses_for_train, n_basis,
    tsvd_method_string, n_presamples, datasource_filename, datasource_channel)
    tsvd_func = tsvd_dict[tsvd_method_string]
    if tsvd_method_string == "TSVDmass3"
        noise_solver = ARMA.ARMASolver(noise_result.model, size(data, 1))
        function tsvd_closure(data, n_basis)
            TSVD_tsvd_mass3(data, n_basis, n_presamples, noise_result.model, noise_solver)
        end
        tsvd_func = tsvd_closure
    elseif tsvd_method_string == "noisemass3"
        function tsvd_closure2(data, n_basis)
            TSVD_tsvd_mass3(data, n_basis, n_presamples, noise_result.autocorr)
        end
        tsvd_func = tsvd_closure2
    end

    std_noise = sqrt(noise_result.autocorr[1]) # the first element in the autocorrelation is the variance
    frac_keep_per_loop = exp(log(frac_keep)/n_loop)
    basis, residual_stds, last_train_inds, train_inds, singular_values = train_loop(data,
        n_pulses_for_train, n_basis, n_presamples, n_loop, frac_keep_per_loop, tsvd_func)
    if tsvd_method_string == "noisemass3"
        projectors, pcovar = computeprojectors(basis,noise_result.autocorr)
    else
        projectors, pcovar = computeprojectors(basis,noise_result.model)
    end
    svdbasis = SVDBasis(
        basis,
        projectors,
        pcovar,
        noise_result)
    sortinds = sortperm(residual_stds)
    percentiles = [10:10:90; 91:99]
    percentile_indicies = sortinds[[round(Int,(p/100)*length(residual_stds)) for p in percentiles]]
    example_pulses = round.(UInt16,data[:,percentile_indicies])
    svdbasis_with_info = SVDBasisWithCreationInfo(
        svdbasis, singular_values, example_pulses,
        residual_stds[percentile_indicies], percentiles,
        n_loop, noise_result.datasource, datasource_filename,
        residual_stds, tsvd_method_string, datasource_channel,
        √(svdbasis.noise_result.autocorr[1])
        )
    return svdbasis,svdbasis_with_info
end

function hdf5save(g::HDF5.DataFile, svdbasis::SVDBasis)
    g["basis"]=svdbasis.basis
    g["projectors"]=svdbasis.projectors
    g["projector_covariance"]=svdbasis.projector_covariance
    Pope.NoiseAnalysis.hdf5save(g_create(g,"noise_result"),svdbasis.noise_result)
end

function hdf5save(g::HDF5.DataFile, x::SVDBasisWithCreationInfo)
    hdf5save(g_create(g,"svdbasis"),x.svdbasis)
    g["singular_values"]=x.singular_values
    g["example_pulses"]=x.example_pulses
    g["std_residuals_of_example_pulses"]=x.std_residuals_of_example_pulses
    g["percentiles_of_sample_pulses"]=x.percentiles_of_sample_pulses
    g["n_loop"]=x.n_loop
    g["noise_model_file"]=x.noise_model_file
    g["pulse_file"]=x.pulse_file
    g["std_residuals"]=x.std_residuals
    g["tsvd_method"]=x.tsvd_method
    g["channel_number"]=x.channel_number
    g["noise_std_dev"]=x.noise_std_dev
end

function hdf5load(T::Type{SVDBasis},g::HDF5.DataFile)
    SVDBasis(
    read(g["basis"]),
    read(g["projectors"]),
    read(g["projector_covariance"]),
    NoiseAnalysis.hdf5load(g["noise_result"])
    )
end

function hdf5load(T::Type{SVDBasisWithCreationInfo},g::HDF5.DataFile)
    SVDBasisWithCreationInfo(
    hdf5load(SVDBasis,g["svdbasis"]),
    read(g["singular_values"]),
    read(g["example_pulses"]),
    read(g["std_residuals_of_example_pulses"]),
    read(g["percentiles_of_sample_pulses"]),
    read(g["n_loop"]),
    read(g["noise_model_file"]),
    read(g["pulse_file"]),
    read(g["std_residuals"]),
    read(g["tsvd_method"]),
    read(g["channel_number"]),
    read(g["noise_std_dev"])
    )
end

function make_basis_one_channel(outputh5, ljhname, noise_filename, frac_keep, n_loop,
    n_pulses_for_train, n_basis, tsvd_method)
    ljh = LJH.ljhopen(ljhname)
    @show LJH.channel(ljh)
    @show ljhname
    @show ljh
    noise_result = NoiseAnalysis.hdf5load(noise_filename,LJH.channel(ljh))
    svdbasis, svdbasiswithcreationinfo = create_basis_one_channel(ljh, noise_result,
        frac_keep,
        n_loop,
        n_pulses_for_train,
        n_basis,
        tsvd_method)
    hdf5save(g_create(outputh5,"$(LJH.channel(ljh))"), svdbasiswithcreationinfo)
end

function make_basis_all_channel(outputh5, ljhdict, noise_filename, frac_keep, n_loop,
    n_pulses_for_train, n_basis, tsvd_method)
    noise_h5 = h5open(noise_filename,"r")
    noise_channels = sort!(parse.(Int,names(noise_h5)))
    ljh_channels = keys(ljhdict)
    bothchannels = sort!(intersect(noise_channels,ljh_channels))
    @show noise_channels
    @show ljh_channels
    @show bothchannels
    for channel_number in bothchannels
        println("\nmaking basis for channel $channel_number")
        ljhname = ljhdict[channel_number]
        try
            make_basis_one_channel(outputh5, ljhname, noise_filename, frac_keep, n_loop,
                n_pulses_for_train, n_basis, tsvd_method)
        catch ex
            println("channel $channel_number failed:")
            println(ex)
        end
    end
    println("Channels in noise_file but not in ljh: $(setdiff(noise_channels,ljh_channels))")
    println("Channels in ljh but not in noise_file: $(setdiff(ljh_channels,noise_channels))")
end


@doc raw"""
    quasi_Newton(M, F, gradF, x)

Perform a quasi Newton iteration for `F` on the manifold `M` starting
in the point `x` using a retraction ``R`` and a vector transport ``T``

The ``k``th iteration consists of
1. Compute the search direction ``η_k = -\mathcal{B}_k [\operatorname{grad}f (x_k)]`` or solve ``\mathcal{H}_k [η_k] = -\operatorname{grad}f (x_k)]``.
2. Determine a suitable stepsize ``α_k`` along the curve ``\gamma(α) = R_{x_k}(α η_k)`` e.g. by using [`WolfePowellLineseach`](@ref).
3. Compute ``x_{k+1} = R_{x_k}(α_k η_k)``.
4. Define ``s_k = T_{x_k, α_k η_k}(α_k η_k)`` and ``y_k = \operatorname{grad}f(x_{k+1}) - T_{x_k, α_k η_k}(\operatorname{grad}f(x_k))``.
5. Compute the new approximate Hessian ``H_{k+1}`` or its inverse ``B_k``.

# Input
* `M` – a manifold ``\mathcal{M}``.
* `F` – a cost function ``F : \mathcal{M} →ℝ`` to minimize.
* `gradF`– the gradient ``\operatorname{grad}F : \mathcal{M} →T_x\mathcal M`` of ``F``.
* `x` – an initial value ``x ∈ \mathcal{M}``.

    stopping_criterion::StoppingCriterion=StopWhenAny(
        StopAfterIteration(max(1000, memory_size)), StopWhenGradientNormLess(10^(-6))
    ),
    return_options=false,

# Optional
* `basis` – (`DefaultOrthonormalBasis()`) basis within the tangent space(s) to represent the
  Hessian (inverse).
* `cautious_update` – (`false`) – whether or not to use
  a [`QuasiNewtonCautiousDirectionUpdate`](@ref)
* `cautious_function` – (`(x) -> x*10^(-4)`) – a monotone increasing function that is zero
  at 0 and strictly increasing at 0 for the cautious update.
* `direction_update` – ([`InverseBFGS`](@ref)`()`) the update rule to use.
* `evaluation` – ([`AllocatingEvaluation`](@ref)) specify whether the gradient works by
   allocation (default) form `gradF(M, x)` or [`MutatingEvaluation`](@ref) in place, i.e.
   is of the form `gradF!(M, X, x)`.
* `initial_operator` – (`Matrix{Float64}(I,n,n)`) initial matrix to use die the
  approximation, where `n=manifold_dimension(M)`, see also `scale_initial_operator`.
* `memory_size` – (`20`) limited memory, number of ``s_k, y_k`` to store. Set to a negative
  value to use a full memory representation
* `retraction_method` – (`ExponentialRetraction()`) a retraction method to use, by default
  the exponential map.
* `scale_initial_operator` - (`true`) scale initial operator with
  ``\frac{⟨s_k,y_k⟩_{x_k}}{\lVert y_k\rVert_{x_k}}`` in the computation
* `step_size` – ([`WolfePowellLineseach`](@ref)`(retraction_method, vector_transport_method)`)
  specify a [`Stepsize`](@ref).
* `stopping_criterion` - (`StopWhenAny(StopAfterIteration(max(1000, memory_size)), StopWhenGradientNormLess(10^(-6))`)
  specify a [`StoppingCriterion`](@ref)
* `vector_transport_method` – (`ParallelTransport()`) a vector transport to use.
* `return_options` – (`false`) – specify whether to return just the result `x` (default) or the complete [`Options`](@ref), e.g. to access recorded values. if activated, the extended result, i.e. the

# Output
* `x_opt` – the resulting (approximately critical) point of the quasi–Newton method
OR
* `options` – the options returned by the solver (see `return_options`)
"""
function quasi_Newton(
    M::AbstractManifold, F::Function, gradF::G, x::P; kwargs...
) where {P,G}
    x_res = allocate(x)
    copyto!(M, x_res, x)
    return quasi_Newton!(M, F, gradF, x_res; kwargs...)
end
@doc raw"""
    quasi_Newton!(M, F, gradF, x; options...)

Perform a quasi Newton iteration for `F` on the manifold `M` starting
in the point `x` using a retraction ``R`` and a vector transport ``T``.

# Input
* `M` – a manifold ``\mathcal{M}``.
* `F` – a cost function ``F: \mathcal{M} →ℝ`` to minimize.
* `gradF`– the gradient ``\operatorname{grad}F : \mathcal{M} → T_x\mathcal M`` of ``F``.
* `x` – an initial value ``x ∈ \mathcal{M}``.

For all optional parameters, see [`quasi_Newton`](@ref).
"""
function quasi_Newton!(
    M::AbstractManifold,
    F::Function,
    gradF::G,
    x::P;
    cautious_update::Bool=false,
    cautious_function::Function=x -> x * 10^(-4),
    retraction_method::AbstractRetractionMethod=ExponentialRetraction(),
    vector_transport_method::AbstractVectorTransportMethod=ParallelTransport(),
    basis::AbstractBasis=DefaultOrthonormalBasis(),
    direction_update::AbstractQuasiNewtonUpdateRule=InverseBFGS(),
    evaluation::AbstractEvaluationType=AllocatingEvaluation(),
    memory_size::Int=20,
    initial_operator::AbstractMatrix=Matrix{Float64}(
        I, manifold_dimension(M), manifold_dimension(M)
    ),
    scale_initial_operator::Bool=true,
    step_size::Stepsize=WolfePowellLineseach(retraction_method, vector_transport_method),
    stopping_criterion::StoppingCriterion=StopWhenAny(
        StopAfterIteration(max(1000, memory_size)), StopWhenGradientNormLess(10^(-6))
    ),
    return_options=false,
    kwargs...,
) where {P,G}
    if memory_size >= 0
        local_dir_upd = QuasiNewtonLimitedMemoryDirectionUpdate(
            direction_update,
            zero_vector(M, x),
            memory_size;
            scale=scale_initial_operator,
            vector_transport_method=vector_transport_method,
        )
    else
        local_dir_upd = QuasiNewtonMatrixDirectionUpdate(
            direction_update,
            basis,
            initial_operator;
            scale=scale_initial_operator,
            vector_transport_method=vector_transport_method,
        )
    end
    if cautious_update == true
        local_dir_upd = QuasiNewtonCautiousDirectionUpdate(
            local_dir_upd; θ=cautious_function
        )
    end

    p = GradientProblem(M, F, gradF; evaluation=evaluation)
    o = QuasiNewtonOptions(
        x,
        get_gradient(p, x),
        local_dir_upd,
        stopping_criterion,
        step_size;
        retraction_method=retraction_method,
        vector_transport_method=vector_transport_method,
    )

    o = decorate_options(o; kwargs...)
    resultO = solve(p, o)

    if return_options
        return resultO
    else
        return get_solver_result(resultO)
    end
end

function initialize_solver!(::GradientProblem, ::QuasiNewtonOptions) end

function step_solver!(p::GradientProblem, o::QuasiNewtonOptions, iter)
    o.gradient = get_gradient(p, o.x)
    η = o.direction_update(p, o)
    α = o.stepsize(p, o, iter, η)
    x_old = deepcopy(o.x)
    retract!(p.M, o.x, o.x, α * η, o.retraction_method)
    β = locking_condition_scale(
        p.M, o.direction_update, x_old, α * η, o.x, o.vector_transport_method
    )
    vector_transport_to!(
        p.M, o.sk, x_old, α * η, o.x, get_update_vector_transport(o.direction_update)
    )
    vector_transport_to!(
        p.M,
        o.gradient,
        x_old,
        o.gradient,
        o.x,
        get_update_vector_transport(o.direction_update),
    )
    o.yk = get_gradient(p, o.x) / β - o.gradient
    update_hessian!(o.direction_update, p, o, x_old, iter)
    return o
end

function locking_condition_scale(
    M::AbstractManifold, ::AbstractQuasiNewtonDirectionUpdate, x_old, v, x, vt
)
    return norm(M, x_old, v) / norm(M, x, vector_transport_to(M, x_old, v, x, vt))
end

@doc raw"""
    update_hessian!(d, p, o, x_old, iter)

update the hessian wihtin the [`QuasiNewtonOptions`](@ref) `o` given a [`Problem`](@ref) `p`
as well as the an [`AbstractQuasiNewtonDirectionUpdate`](@ref) `d` and the last iterate `x_old`.
Note that the current (`iter`th) iterate is already stored in `o.x`.

See also [`AbstractQuasiNewtonUpdateRule`](@ref) for the different rules that are available
within `d`.
"""
update_hessian!(d::AbstractQuasiNewtonDirectionUpdate, ::Any, ::Any, ::Any, ::Any)

function update_hessian!(
    d::QuasiNewtonMatrixDirectionUpdate{InverseBFGS}, p, o, x_old, iter
)
    update_basis!(d.basis, p.M, x_old, o.x, d.vector_transport_method)
    yk_c = get_coordinates(p.M, o.x, o.yk, d.basis)
    sk_c = get_coordinates(p.M, o.x, o.sk, d.basis)
    skyk_c = inner(p.M, o.x, o.sk, o.yk)
    if iter == 1 && d.scale == true
        d.matrix = skyk_c / inner(p.M, o.x, o.yk, o.yk) * d.matrix
    end
    d.matrix =
        (I - sk_c * yk_c' / skyk_c) * d.matrix * (I - yk_c * sk_c' / skyk_c) +
        sk_c * sk_c' / skyk_c
    return d
end

# BFGS update
function update_hessian!(d::QuasiNewtonMatrixDirectionUpdate{BFGS}, p, o, x_old, iter)
    update_basis!(d.basis, p.M, x_old, o.x, d.vector_transport_method)
    yk_c = get_coordinates(p.M, o.x, o.yk, d.basis)
    sk_c = get_coordinates(p.M, o.x, o.sk, d.basis)
    skyk_c = inner(p.M, o.x, o.sk, o.yk)
    if iter == 1 && d.scale == true
        d.matrix = inner(p.M, o.x, o.yk, o.yk) / skyk_c * d.matrix
    end
    d.matrix =
        d.matrix + yk_c * yk_c' / skyk_c -
        d.matrix * sk_c * sk_c' * d.matrix / dot(sk_c, d.matrix * sk_c)
    return d
end

# Inverse DFP update
function update_hessian!(d::QuasiNewtonMatrixDirectionUpdate{InverseDFP}, p, o, x_old, iter)
    update_basis!(d.basis, p.M, x_old, o.x, d.vector_transport_method)
    yk_c = get_coordinates(p.M, o.x, o.yk, d.basis)
    sk_c = get_coordinates(p.M, o.x, o.sk, d.basis)
    skyk_c = inner(p.M, o.x, o.sk, o.yk)
    if iter == 1 && d.scale == true
        d.matrix = inner(p.M, o.x, o.sk, o.sk) / skyk_c * d.matrix
    end
    d.matrix =
        d.matrix + sk_c * sk_c' / skyk_c -
        d.matrix * yk_c * yk_c' * d.matrix / dot(yk_c, d.matrix * yk_c)
    return d
end

# DFP update
function update_hessian!(d::QuasiNewtonMatrixDirectionUpdate{DFP}, p, o, x_old, iter)
    update_basis!(d.basis, p.M, x_old, o.x, d.vector_transport_method)
    yk_c = get_coordinates(p.M, o.x, o.yk, d.basis)
    sk_c = get_coordinates(p.M, o.x, o.sk, d.basis)
    skyk_c = inner(p.M, o.x, o.sk, o.yk)
    if iter == 1 && d.scale == true
        d.matrix = skyk_c / inner(p.M, o.x, o.sk, o.sk) * d.matrix
    end
    d.matrix =
        (I - yk_c * sk_c' / skyk_c) * d.matrix * (I - sk_c * yk_c' / skyk_c) +
        yk_c * yk_c' / skyk_c
    return d
end

# Inverse SR-1 update
function update_hessian!(
    d::QuasiNewtonMatrixDirectionUpdate{InverseSR1}, p, o, x_old, ::Int
)
    update_basis!(d.basis, p.M, x_old, o.x, d.vector_transport_method)
    yk_c = get_coordinates(p.M, o.x, o.yk, d.basis)
    sk_c = get_coordinates(p.M, o.x, o.sk, d.basis)

    # computing the new matrix which represents the approximating operator in the next iteration
    srvec = sk_c - d.matrix * yk_c
    if d.update.r < 0 || abs(dot(srvec, yk_c)) >= d.update.r * norm(srvec) * norm(yk_c)
        d.matrix = d.matrix + srvec * srvec' / (srvec' * yk_c)
    end
    return d
end

# SR-1 update
function update_hessian!(d::QuasiNewtonMatrixDirectionUpdate{SR1}, p, o, x_old, ::Int)
    update_basis!(d.basis, p.M, x_old, o.x, d.vector_transport_method)
    yk_c = get_coordinates(p.M, o.x, o.yk, d.basis)
    sk_c = get_coordinates(p.M, o.x, o.sk, d.basis)

    # computing the new matrix which represents the approximating operator in the next iteration
    srvec = yk_c - d.matrix * sk_c
    if d.update.r < 0 || abs(dot(srvec, sk_c)) >= d.update.r * norm(srvec) * norm(sk_c)
        d.matrix = d.matrix + srvec * srvec' / (srvec' * sk_c)
    end
    return d
end

# Inverse Broyden update
function update_hessian!(
    d::QuasiNewtonMatrixDirectionUpdate{InverseBroyden}, p, o, x_old, ::Int
)
    update_basis!(d.basis, p.M, x_old, o.x, d.vector_transport_method)
    yk_c = get_coordinates(p.M, o.x, o.yk, d.basis)
    sk_c = get_coordinates(p.M, o.x, o.sk, d.basis)
    skyk_c = inner(p.M, o.x, o.sk, o.yk)
    ykBkyk_c = yk_c' * d.matrix * yk_c

    φ = update_broyden_factor!(d, sk_c, yk_c, skyk_c, ykBkyk_c, d.update.update_rule)
    # computing the new matrix which represents the approximating operator in the next iteration
    d.matrix =
        d.matrix - (d.matrix * yk_c * yk_c' * d.matrix) / ykBkyk_c +
        (sk_c * sk_c') / skyk_c +
        φ * ykBkyk_c * (sk_c / skyk_c - (d.matrix * yk_c) / ykBkyk_c) *
        (sk_c / skyk_c - (d.matrix * yk_c) / ykBkyk_c)'
    return d
end

# Broyden update
function update_hessian!(d::QuasiNewtonMatrixDirectionUpdate{Broyden}, p, o, x_old, ::Int)
    update_basis!(d.basis, p.M, x_old, o.x, d.vector_transport_method)
    yk_c = get_coordinates(p.M, o.x, o.yk, d.basis)
    sk_c = get_coordinates(p.M, o.x, o.sk, d.basis)
    skyk_c = inner(p.M, o.x, o.sk, o.yk)
    skHksk_c = sk_c' * d.matrix * sk_c

    φ = update_broyden_factor!(d, sk_c, yk_c, skyk_c, skHksk_c, d.update.update_rule)
    # computing the new matrix which represents the approximating operator in the next iteration
    d.matrix =
        d.matrix - (d.matrix * sk_c * sk_c' * d.matrix) / skHksk_c +
        (yk_c * yk_c') / skyk_c +
        φ * skHksk_c * (yk_c / skyk_c - (d.matrix * sk_c) / skHksk_c) *
        (yk_c / skyk_c - (d.matrix * sk_c) / skHksk_c)'
    return d
end

function update_broyden_factor!(d, sk_c, yk_c, skyk_c, skHksk_c, s::Symbol)
    return update_broyden_factor!(d, sk_c, yk_c, skyk_c, skHksk_c, Val(s))
end

function update_broyden_factor!(d, ::Any, ::Any, ::Any, ::Any, ::Val{:constant})
    return d.update.φ
end

function update_broyden_factor!(d, ::Any, yk_c, skyk_c, skHksk_c, ::Val{:Davidon})
    yk_c_c = d.matrix \ yk_c
    ykyk_c_c = yk_c' * yk_c_c
    if skyk_c <= 2 * (skHksk_c * ykyk_c_c) / (skHksk_c + ykyk_c_c)
        return d.update.φ =
            (skyk_c * (ykyk_c_c - skyk_c)) / (skHksk_c * ykyk_c_c - skyk_c^2)
    else
        return d.update.φ = skyk_c / (skyk_c - skHksk_c)
    end
end

function update_broyden_factor!(d, sk_c, ::Any, skyk_c, ykBkyk_c, ::Val{:InverseDavidon})
    sk_c_c = d.matrix \ sk_c
    sksk_c_c = sk_c' * sk_c_c
    if skyk_c <= 2 * (ykBkyk_c * sksk_c_c) / (ykBkyk_c + sksk_c_c)
        return d.update.φ =
            (skyk_c * (sksk_c_c - skyk_c)) / (ykBkyk_c * sksk_c_c - skyk_c^2)
    else
        return d.update.φ = skyk_c / (skyk_c - ykBkyk_c)
    end
end

function update_basis!(
    b::AbstractBasis, ::AbstractManifold, ::P, ::P, ::AbstractVectorTransportMethod
) where {P}
    return b
end

function update_basis!(
    b::CachedBasis, M::AbstractManifold, x::P, y::P, m::AbstractVectorTransportMethod
) where {P}
    # transport all basis tangent vectors in the tangent space of the next iterate
    for v in b.data
        vector_transport_to!(M, v, x, v, y, m)
    end
    return b
end

# Cautious update
function update_hessian!(
    d::QuasiNewtonCautiousDirectionUpdate{U}, p, o, x_old, iter
) where {U<:AbstractQuasiNewtonDirectionUpdate}
    # computing the bound used in the decission rule
    bound = d.θ(norm(p.M, o.x, o.gradient))
    sk_normsq = norm(p.M, o.x, o.sk)^2
    if sk_normsq != 0 && (inner(p.M, o.x, o.sk, o.yk) / sk_normsq) >= bound
        update_hessian!(d.update, p, o, x_old, iter)
    end
    return d
end

# Limited-memory update
function update_hessian!(
    d::QuasiNewtonLimitedMemoryDirectionUpdate{U}, p, o, x_old, ::Int
) where {U<:InverseBFGS}
    (capacity(d.memory_s) == 0) && return d
    # only transport the first if it does not get overwritten at the end
    start = length(d.memory_s) == capacity(d.memory_s) ? 2 : 1

    for i in start:length(d.memory_s)
        # transport all stored tangent vectors in the tangent space of the next iterate
        vector_transport_to!(
            p.M, d.memory_s[i], x_old, d.memory_s[i], o.x, d.vector_transport_method
        )
        vector_transport_to!(
            p.M, d.memory_y[i], x_old, d.memory_y[i], o.x, d.vector_transport_method
        )
    end

    # add newest
    push!(d.memory_s, o.sk)
    push!(d.memory_y, o.yk)
    return d
end

# all Cautious Limited Memory
function update_hessian!(
    d::QuasiNewtonCautiousDirectionUpdate{QuasiNewtonLimitedMemoryDirectionUpdate{NT,T,VT}},
    p,
    o,
    x_old,
    iter,
) where {NT<:AbstractQuasiNewtonUpdateRule,T,VT<:AbstractVectorTransportMethod}
    # computing the bound used in the decission rule
    bound = d.θ(norm(p.M, x_old, get_gradient(p, x_old)))
    sk_normsq = norm(p.M, o.x, o.sk)^2

    # if the decission rule is fulfilled, the new sk and yk are added
    if sk_normsq != 0 && (inner(p.M, o.x, o.sk, o.yk) / sk_normsq) >= bound
        update_hessian!(d.update, p, o, x_old, iter)
    else
        # the stored vectores are just transported to the new tangent space, sk and yk are not added
        for i in 1:length(d.update.memory_s)
            vector_transport_to!(
                p.M,
                d.update.memory_s[i],
                x_old,
                d.update.memory_s[i],
                o.x,
                d.update.vector_transport_method,
            )
            vector_transport_to!(
                p.M,
                d.update.memory_y[i],
                x_old,
                d.update.memory_y[i],
                o.x,
                d.update.vector_transport_method,
            )
        end
    end
    return d
end

get_solver_result(o::QuasiNewtonOptions) = o.x


# Chase Algorithm
#################

add_srctgt(x::Symbol) = Symbol("src_$(x)") => Symbol("tgt_$(x)")

"""Create a ACSet instance modeling a C-Rel from an instance of an ACSet"""
function crel(x::StructACSet{S}) where {S}
  name = Symbol("Rel_$(typeof(x).name.name)")
  os, hs, _, _, src, tgt = S.parameters
  pres = Presentation(FreeSchema)
  nv = length(os)
  alledge = Symbol[]
  xobs = [Ob(FreeSchema, sym) for sym in vcat(os...,hs...)]
  for x in xobs
    add_generator!(pres, x)
  end
  for (i,h) in enumerate(hs)
    s, t = add_srctgt(h)
    add_generator!(pres, Hom(s, xobs[nv+i], xobs[src[h]]))
    add_generator!(pres, Hom(t, xobs[nv+i], xobs[tgt[h]]))
    append!(alledge, [s,t])
  end
  expr = struct_acset(name, StructACSet, pres, index=alledge)
  eval(expr)
  rel = eval(name)
  J = Base.invokelatest(rel)

  # initialize the C-Rel with x which satisfies the laws
  for o in ob(S)
    add_parts!(J, o, nparts(x, o))
  end
  src, tgt = dom(S), codom(S)
  for (i, d) in enumerate(hom(S))
    hs, ht = add_srctgt(d)
    add_parts!(J, d, nparts(x, src[i]))
    set_subpart!(J, hs, parts(x, src[i]))
    set_subpart!(J, ht, x[d])
  end
  return J
end

"""Convert a C-Rel back to a C-Set, fail if relations are not functional"""
function uncrel(J::StructACSet, I::StructACSet{S}) where {S}
  res = typeof(I)()
  for o in ob(S)
    add_parts!(res, o, nparts(J, o))
  end
  for m in hom(S)
    msrc, mtgt = add_srctgt(m)
    md, mcd = zip(sort(collect(zip(J[msrc], J[mtgt])))...)
    collect(md) == collect(1:length(md)) || error("nonfunctional $J")
    set_subpart!(res, m, mcd)
  end
  return res
end

"""Evaluate a possibly composite composition on a subset of the domain"""
function eval_path(pth, J::StructACSet, xs::Set{Int})::Set{Int}
  pthkind = typeof(pth).parameters[1]  # Warning: uses implementation details
  if pthkind == :id
    return xs
  elseif pthkind == :generator
    args = [pth]
  elseif pthkind == :compose
    args = pth.args
  else
    error(pthkind)
  end
  # Iteratively update X to Y by searching through all (x,y) in R ⊆ X×Y
  for m in map(Symbol, args)
    msrc, mtgt = add_srctgt(m)
    if isempty(xs)
      return xs
    else
      newxs = Set{Int}()
      for (jx, jy) in zip(J[msrc], J[mtgt])
        if jx ∈ xs
          push!(newxs, jy)
        end
      end
      xs = newxs
    end
  end
  return xs
end

"""
Chase algorithm as described by Spivak and Wisnesky in "Fast Left-Kan Extensions
Using The Chase"

Todo: handle attributes
"""
function chase(F::Functor, I::StructACSet{IS}, codomI::StructACSet{S};
               n::Int=8, verbose::Bool=false)::StructACSet{S} where {IS, S}
  J = crel(codomI) # pre-model data structure, initialized with codomI
  η = Dict() # mapping into J from I

  # Further initialize J using the domain instance, I
  for c in generators(dom(F), :Ob)
    η[c] = add_parts!(J, Symbol(Ob(F)[c]), nparts(I, Symbol(c)))
  end
  changed = false # iterate until no change or n is reached
  for counter in 1:n
    if verbose
      println("----- ITERATION $counter ------")
    end

    # Initialize equivalence relations for each component
    eqclasses = Dict([o=>IntDisjointSets(nparts(J, o)) for o in ob(S)])
    changed = false

    # Action α: add new elements when a function is dangling
    fresh = Dict{Symbol, Set{Int}}([o=>Set{Int}() for o in ob(S)])
    for (d, ddom, dcodom) in zip(hom(S), dom(S), codom(S))
      dsrc, dtgt = add_srctgt(d)
      xs = setdiff(parts(J, ddom), J[dsrc] ∪ fresh[ddom])
      for x in xs  # no y such that (x,y) ∈ J(d) and no fresh x
        gx = add_part!(J, dcodom) # add a fresh symbol
        if verbose
          println("α: $d($ddom#$x) ADDED $dcodom $gx")
        end
        add_part!(J, d; Dict([dsrc=>x, dtgt=>gx])...)
        gx_ = push!(eqclasses[dcodom])
        push!(fresh[dcodom], gx)
        gx == gx_ || error("Unexpected mismatch: $eqclasses \n$J")
        changed = true
      end
    end

    # Action βd: add all coincidences induced by D (i.e. fire EGDs)
    for (p, q) in equations(codom(F))
      for x in parts(J, Symbol(dom(p)))
        for px in eval_path(p, J, Set([x]))
          for qx in eval_path(q, J, Set([x]))
            union!(eqclasses[Symbol(codom(p))], px, qx)
            if verbose && px != qx
              println("βF: $p=$q MERGED $(codom(p)) $px = $qx")
            end
            changed |= px != qx # flip iff changed = ⊥ AND px != qx
          end
        end
      end
    end

    # Action βF: add all coincidences induced by F. Naturality square constraint
    for mg in generators(dom(F), :Hom)
      m, Fm, Fmcodom = map(Symbol, [mg, Hom(F)[mg], Ob(F)[codom(mg)]])
      is_id = typeof(Hom(F)[mg]).parameters[1] == :id
      for (x, mx) in enumerate(I[m])
        mα = η[codom(mg)][mx]
        Fms, Fmt = add_srctgt(Fm)
        αx = η[dom(mg)][x]
        αm = is_id ? αx : J[Fmt][incident(J, αx, Fms)[1]]
        union!(eqclasses[Fmcodom], mα, αm)
        if verbose && mα != αm
          println("βF: $mg ($(dom(mg))#$x) MERGED $Fmcodom $mα = $αm")
        end
        changed |= mα != αm # flip iff changed = ⊥ AND mα != αm
      end
    end

    # Action δ: add all coincidences induced by functionality
    for (d, ddom, dcodom) in zip(hom(S), dom(S), codom(S))
      dsrc, dtgt = add_srctgt(d)
      for x in parts(J, ddom)
        eq_ys = collect(Set(J[dtgt][incident(J, x, dsrc)]))
        for (y1, y2) in zip(eq_ys, eq_ys[2:end])
          union!(eqclasses[dcodom], y1, y2)
          if verbose
            println("δ: $d($ddom#$x) MERGED $dcodom $y1 = $y2")
          end
          changed = true
        end
      end
    end

    # Action γ: merge coincidentally equal elements
    μ = Dict{Symbol, Vector{Pair{Int,Int}}}([o=>Pair{Int,Int}[] for o in ob(S)])
    delob = Dict{Symbol, Vector{Int}}([o=>Int[] for o in ob(S)])
    for (o, eq) in pairs(eqclasses)
      eqsets = Dict{Int,Set{Int}}()
      # Recover the equivalence classes from the union-find structure
      for x in parts(J, o)
        eqrep = find_root(eq, x)
        if haskey(eqsets, eqrep)
          push!(eqsets[eqrep], x)
        else
          eqsets[eqrep] = Set([x])
        end
      end
      if verbose
        println("eqset $o $(values(eqsets))")
      end
      # Minimum element is the representative
      for vs in values(eqsets)
        m = minimum(vs)
        delete!(vs, m)
        append!(μ[o], [v=>m for v in vs])
        append!(delob[o], collect(vs))
      end
    end
    # Replace all instances of a class with its representative in J and η
    for (d, ddom, dcodom) in zip(hom(S), dom(S), codom(S))
      dsrc, dtgt = add_srctgt(d)
      isempty(μ[ddom]) || set_subpart!(J, dsrc, replace(J[dsrc], μ[ddom]...))
      isempty(μ[dcodom]) || set_subpart!(J, dtgt, replace(J[dtgt], μ[dcodom]...))
      # remove redundant duplicate relation rows
      seen = Set()
      for (i, st) in enumerate(zip(J[dsrc], J[dtgt]))
        if st ∈ seen
          rem_part!(J, d, i)
        else
          push!(seen, st)
        end
      end
    end
    for (Iob, mapping) in collect(η)
      reps = μ[Symbol(Ob(F)[Iob])]
      if !isempty(reps)
        η[Iob] = replace(mapping, reps...)
      end
    end
    for (o, vs) in collect(delob)
      isempty(vs) || rem_parts!(J, o, sort(vs))
    end
    if !changed
      break # termination condition
    end
  end

  if changed
    error("Chase did not terminate with $n steps: $J")
  else
    return uncrel(J, codomI)
  end
end

# end # module


using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra
using GLM

function time_poisson_response(
    n :: Int64,        # number of cases
    p :: Int64,        # number of predictors (SNPs)
    k :: Int64,        # number of true predictors per group
    debias :: Bool
)
	d = Poisson
	l = canonicallink(d())
	Random.seed!(1111)

	#construct snpmatrix, covariate files, and true model b
	x = simulate_random_snparray(n, p, undef)
	xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
	z = ones(n, 1) # the intercept
	#define true_b
	true_b = zeros(p)
	true_b[1:10] .= collect(0.1:0.1:1.0)
	shuffle!(true_b)
	correct_position = findall(!iszero, true_b)

    #simulate phenotypes (e.g. vector y)
    if d == Normal || d == Poisson || d == Bernoulli
        prob = linkinv.(l, xbm * true_b)
        clamp!(prob, -20, 20)
        y = [rand(d(i)) for i in prob]
    elseif d == NegativeBinomial
        μ = linkinv.(l, xbm * true_b)
        prob = 1 ./ (1 .+ μ ./ nn)
        y = [rand(d(nn, i)) for i in prob] #number of failtures before nn success occurs
    elseif d == Gamma
        μ = linkinv.(l, xbm * true_b)
        β = 1 ./ μ # here β is the rate parameter for gamma distribution
        y = [rand(d(α, i)) for i in β] # α is the shape parameter for gamma
    end
    y = Float64.(y)

	#time the result and return
	result = @timed L0_reg(x, xbm, z, y, 1, k, d(), l, debias=debias, init=false, use_maf=false)

	iter = result[1].iter
   	runtime = result[2]      # seconds
   	memory = result[3] / 1e6 # MB

    return iter, runtime, memory
end

function main(n :: Int64, i :: Int64)
	p = 1000000
	k = 10
	debias = true
	response = "poisson"

	#run result: first run on super small data to force compile, then run result
	iter, runtime, memory = time_poisson_response(100, 100, k, debias)
	iter, runtime, memory = time_poisson_response(n, p, k, debias)

	#write result to corresponding file
	if debias
		dir = response * "_results/"
	else
		dir = response * "_results_nodebias/"
	end
	file = string(n) * "_by_" * string(p) * "_run$i"
	open(dir * file, "w") do f
		write(f, "time (seconds), memory (MB), iteration\n")
		write(f, string(runtime) * ", " * string(memory) * ", " * string(iter))
	end
end

# arguments passed from command line
n = parse(Int64, ARGS[1])
i = parse(Int64, ARGS[2]) #i'th run: time result 5 times so we can average (didn't use benchmarktools cuz takes too long)

main(n, i)


function log_network(
    state::Tuple{AbstractGraph, AbstractArray}, tick_nr::Int64
)
    graph, agent_list = state
    agent_opinion = [a.opinion for a in agent_list]
    agent_perceiv_publ_opinion = [a.perceiv_publ_opinion for a in agent_list]
    agent_indegree = indegree(graph)
    agent_outdegree = outdegree(graph)
    agent_centrality = closeness_centrality(graph)
    agent_cc = local_clustering_coefficient(graph)
    agent_component = (
        [i for i in keys(connected_components(graph)), v in 1:nv(graph)
        if v in connected_components(graph)[i]]
    )

    return DataFrame(
        TickNr = tick_nr,
        AgentID = 1:length(agent_list),
        Opinion = agent_opinion,
        PerceivPublOpinion = agent_perceiv_publ_opinion,
        Indegree = agent_indegree,
        Outdegree = agent_outdegree,
        Centrality = agent_centrality,
        CC = agent_cc,
        Component = agent_component
    )
end

# suppress output of include()
;


"""
```
gensysct(Γ0, Γ1, c, Ψ, Π; ϵ, div)
gensysct(F::LinearAlgebra.GeneralizedSchur, c, Ψ, Π; ϵ, div)
```

Generate state-space solution to canonical-form DSGE model.

System given as
```
Γ0*Dy(t) = Γ1*y(t) + c + Ψ*z(t) + Π*η(t),
```
with z an exogenous variable process and η being endogenously white noise expectational errors.

Returned system is
```
Dy(t) = G1*y(t) + C + impact * z(t)
```

Returned values are
```
G1, C, impact, qt', a, b, z, eu
```

This code is based on a [routine](http://sims.princeton.edu/yftp/gensys/) originally copyright Chris Sims.

Here for completeness, but should not be used for most HACT/HANK models.

Also returned is the qz decomposition, qt'az' = Γ0, qt'bz' = Γ1, with a and b
upper triangular and the system ordered so that all zeros on the diagonal of b are in
the lower right corner, all cases where the real part of bii/aii is greater than or
equal to div appear in the next block above the zeros, and the remaining bii/aii's
all have bii/aii<div .  These elements can be used to construct the full backward and
forward solution.  See the paper \"Solving Linear Rational Expectations Models\",
http://eco-072399b.princeton.edu/yftp/gensys .  Note that if one simply wants the backwar
and forward projection of y on eΨlon, ignoring existence and uniqueness questions, the
projection can be computed by Fourier methods.

If `div` is omitted from argument list, a `div`>1 is calculated.

### Return codes

* `eu[1]==1` for existence
* `eu[2]==1` for uniqueness
* `eu[1]==-1` for existence for white noise η
* `eu==[-2,-2]` for coincident zeros.
"""
function gensysct(Γ0::Matrix{Float64}, Γ1::Matrix{Float64}, c::Array{Float64},
                  Ψ::Matrix{Float64}, Π::Matrix{Float64};
                  check_existence::Bool = true, check_uniqueness::Bool = true,
                  ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0, complex_decomposition::Bool = false)

    F = complex_decomposition ? schurfact(complex(Γ0), complex(Γ1)) : schurfact(real(Γ0), real(Γ1))
    gensysct(F, c, Ψ, Π; check_existence = check_existence, check_uniqueness = check_uniqueness,
             ϵ = ϵ, div = div)
end

# In place version. Changes Γ0, Γ1
function gensysct!(Γ0::Matrix{Float64}, Γ1::Matrix{Float64}, c::Array{Float64},
                   Ψ::Matrix{Float64}, Π::Matrix{Float64};
                   check_existence::Bool = true, check_uniqueness::Bool = true,
                   ϵ::Float64 = sqrt(eps()) * 10,
                   div::Float64 = -1.0, complex_decomposition::Bool = false)

    F = complex_decomposition ? schur!(complex(Γ0), complex(Γ1)) : schur!(real(Γ0), real(Γ1))
    gensysct(F, c, Ψ, Π; check_existence = check_existence, check_uniqueness = check_uniqueness,
             ϵ = ϵ, div = div)
end

# Uses basic schur decomposition
function gensysct(Γ1::Matrix{Float64}, c::Array{Float64}, Ψ::Matrix{Float64}, Π::Matrix{Float64};
                  check_existence::Bool = true, check_uniqueness::Bool = true,
                  ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0, complex_decomposition::Bool = false)
    # Assumes that Γ0 is the identity
    F = complex_decomposition ? schurfact(complex(Γ1)) : schurfact(real(Γ1))
    gensysct(F, c, Ψ, Π; check_existence = check_existence, check_uniqueness = check_uniqueness,
             ϵ = ϵ, div = div)
end

# In place version of above gensysct
function gensysct!(Γ1::Matrix{Float64}, c::Array{Float64}, Ψ::Matrix{Float64}, Π::Matrix{Float64};
                  check_existence::Bool = true, check_uniqueness::Bool = true, ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0, complex_decomposition::Bool = false)
    # Assumes that Γ0 is the identity
    F = complex_decomposition ? schur!(complex(Γ1)) : schur!(real(Γ1))
    gensysct(F, c, Ψ, Π; check_existence = check_existence, check_uniqueness = check_uniqueness,
             ϵ = ϵ, div = div)
end

function gensysct(F::LinearAlgebra.GeneralizedSchur, c::Array{Float64}, Ψ::Matrix{Float64}, Π::Matrix{Float64};
                  check_existence::Bool = true, check_uniqueness::Bool = true, ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0)
    div < 0.0 && (div = new_divct(F; ϵ = ϵ))
    eu = [0, 0]
    a, b = F.S, F.T
    n = size(a, 1)

    for i in 1:n
        if (abs(a[i, i]) < ϵ) && (abs(b[i, i]) < ϵ)
            info("Coincident zeros.  Indeterminacy and/or nonexistence.")
            eu = [-2, -2]
            G1 = Array{Float64, 2}() ;  C = Array{Float64, 1}() ; impact = Array{Float64, 2}()
            a, b, qt, z = FS.S, FS.T, FS.Q, FS.Z
            return G1, C, impact, qt', a, b, z, eu
        end
    end
    movelast = fill(false, n)
    for i in 1:n
        movelast[i] = (real(b[i, i] / a[i, i]) > div) || (abs(a[i, i]) < ϵ)
    end
    nunstab = sum(movelast)
    FS = ordschur!(F, .!movelast)
    a, b, qt, z = FS.S, FS.T, FS.Q, FS.Z

    qt1 = qt[:, 1:(n - nunstab)]
    qt2 = qt[:, (n - nunstab + 1):n]
    a2 = a[(n - nunstab + 1):n, (n - nunstab + 1):n]
    b2 = b[(n - nunstab + 1):n, (n - nunstab + 1):n]
    etawt = qt2' * Π # Ac_mul_B(qt2, Π)
    ~, ueta, deta, veta = decomposition_svdct!(etawt, ϵ = ϵ)
    zwt = qt2' * Ψ # Ac_mul_B(qt2, Ψ)
    bigev, uz, dz, vz = decomposition_svdct!(zwt, ϵ = ϵ)

    # check existence
    if isempty(bigev) || !check_existence
        exist = true
        existx = true
    else
        exist = norm(uz- (ueta * adjoint(ueta)) * uz, 2) < ϵ * n

        zwtx0 = b2 \ zwt
        zwtx = zwtx0
        M = b2 \ a2
        M = scale!(M, 1 / norm(M))
        for i in 2:nunstab
            zwtx = hcat(M * zwtx, zwtx0)
        end
        zwtx = b2 * zwtx
        bigev, ux, dx, vx = decomposition_svdct!(zwtx, ϵ = ϵ)
        existx = norm(ux - ueta * ueta' * ux, 2) < ϵ * n
    end

    etawt1 = qt1' * Π # Ac_mul_B(qt1, Π)
    bigev, ueta1, deta1, veta1 = decomposition_svdct!(etawt1, ϵ = ϵ)

    if nunstab == 0
       eu[1] = 1
    else
        if exist
            eu[1] = -1
        end
    end

    # check uniqueness
    if isempty(veta1) || !check_uniqueness
        eu[2] = 1
    else
        eu[2] = Int64(norm(veta1- (veta * adjoint(veta)) * veta1, 2) < ϵ * n)
    end

    tmat = hcat(eye(n - nunstab), -ueta1 * deta1 * conj(veta1)' * veta * (deta \ ueta'))
    G0 =  vcat(tmat * a, hcat(zeros(nunstab, n - nunstab), eye(nunstab)))
    G1 =  vcat(tmat * b, zeros(nunstab, n))

    G0I = inv(G0)
    G1 = G0I * G1
    usix = (n - nunstab + 1):n
    Busix = b[usix,usix]
    Ausix = a[usix,usix]
    C = G0I * vcat(tmat * qt' * c, (Ausix - Busix) \ (qt2' * c))
    impact = G0I * vcat(tmat * qt' * Ψ, zeros(nunstab, size(Ψ, 2)))
    # C = G0I * vcat(tmat * Ac_mul_B(qt, c), (Ausix - Busix) \ Ac_mul_B(qt2, c))
    # impact = G0I * vcat(tmat * Ac_mul_B(qt, Ψ), zeros(nunstab, size(Ψ, 2)))

    #G1 = z * A_mul_Bc(G1, z)
    G1 = z * G1 * z'
    G1 = real(G1)
    C = real(z * C) .* ones(size(z, 1), 1)
    impact = real(z * impact)

    return G1::Matrix{Float64}, C::Matrix{Float64}, impact::Matrix{Float64}, qt'::Matrix, a::Matrix, b::Matrix, z::Matrix, eu::Vector{Int64}
end

# Port of SeHyoun's schur_solver
## Inputs
# div: number of unstable roots known a priori
## Outputs
# eu[1] = 1 solution exists
#         0 no stable solution
#         -1 solution exists for white noise η
# eu[2] = 1 solution is unique
#         0 solution is not unique
function gensysct(F::LinearAlgebra.Schur, c::Array{Float64}, Ψ::Matrix{Float64}, Π::Matrix{Float64};
                  check_existence::Bool = true, check_uniqueness::Bool = true, ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0)
    eu = [0, 0]
    U, T = F.Z, F.T        # U is unitary matrix, T is Schur matrix
    n = size(U, 1)
    g_eigs = real(F.values) # Get real eigenvalues
    stable_eigs = g_eigs .<= 0 # stable eigenvalues
    nunstab = n - sum(stable_eigs)  # Unstable eigenvalues
    if div > -1
        sort_order = sortperm(g_eigs, rev = true) # g_eigs[sort_order] -> descending sorted g_eigs
        locs = ones(n, 1)
        locs[sort_order[1:div]] = 0 # Set to zero the eigenvalues which are unstable
        FS = ordschur!(F, locs)     # Change F in place, FS creates reference to same object, so no copying
        if nunstab > div
            @warn "<gensysct>: There are more than div number of positive eigenvalues with smallest values:\n"
            println(g_eigs[sort_order][div + 1:nunstab])
        elseif nunstab < div
            @warn "<gensysct>: There are less than div number of positive eigenvalues:\n"
            println(g_eigs[sort_order][nunstab + 1:div])
        end
        nunstab = div
    else
        FS = ordschur!(F, stable_eigs)
        U, T = FS.Z, FS.T
    end

    U1 = U[:, 1:(n - nunstab)]'
    U2 = U[:, (n - nunstab + 1):n]'
    etawt = U2 * Π # Ac_mul_B(U2, Π)
    ~, ueta, deta, veta = decomposition_svdct!(etawt, ϵ = ϵ) # Wrapper function here, see end of script, this checks for bigev and returns ueta, etc., after selecting only the relevant columns

    # check existence
    if check_existence
        zwt = U2 * Ψ   # Ac_mul_B(U2, Ψ)
        bigev, uz, dz, vz = decomposition_svdct!(zwt, ϵ = ϵ)
        if isempty(bigev)
            eu[1] = 1
        else
            eu[1] = norm(uz- (ueta * adjoint(ueta)) * uz, 2) < ϵ * n
        end

        if (eu[1] == 0 && (div == -1))
            @warn "<gensysct>: Solution does not exist"
        end
        impact = real(-Π * veta * (deta \ ueta') * uz * dz * vz' + Ψ)
    else
        eu[1] = 1
        impact = real(-Π * veta * (deta \ ueta') * U2 + Ψ)
    end

    # Check uniqueness
    if check_uniqueness
        etawt1 = U1 * Π # Ac_mul_B(U1, Π)
        ~, ~, deta1, veta1 = decomposition_svdct!(etawt1)
        if isempty(veta1)
            eu[2] = 1
        else
            eu[2] = Int64(norm(veta1 - (veta * adjoint(veta)) * veta1, 2) < ϵ * n)
        end
    end
    I, J, V = SparseArrays.spdiagm_internal(0 => [ones(n - nunstab); zeros(nunstab)])
    diag_m = sparse(I, J, V, n, n)
    G1 = real(U * T * diag_m * U')#spdiagm([ones(n - nunstab); zeros(nunstab)], 0, n, n) * U')
    F = U1[:, 1:nunstab]' * inv(U1[:, nunstab + 1:end]')
    impact = [F * Ψ[nunstab + 1:end, :]; Ψ[nunstab + 1:end, :]]
    C = real(U * c) .* ones(size(U, 1), 1)

    return G1::Matrix{Float64}, C::Matrix{Float64}, impact::Matrix{Float64}, U::Matrix, T::Matrix::Matrix, eu::Vector{Int64}
end

function new_divct(F::LinearAlgebra.GeneralizedSchur; ϵ::Float64 = sqrt(eps()) * 10)
    a, b = F.S, F.T
    n = size(a, 1)
    div = 0.001
    for i in 1:n
        if abs(a[i, i]) > ϵ
            divhat = real(b[i, i] / a[i, i])
            if (ϵ < divhat) && (divhat < div)
                div = 0.5 * divhat
            end
        end
    end
    return div
end

function decomposition_svdct!(A; ϵ::Float64 = sqrt(eps()) * 10)
    Asvd = svd!(A)
    bigev = findall(Asvd.S .> ϵ)
    Au = Asvd.U[:, bigev]
    Ad = diagm(0 => Asvd.S[bigev])
    Av = Asvd.V[:, bigev]
    return bigev, Au, Ad, Av
end

# Additional cases when inputs are sparse matrices
function gensysct(Γ0::SparseMatrixCSC{Float64,Int64}, Γ1::SparseMatrixCSC{Float64,Int64},
                  c::SparseMatrixCSC{Float64,Int64}, Ψ::SparseMatrixCSC{Float64,Int64},
                  Π::SparseMatrixCSC{Float64,Int64};
                  check_existence::Bool = true, check_uniqueness::Bool = true, ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0, complex_decomposition::Bool = false)
    F = complex_decomposition ? schurfact(complex(Matrix{Float64}(Γ0)), complex(Matrix{Float64}(Γ1))) : schurfact(real(Matrix{Float64}(Γ0)), real(Matrix{Float64}(Γ1)))
    gensysct(F, Matrix{Float64}(c), Matrix{Float64}(Ψ), Matrix{Float64}(Π); check_existence = check_existence,
             check_uniqueness = check_uniqueness, ϵ = ϵ, div = div)
end

function gensysct!(Γ0::SparseMatrixCSC{Float64,Int64}, Γ1::SparseMatrixCSC{Float64,Int64},
                   c::SparseMatrixCSC{Float64,Int64}, Ψ::SparseMatrixCSC{Float64,Int64},
                   Π::SparseMatrixCSC{Float64,Int64};
                  check_existence::Bool = true, check_uniqueness::Bool = true, ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0, complex_decomposition::Bool = false)
    F = complex_decomposition ? schur!(complex(Matrix{Float64}(Γ0)), complex(Matrix{Float64}(Γ1))) : schur!(real(Matrix{Float64}(Γ0)), real(Matrix{Float64}(Γ1)))
    gensysct(F, Matrix{Float64}(c), Matrix{Float64}(Ψ), Matrix{Float64}(Π); check_existence = check_existence,
             check_uniqueness = check_uniqueness, ϵ = ϵ, div = div)
end

function gensysct(Γ1::SparseMatrixCSC{Float64,Int64}, c::SparseMatrixCSC{Float64,Int64},
                  Ψ::SparseMatrixCSC{Float64,Int64}, Π::SparseMatrixCSC{Float64,Int64};
                  check_existence::Bool = true, check_uniqueness::Bool = true, ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0, complex_decomposition::Bool = false)
    # Assumes that Γ0 is the identity
    F = complex_decomposition ? schurfact(complex(Matrix{Float64}(Γ1))) : schurfact(real(Matrix{Float64}(Γ1)))
    gensysct(F, Matrix{Float64}(c), Matrix{Float64}(Ψ), Matrix{Float64}(Π); check_existence = check_existence,
             check_uniqueness = check_uniqueness, ϵ = ϵ, div = div)
end

function gensysct!(Γ1::SparseMatrixCSC{Float64,Int64}, c::SparseMatrixCSC{Float64,Int64},
                   Ψ::SparseMatrixCSC{Float64,Int64}, Π::SparseMatrixCSC{Float64,Int64};
                  check_existence::Bool = true, check_uniqueness::Bool = true, ϵ::Float64 = sqrt(eps()) * 10,
                  div::Float64 = -1.0, complex_decomposition::Bool = false)
    # Assumes that Γ0 is the identity
    F = complex_decomposition ? schur!(complex(Matrix{Float64}(Γ1))) : schur!(real(Matrix{Float64}(Γ1)))
    gensysct(F, Matrix{Float64}(c), Matrix{Float64}(Ψ), Matrix{Float64}(Π); check_existence = check_existence,
             check_uniqueness = check_uniqueness, ϵ = ϵ, div = div)
end


# unit test the set up

include("../src/DMD_Util.jl");
include("../src/DMD_Type.jl");

# Generate DMD Synthetic Data
#------------------------------------------------------------------------------
# dimensions
m = 100;    # temporal dimension
n = 100;    # spatial dimension
k = 3;      # number of modes
T = Float64;
# generate data
X, t, at, Bt = simDMD(m, n, k, T; seed=123456);
# specify the loss
function lossFunc(r)
	T = typeof(real(r[0]));
	return T(0.5)*sum(abs2, r)
end
#
function lossGrad(g, r)
	copyto!(g, r);
end
#
# Create object
#------------------------------------------------------------------------------
params = DMDParams(k, X, t, lossFunc, lossGrad);
#
println("create DMD object: OK")


__precompile__()

module FastGaussQuadrature

using Compat, SpecialFunctions

if VERSION ≥ v"0.7-"
    using LinearAlgebra
    cumprod(A::AbstractArray) = Base.cumprod(A, dims=1)
    cumprod(A::AbstractArray, d::Int) = Base.cumprod(A, dims=d)
    sum(A::AbstractArray, n::Int) = Base.sum(A, dims=n)
    sum(A) = Base.sum(A)
    flipdim(A, d) = reverse(A, dims=d)
else
    rmul!(A::AbstractArray, x::Number) = scale!(A, x)
    repeat(A, n, m) = repmat(A, n, m)
    repeat(A...; kwds...) = Base.repeat(A...; kwds...)
    eigen(A) = eig(A)
    const floatmax = realmax
end

export gausslegendre
export gausschebyshev
export gausslaguerre
export gausshermite
export gaussjacobi
export gausslobatto
export gaussradau
export besselroots

import SpecialFunctions: besselj, airyai, airyaiprime

include("gausslegendre.jl")
include("gausschebyshev.jl")
include("gausslaguerre.jl")
include("gausshermite.jl")
include("gaussjacobi.jl")
include("gausslobatto.jl")
include("gaussradau.jl")
include("besselroots.jl")

end


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

using PyCall
using JLD, HDF5

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

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

"This function maps i,j on an Nx by Ny grid to an Nx*Ny 1-D array"
function xy_to_vect(i,j,Nx,Ny)
l = Nx*(j-1) + i
l
end
function nm(f, fp, x; tol=sqrt(eps()))

    ctr, max_steps = 0, 100

    while (abs(f(x)) > tol) && ctr < max_steps
        x = x - f(x) / fp(x)
        ctr = ctr + 1
    end

    ctr >= max_steps ? error("Method did not converge.") : return (x, ctr)

end

"This function computes the integral of the planck function from E1 = h*nu [keV] to infinity at T[keV]"
function B1(E1,T)
    r = (E1./T)

    num_terms = convert(Int32,round(minimum([2+20/minimum(r),1024])))
    output = 0*T
    for n in 1:num_terms
        invn = 1.0/n
        output += invn*exp.(-n*r).*(r.^3 + 3*r.^2*invn .+ 6*r * invn^2 .+ 6*invn^3)
    end
    return output*Planck_int_const
end

"This function gives Bg(T)"
function Bg(E1,E2,T)
    return B1(E1,T) - B2(E2,T)
end

"This function computes the derivative of B1 w.r.t. T"
function DB1DT(y,T)
    r = y./T
    num_terms = convert(Int32,round(minimum([2+20/minimum(r),1024])))
    output = 0*T
    for n in 1:num_terms
        output += exp.(-n*r).*(24*T.^4 .+ 24*n*T.^3 .*y + 12*n^2*T.^2 .*y.^2 + 4*n^3*T.*y.^3 .+ n^4*y.^4)./(n^4*T)
    end
    return output*Planck_int_const
end
"This function computes the derivative of Bg w.r.t. T"
function DBgDT(E1,E2,T)
    return DB1DT(E1,T) .- DB1DT(E2,T)
end
"""
    create_A(D,sigma,Nr,Nz,Lr,Lz; lower_z = "refl",upper_z="refl",upper_r="refl")

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

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

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

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

"This function solves equation 19 from Brunner's paper. It gives E_star_g"
function Estar_g(delta_t,T,Eg,Eg_l,D_func,sigma_func,Q_func,RHS,Nr,Nz,Lr,Lz;
    lower_z = "refl",upper_z="refl",upper_r="refl", LOUD=0)

    sigma_star = c*sigma_func(T) + 1.0/(dt)
    dz = Lz/Nz
    dr = Lr/Nr
    A,b = create_A(c*Dleft,c*Dright, c*Dtop,c*Dbottom, sigma, Nr, Nz, Lr, Lz, lower_z=lower_z,upper_z=upper_z, upper_r=upper_r);
    b += RHS
    return A\b
end #function

"One step of Equation 17"
function Estep(delta_t,T,T_prev,E,E_prev,E_l,D_func_arr,sigma_func_arr,Q_func_arr,group_bounds,EOS,Inv_EOS,Cv,rho,Nr,Nz,Lr,Lz,G; lower_z = "refl",upper_z="refl",upper_r="refl", LOUD=0)
    idt = 1.0/delta_t
    Jee = zeros(Nr*Nz)
    lagged_Emission = 0*Eg
    for g in 1:G
        Jee += sigma_func_arr[g](T).*DBgDT(group_bounds[g],group_bounds[g+1],T)
        lagged_Emission += reshape(sigma_func_arr[g](T),Nr*Nz).*(reshape(Bg(group_bounds[g],group_bounds[g+1],T),Nr*Nz) - E[:,g])*c
    end
    Jee *= c./reshape(Cv(T))
    Jee +=  reshape(rho,Nr,Nz)*idt

    Esteps = zeros(Nr*Nz,G)
    sigma_input = sigma_func_arr[g](T)
    e = EOS(T)
    e_prev = EOS(T_prev)
    for g in 1:G
        RHS = -c./reshape(Cv(T)).*reshape(sigma_func_arr[g](T).*DBgDT(group_bounds[g],group_bounds[g+1],T),Nr,Nz)./Jee
        RHS *= (reshape(rho,Nr,Nz)*idt*(e-e_prev) + reshape(lagged_Emission,Nr,Nz))
        RHS +=  reshape(sigma_func_arr[g](T),Nr,Nz).*reshape(E[:,g],Nr,Nz)*c #this group emission (2nd term on RHS)
        RHS +=  idt*reshape(E[:,g],Nr,Nz) #first term on RHS of 17
        RHS +=  reshape(Q_func(T),Nr,Nz) #source term on RHS of 17
        sigma_input *= (1-c./reshape(Cv(T)).*reshape(sigma_func_arr[g](T).*
                DBgDT(group_bounds[g],group_bounds[g+1],T),Nr,Nz)./Jee)
        sigma_input += idt

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

        A,b = create_A(c*Dleft,c*Dright, c*Dtop,c*Dbottom, sigma_input, Nr, Nz, Lr, Lz, lower_z=lower_z,upper_z=upper_z, upper_r=upper_r);
        b += RHS

        Esteps[:,g] = A\b

    end

    return Esteps
end #function


export compute_fft, compute_cc

@deprecate compute_fft() rfft()
@deprecate compute_cc() correlate()

"""

  compute_fft(R)

Computes windowed rfft of ambient noise data. Returns FFTData structure.

# Arguments
- `R::RawData`: RawData structure
"""
function compute_fft(R::RawData)
    FFT = rfft(R.x,1)
    return FFTData(R.name, R.id,R.loc, R.fs, R.gain, R.freqmin, R.freqmax,
                   R.cc_len, R.cc_step, R.whitened, R.time_norm, R.resp,
                   R.misc, R.notes, R.t, FFT)
end

"""
    compute_cc(FFT1::FFTData, FFT2::FFTData, maxlag::Float64;
               corr_type::String="cross-correlation")

Cross-correlate ambient noise data in the frequency domain.

Cross-correlation can be done using one of three options:

- Cross-correlation: ``C_{AB}(ω) = u_A(ω) u^∗_B(ω)``
- Coherence: ``C_{AB}(ω) = \frac{u_A(ω) u^∗_B(ω)}{∣ u_A(ω) ∣  ∣ u_B(ω) ∣}``
- Deconvolution: ``C_{AB}(ω) = \frac{u_A(ω) u^∗_B(ω)}{∣ u_B(ω) ∣^2}``

# Arguments
- `FFT1::FFTData`: FFTData object of fft'd ambient noise data.
- `FFT2::FFTData`: FFTData object of fft'd ambient noise data.
- `maxlag::Float64`: Maximum lag time (in seconds) in cross-correlation to save,
                     e.g. `maxlag = 20.` will save lag times = -20.:20. s.
- `corr_type::String`: Type of correlation: `cross-correlation`, `coherence` or
                       `deconv`.
"""
function compute_cc(FFT1::FFTData, FFT2::FFTData, maxlag::Float64;
                    corr_type::String="CC")

    comp = FFT1.name[end] * FFT2.name[end]
    # get intersect of dates; return nothing if no intersect
    inter = intersect(FFT1.t,FFT2.t)
    if length(inter) == 0
        throw(ArgumentError("No common windows for $(FFT1.name)-$(FFT2.name) $(FFT1.id)"))
    end

    N = convert(Int,round(FFT1.cc_len * FFT1.fs)) # number of data points
    ind1 = findall(x -> x ∈ inter, FFT1.t)
    ind2 = findall(x -> x ∈ inter, FFT2.t)
    corr = correlate(@views(FFT1.fft[:,ind1]), @views(FFT2.fft[:,ind2]), N,
                     convert(Int,round(maxlag * FFT1.fs)))
    rotated = false

    return CorrData(FFT1, FFT2, comp, rotated, corr_type,
                    maxlag, inter, corr)

end


module Sqlite3Stats

import Distributions
import StatsBase
import SQLite

export StatsBase
export SQLite
export Distributions

function linear_regression(x::Array{Float64, 1}, y::Array{Float64, 1})::Array{Float64, 1}
    meanx = StatsBase.mean(x)
    meany = StatsBase.mean(y)
    beta1 = sum((x .- meanx) .* (y .- meany)) / sum( (x .- meanx) .* (x .- meanx))
    beta0 = meany - beta1 * meanx
    return [beta0, beta1]
end

function register_functions(db::SQLite.DB; verbose::Bool = true)::Nothing

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.quantile(x, 0.25),
        name = "Q1")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.quantile(x, 0.50),
        name = "Q2")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.quantile(x, 0.50),
        name = "MEDIAN")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.quantile(x, 0.75),
        name = "Q3")

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.quantile(x[:,1], x[1,2]),
        name = "QUANTILE", nargs = 2)

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.cov(x[:,1], x[:,2]),
        name = "COV", nargs = 2)

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.cor(x[:,1], x[:,2]),
        name = "COR", nargs = 2)

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.corspearman(x[:,1], x[:,2]),
        name = "CORSPEARMAN", nargs = 2)

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.corkendall(x[:,1], x[:,2]),
        name = "CORKENDALL", nargs = 2)

    # Maximum absolute deviations
    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.maxad(x[:,1], x[:,2]),
        name = "MAXAD", nargs = 2)

    # Mean absolute deviations
    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.meanad(x[:,1], x[:,2]),
        name = "MEANAD", nargs = 2)

    # Mean squared deviations
    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.msd(x[:,1], x[:,2]),
        name = "MSD", nargs = 2)

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.mad(convert(Array{Float64, 1}, x)),
        name = "MAD")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.iqr(convert(Array{Float64, 1}, x)),
        name = "IQR")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.skewness(convert(Array{Float64, 1}, x)),
        name = "SKEWNESS")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.kurtosis(convert(Array{Float64, 1}, x)),
        name = "KURTOSIS")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.geomean(convert(Array{Float64, 1}, x)),
        name = "GEOMEAN")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.harmmean(convert(Array{Float64, 1}, x)),
        name = "HARMMEAN")

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.mode(convert(Array{Float64, 1}, x)),
        name = "MODE")

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.mean(x[:,1], StatsBase.weights(x[:,2])),
        name = "WMEAN", nargs = 2)

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> StatsBase.median(x[:,1], StatsBase.weights(x[:,2])),
        name = "WMEDIAN", nargs = 2)

    SQLite.register(db, [],
        (x,y) -> vcat(x, y),
        x -> StatsBase.entropy(convert(Array{Float64, 1}, x)),
        name = "ENTROPY")

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> linear_regression(x[:,1], x[:,2])[2],
        name = "LINSLOPE", nargs = 2)

    SQLite.register(db, Array{Float64, 2}(undef, (0, 2)),
        (x, a, b) -> vcat(x, [a, b]'),
        x -> linear_regression(x[:,1], x[:,2])[1],
        name = "LININTERCEPT", nargs = 2)

    # qnorm, pnorm, rnorm
    SQLite.register(db, (x, mu, sd) -> Distributions.quantile(Distributions.Normal(mu, sd), x), name = "QNORM")

    SQLite.register(db, (x, mu, sd) -> Distributions.pdf(Distributions.Normal(mu, sd), x), name = "PNORM")

    SQLite.register(db, (mu, sd) -> rand(Distributions.Normal(mu, sd)), name = "RNORM")

    # qt, pt, rt
    SQLite.register(db, (x, dof) -> Distributions.quantile(Distributions.TDist(dof), x), name = "QT")

    SQLite.register(db, (x, dof) -> Distributions.cdf(Distributions.TDist(dof), x), name = "PT")

    SQLite.register(db, (dof) -> rand(Distributions.TDist(dof)), name = "RT")

    # qchisq, pchisq, rchisq
    SQLite.register(db, (x, dof) -> Distributions.quantile(Distributions.Chisq(dof), x), name = "QCHISQ")

    SQLite.register(db, (x, dof) -> Distributions.cdf(Distributions.Chisq(dof), x), name = "PCHISQ")

    SQLite.register(db, (dof) -> rand(Distributions.Chisq(dof)), name = "RCHISQ")

    # qf, pf, rf
    SQLite.register(db, (x, dof1, dof2) -> Distributions.quantile(Distributions.FDist(dof1, dof2), x), name = "QF")

    SQLite.register(db, (x, dof1, dof2) -> Distributions.cdf(Distributions.FDist(dof1, dof2), x), name = "PF")

    SQLite.register(db, (dof1, dof2) -> rand(Distributions.FDist(dof1, dof2)), name = "RF")

    # qpois, ppois, rpois
    SQLite.register(db, (x, lambda) -> Distributions.quantile(Distributions.Poisson(lambda), x), name = "QPOIS")

    SQLite.register(db, (x, lambda) -> Distributions.cdf(Distributions.Poisson(lambda), x), name = "PPOIS")

    SQLite.register(db, (lambda) -> rand(Distributions.Poisson(lambda)), name = "RPOIS")

    # qbinom, pbinom, rbinom
    SQLite.register(db, (x, n, p) -> Distributions.quantile(Distributions.Binomial(n, p), x), name = "QBINOM")

    SQLite.register(db, (x, n, p) -> Distributions.cdf(Distributions.Binomial(n, p), x), name = "PBINOM")

    SQLite.register(db, (n, p) -> rand(Distributions.Binomial(n, p)), name = "RBINOM")

    # qrunif, prunif, runif
    SQLite.register(db, (x, a, b) -> Distributions.quantile(Distributions.Uniform(a, b), x), name = "QUNIF")

    SQLite.register(db, (x, a, b) -> Distributions.cdf(Distributions.Uniform(a, b), x), name = "PUNIF")

    SQLite.register(db, (a, b) -> rand(Distributions.Uniform(a, b)), name = "RUNIF")

    # qexp, pexp, rexp
    SQLite.register(db, (x, theta) -> Distributions.quantile(Distributions.Exponential(theta), x), name = "QEXP")

    SQLite.register(db, (x, theta) -> Distributions.cdf(Distributions.Exponential(theta), x), name = "PEXP")

    SQLite.register(db, (theta) -> rand(Distributions.Exponential(theta)), name = "REXP")

    # qbeta, pbeta, rbeta
    SQLite.register(db, (x, alpha, beta) -> Distributions.quantile(Distributions.Beta(alpha, beta), x), name = "QBETA")

    SQLite.register(db, (x, alpha, beta) -> Distributions.cdf(Distributions.Beta(alpha, beta), x), name = "PBETA")

    SQLite.register(db, (alpha, beta) -> rand(Distributions.Beta(alpha, beta)), name = "RBETA")

    # qcauchy, pcauchy, rcauchy
    SQLite.register(db, (x, mu, sigma) -> Distributions.quantile(Distributions.Cauchy(mu, sigma), x), name = "QCAUCHY")

    SQLite.register(db, (x, mu, sigma) -> Distributions.cdf(Distributions.Cauchy(mu, sigma), x), name = "PCAUCHY")

    SQLite.register(db, (mu, sigma) -> rand(Distributions.Cauchy(mu, sigma)), name = "RCAUCHY")

    return nothing
end

export register_functions

end # module


module DO

using LinearAlgebra

export integrate, DOProblem

include("utils.jl")
include("low_rank_timestepping.jl")
mutable struct DOProblem
    flux # semidiscretization
    factorized # does the flux factorize? true/false
    U0 # initial set of modes n × s
    Z0 # initial set of stochastic coefficients l × s
end

function integrate(prob::DOProblem, ts, σ_min, integration_method = :direct_time_marching)
    U, Z = prob.U0, prob.Z0
    Us, Zs = [U], [Z]
    integrator = getfield(DO, Symbol(string(integration_method,:!)))
    for i in 2:length(ts)
        println(i)
        Δt = ts[i] - ts[i-1]
        if prob.factorized
            L_U, L_Z = prob.flux(U,Z,ts[i])
            N = Δt*normal_direction(U, Z, L_U, L_Z)
            if norm(N) > σ_min
                U, Z = augment_rank(N, U, Z, σ_min)
                println("   rank augmented")
            else
                integrator(Δt, U, Z, L_U, L_Z)
            end
        else
            L = prob.flux(U*Z', ts[i])
            N = Δt*normal_direction(U, Z, L)
            if norm(N) > σ_min
                U, Z = augment_rank(N, U, Z, σ_min)
                println("   rank augmented!")
            else
                integrator(Δt, U, Z, L)
            end
        end
        push!(Us, U)
        push!(Zs, Z)
    end
    return Us, Zs
end

end


using LinearAlgebra
using CUDA

"""
	sim_activity_gpu(J::Matrix{Float32}, g::Float32, h0::Vector{Float32};
		s=0f0, b=nothing, dt=0.1f0, t_bound=1000.0f0, savestep=1,
		minimum_norm=0)

Integrates a tanh rate network with coupling matrix `J` and strength scaling `g`
from initial condition `h0`, returning a `N x NT` matrix of network unit activity.

When `minimum_norm > 0` the integration will be cut off when the network state
falls too close to the origin to prevent floating point underflow. The keyword
arguments `s` and `b` add self-coupling to `J` and add a bias term, respectively.
"""
function sim_activity_gpu(J::Matrix{Float32}, g::Float32, h0::Vector{Float32};
		s=0f0, b=nothing, dt=0.1f0, t_bound=1000.0f0, savestep=1,
		minimum_norm=0)

	N = size(J)[1]
	# Check network size and type consistency
	@assert N == size(J)[2] == length(h0)
	@assert Float32 == typeof(dt) == typeof(t_bound) == typeof(s)

	if b isa Array{Float32} && (size(b) == (N,) || size(b) == (N, 1))
		b = CuArray{Float32}(b)
		use_bias = true
	else
		use_bias = false
	end

	# Discretized simulation length
	maxstep = Int64(t_bound/dt)

	# Allocate memory for output
	out = CuArray{Float32,2}(undef, N, maxstep ÷ savestep)

	h = CuArray{Float32,1}(h0) # Set initial condition
	r = CuArray{Float32,1}(undef, N)

	J = CuArray{Float32,2}(g .* J + s * I)

	step = 0
	while step < maxstep
		step += 1
		map!(tanh, r, h)
		mul!(h, J, r, dt, 1 - dt)
		if use_bias
			h .+= dt .* b
		end
		if step % savestep == 0 # Save every specified time interval
			savept = step ÷ savestep
			out[:,savept] .= h
		end
		if norm(h) < minimum_norm
			D = -I + J .* ((sech.(h)).^2)'
			if all(real(eigvals(D)) .< 0)
				break
			end
		end
	end

	lastindex = step ÷ savestep
	return Matrix{Float32}(view(out, :, 1:lastindex))
end

"""
	sim_lyapunov_gpu(J::Matrix{Float32}, g::Float32, h0::Vector{Float32},
		dh::Vector{Float32}, Q0::Matrix{Float32}; s=0f0, b=nothing, dt=0.1f0,
		t_bound=1000.0f0, t_ONS=1.0f0, wait=0.0f0, return_max_LE=false,
		printstuff=false, scale_limit=100, p_factor=1f-1)

Integrates a tanh network with a given weight matrix and gain parameter from a
specified initial condition, and returns a Lyapunov spectrum calculated along
the trajectory.

Optionally returns the maximum Lyapunov exponent calculated using a direct
numerical simulation of a perturbed trajectory.
"""
function sim_lyapunov_gpu(J::Matrix{Float32}, g::Float32, h0::Vector{Float32},
		dh::Vector{Float32}, Q0::Matrix{Float32}; s=0f0, b=nothing, dt=0.1f0,
		t_bound=1000.0f0, t_ONS=1.0f0, wait=0.0f0, return_max_LE=false,
		printstuff=false, scale_limit=100, p_factor=1f-1)

	N = size(J, 1)
	@assert(size(J) == (length(h0), length(dh)) == size(Q0) == (N,N),
			"Dimensions of the coupling matrix, IC, initial ONS, and " *
			"perturbation must be consistent"
	)
	@assert(Float32 == typeof(s) == typeof(dt) == typeof(t_bound) == typeof(t_ONS)
			== typeof(wait), "Make sure all scalar parameters are the same type"
	)

	maxstep = round(Int64, t_bound / dt)
	startstep = round(Int64, wait / dt)
	tONSstep = round(t_ONS / dt)
	@assert(maxstep % tONSstep == startstep % tONSstep == 0,
		"t_ONS must exactly divide t_bound and wait"
	)

	if b isa Array{Float32} && (size(b) == (N,) || size(b) == (N, 1))
		b = CuArray{Float32,1}(reshape(b, (N,)))
		use_bias = true
	else
		use_bias = false
	end

	# Allocate memory for output
	spectrum = CUDA.zeros(N)

	# Allocate vector for activity and firing rate of current state
	h = CuArray{Float32,1}(h0)
	r = CuArray{Float32,1}(undef, N)

	# Load network structure and ICs to GPU
	J = CuArray{Float32,2}(g .* J + s * I)
	Q = CuArray{Float32,2}(Q0)

	# Allocate memory for Jacobian calculation and orthonormal basis
	sech2(x::Float32)::Float32 = 1 / (CUDA.cosh(x))^2
	sech2h = CuArray{Float32,2}(undef, 1, N)
	D = CuArray{Float32,2}(undef, N, N)
	dQ = CuArray{Float32,2}(undef, N, N)

	if return_max_LE
		max_LE = 0
		p_factor = convert(Float32, p_factor)
		dh = CuArray{Float32,1}(dh ./ norm(dh))
		h2 = h .+ p_factor .* dh
		r2 = CuArray{Float32,1}(undef, N)
	end

	step = 0
	while step < maxstep
		step += 1

		#Evolve the trajectory and the orthonormal basis in the tangent space
		map!(tanh, r, h)
		mul!(h, J, r, dt, 1 - dt)
		if use_bias
			h .+= dt .* b
		end

		map!(sech2, sech2h, h)
		copyto!(D, J)
		D .*= sech2h
		D[diagind(D)] .-= 1

		mul!(dQ, D, Q, dt, 0.0f0)
        Q .+= dQ

		#Evolve the perturbed trajectory and calculate how much it has grown
		if return_max_LE
			map!(tanh, r2, h2)
			mul!(h2, J, r2, dt, 1 - dt)
			if use_bias
				h2 .+= dt .* b
			end

			dh .= (h2 .- h) ./ p_factor
			scale = norm(dh)

			# If the perturbation scale has grown or shrunk too much, then
			# reset it to 1 and track the cumulative sum of log scales
			if scale > scale_limit || scale < 1 / scale_limit || step == startstep
				h2 .= h .+ dh ./ scale .* p_factor
				if step > startstep
					max_LE += log(scale)
				end
			end
		end

		# At each re-orthonormalization step, the log diagonal entries of R are
		# saved in a cumulative sum
		if step % tONSstep == 0
			printstuff && println("Time step: ", step)
			copyto!(dQ, Q)
			F = qr!(dQ)
			diags = view(F.factors, diagind(F.factors))
			if step > startstep  # forcing positive R-diagonal elements
				spectrum .+= log.(abs.(diags))
			end
			fill!(Q, 0)
			map!(sign, view(Q, diagind(Q)), diags)
            lmul!(F.Q, Q)
		end
	end

	# Update the maximum Lyapunov exponent after the simulation has ended so
	# that the full simulation time is used for the calculation
	true_time = dt * (maxstep - startstep)
	if return_max_LE
		max_LE += log(norm((h2 .- h) ./ p_factor))
		return Vector{Float32}(spectrum ./ true_time), max_LE / true_time
	else
		return Vector{Float32}(spectrum ./ true_time)
	end
end


function invariants!(W1::Vector{Float64},W2::Vector{Float64},lf::Vector{Float64},
    lb::Vector{Float64},A::Vector{Float64},Q::Vector{Float64},
    c0::Float64,beta::Float64,rho::Float64,μ::Float64,diff::Float64,
    Qs::Float64,As::Float64)

    # characteristic speeds
    lb[1] = (Q[1]./A[1].*Qs./As .- sqrt.(0.5.*beta./rho).*(A[1].*As).^0.25).*As./Qs;
    lf[1] = (Q[end]./A[end].*Qs./As .+ sqrt.(0.5.*beta./rho).*(A[end].*As).^0.25).*As./Qs;

    # invariants
    W1[:] .= (Q[end-1:end]./A[end-1:end].*Qs./As .+
        4.*(sqrt.(0.5.*beta./rho).*(A[end-1:end].*As).^0.25 .- c0)).*As./Qs;
    W2[:] .= (Q[1:2]./A[1:2].*Qs./As .-
        4.*(sqrt.(0.5.*beta./rho).*(A[1:2].*As).^0.25 .- c0)).*As./Qs;

end


function incident_infections(p1::InfParams, p2::InfParams, S::Int, I1::Vector{Int}, I2::Vector{Int}, R::Int)
    iszero(S) && return 0,0
    infSum1 = 0
    infSum2 = 0
    N = S + sum(I1) + sum(I2) + R
    susceptible_prop = S/N

    for (i, inf) in enumerate(I1)
        d = p1.infectiousness[i]
        k,θ = Distributions.params(d)
        d′ = Gamma(k*susceptible_prop, θ) # conversion from R₀ to Rₜ
        infSum1 += rand(RVsum(d′, inf))
    end

    for (i, inf) in enumerate(I2)
        d = p2.infectiousness[i]
        k,θ = Distributions.params(d)
        d′ = Gamma(k*susceptible_prop, θ) # conversion from R₀ to Rₜ
        infSum2 += rand(RVsum(d′, inf))
    end

    infSum1 = floor(Int, infSum1)
    infSum2 = floor(Int, infSum2)
    inf_sum = infSum1 + infSum2
    overage = inf_sum - S

    if overage > 0
        half_over = div(overage, 2)
        remove1 = min(infSum1, half_over)
        infSum1 -= remove1
        overage -= remove1
        infSum2 -= overage
    end

    return infSum1, infSum2
end

function symptomatic_isolation(params::InfParams, I::Vector{Int})
    isolating = zero(I)
    for (i, inf) in enumerate(I)
        symptomatic_prob = params.symptom_probs[i]
        isolation_prob = symptomatic_prob*(1-params.asymptomatic_prob)*params.symptomatic_isolation_prob
        isolating[i] = rand(Binomial(inf,isolation_prob))
    end
    return isolating
end

function positive_tests(params::InfParams, I::Vector{Int}, T::Matrix{Int}, a::CovidAction)
    pos_tests = zeros(Int, length(params.pos_test_probs))

    for (i, inf) in enumerate(I)
        num_already_tested = sum(@view T[:,i])
        num_tested = floor(Int,(inf-num_already_tested)*a.testing_prop)
        pos_tests[i] = rand(Binomial(num_tested,params.pos_test_probs[i]))
    end
    return pos_tests
end

# Only record number that have taken the test, the number that return postive is
# Binomial dist, such that s' is stochastic on s.
function update_isolations(params::InfParams, I, R, T, a::CovidAction)

    sympt = symptomatic_isolation(params, I) # Number of people isolating due to symptoms
    pos_tests = positive_tests(params, I, T, a)

    sympt_prop = sympt ./ I # Symptomatic Isolation Proportion
    replace!(sympt_prop, NaN=>0.0)

    R += sum(sympt)
    I .-= sympt

    T[end,:] .= pos_tests

    @. T = floor(Int, (1 - sympt_prop)' * T)

    R += sum(@view T[1,:])
    @views I .-= T[1,:]

    @assert all(≥(0), I)

    CovidPOMDPs.shift_test!(T)

    return I, R, T, pos_tests
end

function CovidPOMDPs.sim_step(pomdp::DoubleCovidPOMDP, state::DoubleCovidState, a::CovidAction)
    (;S, I1, I2, R, T1, T2, params, prev_action) = state
    I1 = copy(I1)
    I2 = copy(I2)
    T1 = copy(T1)
    T2 = copy(T2)

    # Update symptomatic and testing-based isolations
    I1, R, T1, pos_tests1 = update_isolations(pomdp.params, I1, R, T1, a)
    I2, R, T2, pos_tests2 = update_isolations(params, I2, R, T2, a)

    # Incident Infections
    R += I1[end]
    R += I2[end]
    I1 = circshift(I1, 1) # shift_inf!(I1)
    I2 = circshift(I2, 1) # shift_inf!(I2)

    inf1, inf2 = incident_infections(pomdp.params, params, S, I1, I2, R)
    I1[1] = inf1
    I2[1] = inf2

    S -= (inf1 + inf2)
    sp = DoubleCovidState(S, I1, I2, R, T1, T2, params, a)

    return sp, inf1+inf2, pos_tests1 .+ pos_tests2
end


using DataStructures
using Pkg
using Plots

using GeometricProblems.LotkaVolterra2d
using GeometricProblems.Diagnostics

# select Plots backend
gr()
# pgfplots()
# plotly()
# plotlyjs()

# create output directory
if !isdir(out_dir * "/" * fig_dir)
    mkpath(out_dir * "/" * fig_dir)
end

# collect simulation parameters
function get_params(Δt, ntime, problem)
    params = OrderedDict()
    params["Δt"] = Δt
    params["nt"] = ntime
    params["problem"] = problem
    params["GeometricIntegrators.jl"] = Pkg.installed()["GeometricIntegrators"]
    params["GeometricProblems.jl"] = Pkg.installed()["GeometricProblems"]
    params
end

function run_lotka_volterra_2d(int::Integrator, ode::Equation, Δt, nt, filename; kwargs...)
    sim = Simulation(ode, int, Δt, "Lotka-Volterra 2d", prefix * filename * ".hdf5", nt; kwargs...)
    sol = run!(sim)
    run_diagnostics(sol, Δt, nt, filename)
end

function run_lotka_volterra_2d(tableau::AbstractTableau, ode::Equation, Δt, nt; integrator=Integrator, filename=nothing)
    int = integrator(ode, tableau, Δt)

    println()
    println("Running ", tableau.name, " (", filename, ") for ", nt, " time steps...")

    run_lotka_volterra_2d(int, ode, Δt, nt, filename)
end

function run_lotka_volterra_2d(basis::Basis, quadrature::Quadrature, ode::Equation, Δt, nt; integrator=Integrator, filename=nothing)
    int = IntegratorCGVI(ode, basis, quadrature, Δt)

    println()
    println("Running ", filename, " for ", nt, " time steps of Δt = ", Δt, "...")

    run_lotka_volterra_2d(int, ode, Δt, nt, filename)
end

function run_diagnostics(sol, Δt, nt, filename)
    if filename ≠ nothing
        ndrift = div(nt, 10)

        if nt > 100000
            npl = nplot
        else
            npl = 1
        end

        if nt ≤ 200
            markersize=5
        else
            markersize=.5
        end

        # plot overview
        plotlotkavolterra2d(sol)
        savefig(out_dir * "/" * fig_dir * "/" * prefix * filename * fig_suff)

        # plot solution
        plotlotkavolterra2dsolution(sol)
        # markersize=markersize
        # xlim(XLIM)
        # ylim(YLIM)
        # xlabel(L"$q_{1} (t)$", labelpad=10, fontsize=20)
        # ylabel(L"$q_{2} (t)$", labelpad=10, fontsize=20)
        savefig(out_dir * "/" * fig_dir * "/" * prefix * filename * "_solution" * fig_suff)

        # plot energy error
        H, ΔH = compute_energy_error(sol.t, sol.q);
        plotenergyerror(sol.t, ΔH)
        savefig(out_dir * "/" * fig_dir * "/" * prefix * filename * "_energy_error" * fig_suff)

        # plot energy drift
        Tdrift, Hdrift = compute_error_drift(sol.t, ΔH, ndrift)
        plotenergydrift(Tdrift, Hdrift)
        savefig(out_dir * "/" * fig_dir * "/" * prefix * filename * "_energy_drift" * fig_suff)

        # plot momentum error
        if isdefined(sol, :p)
            Δp = compute_momentum_error(sol.t, sol.q, sol.p)
            plotmomentumerror(sol.t, Δp)
            savefig(out_dir * "/" * fig_dir * "/" * prefix * filename * "_momentum_error" * fig_suff)
        end

        # plot Lagrange multiplier
        if isdefined(sol, :λ)
            plotlagrangemultiplier(sol.t, sol.λ)
            savefig(out_dir * "/" * fig_dir * "/" * prefix * filename * "_lambda" * fig_suff)
        end
    end
end

function run_tableau_list(tableaus, ode, Δt, ntime)
    for tab in tableaus
        if length(tab) ≥ 4
            nt = tab[4]
        else
            nt = ntime
        end

        if length(tab) ≥ 5
            dt = tab[5]
        else
            dt = Δt
        end

        run_lotka_volterra_2d(tab[2], ode, dt, nt; integrator=tab[1], filename=tab[3])
    end
end

function run_tableau_list_cgvi(tableaus, ode, Δt, ntime)
    for tab in tableaus
        if length(tab) ≥ 4
            nt = tab[4]
        else
            nt = ntime
        end

        if length(tab) ≥ 5
            dt = tab[5]
        else
            dt = Δt
        end

        run_lotka_volterra_2d(tab[1], tab[2], ode, dt, nt; filename=tab[3])
    end
end


using Interpolations
using Test

@testset "IO" begin
    SPACE = " "

    @testset "BSpline" begin
        A = rand(8,20)

        itp = interpolate(A, BSpline(Constant()))
        @test summary(itp) == "8×20 interpolate(::Matrix{Float64}, BSpline(Constant())) with element type Float64" ||
              summary(itp) == "8×20 interpolate(::Array{Float64,2}, BSpline(Constant())) with element type Float64"

        itp = interpolate(A, BSpline(Constant()))
        @test summary(itp) == "8×20 interpolate(::Matrix{Float64}, BSpline(Constant())) with element type Float64" ||
              summary(itp) == "8×20 interpolate(::Array{Float64,2}, BSpline(Constant())) with element type Float64"

        itp = interpolate(A, BSpline(Linear()))
        @test summary(itp) == "8×20 interpolate(::Matrix{Float64}, BSpline(Linear())) with element type Float64" ||
              summary(itp) == "8×20 interpolate(::Array{Float64,2}, BSpline(Linear())) with element type Float64"

        itp = interpolate(A, BSpline(Quadratic(Reflect(OnCell()))))
        @test summary(itp) == "8×20 interpolate(OffsetArray(::Matrix{Float64}, 0:9, 0:21), BSpline(Quadratic(Reflect(OnCell())))) with element type Float64" ||
              summary(itp) == "8×20 interpolate(OffsetArray(::Array{Float64,2}, 0:9, 0:21), BSpline(Quadratic(Reflect(OnCell())))) with element type Float64"

        itp = interpolate(A, (BSpline(Linear()), NoInterp()))
        @test summary(itp) == "8×20 interpolate(::Matrix{Float64}, (BSpline(Linear()), NoInterp())) with element type Float64" ||
              summary(itp) == "8×20 interpolate(::Array{Float64,2}, (BSpline(Linear()), NoInterp())) with element type Float64"

        itp = interpolate!(copy(A), BSpline(Quadratic(InPlace(OnCell()))))
        @test summary(itp) == "8×20 interpolate(::Matrix{Float64}, BSpline(Quadratic(InPlace(OnCell())))) with element type Float64" ||
              summary(itp) == "8×20 interpolate(::Array{Float64,2}, BSpline(Quadratic(InPlace(OnCell())))) with element type Float64"
    end

    @testset "Gridded" begin
        A = rand(20)
        A_x = collect(1.0:2.0:40.0)
        knots = (A_x,)
        itp = interpolate(knots, A, Gridded(Linear()))
        @test summary(itp) == "20-element interpolate((::Vector{Float64},), ::Vector{Float64}, Gridded(Linear())) with element type Float64" ||
              summary(itp) == "20-element interpolate((::Array{Float64,1},), ::Array{Float64,1}, Gridded(Linear())) with element type Float64"

        A = rand(8,20)
        knots = ([x^2 for x = 1:8], [0.2y for y = 1:20])
        itp = interpolate(knots, A, Gridded(Linear()))
        @test summary(itp) == "8×20 interpolate((::Vector{Int64},::Vector{Float64}), ::Matrix{Float64}, Gridded(Linear())) with element type Float64" ||
              summary(itp) == "8×20 interpolate((::Array{Int64,1},::Array{Float64,1}), ::Array{Float64,2}, Gridded(Linear())) with element type Float64"

        itp = interpolate(knots, A, (Gridded(Linear()),Gridded(Constant())))
        @test summary(itp) == "8×20 interpolate((::Vector{Int64},::Vector{Float64}), ::Matrix{Float64}, (Gridded(Linear()), Gridded(Constant()))) with element type Float64" ||
              summary(itp) == "8×20 interpolate((::Array{Int64,1},::Array{Float64,1}), ::Array{Float64,2}, (Gridded(Linear()), Gridded(Constant()))) with element type Float64"

        # issue #260
        A = (1:4)/4
        itp = interpolate((range(0.0, stop=0.3, length=4),), A, Gridded(Linear()))
        io = IOBuffer()
        show(io, MIME("text/plain"), itp)
        str1 = String(take!(io))
        str2 = "4-element interpolate((0.0:0.1:0.3,), ::Array{Float64,1}, Gridded(Linear())) with element type Float64:\n 0.25\n 0.5 \n 0.75\n 1.0"
        str3 = "4-elementinterpolate((0.0:0.1:0.3,),::Vector{Float64},Gridded(Linear()))withelementtypeFloat64:\n 0.25 \n 0.5 \n 0.75\n 1.0"
        @test filter(!isspace, str1) == filter(!isspace, str3) ||
              filter(!isspace, str1) == filter(!isspace, str2)
        io2 = IOBuffer()
        show(io2, itp)
        str1 = String(take!(io2))
        str2 = "4-element interpolate((0.0:0.1:0.3,), ::Array{Float64,1}, Gridded(Linear())) with element type Float64:\n 0.25\n 0.5 \n 0.75\n 1.0"
        str3 = "4-elementinterpolate((0.0:0.1:0.3,),::Vector{Float64},Gridded(Linear()))withelementtypeFloat64:\n 0.25\n 0.5\n 0.75\n 1.0"
        @test filter(!isspace, str1) == filter(!isspace, str3) ||
              filter(!isspace, str1) == filter(!isspace, str2)
    end

    @testset "scaled" begin
        itp = interpolate(1:1.0:10, BSpline(Linear()))
        sitp = scale(itp, -3:.5:1.5)
        @test summary(sitp) == "10-element scale(interpolate(::Vector{Float64}, BSpline(Linear())), (-3.0:0.5:1.5,)) with element type Float64" ||
              summary(sitp) == "10-element scale(interpolate(::Array{Float64,1}, BSpline(Linear())), (-3.0:0.5:1.5,)) with element type Float64"
        io = IOBuffer()
        show(io, MIME("text/plain"), sitp)
        str = String(take!(io))
        @test str == "10-element scale(interpolate(::Vector{Float64}, BSpline(Linear())), (-3.0:0.5:1.5,)) with element type Float64:\n  1.0\n  2.0\n  3.0\n  4.0\n  5.0\n  6.0\n  7.0\n  8.0\n  9.0\n 10.0" ||
              str == "10-element scale(interpolate(::Array{Float64,1}, BSpline(Linear())), (-3.0:0.5:1.5,)) with element type Float64:\n  1.0\n  2.0\n  3.0\n  4.0\n  5.0\n  6.0\n  7.0\n  8.0\n  9.0\n 10.0"
        io2 = IOBuffer()
        show(io2, sitp)
        str2 = String(take!(io2))
        @test str2 == "10-element scale(interpolate(::Vector{Float64}, BSpline(Linear())), (-3.0:0.5:1.5,)) with element type Float64:\n  1.0\n  2.0\n  3.0\n  4.0\n  5.0\n  6.0\n  7.0\n  8.0\n  9.0\n 10.0" ||
              str2 == "10-element scale(interpolate(::Array{Float64,1}, BSpline(Linear())), (-3.0:0.5:1.5,)) with element type Float64:\n  1.0\n  2.0\n  3.0\n  4.0\n  5.0\n  6.0\n  7.0\n  8.0\n  9.0\n 10.0"

        gauss(phi, mu, sigma) = exp(-(phi-mu)^2 / (2sigma)^2)
        testfunction(x,y) = gauss(x, 0.5, 4) * gauss(y, -.5, 2)
        xs = -5:.5:5
        ys = -4:.2:4
        zs = Float64[testfunction(x,y) for x in xs, y in ys]
        itp2 = interpolate(zs, BSpline(Quadratic(Flat(OnGrid()))))
        sitp2 = scale(itp2, xs, ys)
        teststring = "21×41 scale(interpolate(OffsetArray(::Array{Float64,2}, 0:22, 0:42), BSpline(Quadratic(Flat(OnGrid())))), (-5.0:0.5:5.0,$SPACE-4.0:0.2:4.0)) with element type Float64"
        teststring2 = "21×41 scale(interpolate(OffsetArray(::Matrix{Float64}, 0:22, 0:42), BSpline(Quadratic(Flat(OnGrid())))), (-5.0:0.5:5.0,$SPACE-4.0:0.2:4.0)) with element type Float64"
        @test summary(sitp2) == teststring2 ||
              summary(sitp2) == teststring
        io3 = IOBuffer()
        show(io3, sitp2)
        str3 = String(take!(io3))
        @test occursin(teststring2, str3) || occursin(teststring, str3)
    end

    @testset "Monotonic" begin
        A_x = collect(1.0:2.0:40.0)
        A = map(x->x^2, A_x)
        itp = interpolate(A_x, A, FritschButlandMonotonicInterpolation())
        teststring = "20-element interpolate(::Array{Float64,1}, ::Array{Float64,1}, FritschButlandMonotonicInterpolation()) with element type Float64"
        teststring2 = "20-element interpolate(::Vector{Float64}, ::Vector{Float64}, FritschButlandMonotonicInterpolation()) with element type Float64"
        @test summary(itp) == teststring2 || summary(itp) == teststring

        io = IOBuffer()
        show(io, itp)
        str = String(take!(io))
        @test occursin(teststring2, str) || occursin(teststring, str)

        io2 = IOBuffer()
        show(io2, MIME("text/plain"), itp)
        str = String(take!(io2))
        @test occursin(teststring2, str) || occursin(teststring, str)
    end

    @testset "Extrapolation" begin
        A = rand(8,20)

        itpg = interpolate(A, BSpline(Linear()))
        etpg = extrapolate(itpg, Flat())
        teststring = "8×20 extrapolate(interpolate(::Array{Float64,2}, BSpline(Linear())), Flat()) with element type Float64"
        teststring16 = "8×20 extrapolate(interpolate(::Matrix{Float64}, BSpline(Linear())), Flat()) with element type Float64"
        @test summary(etpg) == teststring16 || summary(etpg) == teststring
        io = IOBuffer()
        show(io, etpg)
        str = String(take!(io))
        @test occursin(teststring16, str) || occursin(teststring, str)

        etpf = extrapolate(itpg, NaN)
        teststring2 = "8×20 extrapolate(interpolate(::Array{Float64,2}, BSpline(Linear())), NaN) with element type Float64"
        teststring16 = "8×20 extrapolate(interpolate(::Matrix{Float64}, BSpline(Linear())), NaN) with element type Float64"
        @test summary(etpf) == teststring16 || summary(etpf) == teststring2
        io2 = IOBuffer()
        show(io2, etpf)
        str2 = String(take!(io2))
        @test occursin(teststring16, str2) || occursin(teststring2, str2)
    end

    @testset "Combinations" begin
        xs = 1.0:0.5:2.0
        ys = xs.^3
        interp_cubic = CubicSplineInterpolation(xs, ys)
        io = IOBuffer()
        show(io, MIME("text/plain"), interp_cubic)
        str = String(take!(io))
        @test occursin("3.375", str)
    end

    @testset "Issue #260" begin
        itp = scale(interpolate(zeros(3,3,3),BSpline(Cubic(Free(OnGrid())))),2:4,2:4,2:4)
        io = IOBuffer()
        show(io, MIME("text/plain"), itp)
        str = String(take!(io))
        @test str == "3×3×3 scale(interpolate(OffsetArray(::Array{Float64, 3}, 0:4, 0:4, 0:4), BSpline(Cubic(Free(OnGrid())))), (2:4, 2:4, 2:4)) with element type Float64:\n[:, :, 1] =\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n\n[:, :, 2] =\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n\n[:, :, 3] =\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n 0.0  0.0  0.0" ||
              str == "3×3×3 scale(interpolate(OffsetArray(::Array{Float64,3}, 0:4, 0:4, 0:4), BSpline(Cubic(Free(OnGrid())))), (2:4, 2:4, 2:4)) with element type Float64:\n[:, :, 1] =\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n\n[:, :, 2] =\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n\n[:, :, 3] =\n 0.0  0.0  0.0\n 0.0  0.0  0.0\n 0.0  0.0  0.0"
    end
end # testset