text:
## Matrix for transformations
function transformMaker(coeff, S1Mat, S2Mat; Nc=1)
    NS1 = size(S1Mat)[2]
    NS2 = size(S2Mat)[2]
    if Nc==1
        S0iso = coeff[:,1:2]
        S1iso = transpose(S1Mat*transpose(coeff[:,2+1:2+NS1]))
        S2iso = transpose(S2Mat*transpose(coeff[:,2+NS1+1:2+NS1+NS2]))
    else
        S0iso = coeff[:,1:2*Nc]
        S1MatChan = blockdiag(collect(Iterators.repeated(S1Mat,Nc))...)
        S2MatChan = blockdiag(collect(Iterators.repeated(S2Mat,Nc*Nc))...)
        S1iso = transpose(S1MatChan*transpose(coeff[:,2*Nc+1:2*Nc+Nc*NS1]))
        S2iso = transpose(S2MatChan*transpose(coeff[:,2*Nc+Nc*NS1+1:end]))
    end
    return hcat(S0iso,S1iso,S2iso)
end

function S1_iso_matrix(fhash)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S1 coeff vector to S1iso by
    #   summing over l
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index")
    if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["theta_value"])
    Nj      = length(fhash["j_value"])
    Nf      = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]
    ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    Niso    = Omega ? Nj+2 : Nj+1
    Mat     = zeros(Int32, Niso, Nf)

    # first J elements of iso
    for j = 1:Nj
        for l = 1:Nl
            λ = ψ_ind[j,l]
            Mat[j, λ] = 1
        end
    end

    # Next elements are ϕ, Ω
    I0     = Nj+1
    Mat[I0, ϕ_ind] = 1
    if Omega Mat[I0+1, Ω_ind] = 1 end

    return sparse(Mat)
end

function S2_iso_matrix(fhash)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S2 coeff vector to S2iso by
    #   summing over l1,l2 and fixed Δl.
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index")
    if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["theta_value"])
    Nj      = length(fhash["j_value"])
    Nf      = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]
    ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    Niso    = Omega ? Nj*Nj*Nl+4*Nj+4 : Nj*Nj*Nl+2*Nj+1
    Mat     = zeros(Int32, Niso, Nf*Nf)

    # first J*J*L elements of iso
    for j1 = 1:Nj
        for j2 = 1:Nj
            for l1 = 1:Nl
                for l2 = 1:Nl
                    DeltaL = mod(l1-l2, Nl)
                    λ1     = ψ_ind[j1,l1]
                    λ2     = ψ_ind[j2,l2]

                    Iiso   = j1+Nj*((j2-1)+Nj*DeltaL)
                    Icoeff = λ1+Nf*(λ2-1)
                    Mat[Iiso, Icoeff] = 1
                end
            end
        end
    end

    # Next J elements are λϕ, then J elements ϕλ
    for j = 1:Nj
        for l = 1:Nl
            λ      = ψ_ind[j,l]
            Iiso   = Nj*Nj*Nl+j
            Icoeff = λ+Nf*(ϕ_ind-1)  # λϕ
            Mat[Iiso, Icoeff] = 1

            Iiso   = Nj*Nj*Nl+Nj+j
            Icoeff = ϕ_ind+Nf*(λ-1)  # ϕλ
            Mat[Iiso, Icoeff] = 1
        end
    end

    # Next 1 element is ϕϕ
    I0     = Nj*Nj*Nl+Nj+Nj+1
    Icoeff = ϕ_ind+Nf*(ϕ_ind-1)
    Mat[I0, Icoeff] = 1

    # If the Omega filter exists, add more terms
    if Omega
        # Next J elements are λΩ, then J elements Ωλ
        for j = 1:Nj
            for l = 1:Nl
                λ      = ψ_ind[j,l]
                Iiso   = I0+j
                Icoeff = λ+Nf*(Ω_ind-1)  # λΩ
                Mat[Iiso, Icoeff] = 1

                Iiso   = I0+Nj+j
                Icoeff = Ω_ind+Nf*(λ-1)  # Ωλ
                Mat[Iiso, Icoeff] = 1
            end
        end
        # Next 3 elements are ϕΩ, Ωϕ, ΩΩ
        Iiso   = I0+Nj+Nj
        Mat[Iiso+1, ϕ_ind+Nf*(Ω_ind-1)] = 1
        Mat[Iiso+2, Ω_ind+Nf*(ϕ_ind-1)] = 1
        Mat[Iiso+3, Ω_ind+Nf*(Ω_ind-1)] = 1
    end

    return sparse(Mat)
end

function S1_iso_matrix3d(fhash)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S1 coeff vector to S1iso by
    #   summing over l
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18
    # AKS 2021-Feb-22

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index") & fhash["Omega3d"]

    # unpack fhash
    Nl      = length(fhash["theta_value"])
    Nj      = length(fhash["j_value"])
    Nk      = length(fhash["k_value"])
    Nf      = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]

    # number of iso coefficients
    Niso    = Omega ? (Nj+2)*Nk+3 : (Nj+1)*Nk+1
    Nstep   = Omega ? Nj+2 : Nj+1
    Mat     = zeros(Int32, Niso, Nf)

    # first J elements of iso
    for j = 1:Nj
        for l = 1:Nl
            for k=1:Nk
                λ = ψ_ind[j,l,k]
                Mat[j+(k-1)*Nstep, λ] = 1
            end
        end
    end

    #this is currently how I am storing the phi_k, omega_k
    j = Nj+1
    l = 1
    for k=1:Nk
        λ = ψ_ind[j,l,k]
        Mat[j+(k-1)*Nstep, λ] = 1
    end

    if Omega
        j = Nj+1
        l = 2
        for k=1:Nk
            λ = ψ_ind[j,l,k]
            Mat[j+1+(k-1)*Nstep, λ] = 1
        end
    end

    ### these are the 3d globals

    #phi0
    λ = ψ_ind[Nj+1,1,Nk+1]
    if Omega
        Mat[Niso-2, λ] = 1
    else
        Mat[Niso, λ] = 1
    end

    if Omega
        #omega_0
        λ = ψ_ind[Nj+1,2,Nk+1]
        Mat[Niso-1, λ] = 1

        # omega_3d
        λ = ψ_ind[Nj+1,3,Nk+1]
        Mat[Niso, λ] = 1
    end

    return sparse(Mat)
end

function S2_iso_matrix3d(fhash)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S2 coeff vector to S2iso by
    #   summing over l1,l2 and fixed Δl.
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18
    # AKS 2021-Feb-22

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index") & fhash["Omega3d"]

    # unpack fhash
    L      = length(fhash["theta_value"])
    J      = length(fhash["j_value"])
    K      = length(fhash["k_value"])
    Nf     = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]

    # number of iso coefficients
    NisoL   = Omega ? J*J*K*K*L+2*J*K*K+2*J*K*K+6*J*K : J*J*K*K*L+2*J*K*K+2*J*K
    NnoL    = Omega ? 2*K+3 : K+1
    Mat     = zeros(Int32, NisoL+NnoL^2, Nf*Nf)

    indexes = 1:Nf
    mask = fhash["psi_ind_L"].==0
    NnoLind = indexes[mask]

    #should probably think about reformatting to (Nj+1) so phi
    #phi cross terms are in the right place

    # first J*J*K*K*L elements of iso
    ref = 0
    for j1 = 1:J
        for j2 = 1:J
            for l1 = 1:L
                for l2 = 1:L
                    for k1 = 1:K
                        for k2 = 1:K
                            DeltaL = mod(l1-l2, L)
                            λ1     = ψ_ind[j1,l1,k1]
                            λ2     = ψ_ind[j2,l2,k2]

                            Iiso   = k1+K*((k2-1)+K*((j1-1)+J*((j2-1)+J*DeltaL)))
                            Icoeff = λ1+Nf*(λ2-1)
                            Mat[Iiso, Icoeff] = 1
                        end
                    end
                end
            end
        end
    end
    ref+=J*J*K*K*L

    # Next JKK elements are λϕ, then JKK elements ϕλ

    for j = 1:J
        for l = 1:L
            for k1 = 1:K
                for k2 =1:K
                    λ      = ψ_ind[j,l,k1]
                    Iiso   = ref + k1+K*((k2-1)+K*(j-1))
                    Icoeff = λ+Nf*(ψ_ind[J+1,1,k2]-1)  # λϕ
                    Mat[Iiso, Icoeff] = 1

                    Iiso   = ref+J*K*K+k1+K*((k2-1)+K*(j-1))
                    Icoeff = ψ_ind[J+1,1,k2]+Nf*(λ-1)  # ϕλ
                    Mat[Iiso, Icoeff] = 1
                end
            end
        end
    end

    ref += 2*J*K*K

    # Next JKK elements are λΩ, then JKK elements Ωλ
    if Omega
        for j = 1:J
            for l = 1:L
                for k1 = 1:K
                    for k2 =1:K
                        λ      = ψ_ind[j,l,k1]
                        Iiso   = ref + k1+K*((k2-1)+K*(j-1))
                        Icoeff = λ+Nf*(ψ_ind[J+1,2,k2]-1)  # λϕ
                        Mat[Iiso, Icoeff] = 1

                        Iiso   = ref+J*K*K+k1+K*((k2-1)+K*(j-1))
                        Icoeff = ψ_ind[J+1,2,k2]+Nf*(λ-1)  # ϕλ
                        Mat[Iiso, Icoeff] = 1
                    end
                end
            end
        end
        ref += 2*J*K*K
    end

    # Next 2*J*K elements are ψϕ0 and ϕ0ψ

    ϕ0 = ψ_ind[J+1,1,K+1]
    for j1=1:J
        for l1 = 1:L
            for k1=1:K
                λ1 = ψ_ind[j1,l1,k1]

                Iiso = ref+k1+K*(j1-1)
                Icoeff = λ1+Nf*(ϕ0-1)
                Mat[Iiso, Icoeff] = 1

                Iiso = ref+J*K+k1+K*(j1-1)
                Icoeff = ϕ0+Nf*(λ1-1)
                Mat[Iiso, Icoeff] = 1
            end
        end
    end
    ref += 2*J*K

    if Omega
        Ω0 = ψ_ind[J+1,2,K+1]
        for j1=1:J
            for l1 = 1:L
                for k1=1:K
                    λ1 = ψ_ind[j1,l1,k1]

                    Iiso = ref+k1+K*(j1-1)
                    Icoeff = λ1+Nf*(Ω0-1)
                    Mat[Iiso, Icoeff] = 1

                    Iiso = ref+J*K+k1+K*(j1-1)
                    Icoeff = Ω0+Nf*(λ1-1)
                    Mat[Iiso, Icoeff] = 1
                end
            end
        end
        ref += 2*J*K

        Ω3 = ψ_ind[J+1,3,K+1]
        for j1=1:J
            for l1 = 1:L
                for k1=1:K
                    λ1 = ψ_ind[j1,l1,k1]

                    Iiso = ref+k1+K*(j1-1)
                    Icoeff = λ1+Nf*(Ω3-1)
                    Mat[Iiso, Icoeff] = 1

                    Iiso = ref+J*K+k1+K*(j1-1)
                    Icoeff = Ω3+Nf*(λ1-1)
                    Mat[Iiso, Icoeff] = 1
                end
            end
        end
        ref += 2*J*K
    end

    # take care of the L independent subblocks
    for m1=1:NnoL
        for m2=1:NnoL
            Iiso = NisoL + m1 + NnoL*(m2-1)
            λ1 = NnoLind[m1]
            λ2 = NnoLind[m2]
            Icoeff = λ1+Nf*(λ2-1)
            Mat[Iiso, Icoeff] = 1
        end
    end

    return sparse(Mat)
end

function S1_equiv_matrix(fhash,l_shift)
        # fhash is the filter hash output by fink_filter_hash
        # The output matrix converts an S1 coeff vector to S1iso by
        #   summing over l
        # Matrix is stored in sparse CSC format using SparseArrays.
        # DPF 2021-Feb-18

        # Does hash contain Omega filter?
        Omega   = haskey(fhash, "Omega_index")
        if Omega Ω_ind = fhash["Omega_index"] end

        # unpack fhash
        Nl      = length(fhash["theta_value"])
        Nj      = length(fhash["j_value"])
        Nf      = length(fhash["filt_value"])
        ψ_ind   = fhash["psi_index"]
        ϕ_ind   = fhash["phi_index"]

        # number of iso coefficients
        Niso    = Omega ? Nj+2 : Nj+1
        Mat     = zeros(Int32, Nf, Nf)

        # first J elements of iso
        for j = 1:Nj
            for l = 1:Nl
                λ = ψ_ind[j,l]
                λ1 = ψ_ind[j,mod1(l+l_shift,Nl)]
                Mat[λ1, λ] = 1
            end
        end

        # Next elements are ϕ, Ω
        Mat[ϕ_ind, ϕ_ind] = 1
        if Omega Mat[Ω_ind, Ω_ind] = 1 end

        return sparse(Mat)
