# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/MortarContact2D.jl/blob/master/LICENSE

mutable struct Contact2D <: BoundaryProblem
    master_elements :: Vector{Element}
    rotate_normals :: Bool
    dual_basis :: Bool
    max_distance :: Float64
    iteration :: Int
    contact_state_in_first_iteration :: Symbol
    store_fields :: Vector{String}
end

function Contact2D()
    return Contact2D([], false, true, Inf, 0, :AUTO, [])
end

function FEMBase.add_elements!(::Problem{Contact2D}, ::Any)
    error("use `add_slave_elements!` and `add_master_elements!` to add ",
          "elements to the Contact2D problem.")
end

function FEMBase.get_elements(problem::Problem{Contact2D})
    return [problem.elements; problem.properties.master_elements]
end

function FEMBase.add_slave_elements!(problem::Problem{Contact2D}, elements)
    for element in elements
        push!(problem.elements, element)
    end
end

function FEMBase.add_master_elements!(problem::Problem{Contact2D}, elements)
    for element in elements
        push!(problem.properties.master_elements, element)
    end
end

function FEMBase.get_slave_elements(problem::Problem{Contact2D})
    return problem.elements
end

function FEMBase.get_master_elements(problem::Problem{Contact2D})
    return problem.properties.master_elements
end

function get_slave_dofs(problem::Problem{Contact2D})
    dofs = Int64[]
    for element in get_slave_elements(problem)
        append!(dofs, get_gdofs(problem, element))
    end
    return sort(unique(dofs))
end

function get_master_dofs(problem::Problem{Contact2D})
    dofs = Int64[]
    for element in get_master_elements(problem)
        append!(dofs, get_gdofs(problem, element))
    end
    return sort(unique(dofs))
end

function create_rotation_matrix(element::Element{Seg2}, time::Float64)
    n = element("normal", time)
    R = [0.0 -1.0; 1.0 0.0]
    t1 = R'*n[1]
    t2 = R'*n[2]
    Q1 = [n[1] t1]
    Q2 = [n[2] t2]
    Z = zeros(2, 2)
    Q = [Q1 Z; Z Q2]
    return Q
end

function create_contact_segmentation(problem::Problem{Contact2D}, slave_element::Element{Seg2}, master_elements::Vector, time::Float64; deformed=false)

    result = []

    max_distance = problem.properties.max_distance

    x1 = slave_element("geometry", time)

    if deformed && haskey(slave_element, "displacement")
        x1 = map(+, x1, slave_element("displacement", time))
    end

    slave_midpoint = 1/2*(x1[1] + x1[2])
    slave_length = norm(x1[2]-x1[1])

    for master_element in master_elements

        x2 = master_element("geometry", time)

        if deformed && haskey(master_element, "displacement")
            x2 = map(+, x2, master_element("displacement", time))
        end

        master_midpoint = 1/2*(x2[1] + x2[2])
        master_length = norm(x2[2]-x2[1])
        # charasteristic length
        dist = norm(slave_midpoint - master_midpoint)
        cl = dist/max(slave_length, master_length)
        if cl > max_distance
            continue
        end

        # 3.1 calculate segmentation
        x21, x22 = x2
        xi1a = project_from_master_to_slave(slave_element, x21, time)
        xi1b = project_from_master_to_slave(slave_element, x22, time)
        xi1 = clamp.([xi1a; xi1b], -1.0, 1.0)
        l = 1/2*abs(xi1[2]-xi1[1])
        if isapprox(l, 0.0)
            continue # no contribution in this master element
        end
        push!(result, (master_element, xi1, l))
    end
    return result
end

function FEMBase.assemble_elements!(problem::Problem{Contact2D}, assembly::Assembly,
                                    elements::Vector{Element{Seg2}}, time::Float64)

    props = problem.properties
    props.iteration += 1

    slave_elements = get_slave_elements(problem)

    field_dim = get_unknown_field_dimension(problem)
    field_name = get_parent_field_name(problem)

    # 1. calculate nodal normals for slave element nodes j ∈ S
    normals = calculate_normals([slave_elements...], time; rotate_normals=props.rotate_normals)
    tangents = Dict(j => [t2, -t1] for (j, (t1,t2)) in normals)
    update!(slave_elements, "normal", time => normals)
    update!(slave_elements, "tangent", time => tangents)

    Rn = 0.0

    # 2. loop all slave elements
    for slave_element in slave_elements

        nsl = length(slave_element)
        X1 = slave_element("geometry", time)
        x1 = slave_element("geometry", time)
        if haskey(slave_element, "displacement")
            u1 = slave_element("displacement", time)
            x1 = map(+, x1, u1)
        end
        la1 = slave_element("lambda", time)
        n1 = slave_element("normal", time)
        t1 = slave_element("tangent", time)

        contact_area = 0.0
        contact_error = 0.0
        Q2 = create_rotation_matrix(slave_element, time)

        master_elements = get_master_elements(problem)
        segmentation = create_contact_segmentation(problem, slave_element, master_elements, time)
        if length(segmentation) == 0 # no overlapping in master and slave surfaces with this slave element
            continue
        end

        Ae = Matrix(I, nsl, nsl)
        if props.dual_basis
            De = zeros(nsl, nsl)
            Me = zeros(nsl, nsl)
            for (master_element, xi1, l) in segmentation
                for ip in get_integration_points(slave_element, 3)
                    detJ = slave_element(ip, time, Val{:detJ})
                    w = ip.weight*detJ*l
                    xi = ip.coords[1]
                    xi_s = dot([1/2*(1-xi); 1/2*(1+xi)], xi1)
                    N1 = vec(get_basis(slave_element, xi_s, time))
                    De += w*Matrix(Diagonal(N1))
                    Me += w*N1*N1'
                end
                Ae = De*inv(Me)
            end
        end

        # loop all segments
        for (master_element, xi1, l) in segmentation

            nm = length(master_element)
            X2 = master_element("geometry", time)
            x2 = master_element("geometry", time)
            if haskey(master_element, "displacement")
                u2 = master_element("displacement", time)
                x2 = map(+, x2, u2)
            end

            # 3.3. loop integration points of one integration segment and calculate
            # local mortar matrices
            De = zeros(nsl, nsl)
            Me = zeros(nsl, nsl)
            Ne = zeros(nsl, 2*nsl)
            Te = zeros(nsl, 2*nsl)
            He = zeros(nsl, 2*nsl)
            ce = zeros(nsl)
            ge = zeros(nsl)
            for ip in get_integration_points(slave_element, 3)
                detJ = slave_element(ip, time, Val{:detJ})
                w = ip.weight*detJ*l
                xi = ip.coords[1]
                xi_s = dot([1/2*(1-xi); 1/2*(1+xi)], xi1)
                N1 = vec(get_basis(slave_element, xi_s, time))
                Phi = Ae*N1

                # project gauss point from slave element to master element in direction n_s
                X_s = interpolate(N1, X1) # coordinate in gauss point
                n_s = interpolate(N1, n1) # normal direction in gauss point
                t_s = interpolate(N1, t1) # tangent condition in gauss point
                n_s /= norm(n_s)
                t_s /= norm(t_s)
                xi_m = project_from_slave_to_master(master_element, X_s, n_s, time)
                N2 = vec(get_basis(master_element, xi_m, time))
                X_m = interpolate(N2, X2)

                u_s = interpolate(N1, u1)
                u_m = interpolate(N2, u2)
                x_s = map(+, X_s, u_s)
                x_m = map(+, X_m, u_m)
                la_s = interpolate(Phi, la1)

                # contact virtual work
                De += w*Phi*N1'
                Me += w*Phi*N2'

                # contact constraints
                Ne += w*reshape(kron(N1, n_s, Phi), 2, 4)
                Te += w*reshape(kron(N2, n_s, Phi), 2, 4)
                He += w*reshape(kron(N1, t_s, Phi), 2, 4)
                ge += w*Phi*dot(n_s, x_m-x_s)
                ce += w*N1*dot(n_s, la_s)
                Rn += w*dot(n_s, la_s)

                contact_area += w
                contact_error += 1/2*w*dot(n_s, x_s-x_m)^2
            end

            sdofs = get_gdofs(problem, slave_element)
            mdofs = get_gdofs(problem, master_element)

            # add contribution to contact virtual work
            for i=1:field_dim
                lsdofs = sdofs[i:field_dim:end]
                lmdofs = mdofs[i:field_dim:end]
                add!(problem.assembly.C1, lsdofs, lsdofs, De)
                add!(problem.assembly.C1, lsdofs, lmdofs, -Me)
            end

            # add contribution to contact constraints
            add!(problem.assembly.C2, sdofs[1:field_dim:end], sdofs, Ne)
            add!(problem.assembly.C2, sdofs[1:field_dim:end], mdofs, -Te)
            add!(problem.assembly.D, sdofs[2:field_dim:end], sdofs, He)
            add!(problem.assembly.g, sdofs[1:field_dim:end], ge)
            add!(problem.assembly.c, sdofs[1:field_dim:end], ce)

        end # master elements done

        if "contact area" in props.store_fields
            update!(slave_element, "contact area", time => contact_area)
        end

        if "contact error" in props.store_fields
            update!(slave_element, "contact error", time => contact_error)
        end

    end # slave elements done, contact virtual work ready

    S = sort(collect(keys(normals))) # slave element nodes
    weighted_gap = Dict{Int64, Vector{Float64}}()
    contact_pressure = Dict{Int64, Vector{Float64}}()
    complementarity_condition = Dict{Int64, Vector{Float64}}()
    is_active = Dict{Int64, Int}()
    is_inactive = Dict{Int64, Int}()
    is_slip = Dict{Int64, Int}()
    is_stick = Dict{Int64, Int}()

    la = problem.assembly.la
    # FIXME: for matrix operations, we need to know the dimensions of the
    # final matrices
    ndofs = 2*length(S)
    ndofs = max(ndofs, length(la))
    ndofs = max(ndofs, size(problem.assembly.K, 2))
    ndofs = max(ndofs, size(problem.assembly.C1, 2))
    ndofs = max(ndofs, size(problem.assembly.C2, 2))
    ndofs = max(ndofs, size(problem.assembly.D, 2))
    ndofs = max(ndofs, size(problem.assembly.g, 2))
    ndofs = max(ndofs, size(problem.assembly.c, 2))

    C1 = sparse(problem.assembly.C1, ndofs, ndofs)
    C2 = sparse(problem.assembly.C2, ndofs, ndofs)
    D = sparse(problem.assembly.D, ndofs, ndofs)
    g = Vector(sparse(problem.assembly.g, ndofs, 1)[:])
    c = Vector(sparse(problem.assembly.c, ndofs, 1)[:])

    for j in S
        dofs = [2*(j-1)+1, 2*(j-1)+2]
        weighted_gap[j] = g[dofs]
    end

    state = problem.properties.contact_state_in_first_iteration
    if problem.properties.iteration == 1
        @info("First contact iteration, initial contact state = $state")

        if state == :AUTO
            avg_gap = mean([weighted_gap[j][1] for j in S])
            std_gap = std([weighted_gap[j][1] for j in S])
            if (avg_gap < 1.0e-12) && (std_gap < 1.0e-12)
                state = :ACTIVE
            else
                state = :UNKNOWN
            end
            @info("Average weighted gap = $avg_gap, std gap = $std_gap, automatically determined contact state = $state")
        end

    end

    # active / inactive node detection
    for j in S
        dofs = [2*(j-1)+1, 2*(j-1)+2]
        weighted_gap[j] = g[dofs]

        if length(la) != 0
            p = dot(normals[j], la[dofs])
            t = dot(tangents[j], la[dofs])
            contact_pressure[j] = [p, t]
        else
            contact_pressure[j] = [0.0, 0.0]
        end

        complementarity_condition[j] = contact_pressure[j] - weighted_gap[j]
        if complementarity_condition[j][1] < 0
            is_inactive[j] = 1
            is_active[j] = 0
            is_slip[j] = 0
            is_stick[j] = 0
        else
            is_inactive[j] = 0
            is_active[j] = 1
            is_slip[j] = 1
            is_stick[j] = 0
        end
    end

    if (problem.properties.iteration == 1) && (state == :ACTIVE)
        for j in S
            is_inactive[j] = 0
            is_active[j] = 1
            is_slip[j] = 1
            is_stick[j] = 0
        end
    end

    if (problem.properties.iteration == 1) && (state == :INACTIVE)
        for j in S
            is_inactive[j] = 1
            is_active[j] = 0
            is_slip[j] = 0
            is_stick[j] = 0
        end
    end

    if "weighted gap" in props.store_fields
        update!(slave_elements, "weighted gap", time => weighted_gap)
    end
    if "contact pressure" in props.store_fields
        update!(slave_elements, "contact pressure", time => contact_pressure)
    end
    if "complementarity condition" in props.store_fields
        update!(slave_elements, "complementarity condition", time => complementarity_condition)
    end
    if "active nodes" in props.store_fields
        update!(slave_elements, "active nodes", time => is_active)
    end
    if "inactive nodes" in props.store_fields
        update!(slave_elements, "inactive nodes", time => is_inactive)
    end
    if "stick nodes" in props.store_fields
        update!(slave_elements, "stick nodes", time => is_stick)
    end
    if "slip nodes" in props.store_fields
        update!(slave_elements, "slip nodes", time => is_slip)
    end

    @info("# | A | I | St | Sl | gap | pres | comp")
    for j in S
        str1 = "$j | $(is_active[j]) | $(is_inactive[j]) |  $(is_stick[j]) |  $(is_slip[j]) | "
        str2 = "$(round(weighted_gap[j][1]; digits=3)) | $(round(contact_pressure[j][1]; digits=3)) | $(round(complementarity_condition[j][1]; digits=3))"
        @info(str1 * str2)
    end

    # solve variational inequality

    # remove inactive nodes from assembly
    for j in S
        dofs = [2*(j-1)+1, 2*(j-1)+2]
        n, t = dofs
        if is_inactive[j] == 1
            @debug("$j is inactive, removing dofs $dofs")
            C1[dofs,:] .= 0.0
            C2[dofs,:] .= 0.0
            D[dofs,:] .= 0.0
            g[dofs,:] .= 0.0
        elseif (is_active[j] == 1) && (is_slip[j] == 1)
            @debug("$j is in active/slip, removing tangential constraint $t")
            C2[t,:] .= 0.0
            g[t] = 0.0
            D[t, dofs] .= tangents[j]
        end
    end

    problem.assembly.C1 = C1
    problem.assembly.C2 = C2
    problem.assembly.D = D
    problem.assembly.g = sparsevec(g)

