function NetworkLoad(testsystem)
   PUcorrection = 1
   if testsystem == "AP15busDN"
      usingMatPower = 0
   elseif testsystem == "IEEE33busDN" || testsystem == "IEEE33busDN2" || testsystem == "ISONE8busTN"
      usingMatPower = 0
   elseif testsystem == "R2-25.00-1"
      usingMatPower = 0
      #PUcorrection = 0.01000
   else
      usingMatPower = 1
      PUcorrection = 100
      if testsystem =="case141_OPF"
         PUcorrection = 10
      end
   end

   if usingMatPower == 0
      filename_Node = pwd()"/data/$testsystem/Node.csv"
      filename_Generator = pwd()"/data/$testsystem/Generator.csv"
      filename_Line = pwd()"/data/$testsystem/Line.csv"

      busmat = readcsv(filename_Node, header=true)[1]
      buses = Bus[]
      for i in 1:size(busmat,1)
         bindex = busmat[i,1]
         btype = 0
         Pd = busmat[i,2]
         Qd = busmat[i,3]
         Gs = busmat[i,6]
         Bs = busmat[i,7]
         area = 0
         Vm = 0
         Va = 0
         baseKV = busmat[i,8]
         bzone = 0
         Vmax = busmat[i,4]
         Vmin = busmat[i,5]
         b = Bus(bindex, btype, Pd, Qd, Gs, Bs, area, Vm, Va, baseKV, bzone, Vmax, Vmin)
         push!(buses, b)
      end

      genmat = readcsv(filename_Generator, header=true)[1]
      generators = Generator[]
      for i in 1:size(genmat,1)
            gindex = genmat[i,1]
            gtype = genmat[i,2]
            location = genmat[i,3]
            Pg = genmat[i,4]
            Qg = genmat[i,5]
            Qmax = genmat[i,6]
            Qmin = genmat[i,7]
            Vg = genmat[i,8]
            mBase = genmat[i,9]
            status = genmat[i,10]
            Pmax = genmat[i,11]
            Pmin = genmat[i,12]
            cost = [genmat[i,13], genmat[i,14], genmat[i,15]]
            SUcost = genmat[i,16]
            SDcost = genmat[i,17]
            RU = genmat[i,18]
            RD = genmat[i,19]
            UPtime = genmat[i,20]
            DNtime = genmat[i,21]
            g = Generator(gindex, gtype, location, Pg, Qg, Qmax, Qmin, Vg, mBase, status, Pmax, Pmin, cost, SUcost, SDcost, RU, RD, UPtime, DNtime)
            push!(generators, g)
      end

      branchmat = readcsv(filename_Line, header=true)[1]
      lines = Line[]
      for i in 1:size(branchmat,1)
         lindex = i
         fbus = Int(branchmat[i,2])
         tbus = Int(branchmat[i,1])
         r = branchmat[i,3]
         x = branchmat[i,4]
         b = branchmat[i,5]
         u = branchmat[i,6]
         push!(buses[tbus].children, fbus)#children
         push!(buses[fbus].ancestor, tbus)#ancestor
         push!(buses[tbus].inline, lindex)
         push!(buses[fbus].outline, lindex)
         l = Line(lindex, fbus, tbus, r, x, b, u)
         push!(lines,l)
      end
      for l in 1:length(lines)
         for k in 1:length(lines)
            if lines[l].fbus == lines[k].tbus
               push!(lines[l].ancestor, k)
            elseif lines[l].tbus == lines[k].fbus
               push!(lines[l].children, k)
            end
         end
      end
      datamat = Dict()
      datamat["bus"] = busmat
      datamat["branch"] = branchmat
      datamat["gen"] = genmat
   elseif usingMatPower == 1
      ## MatPowerCases load ### ex) testsystem = "case6ww" /testsystem = "case118"
      mpc = loadcase(testsystem)
      datamat = mpc
      #Sbase = mpc["baseMVA"] * 1e6 # MVA to VA
      #Vbase = mpc["bus"][:,10] * 1e3 # kV to V
      #Zbase = Vbase^2 / Sbase

      buses = Bus[]
      for i in 1:size(mpc["bus"],1)
         bindex = mpc["bus"][i,1]
         btype = mpc["bus"][i,2]
         Pd = mpc["bus"][i,3]
         Qd = mpc["bus"][i,4]
         Gs = mpc["bus"][i,5]
         Bs = mpc["bus"][i,6]
         area = mpc["bus"][i,7]
         Vm = mpc["bus"][i,8]
         Va = mpc["bus"][i,9]
         baseKV = mpc["bus"][i,10]
         bzone = mpc["bus"][i,11]
         Vmax = mpc["bus"][i,12]
         Vmin = mpc["bus"][i,13]
         b = Bus(bindex, btype, Pd, Qd, Gs, Bs, area, Vm, Va, baseKV, bzone, Vmax, Vmin)
         push!(buses, b)
      end

      generators = Generator[]
      for i in 1:size(mpc["gen"],1)
         gindex = i
         gtype = "NotDefined"
         location = mpc["gen"][i,1]
         Pg = mpc["gen"][i,2]
         Qg = mpc["gen"][i,3]
         Qmax = mpc["gen"][i,4]
         Qmin = mpc["gen"][i,5]
         Vg = mpc["gen"][i,6]
         mBase = mpc["gen"][i,7]
         status = mpc["gen"][i,8]
         Pmax = mpc["gen"][i,9]
         Pmin = mpc["gen"][i,10]
         cost = mpc["gencost"][i,5:end]
         SUcost = mpc["gencost"][i,2]
         SDcost = mpc["gencost"][i,3]
         RU = 0
         RD = 0
         UPtime = 0
         DNtime = 0
         g = Generator(gindex, gtype, location, Pg, Qg, Qmax, Qmin, Vg, mBase, status, Pmax, Pmin, cost, SUcost, SDcost, RU, RD, UPtime, DNtime)
         push!(generators, g)
      end

      lines = Line[]
      for i in 1:size(mpc["branch"],1)
         lindex = i
         fbus = Int(mpc["branch"][i,2]) #definition of fbus is based on  {Low, 2013} which is opposite of Matpower
         tbus = Int(mpc["branch"][i,1])
         r = mpc["branch"][i,3]/PUcorrection
         x = mpc["branch"][i,4]/PUcorrection
         b = mpc["branch"][i,5]
         u = 10000
         push!(buses[tbus].children, fbus)#children
         push!(buses[fbus].ancestor, tbus)#ancestor
         push!(buses[tbus].inline, lindex)
         push!(buses[fbus].outline, lindex)
         l = Line(lindex, fbus, tbus, r, x, b, u)
         push!(lines,l)
      end
      #case1
      if testsystem == "case141_OPF"
         for l in 1:6
            lines[l].u = 50
         end
         for l in 7:length(lines)
            lines[l].u = 5.0
         end
      end
      #case2
      if testsystem == "case141_OPF"
         for l in 15:31
            #lines[l].u = 1.4
         end
         for l in 46:51
            #lines[l].u = 1.25
         end
      end

      for l in 1:length(lines)
         for k in 1:length(lines)
            if lines[l].fbus == lines[k].tbus
               push!(lines[l].ancestor, k)
            elseif lines[l].tbus == lines[k].fbus
               push!(lines[l].children, k)
            end
         end
      end
   end
   return buses, lines, generators, datamat
end


#!/usr/bin/env julia

#
#  This is the main file for data assimilation runs.
#
#
# 1. loads up a configuration file,
# 2. obtains data from a WRF model,
# 3. construct covariate vectors
# 3. reads in observations and metadata for a list of stations,
# 4. runs the moisture model and the assimilation mechanism.
#
#

using Calendar
import Calendar.CalendarTime

using Storage
import Storage.setup_tag, Storage.spush, Storage.next_frame, Storage.flush_frame

using Stations
import Stations.Station, Stations.Observation, Stations.load_station_info,
       Stations.load_station_data, Stations.build_observation_data, Stations.register_to_grid, Stations.nearest_grid_point, Stations.obs_value, Stations.obs_station_id

using Kriging
import Kriging.trend_surface_model_kriging

using NetCDF
using WRF

using FM
import FM.FMModel, FM.advance_model, FM.kalman_update

function create_fm_variables(file_name)

    # first extract the dimensions
    nc = NetCDF.open(file_name)
    we_dim = nc.dim["west_east"]
    sn_dim = nc.dim["south_north"]
    t_dim = nc.dim["Time"]
    NetCDF.close(nc)

    var_names = ASCIIString[ "FM1", "FM10", "FM100" ]
    long_names = ["1hr fuel moisture", "10hr fuel moisture", "100hr fuel moisture"]
    for i in 1:3
        if !haskey(nc.vars, var_names[i])
            atts = {"coordinates"=>"XLONG XLAT","units"=>"kg/kg","stagger"=>"","MemoryOrder"=>"XY ","description"=>long_names[i]}
            NetCDF.nccreate(file_name, var_names[i], atts, we_dim, sn_dim, t_dim)
            println("INFO: created variable $(var_names[i])")
        else
            println("INFO: variable $(var_names[i]) already existed")
        end
    end
end

function main(args)

    # the arguments passed to the julia program do not include program name
    if length(args) != 1
        println("Usage: julia run_data_assimilation.jl cfg_file")
        exit(1)
    end

    t1 = Calendar.now()
    println("INFO: run_data_assimilation.jl started on $t1")

    ### Read configuration file and setup the system
    println("INFO: reading in config file $(args[1])")
    cfg = evalfile(args[1])

    # create the output directory if it does not exist
    println("INFO: output directory is ", cfg["output_dir"])
    !ispath(cfg["output_dir"]) && mkdir(cfg["output_dir"])

    # configure Storage mechanism
    Storage.sopen(cfg["output_dir"], "moisture_model_v2_diagnostics.txt", "frame")

    # setup Storage & output policies for interesting quantities
    setup_tag("mt", true, true)

    setup_tag("fm1_model_state_assim", false, false)
    setup_tag("fm1_model_var", false, false)
    setup_tag("fm10_model_state", false, false)
    setup_tag("fm10_model_state_assim", false, false)
    setup_tag("fm10_model_var", false, false)
    setup_tag("fm100_model_state_assim", false, false)
    setup_tag("fm100_model_var", false, false)

    setup_tag("fm10_model_na_state", false, false)
    setup_tag("fm10_model_deltas", false, false)

    setup_tag("kriging_beta", true, true)
    setup_tag("kriging_xtx_cond", false, true)
    setup_tag("kriging_field", false, false)
    setup_tag("kriging_variance", false, false)
    setup_tag("kriging_sigma2_eta", true, true)
    setup_tag("kriging_iters", false, true)
    setup_tag("kriging_subzero_s2_estimates", false, true)
    setup_tag("model_raws_mae", false, true)
    setup_tag("model_raws_mae_assim", false, true)
    setup_tag("model_na_raws_mae", false, true)

    # co-located model/model_na/kriging field/observation
    setup_tag("kriging_obs", false, false)
    setup_tag("kriging_obs_station_ids", false, false)
    setup_tag("kriging_obs_ngp", false, false)
    setup_tag("kriging_errors", true, true)

    setup_tag("kalman_gain_fm10", false, false)

    ### Load WRF output data
    t1 = Calendar.now()
    println("INFO: configuration complete, loading WRF data.")

    # read in data from the WRF output file pointed to by cfg
    w = WRF.load_wrf_data(cfg["wrf_output"], ["HGT"])

    # the terrain height need not be stored for all time points
    WRF.slice_field(w, "HGT")

    # extract WRF fields
    lat, lon = WRF.lat(w), WRF.lon(w)
    dsize = size(lat)
    wtm = WRF.times(w)
    println("INFO: WRF grid size is $(dsize[1]) x $(dsize[2]) and found $(length(wtm)) timepoints.")

    # retrieve equilibria and rain (these are already precomputed)
    Ed, Ew = WRF.field(w, "Ed"), WRF.field(w, "Ew")
    rain = WRF.field(w, "RAIN")
    hgt = WRF.field(w, "HGT")
    T = WRF.interpolated_field(w, "T2")
    P = WRF.interpolated_field(w, "PSFC")

    t2 = Calendar.now()
    println("INFO: WRF output loaded, sliced and diced [$(t2-t1)].")

    ### Load observation data from stations
    io = open(join([cfg["station_info_dir"], cfg["station_info"]], "/"), "r")
    station_ids = filter(x -> x[1] != '#', map(x -> strip(x), readlines(io)))
    close(io)

    # load each station from its info and observation files
    stations = Station[]
    for sid in station_ids
        s = load_station_info(join([cfg["station_info_dir"], string(sid, ".info")], "/"))
        load_station_data(s, join([cfg["station_data_dir"], string(sid, ".obs")], "/"))
	register_to_grid(s, lat, lon)