end

function S2_equiv_matrix(fhash,l_shift)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S2 coeff vector to S2iso by
    #   summing over l1,l2 and fixed Δl.
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index")
    if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["theta_value"])
    Nj      = length(fhash["j_value"])
    Nf      = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]
    ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    Niso    = Omega ? Nj*Nj*Nl+4*Nj+4 : Nj*Nj*Nl+2*Nj+1
    Mat     = zeros(Int32, Nf*Nf, Nf*Nf)

    # first J*J*L elements of iso
    for j1 = 1:Nj
        for j2 = 1:Nj
            for l1 = 1:Nl
                for l2 = 1:Nl
                    λ1     = ψ_ind[j1,l1]
                    λ2     = ψ_ind[j2,l2]
                    λ1_new     = ψ_ind[j1,mod1(l1+l_shift,Nl)]
                    λ2_new     = ψ_ind[j2,mod1(l2+l_shift,Nl)]

                    Icoeff = λ1+Nf*(λ2-1)
                    Icoeff_new = λ1_new+Nf*(λ2_new-1)
                    Mat[Icoeff_new, Icoeff] = 1
                end
            end
        end
    end

    # Next J elements are λϕ, then J elements ϕλ
    for j = 1:Nj
        for l = 1:Nl
            λ      = ψ_ind[j,l]
            Icoeff = λ+Nf*(ϕ_ind-1)  # λϕ

            λ_new      = ψ_ind[j,mod1(l+l_shift,Nl)]
            Icoeff_new = λ_new+Nf*(ϕ_ind-1)  # λϕ

            Mat[Icoeff_new, Icoeff] = 1

            Icoeff = ϕ_ind+Nf*(λ-1)  # ϕλ
            Icoeff_new = ϕ_ind+Nf*(λ_new-1)  # ϕλ
            Mat[Icoeff_new, Icoeff] = 1
        end
    end

    # Next 1 element is ϕϕ
    Icoeff = ϕ_ind+Nf*(ϕ_ind-1)
    Mat[Icoeff, Icoeff] = 1

    # If the Omega filter exists, add more terms
    if Omega
        # Next J elements are λΩ, then J elements Ωλ
        for j = 1:Nj
            for l = 1:Nl
                λ      = ψ_ind[j,l]
                λ_new      = ψ_ind[j,mod1(l+l_shift,Nl)]
                Icoeff = λ+Nf*(Ω_ind-1)  # λΩ
                Icoeff_new = λ_new+Nf*(Ω_ind-1)  # λΩ
                Mat[Icoeff_new, Icoeff] = 1

                Iiso   = I0+Nj+j
                Icoeff = Ω_ind+Nf*(λ-1)  # Ωλ
                Icoeff_new = Ω_ind+Nf*(λ_new-1)  # Ωλ
                Mat[Icoeff_new, Icoeff] = 1
            end
        end
        # Next 3 elements are ϕΩ, Ωϕ, ΩΩ
        Mat[ϕ_ind+Nf*(Ω_ind-1), ϕ_ind+Nf*(Ω_ind-1)] = 1
        Mat[Ω_ind+Nf*(ϕ_ind-1), Ω_ind+Nf*(ϕ_ind-1)] = 1
        Mat[Ω_ind+Nf*(Ω_ind-1), Ω_ind+Nf*(Ω_ind-1)] = 1
    end

    return sparse(Mat)
end

function S1_iso_matrix_gpu(fhash, prec=Float32)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S1 coeff vector to S1iso by
    #   summing over l
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index")
    if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["theta_value"])
    Nj      = length(fhash["j_value"])
    Nf      = fhash["Nf"]
    ψ_ind   = fhash["psi_index"]
    ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    Niso    = Omega ? Nj+2 : Nj+1
    Mat     = zeros(prec, Niso, Nf)

    # first J elements of iso
    for j = 1:Nj
        for l = 1:Nl
            λ = ψ_ind[j,l]
            Mat[j, λ] = 1.0
        end
    end

    # Next elements are ϕ, Ω
    I0     = Nj+1
    Mat[I0, ϕ_ind] = 1.0
    if Omega Mat[I0+1, Ω_ind] = 1.0 end

    return sparse(Mat)
end

function S2_iso_matrix_gpu(fhash, prec=Float32)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S2 coeff vector to S2iso by
    #   summing over l1,l2 and fixed Δl.
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index")
    if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["theta_value"])
    Nj      = length(fhash["j_value"])
    Nf      = fhash["Nf"]
    ψ_ind   = fhash["psi_index"]
    ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    Niso    = Omega ? Nj*Nj*Nl+4*Nj+4 : Nj*Nj*Nl+2*Nj+1
    Mat     = zeros(prec, Niso, Nf*Nf)

    # first J*J*L elements of iso
    for j1 = 1:Nj
        for j2 = 1:Nj
            for l1 = 1:Nl
                for l2 = 1:Nl
                    DeltaL = mod(l1-l2, Nl)
                    λ1     = ψ_ind[j1,l1]
                    λ2     = ψ_ind[j2,l2]

                    Iiso   = j1+Nj*((j2-1)+Nj*DeltaL)
                    Icoeff = λ1+Nf*(λ2-1)
                    Mat[Iiso, Icoeff] = 1.0
                end
            end
        end
    end

    # Next J elements are λϕ, then J elements ϕλ
    for j = 1:Nj
        for l = 1:Nl
            λ      = ψ_ind[j,l]
            Iiso   = Nj*Nj*Nl+j
            Icoeff = λ+Nf*(ϕ_ind-1)  # λϕ
            Mat[Iiso, Icoeff] = 1.0

            Iiso   = Nj*Nj*Nl+Nj+j
            Icoeff = ϕ_ind+Nf*(λ-1)  # ϕλ
            Mat[Iiso, Icoeff] = 1.0
        end
    end

    # Next 1 element is ϕϕ
    I0     = Nj*Nj*Nl+Nj+Nj+1
    Icoeff = ϕ_ind+Nf*(ϕ_ind-1)
    Mat[I0, Icoeff] = 1.0

    # If the Omega filter exists, add more terms
    if Omega
        # Next J elements are λΩ, then J elements Ωλ
        for j = 1:Nj
            for l = 1:Nl
                λ      = ψ_ind[j,l]
                Iiso   = I0+j
                Icoeff = λ+Nf*(Ω_ind-1)  # λΩ
                Mat[Iiso, Icoeff] = 1.0

                Iiso   = I0+Nj+j
                Icoeff = Ω_ind+Nf*(λ-1)  # Ωλ
                Mat[Iiso, Icoeff] = 1.0
            end
        end
        # Next 3 elements are ϕΩ, Ωϕ, ΩΩ
        Iiso   = I0+Nj+Nj
        Mat[Iiso+1, ϕ_ind+Nf*(Ω_ind-1)] = 1.0
        Mat[Iiso+2, Ω_ind+Nf*(ϕ_ind-1)] = 1.0
        Mat[Iiso+3, Ω_ind+Nf*(Ω_ind-1)] = 1.0
    end

    return sparse(Mat)
end

## Assist in convmap

function pre_subdivide(coeff_list; sub_block = 0)
    Nc, Ns = size(coeff_list)
    Nx, Ny = size(coeff_list[3])
    out_maps = zeros(Nx,Ny,Nc-2,Ns)
    for i=1:Ns
        out_maps[:,:,:,i] .= cat(coeff_list[3:end,i]...,dims=3)
    end

    out = zeros(Int(2^(sub_block)),Int(2^(sub_block)),size(out_maps)[3],size(out_maps)[4])
    red = Nx/(2^(sub_block))
    for i=1:Nx
        for j=1:Ny
            out[ceil(Int,i/red),ceil(Int,j/red),:,:] .+=  out_maps[i,j,:,:]
        end
    end
    return out
end

function post_subdivide(coeff_maps; sub_block = 0)
    Nx, Ny, Nc, Ns = size(coeff_maps)

    out = zeros(Int(2^(sub_block)),Int(2^(sub_block)),size(coeff_maps)[3],size(coeff_maps)[4])
    red = Nx/(2^(sub_block))
    for i=1:Nx
        for j=1:Ny
            out[ceil(Int,i/red),ceil(Int,j/red),:,:] .+=  coeff_maps[i,j,:,:]
        end
    end
    return out
end

function sub_strip(coeff_maps, lab; sub_bl_size=64)
    Nx, Ny, Nc, Ns  = size(coeff_maps)
    Nl = size(lab)[1]
    out_stripped = zeros(Nc,Nx*Ny*Ns)
    new_lab = zeros(Nl,Nx*Ny*Ns)
    for k=1:Ns
        for i=1:Nx
            for j=1:Ny
                new_lab[:,i+Nx*(j-1)+Nx*Ny*(k-1)] .= lab[:,k] .+ [0; 0; 0; 0; 0; sub_bl_size*(i-1); sub_bl_size*(j-1)]
                out_stripped[:,i+Nx*(j-1)+Nx*Ny*(k-1)] .= coeff_maps[i,j,:,k]
            end
        end
    end
    return out_stripped, new_lab
end

function renorm(coeff_mats,filter_hash)
    Nc, Ns  = size(coeff_mats)
    Nf = size(filter_hash["filt_index"])[1]

    S1normed = coeff_mats[1:Nf,:]./sum(coeff_mats[1:Nf,:],dims=1)
    S2normed = coeff_mats[Nf+1:end,:]./sum(coeff_mats[Nf+1:end,:],dims=1)

    return vcat(S1normed,S2normed)
end

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
"""
    KGL01(S::T)

Whittaker
"""
function KGL01(S::T) where {T<:NamedTuple}
  return (S.a+S.b+S.c)/((S.a + (S.a+S.b+S.c))/2.0)
end

function KGL02(S::T) where {T<:NamedTuple}
  return KGL01(S).-1.0
end

function KGL03(S::T) where {T<:NamedTuple}
  return (S.b + S.c)/2.0
end

function KGL04(S::T) where {T<:NamedTuple}
  return S.b + S.c
end

function KGL05(S::T) where {T<:NamedTuple}
  return ((S.a+S.b+S.c)^2)/((S.a+S.b+S.c)^2-2*S.b*S.c)-1.0
end

function KGL06(S::T) where {T<:NamedTuple}
  one_over_sum = 1/(2*S.a+S.b+S.c)
  s1 = (S.a+S.b)*log(S.a+S.b)
  s2 = (S.a+S.c)*log(S.a+S.c)
  return log(2*S.a+S.b+S.c)-(one_over_sum*2*S.a*log(2))-(one_over_sum*(s1+s2))
end

function KGL07(S::T) where {T<:NamedTuple}
  return exp(KGL06(S))-1.0
end

function KGL08(S::T) where {T<:NamedTuple}
  return (S.b + S.c) / (2*S.a + S.b + S.c)
end

function KGL09(S::T) where {T<:NamedTuple}
  return KGL08(S)
end

function KGL10(S::T) where {T<:NamedTuple}
  return S.a / (S.a + S.b + S.c)
end

function KGL11(S::T) where {T<:NamedTuple}
  return (2*S.a) / (2*S.a + S.b + S.c)
end

function KGL12(S::T) where {T<:NamedTuple}
  return (2*S.a+S.b+S.c) * ( 1 - (S.a/(S.a+S.b+S.c)))
end

function KGL13(S::T) where {T<:NamedTuple}
  return min(S.b, S.c) / (max(S.b, S.c) + S.a)
end

function KGL14(S::T) where {T<:NamedTuple}
  return 1 - (S.a*(2*S.a+S.b+S.c))/(2*(S.a+S.b)*(S.a+S.c))
end

function KGL15(S::T) where {T<:NamedTuple}
  return (S.b + S.c)/(S.a + S.b + S.c)
end

function KGL16(S::T) where {T<:NamedTuple}
  return KGL15(S)
end

function KGL17(S::T) where {T<:NamedTuple}
  return min(S.b, S.c)/(S.a+S.b+S.c)
