"""
During optimization, a `SolverLog` is used to store convergence information.

A `SolverLog` can log the following:
- `x::Array{Float64,2}`: the history of the decision variables for every iteration.
- `abs_feas_error::Array{Float64,1}`: the absolute feasibility error evolution.
- `abs_opt_error::Array{Float64,1}`: the absolute optimality error evolution.
- `obj_value::Array{Float64,1}`: the evolution of the objective function value.
- `fc_evals::Array{Float64,1}`: the evolution of the number of function evaluations.
- `ga_evals::Array{Float64,1}`: the evolution of the number of gradient evaluations.
- `nStatus::Int`: the solution status return code.

The actual values stored in a `SolverLog` depend on the implementation of the callback used by each solver for every new point.
Using Ipopt, the callback is set with [`Ipopt.setIntermediateCallback`](https://github.com/jump-dev/Ipopt.jl/blob/d9e9176620a9b527a08991a3d41062fa948867f7/test/C_wrapper.jl#L123-L131).
Using Knitro, the callback is set with [`KNITRO.KN_set_newpt_callback`](https://github.com/jump-dev/KNITRO.jl/blob/43de7d38f01d7aa7c418bb24d0c1ab89b61e527a/examples/nlp2.jl#L201-L204).
"""
mutable struct SolverLog
    x::Array{Float64,2}
    abs_feas_error::Array{Float64,1}
    abs_opt_error::Array{Float64,1}
    obj_value::Array{Float64,1}
    fc_evals::Array{Float64,1}
    ga_evals::Array{Float64,1}
    nStatus::Int  # solution status return code

    @doc """
        SolverLog(N)

    Create an empty NLP solver log.

    # Arguments
    - `N::Int`: the total number of decision variables of the NLP.

    See also: [`update!`](@ref)
    """
    function SolverLog(N::Int)
        new(Array{Float64,2}(undef, N, 0),
            Array{Float64,1}(undef, 0),
            Array{Float64,1}(undef, 0),
            Array{Float64,1}(undef, 0),
            Array{Float64,1}(undef, 0),
            Array{Float64,1}(undef, 0),
            -999)
    end
end

"""
    update!(log; <keyword arguments>)

Update a [`SolverLog`](@ref) object.

# Keyword arguments
- `x=zeros(size(log.x, 1))`: an intermediate iterate point.
- `abs_feas_error=0`: the absolute feasibility error at the current point.
- `abs_opt_error=0`: the absolute optimality error at the current point.
- `obj_value=0`: the value of the objective function at the current point.
- `fc_evals=0`: the number of function evaluations requested by the NLP solver so far.
- `ga_evals=0`: the number of gradient evaluations requested by the NLP solver so far.
"""
function update!(log::SolverLog;
                 x=zeros(size(log.x, 1)),
                 abs_feas_error=0,
                 abs_opt_error=0,
                 obj_value=0,
                 fc_evals=0,
                 ga_evals=0)
    log.x = hcat(log.x, x)
    log.abs_feas_error = vcat(log.abs_feas_error, abs_feas_error)
    log.abs_opt_error = vcat(log.abs_opt_error, abs_opt_error)
    log.obj_value = vcat(log.obj_value, obj_value)
    log.fc_evals = vcat(log.fc_evals, fc_evals)
    log.ga_evals = vcat(log.ga_evals, ga_evals)
end

"""
    length(log) -> Integer

Return the length of a [`SolverLog`](@ref).

The length of a [`SolverLog`](@ref) is the same as the number of iterations taken to solve the problem.
"""
function length(log::SolverLog)
    num_iters = size(log.x, 2)
    @assert all(==(num_iters), length.([log.abs_feas_error, log.abs_opt_error, log.obj_value, log.fc_evals, log.ga_evals]))
    return num_iters
end

"""
    save(log, file)

Save a [`SolverLog`](@ref) to an `.npz` file.

See also: [`load!`](@ref)
"""
function save(log::SolverLog, file::String)
    npzwrite(file, Dict(
        "x" => log.x,
        "abs_feas_error" => log.abs_feas_error,
        "abs_opt_error" => log.abs_opt_error,
        "obj_value" => log.obj_value,
        "fc_evals" => log.fc_evals,
        "ga_evals" => log.ga_evals,
        "nStatus" => log.nStatus,
    ))
end

"""
    load!(log, file)

Load an `.npz` file and store it in a [`SolverLog`](@ref) object.

See also: [`save`](@ref)
"""
function load!(log::SolverLog, file::String)
    data = npzread(file)

    log.x = data["x"]
    log.abs_feas_error = data["abs_feas_error"]
    log.abs_opt_error = data["abs_opt_error"]
    log.obj_value = data["obj_value"]
    log.fc_evals = data["fc_evals"]
    log.ga_evals = data["ga_evals"]
    log.nStatus = data["nStatus"]

    return log
end

export
    SolverLog,
    update!,
    length,
    save,
    load!


using Random
import Random: rand, srand

const RAND_MAX = typemax(Int32)

struct LibcRNG <: AbstractRNG end

srand(r::LibcRNG, seed::Integer) = ccall(:srand, Cvoid, (Cuint,), seed)
rand(r::LibcRNG, ::Type{Cint})   = ccall(:rand,  Cint, ())
rand(r::LibcRNG)                 = rand(r, Cint)

struct LibcRNG48 <: AbstractRNG end

srand(r::LibcRNG48, seed::Integer) = ccall(:srand48, Cvoid, (Clong,), seed)
rand(r::LibcRNG48, ::Type{Culong}) = ccall(:lrand48, Cint, ())
rand(r::LibcRNG48, ::Type{Clong})  = ccall(:mrand48, Cint, ())


module covid19abm
using Parameters, Distributions, StatsBase, StaticArrays, Random, Match, DataFrames

@enum HEALTH SUS LAT PRE ASYMP MILD MISO INF IISO HOS ICU REC DED UNDEF

Base.@kwdef mutable struct Human
    idx::Int64 = 0
    health::HEALTH = SUS
    swap::HEALTH = UNDEF
    sickfrom::HEALTH = UNDEF
    age::Int64   = 0    # in years. don't really need this but left it incase needed later
    ag::Int64    = 0
    tis::Int64   = 0   # time in state
    exp::Int64   = 0   # max statetime
    dur::NTuple{4, Int8} = (0, 0, 0, 0)   # Order: (latents, asymps, pres, infs) TURN TO NAMED TUPS LATER
    doi::Int32   = 999   # day of infection.
    iso::Bool = false  ## isolated (limited contacts)
    isovia::Symbol = :null ## isolated via quarantine (:qu), preiso (:pi), intervention measure (:im), or contact tracing (:ct)
    tracing::Bool = false ## are we tracing contacts for this individual?
    tracestart::Int8 = -1 ## when to start tracing, based on values sampled for x.dur
    traceend::Int8 = -1 ## when to end tracing
    tracedby::UInt16 = 0 ## is the individual traced? property represents the index of the infectious person
    tracedxp::UInt16 = 0 ## the trace is killed after tracedxp amount of days
end

Base.@kwdef mutable struct tt
    idx::Int64 = 0
    dur::NTuple{4, Int8} = (lat=0, exp=0, scp=0, asp=0)
end

## default system parameters
@with_kw mutable struct ModelParameters @deftype Float64    ## use @with_kw from Parameters
    β = 0.0
    prov::Symbol = :ontario
    calibration::Bool = false
    modeltime::Int64 = 500
    initialinf::Int64 = 1
    τmild::Int64 = 0 ## days before they self-isolate for mild cases
    fmild::Float64 = 0.0  ## percent of people practice self-isolation
    fsevere::Float64 = 0.0 # fixed at 0.80
    eldq::Float64 = 0.0 ## complete isolation of elderly
    eldqag::Int8 = 5 ## default age group, if quarantined(isolated) is ag 5.
    fasymp::Float64 = 0.50 ## percent going to asymp
    fpre::Float64 = 1.0 ## NOT USED ANYMORE (percent going to presymptomatic)
    fpreiso::Float64 = 0.0 ## percent that is isolated at the presymptomatic stage
    tpreiso::Int64 = 0## preiso is only turned on at this time.
    frelasymp::Float64 = 0.11 ## relative transmission of asymptomatic
    fctcapture::Float16 = 0.0 ## how many symptomatic people identified
    fcontactst::Float16 = 0.0 ## fraction of contacts being isolated/quarantined
    cidtime::Int8 = 0  ## time to identification (for CT) post symptom onset
    cdaysback::Int8 = 0 ## number of days to go back and collect contacts
end

Base.show(io::IO, ::MIME"text/plain", z::Human) = dump(z)

## constants
const HSIZE = 10000
const humans = Array{Human}(undef, HSIZE) # run 100,000
const p = ModelParameters()  ## setup default parameters
const agebraks = @SVector [0:4, 5:19, 20:49, 50:64, 65:99]

export ModelParameters, HEALTH, Human, humans

function runsim(simnum, ip::ModelParameters)
    # function runs the `main` function, and collects the data as dataframes.
    hmatrix = main(ip)
    # get infectors counters
    infectors = _count_infectors()

    ct_numbers = _count_ct_numbers()

    # get simulation age groups
    ags = [x.ag for x in humans] # store a vector of the age group distribution
    all = _collectdf(hmatrix)
    spl = _splitstate(hmatrix, ags)
    ag1 = _collectdf(spl[1])
    ag2 = _collectdf(spl[2])
    ag3 = _collectdf(spl[3])
    ag4 = _collectdf(spl[4])
    ag5 = _collectdf(spl[5])
    insertcols!(all, 1, :sim => simnum); insertcols!(ag1, 1, :sim => simnum); insertcols!(ag2, 1, :sim => simnum);
    insertcols!(ag3, 1, :sim => simnum); insertcols!(ag4, 1, :sim => simnum); insertcols!(ag5, 1, :sim => simnum);
    return (a=all, g1=ag1, g2=ag2, g3=ag3, g4=ag4, g5=ag5, infectors=infectors, ct_numbers=ct_numbers)
end
export runsim

function main(ip::ModelParameters)
    #Random.seed!(sim*726)
    ## datacollection
    # matrix to collect model state for every time step

    # reset the parameters for the simulation scenario
    reset_params(ip)

    hmatrix = zeros(Int64, HSIZE, p.modeltime)
    initialize() # initialize population
    # insert initial infected agents into the model
    # and setup the right swap function.
    if p.calibration
        insert_infected(PRE, p.initialinf, 4)
    else
        insert_infected(LAT, p.initialinf, 4)
    end

    ## save the preisolation isolation parameters
    _fpreiso = p.fpreiso
    p.fpreiso = 0

    # split population in agegroups
    grps = get_ag_dist()

    # start the time loop
    for st = 1:p.modeltime
        # start of day
        if st == p.tpreiso ## time to introduce testing
            p.fpreiso = _fpreiso
        end
        _get_model_state(st, hmatrix) ## this datacollection needs to be at the start of the for loop
        #dyntrans(st)
        dyntrans(st, grps)
        #dyntrans_nosus(st)
        sw = time_update()
        # end of day
    end
    return hmatrix ## return the model state as well as the age groups.
end
export main

reset_params_default() = reset_params(ModelParameters())
function reset_params(ip::ModelParameters)
    # the p is a global const
    # the ip is an incoming different instance of parameters
    # copy the values from ip to p.
    for x in propertynames(p)
        setfield!(p, x, getfield(ip, x))
    end
end
export reset_params, reset_params_default

function _model_check()
    ## checks model parameters before running
    (p.fctcapture > 0 && p.fpreiso > 0) && error("Can not do contact tracing and ID/ISO of pre at the same time.")
    (p.fctcapture > 0 && p.maxtracedays == 0) && error("maxtracedays can not be zero")
end

## Data Collection/ Model State functions
function _get_model_state(st, hmatrix)
    # collects the model state (i.e. agent status at time st)
    for i=1:length(humans)
        hmatrix[i, st] = Int(humans[i].health)
    end
end
export _get_model_state

function _collectdf(hmatrix)
    ## takes the output of the humans x time matrix and processes it into a dataframe
    #_names_inci = Symbol.(["lat_inc", "mild_inc", "miso_inc", "inf_inc", "iiso_inc", "hos_inc", "icu_inc", "rec_inc", "ded_inc"])
    #_names_prev = Symbol.(["sus", "lat", "mild", "miso", "inf", "iiso", "hos", "icu", "rec", "ded"])
    mdf_inc, mdf_prev = _get_incidence_and_prev(hmatrix)
    mdf = hcat(mdf_inc, mdf_prev)
    _names_inc = Symbol.(string.((Symbol.(instances(HEALTH)[1:end - 1])), "_INC"))
    _names_prev = Symbol.(string.((Symbol.(instances(HEALTH)[1:end - 1])), "_PREV"))
    _names = vcat(_names_inc..., _names_prev...)
    datf = DataFrame(mdf, _names)
    insertcols!(datf, 1, :time => 1:p.modeltime) ## add a time column to the resulting dataframe
    return datf
