function update_aux!(edmf::EDMF_PrognosticTKE{N_up}, gm, grid, state, Case, param_set, TS) where {N_up}
    #####
    ##### Unpack common variables
    #####
    kc_surf = kc_surface(grid)
    kf_surf = kf_surface(grid)
    kc_toa = kc_top_of_atmos(grid)
    up = edmf.UpdVar
    en = edmf.EnvVar
    en_thermo = edmf.EnvThermo
    ρ0_f = face_ref_state(state).ρ0
    p0_c = center_ref_state(state).p0
    ρ0_c = center_ref_state(state).ρ0
    α0_c = center_ref_state(state).α0
    g = CPP.grav(param_set)
    c_m = CPEDMF.c_m(param_set)
    KM = center_aux_turbconv(state).KM
    KH = center_aux_turbconv(state).KH
    surface = Case.Sur
    obukhov_length = surface.obukhov_length
    FT = eltype(grid)
    prog_gm = center_prog_grid_mean(state)
    prog_gm_f = face_prog_grid_mean(state)
    aux_up = center_aux_updrafts(state)
    aux_up_f = face_aux_updrafts(state)
    aux_en = center_aux_environment(state)
    aux_en_f = face_aux_environment(state)
    aux_gm = center_aux_grid_mean(state)
    aux_tc_f = face_aux_turbconv(state)
    aux_tc = center_aux_turbconv(state)
    aux_bulk = center_aux_bulk(state)
    prog_en = center_prog_environment(state)
    aux_en_2m = center_aux_environment_2m(state)
    prog_up = center_prog_updrafts(state)
    prog_up_f = face_prog_updrafts(state)
    aux_en_unsat = aux_en.unsat
    aux_en_sat = aux_en.sat
    m_entr_detr = aux_tc.ϕ_temporary
    ∇m_entr_detr = aux_tc.ψ_temporary
    w_up_c = aux_tc.w_up_c
    w_en_c = aux_tc.w_en_c
    wvec = CC.Geometry.WVector
    Δt = TS.dt
    max_area = edmf.max_area
    sde_model = edmf.sde_model

    #####
    ##### center variables
    #####

    @inbounds for k in real_center_indices(grid)
        #####
        ##### Set primitive variables
        #####
        @inbounds for i in 1:N_up
            if is_surface_center(grid, k)
                if prog_up[i].ρarea[k] / ρ0_c[k] >= edmf.minimum_area
                    aux_up[i].θ_liq_ice[k] = edmf.h_surface_bc[i]
                    aux_up[i].q_tot[k] = edmf.qt_surface_bc[i]
                    aux_up[i].area[k] = edmf.area_surface_bc[i]
                else
                    aux_up[i].θ_liq_ice[k] = prog_gm.θ_liq_ice[k]
                    aux_up[i].q_tot[k] = prog_gm.q_tot[k]
                end
            else
                if prog_up[i].ρarea[k] / ρ0_c[k] >= edmf.minimum_area
                    aux_up[i].θ_liq_ice[k] = prog_up[i].ρaθ_liq_ice[k] / prog_up[i].ρarea[k]
                    aux_up[i].q_tot[k] = prog_up[i].ρaq_tot[k] / prog_up[i].ρarea[k]
                    aux_up[i].area[k] = prog_up[i].ρarea[k] / ρ0_c[k]
                else
                    aux_up[i].θ_liq_ice[k] = prog_gm.θ_liq_ice[k]
                    aux_up[i].q_tot[k] = prog_gm.q_tot[k]
                    aux_up[i].area[k] = 0
                end
            end
        end

        #####
        ##### compute bulk
        #####
        aux_bulk.q_tot[k] = 0
        aux_bulk.θ_liq_ice[k] = 0
        aux_bulk.area[k] = sum(i -> aux_up[i].area[k], 1:N_up)
        if aux_bulk.area[k] > 0
            @inbounds for i in 1:N_up
                aux_bulk.q_tot[k] += aux_up[i].area[k] * aux_up[i].q_tot[k] / aux_bulk.area[k]
                aux_bulk.θ_liq_ice[k] += aux_up[i].area[k] * aux_up[i].θ_liq_ice[k] / aux_bulk.area[k]
            end
        else
            aux_bulk.q_tot[k] = prog_gm.q_tot[k]
            aux_bulk.θ_liq_ice[k] = prog_gm.θ_liq_ice[k]
        end
        aux_en.area[k] = 1 - aux_bulk.area[k]
        aux_en.tke[k] = prog_en.ρatke[k] / (ρ0_c[k] * aux_en.area[k])
        aux_en.Hvar[k] = prog_en.ρaHvar[k] / (ρ0_c[k] * aux_en.area[k])
        aux_en.QTvar[k] = prog_en.ρaQTvar[k] / (ρ0_c[k] * aux_en.area[k])
        aux_en.HQTcov[k] = prog_en.ρaHQTcov[k] / (ρ0_c[k] * aux_en.area[k])

        #####
        ##### decompose_environment
        #####
        a_bulk_c = aux_bulk.area[k]
        val1 = 1 / (1 - a_bulk_c)
        val2 = a_bulk_c * val1
        aux_en.q_tot[k] = max(val1 * prog_gm.q_tot[k] - val2 * aux_bulk.q_tot[k], 0) #Yair - this is here to prevent negative QT
        aux_en.θ_liq_ice[k] = val1 * prog_gm.θ_liq_ice[k] - val2 * aux_bulk.θ_liq_ice[k]

        #####
        ##### saturation_adjustment and buoyancy
        #####
        ts_gm = thermo_state_pθq(param_set, p0_c[k], prog_gm.θ_liq_ice[k], prog_gm.q_tot[k])
        aux_tc.θ_virt[k] = TD.virtual_pottemp(ts_gm)

        ts_en = thermo_state_pθq(param_set, p0_c[k], aux_en.θ_liq_ice[k], aux_en.q_tot[k])

        aux_en.T[k] = TD.air_temperature(ts_en)
        aux_en.θ_virt[k] = TD.virtual_pottemp(ts_en)
        aux_en.θ_dry[k] = TD.dry_pottemp(ts_en)
        aux_en.q_liq[k] = TD.liquid_specific_humidity(ts_en)
        aux_en.q_ice[k] = TD.ice_specific_humidity(ts_en)
        rho = TD.air_density(ts_en)
        aux_en.buoy[k] = buoyancy_c(param_set, ρ0_c[k], rho)

        update_sat_unsat(en_thermo, state, k, ts_en)
        aux_en.RH[k] = TD.relative_humidity(ts_en)

        @inbounds for i in 1:N_up
            if aux_up[i].area[k] < edmf.minimum_area && k > kc_surf && aux_up[i].area[k - 1] > 0.0
                qt = aux_up[i].q_tot[k - 1]
                h = aux_up[i].θ_liq_ice[k - 1]
                ts_up = thermo_state_pθq(param_set, p0_c[k], h, qt)
            else
                ts_up = thermo_state_pθq(param_set, p0_c[k], aux_up[i].θ_liq_ice[k], aux_up[i].q_tot[k])
            end
            aux_up[i].q_liq[k] = TD.liquid_specific_humidity(ts_up)
            aux_up[i].q_ice[k] = TD.ice_specific_humidity(ts_up)
            aux_up[i].T[k] = TD.air_temperature(ts_up)
            ρ = TD.air_density(ts_up)
            aux_up[i].buoy[k] = buoyancy_c(param_set, ρ0_c[k], ρ)
            aux_up[i].RH[k] = TD.relative_humidity(ts_up)
        end
        aux_gm.buoy[k] = (1.0 - aux_bulk.area[k]) * aux_en.buoy[k]
        @inbounds for i in 1:N_up
            aux_gm.buoy[k] += aux_up[i].area[k] * aux_up[i].buoy[k]
        end
        @inbounds for i in 1:N_up
            aux_up[i].buoy[k] -= aux_gm.buoy[k]
        end
        aux_en.buoy[k] -= aux_gm.buoy[k]

        #####
        ##### compute bulk thermodynamics
        #####
        aux_bulk.q_liq[k] = 0
        aux_bulk.q_ice[k] = 0
        aux_bulk.T[k] = 0
        aux_bulk.RH[k] = 0
        aux_bulk.buoy[k] = 0
        if a_bulk_c > 0
            @inbounds for i in 1:N_up
                aux_bulk.q_liq[k] += aux_up[i].area[k] * aux_up[i].q_liq[k] / a_bulk_c
                aux_bulk.q_ice[k] += aux_up[i].area[k] * aux_up[i].q_ice[k] / a_bulk_c
                aux_bulk.T[k] += aux_up[i].area[k] * aux_up[i].T[k] / a_bulk_c
                aux_bulk.RH[k] += aux_up[i].area[k] * aux_up[i].RH[k] / a_bulk_c
                aux_bulk.buoy[k] += aux_up[i].area[k] * aux_up[i].buoy[k] / a_bulk_c
            end
        else
            aux_bulk.RH[k] = aux_en.RH[k]
            aux_bulk.T[k] = aux_en.T[k]
        end

        #####
        ##### update_GMV_diagnostics
        #####
        aux_gm.q_liq[k] = (aux_bulk.area[k] * aux_bulk.q_liq[k] + (1 - aux_bulk.area[k]) * aux_en.q_liq[k])
        aux_gm.q_ice[k] = (aux_bulk.area[k] * aux_bulk.q_ice[k] + (1 - aux_bulk.area[k]) * aux_en.q_ice[k])
        aux_gm.T[k] = (aux_bulk.area[k] * aux_bulk.T[k] + (1 - aux_bulk.area[k]) * aux_en.T[k])
        aux_gm.buoy[k] = (aux_bulk.area[k] * aux_bulk.buoy[k] + (1 - aux_bulk.area[k]) * aux_en.buoy[k])

        has_condensate = TD.has_condensate(aux_bulk.q_liq[k] + aux_bulk.q_ice[k])
        aux_bulk.cloud_fraction[k] = if has_condensate && a_bulk_c > 1e-3
            1
        else
            0
        end
    end

    #####
    ##### face variables: diagnose primitive, diagnose env and compute bulk
    #####
    # TODO: figure out why `ifelse` is allocating
    @inbounds for i in 1:N_up
        a_up_bcs = (; bottom = CCO.SetValue(edmf.area_surface_bc[i]), top = CCO.Extrapolate())
        If = CCO.InterpolateC2F(; a_up_bcs...)
        a_min = edmf.minimum_area
        a_up = aux_up[i].area
        @. aux_up_f[i].w = ifelse(If(a_up) >= a_min, max(prog_up_f[i].ρaw / (ρ0_f * If(a_up)), 0), FT(0))
    end
    @inbounds for i in 1:N_up
        aux_up_f[i].w[kf_surf] = edmf.w_surface_bc[i]
    end

    parent(aux_tc_f.bulk.w) .= 0
    a_bulk_bcs = (; bottom = CCO.SetValue(sum(edmf.area_surface_bc)), top = CCO.Extrapolate())
    Ifb = CCO.InterpolateC2F(; a_bulk_bcs...)
    @inbounds for i in 1:N_up
        a_up = aux_up[i].area
        a_up_bcs = (; bottom = CCO.SetValue(edmf.area_surface_bc[i]), top = CCO.Extrapolate())
        Ifu = CCO.InterpolateC2F(; a_up_bcs...)
        @. aux_tc_f.bulk.w += ifelse(Ifb(aux_bulk.area) > 0, Ifu(a_up) * aux_up_f[i].w / Ifb(aux_bulk.area), FT(0))
    end
    # Assuming gm.W = 0!
    @. aux_en_f.w = -Ifb(aux_bulk.area) / (1 - Ifb(aux_bulk.area)) * aux_tc_f.bulk.w

    #####
    #####  diagnose_GMV_moments
    #####
    get_GMV_CoVar(edmf, grid, state, :Hvar, :θ_liq_ice)
    get_GMV_CoVar(edmf, grid, state, :QTvar, :q_tot)
    get_GMV_CoVar(edmf, grid, state, :HQTcov, :θ_liq_ice, :q_tot)

    # TODO - use this inversion in free_convection_windspeed and not compute zi twice
    edmf.zi = get_inversion(grid, state, param_set, surface.Ri_bulk_crit)

    update_surface(Case, grid, state, gm, TS, param_set)
    update_forcing(Case, grid, state, gm, TS, param_set)
    update_radiation(Case, grid, state, gm, TS, param_set)

    compute_pressure_plume_spacing(edmf, param_set)

    #####
    ##### compute_updraft_closures
    #####

    Ic = CCO.InterpolateF2C()
    ∇c = CCO.DivergenceF2C()
    LB = CCO.LeftBiasedC2F(; bottom = CCO.SetValue(FT(0)))

    @inbounds for i in 1:N_up
        # compute ∇m at cell centers
        a_up = aux_up[i].area
        w_up = aux_up_f[i].w
        w_en = aux_en_f.w
        w_gm = prog_gm_f.w
        @. m_entr_detr = a_up * (Ic(w_up) - Ic(w_gm))
        @. ∇m_entr_detr = ∇c(wvec(LB(m_entr_detr)))
        @. w_up_c = Ic(w_up)
        @. w_en_c = Ic(w_en)
        @inbounds for k in real_center_indices(grid)
            # entrainment

            q_cond_up = TD.condensate(TD.PhasePartition(aux_up[i].q_tot[k], aux_up[i].q_liq[k], aux_up[i].q_ice[k]))
            q_cond_en = TD.condensate(TD.PhasePartition(aux_en.q_tot[k], aux_en.q_liq[k], aux_en.q_ice[k]))

            if aux_up[i].area[k] > 0.0

                εδ_model_vars = GeneralizedEntr(;
                    q_cond_up = q_cond_up,
                    q_cond_en = q_cond_en,
                    w_up = w_up_c[k],
                    w_en = w_en_c[k],
                    b_up = aux_up[i].buoy[k],
                    b_en = aux_en.buoy[k],
                    tke = aux_en.tke[k],
                    dMdz = ∇m_entr_detr[k],
                    M = m_entr_detr[k],
                    a_up = aux_up[i].area[k],
                    a_en = aux_en.area[k],
                    R_up = edmf.pressure_plume_spacing[i],
                    RH_up = aux_up[i].RH[k],
                    RH_en = aux_en.RH[k],
                    max_area = max_area,
                    zc_i = grid.zc[k].z,
                    Δt = Δt,
                )

                er = entr_detr(param_set, εδ_model_vars, edmf.entr_closure)
                aux_up[i].entr_sc[k] = er.ε_dyn
                aux_up[i].detr_sc[k] = er.δ_dyn
                # stochastic closure
                stoch_ε = stochastic_closure(param_set, sde_model, Entrainment())
                stoch_δ = stochastic_closure(param_set, sde_model, Detrainment())
                aux_up[i].entr_sc[k] *= stoch_ε
                aux_up[i].detr_sc[k] *= stoch_δ

                aux_up[i].frac_turb_entr[k] = er.ε_turb
                aux_up[i].horiz_K_eddy[k] = er.K_ε
            else
                aux_up[i].entr_sc[k] = 0.0
                aux_up[i].detr_sc[k] = 0.0
                aux_up[i].frac_turb_entr[k] = 0.0
                aux_up[i].horiz_K_eddy[k] = 0.0
            end
        end
    end

    @inbounds for i in 1:N_up

        # pressure
        b_up = aux_up[i].buoy
        a_up = aux_up[i].area
        w_up = aux_up_f[i].w
        w_en = aux_en_f.w
        updraft_top = up.updraft_top[i]
        asp_ratio = 1.0

        b_bcs = (; bottom = CCO.SetValue(b_up[kc_surf]), top = CCO.SetValue(b_up[kc_toa]))
        a_bcs = (; bottom = CCO.SetValue(edmf.area_surface_bc[i]), top = CCO.SetValue(FT(0)))
        w_bcs = (; bottom = CCO.SetValue(wvec(FT(0))), top = CCO.SetValue(wvec(FT(0))))
        bcs = (; a_up = a_bcs, b_up = b_bcs, w_up = w_bcs)

        nh_press_buoy = aux_up_f[i].nh_pressure_b
        nh_press_adv = aux_up_f[i].nh_pressure_adv
        nh_press_drag = aux_up_f[i].nh_pressure_drag

        nh_press_buoy .= nh_pressure_buoy(param_set, a_up, b_up, ρ0_f, asp_ratio, bcs)
        nh_press_adv .= nh_pressure_adv(param_set, updraft_top, a_up, ρ0_f, w_up, bcs)
        nh_press_drag .= nh_pressure_drag(param_set, updraft_top, a_up, ρ0_f, w_up, w_en, bcs)
        @. aux_up_f[i].nh_pressure = nh_press_buoy + nh_press_adv + nh_press_drag
    end

    #####
    ##### compute_eddy_diffusivities_tke
    #####

    # Subdomain exchange term
    Ic = CCO.InterpolateF2C()
    b_exch = center_aux_turbconv(state).b_exch
    parent(b_exch) .= 0
    a_en = aux_en.area
    w_en = aux_en_f.w
    tke_en = aux_en.tke
    @inbounds for i in 1:N_up
        a_up = aux_up[i].area
        w_up = aux_up_f[i].w
        δ_dyn = aux_up[i].detr_sc
        ε_turb = aux_up[i].frac_turb_entr
        @. b_exch +=
            a_up * Ic(w_up) * δ_dyn / a_en * (1 / 2 * (Ic(w_up) - Ic(w_en))^2 - tke_en) -
            a_up * Ic(w_up) * (Ic(w_up) - Ic(w_en)) * ε_turb * Ic(w_en) / a_en
    end

    Shear² = center_aux_turbconv(state).Shear²
    ∂qt∂z = center_aux_turbconv(state).∂qt∂z
    ∂θl∂z = center_aux_turbconv(state).∂θl∂z
    ∂θv∂z = center_aux_turbconv(state).∂θv∂z
    ∂qt∂z_sat = center_aux_turbconv(state).∂qt∂z_sat
    ∂θl∂z_sat = center_aux_turbconv(state).∂θl∂z_sat
    ∂θv∂z_unsat = center_aux_turbconv(state).∂θv∂z_unsat

    ∇0_bcs = (; bottom = CCO.Extrapolate(), top = CCO.Extrapolate())
    If0 = CCO.InterpolateC2F(; ∇0_bcs...)
    u_gm = prog_gm.u
    v_gm = prog_gm.v
    w_en = aux_en_f.w
    # compute shear
    @. Shear² = (∇c(wvec(If0(u_gm))))^2 + (∇c(wvec(If0(v_gm))))^2 + (∇c(wvec(w_en)))^2

    q_tot_en = aux_en.q_tot
    θ_liq_ice_en = aux_en.θ_liq_ice
    θ_virt_en = aux_en.θ_virt
    @. ∂qt∂z = ∇c(wvec(If0(q_tot_en)))
    @. ∂θl∂z = ∇c(wvec(If0(θ_liq_ice_en)))
    @. ∂θv∂z = ∇c(wvec(If0(θ_virt_en)))

    # Second order approximation: Use dry and cloudy environmental fields.
    cf = aux_en.cloud_fraction
    shm = copy(cf)
    parent(shm) .= shrink_mask(vec(cf))

    # Since NaN*0 ≠ 0, we need to conditionally replace
    # our gradients by their default values.
    @. ∂qt∂z_sat = ∇c(wvec(If0(aux_en_sat.q_tot)))
    @. ∂θl∂z_sat = ∇c(wvec(If0(aux_en_sat.θ_liq_ice)))
    @. ∂θv∂z_unsat = ∇c(wvec(If0(aux_en_unsat.θ_virt)))
    for k in real_center_indices(grid)
        if shm[k] == 0
            ∂qt∂z_sat[k] = ∂qt∂z[k]
            ∂θl∂z_sat[k] = ∂θl∂z[k]
            ∂θv∂z_unsat[k] = ∂θv∂z[k]
        end
    end

    @inbounds for k in real_center_indices(grid)

        # buoyancy_gradients
        if edmf.bg_closure == BuoyGradMean()
            # First order approximation: Use environmental mean fields.
            ts_en = TD.PhaseEquil_pTq(param_set, p0_c[k], aux_en.T[k], aux_en.q_tot[k])
            bg_kwargs = (;
                t_sat = aux_en.T[k],
                qv_sat = TD.vapor_specific_humidity(ts_en),
                qt_sat = aux_en.q_tot[k],
                θ_sat = aux_en.θ_dry[k],
                θ_liq_ice_sat = aux_en.θ_liq_ice[k],
                ∂θv∂z_unsat = ∂θv∂z[k],
                ∂qt∂z_sat = ∂qt∂z[k],
                ∂θl∂z_sat = ∂θl∂z[k],
                p0 = p0_c[k],
                en_cld_frac = aux_en.cloud_fraction[k],
                alpha0 = α0_c[k],
            )
            bg_model = EnvBuoyGrad(edmf.bg_closure; bg_kwargs...)

        elseif edmf.bg_closure == BuoyGradQuadratures()

            bg_kwargs = (;
                t_sat = aux_en_sat.T[k],
                qv_sat = aux_en_sat.q_vap[k],
                qt_sat = aux_en_sat.q_tot[k],
                θ_sat = aux_en_sat.θ_dry[k],
                θ_liq_ice_sat = aux_en_sat.θ_liq_ice[k],
                ∂θv∂z_unsat = ∂θv∂z_unsat[k],
                ∂qt∂z_sat = ∂qt∂z_sat[k],
                ∂θl∂z_sat = ∂θl∂z_sat[k],
                p0 = p0_c[k],
                en_cld_frac = aux_en.cloud_fraction[k],
                alpha0 = α0_c[k],
            )
            bg_model = EnvBuoyGrad(edmf.bg_closure; bg_kwargs...)
        end
        bg = buoyancy_gradients(param_set, bg_model)

        # Limiting stratification scale (Deardorff, 1976)
        # compute ∇Ri and Pr
        ∇_Ri = gradient_Richardson_number(bg.∂b∂z, Shear²[k], eps(0.0))
        aux_tc.prandtl_nvec[k] = turbulent_Prandtl_number(param_set, obukhov_length, ∇_Ri)

        ml_model = MinDisspLen(;
            z = grid.zc[k].z,
            obukhov_length = obukhov_length,
            tke_surf = aux_en.tke[kc_surf],
            ustar = surface.ustar,
            Pr = aux_tc.prandtl_nvec[k],
            p0 = p0_c[k],
            ∇b = bg,
            Shear² = Shear²[k],
            tke = aux_en.tke[k],
            b_exch = b_exch[k],
        )

        ml = mixing_length(param_set, ml_model)
        aux_tc.mls[k] = ml.min_len_ind
        aux_tc.mixing_length[k] = ml.mixing_length
        aux_tc.ml_ratio[k] = ml.ml_ratio

        KM[k] = c_m * ml.mixing_length * sqrt(max(aux_en.tke[k], 0.0))
        KH[k] = KM[k] / aux_tc.prandtl_nvec[k]

        aux_en_2m.tke.buoy[k] = -aux_en.area[k] * ρ0_c[k] * KH[k] * bg.∂b∂z
    end

    #####
    ##### compute covarinaces tendendies
    #####
    tke_press = aux_en_2m.tke.press
    w_en = aux_en_f.w
    parent(tke_press) .= 0
    @inbounds for i in 1:N_up
        w_up = aux_up_f[i].w
        nh_press = aux_up_f[i].nh_pressure
        @. tke_press += (Ic(w_en) - Ic(w_up)) * Ic(nh_press)
    end

    compute_covariance_entr(edmf, grid, state, Val(:tke), Val(:w))
    compute_covariance_entr(edmf, grid, state, Val(:Hvar), Val(:θ_liq_ice))
    compute_covariance_entr(edmf, grid, state, Val(:QTvar), Val(:q_tot))
    compute_covariance_entr(edmf, grid, state, Val(:HQTcov), Val(:θ_liq_ice), Val(:q_tot))
    compute_covariance_shear(edmf, grid, state, gm, Val(:tke), Val(:w))
    compute_covariance_shear(edmf, grid, state, gm, Val(:Hvar), Val(:θ_liq_ice))
    compute_covariance_shear(edmf, grid, state, gm, Val(:QTvar), Val(:q_tot))
    compute_covariance_shear(edmf, grid, state, gm, Val(:HQTcov), Val(:θ_liq_ice), Val(:q_tot))
    compute_covariance_dissipation(edmf, grid, state, :tke, param_set)
    compute_covariance_dissipation(edmf, grid, state, :Hvar, param_set)
    compute_covariance_dissipation(edmf, grid, state, :QTvar, param_set)
    compute_covariance_dissipation(edmf, grid, state, :HQTcov, param_set)

    # TODO defined again in compute_covariance_shear and compute_covaraince
    @inbounds for k in real_center_indices(grid)
        aux_en_2m.tke.rain_src[k] = 0
        aux_en_2m.Hvar.rain_src[k] = ρ0_c[k] * aux_en.area[k] * 2 * aux_en.Hvar_rain_dt[k]
        aux_en_2m.QTvar.rain_src[k] = ρ0_c[k] * aux_en.area[k] * 2 * aux_en.QTvar_rain_dt[k]
        aux_en_2m.HQTcov.rain_src[k] = ρ0_c[k] * aux_en.area[k] * aux_en.HQTcov_rain_dt[k]
    end

    get_GMV_CoVar(edmf, grid, state, :tke, :w)

    compute_diffusive_fluxes(edmf, grid, state, gm, Case, TS, param_set)

    # TODO: use dispatch
    if edmf.Precip.precipitation_model == "clima_1m"
        # helper to calculate the rain velocity
        # TODO: assuming gm.W = 0
        # TODO: verify translation
        term_vel_rain = aux_tc.term_vel_rain
        term_vel_snow = aux_tc.term_vel_snow
        prog_pr = center_prog_precipitation(state)

        @inbounds for k in real_center_indices(grid)
            term_vel_rain[k] = CM1.terminal_velocity(param_set, CM1.RainType(), ρ0_c[k], prog_pr.q_rai[k])
            term_vel_snow[k] = CM1.terminal_velocity(param_set, CM1.SnowType(), ρ0_c[k], prog_pr.q_sno[k])
        end
    end

    @inbounds for i in 1:N_up
        up.updraft_top[i] = 0.0

        @inbounds for k in real_center_indices(grid)
            if aux_up[i].area[k] > 1e-3
                up.updraft_top[i] = max(up.updraft_top[i], grid.zc[k])
            end
        end
    end