end


function ecRCCSDpT{T}() where T <: AbstractFloat
    @output "Calling CASCI module...\n"
    # Call CASCI
    cas = Fermi.ConfigurationInteraction.CASCI{T}()
    ecRCCSDpT{T}(cas)
end

function ecRCCSDpT{T}(cas::Fermi.ConfigurationInteraction.CASCI) where T <: AbstractFloat

    @output "Processing CAS data...\n"
    refdet, Casdata = process_cas(cas)
    refwfn = cas.ref

    ecccsd = ecRCCSD{T}(refdet, refwfn, Casdata, CTF())
    drop_occ = Fermi.CurrentOptions["drop_occ"]
    ecRCCSDpT{T}(ecccsd, refwfn.ints, refdet, Casdata["C3dets"], refwfn.ndocc, drop_occ)
end

function ecRCCSDpT{T}(ecccsd::RCCSD, ints::IntegralHelper, ref::Determinant, casT3::Array{Determinant,1}, ndocc::Int, fcn::Int) where T <: AbstractFloat

    @output "\n   • Perturbative Triples Started\n\n"
    @output "T3 within EC active space are going to be skipped\n\n"

    T1 = ecccsd.T1.data
    T2 = ecccsd.T2.data

    Vvvvo = permutedims(ints["OVVV"], (4,2,3,1))
    Vvooo = permutedims(ints["OOOV"], (4,2,1,3))
    Vvovo = permutedims(ints["OOVV"], (3,1,4,2))

    o,v = size(T1)
    Et::T = 0.0

    fo = diag(ints["FOO"])
    fv = diag(ints["FVV"])

    # Pre-allocate Intermediate arrays
    W  = Array{T}(undef, v,v,v)
    V  = Array{T}(undef, v,v,v)

    # Pre-allocate Views
    Vvvvo_4k =    Array{T}(undef,v,v,v)
    Vvooo_2k_3j = Array{T}(undef,v,o)
    Vvooo_2k_3i = Array{T}(undef,v,o)
    Vvooo_2i_3k = Array{T}(undef,v,o)
    Vvooo_2j_3k = Array{T}(undef,v,o)
    T2_1k_2j    = Array{T}(undef,v,v)
    T2_1k_2i    = Array{T}(undef,v,v)
    T2_1i_2k    = Array{T}(undef,v,v)
    T2_1j_2k    = Array{T}(undef,v,v)
    T2_1k       = Array{T}(undef,o,v,v)
    T1_1k       = Array{T}(undef,v)
    Vvovo_2j_4k = Array{T}(undef,v,v)
    Vvovo_2i_4k = Array{T}(undef,v,v)
    Vvvvo_4i    = Array{T}(undef,v,v,v)
    T2_1i       = Array{T}(undef,o,v,v)
    T1_1i       = Array{T}(undef,v)
    Vvvvo_4j    = Array{T}(undef,v,v,v)
    Vvooo_2i_3j = Array{T}(undef,v,o)
    Vvooo_2j_3i = Array{T}(undef,v,o)
    T2_1i_2j    = Array{T}(undef,v,v)
    T2_1j_2i    = Array{T}(undef,v,v)
    T2_1j       = Array{T}(undef,o,v,v)
    Vvovo_2i_4j = Array{T}(undef,v,v)
    T1_1j       = Array{T}(undef,v)
    @output "Computing energy contribution from occupied orbitals:\n"
    for i in 1:o
        iα = ref.α ⊻ 1 << (i+fcn-1)
        @output "→ Orbital {} of {}\n" i o
        Vvvvo_4i .= view(Vvvvo, :,:,:,i)
        T2_1i    .= view(T2, i,:,:,:)
        T1_1i    .= view(T1, i, :)
        for j in 1:i
            jβ = ref.β ⊻ 1 << (j+fcn-1)
            Vvvvo_4j    .= view(Vvvvo, :,:,:,j)
            Vvooo_2i_3j .= view(Vvooo,:,i,j,:)
            Vvooo_2j_3i .= view(Vvooo,:,j,i,:)
            T2_1i_2j    .= view(T2, i,j,:,:)
            T2_1j_2i    .= view(T2, j,i,:,:)
            T2_1j       .= view(T2, j,:,:,:)
            Vvovo_2i_4j .= view(Vvovo, :,i,:,j)
            T1_1j       .= view(T1, j, :)
            δij = Int(i == j)
            for k in 1:j
                kiα = iα ⊻ 1 << (k+fcn-1)
                Vvvvo_4k    .= view(Vvvvo, :,:,:,k)
                Vvooo_2k_3j .= view(Vvooo,:,k,j,:)
                Vvooo_2k_3i .= view(Vvooo,:,k,i,:)
                Vvooo_2i_3k .= view(Vvooo,:,i,k,:)
                Vvooo_2j_3k .= view(Vvooo,:,j,k,:)

                T2_1k_2j .= view(T2, k,j,:,:)
                T2_1k_2i .= view(T2, k,i,:,:)
                T2_1i_2k .= view(T2, i,k,:,:)
                T2_1j_2k .= view(T2, j,k,:,:)

                T2_1k .= view(T2, k,:,:,:)
                T1_1k .= view(T1, k, :)

                Vvovo_2j_4k .= view(Vvovo, :,j,:,k)
                Vvovo_2i_4k .= view(Vvovo, :,i,:,k)

                # Build intermediates W and V for the set i j k
                @tensoropt begin
                    W[a,b,c]  = (Vvvvo_4i[b,d,a]*T2_1k_2j[c,d] - Vvooo_2k_3j[c,l]*T2_1i[l,a,b]  # ijk abc
                              +  Vvvvo_4i[c,d,a]*T2_1j_2k[b,d] - Vvooo_2j_3k[b,l]*T2_1i[l,a,c]  # ikj acb
                              +  Vvvvo_4k[a,d,c]*T2_1j_2i[b,d] - Vvooo_2j_3i[b,l]*T2_1k[l,c,a]  # kij cab
                              +  Vvvvo_4k[b,d,c]*T2_1i_2j[a,d] - Vvooo_2i_3j[a,l]*T2_1k[l,c,b]  # kji cba
                              +  Vvvvo_4j[c,d,b]*T2_1i_2k[a,d] - Vvooo_2i_3k[a,l]*T2_1j[l,b,c]  # jki bca
                              +  Vvvvo_4j[a,d,b]*T2_1k_2i[c,d] - Vvooo_2k_3i[c,l]*T2_1j[l,b,a]) # jik bac

                    V[a,b,c]  = W[a,b,c] + Vvovo_2j_4k[b,c]*T1_1i[a] + Vvovo_2i_4k[a,c]*T1_1j[b] + Vvovo_2i_4j[a,b]*T1_1k[c]
                end

                # Compute Energy contribution
                δjk = Int(j == k)
                for a in 1:v
                    akiα = kiα | 1 << (a+ndocc-1)
                    for b in 1:a
                        bjβ = jβ | 1 << (b+ndocc-1)
                        δab = Int(a == b)
                        for c in 1:b
                            cakiα = akiα | 1 << (c+ndocc-1)
                            _det = Determinant(cakiα, bjβ)
                            if _det in casT3
                                continue
                            end
                            Dd = fo[i] + fo[j] + fo[k] - fv[a] - fv[b] - fv[c]
                            δbc = Int(b == c)
                            @inbounds X = (W[a,b,c]*V[a,b,c] + W[a,c,b]*V[a,c,b] + W[b,a,c]*V[b,a,c] + W[b,c,a]*V[b,c,a] + W[c,a,b]*V[c,a,b] + W[c,b,a]*V[c,b,a])
                            @inbounds Y = (V[a,b,c] + V[b,c,a] + V[c,a,b])
                            @inbounds Z = (V[a,c,b] + V[b,a,c] + V[c,b,a])
                            E = (Y - 2*Z)*(W[a,b,c] + W[b,c,a] + W[c,a,b]) + (Z - 2*Y)*(W[a,c,b]+W[b,a,c]+W[c,b,a]) + 3*X
                            Et += E*(2-δij-δjk)/(Dd*(1+δab+δbc))
                        end
                    end
                end
            end
        end
    end
    @output "Final (T) contribution:   {:15.10f}\n" Et
    @output "ecCCSD(T) energy:         {:15.10f}\n" Et+ecccsd.CorrelationEnergy
    return RCCSDpT{T}(ecccsd, Et)
end


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@with_kw mutable struct ProgressiveHedgingData{T <: AbstractFloat}
    Q::T = 1e10
    δ₁::T = 1.0
    δ₂::T = 1.0
    iterations::Int = 0
end

@with_kw mutable struct ProgressiveHedgingParameters{T <: AbstractFloat}
    ϵ₁::T = 1e-5
    ϵ₂::T = 1e-5
    time_limit::T = Inf
    log::Bool = true
    keep::Bool = true
    offset::Int = 0
    indent::Int = 0
end

"""
    ProgressiveHedgingAlgorithm

Functor object for the progressive-hedging algorithm. Create using the `ProgressiveHedgingSolver` factory function and then pass to a `StochasticPrograms.jl` model.

...
# Parameters
- `τ::AbstractFloat = 1e-6`: Relative tolerance for convergence checks.
- `log::Bool = true`: Specifices if progressive-hedging procedure should be logged on standard output or not.
...
"""
struct ProgressiveHedgingAlgorithm{T <: AbstractFloat,
                                   A <: AbstractVector,
                                   ST <: ScenarioDecompositionStructure,
                                   E <: AbstractProgressiveHedgingExecution,
                                   P <: AbstractPenalization,
                                   PT <: AbstractPenaltyTerm} <: AbstractProgressiveHedging
    structure::ST
    data::ProgressiveHedgingData{T}
    parameters::ProgressiveHedgingParameters{T}

    # Estimate
    ξ::A
    decisions::DecisionMap
    Q_history::A
    primal_gaps::A
    dual_gaps::A

    # Execution
    execution::E

    penalization::P

    # Params
    progress::ProgressThresh{T}

    function ProgressiveHedgingAlgorithm(structure::ScenarioDecompositionStructure,
                                         x₀::AbstractVector,
                                         executer::AbstractExecution,
                                         penalizer::AbstractPenalizer,
                                         penaltyterm::AbstractPenaltyTerm; kw...)
        # Sanity checks
        length(x₀) != num_decisions(structure) && error("Incorrect length of starting guess, has ", length(x₀), " should be ", num_decisions(structure))
        num_subproblems == 0 && error("No subproblems in stochastic program. Cannot run progressive-hedging procedure.")
        n = num_subproblems(structure, 2)
        # Float types
        T = promote_type(eltype(x₀), Float32)
        x₀_ = convert(AbstractVector{T}, copy(x₀))
        A = typeof(x₀_)
        # Structure
        ST = typeof(structure)
        # Penalty term
        PT = typeof(penaltyterm)
        # Execution policy
        execution = executer(T, A, PT)
        E = typeof(execution)
        # Penalization policy
        penalization = penalizer()
        P = typeof(penalization)
        # Algorithm parameters
        params = ProgressiveHedgingParameters{T}(; kw...)

        ph = new{T,A,ST,E,P,PT}(structure,
                                ProgressiveHedgingData{T}(),
                                params,
                                x₀_,
                                structure.decisions[1],
                                A(),
                                A(),
                                A(),
                                execution,
                                penalization,
                                ProgressThresh(T(1.0), 0.0, "$(indentstr(params.indent))Progressive Hedging"))
        # Initialize solver
        initialize!(ph, penaltyterm)
        return ph
    end
end

function show(io::IO, ph::ProgressiveHedgingAlgorithm)
    println(io, typeof(ph).name.name)
    println(io, "State:")
    show(io, ph.data)
    println(io, "Parameters:")
    show(io, ph.parameters)
end

function show(io::IO, ::MIME"text/plain", ph::ProgressiveHedgingAlgorithm)
    show(io, ph)
end

function (ph::ProgressiveHedgingAlgorithm)()
    # Reset timer
    ph.progress.tinit = ph.progress.tlast = time()
    # Start workers (if any)
    start_workers!(ph)
    # Start procedure
    while true
        status = iterate!(ph)
        if status !== nothing
            close_workers!(ph)
            return status
        end
    end
end


using Poly
using ISL
import ISL.API
using Printf