end

function _splitstate(hmatrix, ags)
    #split the full hmatrix into 4 age groups based on ags (the array of age group of each agent)
    #sizes = [length(findall(x -> x == i, ags)) for i = 1:4]
    matx = []#Array{Array{Int64, 2}, 1}(undef, 4)
    for i = 1:length(agebraks)
        idx = findall(x -> x == i, ags)
        push!(matx, view(hmatrix, idx, :))
    end
    return matx
end
export _splitstate

function _get_incidence_and_prev(hmatrix)
    cols = instances(HEALTH)[1:end - 1] ## don't care about the UNDEF health status
    inc = zeros(Int64, p.modeltime, length(cols))
    pre = zeros(Int64, p.modeltime, length(cols))
    for i = 1:length(cols)
        inc[:, i] = _get_column_incidence(hmatrix, cols[i])
        pre[:, i] = _get_column_prevalence(hmatrix, cols[i])
    end
    return inc, pre
end

function _get_column_incidence(hmatrix, hcol)
    inth = Int(hcol)
    timevec = zeros(Int64, p.modeltime)
    for r in eachrow(hmatrix)
        idx = findfirst(x -> x == inth, r)
        if idx !== nothing
            timevec[idx] += 1
        end
    end
    return timevec
end

function _get_column_prevalence(hmatrix, hcol)
    inth = Int(hcol)
    timevec = zeros(Int64, p.modeltime)
    for (i, c) in enumerate(eachcol(hmatrix))
        idx = findall(x -> x == inth, c)
        if idx !== nothing
            ps = length(c[idx])
            timevec[i] = ps
        end
    end
    return timevec
end

function _count_infectors()
    pre_ctr = asymp_ctr = mild_ctr = inf_ctr = 0
    for x in humans
        if x.health != SUS ## meaning they got sick at some point
            if x.sickfrom == PRE
                pre_ctr += 1
            elseif x.sickfrom == ASYMP
                asymp_ctr += 1
            elseif x.sickfrom == MILD || x.sickfrom == MISO
                mild_ctr += 1
            elseif x.sickfrom == INF || x.sickfrom == IISO
                inf_ctr += 1
            else
                error("sickfrom not set right: $(x.sickfrom)")
            end
        end
    end
    return (pre_ctr, asymp_ctr, mild_ctr, inf_ctr)
end

function _count_ct_numbers()
    howmanytrc = 0 # how many traced (num / (inf + mild) ~ percentage)
    howmanyisoct = 0 # how many isolated (of all sick people) (contact tracing)
    howmanyisopi = 0 # how many isolated (of all sick people) (through presymp)
    howmanyiso = 0
    for x in humans
        if x.tracestart > 0
            howmanytrc += 1
        end
        if x.isovia == :ct && x.health != SUS
            howmanyisoct += 1
        end
        if x.isovia == :pi && x.health != SUS
            howmanyisopi += 1
        end
        if x.isovia != :null && x.health != SUS
            howmanyiso += 1
    end
end
    return howmanytrc, howmanyiso, howmanyisoct, howmanyisopi
end
export _collectdf, _get_incidence_and_prev, _get_column_incidence, _get_column_prevalence, _count_infectors, _count_ct_numbers

## initialization functions
function get_province_ag(prov)
    ret = @match prov begin
        :alberta => Distributions.Categorical(@SVector [0.0655, 0.1851, 0.4331, 0.1933, 0.1230])
        :bc => Distributions.Categorical(@SVector [0.0475, 0.1570, 0.3905, 0.2223, 0.1827])
        :canada => Distributions.Categorical(@SVector [0.0540, 0.1697, 0.3915, 0.2159, 0.1689])
        :manitoba => Distributions.Categorical(@SVector [0.0634, 0.1918, 0.3899, 0.1993, 0.1556])
        :newbruns => Distributions.Categorical(@SVector [0.0460, 0.1563, 0.3565, 0.2421, 0.1991])
        :newfdland => Distributions.Categorical(@SVector [0.0430, 0.1526, 0.3642, 0.2458, 0.1944])
        :nwterrito => Distributions.Categorical(@SVector [0.0747, 0.2026, 0.4511, 0.1946, 0.0770])
        :novasco => Distributions.Categorical(@SVector [0.0455, 0.1549, 0.3601, 0.2405, 0.1990])
        :nunavut => Distributions.Categorical(@SVector [0.1157, 0.2968, 0.4321, 0.1174, 0.0380])
        :ontario => Distributions.Categorical(@SVector [0.0519, 0.1727, 0.3930, 0.2150, 0.1674])
        :pei => Distributions.Categorical(@SVector [0.0490, 0.1702, 0.3540, 0.2329, 0.1939])
        :quebec => Distributions.Categorical(@SVector [0.0545, 0.1615, 0.3782, 0.2227, 0.1831])
        :saskat => Distributions.Categorical(@SVector [0.0666, 0.1914, 0.3871, 0.1997, 0.1552])
        :yukon => Distributions.Categorical(@SVector [0.0597, 0.1694, 0.4179, 0.2343, 0.1187])
        :newyork   => Distributions.Categorical(@SVector [0.064000, 0.163000, 0.448000, 0.181000, 0.144000])
        _ => error("shame for not knowing your canadian provinces and territories")
    end
    return ret
end
export get_province_ag

function initialize()
    agedist = get_province_ag(p.prov)
    @inbounds for i = 1:HSIZE
        humans[i] = Human()              ## create an empty human
        x = humans[i]
        x.idx = i
        x.ag = rand(agedist)
        x.age = rand(agebraks[x.ag])
        x.exp = 999  ## susceptible people don't expire.
        x.dur = sample_epi_durations() # sample epi periods
        if rand() < p.eldq && x.ag == p.eldqag   ## check if elderly need to be quarantined.
            x.iso = true
            x.isovia = :qu
        end
    end
end
export initialize

function get_ag_dist()
    # splits the initialized human pop into its age groups
    grps =  map(x -> findall(y -> y.ag == x, humans), 1:length(agebraks))
    return grps
end

function insert_infected(health, num, ag)
    ## inserts a number of infected people in the population randomly
    ## this function should resemble move_to_inf()
    l = findall(x -> x.health == SUS && x.ag == ag, humans)
    if length(l) > 0 && num < length(l)
        h = sample(l, num; replace = false)
        @inbounds for i in h
            x = humans[i]
            if health == PRE
                move_to_pre(x) ## the swap may be asymp, mild, or severe, but we can force severe in the time_update function
            elseif health == LAT
                move_to_latent(x)
            elseif health == INF
                move_to_infsimple(x)
            else
                error("can not insert human of health $(health)")
            end
            x.sickfrom = INF # this will add +1 to the INF count in _count_infectors()... keeps the logic simple in that function.
        end
    end
    return h
end
export insert_infected

function time_update()
    # counters to calculate incidence
    lat=0; pre=0; asymp=0; mild=0; miso=0; inf=0; infiso=0; hos=0; icu=0; rec=0; ded=0;
    for x in humans
        x.tis += 1
        x.doi += 1 # increase day of infection. variable is garbage until person is latent
        if x.tis >= x.exp
            @match Symbol(x.swap) begin
                :LAT  => begin move_to_latent(x); lat += 1; end
                :PRE  => begin move_to_pre(x); pre += 1; end
                :ASYMP => begin move_to_asymp(x); asymp += 1; end
                :MILD => begin move_to_mild(x); mild += 1; end
                :MISO => begin move_to_miso(x); miso += 1; end
                :INF  => begin move_to_inf(x); inf +=1; end
                :IISO => begin move_to_iiso(x); infiso += 1; end
                :HOS  => begin move_to_hospicu(x); hos += 1; end
                :ICU  => begin move_to_hospicu(x); icu += 1; end
                :REC  => begin move_to_recovered(x); rec += 1; end
                :DED  => begin move_to_dead(x); ded += 1; end
                _    => error("swap expired, but no swap set.")
            end
        end
        # run covid-19 functions for other integrated dynamics.
        ct_dynamics(x)
    end
    return (lat, mild, miso, inf, infiso, hos, icu, rec, ded)
end
export time_update

@inline _set_isolation(x::Human, iso) = _set_isolation(x, iso, x.isovia)
@inline function _set_isolation(x::Human, iso, via)
    # a helper setter function to not overwrite the isovia property.
    # a person could be isolated in susceptible/latent phase through contact tracing
    # --> in which case it will follow through the natural history of disease
    # --> if the person remains susceptible, then iso = off
    # a person could be isolated in presymptomatic phase through fpreiso
    # --> if x.iso == true from CT and x.isovia == :ct, do not overwrite
    # a person could be isolated in mild/severe phase through fmild, fsevere
    # --> if x.iso == true from CT and x.isovia == :ct, do not overwrite
    # --> if x.iso == true from PRE and x.isovia == :pi, do not overwrite
    x.iso = iso
    x.isovia == :null && (x.isovia = via)
end

function sample_epi_durations()
    # when a person is sick, samples the
    lat_dist = truncated(LogNormal(log(5.2), 0.1), 4, 7) # truncated between 4 and 7
    pre_dist = truncated(Gamma(1.058, 5/2.3), 0.8, 3)#truncated between 0.8 and 3
    asy_dist = Gamma(5, 1)
    inf_dist = Gamma((3.2)^2/3.7, 3.7/3.2)

    latents = Int.(round.(rand(lat_dist)))
    pres = Int.(round.(rand(pre_dist)))
    latents = latents - pres # ofcourse substract from latents, the presymp periods
    asymps = Int.(ceil.(rand(asy_dist)))
    infs = Int.(ceil.(rand(inf_dist)))
    return (latents, asymps, pres, infs)
end

function move_to_latent(x::Human)
    ## transfers human h to the incubation period and samples the duration
    x.health = LAT
    x.doi = 0 ## day of infection is reset when person becomes latent
    x.tis = 0   # reset time in state
    x.exp = x.dur[1] # get the latent period
    x.swap = rand() < p.fasymp ? ASYMP : PRE
    ## in calibration mode, latent people never become infectious.
    if p.calibration
        x.swap = LAT
        x.exp = 999
    end
end
export move_to_latent

function move_to_asymp(x::Human)
    ## transfers human h to the asymptomatic stage
    x.health = ASYMP
    x.tis = 0
    x.exp = x.dur[2] # get the presymptomatic period
    x.swap = REC
    # x.iso property remains from either the latent or presymptomatic class
    # if x.iso is true, the asymptomatic individual has limited contacts
end
export move_to_asymp

function move_to_pre(x::Human)
    θ = (0.95, 0.9, 0.85, 0.6, 0.2)  # percentage of sick individuals going to mild infection stage
    x.health = PRE
    x.tis = 0   # reset time in state
    x.exp = x.dur[3] # get the presymptomatic period

    if rand() < θ[x.ag]
        x.swap = MILD
    else
        x.swap = INF
    end
    # calculate whether person is isolated
    rand() < p.fpreiso && _set_isolation(x, true, :pi)
end
export move_to_pre

function move_to_mild(x::Human)
    ## transfers human h to the mild infection stage for γ days
    x.health = MILD
    x.tis = 0
    x.exp = x.dur[4]
    x.swap = REC
    # x.iso property remains from either the latent or presymptomatic class
    # if x.iso is true, staying in MILD is same as MISO since contacts will be limited.
    # we still need the separation of MILD, MISO because if x.iso is false, then here we have to determine
    # how many days as full contacts before self-isolation
    # NOTE: if need to count non-isolated mild people, this is overestimate as isolated people should really be in MISO all the time
    #   and not go through the mild compartment
    if x.iso || rand() < p.fmild
        x.swap = MISO
        x.exp = p.τmild
    end
end
export move_to_mild

function move_to_miso(x::Human)
    ## transfers human h to the mild isolated infection stage for γ days
    x.health = MISO
    x.swap = REC
    x.tis = 0
    x.exp = x.dur[4] - p.τmild  ## since tau amount of days was already spent as infectious
    _set_isolation(x, true, :mi)
end
export move_to_miso

function move_to_infsimple(x::Human)
    ## transfers human h to the severe infection stage for γ days
    ## simplified function for calibration/general purposes
    x.health = INF
    x.tis = 0
    x.exp = x.dur[4]
    x.swap = REC
    _set_isolation(x, false, :null)
end