#        println("STATION: $(s.id), $(s.loc), ngp is $(s.ngp) with lat $(lat[s.ngp[1], s.ngp[2]]) and lon $(lon[s.ngp[1], s.ngp[2]]).")
        push!(stations, s)
    end

    # build the observation data from stations
    obs_fm10 = build_observation_data(stations, "FM")
    obs_times = keys(obs_fm10)

    t3 = Calendar.now()
    println("INFO: Station data loaded and preprocessed [$(t3 - t2)].")

    ### Initialize model

    # number of simulated fuel components
    Nf = 3

    # construct initial conditions (FIXME: can we do better here?)
    E = squeeze(0.5 * (Ed[:,:,2] + Ew[:,:,2]), 3)

    # set up parameters
    Q = diagm(cfg["Q"])
    P0 = diagm(cfg["P0"])
    mV = zeros(Float64, dsize)
    pred = zeros(Float64, dsize)
    mresV = zeros(Float64, dsize)
    mid = zeros(Int32, dsize)
    Kg = zeros(Float64, (dsize[1], dsize[2], 9))
    K = zeros(Float64, dsize)
    V = zeros(Float64, dsize)

    # prepare static & time-varying covariates
    cov_ids = cfg["covariates"]
    st_covar_map = [:lon => lon,
                    :lat => lat,
                    :elevation => hgt,
                    :constant => ones(Float64, dsize) ]
    dyn_covar_map = [:temperature => T, :pressure => P, :rain => rain]
    Xd3 = length(cov_ids) + 1
    X = zeros(Float64, (dsize[1], dsize[2], Xd3))
    Xr = zeros(Float64, (dsize[1], dsize[2], Xd3))
    for i in 2:Xd3
        cov_id = cov_ids[i-1]
        if haskey(st_covar_map, cov_id)
            println("INFO: processing static covariate $cov_id.")
            v = st_covar_map[cov_id]
            Xr[:,:,i] = v
        elseif haskey(dyn_covar_map, cov_id)
            println("INFO: found dynamic covariate $(cov_id).")
        else
            error("ERROR: unknown covariate $(cov_id) encountered, fatal.")
        end
    end
    println("INFO: there are $Xd3 covariates (including model state).")

    t1 = Calendar.now()
    println("INFO: starting simulation at $t1 ...")
    dt = (wtm[2] - wtm[1]).millis / 1000
    assim_time_win = cfg["assimilation_time_window"]
    println("INFO: time step from WRF is $dt s, assimilation time window is $assim_time_win s.")

    # construct model grid from fuel parameters
    Tk = [ 1.0, 10.0, 100.0 ]
    models = Array(FMModel, dsize)
    models_na = Array(FMModel, dsize)
    for i in 1:dsize[1]
    	for j in 1:dsize[2]
            geo_loc = (lat[i,j], lon[i,j])
	    models[i,j] = FMModel(geo_loc, Nf, E[i,j], P0, Tk)
	    models_na[i,j] = FMModel(geo_loc, Nf, E[i,j], P0, Tk)
	end
    end

    # create the new variables in the wrfout file for storing fuel moisture
    println("INFO: creating FM variables in wrfout file")
    create_fm_variables(cfg["wrf_output"])
    nc = NetCDF.open(cfg["wrf_output"], NetCDF.NC_WRITE)

    ###  Run the model and data assimilation
    for t in 2:length(wtm)
    	mt = wtm[t]
        spush("mt", mt)

        # run the model update (in parallel if possible)
	for i in 1:dsize[1]
	    for j in 1:dsize[2]
	        advance_model(models[i,j], Ed[i, j, t-1], Ew[i, j, t-1], rain[i, j, t-1], dt, Q)
		advance_model(models_na[i,j], Ed[i, j, t-1], Ew[i, j, t-1], rain[i, j, t-1], dt, Q)
	    end
	end

        # store the model state in an array (and store in output frame)
        fm10_model_state = [ models[i,j].m_ext[2] for i=1:dsize[1], j=1:dsize[2] ]
        fm10_model_na_state = [ models_na[i,j].m_ext[2] for i=1:dsize[1], j=1:dsize[2] ]
        fm10_model_var = [ models[i,j].P[2,2] for i=1:dsize[1], j=1:dsize[2] ]

        spush("fm10_model_state", fm10_model_state)
        spush("fm10_model_na_state", fm10_model_na_state)
        spush("fm10_model_var", fm10_model_var)

        # if observation data for this timepoint is available
        obs_i = Observation[]
        tm_valid_now = filter(x -> abs((mt - x).millis) / 1000.0 <= assim_time_win/2, obs_times)

        # gather all observations
        for tvn in tm_valid_now append!(obs_i, obs_fm10[tvn]) end

        # exclude zero observations - must be sensor failure
        obs_i = filter(x -> obs_value(x) > 0, obs_i)

        # if there are no valid observations, continue with next time step, else run kriging
        if length(obs_i) > 0

            # set the current fm10 model state as the covariate
            X[:,:,1] = fm10_model_state

#            fm10_norm = sum(fm10_model_state.^2)^0.5
            println("INFO: assimilating $(length(obs_i)) observations.")

            # loop over additional covariates
            for i in 2:Xd3
                cov_id = cov_ids[i-1]
                if haskey(st_covar_map, cov_id)
                    # just copy and rescale corresponding static covariate
                    X[:,:,i] = Xr[:,:,i]
                elseif haskey(dyn_covar_map, cov_id)
                    # retrieve the field pointed to by the dynamic covariate id
                    F = dyn_covar_map[cov_id]
                    X[:,:,i] = squeeze(F[:,:,t], 3)
                else
                    error("FATAL: found unknown covariate.")
                end
            end

            # store diagnostic information
            ngp_list = map(x -> nearest_grid_point(x), obs_i)
            stat_ids = map(x -> obs_station_id(x), obs_i)
            m_at_obs = Float64[X[i, j, 1] for (i,j) in  ngp_list]
            m_na_at_obs = Float64[models_na[i,j].m_ext[2] for (i,j) in ngp_list]
            raws = Float64[obs_value(o) for o in obs_i]

            spush("model_raws_mae", mean(abs(m_at_obs - raws)))
            spush("model_na_raws_mae", mean(abs(m_na_at_obs - raws)))

            spush("kriging_obs", raws)
            spush("kriging_obs_station_ids", stat_ids)
            spush("kriging_obs_ngp", ngp_list)

            # compute the kriging estimates and fill in pre-allocated arrays
            trend_surface_model_kriging(obs_i, X, K, V)

            # push diagnostic outputs
            spush("kriging_field", K)
            spush("kriging_variance", V)

            # execute the Kalman update at each grid point
            Kp = zeros(1)
            Vp = zeros(1,1)
            Kg = zeros(Float64, dsize)
            fuel_types = [2]
            for i in 1:dsize[1]
                for j in 1:dsize[2]
                    Kp[1] = K[i,j]
                    Vp[1,1] = V[i,j]
                    Kg[i,j] = kalman_update(models[i,j], Kp, Vp, fuel_types)[1,1]
                end
            end

            # push the fm10 model state after the assimilation
            fm10_model_state = [ models[i,j].m_ext[2] for i=1:dsize[1], j=1:dsize[2] ]
            spush("fm10_model_state_assim", fm10_model_state)

            # retrieve adjustments to time constants and to equilibria
            fm10_adj = zeros(Float64, (6, dsize[1], dsize[2]))
            for i in 1:dsize[1]
                for j in 1:dsize[2]
                    fm10_adj[:,i,j] = models[i,j].m_ext[Nf+1:2*Nf+3]
                end
            end
            fm10_adj_max = [ max(abs(fm10_adj[i,:,:])) for i in 1:6 ]
            spush("fm10_model_deltas", fm10_adj_max)

            # gather model values at ngp points after assimilation
            m_at_obs = Float64[fm10_model_state[i, j] for (i,j) in  ngp_list]
            spush("model_raws_mae_assim", mean(abs(m_at_obs - raws)))

            spush("kalman_gain_fm10", Kg)

            # move to the next storage frame
            next_frame()

        end # if there is anything to assimilate

        # store current assimilated state in the wrfout file
        fm_stor = zeros(Float64, (dsize[1], dsize[2], 1))
        fm_stor[:,:,1] = [ models[i,j].m_ext[1] for i=1:dsize[1], j=1:dsize[2] ]
        NetCDF.putvar(nc, "FM1", [1, 1, t], fm_stor)
        fm_stor[:,:,1] = [ models[i,j].m_ext[2] for i=1:dsize[1], j=1:dsize[2] ]
        NetCDF.putvar(nc, "FM10", [1, 1, t], fm_stor)
        fm_stor[:,:,1] = [ models[i,j].m_ext[3] for i=1:dsize[1], j=1:dsize[2] ]
        NetCDF.putvar(nc, "FM100", [1, 1, t], fm_stor)

    end # for each time point

    # Close down the storage system
    Storage.sclose()

    NetCDF.close(nc)

    t2 = Calendar.now()
    println("INFO: simulation completed at $t2 after $(t2-t1).")

end

main(ARGS)


# cudnnDataType_t
const CUDNN_DATA_FLOAT = Cint(0)
const CUDNN_DATA_DOUBLE = Cint(1)
const CUDNN_DATA_HALF = Cint(2)
const CUDNN_DATA_INT8 = Cint(3)
const CUDNN_DATA_INT32 = Cint(4)
const CUDNN_DATA_INT8x4 = Cint(5)

const Cptr = Ptr{Cvoid}
datatype(::Type{Float32}) = CUDNN_DATA_FLOAT
datatype(::Type{Float64}) = CUDNN_DATA_DOUBLE
datatype(::Type{Float16}) = CUDNN_DATA_HALF
datatype(::Type{Int8}) = CUDNN_DATA_INT8
datatype(::Type{Int32}) = CUDNN_DATA_INT32

# cudnnNanPropagation_t
const CUDNN_NOT_PROPAGATE_NAN = Cint(0)
const CUDNN_PROPAGATE_NAN = Cint(1)

# cudnnMathType_t
const CUDNN_DEFAULT_MATH = Cint(0)
const CUDNN_TENSOR_OP_MATH = Cint(1)

# cudnnDeterminism_t
const CUDNN_NON_DETERMINISTIC = Cint(0)
const CUDNN_DETERMINISTIC = Cint(1)

# cudnnOpTensorOp_t
const CUDNN_OP_TENSOR_ADD = Cint(0)
const CUDNN_OP_TENSOR_MUL = Cint(1)
const CUDNN_OP_TENSOR_MIN = Cint(2)
const CUDNN_OP_TENSOR_MAX = Cint(3)
const CUDNN_OP_TENSOR_SQRT = Cint(4)
const CUDNN_OP_TENSOR_NOT = Cint(5)

# cudnnIndicesType_t
const CUDNN_32BIT_INDICES = Cint(0)
const CUDNN_64BIT_INDICES = Cint(1)
const CUDNN_16BIT_INDICES = Cint(2)
const CUDNN_8BIT_INDICES = Cint(3)

# cudnnPoolingMode_t
const CUDNN_POOLING_MAX = Cint(0)
const CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING = Cint(1)
const CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING = Cint(2)
const CUDNN_POOLING_MAX_DETERMINISTIC = Cint(3)

# cudnnLRNMode_t
const CUDNN_LRN_CROSS_CHANNEL_DIM1 = Cint(0)

# cudnnDivNormMode_t
const CUDNN_DIVNORM_PRECOMPUTED_MEANS = Cint(0)

# cudnnBatchNormMode_t
const CUDNN_BATCHNORM_PER_ACTIVATION = Cint(0)
const CUDNN_BATCHNORM_SPATIAL = Cint(1)
const CUDNN_BATCHNORM_SPATIAL_PERSISTENT = Cint(2)

# cudnnSamplerType_t
const CUDNN_SAMPLER_BILINEAR = Cint(0)

# cudnnCTCLossAlgo_t
const CUDNN_CTC_LOSS_ALGO_DETERMINISTIC = Cint(0)
const CUDNN_CTC_LOSS_ALGO_NON_DETERMINISTIC = Cint(1)


using consonance

push!(LOAD_PATH, "/home/ian/Grad/Research/Disks/DiskJockey/")
push!(LOAD_PATH, "/n/home07/iczekala/DiskJockey/")
push!(LOAD_PATH, "/pool/scout0/DiskJockey/")

using ArgParse

s = ArgParseSettings()
@add_arg_table s begin
    # "--opt1"
    # help = "an option with an argument"
    "--run_index", "-r"
    help = "Output run index"
    arg_type = Int
end

parsed_args = parse_args(ARGS, s)
run_index = parsed_args["run_index"]

imgdir = scratchdir * @sprintf("d%02d/", run_index)

# make the directory where RADMC will make all its files exist
if !ispath(imgdir)
    println("Creating ", imgdir)
    mkdir(imgdir)
end