export run_polyhedral_model

"""
Run polyhedral model on a kernel. Does the following:
creates a context and space for ISL
create an ISL union set representing the domains and instructions in kernel
create an ISL union maps representing access relations in kernel
create an ISL schedule from intial implementation
create dependencies maps by analyzing the accesses and schedule
add loop optimal ordering to dependencies
create a new schedule from all dependencies
create an AST from the schedule
parse AST to julia code
return expression

if debug is set, ISL errors are printed to stderror
if verbose is
    0: no printing (unless debug is on, where ISL errors are printed)
    1: final Julia code is printed
    2: initial C code, loop orderings, new schedule, final C code, and final Julia code printed
    3: all of 2 plus domain, access relations, dependence relations, original schedule, and new schedule constraints printed
"""
function run_polyhedral_model(kernel::LoopKernel; debug=false, verbose=0, tile=-1, thread=false)::Expr
    """
    initialize context and set error printing options to warn/silence
    """
    context = ISL.API.isl_ctx_alloc()
    if debug
        ISL.API.isl_options_set_on_error(context, 0) # warn errors
    else
        ISL.API.isl_options_set_on_error(context, 1) # silence errors
    end

    """
    get the domain (instructions) of the kernel
    """
    instructions = ISL.API.isl_union_set_read_from_str(context, instructions_isl_rep(kernel))
    if verbose >= 3
        println("===DOMAIN===")
        ISL.API.isl_union_set_dump(instructions)
    end

    """
    construct the access relations
    """
    may_read, must_write = accesses_isl_rep(kernel)
    if verbose >= 3
        println("===ACCESS RELATIONS===")
        @show may_read
        @show must_write
    end
    may_read = ISL.API.isl_union_map_read_from_str(context, may_read)
    must_write = ISL.API.isl_union_map_read_from_str(context, must_write)

    """
    construct the original schedule
    """
    schedule = schedule_tree_isl_rep(context, kernel.domains[1], kernel)
    if verbose >= 3
        println("===ORIGINAL SCHEDULE===")
        ISL.API.isl_schedule_dump(schedule)
    end
    if verbose >= 2
        # show original code in C representation
        space = ISL.API.isl_set_read_from_str(context, "{:}")
        build = ISL.API.isl_ast_build_from_context(space)
        ast = ISL.API.isl_ast_build_node_from_schedule(build, ISL.API.isl_schedule_copy(schedule))
        c = ISL.API.isl_ast_node_get_ctx(ast)
        p = ISL.API.isl_printer_to_str(c)
        p = ISL.API.isl_printer_set_output_format(p, 4) # 4 = C code
        q = ISL.API.isl_printer_print_ast_node(p, ast)
        str = ISL.API.isl_printer_get_str(q)
        str = Base.unsafe_convert(Ptr{Cchar}, str)
        println("===ORIGINAL C CODE===")
        println(Base.unsafe_string(str))
    end

    """
    dependency analysis
    """
    # read after write deps
    access = ISL.API.isl_union_access_info_from_sink(ISL.API.isl_union_map_copy(may_read))
    access = ISL.API.isl_union_access_info_set_must_source(access, ISL.API.isl_union_map_copy(must_write))
    access = ISL.API.isl_union_access_info_set_schedule(access, ISL.API.isl_schedule_copy(schedule))
    flow = ISL.API.isl_union_access_info_compute_flow(access)
    raw_deps = ISL.API.isl_union_flow_get_must_dependence(flow)
    # write after read and write after write deps
    access = ISL.API.isl_union_access_info_from_sink(ISL.API.isl_union_map_copy(must_write))
    access = ISL.API.isl_union_access_info_set_must_source(access, must_write)
    access = ISL.API.isl_union_access_info_set_may_source(access, may_read)
    access = ISL.API.isl_union_access_info_set_schedule(access, schedule)
    flow = ISL.API.isl_union_access_info_compute_flow(access)
    waw_deps = ISL.API.isl_union_flow_get_must_dependence(flow)
    war_deps = ISL.API.isl_union_flow_get_may_dependence(flow)
    if verbose >= 3
        println("===DEPENDENCIES (RAW, WAW, WAR)===")
        ISL.API.isl_union_map_dump(raw_deps)
        ISL.API.isl_union_map_dump(waw_deps)
        ISL.API.isl_union_map_dump(war_deps)
    end

    """
    scheduling options
    """
    # maximize band size- better for tiling
    ISL.API.isl_options_set_schedule_maximize_band_depth(context, 1)
    # schedule weakly-connected components together
    ISL.API.isl_options_set_schedule_whole_component(context, 1)
    # avoid coalescing if possible
    ISL.API.isl_options_set_schedule_treat_coalescing(context, 1)
    # step over tile dim
    ISL.API.isl_options_set_tile_scale_tile_loops(context, 1)
    # don't start at 0 for tiles (cleans up code for 1-indexed Julia)
    ISL.API.isl_options_set_tile_shift_point_loops(context, 0)

    """
    construct schedule constraints from dependency analysis
    """
    # dependence constraints
    all_deps = ISL.API.isl_union_map_union(waw_deps, war_deps)
    all_deps = ISL.API.isl_union_map_union(all_deps, raw_deps)
    schedule_constraints = ISL.API.isl_schedule_constraints_on_domain(ISL.API.isl_union_set_copy(instructions))
    schedule_constraints = ISL.API.isl_schedule_constraints_set_validity(schedule_constraints, ISL.API.isl_union_map_copy(all_deps))

    # proximity constraints (keeps loops nested based on dependencies)
    schedule_constraints = ISL.API.isl_schedule_constraints_set_proximity(schedule_constraints, ISL.API.isl_union_map_copy(all_deps))

    # coincidence constraints (things to be scheduled together- doesn't usually do much)
    # schedule_constraints = ISL.API.isl_schedule_constraints_set_coincidence(schedule_constraints, ISL.API.isl_union_map_copy(all_deps))

    if verbose >= 3
        println("===SCHEDULE CONSTRAINTS===")
        ISL.API.isl_schedule_constraints_dump(schedule_constraints)
    end

    """
    compute the new schedule
    """
    schedule = ISL.API.isl_schedule_constraints_compute_schedule(schedule_constraints)
    if verbose >= 3
        println("===NEW SCHEDULE PRE TILING===")
        ISL.API.isl_schedule_dump(schedule)
    end

    """
    reordering of schedule loops
    """
    loop_orderings = get_best_nesting_orders(kernel)
    if verbose >= 2
        println("===LOOP ORDERINGS===")
        @show loop_orderings
    end
    reorder_schedule_loops(kernel, schedule, context, loop_orderings)
    if verbose >= 3
        println("===SCHEDULE AFTER REORDERING===")
        ISL.API.isl_schedule_dump(schedule)
    end

    """
    tiling of band nodes
    """
    if tile != 0
        schedule = tile_schedule(kernel, schedule, context, tile)
    end

    if verbose >= 2
        println("===NEW SCHEDULE===")
        ISL.API.isl_schedule_dump(schedule)
    end
    if verbose >= 3
        println("===DOMAIN AFTER TILING===")
        ISL.API.isl_union_set_dump(ISL.API.isl_schedule_get_domain(schedule))
    end

    """
    construct AST from new schedule
    """
    space = ISL.API.isl_set_read_from_str(context, "{:}")
    build = ISL.API.isl_ast_build_from_context(space)
    ast = ISL.API.isl_ast_build_node_from_schedule(build, schedule)

    if verbose >= 2
        # print C code
        c = ISL.API.isl_ast_node_get_ctx(ast)
        p = ISL.API.isl_printer_to_str(c)
        p = ISL.API.isl_printer_set_output_format(p, 4) # 4 = C code
        q = ISL.API.isl_printer_print_ast_node(p, ast)
        str = ISL.API.isl_printer_get_str(q)
        str = Base.unsafe_convert(Ptr{Cchar}, str)
        println("===FINAL C CODE===")
        println(Base.unsafe_string(str))
    end

    """
    parse back to Julia code
    """
    expr = parse_ast(ast, kernel, thread=thread, outermost=true)
    if verbose >= 1
        println("===FINAL JULIA CODE===")
        @show expr
    end
    return expr
end

"""
POLYHEDRAL MODEL GENERATION FUNCTIONS
Julia -> ISL
"""

"""
helper to turn gensyms into strings without special characters
"""
function sym_to_str(s::Symbol)::String
   str = string(s)
   str = replace.(str, r"#" => "")
   str = replace.(str, r":" => "")
   return str
end

"""
helper to return a unique identifier for a kernel existential quantifier
"""
function unique_identifier()::String
    id = gensym("id")
    return sym_to_str(id)
end

"""
helper to get string rep of params in kernel
"""
function get_params_str(kernel::LoopKernel)
    params = string(kernel.consts)
    params = replace.(params, r":" => "")
    params = replace.(params, r"Symbol" => "")
    return params
end

"""
format a condition on an instruction properly
"""
function format_instruction_cond(cond::Expr)::String
    stringcond = string(cond)
    stringcond = replace.(stringcond, r"%" => " mod ")
    stringcond = replace.(stringcond, r"==" => " = ")
    stringcond = replace.(stringcond, r"&&" => " and ")
    stringcond = replace.(stringcond, r"&" => " and ")
    stringcond = replace.(stringcond, "||" => " or ")
    stringcond = replace.(stringcond, "|" => " or ")
    return stringcond
end

"""
helper to get the domains related to instruction(s) in the ISL format
ex: 0 <= i <= n and 0 <= j <= n
"""
function get_instructions_domains(instructions::Vector{Instruction}, kernel::LoopKernel)::String

    conditions = ""
    count = 1
    for domain in kernel.domains
        for instruction in instructions
            if domain.iname in instruction.dependencies
                if count != 1
                    conditions = string(conditions, " and ")
                end
                count += 1
                instructioncond = ""
                if instruction.cond != :()
                    # instruction has separate domain condition
                    instructioncond = string(" and ", format_instruction_cond(instruction.cond))
                end
                if domain.step != 1
                    identifier = unique_identifier()
                    conditions = string(conditions, @sprintf("exists %s: %s = %s*%s + %s and %s <= %s <= %s", identifier, string(domain.iname), string(domain.step), identifier, string(domain.lowerbound), string(domain.lowerbound), string(domain.iname), string(domain.upperbound)), instructioncond)
                else
                    conditions = string(conditions, @sprintf("%s <= %s <= %s", string(domain.lowerbound), string(domain.iname), string(domain.upperbound)), instructioncond)
                end
            end
        end
    end
    return conditions
end

"""
get the related iteration spaces for an instruction
ex: [i, j]
if primes is true, returns [i', j']
"""
function get_related_domains(instruction::Instruction, kernel::LoopKernel, primes=false)
    count = 1
    domains = "["
    for domain in kernel.domains
        for iname in instruction.dependencies
            if domain.iname == iname
                if count != 1
                    domains = string(domains, ", ")
                end
                count += 1
                domains = string(domains, domain.iname)
                if primes
                    domains = string(domains, "\'")
                end
            end
        end
    end
    domains = string(domains, "]")
    return domains
end

"""
construct the instructions' domains representation in ISL syntax form
ex: (id = :mult) C[i, j] += A[i, k] * B[k, j] in matrix multiplication becomes
[n, m, r, A, B, C] -> {mult[i, j, k] : 0 <= i <= n and 0 <= j <= m and 0 <= k <= r}
"""
function instructions_isl_rep(kernel::LoopKernel)::String
    params = get_params_str(kernel)

    insn_map = "{"
    icount = 1
    for instruction in kernel.instructions
        count = 1
        name = string(sym_to_str(instruction.iname), get_related_domains(instruction, kernel))
        conditions = get_instructions_domains([instruction], kernel)

        if icount != 1
            insn_map = string(insn_map, "; ")
        end
        icount += 1

        if conditions != ""
            insn_map = string(insn_map, name, " : ", conditions)
        else
            insn_map = string(insn_map, name)
        end
    end
    insn_map = string(insn_map, "}")
    if insn_map == "{}"
        insn_map = "{:}"
    end
    return @sprintf("%s -> %s", params, insn_map)
end

"""
helpers for extracting accesses from an expression expr
"""
function expr_accesses(expr::Expr, kernel::LoopKernel)::Vector{Union{Expr, Symbol}}
    deps = []
    if expr.head == :call
        # RHS is operation on multiple terms, add all parts
        for arg in expr.args[2:end]
            append!(deps, expr_accesses(arg, kernel))
        end
    elseif expr.head == :macrocall
        # RHS is macro call on multiple terms, add all parts except macro and line number node
        for arg in expr.args[3:end]
            append!(deps, expr_accesses(arg, kernel))
        end
    else
        push!(deps, expr)
    end
    return deps
end

function expr_accesses(sym::Symbol, kernel::LoopKernel)::Vector{Union{Expr, Symbol}}
    if sym in kernel.consts
        return []
    end
    if sym in [d.iname for d in kernel.domains]
        return []
    end
    return [:($sym[])]
end

function expr_accesses(num::Number, kernel::LoopKernel)::Vector{Union{Expr, Symbol}}
    return []
end

function expr_accesses(l::LineNumberNode, kernel::LoopKernel)::Vector{Union{Expr, Symbol}}
    return []
end