function move_to_inf(x::Human)
    ## transfers human h to the severe infection stage for γ days
    ## for swap, check if person will be hospitalized, selfiso, die, or recover

    # h = prob of hospital, c = prob of icu AFTER hospital
    h = (rand(Uniform(0.02, 0.03)), rand(Uniform(0.02, 0.03)), rand(Uniform(0.28, 0.34)), rand(Uniform(0.28, 0.34)), rand(Uniform(0.60, 0.68)))
    c = (rand(Uniform(0.01, 0.015)), rand(Uniform(0.01, 0.015)), rand(Uniform(0.03, 0.05)), rand(Uniform(0.05, 0.1)), rand(Uniform(0.05, 0.15)))
    mh = [0.01/5, 0.01/5, 0.0135/3, 0.01225/1.5, 0.04/2]     # death rate for severe cases.

    if p.calibration
        h =  (0, 0, 0, 0, 0)
        c =  (0, 0, 0, 0, 0)
        mh = (0, 0, 0, 0, 0)
    end

    time_to_hospital = Int(round(rand(Uniform(2, 5)))) # duration symptom onset to hospitalization

    x.health = INF
    x.swap = UNDEF
    x.tis = 0
    if rand() < h[x.ag]     # going to hospital or ICU but will spend delta time transmissing the disease with full contacts
        x.exp = time_to_hospital
        x.swap = rand() < c[x.ag] ? ICU : HOS
    else ## no hospital for this lucky (but severe) individual
        if rand() < mh[x.ag]
            x.exp = x.dur[4]
            x.swap = DED
        else
            x.exp = x.dur[4]
            x.swap = REC
            if x.iso || rand() < p.fsevere
                x.exp = 1  ## 1 day isolation for severe cases
                x.swap = IISO
            end
        end
    end
    ## before returning, check if swap is set
    x.swap == UNDEF && error("agent I -> ?")
end

function move_to_iiso(x::Human)
    ## transfers human h to the sever isolated infection stage for γ days
    x.health = IISO
    x.swap = REC
    x.tis = 0     ## reset time in state
    x.exp = x.dur[4] - 1  ## since 1 day was spent as infectious
    _set_isolation(x, true, :mi)
end

function move_to_hospicu(x::Human)
    ## to do: don't sample from distribution unless needed
    ## to do: check whether the parameters seyed gave are for mean/std or shape/scale
    #mh = [0.001, 0.001, 0.00135, 0.01225, 0.04]
    #mc = [0.002, 0.002, 0.0027, 0.0245, 0.08]
    mh = [0.01*2, 0.01*2, 0.0135, 0.01225, 0.03*2]
    mc = 2*mh
    #mh = [0, 0, 0, 0.01225*2, 0.04*2]
    #mc = [0, 0, 0, 0.01225*3.5, 0.04*3.5]

    psiH = Int(round(rand(truncated(Gamma(4.5, 2.75), 8, 17))))
    psiC = Int(round(rand(truncated(Gamma(4.5, 2.75), 8, 17)))) + 2
    muH = Int(round(rand(truncated(Gamma(5.3, 2.1), 9, 15))))
    muC = Int(round(rand(truncated(Gamma(5.3, 2.1), 9, 15)))) + 2

    swaphealth = x.swap
    x.health = swaphealth ## swap either to HOS or ICU
    x.swap = UNDEF
    x.tis = 0
    # hospital patients are isolated, but it's irrelevant because they don't have contacts and can't get sick again
    _set_isolation(x, true) # do not set the isovia property here.

    if swaphealth == HOS
        if rand() < mh[x.ag] ## person will die in the hospital
            x.exp = muH
            x.swap = DED
        else
            x.exp = psiH
            x.swap = REC
        end
    end

    if swaphealth == ICU
        if rand() < mc[x.ag] ## person will die in the ICU
            x.exp = muC
            x.swap = DED
        else
            x.exp = psiC
            x.swap = REC
        end
    end
    ## before returning, check if swap is set
    x.swap == UNDEF && error("agent H -> ?")
end

function move_to_dead(h::Human)
    # no level of alchemy will bring someone back to life.
    h.health = DED
    h.swap = UNDEF
    h.tis = 0
    h.exp = 999 ## stay recovered indefinitely
    h.iso = true # a dead person is isolated
    _set_isolation(h, true)  # do not set the isovia property here.
    # isolation property has no effect in contact dynamics anyways (unless x == SUS)
end

function move_to_recovered(h::Human)
    h.health = REC
    h.swap = UNDEF
    h.tis = 0
    h.exp = 999 ## stay recovered indefinitely
    h.iso = false ## a recovered person has ability to meet others
    _set_isolation(h, false)  # do not set the isovia property here.
    # isolation property has no effect in contact dynamics anyways (unless x == SUS)
end

function ct_dynamics(x::Human)
    # main function for ct dynamics
    # turns tracing on if person is infectious and in the right time window
    # applies isolation to all traced contacts
    # turns of isolation for susceptibles > 14 days
    # order of if statements matter here
    !(p.fctcapture > 0) && (return)
    xh = x.health
    xs = x.swap
    dur = x.dur
    doi = x.doi
    #fctcapture::Float16 = 0.0 ## how many of contacts of the infected are we tracing.
    #cidtime::Int8 = 0  ## time to identification (for CT) post symptom onset
    #cdaysback::Int8 = 0 ## number of days to go back and collect contacts
    # tracing::Bool = false ## are we tracing contacts for this individual?
    # tracestart::Int8 = -1 ## when to start tracing, based on values sampled for x.dur
    # traceend::Int8 = -1 ## when to end tracing
    # tracedby::UInt16 = 0 ## is the individual traced? property represents the index of the infectious person
    # tracedxp::UInt16 = 0 ## the trace is killed after tracedxp amount of days

    ## person is newly infectious, calculate tracing numbers
    if xh == LAT && xs != ASYMP && doi == 0
        if rand() < p.fctcapture
            delta = dur[1] + dur[3] + p.cidtime # the latent + presymp + time to identification time
            q = delta - p.cdaysback
            #println("delta = $delta, q = $q")
            x.tracestart = q
            x.traceend = delta
            (x.tracestart > x.traceend) && error("tracestart < traceend")
        end
    end

    ## turn on trace when day of infection == tracestart
    if doi == x.tracestart
        x.tracing = true
    end
    if doi == x.traceend
        ## have to do more work here
        x.tracing = false
        _set_isolation(x, true, :ct)
        alltraced = findall(y -> y.tracedby == x.idx, humans)
        for i in alltraced
            y = humans[i]
            if y.health in (SUS, LAT, PRE, ASYMP, MILD, MISO, INF, IISO)
                _set_isolation(y, true, :ct)
            end
        end
    end

    # a susceptible that was traced through infected is only isolated for 14 days.
    if xh == SUS && x.iso == true && x.isovia == :ct
        x.tracedxp -= 1
        if x.tracedxp == 0
            _set_isolation(x, false)
        end
    end
    return
end
export ct_dynamics

@inline function _get_betavalue(sys_time, xhealth)
    bf = p.β ## baseline PRE
    # values coming from FRASER Figure 2... relative tranmissibilities of different stages.
    if xhealth == ASYMP
        bf = bf * p.frelasymp
    elseif xhealth == MILD || xhealth == MISO
        bf = bf * 0.44
    elseif xhealth == INF || xhealth == IISO
        bf = bf * 0.89
    end
    return bf
end
export _get_betavalue

function dyntrans(sys_time, grps)
    totalinf = 0 # count number of new infected
    ## find all the people infectious
    infs = findall(x -> x.health in (PRE, ASYMP, MILD, MISO, INF, IISO), humans)

    # get all people to meet and their daily contacts to recieve
    # we can sample this at the start of the simulation to avoid everyday
    tomeet = Array{Int64, 1}(undef, HSIZE)
    for i in 1:HSIZE
        x = humans[i]
        idx = x.idx
        ag = x.ag
        cnt = 0
        #if person is isolated, they can recieve only 3 maximum contacts
        if x.iso
            cnt = rand() < 0.5 ? 0 : rand(1:3)
        else
            cnt = rand(nbs[ag])  # expensive operation, try to optimize
        end
        tomeet[i] = cnt
    end

    # go through every infectious person
    for xid in infs
        x = humans[xid]
        xhealth = x.health
        cnts = tomeet[xid]
        # split the counts over age groups
        gpw = Int.(round.(cm[x.ag]*cnts))
        for (i, g) in enumerate(gpw)
            # sample the people from each group
            meet = rand(grps[i], g)
            # go through each person
            for j in meet
                y = humans[j]
                ycnt = tomeet[y.idx]

                if ycnt > 0
                    tomeet[y.idx] = tomeet[y.idx] - 1 # remove a contact
                    # there is a contact to recieve
                     # tracing dynamics

                    if x.tracing
                        if y.tracedby == 0 && rand() < p.fcontactst
                            y.tracedby = x.idx
                            y.tracedxp = 14 ## trace isolation will last for 14 days before expiry
                        end
                    end

                # tranmission dynamics
                    if  y.health == SUS && y.swap == UNDEF
                        beta = _get_betavalue(sys_time, xhealth)
                        if rand() < beta
                            totalinf += 1
                            y.swap = LAT
                            y.exp = y.tis   ## force the move to latent in the next time step.
                            y.sickfrom = xhealth ## stores the infector's status to the infectee's sickfrom
                        end
                    end

                end

            end
        end
    end
    return totalinf
end
export dyntrans

function dyntrans_nosus(sys_time)
    totalinf = 0
    ## find all the people infectious
    infs = findall(x -> x.health in (PRE, ASYMP, MILD, MISO, INF, IISO), humans)
    tomeet = map(1:length(agebraks)) do grp ## will also meet dead people, but ignore for now because it's such a small group
        findall(humans) do h
            c1 = h.ag == grp
            c2 = h.health == SUS ? !(h.iso) : true
            return c1 && c2
        end
    end
    #length(tomeet) <= 5 && return totalinf
    for xid in infs
        x = humans[xid]
        ag = x.ag
        ih = x.health
        if x.iso ## isolated infectious person has limited contacts
            cnt = rand() < 0.5 ? 0 : rand(1:3)
        else
            cnt = rand(nbs[ag])  ## get number of contacts/day
        end
        #cnt >= length(tomeet[ag]) && error("error here")

        # distribute cnt_meet to different groups based on contact matrix.
        # these are not probabilities, but proportions. be careful.
        # going from cnt to the gpw array might remove a contact or two due to rounding.
        gpw = Int.(round.(cm[ag]*cnt))
        #println("cnt: $cnt, gpw: $gpw")
        # let's stratify the human population in it's age groups.
        # this could be optimized by putting it outside the contact_dynamic2 function and passed in as arguments
        # enumerate over the 15 groups and randomly select contacts from each group
        for (i, g) in enumerate(gpw)
            if length(tomeet[i]) > 0
                meet = rand(tomeet[i], g)    # sample 'g' number of people from this group  with replacement
                ## remember, two infected people may meet the same susceptible so its possible that disease can be transferred "twice"
                for j in meet
                    y = humans[j]
                    ## if we are tracing the infected individual, record contacts
                    if x.tracing
                        if y.tracedby == 0
                            y.tracedby = x.idx
                            y.tracedxp = 14 ## trace isolation will last for 14 days before expiry
                        end
                    end
                    if y.health == SUS && y.swap == UNDEF && !y.iso #i.e. a wasted contact
                        bf = p.β ## baseline PRE
                        # values coming from FRASER Figure 2... relative tranmissibilities of different stages.
                        if ih == ASYMP
                            bf = bf * 0.11
                        elseif ih == MILD || ih == MISO
                            bf = bf * 0.44
                        elseif ih == INF || ih == IISO
                            bf = bf * 0.89
                        end
                        if rand() < bf
                            totalinf += 1
                            y.swap = LAT
                            y.exp = y.tis   ## force the move to latent in the next time step.
                            y.sickfrom = ih ## stores the infector's status to the infectee's sickfrom
                        end
                    end
                end
            end
        end
    end
    return totalinf
end
export dyntrans

### old contact matrix
# function contact_matrix()
#     CM = Array{Array{Float64, 1}, 1}(undef, 4)
#     CM[1]=[0.5712, 0.3214, 0.0722, 0.0353]
#     CM[2]=[0.1830, 0.6253, 0.1423, 0.0494]
#     CM[3]=[0.1336, 0.4867, 0.2723, 0.1074]
#     CM[4]=[0.1290, 0.4071, 0.2193, 0.2446]
#     return CM
# end

function contact_matrix()
    # regular contacts, just with 5 age groups.
    #  0-4, 5-19, 20-49, 50-64, 65+
    CM = Array{Array{Float64, 1}, 1}(undef, 5)
    CM[1] = [0.2287, 0.1839, 0.4219, 0.1116, 0.0539]
    CM[2] = [0.0276, 0.5964, 0.2878, 0.0591, 0.0291]
    CM[3] = [0.0376, 0.1454, 0.6253, 0.1423, 0.0494]
    CM[4] = [0.0242, 0.1094, 0.4867, 0.2723, 0.1074]
    CM[5] = [0.0207, 0.1083, 0.4071, 0.2193, 0.2446]
    return CM
end
#
# calibrate for 2.7 r0
# 20% selfisolation, tau 1 and 2.