end


using MathOptInterface
const MOI = MathOptInterface

include("scaled_psd_cone_bridge.jl")

MOI.Utilities.@product_of_sets(
    Cones,
    MOI.Zeros,
    MOI.Nonnegatives,
    MOI.SecondOrderCone,
    ScaledPSDCone,
    MOI.ExponentialCone,
    MOI.DualExponentialCone,
    MOI.PowerCone{T},
    MOI.DualPowerCone{T},
)

const OptimizerCache{T} = MOI.Utilities.GenericModel{
    Cdouble,
    MOI.Utilities.ObjectiveContainer{Cdouble},
    MOI.Utilities.VariablesContainer{Cdouble},
    MOI.Utilities.MatrixOfConstraints{
        Cdouble,
        MOI.Utilities.MutableSparseMatrixCSC{
            Cdouble,
            T,
            MOI.Utilities.ZeroBasedIndexing,
        },
        Vector{Cdouble},
        Cones{Cdouble},
    },
}

function _to_sparse(
    ::Type{T},
    A::MOI.Utilities.MutableSparseMatrixCSC{
        Cdouble,
        T,
        MOI.Utilities.ZeroBasedIndexing,
    },
) where {T}
    return -A.nzval, A.rowval, A.colptr
end

mutable struct MOISolution
    primal::Vector{Float64}
    dual::Vector{Float64}
    slack::Vector{Float64}
    ret_val::Int
    raw_status::String
    objective_value::Float64
    dual_objective_value::Float64
    objective_constant::Float64
    solve_time_sec::Float64
    iterations::Int