"""
construct the access relations map in ISL syntax form

ex: (id = :mult) C[i, j] += A[i, k] * B[k, j] becomes
may_read:
[n, A, B, C] -> {
mult[i, j, k] -> A[k, i] : 0 <= i <= n and 0 <= j <= n and 0 <= k <= n;
mult[i, j, k] -> B[j, k] : 0 <= i <= n and 0 <= j <= n and 0 <= k <= n}
must_write:
[n, A, B, C] -> {mult[i, j, k] -> C[i, j] : 0 <= i <= n and 0 <= j <= n and 0 <= k <= n}
"""
function accesses_isl_rep(kernel::LoopKernel)::Tuple{String, String}
    params = get_params_str(kernel)

    may_read = string(params, " -> {")
    must_write = string(params, " -> {")
    rcount = 1
    wcount = 1
    mcount = 1
    for instruction in kernel.instructions
        ds = get_related_domains(instruction, kernel)
        reads = expr_accesses(instruction.body.args[2], kernel)
        writes = expr_accesses(instruction.body.args[1], kernel)
        head = instruction.body.head
        if head != :(=)
            # read from LHS also
            append!(reads, writes)
        end
        # must writes in LHS
        for part in writes
            if mcount != 1
                must_write = string(must_write, "; ")
            end
            mcount += 1
            must_write = string(must_write, sym_to_str(instruction.iname), ds, " -> ", part)
            # get domains
            conditions = get_instructions_domains([instruction], kernel)
            if conditions != ""
                must_write = string(must_write, " : ", conditions)
            end
        end
        # may reads in RHS and LHS
        for part in reads
            if rcount != 1
                may_read = string(may_read, "; ")
            end
            rcount += 1
            may_read = string(may_read, sym_to_str(instruction.iname), ds, " -> ", part)
            # get domains
            conditions = get_instructions_domains([instruction], kernel)
            if conditions != ""
                may_read = string(may_read, " : ", conditions)
            end
        end
    end

    may_read = string(may_read, "}")
    if may_read == "{}"
        may_read = "{:}"
    end
    must_write = string(must_write, "}")
    if must_write == "{}"
        must_write = "{:}"
    end
    return may_read, must_write
end

"""
Get the loop ordering dependencies for a kernel
For example, to specify the order j, k, i in a loop nest:
    "[n, r, m] -> {out253[i, j, k] -> out253[i', j', k']: j, k, i << j', k', i'}"
"""
function get_loop_ordering_deps_isl(kernel::LoopKernel, loop_orderings::Set{Vector{Symbol}})::String
    params = get_params_str(kernel)

    ordering_deps = "{"
    count = 0

    for order in loop_orderings
        for domain in kernel.domains
            if domain.iname in order
                for instruction in domain.instructions
                    if typeof(instruction) == Instruction
                        # found instruction to use
                        if count != 0
                            ordering_deps = string(ordering_deps, "; ")
                        end
                        count += 1
                        # sym[i, j] -> sym[i', j']
                        ordering_deps = string(ordering_deps, sym_to_str(instruction.iname), get_related_domains(instruction, kernel), " -> ", sym_to_str(instruction.iname), get_related_domains(instruction, kernel, true), ": ")
                        order_string = ""
                        order_string_prime = ""
                        ocount = 0
                        for iname in order
                            if iname in instruction.dependencies
                                if ocount != 0
                                    order_string = string(order_string, ", ")
                                    order_string_prime = string(order_string_prime, ", ")
                                end
                                ocount += 1
                                order_string = string(order_string, iname)
                                order_string_prime = string(order_string_prime, iname, "\'")
                            end
                        end
                        # i, j << i', j'
                        ordering_deps = string(ordering_deps, order_string, " << ", order_string_prime)
                        break
                    end
                end
            end
        end
    end

    ordering_deps = string(ordering_deps, "}")
    return @sprintf("%s -> %s", params, ordering_deps)
end

"""
construct a partial schedule for a grouping of instructions in domains
"""
function partial_schedule_isl(domains::Vector{Domain}, instructions::Vector{Instruction}, kernel::LoopKernel)::String

    mupa = "["

    dcount = 0
    for domain in domains
        set = ""

        icount = 0
        for instruction in instructions
            if domain.iname in instruction.dependencies
                if icount != 0
                    set = string(set, "; ")
                end
                icount += 1
                set = string(set, sym_to_str(instruction.iname), get_related_domains(instruction, kernel), "-> [(", domain.iname, ")]")
            end
        end

        if set != ""
            if dcount != 0
                mupa = string(mupa, ", ")
            end
            dcount += 1
            mupa = string(mupa, "{", set, "}")
        end
    end

    mupa = string(mupa, "]")

    return mupa
end

"""
construct an original schedule for dependency analysis
"""
function schedule_tree_isl_rep(context::Ptr{ISL.API.isl_ctx}, domain::Domain, kernel::LoopKernel, parent_doms=[])::Ptr{ISL.API.isl_schedule}

    new_parent_doms = [domain]
    append!(new_parent_doms, parent_doms)
    schedule = :nothing
    # navigate through original scheduling
    for instruction in domain.instructions
        if typeof(instruction) == Instruction
            # construct instruction schedule
            temp_kern = LoopKernel([instruction], kernel.domains, kernel.args, kernel.argtypes, kernel.consts)
            inst_dom = instructions_isl_rep(temp_kern)
            inst_dom = ISL.API.isl_union_set_read_from_str(context, inst_dom)
            inst_schedule = ISL.API.isl_schedule_from_domain(inst_dom)

            # put in sequence
            if schedule == :nothing
                schedule = inst_schedule
            else
                schedule = ISL.API.isl_schedule_sequence(schedule, inst_schedule)
            end
        else # Domain
            dom_schedule = schedule_tree_isl_rep(context, instruction, kernel, new_parent_doms)
            if schedule == :nothing
                schedule = dom_schedule
            else
                schedule = ISL.API.isl_schedule_sequence(schedule, dom_schedule)
            end
        end
    end

    # insert partial schedule
    insts = Instruction[]
    for instruction in kernel.instructions
        if domain.iname in instruction.dependencies
            push!(insts, instruction)
        end
    end
    partial = partial_schedule_isl([domain], insts, kernel)
    if partial != "[]"
        partial = ISL.API.isl_multi_union_pw_aff_read_from_str(context, partial)
        schedule = ISL.API.isl_schedule_insert_partial_schedule(schedule, partial)
    end

    if schedule == :nothing
        error("Mis-defined initial schedule! Check that instructions are properly nested")
    end

    return schedule
end

"""
PARSING FUNCTIONS
ISL -> Julia
"""

"""
parse an isl AST back into julia code
"""
function parse_ast(ast::Ptr{ISL.API.isl_ast_node}, kernel::LoopKernel; thread=false, outermost=false)::Expr
    node_type = ISL.API.isl_ast_node_get_type(ast)

    if node_type == ISL.API.isl_ast_node_for
        # for loop
        expr = parse_ast_for(ast, kernel, thread=thread, outermost=outermost)

    elseif node_type == ISL.API.isl_ast_node_if
        # if statement
        expr = parse_ast_if(ast, kernel, thread=thread)

    elseif node_type == ISL.API.isl_ast_node_block
        # compound node- multiple elements in block
        expr = parse_ast_block(ast, kernel, thread=thread)

    elseif node_type == ISL.API.isl_ast_node_mark
        # mark in the ast- comment

    elseif node_type == ISL.API.isl_ast_node_user
        # expression statement
        expr = parse_ast_user(ast, kernel, thread=thread)

    elseif node_type == ISL.API.isl_ast_node_error
        # error- quit

    end
end

"""
parse an ast node representing a for loop into an expression
"""
function parse_ast_for(ast::Ptr{ISL.API.isl_ast_node}, kernel::LoopKernel; thread=false, outermost=false)::Expr
    executes_once = ISL.API.isl_ast_node_for_is_degenerate(ast) # if loop executes only once

    # loop qualities
    iter = ISL.API.isl_ast_node_for_get_iterator(ast)
    id = ISL.API.isl_ast_expr_get_id(iter) # id of iterator
    ISL.API.isl_id_free(id)
    name = Base.unsafe_convert(Ptr{Cchar}, ISL.API.isl_id_get_name(id))
    name = Symbol(Base.unsafe_string(name)) # iterator symbol
    lb = parse_ast_expr(ISL.API.isl_ast_node_for_get_init(ast)) # initial condition
    body = parse_ast(ISL.API.isl_ast_node_for_get_body(ast), kernel) # body of loop

    if executes_once == ISL.API.isl_bool_true
        # loop only runs once, just set name = init and run body
        expr = quote
            $name = $init
            $body
        end
        return expr
    else
        step = parse_ast_expr(ISL.API.isl_ast_node_for_get_inc(ast)) # incremental step
        cond = parse_ast_expr(ISL.API.isl_ast_node_for_get_cond(ast)) # final condition
        ub = cond.args[3]
        if cond.args[1] == :(<)
            ub = :($ub - 1)
        elseif cond.args[1] != :(<=)
            error("unexpected conditional operator ", cond.head)
        end

        if outermost && thread
            if step == 1
                # for name = lb:up
                    # body
                expr = quote
                    Base.Threads.@threads for $name = $lb:$ub
                        $body
                    end
                end
                return expr

            else
                # for name = lb:step:up
                    # body
                expr = quote
                    Base.Threads.@threads for $name = $lb:$step:$ub
                        $body
                    end
                end
                return expr
            end
        else
            if step == 1
                # for name = lb:up
                    # body
                expr = quote
                    for $name = $lb:$ub
                        $body
                    end
                end
                return expr

            else
                # for name = lb:step:up
                    # body
                expr = quote
                    for $name = $lb:$step:$ub
                        $body
                    end
                end
                return expr
            end
        end
    end
end

"""
parse an ast node representing an if block into an expression
"""
function parse_ast_if(ast::Ptr{ISL.API.isl_ast_node}, kernel::LoopKernel; thread=false)::Expr
    cond = ISL.API.isl_ast_node_if_get_cond(ast)
    cond = parse_ast_expr(cond)
    if_body = ISL.API.isl_ast_node_if_get_then(ast)
    if_body = parse_ast(if_body, kernel, thread=thread)
    has_else = ISL.API.isl_ast_node_if_has_else(ast)
    if has_else == ISL.API.isl_bool_true
        # if cond
            # if_body
        # else
            # else_body
        # end
        else_body = ISL.API.isl_ast_node_if_get_else(ast)
        else_body = parse_ast(else_body, kernel, thread=thread)
        expr = quote
            if $cond
                $if_body
            else
                $else_body
            end
        end
        return expr
    else
        # if cond
            # if_body
        # end
        expr = quote
            if $cond
                $if_body
            end
        end
        return expr
    end
end

"""
parse an ast node representing a block into an expression
"""
function parse_ast_block(ast::Ptr{ISL.API.isl_ast_node}, kernel::LoopKernel; thread=false)::Expr
    children = ISL.API.isl_ast_node_block_get_children(ast)
    n = ISL.API.isl_ast_node_list_n_ast_node(children)
    exs = []
    for i=0:n-1
        child = ISL.API.isl_ast_node_list_get_at(children, i)
        push!(exs, parse_ast(child, kernel, thread=thread))
    end
    return quote $(exs...) end
end

"""
helpers for replacing symbols in expression
"""
function replace_expr_syms(ex::Expr, sym_dict::Dict{Symbol, Union{Symbol, Number, Expr}})::Expr
    new_args = []
    for arg in ex.args
        push!(new_args, replace_expr_syms(arg, sym_dict))
    end
    new_ex = Expr(ex.head, new_args...)
    return new_ex
end

function replace_expr_syms(sym::Symbol, sym_dict::Dict{Symbol, Union{Symbol, Number, Expr}})
    if haskey(sym_dict, sym)
        return sym_dict[sym]
    end
    return sym
end

function replace_expr_syms(num::Number, sym_dict::Dict{Symbol, Union{Symbol, Number, Expr}})
    return num
end

function replace_expr_syms(l::LineNumberNode, sym_dict::Dict{Symbol, Union{Symbol, Number, Expr}})
    return l
end

"""
parse an ast node representing an expression
"""
function parse_ast_user(ast::Ptr{ISL.API.isl_ast_node}, kernel::LoopKernel; thread=false)::Expr
    expr = ISL.API.isl_ast_node_user_get_expr(ast)
    ISL.API.isl_ast_expr_free(expr)
    first_expr = ISL.API.isl_ast_expr_op_get_arg(expr, 0) # first arg is "function" name (symbol cooresponding to instruction)
    n = ISL.API.isl_ast_expr_get_op_n_arg(expr)
    id = ISL.API.isl_ast_expr_id_get_id(first_expr)
    ISL.API.isl_ast_expr_free(first_expr)
    name = Base.unsafe_convert(Ptr{Cchar}, ISL.API.isl_id_get_name(id)) # name of identifier
    name = Base.unsafe_string(name)
    ISL.API.isl_id_free(id)

    # search for matching instruction
    for instruction in kernel.instructions
        if sym_to_str(instruction.iname) == name
            # found instruction
            # replace original iterators with new iterators
            old_ids = []
            for domain in kernel.domains
                for iname in instruction.dependencies
                    if domain.iname == iname
                        push!(old_ids, iname)
                    end
                end
            end
            new_ids = Dict{Symbol, Union{Symbol, Number, Expr}}()
            for i=1:Int(n)-1
                arg = ISL.API.isl_ast_expr_op_get_arg(expr, i)
                arg = parse_ast_expr(arg)
                new_ids[old_ids[i]] = arg
            end
            ex = replace_expr_syms(instruction.body, new_ids)
            return :(@inbounds $ex)
        end
    end

    return :()