function negative_binomials()
    ## the means/sd here are calculated using _calc_avgag
    means = [10.21, 16.793, 13.7950, 11.2669, 8.0027]
    sd = [7.65, 11.7201, 10.5045, 9.5935, 6.9638]
    totalbraks = length(means)
    nbinoms = Vector{NegativeBinomial{Float64}}(undef, totalbraks)
    for i = 1:totalbraks
        p = 1 - (sd[i]^2-means[i])/(sd[i]^2)
        r = means[i]^2/(sd[i]^2-means[i])
        nbinoms[i] =  NegativeBinomial(r, p)
    end
    return nbinoms
end
const nbs = negative_binomials()
const cm = contact_matrix()
export negative_binomials, contact_matrix, nbs, cm

## internal functions to do intermediate calculations
function _calc_avgag(lb, hb)
    ## internal function to calculate the mean/sd of the negative binomials
    ## returns a vector of sampled number of contacts between age group lb to age group hb
    dists = _negative_binomials_15ag()[lb:hb]
    totalcon = Vector{Int64}(undef, 0)
    for d in dists
        append!(totalcon, rand(d, 10000))
    end
    return totalcon
end
export _calc_avgag

function _negative_binomials_15ag()
    ## negative binomials 15 agegroups
    AgeMean = Vector{Float64}(undef, 15)
    AgeSD = Vector{Float64}(undef, 15)
    #0-4, 5-9, 10-14, 15-19, 20-24, 25-29, 30-34, 35-39, 40-44, 45-49, 50-54, 55-59, 60-64, 65-69, 70+
    AgeMean = [10.21, 14.81, 18.22, 17.58, 13.57, 13.57, 14.14, 14.14, 13.83, 13.83, 12.3, 12.3, 9.21, 9.21, 6.89]
    AgeSD = [7.65, 10.09, 12.27, 12.03, 10.6, 10.6, 10.15, 10.15, 10.86, 10.86, 10.23, 10.23, 7.96, 7.96, 5.83]

    nbinoms = Vector{NegativeBinomial{Float64}}(undef, 15)
    for i = 1:15
        p = 1 - (AgeSD[i]^2-AgeMean[i])/(AgeSD[i]^2)
        r = AgeMean[i]^2/(AgeSD[i]^2-AgeMean[i])
        nbinoms[i] =  NegativeBinomial(r, p)
    end
    return nbinoms
end
export negative_binomials

## references:
# critical care capacity in Canada https://www.ncbi.nlm.nih.gov/pubmed/25888116
end # module end


using DataFrames
using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = page_model()
    include("../src/components/AdaptationCosts.jl")

    adaptationcosts_noneconomic = addadaptationcosts_noneconomic(m)
    adaptationcosts_noneconomic[:y_year_0] = 2015.
    adaptationcosts_noneconomic[:y_year] = Mimi.dim_keys(m.md, :time)
    adaptationcosts_noneconomic[:gdp] = readpagedata(m, "test/validationdata/$valdir/gdp.csv")

    adaptationcosts_economic = addadaptationcosts_economic(m)
    adaptationcosts_economic[:y_year_0] = 2015.
    adaptationcosts_economic[:y_year] = Mimi.dim_keys(m.md, :time)
    adaptationcosts_economic[:gdp] = readpagedata(m, "test/validationdata/$valdir/gdp.csv")

    adaptationcosts_sealevel = addadaptationcosts_sealevel(m)
    adaptationcosts_sealevel[:y_year_0] = 2015.
    adaptationcosts_sealevel[:y_year] = Mimi.dim_keys(m.md, :time)
    adaptationcosts_sealevel[:gdp] = readpagedata(m, "test/validationdata/$valdir/gdp.csv")

    p = load_parameters(m)
    set_leftover_params!(m, p)

    run(m)

    autofac = m[:AdaptiveCostsEconomic, :autofac_autonomouschangefraction]
    atl_economic = m[:AdaptiveCostsEconomic, :atl_adjustedtolerablelevel]
    imp_economic = m[:AdaptiveCostsEconomic, :imp_adaptedimpacts]
    acp_economic = m[:AdaptiveCostsEconomic, :acp_adaptivecostplateau]
    aci_economic = m[:AdaptiveCostsEconomic, :aci_adaptivecostimpact]

    autofac_compare = readpagedata(m, "test/validationdata/autofac_autonomouschangefraction.csv")
    atl_economic_compare = readpagedata(m, "test/validationdata/atl_adjustedtolerableleveloftemprise_market.csv")
    imp_economic_compare = readpagedata(m, "test/validationdata/$valdir/imp_actualreduction_market.csv")
    acp_economic_compare = readpagedata(m, "test/validationdata/acp_adaptivecostplateau_economic.csv")
    aci_economic_compare = readpagedata(m, "test/validationdata/aci_adaptivecostimpact_economic.csv")

    @test autofac ≈ autofac_compare rtol=1e-8
    @test atl_economic ≈ atl_economic_compare rtol=1e-8
    @test imp_economic ≈ imp_economic_compare rtol=1e-8
    @test acp_economic ≈ acp_economic_compare rtol=1e-3
    @test aci_economic ≈ aci_economic_compare rtol=1e-3

    ac_noneconomic = m[:AdaptiveCostsNonEconomic, :ac_adaptivecosts]
    ac_economic = m[:AdaptiveCostsEconomic, :ac_adaptivecosts]
    ac_sealevel = m[:AdaptiveCostsSeaLevel, :ac_adaptivecosts]

    ac_noneconomic_compare = readpagedata(m, "test/validationdata/ac_adaptationcosts_noneconomic.csv")
    ac_economic_compare = readpagedata(m, "test/validationdata/ac_adaptationcosts_economic.csv")
    ac_sealevel_compare = readpagedata(m, "test/validationdata/$valdir/ac_adaptationcosts_sealevelrise.csv")

    @test ac_noneconomic ≈ ac_noneconomic_compare rtol=1e-2
    @test ac_economic ≈ ac_economic_compare rtol=1e-3
    @test ac_sealevel ≈ ac_sealevel_compare rtol=1e-2
end


using DataFlow: Call, constant, inputnode, syntax

const ops = Dict{Symbol,Any}()
include("ops.jl")

# This is to fetch weights when they are stored in
# the constant tensor and not in intializer.
function get_weights(g::Types.Graph)
  temp = Dict{Any, Any}()
  for node in g.node
    if node.op_type == "Constant"
      temp[node.name] = get_array(node.attribute[:value])
    end
  end
  return temp
end

vcall(a...) = vertex(Call(), constant.(a)...)

# Placeholder for array values
weights(f::Types.Model) = weights(f.graph)

"""
Checks location of weights and returns appropriate
values.
Note: Constant weight is deprecated now.
"""
function weights(g::Types.Graph)
  count = 0
  for node in g.node
    if (node.op_type == "Constant")
      count = count + 1
      break
    end
  end
  if (count > 0)
    return get_weights(g)
  end
  return g.initializer
end

function inputs(g::Types.Graph)
  ws = weights(g)
  i = 0
  Dict(x.name => haskey(ws, x.name) ?
        constant(:(weights[$(x.name)])) :
        inputnode(i += 1)
       for x in g.input), i
end

function _graph(g::Types.Graph)
  vs, n = inputs(g)
  for node in g.node
    if node.op_type == "Constant"
      vs[node.output[1]] = ops[Symbol(node.op_type)](node, map(n -> vs[n], node.input)...)
    else
      vs[node.output[1]] = ops[Symbol(node.op_type)](node.attribute, map(n -> vs[n], node.input)...)
    end
  end
  return vs[g.output[1].name], n
end

# Graph Cleanups

ischainable(v) = DataFlow.iscall(v) && all(x -> DataFlow.isconstant(x), v[3:end])
chaindepth(v) = ischainable(v) ? chaindepth(v[2]) + 1 : 0

function _tochain(v, ch)
  ischainable(v) || return v
  if length(v[:]) ≤ 2
    push!(ch, v[1])
  else
    push!(ch, vertex(DataFlow.Lambda(1, vcall(v[1], inputnode(1), v[3:end]...))))
  end
  return _tochain(v[2], ch)
end

function tochain(v)
  ch = []
  v = _tochain(v, ch)
  vcall(vcall(:Chain, reverse(ch)...), v)
end

function chainify(v)
  MacroTools.prewalk(v) do v
    chaindepth(v) > 3 ? tochain(v) : v
  end
end

# Interface
function graph(g::Types.Graph)
  v, n = _graph(g)
  v = chainify(v)
  return vertex(DataFlow.Lambda(n, v)) |> DataFlow.λopen |> DataFlow.λclose
end

"""
Write out the Julia code for the model
"""
code(g::Types.Graph) = graph(g) |> syntax |>
  MacroTools.flatten |> MacroTools.gensym_ids

# function breakcalls(ex)
#   MacroTools.prewalk(ex) do ex
#     iscall(ex) || return ex
#     @capture(ex, f_(args__) | f_.(args__))
#     count(x -> iscall(x), args) ≥ 2 || return ex
#     vars = []
#     args = map(args) do x
#       iscall(x) || return x
#       var = gensym()
#       push!(vars, :($var = $x))
#       return var
#     end
#     ex = isexpr(ex, :call) ? :($f($(args...))) : :($f.($(args...)))
#     :($(vars...); $ex)
#   end
# end


println("Loading packages...")
@time begin
  using Cytof5
  using Random, Distributions
  # TODO: Get rid of this dep
  using RCall
  using BSON
  using ArgParse
  import Cytof5.Model.logger
  include("util.jl")
end
println("Done loading packages.")

# TODO: review
# ARG PARSING

# Define behavior for vector numerical args (input as space delimited string)
ArgParse.parse_item(::Type{Vector{T}}, x::AbstractString) where {T <: Number} = parse.(T, split(x))

function parse_cmd()
  s = ArgParseSettings()

  @add_arg_table s begin
    "--simdat_path"
      arg_type = String
      required = true
    "--MCMC_ITER"
      arg_type = Int
      default = 1000
    "--BURN"
      arg_type = Int
      default = 10000
    "--K_MCMC"
      arg_type = Int
      required = true
    "--L0_MCMC"
      arg_type = Int
      required = true
    "--L1_MCMC"
      arg_type = Int
      required = true

    # Missin mechanism
    "--pBounds"
      arg_type = Vector{Float64}
      default = [.05, .8, .05] # paper
    "--yQuantiles"
      arg_type = Vector{Float64}
      default = [0.0, .25, .5] # paper

    "--RESULTS_DIR"
      arg_type = String
      required = true
    "--SEED"
      arg_type = Int
      default = 0
    "--EXP_NAME"
      arg_type = String
    "--printFreq"
      arg_type = Int
      default = 50
  end

  PARSED_ARGS = parse_args(s)

  if PARSED_ARGS["EXP_NAME"] == nothing
    logger("EXP_NAME not defined. Making default experiment name.")
    MAIN_ARGS = filter(d -> !(d.first in ("RESULTS_DIR", "EXP_NAME")), PARSED_ARGS)
    PARSED_ARGS["EXP_NAME"] = join(["$k$v" for (k, v) in MAIN_ARGS], '_')
  end

  return PARSED_ARGS
end

PARSED_ARGS = parse_cmd()
for (k,v) in PARSED_ARGS
  logger("$k => $v")
end

SIMDAT_PATH = PARSED_ARGS["simdat_path"]

MCMC_ITER = PARSED_ARGS["MCMC_ITER"]
BURN = PARSED_ARGS["BURN"]

K_MCMC = PARSED_ARGS["K_MCMC"]
L0_MCMC = PARSED_ARGS["L0_MCMC"]
L1_MCMC = PARSED_ARGS["L1_MCMC"]
L_MCMC = Dict(0 => L0_MCMC, 1 => L1_MCMC)

# missing mechanism
yQuantiles = PARSED_ARGS["yQuantiles"]
pBounds = PARSED_ARGS["pBounds"]

EXP_NAME = PARSED_ARGS["EXP_NAME"]
SEED = PARSED_ARGS["SEED"]
RESULTS_DIR = PARSED_ARGS["RESULTS_DIR"]
printFreq = PARSED_ARGS["printFreq"]
# END OF ARG PARSING

# CREATE RESULTS DIR
OUTDIR = "$(RESULTS_DIR)/$(EXP_NAME)/"
mkpath(OUTDIR)

# Set random seed
Random.seed!(SEED);

logger("Load simulated data ...");
BSON.@load SIMDAT_PATH simdat
dat = Cytof5.Model.Data(simdat[:y])

logger("Generating priors ...");
@time c = Cytof5.Model.defaultConstants(dat, K_MCMC, L_MCMC,
                                        tau0=10.0, tau1=10.0,
                                        sig2_prior=InverseGamma(3.0, 2.0),
                                        alpha_prior=Gamma(0.1, 10.0),
                                        yQuantiles=yQuantiles, pBounds=pBounds,
                                        similarity_Z=Cytof5.Model.sim_fn_abs(10000),
                                        probFlip_Z=2.0 / (dat.J * K_MCMC),
                                        noisyDist=Normal(0.0, 3.16))