end

function MOISolution()
    return MOISolution(
        Float64[],
        Float64[],
        Float64[],
        0, # SCS_UNFINISHED
        "",
        NaN,
        NaN,
        NaN,
        0.0,
        0,
    )
end

mutable struct Optimizer <: MOI.AbstractOptimizer
    cones::Union{Nothing,Cones{Cdouble}}
    sol::MOISolution
    silent::Bool
    options::Dict{Symbol,Any}
    function Optimizer(; kwargs...)
        optimizer = new(nothing, MOISolution(), false, Dict{Symbol,Any}())
        for (key, value) in kwargs
            MOI.set(optimizer, MOI.RawOptimizerAttribute(String(key)), value)
        end
        return optimizer
    end
end

function MOI.get(::Optimizer, ::MOI.Bridges.ListOfNonstandardBridges)
    return [ScaledPSDConeBridge{Cdouble}]
end

MOI.get(::Optimizer, ::MOI.SolverName) = "SCS"

MOI.is_empty(optimizer::Optimizer) = optimizer.cones === nothing

function MOI.empty!(optimizer::Optimizer)
    optimizer.cones = nothing
    optimizer.sol = MOISolution()
    return
end

###
### MOI.RawOptimizerAttribute
###

function MOI.set(optimizer::Optimizer, param::MOI.RawOptimizerAttribute, value)
    return optimizer.options[Symbol(param.name)] = value
end

function MOI.get(optimizer::Optimizer, param::MOI.RawOptimizerAttribute)
    return optimizer.options[Symbol(param.name)]
end

###
### MOI.Silent
###

MOI.supports(::Optimizer, ::MOI.Silent) = true

function MOI.set(optimizer::Optimizer, ::MOI.Silent, value::Bool)
    optimizer.silent = value
    return
end

MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent

###
### MOI.AbstractModelAttribute
###

function MOI.supports(
    ::Optimizer,
    ::Union{
        MOI.ObjectiveSense,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}},
    },
)
    return true
end

###
### MOI.AbstractVariableAttribute
###

# TODO(odow): support variable primal starts
# function MOI.supports(
#     ::Optimizer,
#     ::MOI.VariablePrimalStart,
#     ::Type{MOI.VariableIndex},
# )
#     return true
# end

###
### MOI.AbstractConstraintAttribute
###

# TODO(odow): support constraint primal and dual starts
# function MOI.supports(
#     ::Optimizer,
#     ::Union{MOI.ConstraintPrimalStart,MOI.ConstraintDualStart},
#     ::Type{<:MOI.ConstraintIndex},
# )
#     return true
# end

function MOI.supports_constraint(
    ::Optimizer,
    ::Type{MOI.VectorAffineFunction{Cdouble}},
    ::Type{
        <:Union{
            MOI.Zeros,
            MOI.Nonnegatives,
            MOI.SecondOrderCone,
            ScaledPSDCone,
            MOI.ExponentialCone,
            MOI.DualExponentialCone,
            MOI.PowerCone{Cdouble},
            MOI.DualPowerCone{Cdouble},
        },
    },
)
    return true
end

function _map_sets(f, ::Type{T}, sets, ::Type{S}) where {T,S}
    F = MOI.VectorAffineFunction{Cdouble}
    cis = MOI.get(sets, MOI.ListOfConstraintIndices{F,S}())
    return T[f(MOI.get(sets, MOI.ConstraintSet(), ci)) for ci in cis]
end

function MOI.optimize!(
    dest::Optimizer,
    src::MOI.Utilities.UniversalFallback{OptimizerCache{T}},
) where {T}
    linear_solver = get(dest.options, :linear_solver, DirectSolver)
    if T != scsint_t(linear_solver)
        cache = MOI.Utilities.UniversalFallback(
            OptimizerCache{scsint_t(linear_solver)}(),
        )
        MOI.copy_to(cache, src)
        return MOI.optimize!(dest, cache)
    end
    # The real stuff starts here.
    MOI.empty!(dest)
    index_map = MOI.Utilities.identity_index_map(src)
    Ab = src.model.constraints
    A = Ab.coefficients

    model_attributes = MOI.get(src, MOI.ListOfModelAttributesSet())
    objective_function_attr =
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Cdouble}}()
    for attr in model_attributes
        if attr != MOI.ObjectiveSense() && attr != objective_function_attr
            throw(MOI.UnsupportedAttribute(attr))
        end
    end
    max_sense = false
    if MOI.ObjectiveSense() in model_attributes
        max_sense = MOI.get(src, MOI.ObjectiveSense()) == MOI.MAX_SENSE
    end
    objective_constant = 0.0
    c = zeros(A.n)
    if objective_function_attr in model_attributes
        obj = MOI.get(src, objective_function_attr)
        objective_constant = MOI.constant(obj)
        for term in obj.terms
            c[term.variable.value] += (max_sense ? -1 : 1) * term.coefficient
        end
    end
    # `model.primal` contains the result of the previous optimization.
    # It is used as a warm start if its length is the same, e.g.
    # probably because no variable and/or constraint has been added.
    if A.n != length(dest.sol.primal)
        resize!(dest.sol.primal, A.n)
        fill!(dest.sol.primal, 0.0)
    end
    if A.m != length(dest.sol.dual)
        resize!(dest.sol.dual, A.m)
        fill!(dest.sol.dual, 0.0)
    end
    if A.m != length(dest.sol.slack)
        resize!(dest.sol.slack, A.m)
        fill!(dest.sol.slack, 0.0)
    end
    # Set starting values and throw error for other variable attributes
    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    MOI.Utilities.pass_attributes(dest, src, index_map, vis_src)
    # Set starting values and throw error for other constraint attributes
    for (F, S) in MOI.get(src, MOI.ListOfConstraintTypesPresent())
        cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
        MOI.Utilities.pass_attributes(dest, src, index_map, cis_src)
    end
    options = copy(dest.options)
    if dest.silent
        options[:verbose] = 0
    end
    dest.cones = deepcopy(Ab.sets)
    sol = scs_solve(
        linear_solver,
        A.m,
        A.n,
        A,
        Ab.constants,
        c,
        Ab.sets.num_rows[1],
        Ab.sets.num_rows[2] - Ab.sets.num_rows[1],
        _map_sets(MOI.dimension, T, Ab, MOI.SecondOrderCone),
        _map_sets(MOI.side_dimension, T, Ab, ScaledPSDCone),
        div(Ab.sets.num_rows[5] - Ab.sets.num_rows[4], 3),
        div(Ab.sets.num_rows[6] - Ab.sets.num_rows[5], 3),
        vcat(
            _map_sets(set -> set.exponent, Float64, Ab, MOI.PowerCone{Cdouble}),
            _map_sets(
                set -> -set.exponent,
                Float64,
                Ab,
                MOI.DualPowerCone{Cdouble},
            ),
        ),
        dest.sol.primal,
        dest.sol.dual,
        dest.sol.slack;
        options...,
    )
    dest.sol = MOISolution(
        sol.x,
        sol.y,
        sol.s,
        sol.ret_val,
        raw_status(sol.info),
        (max_sense ? -1 : 1) * sol.info.pobj,
        (max_sense ? -1 : 1) * sol.info.dobj,
        objective_constant,
        (sol.info.setupTime + sol.info.solveTime) / 1000,
        sol.info.iter,
    )
    return index_map, false