end

"""
parse an ast expression condition
"""
function parse_ast_expr(expr::Ptr{ISL.API.isl_ast_expr})::Union{Symbol, Expr, Number}
    type = ISL.API.isl_ast_expr_get_type(expr)
    if type == ISL.API.isl_ast_expr_id
        id = ISL.API.isl_ast_expr_id_get_id(expr)
        name = Base.unsafe_convert(Ptr{Cchar}, ISL.API.isl_id_get_name(id)) # name of identifier
        name = Base.unsafe_string(name)
        ISL.API.isl_id_free(id)
        return :($(Symbol(name)))
    elseif type == ISL.API.isl_ast_expr_int
        val = ISL.API.isl_ast_expr_int_get_val(expr)
        num = ISL.API.isl_val_get_d(val)
        if Int(num) == num
            num = Int(num)
        end
        ISL.API.isl_val_free(val)
        return :($num)
    else
        if ISL.API.isl_ast_expr_get_op_n_arg(expr) == 1
            arg1 = parse_ast_expr(ISL.API.isl_ast_expr_op_get_arg(expr, 0))
            arg2 = :()
        else
            arg1 = parse_ast_expr(ISL.API.isl_ast_expr_op_get_arg(expr, 0))
            arg2 = parse_ast_expr(ISL.API.isl_ast_expr_op_get_arg(expr, 1))
        end

        op_type = ISL.API.isl_ast_expr_op_get_type(expr)
        if op_type == ISL.API.isl_ast_expr_op_and || op_type == ISL.API.isl_ast_expr_op_and_then
            return :($arg1 && $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_or || op_type == ISL.API.isl_ast_expr_op_or_else
            return :($arg1 || $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_max
            return :(max($arg1, $arg2))
        elseif op_type == ISL.API.isl_ast_expr_op_min
            return :(min($arg1, $arg2))
        elseif op_type == ISL.API.isl_ast_expr_op_minus
            return :(-$arg1)
        elseif op_type == ISL.API.isl_ast_expr_op_add
            return :($arg1 + $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_sub
            return :($arg1 - $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_mul
            return :($arg1 * $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_fdiv_q
            return :(Int(floor($arg1 / $arg2)))
        elseif op_type == ISL.API.isl_ast_expr_op_pdiv_q || op_type == ISL.API.isl_ast_expr_op_div
            return :(Int($arg1 / $arg2))
        elseif op_type == ISL.API.isl_ast_expr_op_pdiv_r || op_type == ISL.API.isl_ast_expr_op_zdiv_r
            return :($arg1 % $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_eq
            return :($arg1 == $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_le
            return :($arg1 <= $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_ge
            return :($arg1 >= $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_lt
            return :($arg1 < $arg2)
        elseif op_type == ISL.API.isl_ast_expr_op_gt
            return :($arg1 > $arg2)
        else
            # invalid
            return :()
        end
    end
end


module Mie

using MPI
using OffsetArrays
using Parameters
using ..Constants
using ..Data
using ..SpecialFunctions

export MieData

const ci = 1.0im

@with_kw mutable struct MieData
    number_eqns::Int64 = 0
    max_order::Int64 = 0
    number_spheres::Int64 = 0
    order::Array{Int64,1} = zeros(Int64, 0)
    offset::Array{Int64,1} = zeros(Int64, 0)
    block::Array{Int64,1} = zeros(Int64, 0)
    block_offset::Array{Int64,1} = zeros(Int64, 0)
    number_field_expansions::Array{Int64,1} = zeros(Int64, 0)
    tm_offset::Array{Int64,1} = zeros(Int64, 0)
    is_optically_active::Array{Bool,1} = fill(false, 0)
    tm_on_file::Array{Bool,1} = fill(false, 0)
    qext::Array{Float64,1} = zeros(0)
    qabs::Array{Float64,1} = zeros(0)
    an::Array{ComplexF64,1} = zeros(ComplexF64, 0)
    cn::Array{ComplexF64,1} = zeros(ComplexF64, 0)
    un::Array{ComplexF64,1} = zeros(ComplexF64, 0)
    vn::Array{ComplexF64,1} = zeros(ComplexF64, 0)
    dn::Array{ComplexF64,1} = zeros(ComplexF64, 0)
    an_inv::Array{ComplexF64,1} = zeros(ComplexF64, 0)
    stored_tm::Array{ComplexF64,1} = zeros(ComplexF64, 0)
end

function miecoefcalc(
    run_params::MSTMParameters,
    xsp::Array{Float64,1},
    ri::OffsetArray{ComplexF64,2},
    hostsphere::Array{Int64,1},
    numberfieldexp::Array{Int64,1},
    qeps::Float64;
    tmatrix::Union{TMaxtrix,Nothing} = nothing,
)::MieData
    nsphere = length(xsp)
    @assert size(ri) == (2, nsphere + 1)
    @assert length(hostsphere) == nsphere
    @assert length(numberfieldexp) == nsphere

    mie = MieData()
    mie.max_order = 0
    mie.number_spheres = nsphere
    mie.order = zeros(Int64, nsphere)
    mie.offset = zeros(Int64, nsphere + 1)
    mie.block = zeros(Int64, nsphere)
    mie.block_offset = zeros(Int64, nsphere + 1)
    mie.number_field_expansions = numberfieldexp
    mie.tm_offset = zeros(Int64, nsphere)
    mie.is_optically_active = fill(false, nsphere)
    mie.tm_on_file = fill(false, nsphere)
    mie.qext = zeros(nsphere)
    mie.qabs = zeros(nsphere)

    # TODO: Need to consider communicators other than COMM_WORLD
    rank = MPI.Comm_rank(MPI.COMM_WORLD)

    ntermstot = 0
    nblktot = 0
    sizetmstore = 0
    rihost = zeros(ComplexF64, 2)

    # TODO: Implement usage of T-Matrix files
    if !isnothing(tmatrix)
    end

    for i in 1:nsphere
        if !mie.tm_on_file[i]
            rihost = ri[:, hostsphere[i]]
            _, _, _, nodrn, _, _, _, _, _, _ = mieoa(xsp[i], parent(ri[:, i]), 0, qeps, ri_medium = parent(rihost))
            mie.order[i] = nodrn
        end
    end

    # Mie order of host sphere should be no lower than that of the inner sphere
    for i in 1:nsphere
        j = hostsphere[i]
        while j != 0
            mie.order[j] = max(mie.order[j], mie.order[i])
            j = hostsphere[j]
        end
    end

    for i in 1:nsphere
        rihost = ri[:, hostsphere[i]]
        mie.is_optically_active[i] = ri[1, i] != ri[2, i]

        # TODO: Implement usage of T-Matrix files
        if mie.tm_on_file[i]
        else
            qext, qsca, qabs, nodrn, _, _, _, _, _, _ =
                mieoa(xsp[i], parent(ri[:, i]), mie.order[i], 0.0, ri_medium = parent(rihost))
            nterms = 4nodrn
            mie.offset[i] = ntermstot
            ntermstot += nterms
            mie.qext[i] = qext
            mie.qabs[i] = qabs
            mie.block_offset[i] = nblktot
            mie.order[i] = nodrn
            mie.block[i] = nodrn * (nodrn + 2) * 2
        end

        nblktot += mie.block[i] * numberfieldexp[i]
        mie.max_order = max(mie.max_order, mie.order[i])
    end

    mie.offset[nsphere + 1] = ntermstot
    mie.block_offset[nsphere + 1] = nblktot
    mie.number_eqns = nblktot

    mie.an = zeros(ComplexF64, ntermstot)
    mie.cn = zeros(ComplexF64, ntermstot)
    mie.un = zeros(ComplexF64, ntermstot)
    mie.vn = zeros(ComplexF64, ntermstot)
    mie.dn = zeros(ComplexF64, ntermstot)
    mie.an_inv = zeros(ComplexF64, ntermstot)

    # TODO: mieoa() is called 3 times for each sphere, can there be any improvements?
    for i in 1:nsphere
        if !mie.tm_on_file[i]
            rihost = ri[:, hostsphere[i]]
            qext, qsca, qabs, nodrn, anp, cnp, unp, vnp, dnp, anpinv =
                mieoa(xsp[i], parent(ri[:, i]), mie.order[i], 0.0, ri_medium = parent(rihost))
            nterms = 4nodrn
            n1 = mie.offset[i] + 1
            n2 = mie.offset[i] + nterms
            mie.an[n1:n2] = reshape(anp, nterms)
            mie.cn[n1:n2] = reshape(cnp, nterms)
            mie.dn[n1:n2] = reshape(dnp, nterms)
            mie.un[n1:n2] = reshape(unp, nterms)
            mie.vn[n1:n2] = reshape(vnp, nterms)
            mie.an_inv[n1:n2] = reshape(anpinv, nterms)
        end
    end

    return mie
end

"""
    mieoa(x, ri, nodr0, qeps, ri_medium)

Optically active Lorenz/Mie coefficients
"""
function mieoa(
    x::Float64,
    ri::Array{ComplexF64,1},
    nodr0::Int64,
    qeps::Float64;
    ri_medium::Union{Array{ComplexF64,1},Nothing} = nothing,
)
    @assert length(ri) == 2
    @assert isnothing(ri_medium) || length(ri_medium) == 2

    rii = zeros(ComplexF64, 2, 2)
    rii[:, 1] = isnothing(ri_medium) ? [1.0 + 0.0im, 1.0 + 0.0im] : ri_medium
    rii[:, 2] = ri
    ribulk = 2.0 ./ (1.0 ./ rii[1, :] .+ 1.0 ./ rii[2, :])
    xri = x * rii
    nstop = if iszero(qeps)
        nodr0
    elseif qeps > 0.0
        floor(Int, x + 4x^(1 / 3)) + 5
    else
        ceil(Int, -qeps)
    end

    Ψ = OffsetArray(zeros(ComplexF64, nstop + 2, 2, 2), 0:(nstop + 1), 1:2, 1:2)
    Ξ = OffsetArray(zeros(ComplexF64, nstop + 2, 2, 2), 0:(nstop + 1), 1:2, 1:2)
    qext1 = zeros(nstop)
    for i in 1:2
        for p in 1:2
            Ψ[:, p, i] = cricbessel(nstop + 1, xri[p, i])
            Ξ[:, p, i] = crichankel(nstop + 1, xri[p, i])
        end
    end

    qext = 0.0
    qsca = 0.0

    Ψ_p = zeros(ComplexF64, 2, 2)
    Ξ_p = zeros(ComplexF64, 2, 2)
    amat = zeros(ComplexF64, 2, 2)
    bmat = zeros(ComplexF64, 2, 2)
    dmat = zeros(ComplexF64, 2, 2)
    gmat = zeros(ComplexF64, 2, 2)
    umat = zeros(ComplexF64, 2, 2)
    vmat = zeros(ComplexF64, 2, 2)
    anp = zeros(ComplexF64, 2, 2, nstop)
    cnp = zeros(ComplexF64, 2, 2, nstop)
    unp = zeros(ComplexF64, 2, 2, nstop)
    vnp = zeros(ComplexF64, 2, 2, nstop)
    dnp = zeros(ComplexF64, 2, 2, nstop)
    anpinv = zeros(ComplexF64, 2, 2, nstop)

    for n in 1:nstop
        for s in 1:2
            for t in 1:2
                Ψ_p[s, t] = Ψ[n - 1, s, t] - Ψ[n, s, t] / xri[s, t] * n
                Ξ_p[s, t] = Ξ[n - 1, s, t] - Ξ[n, s, t] / xri[s, t] * n
            end
        end

        for s in 1:2
            ss = (-1)^s
            for t in 1:2
                st = (-1)^t
                gmat[s, t] =
                    (ss * ribulk[1] + st * ribulk[2]) / (rii[s, 2] * rii[t, 1]) *
                    (Ψ_p[s, 2] * Ξ[n, t, 1] - ss * st * Ψ[n, s, 2] * Ξ_p[t, 1])
                amat[s, t] =
                    (ss * ribulk[1] + st * ribulk[2]) / (rii[s, 2] * rii[t, 1]) *
                    (Ψ_p[s, 2] * Ψ[n, t, 1] - ss * st * Ψ[n, s, 2] * Ψ_p[t, 1])
                umat[s, t] = (ss * ribulk[2] + st * ribulk[2]) / (rii[s, 2] * rii[t, 2]) * ci
                bmat[s, t] =
                    (ss * ribulk[2] + st * ribulk[1]) / (rii[s, 1] * rii[t, 2]) *
                    (Ξ_p[s, 1] * Ψ[n, t, 2] - ss * st * Ξ[n, s, 1] * Ψ_p[t, 2])
                dmat[s, t] = -(ss * ribulk[1] + st * ribulk[1]) / (rii[s, 1] * rii[t, 1]) * ci
                vmat[s, t] =
                    (ss * ribulk[2] + st * ribulk[1]) / (rii[s, 1] * rii[t, 2]) *
                    (Ξ_p[s, 1] * Ξ[n, t, 2] - ss * st * Ξ[n, s, 1] * Ξ_p[t, 2])
            end
        end

        gmatinv = twobytwoinverse(gmat)
        bmatinv = twobytwoinverse(bmat)

        amat = -gmatinv * amat
        umat = -gmatinv * umat
        dmat = -bmatinv * dmat
        vmat = -bmatinv * vmat

        anp[:, :, n] = amat
        dnp[:, :, n] = dmat
        unp[:, :, n] = umat
        vnp[:, :, n] = vmat
        amatinv = twobytwoinverse(amat)
        cnp[:, :, n] = dmat * amatinv
        amatinv = twobytwoinverse(amat)
        anpinv[:, :, n] = amatinv

        fn1 = 2n + 1
        for p in 1:2
            for q in 1:2
                qsca += fn1 * abs(amat[p, q]) * abs(amat[p, q])
            end
            qext1[n] -= fn1 * real(amat[p, p])
        end
        qext += qext1[n]
    end

    nodr = nstop
    if qeps > 0.0
        qextt = qext
        qext = 0.0
        for n in 1:nstop
            qext += qext1[n]
            err = abs(1.0 - qext / qextt)
            if err < qeps || n == nstop
                nodr = n
                break
            end
        end
    end

    qsca *= 2.0 / x / x
    qext *= 2.0 / x / x
    qabs = qext - qsca

    return qext,
    qsca,
    qabs,
    nodr,
    anp[:, :, 1:nodr],
    cnp[:, :, 1:nodr],
    unp[:, :, 1:nodr],
    vnp[:, :, 1:nodr],
    dnp[:, :, 1:nodr],
    anpinv[:, :, 1:nodr]
end

function onemiecoeffmult(mie::MieData, i::Int64, cx::OffsetArray{ComplexF64,3}, mie_coefficent::Char = 'a')
    @assert length(mie.order) >= i
    nodr = mie.order[i]

    @assert size(cx) == (nodr + 2, nodr, 2)

    nterms = 4nodr
    n1 = mie.offset[i] + 1
    n2 = mie.offset[i] + nterms

    an1 = if mie_coefficent == 'a'
        reshape(mie.an[n1:n2], 2, 2, nodr)
    elseif mie_coefficent == 'c'
        reshape(mie.cn[n1:n2], 2, 2, nodr)
    elseif mie_coefficent == 'd'
        reshape(mie.dn[n1:n2], 2, 2, nodr)
    elseif mie_coefficent == 'u'
        reshape(mie.un[n1:n2], 2, 2, nodr)
    else
        reshape(mie.vn[n1:n2], 2, 2, nodr)
    end

    cy = OffsetArray(zeros(ComplexF64, nodr + 2, nodr, 2), 0:(nodr + 1), 1:nodr, 1:2)

    for n in 1:nodr
        for p in 1:2
            for i in n:-1:1
                cy[n + 1, i, p] = an1[p, 1, n] * cx[n + 1, i, 1] + an1[p, 2, n] * cx[n + 1, i, 2]
            end
            for i in 0:n
                cy[i, n, p] = an1[p, 1, n] * cx[i, n, 1] + an1[p, 2, n] * cx[i, n, 2]
            end
        end
    end

    return cy
end

"""
    multmiecoeffmult(mie, nrhs, idir, ain, rhs_list)

Generalized Mie coefficient multiplication.

If idir != 1, we do the transpose.
"""
function multmiecoeffmult(
    mie::MieData,
    nrhs::Int64,
    idir::Int64,
    ain::Array{ComplexF64,1},
    rhs_list::Union{Array{Bool,1},Nothing} = nothing,
)
    neqns = mie.number_eqns
    @assert length(ain) == nrhs * neqns
    @assert isnothing(rhs_list) || length(rhs_list) == nrhs

    if isnothing(rhs_list)
        rhs_list = fill(true, nrhs)
    end

    noffi = 0
    aout = zeros(ComplexF64, nrhs * neqns)
    for i in 1:(mie.number_spheres)
        nodri = mie.order[i]
        nblki = 2nodri * (nodri + 2)
        b11 = noffi + 1

        # TODO: Implement usage of T-Matrix files
        if mie.tm_on_file[i]
        else
            aout_t = OffsetArray(zeros(ComplexF64, nodri + 2, nodri, 2, nrhs), 0:(nodri + 1), 1:nodri, 1:2, 1:nrhs)
            b12 = b11 + nblki * nrhs - 1
            gin_t = OffsetArray(reshape(ain[b11:b12], nodri + 2, nodri, 2, nrhs), 0:(nodri + 1), 1:nodri, 1:2, 1:nrhs)
            nterms = 4nodri
            n1 = mie.offset[i] + 1
            n2 = mie.offset[i] + nterms
            an1 = reshape(mie.an[n1:n2], 2, 2, nodri)
            if mie.number_field_expansions[i] == 1
                for j in 1:nrhs
                    if rhs_list[j]
                        for n in 1:nodri
                            for p in 1:2
                                for q in 1:2
                                    coeff = idir == 1 ? an1[p, q, n] : an1[q, p, n]
                                    for k in n:-1:1
                                        aout_t[n + 1, k, p, j] += gin_t[n + 1, k, q, j] * coeff
                                    end
                                    for k in 0:n
                                        aout_t[k, n, p, j] += gin_t[k, n, q, j] * coeff
                                    end
                                end
                            end
                        end
                    end
                end
                aout[b11:b12] = reshape(aout_t, nblki * nrhs)
            else
                # TODO: It seems that number of field expansions cannot exceed 2. Need to confirm this.

                b21 = noffi + nblki * nrhs + 1
                b22 = b21 + nblki * nrhs - 1
                bin_t =
                    OffsetArray(reshape(ain[b21:b22], nodri + 2, nodri, 2, nrhs), 0:(nodri + 1), 1:nodri, 1:2, 1:nrhs)
                dn1 = reshape(mie.dn[n1:n2], 2, 2, nodri)
                un1 = reshape(mie.un[n1:n2], 2, 2, nodri)
                vn1 = reshape(mie.vn[n1:n2], 2, 2, nodri)
                fout_t = OffsetArray(zeros(ComplexF64, nodri + 2, nodri, 2, nrhs), 0:(nodri + 1), 1:nodri, 1:2, 1:nrhs)
                for j in 1:nrhs
                    if rhs_list[j]
                        for n in 1:nodri
                            for p in 1:2
                                for q in 1:2
                                    coeff1, coeff2, coeff3, coeff4 =
                                        idir == 1 ? (an1[p, q, n], un1[p, q, n], dn1[p, q, n], vn1[p, q, n]) :
                                        (an1[q, p, n], dn1[q, p, n], un1[q, p, n], vn1[q, p, n])
                                    for k in n:-1:1
                                        aout_t[n + 1, k, p, j] +=
                                            gin_t[n + 1, k, q, j] * coeff1 + bin_t[n + 1, k, q, j] * coeff2
                                        fout_t[n + 1, k, p, j] +=
                                            gin_t[n + 1, k, q, j] * coeff3 + bin_t[n + 1, k, q, j] * coeff4
                                    end
                                    for k in 0:n
                                        aout_t[k, n, p, j] += gin_t[k, n, q, j] * coeff1 + bin_t[k, n, q, j] * coeff2
                                        fout_t[k, n, p, j] += gin_t[k, n, q, j] * coeff3 + bin_t[k, n, q, j] * coeff4
                                    end
                                end
                            end
                        end
                    end
                end
                aout[b11:b12] = reshape(aout_t, nblki * nrhs)
                aout[b21:b22] = reshape(fout_t, nblki * nrhs)
            end
        end
        noffi += nblki * nrhs * mie.number_field_expansions[i]
    end

    return aout
end

end # module Mie


const DEMCMC = DifferentialEvolutionMCMC

"""
    sample(model::DEModel, de::DE, n_iter::Int; progress=false, kwargs...)

Samples from the posterior distribution.

# Arguments

- `model`: a model containing likelihood function with data and priors
- `de`: differential evolution object
- `n_iter`: number of iterations or samples

# Keywords

- `progress=false`: show progress of sampler
- `kwargs...`: optional keyword arguments
"""
sample(model::DEModel, de::DE, n_iter::Int; progress=false, kwargs...) = _sample(model::DEModel, de::DE, n_iter::Int; progress, stepfun=step!, kwargs...)

function _sample(model::DEModel, de::DE, n_iter::Int; progress=false, stepfun=step!, kwargs...)
    meter = Progress(n_iter)
    # initialize particles based on prior distribution
    groups = sample_init(model, de, n_iter)
    for iter in 1:n_iter
        de.iter = iter + de.n_initial
        # explicitly pass groups so parallel works
        groups = stepfun(model, de, groups)
        progress ? next!(meter) : nothing
    end
    # convert to chain object
    chain = bundle_samples(model, de, groups, n_iter)
    return chain
end

"""
    sample(model::DEModel, de::DE, ::MCMCThreads, n_iter::Int; progress=false, kwargs...)

Samples from the posterior distribution with each group of particles on a seperarate thread for
the mutation and crossover steps.

# Arguments

- `model`: a model containing likelihood function with data and priors
- `de`: differential evolution object
- `MCMCThreads`: pass MCMCThreads() object to run on multiple threads
- `n_iter`: number of iterations or samples

# Keywords

- `progress=false`: show progress of sampler
- `kwargs...`: optional keyword arguments
"""
function sample(model::DEModel, de::DE, ::MCMCThreads, n_iter::Int; progress=false, kwargs...)
    _sample(model::DEModel, de::DE, n_iter::Int; progress, stepfun=pstep!, kwargs...)
end

"""
    step!(model::DEModel, de::DE, groups)

Perform a single step for DE-MCMC.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: DE-MCMC sampler object
- `groups`: Array of vectors of particles
"""
function step!(model::DEModel, de::DE, groups)
    rand() <= de.α ? migration!(de, groups) : nothing
    groups = update!(model, de, groups)
    store_samples!(de, groups)
    return groups
end

"""
    pstep!(model::DEModel, de::DE, groups)

Perform a single step for DE-MCMC with each particle group on a separate thread.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: DE-MCMC sampler object
- `groups`: Array of vectors of particles
"""
function pstep!(model::DEModel, de::DE, groups)
    rand() <= de.α ? migration!(de, groups) : nothing
    groups = p_update!(model, de, groups)
    store_samples!(de, groups)
    return groups
end

"""
    p_update!(model, de, groups)

Multithreaded update of particles. Particles are updated by block or simultaneously.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `group`: a vector of interacting particles (e.g. chains)
"""
function p_update!(model, de, groups)
    seeds = rand(UInt, length(groups))
    if de.blocking_on(de)
        Threads.@threads for g in 1:de.n_groups
            group = block_update!(model, de, groups[g], seeds[g])
        end
        return groups
    else
        Threads.@threads for g in 1:de.n_groups
            group = mutate_or_crossover!(model, de, groups[g], seeds[g])
        end
        return groups
    end
end

"""
    update!(model, de, groups)(model, de, groups)

Particles are updated by block or simultaneously on a single thread.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `group`: a vector of interacting particles (e.g. chains)
"""
function update!(model, de, groups)
    if de.blocking_on(de)
        return map(g -> block_update!(model, de, g), groups)
    else
        return map(g -> mutate_or_crossover!(model, de, g), groups)
    end
end

function block_update!(model, de, group, seed)
    Random.seed!(seed)
    return block_update!(model, de, group)
end

function block_update!(model, de, group)
    for block_idx in de.blocks
        mutate_or_crossover!(model, de, group, block_idx)
    end
    return group
end

"""
    mutate_or_crossover!(model, de, group, seed::Number)

Selects between mutation and crossover step with probability β.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `group`: a vector of interacting particles (e.g. chains)
- `seed::Number`: RNG seed
"""
function mutate_or_crossover!(model, de, group, seed::Number)
    Random.seed!(seed)
    mutate_or_crossover!(model, de, group)
    return group
end

function mutate_or_crossover!(model, de, group)
    rand() <= de.β ? mutation!(model, de, group) : crossover!(model, de, group)
    return group
end

function mutate_or_crossover!(model, de, group, block_idx)
    rand() <= de.β ? mutation!(model, de, group) : crossover!(model, de, group, block_idx)
    return group
end

"""
    bundle_samples(model::DEModel, de::DE, groups, n_iter)

Converts group particles to a chain object capable of generating convergence diagnostics
and posterior summaries.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `groups`: a vector of groups of particles
- `n_iter`: number of iterations
"""
function bundle_samples(model::DEModel, de::DE, groups, n_iter)
    @unpack burnin, discard_burnin = de
    particles = vcat(groups...)
    Np = length(particles)
    Ns = discard_burnin ? n_iter - burnin : n_iter
    offset = discard_burnin ? burnin : 0
    all_names = get_names(model, particles[1])
    n_parms = length(all_names)
    n_names = length(model.names)
    v = fill(0.0, Ns, n_parms, Np)
    for (c,p) in enumerate(particles)
        for s in 1:Ns
            sΔ = s + offset
            temp = Float64[]
            for ni in 1:n_names
                push!(temp, de.samples[sΔ,ni,c]...)
            end
            push!(temp, p.accept[sΔ], p.lp[sΔ])
            v[s,:,c] = temp'
        end
    end
    chains = Chains(v, all_names, (parameters = [model.names...],
        internals = ["acceptance", "lp"]))
    return chains
end

"""
    sample_init(model::DEModel, de::DE, n_iter)

Creates vectors of particles and samples initial parameter values from priors.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `n_iter`: number of iterations
"""
function sample_init(model::DEModel, de::DE, n_iter)
    de.samples = initialize_samples(de, model, n_iter)
    id = 0
    groups = [[init_particle(de, model, id+=1, n_iter) for p in 1:de.Np]
        for c in 1:de.n_groups]
    return groups
end


module MeasureSystems

#   This file is part of MeasureSystems.jl. It is licensed under the MIT license
#   MeasureSystems Copyright (C) 2021 Michael Reed

import Base: @pure
import UnitSystems
import UnitSystems: UnitSystem, Systems, Constants, Physics, Convert
import UnitSystems: Coupling, measure, unit, universe
export UnitSystems, Measure, measure, cache

# measure

using Measurements
struct Measure{N} end
const measure_cache = Measurement{Float64}[]
@pure measure(::Measure{N}) where N = measure_cache[N]
function cache(M::Measurement{Float64})
    N = findfirst(x->x==M,measure_cache)
    if isnothing(N)
        push!(measure_cache,M)
        N = Base.length(measure_cache)
    end
    return Measure{N}()
end
Base.show(io::IO,M::Measure{N}) where N = show(io,measure(M))

# unit systems

@pure mass(U::UnitSystem,S::UnitSystem) = electronmass(U,S)
@pure electronmass(𝘩,R∞) = αinv^2*R∞*2𝘩/𝘤
@pure electronmass(𝘩,R∞,C::Coupling) = inv(finestructure(C))^2*R∞*2𝘩/𝘤
@pure planckmass(U::UnitSystem,C::Coupling=universe(U)) = electronmass(U,C)/√coupling(C)
@pure planck(U::UnitSystem,C::Coupling=universe(U)) = 2π*planckreduced(U,C)
@pure newton(U::UnitSystem,C::Coupling=universe(U)) = lightspeed(U,C)*planckreduced(U,C)/planckmass(U,C)^2
@pure charge(U::UnitSystem,C::Coupling=universe(U)) = sqrt(2planck(U)/(permeability(U)/finestructure(U))/(lightspeed(U)*rationalization(U)*lorentz(U)^2))

for unit ∈ (:coupling,:finestructure,:electronunit,:protonunit,:protonelectron)
    @eval @pure $unit(C::Coupling) = UnitSystems.$unit(C)
    @eval @pure $unit(U::UnitSystem) = UnitSystems.$unit(universe(U))
end
for unit ∈ (:boltzmann,:planckreduced,:lightspeed,:permeability,:electronmass,:rationalization,:lorentz,:molarmass)
    @eval @pure $unit(U::UnitSystem) = UnitSystems.$unit(U)
end
for unit ∈ (:boltzmann,:planckreduced,:lightspeed,:permeability,:electronmass,:molarmass)
    @eval @pure $unit(U::UnitSystem,C::Coupling) = $unit(U)
end
for unit ∈ Constants
    @eval @pure $unit(U::UnitSystem,S::UnitSystem) = unit($unit(S)/$unit(U))
end
for unit ∈ Convert
    @eval begin
        @pure @inline $unit(v::Real,U::UnitSystem) = $unit(v,U,Metric)
        @pure @inline $unit(v::Real,U::UnitSystem,S::UnitSystem) = (u=$unit(U,S);isone(u) ? v : v/u)
        @pure @inline $unit(v::Real,U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ},S::UnitSystem{kB,ħ,𝘤,μ₀,mₑ}) where {kB,ħ,𝘤,μ₀,mₑ} = v
        @pure @inline $unit(U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ},S::UnitSystem{kB,ħ,𝘤,μ₀,mₑ}) where {kB,ħ,𝘤,μ₀,mₑ} = 1
    end
    if unit ∉ (Constants...,:permittivity,:charge,:magneticflux,:impedance,:conductance)
        @eval @pure @inline $unit(U::UnitSystem) = $unit(U,Metric)
    end
end
for unit ∈ (Systems...,Constants...,Physics...,Convert...)
    unit ∉ (:length,:time) && @eval export $unit
end

# common conversion factors

import UnitSystems: atm,g₀,lbm,slug,ft,ftUS,rankine,kelvin
for CAL ∈ (:cal,:calₜₕ,:cal₄,:cal₁₀,:cal₂₀,:calₘ,:calᵢₜ)
    KCAL = Symbol(:k,CAL)
    @eval import UnitSystems: $CAL, $KCAL
end

# fundamental constants, αinv = (34259-1/4366.8123)/250 # 137.036 exactly?

const mP = measurement("0.00000002176434(24)")
const αinv = measurement("137.035999084(21)")
const μₑᵤ = 1/measurement("1822.888486209(53)")
const μₚᵤ = measurement("1.007276466621(53)")
const R∞ = measurement("10973731.5681601(210)")
const KJ2014 = measurement("483597.8525(30)")*1e9
const RK2014 = measurement("25812.8074555(59)")

import UnitSystems: ΔνCs,Kcd,NA,kB,𝘩,𝘤,𝘦,ħ,Rᵤ,αL,RK1990,KJ1990,ħ1990
const mₑ,μ₀ = electronmass(𝘩,R∞),2𝘩/𝘤/αinv/𝘦^2 # ≈ 4π*(1e-7+5.5e-17), exact charge
const δμ₀,μₚₑ,αG,ħ2014 = μ₀-4π*1e-7,μₚᵤ/μₑᵤ,(mₑ/mP)^2,2/RK2014/KJ2014^2/π
const mₑ1990,mₑ2014 = electronmass(2π*ħ1990,R∞),electronmass(2π*ħ2014,R∞)

# engineering units # Thomson: αL = 1/2

const Universe = Coupling{cache(αG),cache(1/αinv),cache(μₑᵤ),cache(μₚᵤ)}()
const Gauss = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,1,cache(1000mₑ),4π,0.01/𝘤}()
const LorentzHeaviside = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,1,cache(1000mₑ),1,0.01/𝘤}()
const Thomson = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,1,cache(1000mₑ),4π,1/2}()
const Kennelly = UnitSystem{cache(Rᵤ*mₑ/μₑᵤ/0.001),ħ,𝘤,1e-7,cache(mₑ),4π}() # ?
const ESU = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,(100𝘤)^-2,cache(1000mₑ),4π}()
const ESU2019 = UnitSystem{1e7*kB,1e7*ħ,100𝘤,cache(1e3*μ₀/𝘤^2),cache(1000mₑ)}()
const EMU = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,1,cache(1000mₑ),4π}()
const EMU2019 = UnitSystem{1e7*kB,1e7*ħ,100𝘤,cache(1e7*μ₀),cache(1000mₑ)}()
const MTS = UnitSystem{cache(1e6*Rᵤ*mₑ/μₑᵤ),1000ħ,𝘤,4π/1e4,cache(mₑ/1000)}()
const Mixed = UnitSystem{cache(Rᵤ*mₑ/μₑᵤ/0.001),ħ,𝘤,cache(μ₀),cache(mₑ)}()
const Metric = UnitSystem{cache(Rᵤ*mₑ/μₑᵤ/0.001),ħ,𝘤,4π*1e-7,cache(mₑ)}()
const SI1976 = UnitSystem{cache(8.31432mₑ/μₑᵤ/0.001),ħ,𝘤,4π*1e-7,cache(mₑ)}()
const SI2019 = UnitSystem{kB,ħ,𝘤,cache(μ₀),cache(mₑ)}()
const CODATA = UnitSystem{cache(Rᵤ*mₑ2014/μₑᵤ/0.001),cache(ħ2014),𝘤,cache(2RK2014/𝘤/αinv),cache(mₑ2014)}()
const Conventional = UnitSystem{cache(Rᵤ*mₑ1990/μₑᵤ/0.001),ħ1990,𝘤,cache(2RK1990/𝘤/αinv),cache(mₑ1990)}()
const English = UnitSystem{kB*rankine/slug/ft^2,ħ/slug/ft^2,𝘤/ft,4π,cache(mₑ/slug)}()
const EnglishUS = UnitSystem{cache(1000Rᵤ*mₑ/μₑᵤ*rankine/slug/ftUS^2),ħ/slug/ftUS^2,𝘤/ftUS,4π,cache(mₑ/slug)}()

# astronomical units

const GMsun = measurement("1.32712440018(9)")*1e20
const GMearth = measurement("3.986004418(8)")*1e14
const GMjupiter = measurement("1.26686534(9)")*1e17

import UnitSystems: au,LD,day,pc,ly
const GG = 𝘤*ħ/mP^2; const mₛ = GMsun/GG; const Jₛ = mₛ*au^2/day^2; export mₛ,Jₛ,au,day
const IAU = UnitSystem{cache(Rᵤ*mₑ/μₑᵤ/0.001/Jₛ),cache(ħ/day/Jₛ),day*𝘤/au,cache(4π*1e-7*day^2/Jₛ),cache(mₑ/mₛ)}()

# aliased & humorous units

const mf = mass(90/lbm,Metric,English); const Jf = mf*(201.168/14day)^2
const FFF = UnitSystem{cache(1000Rᵤ*mₑ/μₑᵤ*rankine/Jf),cache(ħ/14day/Jf),14day*𝘤/201.168,0,cache(mₑ/mf)}()
const units, US, SI, MKS, temp = UnitSystem, UnitSystem, SI2019, Metric, temperature
const CGS, CGS2019, CGSm, CGSe, HLU = Gauss, EMU2019, EMU, ESU, LorentzHeaviside

# natural units

const Planck = UnitSystem{1,1,1,1,cache(√(4π*αG))}()
const PlanckGauss = UnitSystem{1,1,1,4π,cache(√αG)}()
const Stoney = UnitSystem{1,cache(αinv),1,4π,cache(√(αG*αinv))}()
const Hartree = UnitSystem{1,1,cache(αinv),cache(4π/αinv^2),1}()
const Rydberg = UnitSystem{1,1,cache(2αinv),cache(π/αinv^2),1/2}()
const Schrodinger = UnitSystem{1,1,cache(αinv),cache(4π/αinv^2),cache(√(αG*αinv))}()
const Electronic = UnitSystem{1,cache(αinv),1,4π,1}()
const Natural = UnitSystem{1,1,1,1,1}()
const NaturalGauss = UnitSystem{1,1,1,4π,1}()
const QCD = UnitSystem{1,1,1,1,cache(1/μₚₑ)}()
const QCDGauss = UnitSystem{1,1,1,4π,cache(1/μₚₑ)}()
const QCDoriginal = UnitSystem{1,1,1,cache(4π/αinv),cache(1/μₚₑ)}()

# physical constants

@pure electronmass(U::typeof(Planck),C::Coupling) = sqrt(4π*coupling(C))
@pure electronmass(U::typeof(PlanckGauss),C::Coupling) = sqrt(coupling(C))
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(√(αG*αinv))},C::Coupling) where {kB,ħ,𝘤,μ₀} = sqrt(coupling(C)/finestructure(C))
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(1/μₚₑ)},C::Coupling) where {kB,ħ,𝘤,μ₀} = 1/protonelectron(C)
@pure permeability(U::UnitSystem{kB,ħ,𝘤,cache(4π/αinv^2)},C::Coupling) where {kB,ħ,𝘤} = 4π*finestructure(C)^2
@pure permeability(U::UnitSystem{kB,ħ,𝘤,cache(π/αinv^2)},C::Coupling) where {kB,ħ,𝘤} = π*finestructure(C)^2
@pure lightspeed(U::UnitSystem{kB,ħ,cache(αinv)},C::Coupling) where {kB,ħ} = 1/finestructure(C)
@pure lightspeed(U::UnitSystem{kB,ħ,cache(2αinv)},C::Coupling) where {kB,ħ} = 2/finestructure(C)
@pure planckreduced(U::UnitSystem{kB,cache(αinv)},C::Coupling) where kB = 1/finestructure(C)