cd(imgdir)
println("Now in ", pwd())

files = ["dustopac.inp", "dustkappa_silicate.inp", "stars.inp", "lines.inp", "molecule_co.inp", "wavelength_micron.inp"]
for file in files
    src = homedir * file
    dst = imgdir * file
    cp(src, dst)
end

cp(homedir * "radmc3d.inp.dust", imgdir * "radmc3d.inp")

using constants
using image
using model

function make_image(pars, id::Int)

    write_dust(pars, "", grid)
    write_model(pars, "", grid)

    # RADMC conventions for inclination and PA
    incl = pars.incl # [deg]
    PA = pars.PA # [deg] Position angle runs counter clockwise

    lam0 = cc/230.538e9 * 1e4 # [microns]

    run(`radmc3d mctherm` |> DevNull)
    run(`radmc3d image incl $incl posang $PA npix $npix lambda $lam0` |> DevNull)

    src = "image.out"
    dst = outdir * @sprintf("dimage%04d.out", id)

    cp(src, dst)

end

const global npix = 256 # number of pixels
const global grid = Grid(100, 32, 0.5, 800, true)
write_grid("", grid)

start = run_index * nframes_per_proc + 1
ids = Int[i for i=start:(start + nframes_per_proc)]

for i in ids
    if i <= nframes
        make_image(pars[i], i)
    else
        break
    end
end


VERSION >= v"0.4.0-dev+6521" && __precompile__()

module Images

import Base.Order: Ordering, ForwardOrdering, ReverseOrdering
import Base: ==, .==, +, -, *, /, .+, .-, .*, ./, .^, .<, .>
import Base: atan2, clamp, convert, copy, copy!, ctranspose, delete!, done, eltype,
             fft, float32, float64, get, getindex, haskey, hypot, ifft, length, map, map!,
             maximum, mimewritable, minimum, next, ndims, one, parent, permutedims, reinterpret,
             setindex!, show, showcompact, similar, size, slice, sqrt, squeeze,
             start, strides, sub, sum, write, writemime, zero
# "deprecated imports" are below

using Color, FixedPointNumbers, Compat
import Color: Fractional
if VERSION < v"0.4.0-dev+3275"
    using Base.Graphics
    import Base.Graphics: width, height, Point
else
    using Graphics
    import Graphics: width, height, Point
end
import FixedPointNumbers: ufixed8, ufixed10, ufixed12, ufixed14, ufixed16

using Base.Cartesian
include("compatibility/forcartesian.jl")

# if isdefined(module_parent(Images), :Grid)
#     import ..Grid.restrict
# end

const is_little_endian = ENDIAN_BOM == 0x04030201
immutable TypeConst{N} end  # for passing compile-time constants to functions

include("colortypes.jl")
using .ColorTypes

include("core.jl")
include("map.jl")
include("overlays.jl")
const have_imagemagick = include("ioformats/libmagickwand.jl")
@osx_only include("ioformats/OSXnative.jl")
include("io.jl")
include("labeledarrays.jl")
include("algorithms.jl")
include("connected.jl")
include("edge.jl")

__init__() = LibMagick.init()

function precompile()
    for T in (Uint8, Uint16, Int, Float32, Float64)
        Tdiv = typeof(one(T)/2)
        for N = 2:3
            precompile(restrict!, (Array{Tdiv, N}, Array{T,N}, Int))
            precompile(imfilter, (Array{T,N}, Array{Float64,N}))
            precompile(imfilter, (Array{T,N}, Array{Float64,2}))
            precompile(imfilter, (Array{T,N}, Array{Float32,N}))
            precompile(imfilter, (Array{T,N}, Array{Float32,2}))
        end
    end
    for T in (Float32, Float64)
        for N = 2:3
            precompile(_imfilter_gaussian!, (Array{T,N}, Vector{Float64}))
            precompile(_imfilter_gaussian!, (Array{T,N}, Vector{Int}))
            precompile(imfilter_gaussian_no_nans!, (Array{T,N}, Vector{Float64}))
            precompile(imfilter_gaussian_no_nans!, (Array{T,N}, Vector{Int}))
            precompile(fft, (Array{T,N},))
        end
    end
    for T in (Complex{Float32}, Complex{Float64})
        for N = 2:3
            precompile(ifft, (Array{T,N},))
        end
    end
end

export # types
    Gray,
    GrayAlpha,
    ARGB,
    BGRA,
    RGBA,
    RGB1,
    RGB4,
    YIQ,

    AbstractImage,
    AbstractImageDirect,
    AbstractImageIndexed,
    Image,
    ImageCmap,
    BitShift,
    ClampMin,
    ClampMax,
    ClampMinMax,
    Clamp,
    LabeledArray,
    MapInfo,
    MapNone,
    Overlay,
    OverlayImage,
    ScaleAutoMinMax,
    ScaleMinMax,
    ScaleSigned,
    SliceData,

    # constants
    palette_fire,
    palette_gray32,
    palette_gray64,
    palette_rainbow,

    # core functions
    assert2d,
    assert_scalar_color,
    assert_timedim_last,
    assert_xfirst,
    assert_yfirst,
    colordim,
    colorspace,
    coords,
    coords_spatial,
    copyproperties,
    data,
    dimindex,
    dimindexes,
    getindexim,
    grayim,
    colorim,
    height,
    isdirect,
    isxfirst,
    isyfirst,
    limits,
    maxabsfinite,
    maxfinite,
    minfinite,
    nchannels,
    ncolorelem,
    nimages,
    pixelspacing,
    properties,
    rerange!,
    reslice!,
    restrict,
    sdims,
    size_spatial,
    shareproperties,
    sliceim,
    spacedirections,
    spatialorder,
    spatialpermutation,
    spatialproperties,
    storageorder,
    subim,
    timedim,
    width,
    widthheight,
    raw,

    # io functions
    add_image_file_format,
    imread,
    imwrite,
    loadformat,

    # iterator functions
    first_index,
    iterate_spatial,
    parent,

    # color-related functions
    alphaval,
    blueval,
    greenval,
    redval,
    red,
    green,
    blue,
    hsi2rgb,
    imadjustintensity,
    indexedcolor,
    lut,
    rgb2hsi,
    rgb2ycbcr,
    separate,
    uint32color,
    uint32color!,
    ycbcr2rgb,

    # Scaling of intensity
    sc,
    scale,
    mapinfo,
    uint8sc,
    uint16sc,
    uint32sc,

    # algorithms
    ando3,
    ando4,
    ando5,
    backdiffx,
    backdiffy,
    dilate,
    erode,
    opening,
    closing,
    forwarddiffx,
    forwarddiffy,
    gaussian2d,
    imaverage,
    imcomplement,
    imdog,
    imedge,
    imfilter,
    imfilter!,
    imfilter_fft,
    imfilter_gaussian,
    imfilter_gaussian!,
    imfilter_LoG,
    imgaussiannoise,
    imgradients,
    imlaplacian,
    imlineardiffusion,
    imlog,
    imROF,
    imstretch,
#     imthresh,
    label_components,
    label_components!,
    magnitude,
    magnitude_phase,
    meanfinite,
    ncc,
    orientation,
    padarray,
    phase,
    prewitt,
    sad,
    sadn,
    sobel,
    ssd,
    ssdn,
    thin_edges,
    thin_edges_subpix,
    thin_edges_nonmaxsup,
    thin_edges_nonmaxsup_subpix,

    # phantoms
    shepp_logan

export # Deprecated exports
    ClipMin,
    ClipMax,
    ClipMinMax,
    ScaleInfo,
    climdefault,
    float32sc,
    float64sc,
    ntsc2rgb,
    rgb2gray,
    rgb2ntsc,
    scaleinfo,
    scaleminmax,
    scalesigned

import Base: scale, scale!  # delete when deprecations are removed
@deprecate scaleminmax  ScaleMinMax
@deprecate scaleminmax(img::AbstractArray, min::Real, max::Real)  ScaleMinMax(RGB24, img, min, max)
@deprecate float32sc    float32
@deprecate float64sc    float64
@deprecate uint8sc      ufixed8sc
@deprecate uint16sc(img)  ufixedsc(Ufixed16, img)
@deprecate ClipMin      ClampMin
@deprecate ClipMax      ClampMax
@deprecate ClipMinMax   ClampMinMax
@deprecate climdefault(img) zero(eltype(img)), one(eltype(img))
@deprecate ScaleMinMax{T<:Real}(img::AbstractArray{T}, mn, mx) ScaleMinMax(Ufixed8, img, mn, mx)
@deprecate ScaleMinMax{T<:ColorValue}(img::AbstractArray{T}, mn, mx) ScaleMinMax(RGB{Ufixed8}, img, mn, mx)
@deprecate ntsc2rgb(A::AbstractArray)  convert(Array{RGB}, A)
@deprecate ntsc2rgb(A::AbstractImage)  convert(Image{RGB}, A)
@deprecate rgb2ntsc(A::AbstractArray)  convert(Array{YIQ}, A)
@deprecate rgb2ntsc(A::AbstractImage)  convert(Image{YIQ}, A)
@deprecate scaleinfo    mapinfo
@deprecate scale(mapi::MapInfo, A) map(mapi, A)                # delete imports above when eliminated
@deprecate scale!(dest, mapi::MapInfo, A) map!(mapi, dest, A)  #   "
@deprecate rgb2gray(img::AbstractArray)  convert(Array{Gray}, img)
@deprecate copy(A::AbstractArray, B::AbstractArray) copyproperties(A, B)
@deprecate share(A::AbstractArray, B::AbstractArray) shareproperties(A, B)

const ScaleInfo = MapInfo  # can't deprecate types?

end


using RRIFT
using Test

@testset "in-vivo" begin
    test_study = RRIFT.gbm_study_uids[1]
    # Run download functions twice to check for overwrite protection
    dicom_folders = download_invivo_studies(test_study, destination = "./data/tcga-gbm")
    download_invivo_studies(test_study, destination = "./data/tcga-gbm")
    mask_folder = download_invivo_masks(destination = "./data/tcga-gbm-masks")
    download_invivo_masks(destination = "./data/tcga-gbm-masks")

    mat_file_test = preprocess_dicom_to_mat(destination = "./data/tcga-gbm-test", dicom_folders = dicom_folders, mask_folder = mask_folder)

    mat_folder = "./data/tcga-gbm-mat"
    download_invivo_preprocessed(destination = mat_folder)

    test_fits_a = process_patients(mat_file_test)
    test_fits_b = process_patients(joinpath(mat_folder, test_study * ".mat"))

    thresh = 1e-9
    println("Printing Δ between fits using local file versus downloaded file:")
    for param in (:kt, :kep, :ve, :vp, :kt_rr, :kep_rr, :ve_rr), method in (:tofts, :rrift)
        Δ = maximum(abs.(test_fits_a[param][method][1] .- test_fits_b[param][method][1]))
        println(string(param) * " " * string(method) * ": " * string(Δ))
        @test thresh > Δ
    end

    fits = process_patients(mat_folder)

    ccc_kt = ccc(fits[:kt], lim=(0, 0.2))
    ccc_ve = ccc(fits[:ve], lim=(0, 0.5))
    ccc_vp = ccc(fits[:vp], lim=(0, 0.05))

    ccc_ktrr = ccc(fits[:kt_rr])
    ccc_verr = ccc(fits[:ve_rr])
    ccc_keprr = ccc(fits[:kep_rr])

    println("======")
    println("Printing CCC:")
    println("CCC [kt, ve vp]:")
    println([ccc_kt[:rrift][:tofts], ccc_ve[:rrift][:tofts], ccc_vp[:rrift][:tofts]])
    println("CCC [kt_rr, ve_rr, kep_rr]:")
    println([ccc_ktrr[:rrift][:tofts], ccc_verr[:rrift][:tofts], ccc_keprr[:rrift][:tofts]])
end


using Revise, Parameters, Logging, Test
using CUDA
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils
using ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils
using ..MainLoopKernel

##### clear before next dilatation
threads=(32,5);
blocks =1;
mainArrDims= (67,177,90);
dataBdim = (43,21,17);
metaData = MetaDataUtils.allocateMetadata(mainArrDims,dataBdim);
metaDataDims= size(metaData);
resShmem = CUDA.ones(Bool,(dataBdim[1]+2,dataBdim[2]+2,dataBdim[3]+2 ));
sourceShmem = CUDA.ones(Bool,(dataBdim));
resShmemTotalLength = length(resShmem)
sourceShmemTotalLength= length(sourceShmem)
clearIterResShmemLoop= fld(resShmemTotalLength,threads[1]*threads[2])
clearIterSourceShmemLoop= fld(sourceShmemTotalLength,threads[1]*threads[2])
oldBlockCounter= CUDA.zeros(1)
function clearKernel(oldBlockCounter,resShmem,sourceShmem, clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength, sourceShmemTotalLength)
    locArr=0
    offsetIter=0
    # localOffset=0

    MainLoopKernel.@clearBeforeNextDilatation( clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength, sourceShmemTotalLength)
 return