end

function MOI.optimize!(dest::Optimizer, src::MOI.ModelLike)
    linear_solver = get(dest.options, :linear_solver, DirectSolver)
    T = scsint_t(linear_solver)
    cache = MOI.Utilities.UniversalFallback(OptimizerCache{T}())
    index_map = MOI.copy_to(cache, src)
    MOI.optimize!(dest, cache)
    return index_map, false
end

function MOI.get(optimizer::Optimizer, ::MOI.SolveTimeSec)
    return optimizer.sol.solve_time_sec
end

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return optimizer.sol.raw_status
end

"""
    ADMMIterations()

The number of ADMM iterations completed during the solve.
"""
struct ADMMIterations <: MOI.AbstractModelAttribute end

MOI.is_set_by_optimize(::ADMMIterations) = true

function MOI.get(optimizer::Optimizer, ::ADMMIterations)
    return optimizer.sol.iterations
end

# Implements getter for result value and statuses
# SCS returns one of the following integers:
# -7 SCS_INFEASIBLE_INACCURATE
# -6 SCS_UNBOUNDED_INACCURATE
# -5 SCS_SIGINT
# -4 SCS_FAILED
# -3 SCS_INDETERMINATE
# -2 SCS_INFEASIBLE  : primal infeasible, dual unbounded
# -1 SCS_UNBOUNDED   : primal unbounded, dual infeasible
#  0 SCS_UNFINISHED  : never returned, used as placeholder
#  1 SCS_SOLVED
#  2 SCS_SOLVED_INACCURATE
function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    s = optimizer.sol.ret_val
    @assert -7 <= s <= 2
    if s == -7
        return MOI.ALMOST_INFEASIBLE
    elseif s == -6
        return MOI.ALMOST_DUAL_INFEASIBLE
    elseif s == 2
        return MOI.ALMOST_OPTIMAL
    elseif s == -5
        return MOI.INTERRUPTED
    elseif s == -4
        return MOI.NUMERICAL_ERROR
    elseif s == -3
        return MOI.SLOW_PROGRESS
    elseif s == -2
        return MOI.INFEASIBLE
    elseif s == -1
        return MOI.DUAL_INFEASIBLE
    elseif s == 1
        return MOI.OPTIMAL
    else
        @assert s == 0
        return MOI.OPTIMIZE_NOT_CALLED
    end
end