@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(mₑ)},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),R∞,C)
@pure electronmass(U::UnitSystem{kB,ħ,100𝘤,μ₀,cache(1000mₑ)},C::Coupling) where {kB,ħ,μ₀} = 1000electronmass(SI,C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(mₑ/1000)},C::Coupling) where {kB,ħ,μ₀} = electronmass(SI,C)/1000
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(mₑ2014)},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),R∞,C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(mₑ1990)},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),R∞,C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤/ftUS,μ₀,cache(mₑ/slug)},C::Coupling) where {kB,ħ,μ₀} = electronmass(SI,C)/slug
@pure permeability(U::UnitSystem{kB,ħ,𝘤,cache(μ₀)},C::Coupling) where {kB,ħ,𝘤} = finestructure(C)*2𝘩/𝘤/𝘦^2
@pure permeability(U::typeof(ESU2019),C::Coupling) = 1e3*permeability(SI,C)/𝘤^2
@pure permeability(U::typeof(EMU2019),C::Coupling) = 1e7*permeability(SI,C)
@pure permeability(U::typeof(CODATA),C::Coupling) = 2RK2014*finestructure(C)/𝘤
@pure permeability(U::typeof(Conventional),C::Coupling) = 2RK1990*finestructure(C)/𝘤

@pure molarmass(U::UnitSystem{kB},C::Coupling=universe(U)) = NA*electronmass(U,C)/electronunit(C)
@pure molarmass(U::UnitSystem{1e7*kB},C::Coupling=universe(U)) = 1000molarmass(SI2019,C)
@pure molarmass(U::UnitSystem{1e3*kB},C::Coupling=universe(U)) = molarmass(SI2019,C)/1000
@pure molarmass(U::UnitSystem{cache(boltzmann(MTS))}) = molarmass(CGS)/1e6
@pure molarmass(U::UnitSystem{cache(boltzmann(CGS))}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{cache(boltzmann(FFF))}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{boltzmann(English)},C::Coupling=universe(U)) = 1000molarmass(SI2019,C)
@pure molarmass(U::UnitSystem{cache(boltzmann(EnglishUS))}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{cache(boltzmann(IAU))}) = 1/1000mₛ