end

function KGL18(S::T) where {T<:NamedTuple}
  return (S.b+S.c)/2
end

function KGL19(S::T) where {T<:NamedTuple}
  return ((S.b*S.c)+1)/((S.a+S.b+S.c)^2-(S.a+S.b+S.c)/2)
end

function KGL20(S::T) where {T<:NamedTuple}
  return 1 - (2*S.a)/(2*S.a.+S.b+S.c)
end

function KGL21(S::T) where {T<:NamedTuple}
  return S.a/(S.a+S.c)
end

function KGL22(S::T) where {T<:NamedTuple}
  return min(S.b, S.c)/(min(S.b, S.c) + S.a)
end

function KGL23(S::T) where {T<:NamedTuple}
  return 2*abs(S.b-S.c)/(2*S.a+S.b+S.c)
end

function KGL24(S::T) where {T<:NamedTuple}
  in_par = (2*S.a+S.b+S.c)/(S.a+S.b+S.c)
  return 1-log(in_par)/log(2)
end

ext: jl


text:
import SolveIPlusOrMinusAkronB: IPlusAtKronBWS, solvi_real_eliminate!,
                                QuasiUpperTriangular, solve1!, transformation1,
                                generalized_sylvester_solver!
using Test

n = 3
x = randn(n,n*n)
a = randn(n,n)
b = randn(n,n)
c = randn(n,n)

ws = IPlusAtKronBWS(n, n, n, 2)