function MOI.get(optimizer::Optimizer, attr::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(optimizer, attr)
    value = optimizer.sol.objective_value
    if !MOI.Utilities.is_ray(MOI.get(optimizer, MOI.PrimalStatus()))
        value += optimizer.sol.objective_constant
    end
    return value
end

function MOI.get(optimizer::Optimizer, attr::MOI.DualObjectiveValue)
    MOI.check_result_index_bounds(optimizer, attr)
    value = optimizer.sol.dual_objective_value
    if !MOI.Utilities.is_ray(MOI.get(optimizer, MOI.DualStatus()))
        value += optimizer.sol.objective_constant
    end
    return value
end

function MOI.get(optimizer::Optimizer, attr::MOI.PrimalStatus)
    if attr.result_index > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    elseif optimizer.sol.ret_val in (-3, 1, 2)
        return MOI.FEASIBLE_POINT
    elseif optimizer.sol.ret_val in (-6, -1)
        return MOI.INFEASIBILITY_CERTIFICATE
    end
    return MOI.INFEASIBLE_POINT
end

function MOI.get(
    optimizer::Optimizer,
    attr::MOI.VariablePrimal,
    vi::MOI.VariableIndex,
)
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.sol.primal[vi.value]
end

function MOI.get(
    optimizer::Optimizer,
    attr::MOI.ConstraintPrimal,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.sol.slack[MOI.Utilities.rows(optimizer.cones, ci)]
end

function MOI.get(optimizer::Optimizer, attr::MOI.DualStatus)
    if attr.result_index > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    elseif optimizer.sol.ret_val in (-3, 1, 2)
        return MOI.FEASIBLE_POINT
    elseif optimizer.sol.ret_val in (-7, -2)
        return MOI.INFEASIBILITY_CERTIFICATE
    end
    return MOI.INFEASIBLE_POINT
end

function MOI.get(
    optimizer::Optimizer,
    attr::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.sol.dual[MOI.Utilities.rows(optimizer.cones, ci)]
end

MOI.get(optimizer::Optimizer, ::MOI.ResultCount) = 1


#
# Interface for PIPS-NLP parallel (structured interface)
#

module PipsNlpInterface
Base.include(PipsNlpInterface, "pips_parallel_cfunc.jl")

# using PipsNlpSolver
using StructJuMP, JuMP
using StructJuMPSolverInterface
using Printf
using SparseArrays

import MathProgBase

mutable struct MatStorage
    rowIdx::Vector{Int}
    colIdx::Vector{Int}
    value::Vector{Float64}
    m::Int
    n::Int
    isVal::Bool
    function MatStorage(I,J,m,n)
        return new(I,J,Vector{Float64}(undef, length(I)),m,n,false)
    end
end

# type MatStrStorage
#     rowIdx::Vector{Int}
#     colIdx::Vector{Int}
#     function MatStrStorage()
#         return new(Vector{Int}(),Vector{Int}())
#     end
#     function MatStrStorage(I,J)
#         return new(I,J)
#     end
# end

mutable struct StructJuMPModel <: ModelInterface
    internalModel::JuMP.Model
    status::Int
    n_iter::Int
    prof::Bool
    evaluatorMap::Dict{Int,JuMP.NLPEvaluator}
    iMap::Dict{Int,Pair{Dict{Int,Int},Dict{Int,Int}}}  #eq, ieq, jump->actual used
    jMap::Dict{Pair{Int,Int},Dict{Int,Int}}
    hcMap::Dict{Pair{Int,Int},Dict{Int,Int}}
    hrMap::Dict{Pair{Int,Int},Dict{Int,Int}}
    matJac::Dict{Int,MatStorage}
    matHess::Dict{Int,MatStorage}
    matEqJac::Dict{Pair{Int,Int},MatStorage}
    matIeqJac::Dict{Pair{Int,Int},MatStorage}
    matHessMap::Dict{Pair{Int,Int},MatStorage}
    t_jump::Float64
    t_jac_spconv::Float64
    t_hess_spconv::Float64
    t_hess00_spconv::Float64
    t_eval_g::Float64
    t_itermap_eval_g::Float64
    n_jac_spconv::Int
    n_hess_spconv::Int
    n_hess00_spconv::Int
    n_eval_g::Int

    get_num_scen::Function
    get_sense::Function
    get_status::Function

    get_num_rows::Function
    get_num_cols::Function
    get_num_eq_cons::Function
    get_num_ineq_cons::Function
    get_num_eq_lcons::Function
    get_num_ineq_lcons::Function

    set_status::Function

    set_num_rows::Function
    set_num_cols::Function
    set_num_eq_cons::Function
    set_num_ineq_cons::Function

    str_init_x0::Function
    str_prob_info::Function
    str_eval_f::Function
    str_eval_g::Function
    str_eval_grad_f::Function
    str_eval_jac_g::Function
    str_eval_h::Function
    str_write_solution::Function

    function StructJuMPModel(model::JuMP.Model,prof=false)
        instance = new(model,0,-1,prof,
            Dict{Int,JuMP.NLPEvaluator}(),
            Dict{Int,Pair{Dict{Int,Int},Dict{Int,Int}}}(),
            Dict{Pair{Int,Int},Dict{Int,Int}}(),
            Dict{Pair{Int,Int},Dict{Int,Int}}(),
            Dict{Pair{Int,Int},Dict{Int,Int}}(),
            Dict{Int,MatStorage}(), #jac mat
            Dict{Int,MatStorage}(), #hess mat
            Dict{Pair{Int,Int},MatStorage}(),
            Dict{Pair{Int,Int},MatStorage}(),
            Dict{Pair{Int,Int},MatStorage}(),
            0.0,0.0,0.0,0.0,0.0,0.0,
            0,0,0,0
            )

        initialization(instance)

        instance.get_num_scen = function()
            return num_scenarios(instance.internalModel)
        end
        instance.get_sense = function()
            return getobjectivesense(instance.internalModel)
        end
        instance.get_status = function()
            return instance.status
        end
        instance.get_num_rows = function(id::Int)
            return getNumCons(instance.internalModel,id)
        end
        instance.get_num_cols = function(id::Int)
            return getNumVars(instance.internalModel,id)
        end
        instance.get_num_eq_cons = function(id::Int)
            return length(instance.iMap[id][1])
        end
        instance.get_num_ineq_cons = function(id::Int)
            return length(instance.iMap[id][2])
        end
        instance.get_num_eq_lcons = function()
            return 0
        end
        instance.get_num_ineq_lcons = function()
            return 0
        end
        instance.set_status = function(s::Int)
            instance.status = s
        end
        instance.set_num_rows = function(id::Integer,v::Integer) end
        instance.set_num_cols = function(id::Integer,v::Integer) end
        instance.set_num_eq_cons = function(id::Integer,v::Integer) end
        instance.set_num_ineq_cons = function(id::Integer,v::Integer) end

        instance.str_init_x0 = function(id,x0)
            @assert(id in getLocalBlocksIds(instance.internalModel))
            mm = getModel(instance.internalModel,id)
            nvar = getNumVars(instance.internalModel,id)
            @assert length(x0) == nvar

            for i=1:nvar
                x0[i] = getvalue(Variable(mm,i))
                isnan(x0[i]) ? x0[i]=1.0 : nothing
            end
            # @show x0;
        end

        instance.str_prob_info = function(id,flag,mode,clb,cub,rlb,rub)
            # @show id
            if mode == :Structure
                if flag == 0
                    nn = getNumVars(instance.internalModel,id)
                    mm = getNumCons(instance.internalModel,id)
                else
                    @assert flag == 1
                    nn = getNumVars(instance.internalModel,id)
                    # mm = getNumLinkCons() #TODO: this is not yet supported
                    mm = 0
                end
                # @show nn,mm
                return (nn,mm)
            elseif mode == :Values
                if flag == 0
                    # @show length(clb),length(cub)
                    mm = getModel(instance.internalModel,id)
                    nvar = getNumVars(instance.internalModel,id)
                    @assert length(clb) == nvar
                    @assert length(cub) == nvar
                    array_copy(mm.colUpper, 1, cub, 1, nvar)
                    array_copy(mm.colLower, 1, clb, 1, nvar)
                    lb,ub = JuMP.constraintbounds(mm)
                    # @show lb, ub
                    (eq_idx, ieq_idx) = instance.iMap[id]
                    # @show eq_idx,ieq_idx
                    @assert length(lb) == length(ub)
                    @assert length(eq_idx) + length(ieq_idx) == length(lb)
                    num_eq = length(eq_idx)
                    for i in eq_idx
                        rlb[i[2]] = lb[i[1]]
                        rub[i[2]] = ub[i[1]]
                    end
                    for i in ieq_idx
                        rlb[i[2]+num_eq] = lb[i[1]]
                        rub[i[2]+num_eq] = ub[i[1]]
                    end
                    # @show id, rlb, rub
                else
                    @assert flag == 1
                    @assert length(rlb) == length(rub) == 0 #TODO: not yet support linking constraints syntax
                end
            else
                @assert false mode
            end
        end

        instance.str_eval_f = function(id,x0,x1)
            # @printf("#*************** eval_f - %d \n", id)
            # @show x0, x1
            # x0, x1 = load_x("pips", instance.n_iter)

            e =  instance.evaluatorMap[id]

            @timing instance.prof tic()
            obj = MathProgBase.eval_f(e,build_x(instance.internalModel,id,x0,x1))
            @timing instance.prof instance.t_jump += toq()
            # @show obj
            return obj

        end

        instance.str_eval_g = function(id::Int,x0::Vector{Float64},x1::Vector{Float64}, new_eq_g::Vector{Float64}, new_inq_g::Vector{Float64})
            # x0, x1 = load_x("pips", instance.n_iter)
            @timing instance.prof tic()

            e = instance.evaluatorMap[id]
            g = Vector{Float64}(undef, getNumCons(instance.internalModel,id))

            @timing instance.prof tic()
            MathProgBase.eval_g(e,g,build_x(instance.internalModel,id,x0,x1))
            @timing instance.prof instance.t_jump += toq()

            p = instance.iMap[id]
            eq_idx = p[1]
            ieq_idx = p[2]
            @assert length(new_eq_g) == length(eq_idx)
            @assert length(ieq_idx) == length(new_inq_g)

            @timing instance.prof tic()
            for i in eq_idx
                new_eq_g[i[2]] = g[i[1]]
                # @assert !haskey(ieq_idx,i[1])
            end
            for i in ieq_idx
                # @assert !haskey(eq_idx,i[1])
                new_inq_g[i[2]] = g[i[1]]
            end
            @timing instance.prof instance.t_itermap_eval_g+=toq()

            # @printf("#********  str_eval_g - %d \n", id)
            # @show x0, x1
            # @show new_eq_g, new_inq_g
            @timing instance.prof begin
                instance.t_eval_g += toq()
                instance.n_eval_g += 1
            end
        end

        instance.str_eval_grad_f = function(rowid, colid, x0, x1, new_grad_f)
            # x0, x1 = load_x("pips", instance.n_iter)
            m = instance.internalModel
            @assert rowid >= colid
            @assert sum(new_grad_f) == 0.0
            e = instance.evaluatorMap[rowid]
            x = build_x(m,rowid,x0,x1)
            g = Vector{Float64}(undef, length(x))

            @timing instance.prof tic()

            MathProgBase.eval_grad_f(e,g,x)
            @timing instance.prof instance.t_jump += toq()

            @assert length(g) == MathProgBase.numvar(getModel(m,rowid))
            @assert length(new_grad_f) == getNumVars(m,colid)

            var_idx_map = get_jac_col_idx_map(m,rowid,colid,instance.jMap)
            for i in var_idx_map
                new_grad_f[i[2]] = g[i[1]]
            end

            # @printf("#**********  str_eval_grad_f, %d %d \n", rowid, colid)
            # @show x0, x1, new_grad_f
            # @printf("grad%d%d=new_grad_f; \n",rowid,colid)
        end

        instance.str_eval_jac_g = function(rowid, colid, flag, x0 , x1, mode, e_rowidx, e_colptr, e_values, i_rowidx, i_colptr, i_values)
            # @printf("#**********  str_eval_jac_g  - %s - %d - %d %d \n", mode, flag, rowid, colid);
            # @show x0, x1
            # x0, x1 = load_x("pips",instance.n_iter)

            m = instance.internalModel
            # @show instance
            # @show instance.jMap
            jMap = instance.jMap
            matJac = instance.matJac
            matEqJac = instance.matEqJac
            matIeqJac = instance.matIeqJac
            @assert rowid<=num_scenarios(m) && colid <= num_scenarios(m)
            @assert rowid >= colid
            if(mode == :Structure)
                if !haskey(matJac, rowid)
                    e = instance.evaluatorMap[rowid]

                    @timing instance.prof tic()
                    jac_I,jac_J = MathProgBase.jac_structure(e)
                    @timing instance.prof instance.t_jump += toq()

                    mm = getModel(m,rowid)
                    matJac[rowid] = MatStorage(jac_I,jac_J, MathProgBase.numconstr(mm),MathProgBase.numvar(mm))
                    @assert length(jac_I) == length(matJac[rowid].value)
                end
                mat = matJac[rowid]
                jac_I = mat.rowIdx
                jac_J = mat.colIdx

                p = Pair{Int,Int}(rowid,colid)
                if !haskey(matEqJac,p)
                    @assert !haskey(matIeqJac,p)
                    (eq_idx, ieq_idx) = instance.iMap[rowid]
                    var_idx = get_jac_col_idx_map(m,rowid,colid,jMap)
                    # @show var_idx

                    eq_jac_I = Vector{Int}()
                    ieq_jac_I = Vector{Int}()
                    eq_jac_J = Vector{Int}()
                    ieq_jac_J = Vector{Int}()
                    for i = 1:length(jac_I)
                        ii = jac_I[i]
                        jj = jac_J[i]
                        if haskey(var_idx,jj)
                            if haskey(eq_idx,ii)
                                push!(eq_jac_I, eq_idx[ii])
                                push!(eq_jac_J, var_idx[jj])
                            else
                                @assert haskey(ieq_idx,ii)
                                push!(ieq_jac_I, ieq_idx[ii])
                                push!(ieq_jac_J, var_idx[jj])
                            end
                        end
                    end
                    matEq = MatStorage(eq_jac_I,eq_jac_J,length(eq_idx), getNumVars(m,colid))
                    matIeq = MatStorage(ieq_jac_I,ieq_jac_J,length(ieq_idx), getNumVars(m,colid))
                    matEqJac[p] = matEq
                    matIeqJac[p] = matIeq
                end
                matEq = matEqJac[p]
                matIeq = matIeqJac[p]
                fill!(matEq.value,1.0)
                fill!(matIeq.value,1.0)
                @assert matEq.m + matIeq.m == mat.m

                # @show eq_jac_I
                # @show eq_jac_J

                @timing instance.prof tic()
                eq_jac = sparse(matEq.rowIdx,matEq.colIdx,matEq.value,matEq.m, matEq.n)
                @timing instance.prof begin
                    instance.t_jac_spconv += toq()
                    instance.n_jac_spconv += 1
                end
                # @show ieq_jac_I
                # @show ieq_jac_J
                @timing instance.prof tic()
                ieq_jac = sparse(matIeq.rowIdx,matIeq.colIdx,matIeq.value,matIeq.m, matIeq.n)
                @timing instance.prof begin
                    instance.t_jac_spconv += toq()
                    instance.n_jac_spconv += 1
                end

                nnzEqJac = length(eq_jac.nzval)
                nnzIeqJac = length(ieq_jac.nzval)

                return nnzEqJac,nnzIeqJac
            elseif(mode == :Values)
                p = Pair{Int,Int}(rowid,colid)
                @assert haskey(matEqJac,p)
                @assert haskey(matIeqJac,p)
                @assert haskey(matJac,rowid)
                mat = matJac[rowid]
                jac_I = mat.rowIdx
                jac_J = mat.colIdx
                jac_g = mat.value
                if(!mat.isVal)
                    e = instance.evaluatorMap[rowid]

                    @timing instance.prof tic()
                    MathProgBase.eval_jac_g(e,jac_g,build_x(m,rowid,x0,x1))
                    @timing instance.prof instance.t_jump += toq()

                    mat.isVal = true
                end

                (eq_idx, ieq_idx) = instance.iMap[rowid]
                var_idx = get_jac_col_idx_map(m,rowid,colid,jMap)

                matEq =  matEqJac[p]
                matIeq = matIeqJac[p]
                eq_jac_g = matEq.value
                eqi = 1
                ieq_jac_g = matIeq.value
                ieqi = 1

                for i = 1:length(jac_I)
                    ii = jac_I[i]
                    jj = jac_J[i]
                    vv = jac_g[i]
                    # @show ii, jj, vv
                    if haskey(var_idx,jj)
                        if haskey(eq_idx,ii)
                            eq_jac_g[eqi] = vv
                            eqi += 1
                        else
                            @assert haskey(ieq_idx,ii)
                            ieq_jac_g[ieqi] = vv
                            ieqi += 1
                        end
                    end
                end

                if(length(eq_jac_g) != 0)

                    @timing instance.prof tic()
                    eq_jac = sparse(matEq.rowIdx, matEq.colIdx, matEq.value, matEq.m, matEq.n, keepzeros=true)
                    @timing instance.prof begin
                        instance.t_jac_spconv += toq()
                        instance.n_jac_spconv += 1
                    end
                    # @printf("em=%d; en=%d;\n", length(eq_idx), getNumVars(m,colid))
                    # @show eq_jac_I, eq_jac_J, eq_jac_g
                    # @printf("ejac%d%d=sparse(eq_jac_I,eq_jac_J,eq_jac_g,em,en); \n",rowid,colid)

                    array_copy(eq_jac.rowval,1,e_rowidx,1,length(eq_jac.rowval))
                    array_copy(eq_jac.colptr,1,e_colptr,1,length(eq_jac.colptr))
                    array_copy(eq_jac.nzval, 1,e_values,1,length(eq_jac.nzval))
                    # convert_to_c_idx(e_rowidx)
                    # convert_to_c_idx(e_colptr)

                    filename = string("jaceq_",rowid,"_",colid)
                    write_mat_to_file(filename,eq_jac)
                    # @show eq_jac
                    convert_to_c_idx(e_rowidx)
                    convert_to_c_idx(e_colptr)
                end

                if(length(ieq_jac_g) != 0)
                    @timing instance.prof tic()
                    ieq_jac = sparse(matIeq.rowIdx, matIeq.colIdx, matIeq.value, matIeq.m, matIeq.n, keepzeros=true)
                    @timing instance.prof begin
                        instance.t_jac_spconv += toq()
                        instance.n_jac_spconv += 1
                    end
                    # @printf("im=%d; in=%d;\n", length(ieq_idx), getNumVars(m,colid))
                    # @show ieq_jac_I, ieq_jac_J, ieq_jac_g
                    # @printf("ijac%d%d=sparse(ieq_jac_I,ieq_jac_J,ieq_jac_g,im,in); \n",rowid,colid)
                    # @printf("jac%d%d=vcat(ejac%d%d,ijac%d%d) \n",rowid,colid,rowid,colid,rowid,colid)

                    array_copy(ieq_jac.rowval,1,i_rowidx,1,length(ieq_jac.rowval))
                    array_copy(ieq_jac.colptr,1,i_colptr,1,length(ieq_jac.colptr))
                    array_copy(ieq_jac.nzval, 1,i_values,1,length(ieq_jac.nzval))
                    # convert_to_c_idx(i_rowidx)
                    # convert_to_c_idx(i_colptr)

                    filename = string("jacieq_",rowid,"_",colid)
                    write_mat_to_file(filename,ieq_jac)
                    # @show ieq_jac
                    convert_to_c_idx(i_rowidx)
                    convert_to_c_idx(i_colptr)
                end

                if(flag == 2)
                    mat.isVal = false
                end
            else
                @assert false mode
            end

            # @show "end str_eval_jac_g"
        end

        instance.str_eval_h = function(rowid, colid, flag, x0, x1, obj_factor, lambda, mode, rowidx, colptr, values)
            # @printf("#**********  str_eval_h - %s - %d  - %d %d  %f \n", mode, flag, rowid, colid, obj_factor)
            # @show x0, x1
            # @show lambda
            # x0, x1 = load_x("pips", instance.n_iter)

            # @show x0, x1
            # @show lambda
            m = instance.internalModel
            matHess = instance.matHess
            matHessMap = instance.matHessMap
            hcMap = instance.hcMap
            hrMap = instance.hrMap
            @assert rowid<=num_scenarios(m) && colid <=num_scenarios(m)
            if(mode == :Structure)
                if !haskey(matHess,colid)
                    e = instance.evaluatorMap[colid]
                    mm = getModel(m,rowid)

                    @timing instance.prof tic()
                    (h_J,h_I) = MathProgBase.hesslag_structure(e) # upper trangular
                    @timing instance.prof instance.t_jump += toq()

                    matHess[colid] = MatStorage(h_I,h_J,MathProgBase.numvar(mm),MathProgBase.numvar(mm))
                end
                mat = matHess[colid]
                h_I = mat.rowIdx
                h_J = mat.colIdx

                p = Pair{Int,Int}(rowid,colid)
                if !haskey(matHessMap,p)
                    if rowid == colid  #diagonal
                        col_var_idx, row_var_idx = get_h_col_idx_map(m,rowid,rowid,hcMap,hrMap)
                        new_h_I = Vector{Int}()
                        new_h_J = Vector{Int}()
                        for i = 1:length(h_I)
                            ii = h_I[i]
                            jj = h_J[i]
                            if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                                new_ii = row_var_idx[ii]
                                new_jj = col_var_idx[jj]
                                if new_ii < new_jj
                                    push!(new_h_I,new_ii)
                                    push!(new_h_J,new_jj)
                                else
                                    push!(new_h_I,new_jj)
                                    push!(new_h_J,new_ii)
                                end
                            end
                        end
                        matP = MatStorage(new_h_I, new_h_J, getNumVars(m,rowid), getNumVars(m,rowid))
                        matHessMap[p] = matP
                        fill!(matP.value,1.0)
                    elseif rowid == 0  && colid != 0 #border corss hessian
                        col_var_idx, row_var_idx = get_h_col_idx_map(m,rowid,colid,hcMap,hrMap)
                        new_h_I = Vector{Int}()
                        new_h_J = Vector{Int}()
                        for i = 1:length(h_I)
                            ii = h_I[i]
                            jj = h_J[i]
                            if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                                new_ii = row_var_idx[ii]
                                new_jj = col_var_idx[jj]
                                push!(new_h_I,new_ii)
                                push!(new_h_J,new_jj)
                            end
                        end
                        matP = MatStorage(new_h_I, new_h_J, getNumVars(m,colid), getNumVars(m,rowid))
                        matHessMap[p] = matP
                        fill!(matP.value,1.0)
                        # @show h_I, h_J
                        # @show col_var_idx
                        # @show row_var_idx
                        # @show new_h_I, new_h_J
                        # @show matP

                        #structure for root contribution
                        col_var_idx,row_var_idx = get_h_col_idx_map(m,colid,rowid,hcMap,hrMap)
                        # @show col_var_idx
                        # @show row_var_idx
                        new_h_I = Vector{Int}()
                        new_h_J = Vector{Int}()
                        for i = 1:length(h_I)
                            ii = h_I[i]
                            jj = h_J[i]
                            if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                                new_ii = row_var_idx[ii]
                                new_jj = col_var_idx[jj]
                                if new_ii < new_jj
                                    push!(new_h_I,new_ii)
                                    push!(new_h_J,new_jj)
                                else
                                    push!(new_h_I,new_jj)
                                    push!(new_h_J,new_ii)
                                end
                            end
                        end
                        # @show h0_I,h0_J
                        p0 = Pair{Int,Int}(colid,rowid)
                        matP0 = MatStorage(new_h_I, new_h_J, getNumVars(m,0),getNumVars(m,0))
                        matHessMap[p0] = matP0
                    else
                        @assert (rowid !=0 && colid == 0)
                        @assert false
                    end
                end
                matP = matHessMap[p]

                @timing instance.prof tic()
                laghess = sparse(matP.rowIdx,matP.colIdx, matP.value, matP.m, matP.n)
                @timing instance.prof begin
                    instance.t_hess_spconv += toq()
                    instance.n_hess_spconv += 1
                end
                return length(laghess.nzval)
            elseif(mode == :Values)
                p = Pair{Int,Int}(rowid,colid)
                if rowid == colid
                    @assert haskey(matHess,rowid)
                    mat = matHess[rowid]
                    if !mat.isVal
                        e = instance.evaluatorMap[rowid]
                        (eq_idx, ieq_idx) = instance.iMap[rowid]
                        numeq = length(eq_idx)
                        lam_new = Vector{Float64}(undef, length(lambda))
                        for i in eq_idx
                            lam_new[i[1]] = lambda[i[2]]
                        end
                        for i in ieq_idx
                            lam_new[i[1]] = lambda[i[2]+numeq]
                        end
                        x = build_x(m,rowid,x0,x1)
                        # @show x,lam_new
                        @timing instance.prof tic()
                        MathProgBase.eval_hesslag(e,mat.value,x,obj_factor,lam_new)
                        @timing instance.prof instance.t_jump += toq()

                        mat.isVal = true
                    end
                    h_I = mat.rowIdx
                    h_J = mat.colIdx
                    h = mat.value
                    # @show x0,x1
                    # @show x
                    # @show lambda
                    # @show h_I
                    # @show h_J
                    # @show h
                    col_var_idx,row_var_idx = get_h_col_idx_map(m,rowid, colid,hcMap,hrMap)
                    # @show col_var_idx
                    # @show row_var_idx
                    matP = matHessMap[p]
                    matPi = 1
                    for i = 1:length(h_I)
                        ii = h_I[i]
                        jj = h_J[i]
                        if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                            matP.value[matPi] =  h[i]
                            matPi += 1
                        end
                    end
                    # @show new_h_I
                    # @show new_h_J
                    # @show new_h
                    @timing instance.prof tic()
                    str_laghess = sparse(matP.rowIdx,matP.colIdx,matP.value,matP.m,matP.n,keepzeros = true)
                    @timing instance.prof begin
                        instance.t_hess_spconv += toq()
                        instance.n_hess_spconv += 1
                    end
                    # @printf("m=%d;n=%d; \n",getNumVars(m,rowid),getNumVars(m,rowid))
                    # @show new_h_I, new_h_J, new_h
                    # @printf(" hess%d%d=sparse(new_h_I,new_h_J,new_h, m, n); \n",rowid,colid)
                    # @show str_laghess
                    array_copy(str_laghess.rowval,1,rowidx,1,length(str_laghess.rowval))
                    array_copy(str_laghess.colptr,1,colptr,1,length(str_laghess.colptr))
                    array_copy(str_laghess.nzval, 1,values,1,length(str_laghess.nzval))
                elseif rowid == 0  && colid != 0 #border corss hessian
                    @assert haskey(matHess,colid)
                    mat = matHess[colid]
                    if !mat.isVal
                        e = instance.evaluatorMap[colid]
                        x = build_x(m,colid,x0,x1)
                        (eq_idx, ieq_idx) = instance.iMap[colid]
                        numeq = length(eq_idx)
                        lam_new = Vector{Float64}(undef, length(lambda))
                        for i in eq_idx
                            lam_new[i[1]] = lambda[i[2]]
                        end
                        for i in ieq_idx
                            lam_new[i[1]] = lambda[i[2]+numeq]
                        end

                        @timing instance.prof tic()
                        MathProgBase.eval_hesslag(e,mat.value,x,obj_factor,lam_new)
                        @timing instance.prof instance.t_jump += toq()

                        mat.isVal = true
                    end
                    h_I = mat.rowIdx
                    h_J = mat.colIdx
                    h = mat.value
                    # @show h_I, h_J
                    # @show h

                    col_var_idx,row_var_idx = get_h_col_idx_map(m,0, colid,hcMap,hrMap)
                    matP = matHessMap[p]
                    matPi = 1
                    for i = 1:length(h_I)
                        ii = h_I[i]
                        jj = h_J[i]
                        if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                            matP.value[matPi] = h[i]
                            matPi += 1
                        end
                    end
                    @timing instance.prof tic()
                    str_laghess = sparse(matP.rowIdx,matP.colIdx,matP.value,matP.m,matP.n,keepzeros = true)
                    @timing instance.prof begin
                        instance.t_hess_spconv += toq()
                        instance.n_hess_spconv += 1
                    end
                    # @printf("m=%d;n=%d; \n",getNumVars(m,rowid),getNumVars(m,rowid))
                    # @show new_h_I, new_h_J, new_h
                    # @printf(" hess%d%d=sparse(new_h_I,new_h_J,new_h, m, n); \n",rowid,colid)
                    # @show str_laghess
                    array_copy(str_laghess.rowval,1,rowidx,1,length(str_laghess.rowval))
                    array_copy(str_laghess.colptr,1,colptr,1,length(str_laghess.colptr))
                    array_copy(str_laghess.nzval, 1,values,1,length(str_laghess.nzval))
                elseif rowid !=0 && colid == 0 #root contribution
                    @assert flag == 2
                    @assert haskey(matHess,rowid)
                    mat = matHess[rowid]
                    if !mat.isVal
                        e = instance.evaluatorMap[rowid]
                        x = build_x(m,rowid,x0,x1)
                        (eq_idx, ieq_idx) = instance.iMap[rowid]
                        numeq = length(eq_idx)
                        lam_new = Vector{Float64}(undef, length(lambda))
                        for i in eq_idx
                            lam_new[i[1]] = lambda[i[2]]
                        end
                        for i in ieq_idx
                            lam_new[i[1]] = lambda[i[2]+numeq]
                        end
                        @timing instance.prof tic()
                        MathProgBase.eval_hesslag(e,mat.value,x,obj_factor,lam_new)
                        @timing instance.prof instance.t_jump += toq()

                        mat.isVal = true
                    end
                    h_I = mat.rowIdx
                    h_J = mat.colIdx
                    h = mat.value
                    # @show h_I, h_J
                    # @show h
                    col_var_idx,row_var_idx = get_h_col_idx_map(m, rowid, 0,hcMap,hrMap)
                    matP = matHessMap[p]
                    matPi = 1
                    for i = 1:length(h_I)
                        ii = h_I[i]
                        jj = h_J[i]
                        if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                            matP.value[matPi] = h[i]
                            matPi += 1
                        end
                    end
                    matP0 = matHess[0]
                    h0_I = matP0.rowIdx
                    h0_J = matP0.colIdx
                    @assert matP0.m == matP.m == matP0.n == matP.n == getNumVars(m,0)

                    # @show h0_I,h0_J
                    @timing instance.prof tic()
                    str_laghess = sparse([matP.rowIdx;h0_I], [matP.colIdx;h0_J], [matP.value;zeros(Float64,length(h0_I))], matP.m, matP.n, keepzeros=true)
                    @timing instance.prof begin
                        instance.t_hess00_spconv += toq()
                        instance.n_hess00_spconv += 1
                    end
                    # new_h_I = [new_h_I;h0_I]
                    # new_h_J = [new_h_J;h0_J]
                    # new_h = [new_h;zeros(Float64,length(h0_I))]
                    # @printf("m=%d; n=%d; \n",getNumVars(m,0),getNumVars(m,0))
                    # @show new_h_I, new_h_J, new_h
                    # @printf(" hess%d%d=sparse(new_h_I,new_h_J,new_h, m, n); \n",rowid,colid)
                    array_copy(str_laghess.rowval,1,rowidx,1,length(str_laghess.rowval))
                    array_copy(str_laghess.colptr,1,colptr,1,length(str_laghess.colptr))
                    array_copy(str_laghess.nzval, 1,values,1,length(str_laghess.nzval))
                else
                    @assert false
                end
                # filename = string("hess_",rowid,"_",colid)
                # write_mat_to_file(filename,str_laghess)

                convert_to_c_idx(rowidx)
                convert_to_c_idx(colptr)
                if flag == 2
                    mat.isVal = false
                end
                # @show values
            else
                @assert false mode
            end
            # @show "end str_eval_h"
        end

        instance.str_write_solution = function(id, x, y_eq, y_ieq)
            # @show id, x, y_eq, y_ieq
            @assert id in getLocalBlocksIds(instance.internalModel)
            @assert length(x) == instance.get_num_cols(id)
            @assert length(y_eq) == instance.get_num_eq_cons(id)
            @assert length(y_ieq) == instance.get_num_ineq_cons(id)

            #write back the primal to JuMP
            mm = getModel(instance.internalModel,id)
            for i = 1:length(x)
                setvalue(Variable(mm,i), x[i])
            end
        end

        return instance
    end
end

# function show_time(m::ModelInterface)
#     t_model_evaluation = m.t_sj_init +
#     m.t_sj_init_x0 +
#     m.t_prob_info +
#     m.t_eval_f +
#     m.t_eval_g+
#     m.t_eval_grad_f+
#     m.t_eval_jac_g+
#     m.t_eval_h +
#     m.t_write_solution;
#     return t_model_time
# end

#######
# Linking with PIPS Julia Structure interface
######

function structJuMPSolve(model; with_prof=false, suppress_warmings=false,kwargs...)
    # @show "solve"
    t_sj_lifetime = 0.0
    @timing with_prof tic()

    # MPI.Init() ＃initialize in model loading

    comm = getStructure(model).mpiWrapper.comm
    # @show "[$(MPI.Comm_rank(comm))/$(MPI.Comm_size(comm))] create problem "

    t_sj_model_init = 0.0
    @timing with_prof tic()

    prob = PipsNlpSolver.createProblemStruct(comm, StructJuMPModel(model,with_prof), with_prof)

    @timing with_prof t_sj_model_init += toq()

    # @show "end createStructJuMPPipsNlpProblem"

    t_sj_solver_total = 0.0
    @timing with_prof tic()

    status = PipsNlpSolver.solveProblemStruct(prob)

    @timing with_prof t_sj_solver_total += toq()

    @timing with_prof t_sj_lifetime += toq()

    # solver_time = solver_total - modeling_time
    # if(0==MPI.Comm_rank(MPI.COMM_WORLD))
    #   @printf "Total time %.4f (initialization=%.3f modelling=%.3f solver=%.3f) (in sec)\n" t_total prob.model.t_sj_init modeling_time solver_time
    # end
    @timing with_prof begin
        mid, nprocs = getMyRank()
        # bname = string(split(ARGS[1],"/")[2],"_",num_scenarios(model))
        # run(`mkdir -p ./out/$bname`)
        # fname = string("./out/",bname,"/",bname,"_",nprocs,".",mid,".jl.txt")
        # # @show bname, fname
        # tfile = open(fname, "w")
        s1 = @sprintf("[%d/%d] [ t_sj_model_init %f t_sj_solver_total %f  t_sj_lifetime %f ] \n", mid, nprocs, t_sj_model_init, t_sj_solver_total, t_sj_lifetime)
        s2 = @sprintf("[%d/%d] [ t_jl_str_total %f t_jl_eval_total %f  ] \n", mid, nprocs, prob.t_jl_str_total, prob.t_jl_eval_total)
        s3 = @sprintf("[%d/%d] [ t_jac_spconv %f n_jac_spconv %d  ] \n", mid, nprocs, prob.model.t_jac_spconv, prob.model.n_jac_spconv)
        s4 = @sprintf("[%d/%d] [ t_hess_spconv %f n_hess_spconv %d  ] \n", mid, nprocs, prob.model.t_hess_spconv, prob.model.n_hess_spconv)
        s5 = @sprintf("[%d/%d] [ t_hess00_spconv %f n_hess00_spconv %d  ] \n", mid, nprocs, prob.model.t_hess00_spconv,prob.model.n_hess00_spconv)
        s6 = @sprintf("[%d/%d] [ t_jump %f  ] \n", mid, nprocs, prob.model.t_jump)
        s7 = @sprintf("[%d/%d] [ t_eval_g %f t_itermap_eval_g %f n_eval_g %d ] \n", mid, nprocs, prob.model.t_eval_g, prob.model.t_itermap_eval_g, prob.model.n_eval_g)
        if(mid == 0)
            @printf("%s",s1)
            @printf("%s",s2)
            @printf("%s",s3)
            @printf("%s",s4)
            @printf("%s",s5)
            @printf("%s",s6)
            @printf("%s",s7)
        end
        # @printf(tfile, "%s", s1)
        # @printf(tfile, "%s", s2)
        # @printf("%s",s3)
        # @printf("%s",s4)
        # @printf("%s",s5)
        # @printf("%s",s6)
        # @printf("%s",s7)
        # close(tfile)
        # n1 = string("./out/",nprocs,".",mid,".c.txt")
        # n2 = string("./out/",bname,"/",bname,"_",nprocs,".",mid,".c.txt")
#        run(`mv $n1 $n2`)
    end

    return PIPSRetCodeToSolverInterfaceCode[status]
end

# function load_x(subdir,iter)
#     @printf("load x from ./%s/x%d \n",subdir,iter)
#     x = readdlm(string("./",subdir,"/x",iter))
#     x0 = x[1:78]
#     x1 = x[79:144]
#     return x0, x1
# end

# function write_x0(subdir,iter,x)
#     @printf("writing x to ./%s/x0_%d \n",subdir,iter)
#     run(`mkdir -p ./$subdir`)
#     writedlm(string("./",subdir,"/x0_",iter),x,",")
# end
# function write_x1(subdir,iter,x)
#     @printf("writing x to ./%s/x1_%d \n",subdir,iter)
#     run(`mkdir -p ./$subdir`)
#     writedlm(string("./",subdir,"/x1_",iter),x,",")
# end

function initialization(instance::StructJuMPModel)
    iMap = instance.iMap
    model = instance.internalModel
    evaluatorMap = instance.evaluatorMap
    @assert length(iMap) == length(evaluatorMap) == 0
    for id in getLocalBlocksIds(model)
        mm = getModel(model,id)
        e = JuMP.NLPEvaluator(mm)
        MathProgBase.initialize(e,[:Grad,:Jac,:Hess])
        push!(evaluatorMap,id=>e)
        eq_idx = Dict{Int,Int}()
        ieq_idx = Dict{Int,Int}()
        push!(iMap,id=>Pair(eq_idx,ieq_idx))
        lb,ub=JuMP.constraintbounds(mm)

        for i =1:length(lb)
            if lb[i] == ub[i]
                eq_idx[i] = length(eq_idx) + 1
            else
                ieq_idx[i] =length(ieq_idx) + 1 #remember to offset length(eq_idx)
            end
        end
    end
end

function get_jac_col_idx_map(m,rowid,colid,map)
    p = Pair{Int,Int}(rowid,colid)
    if haskey(map,p)
        idx_map = map[p]
    else
        idx_map = get_jac_col_var_idx(m,rowid, colid)
        push!(map,p=>idx_map)
    end
    return idx_map
end

function get_jac_col_var_idx(m,rowid, colid)  #this method customerized for no linking constraint presented
    # @show "get_jac_col_var_idx",rowid,colid
    idx_map = Dict{Int,Int}() #dummy (jump) -> actual used
    if rowid == colid
        nvar = getNumVars(m,rowid)
        for i = 1:nvar
            idx_map[i] = i
        end
    else
        @assert rowid!=0 && rowid != colid
        mm = getModel(m,rowid)
        othermap = getStructure(mm).othermap
        for p in othermap
            pidx = p[1].col
            cidx = p[2].col
            idx_map[cidx] = pidx
        end
    end
    # @show idx_map
    return idx_map
end

function get_h_col_idx_map(m,rowid,colid,cmap,rmap)
    p = Pair{Int,Int}(rowid,colid)
    if haskey(cmap,p)
        cidx_map = cmap[p]
        ridx_map = rmap[p]
    else
        cidx_map,ridx_map = get_h_var_idx(m,rowid,colid)
        push!(cmap,p=>cidx_map)
        push!(rmap,p=>ridx_map)
    end
    return cidx_map, ridx_map
end

function get_h_var_idx(m,rowid, colid)
    # @show "get_h_var_idx",rowid,colid
    col_idx_map = Dict{Int,Int}() #dummy (jump) -> actual used
    row_idx_map = Dict{Int,Int}()
    if rowid == colid
        nvar = getNumVars(m,rowid) #need to place model variable in front of non model variable.
        for i = 1:nvar
            col_idx_map[i] = i
            row_idx_map[i] = i
        end
    elseif rowid == 0  && colid != 0 #border
        mm = getModel(m,colid)
        othermap = getStructure(mm).othermap
        for p in othermap
            pidx = p[1].col
            cidx = p[2].col
            col_idx_map[cidx] = pidx
        end
        for i = 1:getNumVars(m,colid)
            row_idx_map[i] = i
        end
    elseif colid == 0 && rowid != 0 #root contrib.
        mm = getModel(m,rowid)
        othermap = getStructure(mm).othermap
        for p in othermap
            pidx = p[1].col
            cidx = p[2].col
            col_idx_map[cidx] = pidx
            row_idx_map[cidx] = pidx
        end
    else
        @assert false rowid colid
    end
    # @show col_idx_map,row_idx_map
    return col_idx_map,row_idx_map
end

KnownSolvers["PipsNlp"] = PipsNlpInterface.structJuMPSolve

end


"""
```
Li2020{T} <: AbstractNLCTFPModel{T}
```

The `Li2020` type defines the structure of the model in Li, Wenhao (2020), "Public Liquidity and Financial Crises."

### Fields

#### Parameters
* `parameters::Vector{AbstractParameter}`: Vector of all time-invariant model
  parameters.

* `keys::OrderedDict{Symbol,Int}`: Maps human-readable names for all model
  parameters to their indices in `parameters`.

#### Inputs to Measurement and Equilibrium Functional Equations

The following fields are dictionaries that map human-readable names to indices.

* `state_variables::OrderedDict{Symbol,Int}`: Maps each state variable to its dimension on a Cartesian grid

* `endogenous_variables::OrderedDict{Symbol,Int}`: Maps endogenous variables
    calculated during the solution of functional equations to an index

* `exogenous_shocks::OrderedDict{Symbol,Int}`: Maps each shock to a column in
  the measurement equation

* `observables::OrderedDict{Symbol,Int}`: Maps each observable to a row in the
  model's measurement equation matrices.

* `pseudo_observables::OrderedDict{Symbol,Int}`: Maps each pseudo-observable to
  a row in the model's pseudo-measurement equation matrices.

#### Model Specifications and Settings

* `spec::String`: The model specification identifier, \"an_schorfheide\", cached
  here for filepath computation.

* `subspec::String`: The model subspecification number, indicating that some
  parameters from the original model spec (\"ss0\") are initialized
  differently. Cached here for filepath computation.

* `settings::Dict{Symbol,Setting}`: Settings/flags that affect computation
  without changing the economic or mathematical setup of the model.

* `test_settings::Dict{Symbol,Setting}`: Settings/flags for testing mode

#### Other Fields

* `rng::MersenneTwister`: Random number generator. Can be is seeded to ensure
  reproducibility in algorithms that involve randomness (such as
  Metropolis-Hastings).

* `testing::Bool`: Indicates whether the model is in testing mode. If `true`,
  settings from `m.test_settings` are used in place of those in `m.settings`.

* `observable_mappings::OrderedDict{Symbol,Observable}`: A dictionary that
  stores data sources, series mnemonics, and transformations to/from model units.
  DSGE.jl will fetch data from the Federal Reserve Bank of
  St. Louis's FRED database; all other data must be downloaded by the
  user. See `load_data` and `Observable` for further details.

* `pseudo_observable_mappings::OrderedDict{Symbol,PseudoObservable}`: A
  dictionary that stores names and transformations to/from model units. See
  `PseudoObservable` for further details.
"""
mutable struct Li2020{T} <: AbstractNLCTFPModel{T}
    parameters::ParameterVector{T}                         # vector of all time-invariant model parameters
    keys::OrderedDict{Symbol,Int}                          # human-readable names for all the model
                                                           # parameters and steady-states
    state_variables::OrderedDict{Symbol,Int}               # dimension number of state variable

    functional_variables::OrderedDict{Symbol,Int}
    derivatives::OrderedDict{Symbol,Int}
    endogenous_variables::OrderedDict{Symbol,Int}
    exogenous_shocks::OrderedDict{Symbol,Int}
    observables::OrderedDict{Symbol,Int}
    pseudo_observables::OrderedDict{Symbol,Int}

    spec::String                                           # Model specification number (eg "m990")
    subspec::String                                        # Model subspecification (eg "ss0")
    settings::Dict{Symbol,Setting}                         # Settings/flags for computation
    test_settings::Dict{Symbol,Setting}                    # Settings/flags for testing mode
    rng::MersenneTwister                                   # Random number generator
    testing::Bool                                          # Whether we are in testing mode or not

    observable_mappings::OrderedDict{Symbol, Observable}
    pseudo_observable_mappings::OrderedDict{Symbol, PseudoObservable}
end

description(m::Li2020) = "Julia implementation of model defined in 'Public Liquidity and Financial Crises' (2020) by Wenhao Li: Li2020, $(m.subspec)"

"""
`init_model_indices!(m::Li2020)`

Arguments:
`m:: Li2020`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and observables.
"""
function init_model_indices!(m::Li2020)
    # State_Variables
    state_variables = collect([:w])

    # Exogenous shocks
    exogenous_shocks = collect([:K_sh, :N_sh]) # capital shock K, liquidity shock N

    # Variables for functional equations
    functional_variables = collect([:p, :Q̂, :xg])

    # Derivatives of variables
    init_derivatives!(m, Dict{Symbol, Vector{Int}}(:p => standard_derivs(1)), state_variables)
    derivatives = keys(get_derivatives(m))

    # Endogenous variables
    endogenous_variables = collect([:ψ, :xK, :yK, :yg, :σp, :σ, :σh, :σw, :μR_rd, :rd_rg, :rd_rg_H, :μb_μh, :μw, :μp, :μK, :μR, :rd, :rg, :rd_rf,
                                    :μb, :μh, :invst, :lvg, :κp, :κb, :κd, :κh, :κfs, :firesale_jump, :κw, :liq_prem, :bank_liq_frac,
                                    :δ_x, :indic, :rf, :rh, :K_growth, :κK])

    # Observables
    observables = keys(m.observable_mappings)

    # Pseudo-observables
    pseudo_observables = keys(m.pseudo_observable_mappings)

    for (i,k) in enumerate(state_variables);        m.state_variables[k]      = i end
    for (i,k) in enumerate(exogenous_shocks);       m.exogenous_shocks[k]     = i end
    for (i,k) in enumerate(endogenous_variables);   m.endogenous_variables[k] = i end
    for (i,k) in enumerate(functional_variables);   m.functional_variables[k] = i end
    for (i,k) in enumerate(derivatives);            m.derivatives[k]          = i end
    for (i,k) in enumerate(observables);            m.observables[k]          = i end
    for (i,k) in enumerate(pseudo_observables);     m.pseudo_observables[k]   = i end
end

function Li2020(subspec::String = "ss0";
                custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                testing = false)

    # Model-specific specifications
    spec               = split(basename(@__FILE__),'.')[1]
    subspec            = subspec
    settings           = Dict{Symbol,Setting}()
    test_settings      = Dict{Symbol,Setting}()
    rng                = MersenneTwister(0)

    # initialize empty model
    m = Li2020{Float64}(
            # model parameters and steady state values
            Vector{AbstractParameter{Float64}}(), OrderedDict{Symbol,Int}(),

            # model indices
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),

            spec,
            subspec,
            settings,
            test_settings,
            rng,
            testing,
            OrderedDict{Symbol,Observable}(),
            OrderedDict{Symbol,PseudoObservable}())

    # Set settings
    model_settings!(m)

    # default_test_settings!(m)
    for custom_setting in values(custom_settings)
        m <= custom_setting
    end

    # Set observable and pseudo-observable transformations
    # init_observable_mappings!(m)
    # init_pseudo_observable_mappings!(m)

    # Initialize parameters
    init_parameters!(m)

    init_model_indices!(m)
    init_subspec!(m)

    return m
end

"""
```
init_parameters!(m::Li2020)
```
Initializes the model's parameters.
"""
function init_parameters!(m::Li2020)
    # Productivity
    m <= parameter(:AH, 0.15, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(),
                   Uniform(0., 1e3), fixed = false, description = "Banker's productivity")
    m <= parameter(:AL, 0.125, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Household's productivity")

    # Bank run and fire sales
    m <= parameter(:λ, 0.025, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.),
                   fixed = false, description = "Probability of liquidity shock")
    m <= parameter(:β, 0.25, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Probability of running on a deposit account")
    m <= parameter(:α, 0.21, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Illiquidity discount on capital during fire sale")
    m <= parameter(:ϵ, 1e-3, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Fraction of remaining wealth for bankers after bankruptcy")
    m <= parameter(:θ, 1e-5, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Probability of non-zero exposure to liquidity shock")

    # Shadow value of liquidity
    m <= parameter(:π, 0.19, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Illiquidity discount on illiquid safe asset")

    # Macro parameters
    m <= parameter(:δ, 0.1, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.),
                   fixed = false, description = "Depreciation rate of capital")
    m <= parameter(:ρ, 0.04, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Discount rate")
    m <= parameter(:σK, 0.033, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3), fixed = false,
                   description = "Volatility of capital shocks")
    m <= parameter(:χ, 3., (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3), fixed = false,
                   description = "Adjustment cost of internal investment function")

    # Stationary distribution
    m <= parameter(:η, 0.05, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Rate of retirement for bankers")
end

"""
```
model_settings!(m::Li2020)
```
creates the model's settings.
"""
function model_settings!(m::Li2020)

    # Investment functions
    m <= Setting(:Φ, quadratic_investment, "Internal investment function")
    m <= Setting(:μK, growth_quadratic_investment, "Rate of capital growth")
    m <= Setting(:∂Φ, derivative_quadratic_investment_li2020, "Derivative of internal investment function")

    # Numerical settings for grid
    m <= Setting(:N, 100, "Grid size")
    m <= Setting(:stategrid_method, :exponential, "Type of grid to construct for state variables")
    m <= Setting(:stategrid_dimensions, OrderedDict{Symbol, Tuple{Float64, Float64, Int}}(:w => (1e-3, 1., get_setting(m, :N))),
                 "Information about the dimensions of the state space")
    m <= Setting(:stategrid_splice, 0.2, "Li (2020) constructs the grid in two parts. This value is where the first half stops.")

    # Numerical settings for no jump equilibrium
    m <= Setting(:boundary_conditions, OrderedDict{Symbol, Vector{Float64}}(:p => [0.; 0.]),
                 "Boundary conditions for differential equations.")
    m <= Setting(:ode_reltol, 1e-4, "Relative tolerance for ODE integration")
    m <= Setting(:ode_abstol, 1e-12, "Absolute tolerance for ODE integration")
    m <= Setting(:essentially_one, 0.999,
                "If ψ is larger than this value, then we consider it essentially one for some numerical purposes.")
    m <= Setting(:ode_integrator, DP5(), "Numerical ODE integrator for no-jump solution") # DP5 is Matlab's ode45

    # Numerical settings for functional iteration
    m <= Setting(:tol, 1e-4, "Tolerance for functional iteration")
    m <= Setting(:learning_rate, 0.4, "Learning rate for the update of functional variables after each loop")
    m <= Setting(:max_iter, 40, "Maximum number of loops during functional iteration")
    m <= Setting(:error_method, :total_error, "Method for calculating error at the end of each loop during functional iteration")
    m <= Setting(:v₀, 3e-8, "Parameter for damping function")
    m <= Setting(:damping_function, x -> get_setting(m, :v₀) ./ x, "Dampling function to avoid corners")
    m <= Setting(:p₀_perturb, 1e-14, "Perturbation of boundary condition for q at w = 0")
    m <= Setting(:κp_grid, Vector{Float64}(undef, 0),
                 "Vector of guesses for κp used during iteration for new κp and xK within each functional loop")
    m <= Setting(:p_interpolant, Gridded(Linear()), "Interpolant method for fitted p.")
    m <= Setting(:xK_interpolant, Gridded(Linear()), "Interpolant method for xK after solving xK for each guess of κp.")
    m <= Setting(:xg_interpolant, Gridded(Linear()), "Interpolant method for xg after solving xK and κp.")
    m <= Setting(:κp_interpolant, Gridded(Linear()), "Interpolant method for κp after solving xK and κp.")
    m <= Setting(:Q̂_interpolant, Gridded(Linear()), "Interpolant method for Q̂.")
    m <= Setting(:inside_iteration_nlsolve_tol, 1e-6, "Tolerance for nlsolve in the inside iteration")
    m <= Setting(:inside_iteration_nlsolve_max_iter, 400, "Maximum number of iterations for nlsolve for the inside iteration")
    m <= Setting(:xg_tol, 2e-3, "Lower but permissible tolerance for xg fixed point in the inside iteration")
    m <= Setting(:yg_tol, 1e-5, "Lowest permissible value for yg")
    m <= Setting(:p_tol, 1e-8, "Tolerance when calculating p after completing an inside iteration")
    m <= Setting(:firesale_bound, 0.99, "Upper bound for size of firesale jumps that do not wipe out a banker's net worth")
    m <= Setting(:firesale_interpolant, Gridded(Linear()), "Interpolant method for the firesale value.")
    m <= Setting(:N_GH, 10, "Number of nodes for Gauss-Hermite quadrature.")
    m <= Setting(:Q̂_tol, 1e-5, "Tolerance for approximation of Q̂ via Gauss-Hermite quadrature")
    m <= Setting(:Q̂_max_it, 1000, "Maximum number of iterations for approximation of Q̂.")
    m <= Setting(:dt, 1. / 12., "Simulation interval for Q̂ approximation as a fraction of a 1 year")

    # Other settings for initialization
    m <= Setting(:nojump_parameters, [:ρ, :AH, :AL, :σK, :χ, :δ, :η], "Keys of parameters used when solving the no-jump equilibrium.")

    # Calibration targets
    m <= Setting(:avg_gdp, 0.145, "Average GDP")
    m <= Setting(:liq_gdp_ratio, 0.39, "Average Liquidity/GDP ratio in the data")
    m <= Setting(:gov_bond_gdp_level, 0.3, "Ratio of Government Bonds to GDP")
end