@pure luminousefficacy(U::UnitSystem{1}) = 1
@pure luminousefficacy(U::UnitSystem) = power(Kcd,SI2019,U)

@pure universe(::typeof(FFF)) = Universe
for unit ∈ UnitSystems.Systems
    if unit ∉ (:Natural,:NaturalGauss)
        @eval universe(::typeof($unit)) = Universe
    end
end

const dir = dirname(pathof(UnitSystems))
include("$dir/kinematic.jl")
include("$dir/electromagnetic.jl")
include("$dir/thermodynamic.jl")
include("$dir/physics.jl")
include("$dir/systems.jl")

end # module


#####
##### Tendency specification
#####

#####
##### Sources
#####

eq_tends(pv::AbstractPrognosticVariable, m::AtmosModel, tt::Source) =
    (m.source[pv]..., eq_tends(pv, turbconv_model(m), tt)...)

#####
##### First order fluxes
#####

eq_tends(::Mass, ::Anelastic1D, ::Flux{FirstOrder}) = ()

eq_tends(::Mass, ::Compressible, ::Flux{FirstOrder}) = (Advect(),)

# Mass
eq_tends(pv::Mass, atmos::AtmosModel, tt::Flux{FirstOrder}) =
    (eq_tends(pv, compressibility_model(atmos), tt))