index = [n,n]
j = 1
depth = 2
t = QuasiUpperTriangular(lu(a\b).U)
s = QuasiUpperTriangular(lu(c).U)
d = zeros(n^(depth+1))
d[1:n*n] = ones(n*n)
d_orig = copy(d)
nd = n*n^(depth-1)
drange1 = 1:nd
r = 1.0
solvi_real_eliminate!(1, n, nd, drange1, depth - 1, r, t, s, d, ws)
@test ws.work1[drange1] ≈ kron(s', t)*d_orig[drange1]
kst = kron(s', t)
d_target = -kron(s[1, 2:n], kst*d_orig[drange1])
@test d[nd + 1:end] ≈ d_target

index = [1,2]
drange1 = drange1 .+ nd
d = zeros(n^(depth+1))
d[drange1] = ones(nd)
d_orig = copy(d)
r = 1.0

@show drange1
solvi_real_eliminate!(2, n, nd, drange1, depth - 1, r, t, s, d, ws)
d_target = -kron(s[2, 3:n], kst*d_orig[drange1])
@show ws.work1
@show d
@test d[2*nd + 1:end] ≈ d_target

aa = a\b
t = lu(aa).U
s = lu(c).U

r = 1.0
d = randn(n)
d_orig = copy(d)
depth = 0
s2 = QuasiUpperTriangular(s*s)
t2 = QuasiUpperTriangular(t*t)

t = QuasiUpperTriangular(t)
s = QuasiUpperTriangular(s)

solve1!(r, depth,t, t2, s, s2, d, ws)
d_target = (I(n) + t)\d_orig
@test d_target ≈ d

r = 1.0
d = randn(n*n)
d_orig = copy(d)
depth = 1
index = zeros(Int64,depth)
d_block_number = n^depth
solve1!(r, depth,t, t2, s, s2, d, ws)

d_target = (I(n^2) + kron(s,t))\d_orig

@test d ≈ d_target

r = 1.0
d = randn(n*n*n)
d_orig = copy(d)
depth = 2
index = zeros(Int64,depth)
d_block_number = n^depth
solve1!(r, depth,t, t2, s, s2, d, ws)

d_target = (I(n^3) + kron(kron(s,s),t))\d_orig

@test d ≈ d_target

a = 0.5
b1 = 0.1
b2 = 1.3
depth = 2
tt = QuasiUpperTriangular(t)
ss = QuasiUpperTriangular(s)
d = randn(2*n^depth)
d_orig = copy(d)
transformation1(a, b1, b2, depth, tt, ss, d, ws)
d_target = d_orig + kron([a -b1; -b2 a],kron(s,t))*d_orig

@test d ≈ d_target

d = a*x + b*x*kron(c,c)
generalized_sylvester_solver!(a, b, c, d, 2, ws)

@test d ≈ x


text:
##################################################
## Classical derivative-based, iterative, root-finding algorithms: Newton and Halley
## Historic, we have derivative free versions of similar order

## Newton
"""

    Roots.Newton()

Implements Newton's [method](http://tinyurl.com/b4d7vls): `x_n1 = xn -
f(xn)/f'(xn)`.  This is a quadratically converging method requiring
one derivative. If a derivative is not specified, the `ForwardDiff` package
will be used, as applicable.

Unlike other methods, this method accepts complex inputs.
"""
struct Newton <: UnivariateZeroMethod
end

function callable_function(method::Newton, f::Tuple)
    length(f) == 1 && return FirstDerivative(f[1], D(f[1]))
    FirstDerivative(f[1], f[2])
end
callable_function(method::Newton, f::Any) = FirstDerivative(f, D(f))

function update_state(method::Newton, fs, o::UnivariateZeroState{T}, options::UnivariateZeroOptions) where {T}
    xn = o.xn1
    fxn = o.fxn1
    fpxn = fs.fp(xn)

    if isissue(fpxn)
        o.stopped=true
        return
    end

    xn1 = xn - fxn / fpxn
    fxn1 = fs.f(xn1)
    incfn(o)

    o.xn0, o.xn1 = xn, xn1
    o.fxn0, o.fxn1 = fxn, fxn1

    incsteps(o)

end

## extra work to allow for complex values

## extra work to allow for complex values
function derivative_free_setup(method::Newton, fs::CallableFunction, x0::T;
                                  bracket=missing,
                                  xabstol=zero(T), xreltol=zero(T),
                                  abstol=4*eps(T), reltol=4*eps(T),
                                  maxevals=40, maxfnevals=typemax(Int),
                                  verbose::Bool=false) where {T<:AbstractFloat}
    x = float(x0)

    prob = UnivariateZeroProblem(fs, x, bracket)
    options = UnivariateZeroOptions(xabstol, xreltol, abstol, reltol,  maxevals, maxfnevals, verbose)
    prob, options
end

function derivative_free_setup(method::Newton, fs::CallableFunction, x0::Complex{T};
                                  bracket=missing,
                                  xabstol=zero(T), xreltol=zero(T),
                                  abstol=4*eps(T), reltol=4*eps(T),
                                  maxevals=40, maxfnevals=typemax(Int),
                                  verbose::Bool=false) where {T<:AbstractFloat}
    x = float(x0)
    bracket = missing     # bracket makes no sense for complex input, but one is expected

    prob = UnivariateZeroProblem(fs, x, bracket)
    options = UnivariateZeroOptions(xabstol, xreltol, abstol, reltol,  maxevals, maxfnevals, verbose)
    prob, options
end

"""

Implementation of Newton's method: `x_n1 = x_n - f(x_n)/ f'(x_n)`

Arguments:

* `f::Function` -- function to find zero of

* `fp::Function=D(f)` -- derivative of `f`. Defaults to automatic derivative

* `x0::Number` -- initial guess. For Newton's method this may be complex.

Keyword arguments:

* `ftol`. Stop iterating when |f(xn)| <= max(1, |xn|) * ftol.

* `xtol`. Stop iterating when |xn+1 - xn| <= xtol + max(1, |xn|) * xtolrel

* `xtolrel`. Stop iterating when |xn+1 - xn| <= xtol + max(1, |xn|) * xtolrel

* `maxeval`. Stop iterating if more than this many steps, throw error.

* `maxfneval`. Stop iterating if more than this many function calls, throw error.

* `verbose::Bool=false` Set to `true` to see trace.

"""
newton(f, x0; kwargs...) = find_zero(f, x0, Newton(); kwargs...)
newton(f, fp, x0; kwargs...) = find_zero((f, fp), x0, Newton(); kwargs...)

## Halley

"""

    Roots.Halley()

Implements Halley's [method](http://tinyurl.com/yd83eytb),
`x_n1 = xn - (2 f(xn)*f'(xn)) / (2 f'(xn)^2 - f(xn) * f''(xn))`.
This method is cubically converging, but requires more function calls per step than
other methods.
"""
struct Halley <: UnivariateZeroMethod
end

function callable_function(method::Halley, f::Tuple)
    length(f) == 1 && return SecondDerivative(f[1], D(f[1]), D(f[1],2))
    length(f) == 2 && return SecondDerivative(f[1], f[2], D(f[2],1))
    SecondDerivative(f[1], f[2], f[3])
end
callable_function(method::Halley, f) = SecondDerivative(f, D(f), D(f, 2))

function update_state(method::Halley, fs, o::UnivariateZeroState{T}, options::UnivariateZeroOptions) where {T}
    xn = o.xn1
    fxn = o.fxn1
    fpxn = fs.fp(xn); incfn(o)
    fppxn = fs.fpp(xn); incfn(o)

    xn1 = xn - 2fxn*fpxn / (2*fpxn*fpxn - fxn * fppxn)
    fxn1 = fs.f(xn1); incfn(o)

    o.xn0, o.xn1 = xn, xn1
    o.fxn0, o.fxn1 = fxn, fxn1
    incsteps(o)
end

"""

Implementation of Halley's method. `xn1 = xn - 2f(xn)*f'(xn) / (2*f'(xn)^2 - f(xn) * f''(xn))`

Arguments:

* `f::Function` -- function to find zero of

* `fp::Function=D(f)` -- derivative of `f`. Defaults to automatic derivative

* `fpp:Function=D(f,2)` -- second derivative of `f`.

* `x0::Real` -- initial guess

Keyword arguments:

* `ftol`. Stop iterating when |f(xn)| <= max(1, |xn|) * ftol.

* `xtol`. Stop iterating when |xn+1 - xn| <= xtol + max(1, |xn|) * xtolrel

* `xtolrel`. Stop iterating when |xn+1 - xn| <= xtol + max(1, |xn|) * xtolrel

* `maxeval`. Stop iterating if more than this many steps, throw error.

* `verbose::Bool=false` Set to `true` to see trace.

"""
halley(f,  x0; kwargs...) = find_zero(f, x0, Halley(); kwargs...)
halley(f, fp, x0; kwargs...) = find_zero((f, fp), x0, Halley(); kwargs...)
halley(f, fp, fpp, x0; kwargs...) = find_zero((f, fp, fpp), x0, Halley(); kwargs...)


text:
#!/usr/bin/env julia

export pe20

pe20(n::Integer=big(100))=factorial(n)|>digits|>sum

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe20()|>print
end

max_stars_repo_path: 20. Factorial digit sum/20.jl

max_stars_repo_name: miRoox/ProjectEuler

max_issues_repo_path: 20. Factorial digit sum/20.jl

max_issues_repo_name: miRoox/ProjectEuler

max_forks_repo_path: 20. Factorial digit sum/20.jl

max_forks_repo_name: miRoox/ProjectEuler


text:
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

using Distributed
using DiffEqParamEstim
using DifferentialEquations
using Statistics
using Random
using Optim
using JLSO
using Logging
using LoggingExtras
import StatsBase: sample
import Distributions: Uniform
import Dates: now

using Revise
includet("./problem_set.jl")
includet("./objective_function.jl")
includet("./utils.jl")
import .ProblemSet: get_problem, get_problem_key, DEProblem
import .ObjectiveFunction: data_shooting, single_shooting, tikhonov
import .Utils: rand_guess, add_noise, nmse, transpose_vector

Random.seed!(1234)

const RT = 0.75
const MAXT = 10^4
const M = 30
const MIN_LOG = Logging.Info

"""
Performs optimization and returns a vector with
information about test error (first vector position),
how much time was spent (second position) and
estimated parameters (third position)
"""
function optim_res(obj_fun::Function,
					testing_set::Vector{ProblemSet.DEProblem},
					p0::Vector{T})::Vector{Vector{T}} where T
	g_t_lim = MAXT
	f_tol = 10^-12
	x_tol = 10^-6
	iter = 10^8

	SAMIN_options = Optim.Options(x_tol=x_tol, f_tol=f_tol,
								iterations=iter, time_limit=g_t_lim)

	# Local method disabled since it does not improve accuracy
	# t_lim = 1

	# Grad_options = Optim.Options(x_tol=x_tol, f_tol=f_tol,
	# 							iterations=iter, time_limit=t_lim)

	inner_optimizer = Optim.LBFGS()

	lb = testing_set[1].bounds[1]
	ub = testing_set[1].bounds[2]
	elapsed_time = zero(T)
	phi_est = deepcopy(p0)
	nrmse = Inf64

	try
		elapsed_time += @elapsed res_obj = Optim.optimize(obj_fun,
									lb,ub,
									p0,
									Optim.SAMIN(verbosity=0, rt=RT), SAMIN_options)

		# Local method disabled since it does not improve accuracy
		#elapsed_time += @elapsed res_obj = Optim.optimize(obj_fun,
		#							lb,ub,
		#							res_obj.minimizer[1:length(p0)],
		#							Fminbox(inner_optimizer),
		#							Grad_options,
		#							autodiff = :forward)

		phi_est = res_obj.minimizer[1:length(p0)]
	catch e
		bt = backtrace()
		msg = sprint(showerror, e, bt)
		@warn """Error!
				   There was an error during the optimization step.""" msg
	end

	try
		generate_estimated_data = function (p)
			tspan = (p.t[1], p.t[end])
			ode_prob = ODEProblem(p.fun, p.data[1], tspan, phi_est)
			ode_sol  = solve(ode_prob, AutoVern7(Rodas5()), saveat=p.t)
			ode_sol.u
		end

		generate_then_transpose(p) = (p
									  |> generate_estimated_data
									  |> transpose_vector)
		data_estimated_set = map(generate_then_transpose, testing_set)

		get_then_transpose(p) = p.data |> transpose_vector
		data_testing_set = map(get_then_transpose, testing_set)

		testing_estimated_pairs = zip(data_testing_set, data_estimated_set)

		nmse_state_vectors(dataset_pairs) =  map(state_pairs -> nmse(state_pairs...),
												 zip(dataset_pairs...))

		total_nmse = map(nmse_state_vectors, testing_estimated_pairs)
		# (mean) Normalized Root Mean Squared Error
		nrmse = mean(total_nmse) |> mean |> sqrt
		@debug("NRMSE",
			   Symbol(testing_set[1].fun),
			   phi_est,
			   nrmse,
			   data_estimated_set,
			   data_testing_set,
			   testing_estimated_pairs)
	catch e
		bt = backtrace()
		msg = sprint(showerror, e, bt)
		@warn """Error!
				   There was an error during the NRMSE calculation step.""" msg
	end

	if elapsed_time == zero(T)
		elapsed_time == NaN64
	end

	return [[nrmse], [elapsed_time], phi_est]
end

function cv_optimize(training_set::Vector{ProblemSet.DEProblem},
					testing_set::Vector{ProblemSet.DEProblem},
					p0::Vector{T},
					k::Int64,
					f::Function,
					phi_prior::Vector{T},
					m::Int64,
					w::Array)::Array{Array{T}} where T

	results = [[Inf64], [NaN64], p0]
	#lambda_arr = [1e-2,1e-1,0.0,1e0,1e1,1e2]
	lambda_arr = [0.0,1e-3,1e-2,1e-1,1e0]
	num_lambdas = length(lambda_arr)
	best_lambda = lambda_arr[1]

	bounds = training_set[1].bounds

	elapsed_time = zero(T)

	# Lamdba selection via
	# Cross Validation
	@debug """CV
				Initial guess:\n$(p0)
				Starting Cross Validation!
				Preprocessing starting
				..."""

	if sum(w) != 0.0
		# K-Fold Cross Validation

		num_training_sets = length(training_set)
		num_elem_folds = convert(Int64, (num_training_sets/k))

		cv_folds = sample(1:num_training_sets, (k,num_elem_folds), replace=false)

		num_elem_train_slice = convert(Int64, (num_training_sets - num_elem_folds))

		# Pre - Allocate and prepare necessary data
		partial_res = [[Inf64], [NaN64], p0]
		partial_res_arr = [[[Vector{Float64}(undef, 1)
							 for _ in 1:3] for _ in 1:k]
						   for _ in 1:num_lambdas]

		training_set_folds = [Vector{ProblemSet.DEProblem}(undef,num_elem_train_slice)
	 						  for _ in 1:k]
		training_set_arr = [[Vector{ProblemSet.DEProblem}(undef,num_elem_train_slice)
	 						 for _ in 1:k]
							for _ in 1:num_lambdas]

		hold_out_set_folds = [Vector{ProblemSet.DEProblem}(undef,num_elem_folds)
							  for _ in 1:k]
		hold_out_set_arr = [[Vector{ProblemSet.DEProblem}(undef,num_elem_folds)
							 for _ in 1:k]
							for _ in 1:num_lambdas]

		obj_fun_folds = Vector{Function}(undef,k)
		obj_fun_arr = [Vector{Function}(undef,k) for _ in 1:num_lambdas]

		for i in 1:num_lambdas
			for j in 1:k
				hold_out = cv_folds[j,:]
				slice = filter(x -> !(x in hold_out), 1:num_training_sets)

				training_sliced = training_set[slice]
				training_set_folds[j] = training_sliced

				hold_out_set = training_set[hold_out]
				hold_out_set_folds[j] = hold_out_set

				obj_fun = function (x)
					total_error = zero(T)
					@inbounds for i in 1:length(training_sliced)
						total_error += f(x,
										 training_sliced[i].data,
										 training_sliced[i].t,
										 training_sliced[i].fun)
					end
					total_error += tikhonov(lambda_arr[i], x, phi_prior, w)
					return total_error
				end
				obj_fun_folds[j] = obj_fun
			end
			obj_fun_arr[i] = obj_fun_folds
			training_set_arr[i] = training_set_folds
			hold_out_set_arr[i] = hold_out_set_folds
		end

		@debug """Preprocessing done!
					Async CV starting
					..."""

		# Async CV
		@sync for i in 1:num_lambdas
			for j in 1:k
				Threads.@spawn partial_res_arr[i][j] = optim_res(obj_fun_arr[i][j],
																 hold_out_set_arr[i][j],
																 p0)
			end
		end

		# Post - Process results
		@debug """Async CV done for $(method)!
					Evaluating results:"""
		fold_error_mean = zero(T)
		lambda_hold_errors = Vector{Float64}(undef,num_lambdas)
		for i in 1:num_lambdas
			fold_error_mean = 0.0
			@debug """Results for $(lambda_arr[i])"""
			for j in 1:k
				pres = partial_res_arr[i][j]
				fold_error_mean += pres[1][1]
				elapsed_time += pres[2][1]
				@debug """Error $(j)""" pres[1][1]
			end
			fold_error_mean /= k
			@debug """Mean error on $(lambda_arr[i])""" fold_error_mean
			lambda_hold_errors[i] = fold_error_mean
		end

		elapsed_time /= num_lambdas*k
		best_lambda = lambda_arr[argmin(lambda_hold_errors)]

		@debug """Done!
					Best lambda:\n$(best_lambda)
					Optimizing on whole dataset.
					..."""
	end

	final_obj_fun = function (x)
		total_error = zero(T)
		@inbounds for i in 1:length(training_set)
			total_error += f(x,
							 training_set[i].data,
							 training_set[i].t,
							 training_set[i].fun)
		end
		total_error += tikhonov(best_lambda, x, phi_prior, w)
		return total_error
	end

	try
		partial_res = optim_res(final_obj_fun, testing_set, p0)
		results = [partial_res[1],
				   # [(partial_res[2][1]+elapsed_time)/2.0],
				   partial_res[2],
				   partial_res[3]]

		@debug """Done!
					Test Error:\n$(results[1])
					Estimated parameters:\n$(results[3])\n"""
	catch e
		bt = backtrace()
		msg = sprint(showerror, e, bt)
		@warn """Error!
				   There was an error during the call to optim_res.""" msg
	end

	@info """Results.
			  Problem:\t$(Symbol(training_set[1].fun))
			  Method:\t$(Symbol(f))
			  Lambdas:\t$(lambda_arr)
			  Initial phi:\t$(p0)
			  CV Errors:\t$(@isdefined(lambda_hold_errors) ? lambda_hold_errors : false)
			  Best lambda:\t$(best_lambda)
			  Final Errors:\t$(results[1][1])
			  Final phi:\t$(results[3])"""
	return results
end

function get_results(method::String,
					training_set::Vector{ProblemSet.DEProblem},
					testing_set::Vector{ProblemSet.DEProblem},
	 				p0::Vector{T})::Array{Array{T}} where T

	k = convert(Int64, length(training_set)/1)

	add_time = zero(T)

	if method == "SS"
		f = single_shooting
		phi_prior = zeros(T,length(p0))
		m = length(phi_prior)
		w = zeros(T,m)
	elseif method == "SSR"
		f = single_shooting
		phi_prior = zeros(T,length(p0))
		m = length(phi_prior)
		w = ones(T,m)

		partial_res = cv_optimize(training_set,testing_set,
										p0,k,f,phi_prior,m,w)
		add_time = partial_res[2][1]
		phi_prior = partial_res[3]
		p0 = phi_prior
	elseif method == "DS"
		f = data_shooting
		phi_prior = zeros(T,length(p0))
		m = length(phi_prior)
		w = zeros(T,m)
	elseif method == "DSS"
		f = data_shooting
		phi_prior = zeros(T,length(p0))
		m = length(phi_prior)
		w = ones(T,m)
		partial_res = cv_optimize(training_set,testing_set,
										p0,k,f,phi_prior,m,w)
		add_time = partial_res[2][1]

		f = single_shooting
		phi_prior = partial_res[3]
		p0 = phi_prior

		# Check the model error.
		# If the model is good, then use the
		# best case strategy.
		# Otherwise, use medium case strategy.
		if partial_res[1][1] < 0.3
			w = ones(T,m)./phi_prior
		else
			w = ones(T,m)
		end
	end

	results = cv_optimize(training_set,testing_set,
							p0,k,f,phi_prior,m,w)
	results[2][1] += add_time

	return results
end

function experiment(p_num::Int64,samples::AbstractArray{<:Int},
					noise_levels::AbstractArray{<:AbstractFloat},
					methods::Array{<:String},
					dir::String)::Dict

	all_methods = ["DS","SS","SSR","DSS"]

	results::Dict = Dict()
	prob_key::String = get_problem_key(p_num)
	problem::ProblemSet.DEProblem = get_problem(prob_key)

	fun::Function = problem.fun
	phi::Array = problem.phi
	bounds::Vector = problem.bounds
	ini_cond::Array = problem.data[1]
	t::AbstractArray = problem.t

	# Minimum number of data points
	#min_data = round(length(phi)/length(ini_cond),RoundUp)
	#data_sams = convert.(eltype(sams[1]),sams.*min_data)
	data_sams = samples

	num_reps = M

	results_final = Dict()
	for sam in data_sams
		results_final[Symbol(sam)] = Dict()
		for noise in noise_levels
			results_final[Symbol(sam)][Symbol(noise)] = Dict()
			for m in methods
				results_final[Symbol(sam)][Symbol(noise)][Symbol(m)] =
					[[[NaN64] for _ in 1:3]
					for _ in 1:num_reps]
			end
		end
	end

	num_trainings = 5
	num_tests = 10

	training_set_dict = Dict()
	testing_set_dict = Dict()
	guess_dict = Dict()
	for sam in data_sams
		training_set_dict[sam] = Dict()
		testing_set_dict[sam] = Dict()
		guess_dict[sam] = Dict()
		for noise in noise_levels
			training_set_dict[sam][noise] = [Vector{ProblemSet.DEProblem}(undef,num_trainings)
											for _ in 1:num_reps]
			testing_set_dict[sam][noise] = [Vector{ProblemSet.DEProblem}(undef,num_tests)
											for _ in 1:num_reps]
			guess_dict[sam][noise] = [Vector{Float64}(undef,length(phi))
									for _ in 1:num_reps]
		end
	end

	var_ini_cond = 0.01

	for sam in data_sams
		for noise in noise_levels
			# Artificial Data
			_t = range(t[1], stop=t[end], length=sam)
			tspan = (t[1], t[end])
			ode_prob = ODEProblem(fun, ini_cond, tspan, phi)
			ode_sol  = solve(ode_prob, AutoVern7(Rodas5()), saveat=_t)
			data = ode_sol.u
			#data_plot = plot(t,data')
			#display(data_plot)

			guess_arr = [rand_guess(bounds) for _ in 1:num_reps]
			guess_dict[sam][noise] .= guess_arr

			training_set_arr = [[ProblemSet.DEProblem(problem.fun, problem.phi,
								problem.bounds, add_noise(data,noise), _t)
				 			for _ in 1:num_trainings] for _ in 1:num_reps]
			training_set_dict[sam][noise] .= training_set_arr

			testing_set_arr = []
			for r in 1:num_reps
				testing_set_arr_partial = []
				for j in 1:num_tests
					_t = range(t[1], stop=t[end], length=sam)
					tspan = (t[1], t[end])
					ode_prob = ODEProblem(fun, add_noise(ini_cond,var_ini_cond), tspan, phi)
					ode_sol  = solve(ode_prob, AutoVern7(Rodas5()), saveat=_t)
					data = ode_sol.u

					push!(testing_set_arr_partial,ProblemSet.DEProblem(problem.fun, problem.phi,
										problem.bounds, add_noise(data,noise), _t))
				end
				push!(testing_set_arr,testing_set_arr_partial)
			end
			testing_set_dict[sam][noise] .= testing_set_arr
		end
	end #samples loop

	@sync for sam in data_sams
		for noise in noise_levels
			for m in methods
				for i in 1:num_reps
					Threads.@spawn results_final[Symbol(sam)][Symbol(noise)][Symbol(m)][i] =
											get_results(m,
											training_set_dict[sam][noise][i],
											testing_set_dict[sam][noise][i],
											guess_dict[sam][noise][i])

				end
			end
		end
	end

	return results_final
end

function run_experiments(problems::Vector{<:Int},
							samples::Vector{<:Int},
							noise_levels::Vector{<:AbstractFloat},
							methods::Vector{String},
							dir::String)::Nothing

	# Temporary directory for saving
	# serialized intermediary results
	mkdir(joinpath(dir,"tmp"))

	results = Vector(undef,length(problems))
	for i in 1:length(problems)
		p = problems[i]
		p_name = get_problem_key(p)
		result = experiment(p,samples,noise_levels,methods,dir)
		results[i] = Pair(Symbol(p_name), result)
		try
			JLSO.save(joinpath(dir,"tmp",p_name*".jlso"), results[i])
		catch e
			bt = backtrace()
			msg = sprint(showerror, e, bt)
			@warn """Error!
						Error during the call to JLSO.save for $(results[i])."""
		end
	end

	try
		JLSO.save(joinpath(dir,"experiment_results.jlso"), results...)
		rm(joinpath(dir,"tmp"), recursive=true)
	catch e
		bt = backtrace()
		msg = sprint(showerror, e, bt)
		@warn """Error!
				   Error when joining all the results.""" msg
	end

end

function main(args::Array{<:String})::Nothing
	path = string(args[1])
	problems = eval(Meta.parse(args[2]))
	samples = eval(Meta.parse(args[3]))
	noise_level = eval(Meta.parse(args[4]))

	methods = ["DS","SS","SSR","DSS"]

	main_file_logger = MinLevelLogger(FileLogger(joinpath(path,"general.log")),
									  MIN_LOG)
	console_logger = ConsoleLogger()
	main_logger = TeeLogger(console_logger, main_file_logger)
	global_logger(main_logger)

	@info "Program started." now()
	@info "Number of threads." Threads.nthreads()

	minor_logger = FileLogger(joinpath(path,"info.log"))
	minor_and_global_logger = TeeLogger(minor_logger, global_logger())

	with_logger(minor_and_global_logger) do
		@info "Arguments." args
		@info "Methods." methods
		@info "Cooling rate." RT
		@info "Maximum Time." MAXT
		@info "Number of runs." M
	end

	elapsed_time = @elapsed run_experiments(problems,samples,
										noise_level,methods,path)

	with_logger(minor_and_global_logger) do
		@info "Program finished." now()
		@info "Elapsed time." elapsed_time
	end

end

main(ARGS)

max_stars_repo_name: atgmello/DEParamEst

max_issues_repo_name: atgmello/DEParamEst

max_forks_repo_name: atgmello/DEParamEst


text:
using GPMTD
using LinearAlgebra
using SparseProbVec
using PDMats
using Random123
using Statistics

using BSON
using BSON: @load
whichdat = "postsim_Ricker_lag2only_n50_L2_427191"
BSON.@load "postsim/$(whichdat).bson" model sims init mesg
# BSON.@load "postsimB/$(whichdat).bson" model sims init mesg

using Plotly
using StatsBase

nsim = length(sims)

iter_use = 1:nsim
iter_use = range(1, nsim, step=2)
iter_use = 1900:nsim

# counts of mixture component membership
sims_Scounts = permutedims( hcat( [ StatsBase.counts( sims[i][:ζ], 0:model.L ) for i in iter_use ]... ) )
mean(sims_Scounts ./ model.n, dims=1)

tmp = hcat( [ sims[ii][:ζ] for ii in iter_use ]... )
sims_Scount_obs = permutedims(hcat([ StatsBase.counts( tmp[i,:], 0:model.L ) for i = 1:model.n ]...) ./ length(iter_use))

# log-likelihood
sims_llik = [ sims[i][:llik] for i in iter_use ]
plot(sims_llik)

## parameter inferences
j = 1 # 0 for intercept
lam = [ exp(sims[i][:lλ][j+1]) for i in iter_use ]
plot(lam)

for j = 0:model.L
    lam = [ exp(sims[i][:lλ][j+1]) for i in iter_use ]
    println("lam $(j) mean ", round(mean(lam), digits=3))
    println("lam $(j) q025 ", round(quantile(lam, 0.025), digits=4))
    println("lam $(j) q975 ", round(quantile(lam, 0.975), digits=4), "\n")
end

int_mu = [ sims[i][:intercept].μ for i in iter_use ]
plot(int_mu)

int_sig2 = [ sims[i][:intercept].σ2 for i in iter_use ]
plot(int_sig2)

j = 1
mixc_mu = [ sims[i][:mixcomps][j][:μ] for i in iter_use ]
mean(mixc_mu)
std(mixc_mu)
plot(mixc_mu)

mixc_sig2 = [ sims[i][:mixcomps][j][:σ2] for i in iter_use ]
plot(mixc_sig2)
mean(mixc_sig2)
plot(sqrt.(mixc_sig2))
mean(sqrt.(mixc_sig2))

j = 1
mixc_kap = [ sims[i][:mixcomps][j][:κ] for i in iter_use ]
plot(mixc_kap)

mixc_ls = [ sims[i][:mixcomps][j][:corParams].lenscale for i in iter_use ]
plot(mixc_ls)

covstat = [ sims[i][:mixcomps][j][:κ]*sims[i][:mixcomps][j][:σ2]/(sims[i][:mixcomps][j][:corParams].lenscale^2.5) for i in iter_use ]
plot(covstat)

j = 2
kap_df = [ sims[i][:mixcomps][j][:κ_hypers].κ_ν for i in iter_use ]
plot(kap_df)

j = 2
kap0 = [ sims[i][:mixcomps][j][:κ_hypers].κ0 for i in iter_use ]
plot(kap0)

j = 2
ls_df = [ sims[i][:mixcomps][j][:corHypers].lenscale_ν for i in iter_use ]
plot(ls_df)

j = 2
ls0 = [ sims[i][:mixcomps][j][:corHypers].lenscale0 for i in iter_use ]
plot(ls0)

j = 2
ord = sortperm(model.X[:,j])
F = hcat([ sims[i][:mixcomps][j][:fx] .+ sims[i][:mixcomps][j][:μ] for i in iter_use ]...)
Fmean = mean(F, dims=2)[:,1]
Fq025 = [ quantile(F[i,:], 0.025) for i = 1:size(F,1)]
Fq975 = [ quantile(F[i,:], 0.975) for i = 1:size(F,1)]

trace_mean = scatter(Dict("x" => model.X[ord,j], "y" => Fmean[ord], "name" => "mean",
    "mode" => "both", "type" => "scatter", "marker" => Dict("color" => "black", "line" => Dict("color" => "black"))))
trace_05 = scatter(Dict("x" => model.X[ord,j], "y" => Fq025[ord], "name" => "Q025",
    "mode" => "both", "type" => "scatter", "marker" => Dict("color" => "red", "line" => Dict("color" => "red"))))
trace_95 = scatter(Dict("x" => model.X[ord,j], "y" => Fq975[ord], "name" => "Q975",
    "mode" => "both", "type" => "scatter", "marker" => Dict("color" => "orange", "line" => Dict("color" => "orange"))))

plot([trace_mean, trace_05, trace_95])

plot(scatter(Dict("x" => X[ord,j], "y" => F[ord,50],
    "mode" => "both", "type" => "scatter")))

### inferences on a grid of x values

include("Rfuns.jl")

n_star = 100
minx = minimum(model.X)
maxx = maximum(model.X)
ranx = maxx - minx

X_star_vals = collect(range(minx - 0.1*ranx, length=n_star,
                        stop=(maxx + 0.1*ranx)))
# X_star = rcopy(xpndgrid([X_star_vals for k = 1:model.K]...))
X_star_vals = collect(range(2.0, length=n_star,
                        stop=11.5))

γindx = [2]

# iter_use = sort(findall([ sims[ii][:γ] == γ_use for ii in 1:length(sims)]))
iter_use = 1:nsim
# iter_use = 1900:2000
iter_use = range(1, nsim, step=2)

Kuse = length(γindx)
X_star00 = rcopy(xpndgrid([X_star_vals for k = 1:Kuse]...))

# X_star0 = zeros(Float64, size(X_star00,1), model.K)
X_star0 = fill(mean(model.y), size(X_star00,1), model.L)
# X_star0 = rand(size(X_star00,1), model.L) .* ranx .+ minx # draw uniformly over the range

X_star0[:,γindx] = deepcopy(X_star00)
X_star = deepcopy(X_star0)

# transition mean function
@time Ey = getEyPy(sims[iter_use], X_star, model.X, model.D)

mean_Ey = mean(Ey, dims=2)[:,1]

q025_Ey = [quantile(Ey[jj,:], 0.025) for jj = 1:size(Ey,1)]
q975_Ey = [quantile(Ey[jj,:], 0.975) for jj = 1:size(Ey,1)]

## one dimension
size(Ey)
plotR(X_star_vals, mean_Ey, ylab="", xlab="",
    ylim=[minimum(q025_Ey), maximum(q975_Ey)], lwd=2, typ="l")
plotR(X_star_vals, mean_Ey, ylab="", xlab="",
    ylim=[0.0, 12.3], lwd=2, typ="l")
plotR(X_star_vals, mean_Ey, ylab="", xlab="",
    ylim=[minimum(model.y)-0.05*ranx, maximum(model.y)+0.05*ranx], lwd=2, typ="l")
polygonR(vcat(X_star_vals, reverse(X_star_vals)), vcat(q025_Ey, reverse(q975_Ey)), col="gray85", border="NA")
pointsR(model.X[:,γindx[1]], model.y)

linesR(X_star_vals, mean_Ey, lwd=2)

R"title(ylab='y')"
R"title(xlab='y lag')"

R"dev.off()"

### two dimensions
# if extreme values mess up graph
# mean_Ey[findall(mean_Ey .< -1.0)] *= 0.0

# hcat(X_star[:,1], X_star[:,2], mean_Ey)
mean_Ey_mat = Matrix(reshape(mean_Ey, fill(n_star,Kuse)...))

q025_Ey_mat = Matrix(reshape(q025_Ey, fill(n_star,Kuse)...))
q975_Ey_mat = Matrix(reshape(q975_Ey, fill(n_star,Kuse)...))

trace1 = surface(Dict(
  :x => X_star_vals,
  :y => X_star_vals,
  :z => q025_Ey_mat,
  :colorscale => "Viridis",
  :opacity => 0.7,
  :showscale => false,
  :type => "surface"
))

trace2 = surface(Dict(
  :x => X_star_vals,
  :y => X_star_vals,
  :z => mean_Ey_mat,
  # :colorbar => Dict(
  #   :ticklen => 2,
  #   :title => ""
  # ),
  :colorscale => "Viridis",
  :showscale => false,
  :type => "surface"
))

trace3 = surface(Dict(
  :x => X_star_vals,
  :y => X_star_vals,
  :z => q975_Ey_mat,
  :colorscale => "Viridis",
  :opacity => 0.7,
  :showscale => false,
  :type => "surface"
))

trace0 = scatter3d(Dict(
    :x => model.X[:,γindx[1]],
    :y => model.X[:,γindx[2]],
    :z => model.y,
    :marker => Dict(
        :color => "rgba(0,0,0,0.70)",
        :size => 2
    ),
    :mode => "markers"
))

tracetrue = surface(Dict(
  :x => X_star_vals,
  :y => X_star_vals,
  :z => truemean_mat,
  # :colorscale => "Viridis",
  :opacity => 0.7,
  :showscale => false,
  :type => "surface"
))

data = [trace2, trace0]
data = [trace1, trace2, trace3, trace0#=, tracetest=#]
data = [tracetrue, trace2, trace0]
data = [trace0]

layout = Layout(Dict(
  :hovermode => "closest",
  :margin => Dict(
    :r => 10,
    :t => 25,
    :b => 40,
    :l => 60
  ),
  :scene => Dict(
    :xaxis => Dict(:title => "log y[t-$(γindx[1])]",
                   :range => [-1.5, 3.0] #=,
                   :type => "log" #=,
                   :autorange => true =# =#),
    :yaxis => Dict(:title => "log y[t-$(γindx[2])]",
                    :range => [-1.5, 3.0] #=,
                    :type => "log" #=,
                    :autorange => true =# =#),
    :zaxis => Dict(:title => "log y[t]",
                    :range => [-1.5, 3.0] #=,
                    :type => "log" #=,
                    :autorange => true =# =#),
  ),
  :showlegend => false
))

p = plot(data, layout)

plot(data)

## transition density
n_y = 150
y_grid = collect(range(minx - 0.3*ranx, length=n_y,
                        stop=(maxx + 0.3*ranx)))

# number of distinct x locations at which to evaluate
n_star_Py = 3

# x values
X_star_Py = deepcopy(X_star[1:n_star_Py,:])
X_star_Py[1, γindx[1]] = 2.0
X_star_Py[2, γindx[1]] = 5.0
X_star_Py[3, γindx[1]] = 7.0

X_star_Py = deepcopy(X_star[1:n_star_Py,:])
X_star_Py[1, γindx[1]] = 50.0
X_star_Py[2, γindx[1]] = 66.0
X_star_Py[3, γindx[1]] = 80.0

X_star_Py = deepcopy(X_star[1:n_star_Py,:])
X_star_Py[1, γindx[1]:γindx[2]] = [10.0, 2.0]
X_star_Py[2, γindx[1]:γindx[2]] = [1.0, 1.0]
X_star_Py[3, γindx[1]:γindx[2]] = [2.0, 7.0]

X_star_Py = deepcopy(X_star[1:n_star_Py,:])
X_star_Py[1, γindx[1]:γindx[2]] = [2.3, 0.8]
X_star_Py[2, γindx[1]:γindx[2]] = [1.0, -0.5]
X_star_Py[3, γindx[1]:γindx[2]] = [1.0, 1.6]
X_star_Py[4, γindx[1]:γindx[2]] = [1.7, 0.7]
X_star_Py[5, γindx[1]:γindx[2]] = [-0.75, 2.3]

X_star_Py

# transition mean and densities
Ey0, Py = getEyPy(sims[iter_use], X_star_Py,
            model.X, model.D; densout=true, y_grid=y_grid)

mean_Py = mean(Py, dims=3)[:,:,1]
q025_Py = [quantile(Py[i,j,:], 0.025) for i = 1:n_y, j = 1:n_star_Py]
# p50_Py = [quantile(Py[i,j,:], 0.5) for i = 1:n_y, j = 1:n_star_Py]
q975_Py = [quantile(Py[i,j,:], 0.975) for i = 1:n_y, j = 1:n_star_Py]

xind = 2
X_star_Py[xind,:]

R"par(mar=c(4,4,2,1)+0.1)"
plotR(y_grid, mean_Py[:,xind],
    main="y[t-1] = $(round(X_star_Py[xind,γindx[1]], digits=2))",
    ylab="density", xlab="y[t] (minutes)", typ="l", ylim=[0.0, maximum(q975_Py[:,xind])], lty=1)
polygonR(vcat(y_grid, reverse(y_grid)), vcat(q025_Py[:,xind], reverse(q975_Py[:,xind])), col="gray90", border="NA")
linesR(y_grid, mean_Py[:,xind], lwd=2)
R"dev.off()"

R"par(mar=c(4,4,2,1)+0.1)"
plotR(y_grid, mean_Py[:,xind],
    main="y[t-1] = $(round(X_star_Py[xind,γindx[1]], digits=2))",
    ylab="density", xlab="y[t] (minutes)", typ="l", ylim=[0.0, maximum(q975_Py[:,xind])], lty=1)
polygonR(vcat(y_grid, reverse(y_grid)), vcat(q025_Py[:,xind], reverse(q975_Py[:,xind])), col="gray90", border="NA")
linesR(y_grid, mean_Py[:,xind], lwd=2)
R"dev.off()"

R"par(mar=c(4,4,2,1)+0.1)"
plotR(y_grid, mean_Py[:,xind],
    main="y[t-1] = $(round(X_star_Py[xind,γindx[1]], digits=2)), y[t-2] = $(round(X_star_Py[xind,γindx[2]], digits=2))",
    ylab="density", xlab="y[t] (minutes)", typ="l", ylim=[0.0, maximum(q975_Py[:,xind])], lty=1)
polygonR(vcat(y_grid, reverse(y_grid)), vcat(q025_Py[:,xind], reverse(q975_Py[:,xind])), col="gray90", border="NA")
linesR(y_grid, mean_Py[:,xind], lwd=2)
R"dev.off()"


text:
export QRDQNLearner, quantile_huber_loss

using Zygote: ignore, dropgrad
using Random: GLOBAL_RNG, AbstractRNG
using StatsBase: mean

function quantile_huber_loss(ŷ, y; κ=1.0f0)
    N, B = size(y)
    Δ = reshape(y, N, 1, B) .- reshape(ŷ, 1, N, B)
    abs_error = abs.(Δ)
    quadratic = min.(abs_error, κ)
    linear = abs_error .- quadratic
    huber_loss = 0.5f0 .* quadratic .* quadratic .+ κ .* linear

    cum_prob = ignore() do
        send_to_device(device(y), range(0.5f0 / N; length=N, step=1.0f0 / N))
    end
    loss = dropgrad(abs.(cum_prob .- (Δ .< 0))) .* huber_loss
    mean(sum(loss; dims=1))
end

Base.@kwdef mutable struct QRDQNLearner{A<:Approximator{<:TwinNetwork}} <: AbstractLearner
    approximator::A
    n_quantile::Int
    loss_func::Any = quantile_huber_loss
    γ::Float32 = 0.99f0
    rng::AbstractRNG = GLOBAL_RNG
    # for recording
    loss::Float32 = 0.0f0
end

Functors.functor(x::QRDQNLearner) = (; approximator=x.approximator), y -> @set x.approximator = y.approximator

(L::QRDQNLearner)(s::AbstractArray) = vec(mean(reshape(L.approximator(s), L.n_quantile, :), dims=1))

function RLBase.optimise!(learner::QRDQNLearner, batch::NamedTuple)
    A = learner.approximator
    Q = A.model.source
    Qₜ = A.model.target
    γ = learner.γ
    N = learner.n_quantile
    loss_func = learner.loss_func

    s, s′, a, r, t = map(x -> batch[x], SS′ART)
    batch_size = length(r)
    a = CartesianIndex.(a, 1:batch_size)

    target_quantiles = reshape(Qₜ(s′), N, :, batch_size)
    qₜ = dropdims(mean(target_quantiles; dims=1); dims=1)
    aₜ = dropdims(argmax(qₜ, dims=1); dims=1)
    @views target_quantile_aₜ = target_quantiles[:, aₜ]
    y = reshape(r, 1, batch_size) .+ γ .* reshape(1 .- t, 1, batch_size) .* target_quantile_aₜ

    gs = gradient(params(A)) do
        q = reshape(Q(s), N, :, batch_size)
        @views ŷ = q[:, a]

        loss = loss_func(ŷ, y)

        ignore() do
            learner.loss = loss
        end
        loss
    end

    optimise!(A, gs)
end


text:
struct ArrayReference
    array::Symbol
    ref::Vector{Union{Symbol,Int}}
    loaded::Base.RefValue{Bool}
    ptr::Symbol
end
function ArrayReference(array::Symbol, ref, loaded)
    ArrayReference(
        array, ref, loaded,
        Symbol("##vptr##_", array)
    )
end
ArrayReference(array::Symbol, ref) = ArrayReference(array, ref, Ref{Bool}(false))
function ArrayReference(
    array::Symbol,
    ref::AbstractVector
)
    ArrayReference(array, ref, Ref{Bool}(false))
end
function Base.hash(x::ArrayReference, h::UInt)
    @inbounds for n ∈ eachindex(x)
        h = hash(x.ref[n], h)
    end
    hash(x.array, h)
end
function loopdependencies(ref::ArrayReference)
    ld = Symbol[]
    for r ∈ ref.ref
        r isa Symbol && push!(ld, r)
    end
    ld
end
function Base.isequal(x::ArrayReference, y::ArrayReference)
    x.array === y.array || return false
    nrefs = length(x.ref)
    nrefs == length(y.ref) || return false
    all(n -> x.ref[n] === y.ref[n], 1:nrefs)
    # for n ∈ 1:nrefs
        # x.ref[n] === y.ref[n] || return false
    # end
    # true
end

Base.:(==)(x::ArrayReference, y::ArrayReference) = isequal(x, y)

function ref_from_expr(ex, offset1::Int = 0, offset2 = 0)
    ArrayReference( ex.args[1 + offset1], @view(ex.args[2 + offset2:end]), Ref(false) )
end
ref_from_ref(ex::Expr) = ref_from_expr(ex, 0, 0)
ref_from_getindex(ex::Expr) = ref_from_expr(ex, 1, 1)
ref_from_setindex(ex::Expr) = ref_from_expr(ex, 1, 2)
function ArrayReference(ex::Expr)
    ex.head === :ref ? ref_from_ref(ex) : ref_from_getindex(ex)
end
function Base.:(==)(x::ArrayReference, y::Expr)
    if y.head === :ref
        isequal(x, ref_from_ref(y))
    elseif y.head === :call && first(y.args) === :getindex
        isequal(x, ref_from_getindex(y))
    else
        false
    end
end
Base.:(==)(x::ArrayReference, y) = false

# Avoid memory allocations by accessing this
const NOTAREFERENCE = ArrayReference(Symbol(""), Union{Symbol,Int}[], Ref(false))

@enum OperationType begin
    constant
    memload
    compute
    memstore
end

# const ID = Threads.Atomic{UInt}(0)

# TODO: can some computations be cached in the operations?
"""
"""
struct Operation
    identifier::Int
    variable::Symbol
    elementbytes::Int
    instruction::Instruction
    node_type::OperationType
    dependencies::Vector{Symbol}
    reduced_deps::Vector{Symbol}
    parents::Vector{Operation}
    ref::ArrayReference
    mangledvariable::Symbol
    function Operation(
        identifier::Int,
        variable,
        elementbytes,
        instruction,
        node_type,
        dependencies = Symbol[],
        reduced_deps = Symbol[],
        parents = Operation[],
        ref::ArrayReference = NOTAREFERENCE
    )
        new(
            identifier, variable, elementbytes, instruction, node_type,
            convert(Vector{Symbol},dependencies),
            convert(Vector{Symbol},reduced_deps),
            convert(Vector{Operation},parents),
            ref,
            Symbol("##", variable, :_)
        )
    end
end

 # negligible save on allocations for operations that don't need these (eg, constants).
const NODEPENDENCY = Union{Symbol,Int}[]
const NOPARENTS = Operation[]

function Base.show(io::IO, op::Operation)
    if isconstant(op)
        if op.instruction === LOOPCONSTANT
            print(io, Expr(:(=), op.variable, 0))
        else
            print(io, Expr(:(=), op.variable, op.instruction.instr))
        end
    elseif isload(op)
        print(io, Expr(:(=), op.variable, Expr(:ref, op.ref.array, op.ref.ref...)))
    elseif iscompute(op)
        print(io, Expr(:(=), op.variable, Expr(op.instruction, name.(parents(op))...)))
    elseif isstore(op)
        print(io, Expr(:(=), Expr(:ref, op.ref.array, op.ref.ref...), name(first(parents(op)))))
    end
end

function isreduction(op::Operation)
    ((op.node_type == compute) || (op.node_type == memstore)) && length(op.reduced_deps) > 0
    # (op.node_type == memstore) && (length(op.symbolic_metadata) < length(op.dependencies))# && issubset(op.symbolic_metadata, op.dependencies)
end
isload(op::Operation) = op.node_type == memload
iscompute(op::Operation) = op.node_type == compute
isstore(op::Operation) = op.node_type == memstore
isconstant(op::Operation) = op.node_type == constant
accesses_memory(op::Operation) = isload(op) | isstore(op)
elsize(op::Operation) = op.elementbytes
dependson(op::Operation, sym::Symbol) = sym ∈ op.dependencies
parents(op::Operation) = op.parents
# children(op::Operation) = op.children
loopdependencies(op::Operation) = op.dependencies
reduceddependencies(op::Operation) = op.reduced_deps
identifier(op::Operation) = op.identifier + 1
name(op::Operation) = op.variable
instruction(op::Operation) = op.instruction

refname(op::Operation) = op.ref.ptr
"""
    mvar = mangledvar(op)

Returns the mangled variable name, for use in the produced expressions.
These names will be further processed if op is tiled and/or unrolled.

```julia
    if tiled ∈ loopdependencies(op) # `suffix` is tilenumber
        mvar = Symbol(op, suffix, :_)
    end
    if unrolled ∈ loopdependencies(op) # `u` is unroll number
        mvar = Symbol(op, u)
    end
```
"""
mangledvar(op::Operation) = op.mangledvariable

"""
Returns `0` if the op is the declaration of the constant outerreduction variable.
Returns `n`, where `n` is the constant declarations's index among parents(op), if op is an outter reduction.
Returns `-1` if not an outerreduction.
"""
function isouterreduction(op::Operation)
    if isconstant(op) # equivalent to checking if length(loopdependencies(op)) == 0
        op.instruction === LOOPCONSTANT ? 0 : -1
    elseif iscompute(op)
        var = op.variable
        for (n,opp) ∈ enumerate(parents(op))
            opp.variable === var && opp.instruction === LOOPCONSTANT && return n
        end
        -1
    else
        -1
    end
end

# function hasintersection(s1::Set{T}, s2::Set{T}) where {T}
    # for x ∈ s1
        # x ∈ s2 && return true
    # end
    # false
# end

# function symposition(op::Operation, sym::Symbol)
    # findfirst(s -> s === sym, op.symbolic_metadata)
# end
# function stride(op::Operation, sym::Symbol)
    # @assert accesses_memory(op) "This operation does not access memory!"
    # # access stride info?
    # op.numerical_metadata[symposition(op,sym)]
# end


text:
timestep=1.0e-4
gravity=-9.81

T=Float64

# Parameters
height = 1.0
num_balls = 2
rod_separation = 0.1

mass_connector = 1.0
length_connector = (num_balls - 1) * rod_separation
radius_connector = 0.025

mass_rod = 0.001
length_rod = 0.5
radius_rod = 0.01

mass_ball = 1.0
radius_ball = 0.5 * rod_separation

# Bodies
origin = Origin{T}()
connector = Capsule(radius_connector, length_connector, mass_connector,
    axis_offset=Quaternion(0.707107, 0.707107, 0.0,  0.0),
    name=:connector)

rods = [Capsule(radius_rod, length_rod, mass_rod,
    name=Symbol("rod_$i")) for i = 1:num_balls]
balls = [Sphere(radius_ball, mass_ball,
    name=Symbol("ball_$i")) for i = 1:num_balls]

# Joints
joint_axis = [1.0; 0.0; 0.0]
damper = 0.0
joint_world_connector = JointConstraint(Fixed(origin, connector,
    parent_vertex=[0.0; 0.0; height]))

joint_connector_rods = [JointConstraint(Revolute(connector, rod, joint_axis;
    damper=damper,
    parent_vertex=[0.0; (i - 1) * rod_separation - 0.5 * rod_separation * (num_balls - 1);  0.0],
    child_vertex=[0.0; 0.0; 0.5 * length_rod]),
    name=Symbol("connector_rod_$i")) for (i, rod) in enumerate(rods)]

# i = 1
# fixed_joint = JointConstraint(Fixed(connector, rods[i];
#         parent_vertex=[0.0; (i - 1) * rod_separation - 0.5 * rod_separation * (num_balls - 1);  0.0],
#         child_vertex=[0.0; 0.0; 0.5 * length_rod]),
#         name=Symbol("connector_rod_$i"))

# joint_connector_rods = [fixed_joint, joint_connector_rods[2]]

joint_rods_balls = [JointConstraint(Fixed(rods[i], balls[i],
    parent_vertex=[0.0; 0.0; -0.5 * length_rod],
    child_vertex=[0.0; 0.0; 0.0])) for i = 1:num_balls]

# Contacts
friction_coefficient = 0.5

collision = SphereSphereCollision{Float64,2,3,6}(
    [0.0; 0.0; 0.0],
    [0.0; 0.0; 0.0],
    radius_ball,
    radius_ball)

friction_parameterization = SA{Float64}[
        1.0  0.0
        0.0  1.0
]

body_body = NonlinearContact{Float64,8}(friction_coefficient, friction_parameterization, collision)

# collision = SphereSphereCollision{Float64,0,3,0}(
#     [0.0; 0.0; 0.0],
#     [0.0; 0.0; 0.0],
#     radius_ball,
#     radius_ball)

# friction_parameterization = szeros(Float64, 0, 2)

# body_body = ImpactContact{Float64,8}(friction_parameterization, collision)

bodies = [connector, rods..., balls...]
joints = [joint_world_connector, joint_connector_rods..., joint_rods_balls...]
contacts = [ContactConstraint((body_body, bodies[i + 1 + num_balls].id, bodies[i + 1 + num_balls + 1].id), name=Symbol("body_body_$i")) for i = 1:(num_balls-1)]

mech = Mechanism(origin, bodies, joints, contacts,
    gravity=gravity,
    timestep=timestep)

minimal_dimension(mech)
set_minimal_state!(mech, [0.25; 0.0; 0.0; 0.0; 0.0; 0.0])

@show adjacency_matrix(mech.joints, mech.bodies, mech.contacts)

distance(mech.contacts[1].model.collision,
    mech.bodies[1 + num_balls + 1].state.x2, mech.bodies[1 + num_balls + 1].state.q2,
    mech.bodies[1 + num_balls + 2].state.x2, mech.bodies[1 + num_balls + 2].state.q2,)

contact_point(:parent, mech.contacts[1].model.collision,
    mech.bodies[1 + num_balls + 1].state.x2, mech.bodies[1 + num_balls + 1].state.q2,
    mech.bodies[1 + num_balls + 2].state.x2, mech.bodies[1 + num_balls + 2].state.q2,)

contact_point(:child, mech.contacts[1].model.collision,
    mech.bodies[1 + num_balls + 1].state.x2, mech.bodies[1 + num_balls + 1].state.q2,
    mech.bodies[1 + num_balls + 2].state.x2, mech.bodies[1 + num_balls + 2].state.q2,)

contact_normal(mech.contacts[1].model.collision,
    mech.bodies[1 + num_balls + 1].state.x2, mech.bodies[1 + num_balls + 1].state.q2,
    mech.bodies[1 + num_balls + 2].state.x2, mech.bodies[1 + num_balls + 2].state.q2,)

contact_tangent(mech.contacts[1].model.collision,
    mech.bodies[1 + num_balls + 1].state.x2, mech.bodies[1 + num_balls + 1].state.q2,
    mech.bodies[1 + num_balls + 2].state.x2, mech.bodies[1 + num_balls + 2].state.q2,)

storage = simulate!(mech, 1.0,
    record=true,
    opts=SolverOptions(rtol=1.0e-6, btol=1.0e-6, verbose=true))

vis = Visualizer()
render(vis)
visualize(mech, storage,
    vis=vis)

# ## Visualize
# vis = Visualizer()
# render(vis)
build_robot(mech, vis=vis)
set_robot(vis, mech, get_maximal_state(mech), show_contact=false)


text:
using Statistics
using NCDatasets
using Test

using Oceananigans
using Oceananigans.Diagnostics
using Oceananigans.Fields
using Oceananigans.OutputWriters

using Dates: Millisecond
using Oceananigans: write_output!
using Oceananigans.BoundaryConditions: PBC, FBC, ZFBC, ContinuousBoundaryFunction
using Oceananigans.TimeSteppers: update_state!

include("utils_for_runtests.jl")

archs = test_architectures()

#####
##### WindowedTimeAverage tests
#####

function instantiate_windowed_time_average(model)

    set!(model, u = (x, y, z) -> rand())

    u, v, w = model.velocities

    u₀ = similar(interior(u))
    u₀ .= interior(u)

    wta = WindowedTimeAverage(model.velocities.u, schedule=AveragedTimeInterval(10, window=1))

    return all(wta(model) .== u₀)
end

function time_step_with_windowed_time_average(model)

    model.clock.iteration = 0
    model.clock.time = 0.0

    set!(model, u=0, v=0, w=0, T=0, S=0)

    wta = WindowedTimeAverage(model.velocities.u, schedule=AveragedTimeInterval(4, window=2))

    simulation = Simulation(model, Δt=1.0, stop_time=4.0)
    simulation.diagnostics[:u_avg] = wta
    run!(simulation)

    return all(wta(model) .== interior(model.velocities.u))
end

#####
##### Dependency adding tests
#####

function dependencies_added_correctly!(model, windowed_time_average, output_writer)

    model.clock.iteration = 0
    model.clock.time = 0.0

    simulation = Simulation(model, Δt=1.0, stop_iteration=1)
    push!(simulation.output_writers, output_writer)
    run!(simulation)

    return windowed_time_average ∈ values(simulation.diagnostics)
end

function test_dependency_adding(model)

    windowed_time_average = WindowedTimeAverage(model.velocities.u, schedule=AveragedTimeInterval(4, window=2))

    output = Dict("time_average" => windowed_time_average)
    attributes = Dict("time_average" => Dict("longname" => "A time average",  "units" => "arbitrary"))
    dimensions = Dict("time_average" => ("xF", "yC", "zC"))

    # JLD2 dependencies test
    jld2_output_writer = JLD2OutputWriter(model, output, schedule=TimeInterval(4), dir=".", prefix="test", force=true)

    @test dependencies_added_correctly!(model, windowed_time_average, jld2_output_writer)

    # NetCDF dependency test
    netcdf_output_writer = NetCDFOutputWriter(model, output,
                                              schedule = TimeInterval(4),
                                              filepath = "test.nc",
                                              output_attributes = attributes,
                                              dimensions = dimensions)

    @test dependencies_added_correctly!(model, windowed_time_average, netcdf_output_writer)

    rm("test.nc")

    return nothing
end

#####
##### Test time averaging of output
#####

function test_windowed_time_averaging_simulation(model)

    jld_filename1 = "test_windowed_time_averaging1"
    jld_filename2 = "test_windowed_time_averaging2"

    model.clock.iteration = model.clock.time = 0
    simulation = Simulation(model, Δt=1.0, stop_iteration=0)

    jld2_output_writer = JLD2OutputWriter(model, model.velocities,
                                          schedule = AveragedTimeInterval(π, window=1),
                                          prefix = jld_filename1,
                                          force = true)

    # https://github.com/Alexander-Barth/NCDatasets.jl/issues/105
    nc_filepath1 = "windowed_time_average_test1.nc"
    nc_outputs = Dict(string(name) => field for (name, field) in pairs(model.velocities))
    nc_output_writer = NetCDFOutputWriter(model, nc_outputs, filepath=nc_filepath1,
                                          schedule = AveragedTimeInterval(π, window=1))

    jld2_outputs_are_time_averaged = Tuple(typeof(out) <: WindowedTimeAverage for out in jld2_output_writer.outputs)
      nc_outputs_are_time_averaged = Tuple(typeof(out) <: WindowedTimeAverage for out in values(nc_output_writer.outputs))

    @test all(jld2_outputs_are_time_averaged)
    @test all(nc_outputs_are_time_averaged)

    # Test that the collection does *not* start when a simulation is initialized
    # when time_interval ≠ time_averaging_window
    simulation.output_writers[:jld2] = jld2_output_writer
    simulation.output_writers[:nc] = nc_output_writer

    run!(simulation)

    jld2_u_windowed_time_average = simulation.output_writers[:jld2].outputs.u
    nc_w_windowed_time_average = simulation.output_writers[:nc].outputs["w"]

    @test !(jld2_u_windowed_time_average.schedule.collecting)
    @test !(nc_w_windowed_time_average.schedule.collecting)

    # Test that time-averaging is finalized prior to output even when averaging over
    # time_window is not fully realized. For this, step forward to a time at which
    # collection should start. Note that time_interval = π and time_window = 1.0.
    simulation.Δt = 1.5
    simulation.stop_iteration = 2
    run!(simulation) # model.clock.time = 3.0, just before output but after average-collection.

    @test jld2_u_windowed_time_average.schedule.collecting
    @test nc_w_windowed_time_average.schedule.collecting

    # Step forward such that time_window is not reached, but output will occur.
    simulation.Δt = π - 3 + 0.01 # ≈ 0.15 < 1.0
    simulation.stop_iteration = 3
    run!(simulation) # model.clock.time ≈ 3.15, after output

    @test jld2_u_windowed_time_average.schedule.previous_interval_stop_time ==
        model.clock.time - rem(model.clock.time, jld2_u_windowed_time_average.schedule.interval)

    @test nc_w_windowed_time_average.schedule.previous_interval_stop_time ==
        model.clock.time - rem(model.clock.time, nc_w_windowed_time_average.schedule.interval)

    # Test that collection does start when a simulation is initialized and
    # time_interval == time_averaging_window
    model.clock.iteration = model.clock.time = 0

    simulation.output_writers[:jld2] = JLD2OutputWriter(model, model.velocities,
                                                        schedule = AveragedTimeInterval(π, window=π),
                                                          prefix = jld_filename2,
                                                           force = true)

    nc_filepath2 = "windowed_time_average_test2.nc"
    nc_outputs = Dict(string(name) => field for (name, field) in pairs(model.velocities))
    simulation.output_writers[:nc] = NetCDFOutputWriter(model, nc_outputs, filepath=nc_filepath2,
                                                        schedule=AveragedTimeInterval(π, window=π))

    run!(simulation)

    @test simulation.output_writers[:jld2].outputs.u.schedule.collecting
    @test simulation.output_writers[:nc].outputs["w"].schedule.collecting

    rm(nc_filepath1)
    rm(nc_filepath2)
    rm(jld_filename1 * ".jld2")
    rm(jld_filename2 * ".jld2")

    return nothing
end

#####
##### Run output writer tests!
#####

@testset "Output writers" begin
    @info "Testing output writers..."

    for arch in archs
        # Some tests can reuse this same grid and model.
        topo = (Periodic, Periodic, Bounded)
        grid = RegularRectilinearGrid(topology=topo, size=(4, 4, 4), extent=(1, 1, 1))
        model = NonhydrostaticModel(architecture=arch, grid=grid)

        @testset "WindowedTimeAverage [$(typeof(arch))]" begin
            @info "  Testing WindowedTimeAverage [$(typeof(arch))]..."
            @test instantiate_windowed_time_average(model)
            @test time_step_with_windowed_time_average(model)
            @test_throws ArgumentError AveragedTimeInterval(1.0, window=1.1)
        end
    end

    include("test_netcdf_output_writer.jl")
    include("test_jld2_output_writer.jl")
    include("test_checkpointer.jl")

    for arch in archs
        topo =(Periodic, Periodic, Bounded)
        grid = RegularRectilinearGrid(topology=topo, size=(4, 4, 4), extent=(1, 1, 1))
        model = NonhydrostaticModel(architecture=arch, grid=grid)

        @testset "Dependency adding [$(typeof(arch))]" begin
            @info "    Testing dependency adding [$(typeof(arch))]..."
            test_dependency_adding(model)
        end

        @testset "Time averaging of output [$(typeof(arch))]" begin
            @info "    Testing time averaging of output [$(typeof(arch))]..."
            test_windowed_time_averaging_simulation(model)
        end
    end
end


text:
export bracket, inverse_uk, inverse, inverse!

#Expands the range of valued searched geometrically until a root is bracketed
# Julia's adaptation of zbrac Numerical Recipes in Fortran 77 p.345
function bracket(f, a, b)
    F = 1.6
    Niter = 100
    outa = deepcopy(a)
    outb = deepcopy(b)
    counter = 0
    # Root is not bracketed
    while f(outa)*f(outb)>0 && counter<Niter
        Δ = F*deepcopy(outa-outb)
        if abs(f(outa))<abs(f(outb))
        outa += Δ
        else
        outb += Δ
        end
        counter += 1
    end
    return outa, outb
    @assert counter<=(Niter-1) "Maximal number of iterations reached"
end

function inverse_uk(u::uk, x, κ; z0::Real=0.0)
    if u.p==-1
    # Identity transformation
        return x
    elseif u.p==0
    # Solve u(z)= az + b = x for z
        return (x-u.coeffk[1])/u.coeffk[2]
    else
        zlim = (u.ξk[1]-κ*u.σk[1], u.ξk[end]+κ*u.σk[end])
        #Ensure that the zero is bracketed
        zlim = bracket(z->u(z)-x, zlim[1], zlim[2])
        # Roots.Brent()
        return find_zero(z->u(z)-x, zlim, Roots.Brent())
    end
end

# y is the observation of size ny
function inverse(x::AbstractVector{Float64}, F, S::RadialMap, ystar; apply_rescaling::Bool=true)
    @get S (Nx, p, κ)
    ny = size(ystar,1)
    x[1:ny] .= ystar
    # Recursive 1D root-finding
    # Start optimization from the a priori component

    if apply_rescaling == true
            transform!(S.L, x)
    end

    @inbounds for i=ny+1:Nx
        Ui = S.C[i]
        uk_i = component(Ui, i)
        x[i] = inverse_uk(uk_i, F[i] - off_diagonal(Ui, view(x,1:i-1)), κ)
    end

    if apply_rescaling == true
            itransform!(S.L, x)
    end
end

# y is the observation of size ny
function inverse(x::AbstractVector{Float64}, F, S::SparseRadialMap, ystar; apply_rescaling::Bool=true)
    @get S (Nx, p, κ)
    ny = size(ystar,1)
    x[1:ny] .= ystar

    if apply_rescaling == true
        transform!(S.L, x)
    end

    # Recursive 1D root-finding
    # Start optimization from the a priori component
    @inbounds for i=ny+1:Nx
        Ui = S.C[i]
        uk_i = component(Ui, i)
        x[i] = inverse_uk(uk_i, F[i] - off_diagonal(Ui, view(x,1:i-1)), κ)
    end

    if apply_rescaling == true
        transform!(S.L, x)
    end
end

function inverse!(X::Array{Float64,2}, F, S::SparseRadialMap, ystar::AbstractVector{Float64}; apply_rescaling::Bool=true, start::Int64=1)
    @get S (Nx, p, κ)
    Nx = S.Nx
    NxX, Ne = size(X)

    Ny = size(ystar,1)
    @assert NxX == Nx
    @assert 1 <= Ny < Nx
    @assert size(F) == (Nx, Ne)

    @view(X[1:Ny,:]) .= ystar

    if apply_rescaling == true
        transform!(S.L, X)
    end

    # Recursive 1D root-finding
    # Start optimization from the a priori component
    @inbounds for i=Ny+1:Nx
        Ui = S.C[i]
        if Ui.p[end] != -1
            uk_i = component(Ui, i)
            for j=1:Ne
                X[i,j] = inverse_uk(uk_i, F[i,j] - off_diagonal(Ui, view(X,1:i-1,j)), κ)
            end
        end
    end

    if apply_rescaling == true
        itransform!(S.L, X)
    end
end


text:
@testset "utils" begin
    m = [
        1 2 3 4 5;
        6 7 8 9 10
    ]

    @test QuantumStateDB.hexbytes(m) ==
        "0100000000000000" * "0600000000000000" *
        "0200000000000000" * "0700000000000000" *
        "0300000000000000" * "0800000000000000" *
        "0400000000000000" * "0900000000000000" *
        "0500000000000000" * "0a00000000000000"

    @test QuantumStateDB.hexbytes_str(m) == "\\x" *
        "0100000000000000" * "0600000000000000" *
        "0200000000000000" * "0700000000000000" *
        "0300000000000000" * "0800000000000000" *
        "0400000000000000" * "0900000000000000" *
        "0500000000000000" * "0a00000000000000"

    @test reshape(QuantumStateDB.hexbytes2array(Int)(
        reinterpret(UInt8, reshape(m, :))
    ), 2, 5) == m
end


text:
__precompile__(false)

module ReplMaker

import REPL
import REPL.LineEdit

export initrepl, enter_mode!

"""
```
         initrepl(parser::Function;
                  prompt_text = "myrepl> ",
                  prompt_color = :blue,
                  start_key = ')',
                  repl = Base.active_repl,
                  mode_name = :mylang,
                  valid_input_checker::Function = (s -> true),
                  keymap::Dict = REPL.LineEdit.default_keymap_dict,
                  completion_provider = REPL.REPLCompletionProvider(),
                  sticky = true,
                  startup_text=true
                  )
```
creates a custom repl mode which takes in code and parses it according to whatever parsing function you
provide in the argument `parser`. Choose which key initializes the repl mode with `start_key`, the name of
your repl mode with `mode_name` and optionally provide a function which checks if a given repl input is valid
before parsing with `valid_input_checker`. Autocompletion options are supplied through the argument
`completion_provider` which defaults to the standard julia REPL TAB completions.
"""
function initrepl(parser::Function;
                  prompt_text = "myrepl> ",
                  prompt_color = :blue,
                  start_key = ')',
                  repl = Base.active_repl,
                  mode_name = :mylang,
                  valid_input_checker::Function = (s -> true),
                  keymap::Dict = REPL.LineEdit.default_keymap_dict,
                  completion_provider = REPL.REPLCompletionProvider(),
                  sticky_mode = true,
                  startup_text=true
                  )

    color = Base.text_colors[prompt_color]

    julia_mode = repl.interface.modes[1]
    prefix = repl.hascolor ? color : ""
    suffix = repl.hascolor ? (repl.envcolors ? Base.input_color : repl.input_color()) : ""

    lang_mode = LineEdit.Prompt(prompt_text;
                                prompt_prefix    = prefix,
                                prompt_suffix    = suffix,
                                keymap_dict      = keymap,
                                on_enter         = valid_input_checker,
                                complete         = completion_provider,
                                sticky           = sticky_mode
                                )

    lang_mode.on_done = REPL.respond(parser, repl, lang_mode)

    push!(repl.interface.modes, lang_mode)

    hp                     = julia_mode.hist
    hp.mode_mapping[mode_name |> Symbol] = lang_mode
    lang_mode.hist         = hp

    search_prompt, skeymap = LineEdit.setup_search_keymap(hp)

    mk = REPL.mode_keymap(julia_mode)

    lang_keymap = Dict(
    start_key => (s, args...) ->
      if isempty(s) || position(LineEdit.buffer(s)) == 0
        enter_mode!(s, lang_mode)
      else
        LineEdit.edit_insert(s, start_key)
      end
    )

    lang_mode.keymap_dict = LineEdit.keymap(Dict[
        skeymap,
        mk,
        LineEdit.history_keymap,
        LineEdit.default_keymap,
        LineEdit.escape_defaults,
    ])
    if start_key in keys(julia_mode.keymap_dict)
        @warn "REPL key '$start_key' overwritten."
    end

    julia_mode.keymap_dict = LineEdit.keymap_merge(julia_mode.keymap_dict, lang_keymap)

    startup_text && println("REPL mode $mode_name initialized. Press $start_key to enter and backspace to exit.")

  lang_mode
end

"""
an extension of `REPL.LineEdit.transition` to achieve handy use.

When
- giving 3 arguments(action, mistate, mode),
  it equals to `REPL.LineEdit.transition`.

- giving 2 arguments(mistate, mode),
  there is a default 'action' copying the
  current buffer to the target REPL mode.

- giving 1 argument(mode),
  the action is the same as the case of 2 arguments,
  and 'mistate' is default to be `Base.active_repl.mistate`
"""
function enter_mode!(f, s :: LineEdit.MIState, lang_mode)
  LineEdit.transition(f, s, lang_mode)
end

function enter_mode!(s :: LineEdit.MIState, lang_mode)
  function default_trans_action()
    buf = copy(LineEdit.buffer(s))
    LineEdit.state(s, lang_mode).input_buffer = buf
  end
  enter_mode!(default_trans_action, s, lang_mode)
end

function enter_mode!(lang_mode)
  enter_mode!(Base.active_repl.mistate, lang_mode)
end

end