Cytof5.Model.printConstants(c)
println("dat.I: $(dat.I)")
println("dat.J: $(dat.J)")
println("dat.N: $(dat.N)")

# Plot missing mechanism
logger("Plot missing mechanism")
util.plotPdf("$(OUTDIR)/prob_miss.pdf")
R"par(mfrow=c($(dat.I), 1))"
for i in 1:dat.I
  util.plotProbMiss(c.beta, i)
end
R"par(mfrow=c(1,1))"
util.devOff()

logger("Generating initial state ...");
# @time init = Cytof5.Model.genInitialState(c, dat)
logger("use smart init ...")
@time init = Cytof5.Model.smartInit(c, dat)

# Plot initial Z
util.plotPdf("$(OUTDIR)/Z_init.pdf")
addGridLines(J::Int, K::Int, col="grey") = util.abline(v=(1:K) .+ .5, h=(1:J) .+ .5, col=col)
util.myImage(init.Z, xlab="Features", ylab="Markers", addL=false, f=Z->addGridLines(dat.J, c.K))
util.devOff()

# Fit Model
nsamps_to_thin(nsamps::Int, nmcmc::Int) = max(1, div(nmcmc, nsamps))

@time out, lastState, ll, metrics, dden=
  Cytof5.Model.cytof5_fit(init, c, dat,
                          monitors=[[:Z, :lam, :W,
                                     :sig2, :delta,
                                     :alpha, :v,
                                     :eta, :eps],
                                    [:y_imputed, :gam]],
                          thins=[2, nsamps_to_thin(10, MCMC_ITER)],
                          nmcmc=MCMC_ITER, nburn=BURN,
                          computeLPML=true, computeDIC=true,
                          computedden=true, thin_dden=nsamps_to_thin(200, MCMC_ITER),
                          use_repulsive=false,
                          joint_update_Z=true,
                          printFreq=10, flushOutput=true)

logger("Saving Data ...");
println("length of dden: $(length(dden))")
# @save "$(OUTDIR)/output.jld2" out dat ll lastState c dat metrics
BSON.@save "$(OUTDIR)/output.bson" out ll lastState c metrics init dden simdat

logger("MCMC Completed.");

#= Test
julia --color=yes sim.jl --I=3 --J=32 --N_factor=100 --K=8 --K_MCMC=10 --L0_MCMC=5 --L1_MCMC=5 \
      --MCMC_ITER=10 --BURN=10 --RESULTS_DIR=results --EXP_NAME=bla
=#


mutable struct TimeDerivatives{N,MD,ED}
    accelerations::Vector{SVectorF{N}}
    mass_derivatives::MD
    energy_derivatives::ED

    function TimeDerivatives(
        velocities::Velocities{N},
        mass_component::MassComponent{MD,<:Kernel{N}},
        energy_component::EnergyComponent{ED},
    ) where {N,MD,ED}
        new{N,MD,ED}(
            initderivatives(velocities),
            initderivatives(mass_component),
            initderivatives(energy_component),
        )
    end
end

"""
    initderivatives!(
        derivatives::TimeDerivatives{N,MD,ED},
        velocities::Velocities{N},
        mass_component::MassComponent{MD,<:Kernel{N}},
        energy_component::EnergyComponent{ED},
    )

Resets the given fluid derivatives.
"""
function initderivatives!(
    derivatives::TimeDerivatives{N,MD,ED},
    velocities::Velocities{N},
    mass_component::MassComponent{MD,<:Kernel{N}},
    energy_component::EnergyComponent{ED},
) where {N,MD,ED}
    initderivatives!(derivatives.accelerations, velocities)
    initderivatives!(derivatives.mass_derivatives, mass_component)
    initderivatives!(derivatives.energy_derivatives, energy_component)
end

"""
    updatefluid!(
        time_derivatives::TimeDerivatives{N,MD,ED},
        mass_component::MassComponent{MD,<:Kernel{N}},
        energy_component::EnergyComponent{ED},
        diffusion_component::DiffusionComponent,
        pressures::Pressures,
        velocities::Velocities{N},
        positions::Vector{SVectorF{N}},
        eos::EquationOfState,
        boundary_component::BoundaryComponent,
    )

Update the state of the fluid variables and derivatives according to the
current particle positions.
"""
function updatefluid!(
    time_derivatives::TimeDerivatives{N,MD,ED},
    mass_component::MassComponent{MD,<:Kernel{N}},
    energy_component::EnergyComponent{ED},
    diffusion_component::DiffusionComponent,
    pressures::Pressures,
    velocities::Velocities{N},
    positions::Vector{SVectorF{N}},
    eos::EquationOfState,
    boundary_component::BoundaryComponent{N},
) where {N,MD,ED}
    updatemasses!(mass_component, positions)
    updatepressures!(pressures, mass_component, energy_component, eos)
    updatediffusion!(diffusion_component, energy_component, eos)
    updateboundaries!(
        boundary_component,
        positions,
        velocities,
        mass_component,
        pressures,
    )

    initderivatives!(
        time_derivatives,
        velocities,
        mass_component,
        energy_component,
    )

    n::UInt = length(positions)
    for i = 1:n-1, j = i:n
        rᵢⱼ = positions[j] - positions[i]
        r = norm(rᵢⱼ)
        hᵢ_hⱼ = get_kernel_widths(mass_component, (i, j))
        qᵢ_qⱼ = normalize_distance(mass_component, r, hᵢ_hⱼ)

        if r > 0 && nonzeroat(get_kernel(mass_component), qᵢ_qⱼ)
            r̂ᵢⱼ = rᵢⱼ / r
            updatederivatives!(
                time_derivatives,
                (i, j),
                (rᵢⱼ, r̂ᵢⱼ, r),
                hᵢ_hⱼ,
                qᵢ_qⱼ,
                velocities,
                pressures,
                mass_component,
                energy_component,
                diffusion_component,
            )
        end
    end

    boundary_positions = get_positions(boundary_component)
    nb::UInt = length(boundary_positions)
    for i = 1:n, j = 1:nb
        rᵢⱼ = boundary_positions[j] - positions[i]
        r = norm(rᵢⱼ)
        hᵢ = get_kernel_width(mass_component, i)
        qᵢ = normalize_distance(mass_component, r, hᵢ)

        if r > 0 && nonzeroat(get_kernel(mass_component), qᵢ)
            r̂ᵢⱼ = rᵢⱼ / r
            updatederivatives_boundary!(
                time_derivatives,
                (i, j),
                (rᵢⱼ, r̂ᵢⱼ, r),
                hᵢ,
                qᵢ,
                velocities,
                pressures,
                mass_component,
                energy_component,
                diffusion_component,
                boundary_component,
            )
        end
    end
end

"""
    updatederivatives!(
        time_derivatives::TimeDerivatives{N,MD,ED},
        i_j::Tuple{Unsigned,Unsigned},
        (rᵢⱼ, r̂ᵢⱼ, r)::Tuple{SVectorF{N},SVectorF{N},Number},
        hᵢ_hⱼ::Union{Number,Tuple{Number,Number}},
        qᵢ_qⱼ::Union{Number,Tuple{Number,Number}},
        velocities::Velocities{N},
        pressures::Pressures,
        mass_component::MassComponent{MD,<:Kernel{N}},
        energy_component::EnergyComponent{ED},
        diffusion_component::DiffusionComponent,
    )

Update time derivatives for particles `i` and `j`.
"""
function updatederivatives!(
    time_derivatives::TimeDerivatives{N,MD,ED},
    i_j::Tuple{Unsigned,Unsigned},
    (rᵢⱼ, r̂ᵢⱼ, r)::Tuple{SVectorF{N},SVectorF{N},Number},
    hᵢ_hⱼ::Union{Number,Tuple{Number,Number}},
    qᵢ_qⱼ::Union{Number,Tuple{Number,Number}},
    velocities::Velocities{N},
    pressures::Pressures,
    mass_component::MassComponent{MD,<:Kernel{N}},
    energy_component::EnergyComponent{ED},
    diffusion_component::DiffusionComponent,
) where {N,MD,ED}
    vᵢ, vⱼ = get_velocities(velocities, i_j)
    vᵢⱼ = vⱼ - vᵢ
    vᵣ = vᵢⱼ ⋅ r̂ᵢⱼ

    ρᵢ_ρⱼ = get_mass_densities(mass_component, i_j)
    Pᵢ_Pⱼ = get_pressures(pressures, i_j)

    ddr_W_r_hᵢ_hⱼ = ddr(get_kernel(mass_component), qᵢ_qⱼ, hᵢ_hⱼ)

    m_ddr_W_r_hᵢ_hⱼ = compute_m_ddr_W_r_h(i_j, ddr_W_r_hᵢ_hⱼ, mass_component)

    Δaᵢ_pressure, Δaⱼ_pressure =
        compute_Δa_pressure(ρᵢ_ρⱼ, Pᵢ_Pⱼ, m_ddr_W_r_hᵢ_hⱼ)

    Δa_diffusion = compute_Δa_diffusion(
        r,
        vᵣ,
        hᵢ_hⱼ,
        ρᵢ_ρⱼ,
        ddr_W_r_hᵢ_hⱼ,
        mass_component,
        diffusion_component,
    )

    updatederivatives!(
        time_derivatives.mass_derivatives,
        i_j,
        vᵣ,
        m_ddr_W_r_hᵢ_hⱼ,
        mass_component,
    )
    updatederivatives!(
        time_derivatives.energy_derivatives,
        i_j,
        vᵣ,
        (Δaᵢ_pressure, Δaⱼ_pressure, Δa_diffusion),
        energy_component,
    )
    updatederivatives!(
        time_derivatives.accelerations,
        i_j,
        r̂ᵢⱼ,
        (Δaᵢ_pressure, Δaⱼ_pressure, Δa_diffusion),
        velocities,
    )
end

"""
    updatederivatives_boundary!(
        time_derivatives::TimeDerivatives{N,MD,ED},
        (i, j)::Tuple{Unsigned,Unsigned},
        (rᵢⱼ, r̂ᵢⱼ, r)::Tuple{SVectorF{N},SVectorF{N},Number},
        hᵢ::Number,
        qᵢ::Number,
        velocities::Velocities{N},
        pressures::Pressures,
        mass_component::MassComponent{MD,<:Kernel{N}},
        energy_component::EnergyComponent{ED},
        diffusion_component::DiffusionComponent,
        boundary_component::BoundaryComponent{N},
    )

Update time derivatives for particle `i` due to boundary particle `j`.
"""
function updatederivatives_boundary!(
    time_derivatives::TimeDerivatives{N,MD,ED},
    (i, j)::Tuple{Unsigned,Unsigned},
    (rᵢⱼ, r̂ᵢⱼ, r)::Tuple{SVectorF{N},SVectorF{N},Number},
    hᵢ::Number,
    qᵢ::Number,
    velocities::Velocities{N},
    pressures::Pressures,
    mass_component::MassComponent{MD,<:Kernel{N}},
    energy_component::EnergyComponent{ED},
    diffusion_component::DiffusionComponent,
    boundary_component::BoundaryComponent{N},
) where {N,MD,ED}
    vᵢ = get_velocity(velocities, i)
    vⱼ = get_velocity(boundary_component, j)
    vᵢⱼ = vⱼ - vᵢ
    vᵣ = vᵢⱼ ⋅ r̂ᵢⱼ

    ρᵢ = get_mass_density(mass_component, i)
    ρⱼ = compute_mass_density(boundary_component, energy_component, eos, i, j)

    Pᵢ = get_pressure(pressures, i)
    Pⱼ = get_pressure(boundary_component, j)

    ddr_W_r_hᵢ = ddr(get_kernel(mass_component), qᵢ, hᵢ)

    m_ddr_W_r_hᵢ = compute_m_ddr_W_r_h(i, ddr_W_r_hᵢ, mass_component)

    Δaᵢ_pressure, Δaⱼ_pressure =
        compute_Δa_pressure((ρᵢ, ρⱼ), (Pᵢ, Pⱼ), m_ddr_W_r_hᵢ)

    Δa_diffusion = compute_Δa_diffusion(
        r,
        vᵣ,
        hᵢ,
        (ρᵢ, ρⱼ),
        ddr_W_r_hᵢ,
        mass_component,
        diffusion_component,
    )

    updatederivative!(
        time_derivatives.mass_derivatives,
        i,
        vᵣ,
        ddr_W_r_hᵢ,
        mass_component,
    )
    updatederivative!(
        time_derivatives.energy_derivatives,
        i,
        vᵣ,
        (Δaᵢ_pressure, Δaⱼ_pressure, Δa_diffusion),
        energy_component,
    )
    updatederivative!(
        time_derivatives.accelerations,
        i,
        r̂ᵢⱼ,
        (Δaᵢ_pressure, Δaⱼ_pressure, Δa_diffusion),
        velocities,
    )