# Momentum
eq_tends(pv::Momentum, ::Compressible, ::Flux{FirstOrder}) =
    (PressureGradient(),)

eq_tends(pv::Momentum, ::Anelastic1D, ::Flux{FirstOrder}) = ()

eq_tends(pv::Momentum, m::AtmosModel, tt::Flux{FirstOrder}) =
    (Advect(), eq_tends(pv, compressibility_model(m), tt)...)

# Energy
eq_tends(::Energy, m::TotalEnergyModel, tt::Flux{FirstOrder}) =
    (Advect(), Pressure())

eq_tends(::ρθ_liq_ice, m::θModel, tt::Flux{FirstOrder}) = (Advect(),)

# TODO: make radiation aware of which energy formulation is used:
# eq_tends(pv::PV, m::AtmosModel, tt::Flux{FirstOrder}) where {PV <: AbstractEnergyVariable} =
#     (eq_tends(pv, energy_model(m), tt)..., eq_tends(pv, energy_model(m), radiation_model(m), tt)...)
eq_tends(pv::AbstractEnergyVariable, m::AtmosModel, tt::Flux{FirstOrder}) = (
    eq_tends(pv, energy_model(m), tt)...,
    eq_tends(pv, radiation_model(m), tt)...,
)

# AbstractMoistureVariable
eq_tends(::AbstractMoistureVariable, ::AtmosModel, ::Flux{FirstOrder}) =
    (Advect(),)

# AbstractPrecipitationVariable
eq_tends(
    pv::AbstractPrecipitationVariable,
    m::AtmosModel,
    tt::Flux{FirstOrder},
) = (eq_tends(pv, precipitation_model(m), tt)...,)

# Tracers
eq_tends(pv::Tracers{N}, ::AtmosModel, ::Flux{FirstOrder}) where {N} =
    (Advect(),)

#####
##### Second order fluxes
#####

eq_tends(
    ::Union{Mass, Momentum, AbstractMoistureVariable},
    ::DryModel,
    ::Flux{SecondOrder},
) = ()
eq_tends(
    ::Union{Mass, Momentum, AbstractMoistureVariable},
    ::AbstractMoistureModel,
    ::Flux{SecondOrder},
) = (MoistureDiffusion(),)

# Mass
eq_tends(pv::Mass, m::AtmosModel, tt::Flux{SecondOrder}) =
    (eq_tends(pv, moisture_model(m), tt)...,)

# Momentum
eq_tends(pv::Momentum, m::AtmosModel, tt::Flux{SecondOrder}) = (
    ViscousStress(),
    eq_tends(pv, moisture_model(m), tt)...,
    eq_tends(pv, turbconv_model(m), tt)...,
    eq_tends(pv, hyperdiffusion_model(m), tt)...,
)

# Energy
eq_tends(::Energy, m::TotalEnergyModel, tt::Flux{SecondOrder}) =
    (ViscousFlux(), DiffEnthalpyFlux())

eq_tends(::ρθ_liq_ice, m::θModel, tt::Flux{SecondOrder}) = (ViscousFlux(),)

eq_tends(pv::AbstractEnergyVariable, m::AtmosModel, tt::Flux{SecondOrder}) = (
    eq_tends(pv, energy_model(m), tt)...,
    eq_tends(pv, turbconv_model(m), tt)...,
    eq_tends(pv, hyperdiffusion_model(m), tt)...,
)

# AbstractMoistureVariable
eq_tends(pv::AbstractMoistureVariable, m::AtmosModel, tt::Flux{SecondOrder}) = (
    eq_tends(pv, moisture_model(m), tt)...,
    eq_tends(pv, turbconv_model(m), tt)...,
    eq_tends(pv, hyperdiffusion_model(m), tt)...,
)

# AbstractPrecipitationVariable
eq_tends(
    pv::AbstractPrecipitationVariable,
    m::AtmosModel,
    tt::Flux{SecondOrder},
) = (eq_tends(pv, precipitation_model(m), tt)...,)

# Tracers
eq_tends(pv::Tracers{N}, m::AtmosModel, tt::Flux{SecondOrder}) where {N} =
    (eq_tends(pv, tracer_model(m), tt)...,)


# convert CorrData to SeisData
using SeisIO, SeisNoise, Dates, Logging
export corr2seis
""" Example of CorrData

CorrData with 1 Corrs
      NAME: "YW.FACR..BHZ.YW.FACR..BHZ"
        ID: "2010-01-01"
       LOC: 0.0 N, 0.0 E, 0.0 m
      COMP: "ZZ"
   ROTATED: false
 CORR_TYPE: "cross-correlation"
        FS: 5.0
      GAIN: 1.0
   FREQMIN: 0.1
   FREQMAX: 2.4
    CC_LEN: 3600
   CC_STEP: 1800
  WHITENED: false
 TIME_NORM: false
      RESP: c = 0.0, 0 zeros, 0 poles
      MISC: 2 entries
     NOTES: 3 entries
    MAXLAG: 100.0
         T: 2010-01-01T00:00:00.000            …
      CORR: 1001×1 Array{Float32,2}
"""

""" Example of CorrData.misc
Dict{String,Any} with 27 entries:
  "e"        => 86399.8
  "dist"     => 0.0
  "scale"    => 1.0
  "location" => Dict("YW.FACO..BHZ"=>0.0 N, 0.0 E, 0.0 m,"TA.G05D..BHZ"=>0.0 N, 0.0 E, 0.0 m)
  "iztype"   => 9
  "khole"    => "       \xb0"
  "lcalda"   => 0
  "nzhour"   => 0
  "nvhdr"    => 6
  "nzjday"   => 40
  "kcmpnm"   => "BHZ\0\0\0\0\0"
  "depmin"   => -7.54567e-8
  "iftype"   => 1
  "depmax"   => 9.53069e-8
  "nzmin"    => 0
  "b"        => 0.0
  "delta"    => 0.2
  "npts"     => 432000
  "kstnm"    => "G05D\0\0\0\0"
  "nzyear"   => 2010
  "nzmsec"   => 0
  "leven"    => 1
  "lpspol"   => 1
  "depmen"   => -7.11344e-16
  "nzsec"    => 0
  "lovrok"   => 1
  "knetwk"   => "TA\0\0\0\0\0\0"
"""

""" Example of SeisData

SeisChannel with 432000 samples
    ID: YW.FACO..BHZ
  NAME:
   LOC: 0.0 N, 0.0 E, 0.0 m
    FS: 5.0
  GAIN: 1.0
  RESP: c = 0.0, 0 zeros, 0 poles
 UNITS:
   SRC: /Users/xiaotaoyang/Work/Cascadia/Raw/Event_2010_001/YW.FACO.BHZ.sac
  MISC: 0 entries
 NOTES: 1 entries
     T: 2010-01-01T00:00:00.010 (0 gaps)
     X: +1.595e-18
        +1.421e-14
            ...
        +6.933e-11
        (nx = 432000)

"""
# richmisc: true if include location
function corr2seis(C::CorrData)
    #get the size of CorrData, to determine how many channels needed.

    #create an empty SeisData object with the right size.
    S = SeisData(size(C.corr,2))

    #loop through the CorrData (corr) and pass the header values to SeisData object
    # in loop, save customzed headers to misc dictionary.
    # save freqmin and freqmax twice:
    # 1. use real field names. 2. use USER0 and USER1, to be saved to SAC for later use.
    stemp = split(C.name,".")
    evname = join([stemp[1],stemp[2],stemp[3],stemp[4]],".")  #use the first entry in location as the event name
    stname = join([stemp[5],stemp[6],stemp[7],stemp[8]],".")
    evla = 0.0
    evlo = 0.0
    evdp = 0.0
    stla = 0.0
    stlo = 0.0
    stel = 0.0
    if in("location",collect(keys(C.misc)))
        # println(C.name,": ",collect(keys(C.misc["location"])))
        # evname = collect(keys(C.misc["location"]))[1]  #use the first entry in location as the event name
        # stname = collect(keys(C.misc["location"]))[2]
        evla = C.misc["location"][evname].lat
        evlo = C.misc["location"][evname].lon
        evdp = C.misc["location"][evname].el
        stla = C.misc["location"][stname].lat
        stlo = C.misc["location"][stname].lon
        stel = C.misc["location"][stname].el
    else
        @warn "$(C.name) : 'location' field not defined in misc. Use default locations (0.0)."
    end

    for i = 1:S.n
        SC = SeisChannel();
        # for tv in C.corr[:,i]
        #     push!(SC.x,tv)
        # end
        SC.x = C.corr[:,i]
        SC.id = C.id
        SC.name = C.name
        setfield!(SC.loc,:lat,stla)
        setfield!(SC.loc,:lon,stlo)
        setfield!(SC.loc,:dep,stel)
        # = C.loc
        SC.fs = C.fs
        SC.gain = C.gain
        SC.resp = C.resp
        SC.src = C.name
        SC.t = [1 Int64(C.t[i]*1e6);size(C.corr,1) 0]
        SC.notes = C.notes
        # TO-DO get misc values
        SC.misc = C.misc

        #update some values in misc for SAC headers.
        SC.misc["b"] = -1.0*C.maxlag
        SC.misc["e"] = C.maxlag
        SC.misc["delta"] = 1.0/C.fs
        SC.misc["npts"] = size(C.corr,1)
        SC.misc["depmin"] = minimum(C.corr[:,i])
        SC.misc["depmax"] = maximum(C.corr[:,i])
        SC.misc["depmen"] = sum(C.corr[:,i])/size(C.corr,1)
        SC.misc["kevnm"] = evname
        SC.misc["evla"] = evla
        SC.misc["evlo"] = evlo
        SC.misc["evdp"] = evdp
        SC.misc["stla"] = stla
        SC.misc["stlo"] = stlo
        SC.misc["stel"] = stel
        SC.misc["nzyear"] = Dates.year(u2d(C.t[i]))
        SC.misc["nzjday"] = Dates.date2epochdays(Date(u2d(C.t[i]))) - Dates.date2epochdays(Date(string(Dates.year(u2d(C.t[i])))*"-01-01")) + 1
        SC.misc["nzhour"] = Dates.hour(u2d(C.t[i]))
        SC.misc["nzmin"] = Dates.minute(u2d(C.t[i]))
        SC.misc["nzsec"] = Dates.second(u2d(C.t[i]))
        SC.misc["nzmsec"] = Dates.millisecond(u2d(C.t[i]))
        SC.misc["comp"] = C.comp
        SC.misc["rotated"] = C.rotated
        SC.misc["corr_type"] =  C.corr_type
        # SC.misc["kuser0"] =  C.corr_type
        SC.misc["freqmin"] = C.freqmin
        SC.misc["freqmax"] = C.freqmax
        # SC.misc["user0"] = C.freqmin
        # SC.misc["user1"] = C.freqmax
        SC.misc["cc_len"] = C.cc_len
        SC.misc["cc_step"] = C.cc_step
        # SC.misc["user2"] = C.cc_len
        # SC.misc["user3"] = C.cc_step
        # SC.misc["whitened"] = C.whitened
        # SC.misc["time_norm"] = C.time_norm
        # if C.whitened
        #     SC.misc["kuser1"] =  "NOWHITND"
        # else
        #     SC.misc["kuser1"] =  "WHITENED"
        # end
        # if C.time_norm
        #     SC.misc["kuser2"] =  "TIMENORM"
        # else
        #     SC.misc["kuser2"] =  "NOTIMENM"
        # end
        if C.whitened
            SC.misc["whitened"] =  "WHITENED"
        else
            SC.misc["whitened"] =  "NOWHITND"
        end
        if typeof(C.time_norm) == Bool
            if C.time_norm
                SC.misc["time_norm"] =  "TIMENORM"
            else
                SC.misc["time_norm"] =  "NOTIMENM"
            end
        else
            SC.misc["time_norm"] =  C.time_norm
        end
        S[i] = SC
    end

    return(S)
end