end

@cuda threads=threads blocks=blocks clearKernel(oldBlockCounter,resShmem,sourceShmem, clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength, sourceShmemTotalLength)
@test sum(resShmem)==0
@test sum(sourceShmem)==0

###########iterateOverWorkQueue
using Revise, Parameters, Logging, Test
using CUDA
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils
using ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils
using ..MainLoopKernel

numberInWorkQueue = 2777
fpTotal=2000
fnTotal=1000
workQueaue = CUDA.zeros(Int64,4,4000)#allocateWorkQueue(fpTotal,fnTotal)
threads=(32,5);
blocks =17;
workQueauecounter= CuArray([numberInWorkQueue])
goldToBeDilatated= CUDA.ones(Bool,1)
segmToBeDilatated= CUDA.ones(Bool,1)
shmemSumLengthMaxDiv4= fld(36*14,4)*4
indexesArr = CUDA.zeros(Int64,1000000)
singleVal = CUDA.zeros(UInt32,1)

for i in 1:length(workQueaue)
    workQueaue[i]= i#using linear index
end
for i in 1 :(fpTotal+fnTotal)
    for j in 1:3
        workQueaue[j,i]= i*j
    end
    workQueaue[4,i]=1 # always gold pass
end
# 70*4
# for i in 1:length(workQueaue)
#     workQueaue[i]= mod(i,255)#using linear index
# end
# for i in 1 :(fpTotal+fnTotal)
#     for j in 1:3
#         workQueaue[j,i]= mod(i*j,255)
#     end
#     workQueaue[4,i]=1 # always gold pass
# end

# for i in 1:length(workQueaue)
#     workQueaue[i]=i#using linear index
# end
# for i in 1 :((fpTotal+fnTotal)*4)
#     workQueaue[i]=i
#     # for j in 1:3
#     #     workQueaue[j,i]=i*j
#     # end
#     # workQueaue[4,i]=1 # always gold pass
# end
Int64.(Array(workQueaue[:,1:8]))

# workQueaue[*4+4]
function iterateOverWorkQueueKernel(indexesArr,workQueauecounter,workQueaue,goldToBeDilatated, segmToBeDilatated,shmemSumLengthMaxDiv4,singleVal)
    shmemSum =  @cuStaticSharedMem(UInt32,(36,14)) # we need this additional spots
    MainLoopKernel.@iterateOverWorkQueue(workQueauecounter,workQueaue,goldToBeDilatated, segmToBeDilatated,shmemSumLengthMaxDiv4,begin

    @ifXY 1 1 CUDA.@atomic singleVal[1]+=1
    end)

 return
end

@cuda threads=threads blocks=blocks iterateOverWorkQueueKernel(indexesArr,workQueauecounter,workQueaue,goldToBeDilatated, segmToBeDilatated,shmemSumLengthMaxDiv4,singleVal)
@test Int64(singleVal[1])==numberInWorkQueue

aa = (cld(33,2))
bb = shmemSumLengthMaxDiv4

fld((cld(33,2)),shmemSumLengthMaxDiv4)

(numberInWorkQueue*4)-sum(indexesArr)

coords = CartesianIndices(indexesArr[1:numberInWorkQueue*4])
aaa = sort(filter(coo-> indexesArr[coo]==0, coords)) #3128


export MPIElasticPropagatorParams, MPIElasticSource, MPIElasticReceiver, MPIElasticPropagator, compute_PML_Params!,
        MPIElasticPropagatorSolver, MPISimulatedObservation!, MPIElasticReceiver, extract_local_patch
@with_kw mutable struct MPIElasticPropagatorParams
    # number of grids along x,y axis and time steps
    NX::Int64 = 0
    NY::Int64 = 0
    M::Int64 = 0
    N::Int64 = 0
    n::Int64 = 300
    NSTEP::Int64 = 2000 * 2

    # size of grid cell and time step
    DELTAX::Float64 = 10.
    DELTAY::Float64 = 10.
    DELTAT::Float64 = 2.e-3 / 2

    # PML boundary conditon
    f0::Float64 = 5. # source frequency
    vp_ref::Float64 = 2000.
    USE_PML_XMIN::Bool = true
    USE_PML_XMAX::Bool = true
    USE_PML_YMIN::Bool = true
    USE_PML_YMAX::Bool = true
    NPOINTS_PML::Int64 = 12
    NPOWER::Float64 = 2.
    K_MAX_PML::Float64 = 1.
    ALPHA_MAX_PML::Float64 = 2. * π * ( f0/2.)
    Rcoef::Float64 = 0.001 # Relative reflection coefficient

    # Auxilliary Data
    ax::Union{Missing, Array{Float64, 2}} = missing
    bx::Union{Missing, Array{Float64, 2}} = missing
    kx::Union{Missing, Array{Float64, 2}} = missing
    ay::Union{Missing, Array{Float64, 2}} = missing
    by::Union{Missing, Array{Float64, 2}} = missing
    ky::Union{Missing, Array{Float64, 2}} = missing

    # display params
    IT_DISPLAY::Int64 = 0

    # mpi
    II::Int64 = -1
    JJ::Int64 = -1

    # indices
    i_1j::Union{Missing, Array{Int64, 1}} = missing
    i_2j::Union{Missing, Array{Int64, 1}} = missing
    i1j::Union{Missing, Array{Int64, 1}} = missing
    i2j::Union{Missing, Array{Int64, 1}} = missing
    ij_1::Union{Missing, Array{Int64, 1}} = missing
    ij_2::Union{Missing, Array{Int64, 1}} = missing
    ij1::Union{Missing, Array{Int64, 1}} = missing
    ij2::Union{Missing, Array{Int64, 1}} = missing
    ij::Union{Missing, Array{Int64, 1}} = missing
    i1j1::Union{Missing, Array{Int64, 1}} = missing

end

"""
    MPIElasticSource

The input is the global index.

`srci`, `srcj`: nsrc
`srcv`: nt × nsrc, each row corresponds to a source time function
"""
mutable struct MPIElasticSource
    srci::Union{Missing, PyObject}
    srcj::Union{Missing, PyObject}
    srctype::Union{Missing,PyObject}
    srcv::Union{Missing, PyObject}
    function MPIElasticSource(param::MPIElasticPropagatorParams, srci::Array{Int64,1},
         srcj::Array{Int64,1}, srctype::Array{Int64, 1}, srcv::Union{Array{Float64,2}, PyObject})
        II, JJ = param.II, param.JJ
        nsrc = length(srci)
        srcv = constant(srcv)
        srcv = srcv'
        local_srci = Int64[]
        local_srcj = Int64[]
        local_srcv = PyObject[]
        local_srctype = Int64[]
        for i = 1:nsrc
            if ((II - 1) * param.n + 1 <= srci[i] <= II * param.n) &&
                    ((JJ - 1) * param.n + 1 <= srcj[i] <= JJ * param.n)
                    push!(local_srci, srci[i] - (II - 1) * param.n + 2)
                    push!(local_srcj, srcj[i] - (JJ - 1) * param.n + 2)
                    push!(local_srcv, srcv[i])
                    push!(local_srctype, srctype[i])
            end
        end
        local_srcv = hcat(local_srcv...)
        if length(local_srci)>0
            local_srci, local_srcj, local_srctype, local_srcv = convert_to_tensor([local_srci, local_srcj, local_srctype, local_srcv], [Int64, Int64, Int64, Float64])
            new(local_srci, local_srcj, local_srctype, local_srcv)
        else
            new(missing, missing, missing, missing)
        end
    end
end

mutable struct MPIElasticReceiver
    rcvi::Union{Missing,Array{Int64, 1}}
    rcvj::Union{Missing,Array{Int64, 1}}
    rcvtype::Union{Missing, Array{Int64, 1}}
    rcvv::Union{Missing, PyObject}
    function MPIElasticReceiver(param::MPIElasticPropagatorParams, rcvi::Array{Int64,1},
         rcvj::Array{Int64,1}, rcvtype::Array{Int64, 1})
        II, JJ = param.II, param.JJ
        nrcv = length(rcvi)
        local_rcvi = Int64[]
        local_rcvj = Int64[]
        local_rcvtype = Int64[]
        for i = 1:nrcv
            if ((II - 1) * param.n + 1 <= rcvi[i] <= II * param.n) &&
                    ((JJ - 1) * param.n + 1 <= rcvj[i] <= JJ * param.n)
                    push!(local_rcvi, rcvi[i] - (II - 1) * param.n)
                    push!(local_rcvj, rcvj[i] - (JJ - 1) * param.n)
                    push!(local_rcvtype, rcvtype[i])
            end
        end
        if length(local_rcvi)>0
            new(local_rcvi, local_rcvj, local_rcvtype, missing)
        else
            new(missing, missing, missing, missing)
        end
    end
end

"""
    ElasticPropagator

A structure holding all the information for acoustic simulation
"""
mutable struct MPIElasticPropagator
    param::MPIElasticPropagatorParams
    src::MPIElasticSource
    vx::PyObject
    vy::PyObject
    sigmaxx::PyObject
    sigmayy::PyObject
    sigmaxy::PyObject
    mem::PyObject
end

function get_mpi_id2(a, b, n)
    idx = Int64[]
    for i = 1:length(a)
        for j = 1:length(b)
            push!(idx, (n+4)*(a[i]-1)+b[j])
        end
    end
    idx
end

function extract_local_patch(param::MPIElasticPropagatorParams, v::Union{PyObject, Array{Float64,2}}; deps = missing, tag = 0)
    v = v[(param.II - 1) * param.n + 1:param.II * param.n, (param.JJ - 1) * param.n + 1:param.JJ * param.n]
    v = mpi_halo_exchange2(v, param.M, param.N, deps = deps, tag = tag)
end