end

"""
    evolvevariables!(
        velocities::Velocities{N},
        mass_component::MassComponent{MD},
        energy_component::EnergyComponent{ED},
        time_derivatives::TimeDerivatives{N,MD,ED},
        Δt::Number,
    )

Evolve the components with the given derivatives and time step.
"""
function evolvevariables!(
    velocities::Velocities{N},
    mass_component::MassComponent{MD},
    energy_component::EnergyComponent{ED},
    time_derivatives::TimeDerivatives{N,MD,ED},
    Δt::Number,
) where {N,MD,ED}
    evolvevariables!(velocities, time_derivatives.accelerations, Δt)
    evolvevariables!(mass_component, time_derivatives.mass_derivatives, Δt)
    evolvevariables!(energy_component, time_derivatives.energy_derivatives, Δt)
end

"""
    evolvevariables!(
        velocities::Velocities{N},
        mass_component::MassComponent{MD},
        energy_component::EnergyComponent{ED},
        Δt::Number,
        time_derivatives_and_weights::Tuple{TimeDerivatives{N,MD,ED},Float}...,
    )

Evolve the components with the given time step and weighed sets of derivatives.
"""
function evolvevariables!(
    velocities::Velocities{N},
    mass_component::MassComponent{MD},
    energy_component::EnergyComponent{ED},
    Δt::Number,
    time_derivatives_and_weights::Tuple{TimeDerivatives{N,MD,ED},Float}...,
) where {N,MD,ED}
    for (time_derivatives, weight) in time_derivatives_and_weights
        weighted_Δt = weight * Δt
        evolvevariables!(
            velocities,
            mass_component,
            energy_component,
            time_derivatives,
            weighted_Δt,
        )
    end
end

"""
    evolvepositions!(
        positions::Vector{SVectorF{N}},
        Δt::Number,
        velocities_and_weights::Tuple{Velocities{N},Number}...,
    )

Evolves the positions with the given time step weighed sets of velocities.
"""
function evolvepositions!(
    positions::Vector{SVectorF{N}},
    Δt::Number,
    velocities_and_weights::Tuple{Velocities{N},Number}...,
) where {N}
    for (velocities, weight) in velocities_and_weights
        weighted_Δt = weight * Δt
        evolvepositions!(positions, velocities, weighted_Δt)
    end
end


export reconstruction

include("RecoParameters.jl")
include("DirectReconstruction.jl")
include("IterativeReconstruction.jl")

"""
    reconstruction(acqData::AcquisitionData, recoParams::Dict)

Performs image reconstruction of an AcquisitionData object.
Parameters are specified in a dictionary.

Reconstruction types are specified by the symbol `:reco`.
Valid reconstruction names are:
* :direct - direct Fourier reconstruction
* :standard           - iterative reconstruction for all contrasts, coils & slices independently
* :multiEcho          - iterative joint reconstruction of all echo images
* :multiCoil          - SENSE-type iterative reconstruction
* :multiCoilMultiEcho - SENSE-type iterative reconstruction of all echo images
"""
function reconstruction(acqData::AcquisitionData, recoParams::Dict)
  # check dimensionality of encoding
  encodingDims = dims(trajectory(acqData))
  if encodingDims==3 && numSlices(acqData)>1
    @error "reconstruction of multiple 3d-encoded volumina is not yet supported"
  end

  # load reconstruction parameters
  recoParams = merge(defaultRecoParams(), recoParams)

  # iterative reco
  reconSize, weights, sparseTrafo, reg, normalize, encOps, solvername, senseMaps = setupIterativeReco(acqData, recoParams)
  if recoParams[:reco] == "standard"
    return reconstruction_simple(acqData, reconSize[1:encodingDims], reg, sparseTrafo, weights, solvername, normalize, encOps, recoParams)
  elseif recoParams[:reco] == "multiEcho"
    return reconstruction_multiEcho(acqData, reconSize[1:encodingDims], reg, sparseTrafo, weights, solvername, normalize, encOps, recoParams)
  elseif recoParams[:reco] == "multiCoil"
    return reconstruction_multiCoil(acqData, reconSize[1:encodingDims], reg, sparseTrafo, weights, solvername, senseMaps, normalize, encOps, recoParams)
  elseif recoParams[:reco] == "multiCoilMultiEcho"
    return reconstruction_multiCoilMultiEcho(acqData, reconSize[1:encodingDims], reg, sparseTrafo, weights, solvername, senseMaps, normalize, encOps, recoParams)
  end

  # direct reco
  if recoParams[:reco] != "direct"
    @error "reco modell $(recoParams[:reco]) not found"
  end
  reconSize, weights, cmap = setupDirectReco(acqData, recoParams)
  return reconstruction_direct(acqData, reconSize[1:encodingDims], weights, cmap)
end

"""
    reconstruction(acqData::AcquisitionData, recoParams::Dict,filename::String; force=false)

performs the same image reconstrucion as `reconstruction(acqData::AcquisitionData, recoParams::Dict)`
and saves the image in a file with name `filename`.
If `force=false`, the reconstructed image is loaded from the the file `filename` if the latter is
present.
"""
function reconstruction(acqData::AcquisitionData, recoParams::Dict, filename::String;
                        force=false)
  if !force && isfile(filename)
    return recoImage( RecoFileIBI(filename) )
  else
    I = reconstruction(acqData, recoParams)
    saveasRecoFile(filename, I, recoParams)
    return I
  end
end


# Routines for observations of light sources in particular images,
# rather than for sources in the abstract, in physical units,
# and rather than for images alone (that's image_model.jl).

"""
Attributes of the patch of sky surrounding a single
celestial object in a single image.

Attributes:
  - center: The approximate source location in world coordinates
  - radius_pix: The width of the influence of the object in pixel coordinates
  - psf: The point spread function in this region of the sky
  - wcs_jacobian: The jacobian of the WCS transform in this region of the
                  sky for each band
  - pixel_center: The pixel location of center in each band.
"""
immutable SkyPatch
    center::Vector{Float64}
    radius_pix::Float64

    psf::Vector{PsfComponent}
    wcs_jacobian::Matrix{Float64}
    pixel_center::Vector{Float64}
end

"""
Initialize a SkyPatch from an existing SkyPatch and a new PSF.
"""
SkyPatch(patch::SkyPatch, psf::Vector{PsfComponent}) =
    SkyPatch(patch.center, patch.radius_pix, psf, patch.wcs_jacobian,
             patch.pixel_center)

"""Centers of patches in pixel coordinates"""
patch_ctrs_pix(patches::Vector{SkyPatch}) = [p.pixel_center for p in patches]

"""Radii of patches in pixel coordinates"""
patch_radii_pix(patches::Vector{SkyPatch}) = [p.radius_pix for p in patches]

"""
A fast function to determine which sources might belong to which tiles.

Args:
  - tiles: A TiledImage
  - patch_ctrs: Vector of length-2 vectors, giving pixel center of each patch.
  - patch_radii_px: Radius of each patch.

Returns:
  - An array (over tiles) of a vector of candidate
    source patches.  If a patch is a candidate, it may be within the patch
    radius of a point in the tile, though it might not.
"""
function local_source_candidates(tile::ImageTile,
                                 patch_ctrs::Vector{Vector{Float64}},
                                 patch_radii_px::Vector{Float64})
    @assert length(patch_ctrs) == length(patch_radii_px)

    ret = Int[]

    # Find the patches that are less than the radius plus diagonal from the
    # center of the tile.  These are candidates for having some
    # overlap with the tile.
    tile_center = (mean(tile.h_range), mean(tile.w_range))
    h_width, w_width = size(tile.pixels)
    tile_diag = (0.5 ^ 2) * (h_width ^ 2 + w_width ^ 2)

    for s in 1:length(patch_ctrs)
        patch_dist = (tile_center[1] - patch_ctrs[s][1])^2
                    + (tile_center[2] - patch_ctrs[s][2])^2
        if patch_dist <= (tile_diag + patch_radii_px[s])^2
            push!(ret, s)
        end
    end

    return ret
end

"""
Args:
  - tile: An ImageTile (containing tile coordinates)
  - patch_ctrs: Vector of length-2 vectors, giving pixel center of each patch.
  - patch_radii_px: Radius of each patch.

Returns:
  - A vector of source ids (from 1 to length(patches)) that influence
    pixels in the tile.  A patch influences a tile if
    there is any overlap in their squares of influence.
"""
function get_local_sources(tile::ImageTile,
                           patch_ctrs::Vector{Vector{Float64}},
                           patch_radii_px::Vector{Float64})
    @assert length(patch_ctrs) == length(patch_radii_px)
    tile_source_map = Int[]
    tile_ctr = (mean(tile.h_range), mean(tile.w_range))
    h_width, w_width = size(tile.pixels)

    for i in eachindex(patch_ctrs)
        patch_ctr = patch_ctrs[i]
        patch_r = patch_radii_px[i]

        # This is a "ball" in the infinity norm.
        if ((abs(tile_ctr[1] - patch_ctr[1]) < patch_r + 0.5 * h_width) &&
            (abs(tile_ctr[2] - patch_ctr[2]) < patch_r + 0.5 * w_width))
            push!(tile_source_map, i)
        end
    end

    tile_source_map
end

"""
Get the sources associated with each tile in a TiledImage.

Args:
  - tiled_image: A TiledImage
  - patch_ctrs: Vector of length-2 vectors, giving pixel center of each patch.
  - patch_radii_px: Radius of each patch.
Returns:
  - An array (same dimensions as the tiles) of vectors of indices
    into patches indicating which patches are affected by any pixels
    in the tiles.
"""
function get_sources_per_tile(image_tiles::Matrix{ImageTile},
                              patch_ctrs::Vector{Vector{Float64}},
                              patch_radii_px::Vector{Float64})
    out = similar(image_tiles, Vector{Int})

    HH, WW = size(image_tiles)
    for ww in 1:WW, hh in 1:HH
        cands = local_source_candidates(image_tiles[hh, ww],
                                        patch_ctrs,
                                        patch_radii_px)
        # get indicies in cands that truly overlap the tile.
        idx = get_local_sources(image_tiles[hh, ww],
                                patch_ctrs[cands],
                                patch_radii_px[cands])
        out[hh, ww] = cands[idx]
    end

    return out
end

function choose_patch_radius(ce::CatalogEntry,
                             b::Int64,
                             psf_width::AbstractFloat,
                             epsilon::AbstractFloat;
                             width_scale=1.0)
    # The galaxy scale is the point with half the light -- if the light
    # were entirely in a univariate normal, this would be at 0.67 standard
    # deviations.  We are being a bit conservative here.
    obj_width =
      ce.is_star ? psf_width: width_scale * ce.gal_scale / 0.67 + psf_width

    flux = ce.is_star ? ce.star_fluxes[b] : ce.gal_fluxes[b]
    @assert flux > 0.

    # Choose enough pixels that the light is either 90% of the light
    # would be captured from a 1d gaussian or 5% of the sky noise,
    # whichever is a larger radius.
    pdf_90 = exp(-0.5 * (1.64)^2) / (sqrt(2pi) * obj_width)
    pdf_target = min(pdf_90, epsilon / (20 * flux))
    rhs = log(pdf_target) + 0.5 * log(2pi) + log(obj_width)
    sqrt(-2 * (obj_width ^ 2) * rhs)
end

"""
Choose a reasonable patch radius based on the catalog.
TODO: Select this by rendering the object and solving an optimization
problem.

Args:
  - pixel_center: The pixel location of the object
  - ce: The catalog entry for the object
  - psf: The psf at this location
  - img: The image
  - width_scale: A multiple of standard deviations to use
  - max_radius: The maximum radius in pixels.

Returns:
  - A radius in pixels chosen from the catalog entry.
"""
function choose_patch_radius(
            pixel_center::Vector{Float64},
            ce::CatalogEntry,
            psf::Array{PsfComponent},
            img::TiledImage;
            width_scale=1.0,
            max_radius=100)
    psf_width = Model.get_psf_width(psf, width_scale=width_scale)

    # Get the average sky noise around a source
    close_tile = get_containing_tile(pixel_center, img)
    epsilon = mean(close_tile.epsilon_mat)

    radius_req = choose_patch_radius(ce, img.b, psf_width, epsilon;
                                        width_scale=width_scale)
    min(radius_req, max_radius)
end


using HDF5: attributes, create_group, h5open
using HDF5: File as h5file
using NeXLCore

"""
   Base.write(h5::HDF5.File, path::AbstractString, mats::AbstractArray{<:Material})

Write an AbstractArray{<:Material} to an HDF5 file.
"""
function Base.write(h5::h5file, path::AbstractString, mats::AbstractArray{<:Material})
    els=sort(collect(mapreduce(m->keys(m), union, mats)))
    data = zeros(Float64, length(els), size(mats)...)
    for (i, el) in enumerate(els), ci in CartesianIndices(mats)
        data[i, ci] = NeXLUncertainties.value(mats[ci][el])
    end
    g = create_group(h5, path)
    attributes(g)["elements"] = map(el->symbol(el), els)
    g["massfractions", compress=3] = data
    mats
end

"""
   Base.read(h5::HDF5.File, path::AbstractString, ::Type{Material})

Read an AbstractArray{<:Material} from an HDF5 file.
"""
function Base.read(h5::h5file, path::AbstractString, ::Type{Material})
    @assert haskey(h5, path)
    g = h5[path]
    els = map(el->parse(Element,el), read(attributes(g)["elements"]))
    tmp = g["massfractions"]
    data = zeros(eltype(tmp), size(tmp)...)
    copy!(data, tmp[axes(tmp)...])
    tmp = nothing
    @assert length(els) == size(data,1) "The number of elements must match the depth of the mass-fractions data."
    return map(CartesianIndices(size(data)[2:end])) do ci
        m = data[:, ci]
        material("M$(ci.I)", Dict(els[i] => m[i] for i in filter(j -> m[j] > 0, eachindex(m))))
    end
end


# compute the full right side of the BSEs for all channels
function compute_Γ!(
    Λ      :: Float64,
    r      :: Reduced_lattice,
    m      :: Mesh,
    a1     :: Action,
    a2     :: Action,
    tbuffs :: Vector{NTuple{3, Matrix{Float64}}},
    temps  :: Vector{Array{Float64, 3}},
    corrs  :: Array{Float64, 3},
    eval   :: Int64,
    Γ_tol  :: NTuple{2, Float64}
    )      :: Nothing

    @sync begin
        for w1 in 1 : m.num_Ω
            # compute q1
            Threads.@spawn begin
                compute_channel_s_BSE!(Λ, 1, w1, -1, -1, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                compute_channel_t_BSE!(Λ, 1, w1, -1, -1, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                compute_channel_u_BSE!(Λ, 1, w1, -1, -1, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
            end

            for w3 in 1 : m.num_ν
                # compute q2_1
                Threads.@spawn begin
                    compute_channel_s_BSE!(Λ, 2, w1, w3, -1, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_t_BSE!(Λ, 2, w1, w3, -1, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_u_BSE!(Λ, 2, w1, w3, -1, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                end

                # compute q2_2
                Threads.@spawn begin
                    compute_channel_s_BSE!(Λ, 3, w1, -1, w3, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_t_BSE!(Λ, 3, w1, -1, w3, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_u_BSE!(Λ, 3, w1, -1, w3, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                end

                for w2 in w3 : m.num_ν
                    # compute q3
                    Threads.@spawn begin
                        compute_channel_s_BSE!(Λ, 4, w1, w2, w3, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_t_BSE!(Λ, 4, w1, w2, w3, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_u_BSE!(Λ, 4, w1, w2, w3, r, m, a1, a2, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    end
                end
            end
        end
    end

    symmetrize!(r, a2)

    return nothing
end

# compute the full right side of the Katanin truncated flow equations for all channels
function compute_dΓ_1l!(
    Λ      :: Float64,
    r      :: Reduced_lattice,
    m      :: Mesh,
    a      :: Action,
    da     :: Action,
    tbuffs :: Vector{NTuple{3, Matrix{Float64}}},
    temps  :: Vector{Array{Float64, 3}},
    corrs  :: Array{Float64, 3},
    eval   :: Int64,
    Γ_tol  :: NTuple{2, Float64}
    )      :: Nothing

    # compute boundary corrections
    compute_corrs_kat!(Λ, m, a, da, corrs, Γ_tol)

    @sync begin
        for w1 in 1 : m.num_Ω
            # compute q1
            Threads.@spawn begin
                compute_channel_s_kat!(Λ, 1, w1, -1, -1, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                compute_channel_t_kat!(Λ, 1, w1, -1, -1, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                compute_channel_u_kat!(Λ, 1, w1, -1, -1, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
            end

            for w3 in 1 : m.num_ν
                # compute q2_1
                Threads.@spawn begin
                    compute_channel_s_kat!(Λ, 2, w1, w3, -1, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_t_kat!(Λ, 2, w1, w3, -1, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_u_kat!(Λ, 2, w1, w3, -1, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                end

                # compute q2_2
                Threads.@spawn begin
                    compute_channel_s_kat!(Λ, 3, w1, -1, w3, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_t_kat!(Λ, 3, w1, -1, w3, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_u_kat!(Λ, 3, w1, -1, w3, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                end

                for w2 in w3 : m.num_ν
                    # compute q3
                    Threads.@spawn begin
                        compute_channel_s_kat!(Λ, 4, w1, w2, w3, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_t_kat!(Λ, 4, w1, w2, w3, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_u_kat!(Λ, 4, w1, w2, w3, r, m, a, da, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    end
                end
            end
        end
    end

    symmetrize!(r, da)

    return nothing
end

# compute the full right side of the two loop truncated flow equations for all channels
function compute_dΓ_2l!(
    Λ      :: Float64,
    r      :: Reduced_lattice,
    m      :: Mesh,
    a      :: Action,
    da     :: Action,
    da_l   :: Action,
    tbuffs :: Vector{NTuple{3, Matrix{Float64}}},
    temps  :: Vector{Array{Float64, 3}},
    corrs  :: Array{Float64, 3},
    eval   :: Int64,
    Γ_tol  :: NTuple{2, Float64}
    )      :: Nothing

    # compute one loop
    compute_dΓ_1l!(Λ, r, m, a, da, tbuffs, temps, corrs, eval, Γ_tol)

    # compute boundary corrections
    compute_corrs!(Λ, m, a, corrs, Γ_tol)

    @sync begin
        for w1 in 1 : m.num_Ω
            # compute q1
            Threads.@spawn begin
                compute_channel_s_left!(Λ, 1, w1, -1, -1, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                compute_channel_t_left!(Λ, 1, w1, -1, -1, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                compute_channel_u_left!(Λ, 1, w1, -1, -1, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
            end

            for w3 in 1 : m.num_ν
                # compute q2_1
                Threads.@spawn begin
                    compute_channel_s_left!(Λ, 2, w1, w3, -1, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_t_left!(Λ, 2, w1, w3, -1, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_u_left!(Λ, 2, w1, w3, -1, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                end

                # compute q2_2
                Threads.@spawn begin
                    compute_channel_s_left!(Λ, 3, w1, -1, w3, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_t_left!(Λ, 3, w1, -1, w3, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_u_left!(Λ, 3, w1, -1, w3, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                end

                for w2 in 1 : m.num_ν
                    # compute q3
                    Threads.@spawn begin
                        compute_channel_s_left!(Λ, 4, w1, w2, w3, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_t_left!(Λ, 4, w1, w2, w3, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_u_left!(Λ, 4, w1, w2, w3, r, m, a, da, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    end
                end
            end
        end
    end

    symmetrize_add_to!(r, da_l, da)

    return nothing
end

# compute the full right side of the multiloop truncated flow equations for all channels
function compute_dΓ_ml!(
    Λ       :: Float64,
    r       :: Reduced_lattice,
    m       :: Mesh,
    loops   :: Int64,
    a       :: Action,
    da      :: Action,
    da_l    :: Action,
    da_c    :: Action,
    da_temp :: Action,
    da_Σ    :: Action,
    tbuffs  :: Vector{NTuple{3, Matrix{Float64}}},
    temps   :: Vector{Array{Float64, 3}},
    corrs   :: Array{Float64, 3},
    eval    :: Int64,
    Γ_tol   :: NTuple{2, Float64}
    )       :: Nothing

    # compute two loop
    compute_dΓ_2l!(Λ, r, m, a, da, da_l, tbuffs, temps, corrs, eval, Γ_tol)

    # update temporary buffer and reset terms for self energy corrections
    reset_Γ!(da_temp)
    reset_Γ!(da_Σ)
    symmetrize_add_to!(r, da_l, da_temp)

    for loop in 3 : loops
        @sync begin
            for w1 in 1 : m.num_Ω
                # compute q1
                Threads.@spawn begin
                    compute_channel_s_central!(Λ, 1, w1, -1, -1, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_t_central!(Λ, 1, w1, -1, -1, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_u_central!(Λ, 1, w1, -1, -1, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                end

                for w3 in 1 : m.num_ν
                     # compute q2_1
                    Threads.@spawn begin
                        compute_channel_s_central!(Λ, 2, w1, w3, -1, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_t_central!(Λ, 2, w1, w3, -1, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_u_central!(Λ, 2, w1, w3, -1, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    end

                     # compute q2_2
                    Threads.@spawn begin
                        compute_channel_s_central!(Λ, 3, w1, -1, w3, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_t_central!(Λ, 3, w1, -1, w3, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_u_central!(Λ, 3, w1, -1, w3, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    end

                    for w2 in w3 : m.num_ν
                         # compute q3
                        Threads.@spawn begin
                            compute_channel_s_central!(Λ, 4, w1, w2, w3, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                            compute_channel_t_central!(Λ, 4, w1, w2, w3, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                            compute_channel_u_central!(Λ, 4, w1, w2, w3, r, m, a, da_l, da_c, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        end
                    end
                end
            end
        end

        @sync begin
            for w1 in 1 : m.num_Ω
                # compute q1
                Threads.@spawn begin
                    compute_channel_s_left!(Λ, 1, w1, -1, -1, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_t_left!(Λ, 1, w1, -1, -1, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    compute_channel_u_left!(Λ, 1, w1, -1, -1, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                end

                for w3 in 1 : m.num_ν
                    # compute q2_1
                    Threads.@spawn begin
                        compute_channel_s_left!(Λ, 2, w1, w3, -1, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_t_left!(Λ, 2, w1, w3, -1, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_u_left!(Λ, 2, w1, w3, -1, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    end

                    # compute q2_2
                    Threads.@spawn begin
                        compute_channel_s_left!(Λ, 3, w1, -1, w3, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_t_left!(Λ, 3, w1, -1, w3, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        compute_channel_u_left!(Λ, 3, w1, -1, w3, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                    end

                    # compute q3
                    for w2 in 1 : m.num_ν
                        Threads.@spawn begin
                            compute_channel_s_left!(Λ, 4, w1, w2, w3, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                            compute_channel_t_left!(Λ, 4, w1, w2, w3, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                            compute_channel_u_left!(Λ, 4, w1, w2, w3, r, m, a, da_temp, da_l, tbuffs[Threads.threadid()], temps[Threads.threadid()], corrs, eval, Γ_tol)
                        end
                    end
                end
            end
        end

        # update temporary buffer
        symmetrize!(r, da_c)
        replace_with_Γ!(da_temp, da_c)
        symmetrize_add_to!(r, da_l, da_temp)

        # update self energy corrections and flow
        add_to_Γ!(da_c, da_Σ)
        add_to_Γ!(da_temp, da)
    end

    return nothing
end


#SinglePT
    for op in [:helmholtz, :gibbs, :internal_energy, :enthalpy,:cp,:cv,:volume,:entropy]
        mol_op_impl = Symbol(:mol_,op,:_impl)
        mass_op_impl = Symbol(:mass_,op,:_impl)
        total_op_impl = Symbol(:total_,op,:_impl)
        mol_op = Symbol(:mol_,op)
        mass_op = Symbol(:mass_,op)
        total_op = Symbol(:total_,op)
        if op == :volume
            _unit = u"m^3/kg"
            mol_unit = u"m^3/mol"
            total_unit = u"m^3"

        elseif op in (:cv,:cp,:entropy)
            _unit = u"J/(kg*K)"
            mol_unit = u"J/(mol*K)"
            total_unit = u"J/(K)"
        else
            _unit = u"J/(kg)"
            mol_unit = u"J/(mol)"
            total_unit = u"J"
        end

        @eval begin

            #dispatch basic mass impl to the correct region, P,T
            function $mass_op_impl(mt::SinglePT,model::IndustrialWater,p,t)
                id = region_id(mt,model,p,t)
                return $mass_op_impl(mt,IF97Region{id}(),p,t)
            end

            function $mass_op(model::IndustrialWater,st::ThermodynamicState,unit=$_unit)
                return $mass_op(state_type(st),model,st,unit)
            end
            # P T impl
            function $mass_op(mt::SinglePT,model::IndustrialWater,st::ThermodynamicState,unit)
                p = pressure(FromState(),st)
                t = temperature(FromState(),st)
                res = $mass_op_impl(mt,model,p,t)
                return convert_unit($_unit,unit,res)
            end

            #mol op
            function $mol_op(model::IndustrialWater,st::ThermodynamicState,unit=$mol_unit)
                prod = molar_mass(FromState(),st,u"kg/mol",molecular_weight(model))
                res =  $mass_op(state_type(st),model,st,$_unit)*prod
                return convert_unit($mol_unit,unit,res)
            end
        end

        if !(op in (:cv,:cp))
            #total ops
            @eval begin
                function $total_op(model::IndustrialWater,st::ThermodynamicState,unit=$total_unit)
                    prod = mass(FromState(),st,u"kg",molecular_weight(model))
                    res =  $mass_op(mt,model,st,unit)*prod
                    return convert_unit($total_unit,unit,res)
                end
            end
        end

    if op != :enthalpy
        @eval begin
            #if not enthalpy, define PH impl
            function $mass_op_impl(mt::SinglePH,model::IndustrialWater,p,h)
                id = region_id(mt,model,p,h)
                return $mass_op_impl(mt,IF97Region{id}(),p,h)
            end

            function $mass_op(mt::SinglePH,model::IndustrialWater,st::ThermodynamicState,unit)
                p = pressure(FromState(),st)
                h = mass_enthalpy(FromState(),st)
                res = $mass_op_impl(mt,model,p,h)
                return convert_unit($_unit,unit,res)
            end
        end
    end

    if op != :entropy
        @eval begin
            #if not entropy, define PS impl
            function $mass_op_impl(mt::SinglePS,model::IndustrialWater,p,s)
                id = region_id(mt,model,p,h)
                t = temperature_impl(mt,IF97Region{id}(),p,s) #transform to SinglePT
                _mt = QuickStates.pt()
                return $mass_op_impl(_mt,model,p,t)
            end

            function $mass_op(mt::SinglePS,model::IndustrialWater,st::ThermodynamicState,unit)
                p = pressure(FromState(),st)
                s = mass_entropy(FromState(),st)
                res = $mass_op_impl(mt,model,p,s)
                return convert_unit($_unit,unit,res)
            end
        end
    end
end

function sound_speed(model::IndustrialWater,st::ThermodynamicState,unit=u"m/s")
    return sound_speed(state_type(st),model,st,unit)
end

function sound_speed(mt::SinglePT,model::IndustrialWater,st::ThermodynamicState,unit)
    p = pressure(FromState(),st)
    t = temperature(FromState(),st)
    id = region_id(mt,model,p,t)
    res = sound_speed_impl(mt,IF97Region{id}(),p,t)
    return convert_unit(u"m/s",unit,res)
end

function sound_speed(mt::SinglePH,model::IndustrialWater,st::ThermodynamicState,unit)
    p = pressure(FromState(),st)
    h = mass_enthalpy(FromState(),st)
    id = region_id(mt,model,p,t)
    res = sound_speed_impl(mt,IF97Region{id}(),p,h)
    return convert_unit(u"m/s",unit,res)
end

    #==
for op in (:helmholtz, :gibbs, :internal_energy, :enthalpy)
    mol_op_impl = Symbol(:mol_,op,:_impl)
    mass_op_impl = Symbol(:mass_,op,:_impl)
    total_op_impl = Symbol(:total_,op,:_impl)
    mol_op = Symbol(:mol_,op)
    mass_op = Symbol(:mass_,op)
    total_op = Symbol(:total_,op)
    for spec in [:SinglePT,:SinglePH,:SinglePS]
    @eval begin
        function $mass_op_impl(mt,model::HelmholtzModel,p,t)
            id = region_id(mt,model,p,t)
            return $mass_op_impl()
        end

        function ThermoState.$mol_op(model::HelmholtzModel,st::ThermodynamicState,unit=u"J/mol")
            return $mol_op(state_type(st),model,st,unit)
        end

        function $mol_op(mt::SingleVT,model::HelmholtzModel,st::ThermodynamicState,unit = u"J/mol")
            mw = molecular_weight(model)
            v = mol_volume(FromState(),st,u"m^3/mol",mw)
            t = temperature(FromState(),st)
            val = $mol_op_impl(mt,model,v,t)
            return convert_unit(u"J/mol",unit,val)
        end

    end
end

==#
function mass_density(model::IndustrialWater,st::ThermodynamicState,unit=u"kg/(m^3)")
    return mass_density(state_type(st),model,st,unit)
end

function mol_density(model::IndustrialWater,st::ThermodynamicState,unit=u"mol/(m^3)")
    return mol_density(state_type(st),model,st,unit)
end

function mass_density(mt,model::IndustrialWater,st::ThermodynamicState,unit)
    res = mass_volume(model,st)
    rho = one(res)/res
    return convert_unit(u"kg/(m^3)",unit,rho)
end

function mol_density(mt,model::IndustrialWater,st::ThermodynamicState,unit)
    res = mol_volume(model,st)
    rho = one(res)/res
    return convert_unit(u"mol/(m^3)",unit,rho)
end


#######
# Here we store the parameters corresponding to the non-SQUID-like geometry (`mask = false`)
######

p = Params(scale = 10, Δ=.0, Ln=200, Ls=0, W = 320, U = 0.0005, μSL = .8, μSR=.8,
    μN = 0*-0.0034, V_perp =  0., V_par = 1e-6, B_par = 0, B_perp = 0., Vu = 0.,
    edges = :armchair_zigzag, smooth = false, dir = :x, U_dimer = 0.0, α = 0., smoothμ = false,
    localselfenergy = false, analyticalselfenergy = false, normal=true);

p_λarm  = reconstruct(p, λ = 0.03, edges = :armchair_zigzag) # edge regime
p_λzig  = reconstruct(p, λ = 0.03, edges = :zigzag_armchair) # edge regime

p_λarm_gap  = reconstruct(p, λ = 0.03, α = 0.02, V_par = 0.005,
    edges = :armchair_zigzag) # gapped regime
p_λzig_gap  = reconstruct(p, λ = 0.03, α = 0.02, V_par = 0.005,
    edges = :zigzag_armchair) # gapped regime

p_λarm_bulk = reconstruct(p, λ = 0.03, α = 0.02, V_par = 0.005,
    edges = :armchair_zigzag, μN = .1) # bulk regime
p_λzig_bulk = reconstruct(p, λ = 0.03, α = 0.02, V_par = 0.005,
    edges = :zigzag_armchair, μN = .1) # bulk regime

p_λarm_gap_ls  = reconstruct(p_λarm_gap, μSL = 0.00001, μSR = 0.00001,
    Ls = 100, Δ = 1, normal = false) # gapped regime
p_λarm_bulk_ls  = reconstruct(p_λarm_bulk, μSL = 8, μSR = 8,
    Ls = 100, Ln = 0, Δ = 1, normal = false) # bulk regime
p_λarm_ls  = reconstruct(p_λarm, μSL = 0.00001, μSR = 0.00001,
    Ls = 100, Δ = 1, normal = false) # edge regime
pa_inv_smoothμ = reconstruct(p, smoothμ = true, λ = 0.02,
    Ls = 0, Ln = 200, W = 200, scale = 20, Δxμ = 5, λμ = 5, λμbarrier = 2, Δμbarrier = -2)

p_test2 = reconstruct(p_λarm_gap_ls, scale = 1, Ls = 1, Ln = 2, W = 5)
p_test4 = reconstruct(p_test2, scale = 1, Ls = 2, Ln = 4, W = 20, μN = .5);
p_test5 = reconstruct(p_test2, scale = 1, Ls = 4, Ln = 20, W = 20, μN = .5);

pa_inv_smoothμ_ls = reconstruct(pa_inv_smoothμ, Ls = 50, μSL = .8,
     μSR = .8, Δ = 0.004, normal = false)
pa_inv_smoothμ_ls_trivgap = reconstruct(pa_inv_smoothμ, smoothμ = false, U = 0.04)
pa_inv_smoothμ_ls_bulkconduction =
    reconstruct(pa_inv_smoothμ_ls, smoothμ = false, μN = 0.015)

####### Flux to B conversion
Φ0 = 2067.8338 # T*nm^2
function Bprange(p, Φrange)
    Bperp = collect(Φrange./(p.Ln * p.W)) #conversion from flux to Bfield (size dependent)
    Bcriteria(Bperp, p.scale)
    return Bperp
end

"We need to make sure that we are not in the Hofstadter regime,
    for the scaled and non scaled system"
function Bcriteria(B::Array, scale; threshold = 50.)
    ħc = ustrip(u"eV*nm",ħ*c_0)
    a0 = 0.246    #nm
    Bmax = findmax(abs.(B))[1]
    l_B = ħc/Bmax
    l_scale = scale * a0
    println(l_B/ l_scale)
    println(l_B/a0)
    l_B/ l_scale > threshold ? nothing : @warn("lB ≈ a0*scale, decrease scaling")
    l_B/ a0 > threshold ? nothing : @warn("lB ≈ a0, increase BLG area
        (we are in the Hofstadter regime)")
end


using Metalhead, Test
using Flux

@testset "MLPMixer" begin
    @test size(MLPMixer()(rand(Float32, 256, 256, 3, 2))) == (1000, 2)
    @test_skip gradtest(MLPMixer(), rand(Float32, 256, 256, 3, 2))
end


struct SemiconductorPoint{T}
    Data::SemiconductorData
    p::T
end

SemiconductorPoint1D(Data::SemiconductorData,p::Real=0.0) = SemiconductorPoint(Data,p)

function SemiconductorPoint1D(Data::SemiconductorData,p::Real,ψ,φfₑ=0.0*Eunit,φfₚ=0.0*Eunit)
    D = SemiconductorData(Data,ψ,φfₑ,φfₚ)
    SemiconductorPoint1D(D,p)
end

function SemiconductorPoint1D(S1D::SemiconductorPoint,p)
    SemiconductorPoint1D(S1D.Data,p)
end

Eg(S::SemiconductorPoint) = Eg(S.Data)
Ec(S::SemiconductorPoint) = Ec(S.Data)
Ev(S::SemiconductorPoint) = Ev(S.Data)
Ef(S::SemiconductorPoint) = Ef(S.Data)
Ei(S::SemiconductorPoint) = Ei(S.Data)
#= n(S::SemiconductorPoint) = n(S.Data)
p(S::SemiconductorPoint) = p(S.Data) =#
n(S::SemiconductorPoint,ψ=0.0*Eunit,φfₑ=0.0*Eunit) = n(S.Data,ψ,φfₑ)
p(S::SemiconductorPoint,ψ=0.0*Eunit,φfₚ=0.0*Eunit) = p(S.Data,ψ,φfₚ)
Na(S::SemiconductorPoint) = Na(S.Data)
Nd(S::SemiconductorPoint) = Nd(S.Data)
N(S::SemiconductorPoint) = N(S.Data)
N(S::SemiconductorPoint,ψ) = N(S.Data,ψ)
P(S::SemiconductorPoint) = P(S.Data)
P(S::SemiconductorPoint,ψ) = P(S.Data,ψ)
Pt(S::SemiconductorPoint) = Pt(S.Data)
Nt(S::SemiconductorPoint) = Nt(S.Data)
ni(S::SemiconductorPoint) = ni(S.Data)
φf(S::SemiconductorPoint) = φf(S.Data)
r(S::SemiconductorPoint) = S.p
ψ(S::SemiconductorPoint) = S.Data.ψ
φfₑ(S::SemiconductorPoint) = S.Data.φfₑ
φfₚ(S::SemiconductorPoint) = S.Data.φfₚ
ϵᵣ(S::SemiconductorPoint) = ϵᵣ(S.Data)
ξ(S::SemiconductorPoint) = ξ(S.Data,ψ(S))
ξ(S::SemiconductorPoint,ψ) = ξ(S.Data,ψ)
ξ(S::SemiconductorPoint,ψ,ϕ) = ξ(S.Data,ψ,ϕ)
ρ(S::SemiconductorPoint) = ρ(S.Data)
ρ(S::SemiconductorPoint,ψ) = ρ(S.Data,ψ)
Δχ(S::SemiconductorPoint,δχ) = Δχ(S.Data,δχ)
χ(S::SemiconductorPoint) = χ(S.Data)
Temperature(S::SemiconductorPoint) = Temperature(S.Data)
Nv(S::SemiconductorPoint) = Nv(S.Data)
Nc(S::SemiconductorPoint) = Nc(S.Data)
DistFun(S::SemiconductorPoint,E) = DistFun(S.Data,E)
C(S::SemiconductorPoint) = C(S.Data)

function applyPsi(S::SemiconductorPoint,ψ)
    D = S.Data
    _p = S.p
    _φfₚ = φfₚ(S)
    _φfₑ = φfₑ(S)
    return SemiconductorPoint1D(D,_p,ψ,_φfₑ,_φfₚ)
end

function applyEfe(S::SemiconductorPoint,φfₑ)
    D = S.Data
    _p = S.p
    psi = ψ(S)
    _φfₚ = φfₚ(S)
    return SemiconductorPoint1D(D,_p,psi,φfₑ,_φfₚ)
end

function applyEfp(S::SemiconductorPoint,φfₚ)
    D = S.Data
    _p = S.p
    psi = ψ(S)
    _φfₑ = φfₑ(S)
    return SemiconductorPoint1D(D,_p,psi,_φfₑ,φfₚ)
end