function compute_PML_Params!(param::MPIElasticPropagatorParams)
    NX, NY = param.NX, param.NY
    n = param.n
    r = mpi_rank()
    param.M = div(NX,n)
    param.N = div(NY,n)
    param.II, param.JJ = div(r, param.N)+1, mod(r, param.N)+1
    if mpi_size() != param.M * param.N
        error("MPI Size Mismatch: mpi_size = $(mpi_size()), M = $(param.M), N = $(param.N)")
    end
    @assert NX>0 && NY>0
    @assert mod(NX, n)==0 && mod(NY, n)==0

    # ID
    k1 = 4:n+3
    k2 = 5:n+4
    k_1 = 2:n+1
    k_2 = 1:n
    kk = 3:n+2
    param.i_1j = get_mpi_id2(k_1, kk, n)
    param.i_2j = get_mpi_id2(k_2, kk, n)
    param.i1j = get_mpi_id2(k1, kk, n)
    param.i2j = get_mpi_id2(k2, kk, n)
    param.ij_1 = get_mpi_id2(kk, k_1, n)
    param.ij_2 = get_mpi_id2(kk, k_2, n)
    param.ij1 = get_mpi_id2(kk, k1, n)
    param.ij2 = get_mpi_id2(kk, k2, n)
    param.ij = get_mpi_id2(kk, kk, n)
    param.i1j1 = get_mpi_id2(k1, k1, n)

    ## define profile of absorption in PML region
    NPOINTS_PML = param.NPOINTS_PML
    NPOWER = param.NPOWER
    NX = param.NX
    NY = param.NY
    DELTAX = param.DELTAX
    DELTAY = param.DELTAY
    DELTAT = param.DELTAT
    K_MAX_PML = param.K_MAX_PML
    ALPHA_MAX_PML = param.ALPHA_MAX_PML
    USE_PML_XMIN = param.USE_PML_XMIN
    USE_PML_XMAX = param.USE_PML_XMAX
    USE_PML_YMIN = param.USE_PML_YMIN
    USE_PML_YMAX = param.USE_PML_YMAX
    Rcoef = param.Rcoef

    thickness_PML_x = NPOINTS_PML * DELTAX
    thickness_PML_y = NPOINTS_PML * DELTAY

    # reflection coefficient
    d0_x = - (NPOWER + 1) * param.vp_ref * log(Rcoef) / (2. * thickness_PML_x)
    d0_y = - (NPOWER + 1) * param.vp_ref * log(Rcoef) / (2. * thickness_PML_y)

    xoriginleft = thickness_PML_x
    xoriginright = (NX-0.5)*DELTAX - thickness_PML_x

    a_x = zeros(n)
    a_y = zeros(n)
    a_x_half = zeros(n)
    a_y_half = zeros(n)
    b_x = zeros(n)
    b_y = zeros(n)
    b_x_half = zeros(n)
    b_y_half = zeros(n)
    d_x = zeros(n)
    d_y = zeros(n)
    d_x_half = zeros(n)
    d_y_half = zeros(n)
    K_x = ones(n)
    K_y = ones(n)
    K_x_half = ones(n)
    K_y_half = ones(n)
    alpha_x = zeros(n)
    alpha_y = zeros(n)
    alpha_x_half = zeros(n)
    alpha_y_half = zeros(n)

    for i = 1:n
        xval = DELTAX * (i-1) + (param.II - 1) * param.n * DELTAX
        if (USE_PML_XMIN)

            abscissa_in_PML = xoriginleft - xval
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_x
                d_x[i] = d0_x * abscissa_normalized^NPOWER
                K_x[i] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_x[i] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_x[i] = ALPHA_MAX_PML * abscissa_normalized
            end

            abscissa_in_PML = xoriginleft - (xval + DELTAX/2.)
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_x
                d_x_half[i] = d0_x * abscissa_normalized^NPOWER
                K_x_half[i] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_x_half[i] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_x_half[i] = ALPHA_MAX_PML * abscissa_normalized
            end

        end

        if (USE_PML_XMAX)
            abscissa_in_PML = xval - xoriginright
            if (abscissa_in_PML >= 0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_x
                d_x[i] = d0_x * abscissa_normalized^NPOWER
                K_x[i] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_x[i] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_x[i] = ALPHA_MAX_PML * abscissa_normalized
            end

            abscissa_in_PML = xval + DELTAX/2.0 - xoriginright
            if (abscissa_in_PML >= 0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_x
                d_x_half[i] = d0_x * abscissa_normalized^NPOWER
                K_x_half[i] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_x_half[i] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_x_half[i] = ALPHA_MAX_PML * abscissa_normalized
            end
        end

        if (alpha_x[i] < 0)
            alpha_x[i] = 0
        end
        if (alpha_x_half[i] < 0)
            alpha_x_half[i] = 0
        end

        b_x[i] = exp(- (d_x[i] / K_x[i] + alpha_x[i]) * DELTAT)
        b_x_half[i] = exp(- (d_x_half[i] / K_x_half[i] + alpha_x_half[i]) * DELTAT)

        if (abs(d_x[i]) > 1e-6)
            a_x[i] = d_x[i] * (b_x[i] - 1.) / (K_x[i] * (d_x[i] + K_x[i] * alpha_x[i]))
        end
        if (abs(d_x_half[i]) > 1e-6)
            a_x_half[i] = d_x_half[i] * (b_x_half[i] - 1.) / (K_x_half[i] * (d_x_half[i] + K_x_half[i] * alpha_x_half[i]))
        end
    end

    yoriginbottom = thickness_PML_y
    yorigintop = (NY-0.5)*DELTAY - thickness_PML_y

    for j = 1:n

        yval = DELTAY * (j-1) + (param.JJ - 1) * param.n * DELTAX
        if param.USE_PML_YMIN
            abscissa_in_PML = yoriginbottom - yval
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_y
                d_y[j] = d0_y * abscissa_normalized^NPOWER
                K_y[j] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_y[j] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_y[j] = ALPHA_MAX_PML * abscissa_normalized
            end
            abscissa_in_PML = yoriginbottom - (yval + DELTAY/2.)
            if abscissa_in_PML >= 0
                abscissa_normalized = abscissa_in_PML / thickness_PML_y
                d_y_half[j] = d0_y * abscissa_normalized^NPOWER
                K_y_half[j] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_y_half[j] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_y_half[j] = ALPHA_MAX_PML * abscissa_normalized
            end
        end

        if param.USE_PML_YMAX
            abscissa_in_PML = yval - yorigintop
            if abscissa_in_PML >= 0
                abscissa_normalized = abscissa_in_PML / thickness_PML_y
                d_y[j] = d0_y * abscissa_normalized^NPOWER
                K_y[j] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_y[j] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_y[j] = ALPHA_MAX_PML * abscissa_normalized
            end

            abscissa_in_PML = yval + DELTAY/2. - yorigintop
            if abscissa_in_PML >= 0
                abscissa_normalized = abscissa_in_PML / thickness_PML_y
                d_y_half[j] = d0_y * abscissa_normalized^NPOWER
                K_y_half[j] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_y_half[j] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_y_half[j] = ALPHA_MAX_PML * abscissa_normalized
            end
        end

        b_y[j] = exp(- (d_y[j] / K_y[j] + alpha_y[j]) * DELTAT)
        b_y_half[j] = exp(- (d_y_half[j] / K_y_half[j] + alpha_y_half[j]) * DELTAT)
        if abs(d_y[j]) > 1e-6
            a_y[j] = d_y[j] * (b_y[j] - 1.) / (K_y[j] * (d_y[j] + K_y[j] * alpha_y[j]))
        end
        if (abs(d_y_half[j]) > 1e-6)
            a_y_half[j] = d_y_half[j] * (b_y_half[j] - 1.) / (K_y_half[j] * (d_y_half[j] + K_y_half[j] * alpha_y_half[j]))
        end
    end
    # @show norm(d_x-d_y)

    ax = [a_x'; a_x_half']
    bx = [b_x'; b_x_half']
    kx = [K_x'; K_x_half']
    ay = [a_y'; a_y_half']
    by = [b_y'; b_y_half']
    ky = [K_y'; K_y_half']

    param.ax = ax
    param.bx = bx
    param.kx = kx
    param.ay = ay
    param.by = by
    param.ky = ky
    param
end

function MPIElasticPropagatorSolver(param::MPIElasticPropagatorParams, src::MPIElasticSource,
    ρ::Union{PyObject, Array{Float64, 2}}, λ::Union{PyObject, Array{Float64, 2}}, μ::Union{PyObject, Array{Float64, 2}};
    tag_offset::Int64 = 0, dep = nothing)
    ρ = reshape(convert_to_tensor(ρ),(-1,))
    λ = reshape(convert_to_tensor(λ),(-1,))
    μ = reshape(convert_to_tensor(μ),(-1,))
    n = param.n
    @assert size(ρ)==size(λ)==size(μ)==((n+4)^2,)
    @assert param.II>0
    if !ismissing(src.srcv)
        @assert size(src.srcv, 2)==length(src.srci)==length(src.srcj)==length(src.srctype)
        @assert size(src.srcv, 1)==param.NSTEP
    end
    NSTEP = param.NSTEP
    vx_arr = TensorArray(NSTEP+1);
    vx_arr = write(vx_arr, 1, constant(zeros(n^2)))
    vy_arr = TensorArray(NSTEP+1);
    vy_arr = write(vy_arr, 1, constant(zeros(n^2)))
    sigmaxx_arr = TensorArray(NSTEP+1);
    sigmaxx_arr = write(sigmaxx_arr, 1, constant(zeros(n^2)))
    sigmayy_arr = TensorArray(NSTEP+1);
    sigmayy_arr = write(sigmayy_arr, 1, constant(zeros(n^2)))
    sigmaxy_arr = TensorArray(NSTEP+1);
    sigmaxy_arr = write(sigmaxy_arr, 1, constant(zeros(n^2)))
    mem_arr = [TensorArray(NSTEP+1) for i = 1:8]
    for i = 1:8
        mem_arr[i] = write(mem_arr[i], 1, constant(zeros(n^2)))
    end
    ax,bx,kx,ay,by,ky = param.ax,param.bx,param.kx,param.ay,param.by,param.ky
    function condition(i, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arr...)
        i<=NSTEP+1
    end

    if isnothing(dep)
        dep = 0.0
    end

    function body(i, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arr...)
        mem = Array{PyObject}(undef, 8)
        σxx,σyy,σxy,vx,vy =
            read(sigmaxx_arr, i-1), read(sigmayy_arr, i-1), read(sigmaxy_arr, i-1),
            read(vx_arr, i-1), read(vy_arr, i-1)
        for k = 1:8
            m = read(mem_arr[k], i-1)
            mem[k] = _reshape_and_halo_exchange(m, param, 18*i+k-1+tag_offset, k>1 ? mem[k-1][1] + dep : λ[1] + ρ[1] + μ[1] + vy[1] + σxy[1] + dep)
        end

        mpi_idx = 18*i+8+tag_offset

        halo = (op, deps, mpi_idx)->begin
            op = _reshape_and_halo_exchange(op, param, mpi_idx, deps)
            mpi_idx = mpi_idx + 1
            op, mpi_idx
        end
        vx, mpi_idx = halo(vx, mem[8][1], mpi_idx)
        vy, mpi_idx = halo(vy, vx[1], mpi_idx)
        σxx, mpi_idx = halo(σxx, vy[1], mpi_idx)
        σyy, mpi_idx = halo(σyy, σxx[1], mpi_idx)
        σxy, mpi_idx = halo(σxy, σyy[1], mpi_idx)

        srci,srcj,srctype,srcv = ElasticSourceAtTimeT(src, i-1)
        σxx, σyy, σxy,vx, vy, mem, mpi_idx = one_step(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
            ax, bx, ay, by, kx, ky, srci, srcj, srcv, srctype, mpi_idx)

        sigmaxx_arr,sigmayy_arr,sigmaxy_arr,vx_arr,vy_arr =
            write(sigmaxx_arr, i, σxx), write(sigmayy_arr, i, σyy), write(sigmaxy_arr, i, σxy),
                    write(vx_arr, i, vx), write(vy_arr, i, vy)

        mem_arr_ = Array{PyObject}(undef, 8)
        for k = 1:8
            mem_arr_[k] = write(mem_arr[k], i, mem[k])
        end

        if param.IT_DISPLAY>0
            op = tf.cond(tf.equal(tf.math.floormod(i,param.IT_DISPLAY),1), ()->tf.print("time step = ", i-1, " out of $NSTEP"), ()->tf.no_op())
            i = bind(i, op)
        end

        i+1, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arr_...
    end
    i = constant(2, dtype=Int32)
    _, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arrs =
        while_loop(condition, body, [i, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arr...])

    vx = reshape(stack(vx_arr), (param.NSTEP+1, n, n))
    vy = reshape(stack(vy_arr), (param.NSTEP+1, n, n))
    σxx = reshape(stack(sigmaxx_arr), (param.NSTEP+1, n, n))
    σyy = reshape(stack(sigmayy_arr), (param.NSTEP+1, n, n))
    σxy = reshape(stack(sigmaxy_arr), (param.NSTEP+1, n, n))
    MPIElasticPropagator(
        param, src, vx, vy, σxx, σyy, σxy, mem_arrs
    )
end

function _reshape_and_print(o, m, n, info)
    o = reshape(o, (m, n))
    o = print_tensor(o, info)
    reshape(o, (m*n,))
end

function _reshape_and_halo_exchange(o, param, tag, deps)
    @assert size(o)==(param.n^2,)
    o = reshape(o, (param.n, param.n))
    out = mpi_halo_exchange2(o, param.M, param.N, tag = tag, deps= deps)
    reshape(out, ((param.n+4)^2,))
end

function fw1(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
        ax, bx, ay, by, kx, ky, mpi_idx)
    NX, NY = param.NX, param.NY
    n = param.n
    DELTAX, DELTAY, DELTAT = param.DELTAX, param.DELTAY, param.DELTAT
    i_1j = param.i_1j
    i_2j = param.i_2j
    i1j = param.i1j
    i2j = param.i2j
    ij_1 = param.ij_1
    ij_2 = param.ij_2
    ij1 = param.ij1
    ij2 = param.ij2
    ij = param.ij

    # λ_ = 0.5*(λ[i1j]+λ[ij])
    # μ_ = 0.5*(μ[i1j]+μ[ij])
    λ_ = λ[ij]
    μ_ = μ[ij]
    λμ = λ_ + 2μ_
    value_dvx_dx = (27vx[i1j]-27vx[ij]-vx[i2j]+vx[i_1j])/(24*DELTAX)
    value_dvy_dy = (27*vy[ij]-27*vy[ij_1]-vy[ij1]+vy[ij_2]) / (24*DELTAY)
    # broadcast
    mem[1] =  makevector(mem[1], ij, adbroadcast(bx[2,:], mem[1][ij], 1) + adbroadcast(ax[2,:], value_dvx_dx, 1));
    mem[2] =  makevector(mem[2], ij, adbroadcast(by[1,:], mem[2][ij], 2) + adbroadcast(ay[1,:], value_dvy_dy, 2));
    value_dvx_dx = adbroadcast(value_dvx_dx, kx[2,:], 3) + mem[1][ij]
    value_dvy_dy = adbroadcast(value_dvy_dy, ky[1,:], 4) + mem[2][ij]
    σxx = scatter_add_op(σxx, ij, (λμ * value_dvx_dx + λ_ * value_dvy_dy) * DELTAT )
    σyy = scatter_add_op(σyy, ij, (λμ * value_dvy_dy + λ_ * value_dvx_dx) * DELTAT )

    σxx = _reshape_and_halo_exchange(σxx[ij], param, mpi_idx, σxy[1])
    σyy = _reshape_and_halo_exchange(σyy[ij], param, mpi_idx+1, σxx[1])

    return vx, vy, σxx, σyy,σxy, mem, mpi_idx+2
end

function fw2(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
    ax, bx, ay, by, kx, ky, mpi_idx)
    NX, NY = param.NX, param.NY
    n = param.n
    DELTAX, DELTAY, DELTAT = param.DELTAX, param.DELTAY, param.DELTAT
    i_1j = param.i_1j
    i_2j = param.i_2j
    i1j = param.i1j
    i2j = param.i2j
    ij_1 = param.ij_1
    ij_2 = param.ij_2
    ij1 = param.ij1
    ij2 = param.ij2
    ij = param.ij

    # μ_ = 0.5*(μ[ij]+μ[ij1])
    μ_ = μ[ij]
    value_dvy_dx = (27vy[ij]-27vy[i_1j]-vy[i1j]+vy[i_2j])/(24*DELTAX)
    value_dvx_dy = (27*vx[ij1]-27*vx[ij]-vx[ij2]+vx[ij_1]) / (24*DELTAY)

    # broadcast
    mem[3] = makevector(mem[3], ij,  adbroadcast(bx[1,:], mem[3][ij], 1) + adbroadcast(ax[1, :], value_dvy_dx, 1))
    mem[4] = makevector(mem[4], ij,  adbroadcast(by[2,:], mem[4][ij], 2) + adbroadcast(ay[2, :], value_dvx_dy, 2));

    value_dvy_dx = adbroadcast(value_dvy_dx, kx[1,:], 3) + mem[3][ij]
    value_dvx_dy = adbroadcast(value_dvx_dy, ky[1,:], 4) + mem[4][ij]

    σxy = scatter_add_op(σxy, ij, μ_*(value_dvy_dx + value_dvx_dy) * DELTAT )

    σxy = _reshape_and_halo_exchange(σxy[ij], param, mpi_idx, σyy[1])

    return vx, vy, σxx, σyy, σxy,mem, mpi_idx+1
end

function fw3(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
    ax, bx, ay, by, kx, ky, mpi_idx)
    NX, NY = param.NX, param.NY
    n = param.n
    DELTAX, DELTAY, DELTAT = param.DELTAX, param.DELTAY, param.DELTAT
    i_1j = param.i_1j
    i_2j = param.i_2j
    i1j = param.i1j
    i2j = param.i2j
    ij_1 = param.ij_1
    ij_2 = param.ij_2
    ij1 = param.ij1
    ij2 = param.ij2
    ij = param.ij

    value_dsigmaxx_dx = (27*σxx[ij]-27*σxx[i_1j]-σxx[i1j]+σxx[i_2j]) / (24*DELTAX);
    value_dsigmaxy_dy = (27*σxy[ij]-27*σxy[ij_1]-σxy[ij1]+σxy[ij_2]) / (24*DELTAY);

    mem[5] = makevector(mem[5],ij, adbroadcast(bx[1,:], mem[5][ij], 1) + adbroadcast(ax[1,:], value_dsigmaxx_dx, 1))
    mem[6] = makevector(mem[6],ij, adbroadcast(by[1,:], mem[6][ij], 2) + adbroadcast(ay[1,:], value_dsigmaxy_dy, 2))

    value_dsigmaxx_dx = adbroadcast(value_dsigmaxx_dx, kx[1,:], 3) + mem[5][ij]
    value_dsigmaxy_dy = adbroadcast(value_dsigmaxy_dy, ky[1,:], 4) + mem[6][ij]

    vx = scatter_add_op(vx, ij, (value_dsigmaxx_dx + value_dsigmaxy_dy) * DELTAT / ρ[ij])
    vx = _reshape_and_halo_exchange(vx[ij], param, mpi_idx, σxy[1])

    return vx, vy, σxx, σyy, σxy, mem, mpi_idx+1
end

function fw4(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
    ax, bx, ay, by, kx, ky, mpi_idx)
    NX, NY, n = param.NX, param.NY, param.n
    DELTAX, DELTAY, DELTAT = param.DELTAX, param.DELTAY, param.DELTAT
    i_1j = param.i_1j
    i_2j = param.i_2j
    i1j = param.i1j
    i2j = param.i2j
    ij_1 = param.ij_1
    ij_2 = param.ij_2
    ij1 = param.ij1
    ij2 = param.ij2
    ij = param.ij
    i1j1 = param.i1j1

    # ρ_ = 0.25 * (ρ[ij] + ρ[i1j] + ρ[i1j1] + ρ[ij1]);
    ρ_ = ρ[ij]

    value_dsigmaxy_dx = (27*σxy[i1j]-27*σxy[ij]-σxy[i2j]+σxy[i_1j]) / (24*DELTAX);
    value_dsigmayy_dy = (27*σyy[ij1]-27*σyy[ij]-σyy[ij2]+σyy[ij_1]) / (24*DELTAY);

    mem[7] = makevector(mem[7], ij, adbroadcast(bx[2,:], mem[7][ij], 1) + adbroadcast(ax[2,:] , value_dsigmaxy_dx, 1))
    mem[8] = makevector(mem[8], ij, adbroadcast(by[2,:], mem[8][ij], 2) + adbroadcast(ay[2,:] , value_dsigmayy_dy, 2))

    value_dsigmaxy_dx = adbroadcast(value_dsigmaxy_dx, kx[2,:], 3) + mem[7][ij]
    value_dsigmayy_dy = adbroadcast(value_dsigmayy_dy, ky[2,:], 4) + mem[8][ij]

    vy = scatter_add_op( vy, ij, (value_dsigmaxy_dx + value_dsigmayy_dy) * DELTAT / ρ_)
    vy = _reshape_and_halo_exchange(vy[ij], param, mpi_idx, vx[1])
    return vx, vy, σxx, σyy, σxy,mem, mpi_idx+1
end

function one_step(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
    ax, bx, ay, by, kx, ky, srci, srcj, srcv, srctype, mpi_idx)
    n = param.n
    @assert size(ρ)==size(λ)==size(μ)==size(vx)==size(vy)==size(σxx)==size(σyy)==size(σxy)==((n+4)^2,)

    add_source = load_op_and_grad("$(@__DIR__)/../deps/CustomOps/build/libADSeismic","add_source", multiple=true)
    vx, vy, σxx, σyy, σxy,mem, mpi_idx = fw1(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
                    ax, bx, ay, by, kx, ky, mpi_idx)
    vx, vy, σxx, σyy, σxy,mem, mpi_idx = fw2(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
                    ax, bx, ay, by, kx, ky, mpi_idx)
    vx, vy, σxx, σyy, σxy,mem, mpi_idx = fw3(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
                    ax, bx, ay, by, kx, ky, mpi_idx)
    vx, vy, σxx, σyy, σxy,mem, mpi_idx = fw4(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
                    ax, bx, ay, by, kx, ky, mpi_idx)
    if !ismissing(srci)
        σxx, σyy, σxy,vx, vy  =  add_source(σxx, σyy, σxy,vx, vy, constant(srci),constant(srcj),constant(srctype),constant(n+2),constant(n+2),srcv)
    end
    ij = param.ij
    σxx, σyy, σxy,vx, vy  = σxx[ij], σyy[ij], σxy[ij],vx[ij],vy[ij]
    for i = 1:8
        mem[i] = mem[i][ij]
    end
    return σxx, σyy, σxy,vx, vy, mem, mpi_idx
end

function ElasticSourceAtTimeT(src::MPIElasticSource,i::PyObject)
    if ismissing(src.srci)
        return missing, missing, missing, missing
    else
        return src.srci, src.srcj, src.srctype, src.srcv[i]
    end
end

"""
    SimulatedObservation!(ap::ElasticPropagator, rcv::ElasticReceiver)

Extract and save simulated displacement u into rcv::ElasticReceiver.
"""
function MPISimulatedObservation!(ep::MPIElasticPropagator, rcv::MPIElasticReceiver)
    if ismissing(rcv.rcvi)
        op = group([ep.vx, ep.vy, ep.sigmaxx, ep.sigmayy, ep.sigmaxy])
        rcv.rcvv = op
        return
    end

    N = (ep.param.n)*(ep.param.n)
    get_receive = load_op_and_grad("$(@__DIR__)/../deps/CustomOps/build/libADSeismic", "get_receive")
    u = get_receive(tf.reshape(ep.vx, (ep.param.NSTEP+1, N)),
                    tf.reshape(ep.vy, (ep.param.NSTEP+1, N)),
                    tf.reshape(ep.sigmaxx, (ep.param.NSTEP+1, N)),
                    tf.reshape(ep.sigmayy, (ep.param.NSTEP+1, N)),
                    tf.reshape(ep.sigmaxy, (ep.param.NSTEP+1, N)),
                    convert_to_tensor(rcv.rcvi), convert_to_tensor(rcv.rcvj), convert_to_tensor(rcv.rcvtype),
                        constant(ep.param.n-2), constant(ep.param.n-2))
    rcv.rcvv = u
end


using Gadgetron
using Test

@testset "ImageSerialization" begin

    meta = MRD.MetaDict()
    meta["penguin"] = 2.0
    meta["dinosaur"] = [2,3,4]
    img = MRD.Image(MRD.ImageHeader(version=4, acquisition_time_stamp=42), zeros(Float32, 2,3,4,1),meta)
    io = IOBuffer()
    MRD.write(io,img)
    seekstart(io)
    img3 = MRD.read(io,MRD.Image)

    @test img.header == img3.header
    @test img.data == img3.data
    @test img.meta == img3.meta
end

@testset "ImageSize" begin

    img = MRD.Image(MRD.ImageHeader(),zeros(Float32,1,1,1,1))
    header = MRD.RawImageHeader(img)

    io = IOBuffer()

    MRD.write(io,header)

    @test position(io) == 198
    meta = MRD.MetaDict()
    MRD.write(io,meta)
    @test position(io) == (198+8)

    io = IOBuffer()

end


import DiffEqBase
export DiffEqJLSolver, DiffEqJLIMEXSolver

abstract type AbstractDiffEqJLSolver <: AbstractODESolver end

"""
    DiffEqJLSolver(f, RKA, RKB, RKC, Q; dt, t0 = 0)

This is a time stepping object for explicitly time stepping the differential
equation given by the right-hand-side function `f` with the state `Q`, i.e.,

```math
  \\dot{Q} = f(Q, t)
```

via a DifferentialEquations.jl DEAlgorithm, which includes support
for OrdinaryDiffEq.jl, Sundials.jl, and more.
"""
mutable struct DiffEqJLSolver{I} <: AbstractDiffEqJLSolver
    integ::I
    steps::Int

    function DiffEqJLSolver(
        rhs!,
        alg,
        Q,
        args...;
        t0 = 0,
        p = nothing,
        kwargs...,
    )
        prob = DiffEqBase.ODEProblem(
            (du, u, p, t) -> rhs!(du, u, p, t; increment = false),
            Q,
            (float(t0), typemax(typeof(float(t0)))),
            p,
        )
        integ = DiffEqBase.init(
            prob,
            alg,
            args...;
            adaptive = false,
            save_everystep = false,
            save_start = false,
            save_end = false,
            kwargs...,
        )
        new{typeof(integ)}(integ, 0)
    end
end

"""
    DiffEqJLSolver(f, RKA, RKB, RKC, Q; dt, t0 = 0)

This is a time stepping object for explicitly time stepping the differential
equation given by the right-hand-side function `f` with the state `Q`, i.e.,

```math
  \\dot{Q} = f_I(Q, t) + f_E(Q, t)
```

via a DifferentialEquations.jl DEAlgorithm, which includes support
for OrdinaryDiffEq.jl, Sundials.jl, and more.
"""
mutable struct DiffEqJLIMEXSolver{I} <: AbstractDiffEqJLSolver
    integ::I
    steps::Int

    function DiffEqJLIMEXSolver(
        rhs!,
        rhs_implicit!,
        alg,
        Q,
        args...;
        t0 = 0,
        p = nothing,
        kwargs...,
    )
        prob = DiffEqBase.SplitODEProblem(
            (du, u, p, t) -> rhs_implicit!(du, u, p, t; increment = false),
            (du, u, p, t) -> rhs!(du, u, p, t; increment = false),
            Q,
            (float(t0), typemax(typeof(float(t0)))),
            p,
        )
        integ = DiffEqBase.init(
            prob,
            alg,
            args...;
            adaptive = false,
            save_everystep = false,
            save_start = false,
            save_end = false,
            kwargs...,
        )

        new{typeof(integ)}(integ, 0)
    end
end

gettime(solver::AbstractDiffEqJLSolver) = solver.integ.t
getdt(solver::AbstractDiffEqJLSolver) = solver.integ.dt
updatedt!(solver::AbstractDiffEqJLSolver, dt) =
    DiffEqBase.set_proposed_dt!(solver.integ, dt)
updatetime!(solver::AbstractDiffEqJLSolver, t) =
    DiffEqBase.set_t!(solver.integ, t)
isadjustable(solver::AbstractDiffEqJLSolver) = true # Is this isadaptive? Or something different?

"""
    ODESolvers.general_dostep!(Q, solver::AbstractODESolver, p,
                               timeend::Real, adjustfinalstep::Bool)

Use the solver to step `Q` forward in time from the current time, to the time
`timeend`. If `adjustfinalstep == true` then `dt` is adjusted so that the step
does not take the solution beyond the `timeend`.
"""
function general_dostep!(
    Q,
    solver::AbstractDiffEqJLSolver,
    p,
    timeend::Real;
    adjustfinalstep::Bool,
)
    integ = solver.integ

    if first(integ.opts.tstops) !== timeend
        DiffEqBase.add_tstop!(integ, timeend)
    end
    dostep!(Q, solver, p, time)
    solver.integ.t
end

function dostep!(
    Q,
    solver::AbstractDiffEqJLSolver,
    p,
    time,
    slow_δ = nothing,
    slow_rv_dQ = nothing,
    in_slow_scaling = nothing,
)

    integ = solver.integ
    integ.p = p # Can this change?

    rv_Q = realview(Q)
    if integ.u != Q
        integ.u .= Q
        DiffEqBase.u_modified!(integ, true)
        # Will time always be correct?
    end

    DiffEqBase.step!(integ)
    rv_Q .= solver.integ.u
end

function DiffEqJLConstructor(alg)
    constructor =
        (F, Q; dt = 0, t0 = 0) -> DiffEqJLSolver(F, alg, Q; t0 = t0, dt = dt)
    return constructor
end


# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.ContactDetectionMPR (Modia3D/contactDetection/ContactDetectionMPR/_module.jl)
#
using LinearAlgebra
EYE3() = Matrix(1.0I,3,3)

# returns true, if AABB's are touching
AABB_touching(aabb1::Basics.BoundingBox, aabb2::Basics.BoundingBox) = aabb1.x_max >= aabb2.x_min && aabb1.x_min <= aabb2.x_max &&
    aabb1.y_max >= aabb2.y_min && aabb1.y_min <= aabb2.y_max &&
    aabb1.z_max >= aabb2.z_min && aabb1.z_min <= aabb2.z_max

function Composition.initializeContactDetection!(world::Composition.Object3D, scene::Composition.Scene)::Nothing
    if typeof(scene.options.contactDetection) == Modia3D.ContactDetectionMPR_handler
        ch::Modia3D.ContactDetectionMPR_handler = scene.options.contactDetection
        ch.contactPairs = Composition.ContactPairs(world, scene.AABB, scene.superObjs, scene.allowedToMove,
            scene.options.visualizeBoundingBox,
            scene.options.nVisualContSupPoints, ch.visualizeContactPoints,
            ch.visualizeSupportPoints, ch.defaultContactSphereDiameter)
        if ch.contactPairs.nz == 0
            Composition.closeContactDetection!(ch)
            scene.collide = false
                @warn "... From Modia3D collision handler: Collision handling switched off, since no contacts can take place (nz=0).\n" *
                    "... You might need to set canCollide=true at joints.\n"
            return
        end
        @assert(ch.contactPairs.ne > 0)
        @assert(ch.contactPairs.nz > 0)
    end
end

function Composition.setComputationFlag(ch::Composition.ContactDetectionMPR_handler)
    ch.distanceComputed = false
end

# This function performs (a) a broad phase to determine which
# shapes are potentially in contact to each other, (b) computes
# the distances of these shapes in a narrow phase and (c) selects
# the nz shape pairs with the smallest distances and orders them
# according to their distances in O(n_n log(n_z)) operations.
# This function is called after initialization and after state events.
function Composition.selectContactPairsWithEvent!(sim, scene::Composition.Scene, ch::Composition.ContactDetectionMPR_handler, world::Composition.Object3D)
    empty!(ch.contactDict)
    ch.noContactMinVal =  42.0
    selectContactPairs!(sim,scene,ch,world,true)
    ch.contactPairs.nzContact = length(ch.contactDict)
    return nothing
end

# This function performs (a) a broad phase to determine which
# shapes are potentially in contact to each other, (b) computes
# the distances of these shapes in a narrow phase and (c) selects
# the nz shape pairs with the smallest distances and orders them
# according to their distances in O(n_n log(n_z)) operations.
# This function is called for contact detection during simulation.
# It also checks if the contact pairs (contact == true), from a previous call of selectContactPairsWithEvent!(...)
# are still true.
function Composition.selectContactPairsNoEvent!(sim, scene::Composition.Scene, ch::Composition.ContactDetectionMPR_handler, world::Composition.Object3D)
    ch.noContactMinVal =  42.0
    selectContactPairs!(sim,scene,ch,world,false)
    return nothing
end

function selectContactPairs!(sim, scene::Composition.Scene, ch::Composition.ContactDetectionMPR_handler, world::Composition.Object3D, hasEvent::Bool)
    if !ch.distanceComputed
        computeDistances(scene, ch, world, false, hasEvent)
    end

    # Store z in simulation engine
    if !isnothing(sim)
        simh = sim.eventHandler
        # z[1] ... zero crossing function from contact to no contact
        # max. distance (= min. penetration) of active contact pairs
        if isempty(ch.contactDict)
            @inbounds simh.z[1] = -42.0
        else
            (pair, key)  = findmax(ch.contactDict)
            @inbounds simh.z[1] = pair.distanceWithHysteresis
        end
        # z[2] ... zero crossing function from no contact to contact
        # min. distance of inactive contact pairs
        @inbounds simh.z[2] = ch.noContactMinVal
    end
    ch.distanceComputed = true
    return nothing
end

# This function performs (a) a broad phase to determine which
# shapes are potentially in contact to each other, (b) computes
# the distances of these shapes in a narrow phase and (c)
# stores the distances of the contact pairs selected by the last
# call of function selectContactPairsWithEvent!(C) in z in O(n log(nz))
# operations. This function is called whenever the integrator
# requests a new zero-crossing function evaluation
function Composition.getDistances!(scene::Composition.Scene, ch::Composition.ContactDetectionMPR_handler, world::Composition.Object3D)
    if !ch.distanceComputed
        computeDistances(scene, ch, world, true, false)
        ch.distanceComputed = true
        #println("Abstand wird berechnet, ", length(ch.contactDict) )
    end
    return nothing
end

function computeDistances(scene::Composition.Scene, ch::Composition.ContactDetectionMPR_handler, world::Composition.Object3D, phase2::Bool, hasEvent::Bool=true)
    noCPairs  = scene.noCPairs
    superObjs = scene.superObjs
    AABB = scene.AABB
    if length(superObjs) > 1
        k = 0
        @inbounds for i = 1:length(superObjs)
            superObj = superObjs[i].superObjCollision.superObj
            for j = 1:length(superObj)
                k = k +1
                obj = superObj[j]
                AABB[i][j] = Modia3D.boundingBox!(obj,
                    AABB[i][j]; tight=false, scaleFactor=0.01)
                if scene.options.visualizeBoundingBox
                    updateVisualBoundingBox!(world.AABBVisu[k], AABB[i][j])
                end
            end
        end

        # counter
        # is: actual super - object
        # js: subsequent super - object
        # i: Object3D of is_th super - object
        # j: Object3D of js_th super - object
        n = 0
        @inbounds for is = 1:length(superObjs)
            actSuperObj = superObjs[is].superObjCollision.superObj
            if !isempty(actSuperObj)
                actSuperAABB = AABB[is]
                for i = 1:length(actSuperObj)
                    actObj = actSuperObj[i]       # determine contact from this Object3D with all Object3Ds that have larger indices
                    actAABB = actSuperAABB[i]
                    for js = is+1:length(superObjs)
                        if !(js in noCPairs[is])    # PairID is not in objects which cant collide
                            nextSuperObj = superObjs[js].superObjCollision.superObj
                            nextSuperAABB = AABB[js]
                            for j = 1:length(nextSuperObj)
                                n += 1
                                nextObj  = nextSuperObj[j]
                                nextAABB = nextSuperAABB[j]
                                pairID = computePairID(scene, actObj, nextObj, is,i,js,j)

                                storeDistancesForSolver!(world, pairID, ch, actObj, nextObj,
                                actAABB, nextAABB, phase2, hasEvent)
    end; end; end; end; end; end; end
    return nothing
end

const zEps  = 1.e-8
const zEps2 = 2*zEps

function pushCollisionPair!(ch, contact, distanceWithHysteresis, pairID, contactPoint1,
    contactPoint2, contactNormal, actObj, nextObj, hasevent, supportPointsDefined, support1A, support2A,
    support3A, support1B, support2B, support3B)::Nothing
    if contact
        push!(ch.contactDict, pairID => Modia3D.ContactPair(contactPoint1,contactPoint2,
            contactNormal, actObj, nextObj, distanceWithHysteresis, supportPointsDefined,
            support1A, support2A, support3A, support1B, support2B, support3B))
    else
        if ch.noContactMinVal > distanceWithHysteresis
            ch.noContactMinVal = distanceWithHysteresis
    end; end
    return nothing
end

function storeDistancesForSolver!(world::Composition.Object3D, pairID::Composition.PairID,
    ch::Composition.ContactDetectionMPR_handler, actObj::Composition.Object3D, nextObj::Composition.Object3D, actAABB::Basics.BoundingBox, nextAABB::Basics.BoundingBox, phase2::Bool, hasEvent::Bool)

    # If object pairs are already in contact, always perform narrow phase!!
    hasContact = haskey(ch.contactDict, pairID)

    # Broad phase
    if hasContact || AABB_touching(actAABB, nextAABB) # AABB's are overlapping
        # narrow phase
        (distanceOrg, contactPoint1, contactPoint2, contactNormal, supportPointsDefined,
        support1A, support1B, support2A, support2B, support3A, support3B) =
            mpr(ch, actObj, nextObj)

    else # AABB's are not overlapping
        (distanceOrg, contactPoint1, contactPoint2, contactNormal, supportPointsDefined,
        support1A, support1B, support2A, support2B, support3A, support3B) =
            computeDistanceBetweenAABB(actAABB, nextAABB)
      #  println("AABB not overlapping")
    end

    contact                = hasEvent ? distanceOrg < -zEps : hasContact
    distanceWithHysteresis = contact  ? distanceOrg : distanceOrg + zEps2
    #  println("mpr dist: ", distanceWithHysteresis, " ", Modia3D.fullName(actObj), " ", Modia3D.fullName(nextObj))

    if hasEvent
        # At event instant
        pushCollisionPair!(ch, contact, distanceWithHysteresis, pairID,
            contactPoint1, contactPoint2, contactNormal, actObj,nextObj, hasEvent,
            supportPointsDefined, support1A, support2A, support3A, support1B, support2B, support3B)

    elseif !phase2
        # Integrator requires zero crossing function
        if contact
            # Collision pair had contact since the last event
            Modia3D.updateContactPair!(ch.contactDict[pairID], contactPoint1, contactPoint2,
                contactNormal, actObj, nextObj, distanceWithHysteresis, supportPointsDefined,
                support1A, support2A, support3A, support1B, support2B, support3B)
        else
            # Collision pair did not have contact since the last event
            pushCollisionPair!(ch, contact, distanceWithHysteresis, pairID,
                contactPoint1, contactPoint2, contactNormal, actObj, nextObj, hasEvent,
                supportPointsDefined, support1A, support2A, support3A, support1B, support2B, support3B)
        end

    else
        # Other model evaluation (only getDistances!())
        if contact
            # Collision pair had contact since the last event
            Modia3D.updateContactPair!(ch.contactDict[pairID], contactPoint1, contactPoint2,
                contactNormal, actObj, nextObj, distanceWithHysteresis, supportPointsDefined,
                support1A, support2A, support3A, support1B, support2B, support3B)
        else
            if ch.noContactMinVal > distanceWithHysteresis
                ch.noContactMinVal = distanceWithHysteresis
    end; end; end
    return nothing
end

function computeDistanceBetweenAABB(actAABB::Basics.BoundingBox, nextAABB::Basics.BoundingBox)
    xd = computeDistanceOneAxisAABB(actAABB.x_min, actAABB.x_max, nextAABB.x_min, nextAABB.x_max)
    yd = computeDistanceOneAxisAABB(actAABB.y_min, actAABB.y_max, nextAABB.y_min, nextAABB.y_max)
    zd = computeDistanceOneAxisAABB(actAABB.z_min, actAABB.z_max, nextAABB.z_min, nextAABB.z_max)
    distance = sqrt(xd^2 + yd^2 + zd^2)
    return (distance, Modia3D.ZeroVector3D, Modia3D.ZeroVector3D, Modia3D.ZeroVector3D, false, Modia3D.ZeroVector3D, Modia3D.ZeroVector3D, Modia3D.ZeroVector3D, Modia3D.ZeroVector3D, Modia3D.ZeroVector3D, Modia3D.ZeroVector3D)
end

function computeDistanceOneAxisAABB(A_min, A_max, B_min, B_max)
    if A_max < B_min
        return B_min - A_max
    elseif A_min > B_max
        return A_min - B_max
    else
        return 0.0
end; end

function Composition.closeContactDetection!(ch::Composition.ContactDetectionMPR_handler)
    empty!(ch.lastContactDict)
    empty!(ch.contactDict)
    ch.noContactMinVal =  42.0
end

function updateVisualBoundingBox!(obj::Composition.Object3D, aabb::Basics.BoundingBox)
    box::Shapes.Box = obj.shape
    box.lengthX = abs(aabb.x_max - aabb.x_min)
    box.lengthY = abs(aabb.y_max - aabb.y_min)
    box.lengthZ = abs(aabb.z_max - aabb.z_min)

    obj.r_abs = SVector((aabb.x_max + aabb.x_min)/2.0,
                        (aabb.y_max + aabb.y_min)/2.0,
                        (aabb.z_max + aabb.z_min)/2.0)
end


"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	load_cap_reserve_margin(setup::Dict, path::AbstractString, inputs_crm::Dict)

Function for reading input parameters related to planning reserve margin constraints
"""
function load_cap_reserve_margin(setup::Dict, path::AbstractString, inputs_crm::Dict)
    # Definition of capacity reserve margin (crm) by locational deliverability area (LDA)
    inputs_crm["dfCapRes"] = DataFrame(
        CSV.File(joinpath(path, "Capacity_reserve_margin.csv"), header = true),
        copycols = true,
    )

    # Ensure float format values:

    # Identifying # of planning reserve margin constraints for the system
    res = count(s -> startswith(String(s), "CapRes"), names(inputs_crm["dfCapRes"]))
    first_col = findall(s -> s == "CapRes_1", names(inputs_crm["dfCapRes"]))[1]
    last_col = findall(s -> s == "CapRes_$res", names(inputs_crm["dfCapRes"]))[1]
    inputs_crm["dfCapRes"] = Matrix{Float64}(inputs_crm["dfCapRes"][:, first_col:last_col])
    inputs_crm["NCapacityReserveMargin"] = res

    println("Capacity_reserve_margin.csv Successfully Read!")

    return inputs_crm
end

@doc raw"""
	load_cap_reserve_margin_trans(setup::Dict, inputs_crm::Dict, network_var::DataFrame)

Function for reading input parameters related to participation of transmission imports/exports in capacity reserve margin constraint.
"""
function load_cap_reserve_margin_trans(setup::Dict, inputs_crm::Dict, network_var::DataFrame)

    res = inputs_crm["NCapacityReserveMargin"]

    first_col_trans = findall(s -> s == "CapRes_1", names(network_var))[1]
    last_col_trans = findall(s -> s == "CapRes_$res", names(network_var))[1]
    dfTransCapRes = network_var[:, first_col_trans:last_col_trans]
    inputs_crm["dfTransCapRes"] = Matrix{Float64}(dfTransCapRes[completecases(dfTransCapRes), :])

    first_col_trans_derate = findall(s -> s == "DerateCapRes_1", names(network_var))[1]
    last_col_trans_derate = findall(s -> s == "DerateCapRes_$res", names(network_var))[1]
    dfDerateTransCapRes = network_var[:, first_col_trans_derate:last_col_trans_derate]
    inputs_crm["dfDerateTransCapRes"] = Matrix{Float64}(dfDerateTransCapRes[completecases(dfDerateTransCapRes), :])

    first_col_trans_excl = findall(s -> s == "CapRes_Excl_1", names(network_var))[1]
    last_col_trans_excl = findall(s -> s == "CapRes_Excl_$res", names(network_var))[1]
    dfTransCapRes_excl = network_var[:, first_col_trans_excl:last_col_trans_excl]
    inputs_crm["dfTransCapRes_excl"] = Matrix{Float64}(dfTransCapRes_excl[completecases(dfTransCapRes_excl), :])

    return inputs_crm
end


"""
    absorb

Function for constructing the `FunctionTerm{typeof(absorb)}` used in decompose.
"""
function absorb end
"""
    decompose(f::FormulaTerm,
              data,
              contrasts::Dict{Symbol},
              wts::Union{Nothing,Symbol},
              panel::Union{Nothing,Symbol},
              time::Union{Nothing,Symbol},
              estimator::Type{<:Union{EconometricsModel,ModelEstimator}},
              vce::VCE)

Decomposes the arguments passed to a fit(::EconometricModel) into its components.

# Returns

- data
- exogenous::FormulaTerm
- iv::FormulaTerm
- estimator::ModelEstimator
- X::Matrix{Float64}
- y::VecOrMat{<:Union{Bool,Int,Float64}}
- z::Vector{Float64}
- Z::Matrix{Float64}
- wts::FrequencyWeights
"""
function decompose(f::FormulaTerm,
                   data,
                   contrasts::Dict{Symbol},
                   wts::Union{Nothing,Symbol},
                   panel::Union{Nothing,Symbol},
                   time::Union{Nothing,Symbol},
                   estimator::Type{<:Union{EconometricsModel,ModelEstimator}},
                   vce::VCE)
    rhs = isa(f.rhs, Tuple) ? collect(f.rhs) : [f.rhs]
    absorbed = findall(t -> isa(t, FunctionTerm{typeof(absorb)}), rhs)
    # Conditions based on panel and temporal indices and absorbed features
    if isa(estimator, Type{<:BetweenEstimator}) && isnothing(panel)
        throw(ArgumentError("The between estimator requires the panel identifier."))
    elseif isa(estimator, Type{<:RandomEffectsEstimator}) && (isnothing(panel) || isnothing(time))
        throw(ArgumentError("The random effects estimator requires the panel and temporal identifiers."))
    elseif (isa(estimator, Type{<:BetweenEstimator}) ||
            isa(estimator, Type{<:RandomEffectsEstimator})) && !isempty(absorbed)
        throw(ArgumentError("Absorbing features is only implemented for least squares."))
    end
    pns = convert(Vector{Symbol}, filter!(!isnothing, union(termvars(f), [panel, time, wts])))
    data = select(data, pns...)
    if isa(wts, Symbol)
		wts = ifelse.(getproperty(data, wts) .≤ 0, missing, getproperty(data, wts))
    end
	categorical_variables = Tables.schema(data) |>
		(x -> [name for (name, type) ∈ zip(x.names, x.types) if
			isa(type, Type{<:Union{AbstractString, CategoricalValue}})])
	data_categorical =
		zip(categorical_variables,
			(levels(col), isordered(col)) for
		    col ∈ eachcolumn(select(data, categorical_variables...))) |>
		Dict
	data = materializer(data)(row for row ∈ rows(data) if
        all(pn -> !ismissing(getproperty(row, pn)), propertynames(row)))
    for cn ∈ categorical_variables
		col = getproperty(data, cn)
		levels!(col, intersect(data_categorical[cn][1],
							   unique(col)))
		ordered!(col, data_categorical[cn][2])
	end
    if isempty(absorbed)
        absorbed = Vector{Vector{Vector{Int}}}()
    elseif length(absorbed) == 1
        vce ≥ HC2 &&
            throw(ArgumentError("When absorbing features, only HC2-HC4 are unavailable."))
        absorbed = rhs[absorbed[1]]
        if isa(absorbed, AbstractVector)
            absorbed = unique(flatten(termvars(t) for t ∈ absorbed))
        else
            absorbed = termvars(absorbed)
        end
        absorbed = [ [ findall(isequal(l), getproperty(data, dim)) for l ∈ levels(getproperty(data, dim)) ] for dim ∈ absorbed ]
    else
        throw(ArgumentError("There can only be at most one absorb term"))
    end
    exo_rhs = filter(t -> !(isa(t, FormulaTerm) || isa(t, FunctionTerm{typeof(absorb)})), rhs)
    exogenous = FormulaTerm(f.lhs,
                            isempty(exo_rhs) ? Tuple([ConstantTerm(1)]) : Tuple(exo_rhs))
    @assert length(terms(exogenous.lhs)) == 1 "Response should be a single variable"
    iv = findall(t -> isa(t, FormulaTerm), rhs)
    if length(iv) == 0
        iv = @eval(@formula(0 ~ 0))
    elseif length(iv) == 1
        iv = rhs[iv[1]]
    else
        throw(ArgumentError("Formula syntax is response ~ exogenous + (endogenous ~ instruments)"))
    end
    sc = schema(data, contrasts)
    f = apply_schema(f, sc)
    exogenous = apply_schema(exogenous, sc, EconometricsModel)
    iv = apply_schema(iv, sc)
    # Conditions based on response type and instrumental variables
    if isa(f, FormulaTerm{<:CategoricalTerm})
        isa(iv.lhs, InterceptTerm{false}) ||
            throw(ArgumentError("Instrumental variables are only implemented for linear models."))
        vce == OIM ||
            throw(ArgumentError("Robust variance covariance estimators only available for continous response models."))
        if isordered(getproperty(data, f.lhs.sym))
            any(x -> isa(x, InterceptTerm{false}), f.rhs.terms) &&
                throw(ArgumentError("This estimator requires an intercept term."))
            estimator = OrdinalResponse(f.lhs.contrasts)
        else
            estimator = NominalResponse(f.lhs.contrasts)
        end
    elseif isa(estimator, Type{<:RandomEffectsEstimator}) &&
           any(x -> isa(x, InterceptTerm{false}), f.rhs.terms)
        throw(ArgumentError("The random effects estimator requires an intercept term."))
    end
    if isa(exogenous.lhs, CategoricalTerm)
        y = getproperty(data, exogenous.lhs.sym)
        X = modelcols(exogenous.rhs, data)
        if isordered(y)
            X = X[:,2:end]
        end
    else
        y, X = modelcols(exogenous, data)
    end
    if !isa(iv.lhs, InterceptTerm)
        z, Z = modelcols(iv, data)
        z = isa(z, Tuple) ? hcat(z...) : convert(Matrix{Float64}, reshape(z, length(z), 1))
    else
        z = zeros(0, 0)
        Z = zeros(0, 0)
    end
	wts = isnothing(wts) ?
		  FrequencyWeights(ones(length(y))) :
		  FrequencyWeights(collect(wts))
    if isa(estimator, Type{<:RandomEffectsEstimator})
        panel = (panel,
                 [ findall(isequal(l), getproperty(data, panel)) for l ∈ levels(getproperty(data, panel)) ])
        time = (time, [ findall(isequal(l), getproperty(data, time)) for l ∈ levels(getproperty(data, time)) ])
        estimator = RandomEffectsEstimator(panel, time, X, y, z, Z, FrequencyWeights(wts))
    elseif isa(estimator, Type{<:BetweenEstimator})
        estimator = BetweenEstimator(panel, [ findall(isequal(l), getproperty(data, panel)) for l ∈ levels(getproperty(data, panel)) ])
    elseif isa(estimator, Type{<:EconometricModel}) || isa(estimator, Type{<:ContinuousResponse})
        estimator = ContinuousResponse(absorbed)
    end
    data, exogenous, iv, estimator, X, y, z, Z, wts
end

remove_intercept(f::AbstractString) =
    occursin(" ~ 1 + ", f) ? replace(f, " ~ 1 + " => " ~ ") : f
add_intercept(f::AbstractString) =
    !occursin(r" ~ -?[0-1] ?", f) ? replace(f, r"(^.*?) ~ " => s"\1 ~ 1 + ") : f
clean_lhs(obj::AbstractTerm) = string(obj)
clean_lhs(obj::FunctionTerm) = string(obj)[3:end - 1]
clean_rhs(obj::Tuple) = mapreduce(clean_rhs, (x, y) -> "$x + $y", obj)
clean_rhs(obj::AbstractTerm) = string(obj)
clean_rhs(obj::InteractionTerm) = mapreduce(clean_rhs, (x, y) -> "$x & $y", obj.terms)
clean_rhs(obj::FunctionTerm) = string(obj)[2:end]
clean_rhs(obj::FormulaTerm) = string("(", obj.lhs, " ~ ", clean_rhs(obj.rhs), ")")
clean_rhs(obj::FunctionTerm{typeof(absorb)}) = string(obj)[3:end - 1]
clean_fm(obj::EconometricsModel) =
    obj.f |>
    (f -> string(clean_lhs(f.lhs), " ~ ", clean_rhs(f.rhs))) |>
    (f -> isa(obj, EconometricModel{<:Union{RandomEffectsEstimator,
                                            OrdinalResponse}}) ?
          remove_intercept(f) : add_intercept(f)) |>
    (f -> "Formula: $f")