text:
if ! ( :EMOM in names(Main) )
    include(joinpath(@__DIR__, "EMOM.jl"))
end

if ! ( :DataManager in names(Main) )
    include(joinpath(@__DIR__, "..", "..", "share", "DataManager.jl"))
end

if ! ( :DataManager in names(Main) )
    include(joinpath(@__DIR__, "..", "..", "share", "ModelClockSystem.jl"))
end

if ! ( :DataManager in names(Main) )
    include(joinpath(@__DIR__, "..", "..", "share", "LogSystem.jl"))
end

if ! ( :Parallelization in names(Main) )
    include(joinpath(@__DIR__, "..", "..", "share", "Parallelization.jl"))
end

module ENGINE_EMOM

    using MPI
    using Dates
    using Formatting
    using NCDatasets
    using JSON

    using ..EMOM
    using ..DataManager
    using ..Parallization
    using ..ModelClockSystem
    using ..LogSystem

    include(joinpath(@__DIR__, "..", "..", "share", "AppendLine.jl"))

    name = "EMOM"

    mutable struct EMOM_DATA
        casename    :: AbstractString
        mb          :: EMOM.ModelBlock
        clock       :: ModelClock

        x2o         :: Union{Dict, Nothing}
        o2x         :: Union{Dict, Nothing}

        config     :: Dict

        recorders   :: Union{Dict, Nothing}

        jdi        :: JobDistributionInfo

        sync_data   :: Dict
    end

    function init(;
        casename     :: AbstractString,
        clock        :: ModelClock,
        config      :: Dict,
        read_restart :: Bool,
    )

        comm = MPI.COMM_WORLD
        rank = MPI.Comm_rank(comm)
        comm_size = MPI.Comm_size(comm)

        if comm_size == 1
            throw(ErrorException("Need at least 2 processes to work"))
        end

        is_master = ( rank == 0 )

        local master_mb = nothing
        local master_ev = nothing

        if is_master
            archive_list_file = joinpath(config["DRIVER"]["caserun"], config["DRIVER"]["archive_list"])
            if isfile(archive_list_file)
                writeLog("File {:s} already exists. Remove it.", archive_list_file)
                rm(archive_list_file)
            end
        end

        if is_master

            cfg_desc = EMOM.getConfigDescriptor()

            misc_config = EMOM.validateConfigEntries(config["MODEL_MISC"], cfg_desc["MODEL_MISC"])
            core_config = EMOM.validateConfigEntries(config["MODEL_CORE"], cfg_desc["MODEL_CORE"])

            # If `read_restart` is true then read restart file: config["rpointer_file"]
            # If not then initialize ocean with default profile if `initial_file`
            # is "nothing", with `init_file` if it is nonempty.

            if read_restart

                println("`read_restart` is on. Look for rpointer file...")

                rpointer_file = joinpath(config["DRIVER"]["caserun"], config["MODEL_MISC"]["rpointer_file"])

                if !isfile(rpointer_file)
                    throw(ErrorException(format("File {:s} does not exist!", rpointer_file)))
                end

                writeLog("Reading rpointer file {:s}", rpointer_file)

                snapshot_filename = ""
                open(rpointer_file, "r") do file
                    snapshot_filename  = chomp(readline(file))
                end

                if !isfile(snapshot_filename)
                    throw(ErrorException(format("Initial file \"{:s}\" does not exist!", snapshot_filename)))
                end

                master_mb = EMOM.loadSnapshot(snapshot_filename; overwrite_config=core_config)
            else

                init_file = misc_config["init_file"]

                if init_file == "BLANK_PROFILE"

                    writeLog("`init_file` == '$(init_file)'. Initialize an empty ocean.")
                    master_ev = EMOM.Env(core_config, verbose=is_master)
                    master_mb = EMOM.ModelBlock(master_ev; init_core=false)

                    master_mb.fi.sv[:TEMP] .= rand(size(master_mb.fi.sv[:TEMP])...)
                    master_mb.fi.sv[:SALT] .= rand(size(master_mb.fi.sv[:TEMP])...)

                elseif init_file != ""

                    println("Initial ocean with profile: ", init_file)
                    println("Initial ocean with domain file: ", core_config["domain_file"])
                    master_mb = EMOM.loadSnapshot(init_file; overwrite_config=core_config)

                else

                    throw(ErrorException("`init_file` cannot be empty string"))

                end
            end

            master_ev = master_mb.ev
        end

        if is_master
            if length(master_ev.cdata_varnames) != 0
                writeLog("The following datastream variables are used: {:s}.", join(master_ev.cdata_varnames, ", "))
            else
                writeLog("No datastream is used.")
            end
        end

        jdi = nothing
        master_ev_config = nothing
        if is_master
            jdi = JobDistributionInfo(nworkers = comm_size - 1, Ny = master_ev.Ny; overlap=3)
            master_ev_config = master_ev.config
        end
        jdi = MPI.bcast(jdi, 0, comm)
        master_ev_config = MPI.bcast(master_ev_config, 0, comm)

        # Second, create ModelBlocks based on ysplit_info
        if is_master
            my_ev = master_ev
            my_mb = master_mb
        else
            my_ev          = EMOM.Env(master_ev_config; sub_yrng = getYsplitInfoByRank(jdi, rank).pull_fr_rng)
            my_mb          = EMOM.ModelBlock(my_ev; init_core = true)
        end

        MPI.Barrier(comm)
        # Third, register all the variables.
        # Fourth, weaving MPI sending relation

        x2o = nothing
        o2x = nothing

        if is_master

            empty_arr_sT = zeros(Float64, 1, my_ev.Nx, my_ev.Ny)
            x2o = Dict(
                "SWFLX"       => my_mb.fi.SWFLX,
                "NSWFLX"      => my_mb.fi.NSWFLX,
                "TAUX_east"   => my_mb.fi.TAUX_east,
                "TAUY_north"  => my_mb.fi.TAUY_north,
                "FRWFLX"      => copy(empty_arr_sT),
                "VSFLX"       => my_mb.fi.VSFLX,
            )

            o2x = Dict(
                "SST"         => my_mb.fi.sv[:SST],
                "Q_FRZMLTPOT" => my_mb.fi.Q_FRZMLTPOT,
                "mask"        => my_mb.ev.topo.sfcmask_sT,
            )

        end
        # Synchronizing Data
        sync_data_list = Dict(

            # These forcings are syned from
            # master to slave before each
            # timestep starts.
            :forcing => (
                "SWFLX",
                "NSWFLX",
                "TAUX_east",
                "TAUY_north",
            ),

            # These states will be synced from
            # slave to master right after
            # completing each timestep
            :output_state   => (
                "TEMP",
                "SALT",
                "UVEL",
                "VVEL",
                "WVEL",
                "CHKTEMP",
                "CHKSALT",
                "TAUX",
                "TAUY",
                "ADVT",
                "ADVS",
                "VDIFFT",
                "VDIFFS",
                "WKRSTT",
                "WKRSTS",
                "HMXL",
                "Q_FRZMLTPOT",
                "Q_FRZMLTPOT_NEG",
                "Q_FRZHEAT",
                "Q_LOST",
                "QFLXT",
                "QFLXS",
                "Ks_H_U",
                "Ks_H_V",
            ),

            # These states are synced from
            # master to slave after :output_state
            # are syned
            :thermo_state   => (
                "TEMP",
                "SALT",
            ),

            :qflx_direct   => (
                "TEMP",
                "SALT",
                "HMXL",
            ),

            # These states are syned from
            # slave to master and master to slabe
            # during the substeps.
            :intm_state => (
                "INTMTEMP",
                "INTMSALT",
            )

        )

        sync_data = Dict()
        for (k, l) in sync_data_list
            sync_data[k] = Array{DataUnit, 1}(undef, length(l))
            for (n, varname) in enumerate(l)
                sync_data[k][n] = my_mb.dt.data_units[varname]
            end
        end

        MD = EMOM_DATA(
            casename,
            my_mb,
            clock,
            x2o,
            o2x,
            config,
            nothing,
            jdi,
            sync_data,
        )

        if is_master

            activated_record = []
            MD.recorders = Dict()
            complete_variable_list = EMOM.getDynamicVariableList(my_mb; varsets=["ALL",])

#            additional_variable_list = EMOM.getVariableList(ocn, :COORDINATE)

            for rec_key in ["daily_record", "monthly_record"]

                activated = false

                println("# For record key: " * string(rec_key))

                varnames = Array{String}(undef, 0)
                varsets  = Array{String}(undef, 0)

                for v in misc_config[rec_key]
                    m = match(r"\{(?<varset>[a-zA-Z_]+)\}", v)
                    if m != nothing
                        push!(varsets, m[:varset])
                    else
                        push!(varnames, v)
                    end
                end

                #println("varnames: ", varnames)
                rec_varnames = EMOM.getDynamicVariableList(my_mb; varnames=varnames, varsets=varsets) |> keys |> collect
                println("Record varnames: ", rec_varnames)

                add_varnames = EMOM.getCompleteVariableList(my_mb, :static) |> keys |> collect

                #println("Additional varanames: ", add_varnames)

                #=
                if typeof(misc_config[rec_key]) <: Symbol
                    misc_config[rec_key] = EMOM.getVariableList(my_mb, misc_config[rec_key]) |> keys |> collect
                end

                # Qflux_finding mode requires certain output
                if my_ev.config[:Qflx_finding] == :on
                    append!(misc_config[rec_key], EMOM.getVariableList(my_mb, :QFLX_FINDING) |> keys )
                end

                # Load variables information as a list
                for varname in misc_config[rec_key]

                    println(format("Request output variable: {:s}", varname))
                    if haskey(complete_variable_list, varname)
                        println(format("Using varaible: {:s}", varname))
                        push!(var_list, varname)#, complete_variable_list[varname]... ) )
                    else
                        throw(ErrorException("Unknown varname in " * string(rec_key) * ": " * varname))
                    end
                end
                =#

                if length(rec_varnames) != 0

                    MD.recorders[rec_key] = Recorder(
                        my_mb.dt,
                        rec_varnames,
                        EMOM.var_desc;
                        other_varnames=add_varnames,
                    )

                    push!(activated_record, rec_key)
                end
            end

            # Record (not output) happens AFTER the simulation.
            # Output of the current simulation happens at the
            # BEGINNING of the next simulation.
            #
            # Reason 1:
            # CESM does not simulate the first day of a `continue` run.
            # The first day has been simulated which is the last day of
            # the last run which is exactly the restart file. This is
            # also why we have to call archive_record! function in the
            # end of initialization.
            #
            # Reason 2:
            # Output happens at the beginning the next simulation. By
            # doing this we can get rid of the problem of deciding which
            # day is the end of month.
            #
            # This is also the way CAM chooses to do detect the end of
            # current month.
            # See: http://www.cesm.ucar.edu/models/cesm1.0/cesm/cesmBbrowser/html_code/cam/time_manager.F90.html
            #      is_end_curr_month
            #

            # Must create the record file first because the
            # run of the first day is not called in CESM
            if misc_config["enable_archive"]

                function begOfNextMonth(t::AbstractCFDateTime)
                    y = Dates.year(t)
                    m = Dates.month(t)
                    return typeof(t)(y, m, 1) + Month(1)
                end

                if "daily_record" in activated_record
                    recorder_day = MD.recorders["daily_record"]
                    addAlarm!(
                        clock,
                        "[Daily] Create daily output file.",
                        clock.time,
                        2;
                        callback = function (clk, alm)
                            createRecordFile!(MD, "h1.day", recorder_day)
                        end,
                        recurring = begOfNextMonth,
                    )

                    addAlarm!(
                        clock,
                        "[Daily] Daily output",
                        clock.time,
                        1;
                        callback = function (clk, alm)
                            record!(recorder_day)
                            avgAndOutput!(recorder_day) # This is important
                        end,
                        recurring = Day(1),
                    )

                end

                if "monthly_record" in activated_record

                    # Design alarm such that
                    # (1) Create output file first
                    # (2) Record the initial condition
                    # (3) Record simulation after one day is stepped
                    # (4) If it is the first day of next month
                    #     (i)   avg and output data
                    #     (ii)  create next monthly file
                    #     (iii) record this step in the new file in (ii)
                    #

                    recorder_mon = MD.recorders["monthly_record"]

                    addAlarm!(
                        clock,
                        "[Monthly] Create monthly output file.",
                        clock.time, # Rings immediately
                        2;
                        callback = function (clk, alm)
                            createRecordFile!(MD, "h0.mon", recorder_mon)
                        end,
                        recurring = begOfNextMonth,
                    )

                    addAlarm!(
                        clock,
                        "[Monthly] Daily accumulation using record!",
                        clock.time, # Remember we need to record the first one. So alarm rings immediately.
                        1;
                        callback = function (clk, alm)
                            record!(recorder_mon)
                        end,
                        recurring = Day(1),
                    )

                    addAlarm!(
                        clock,
                        "[Monthly] Average and output monthly data.",
                        begOfNextMonth(clock.time), # Start from next month
                        3;  # Higher priority so it outputs data before creating next new monthly file
                        callback = function (clk, alm)
                            avgAndOutput!(recorder_mon)
                        end,
                        recurring = begOfNextMonth,
                    )

                end

            end

        end

        MPI.Barrier(comm)

        syncField!(
            MD.sync_data[:thermo_state],
            MD.jdi,
            :M2S,
            :BLOCK,
        )

        return MD

    end

    function run!(
        MD            :: EMOM_DATA;
        Δt            :: Second,
        write_restart :: Bool,
    )

        comm = MPI.COMM_WORLD
        rank = MPI.Comm_rank(comm)
        is_master = rank == 0

        Δt_float = Float64(Δt.value)

        substeps = MD.mb.ev.config["substeps"]
        Δt_substep = Δt_float / substeps

        syncField!(
            MD.sync_data[:forcing],
            MD.jdi,
            :M2S,
            :BLOCK,
        )

        if ! is_master

            EMOM.reset!(MD.mb.co.wksp)
            EMOM.updateDatastream!(MD.mb, MD.clock)
            EMOM.checkBudget!(MD.mb, Δt_float; stage=:BEFORE_STEPPING)

            Δz_min = minimum(view(MD.mb.ev.gd.Δz_T, :, 1, 1))
            EMOM.setupForcing!(MD.mb; w_max = Δz_min / Δt_substep * 0.9)

        end

        if ! is_master
            # Because substep in stepAdvection updates _INTMX_
            # based on _INTMX_ itself, we need to copy it first.
            MD.mb.tmpfi._INTMX_ .= MD.mb.fi._X_
        end
        for substep = 1:substeps

            if ! is_master
                EMOM.reset!(MD.mb.co.wksp)
                EMOM.stepAdvection!(MD.mb, Δt_substep)
                EMOM.checkBudget!(MD.mb, Δt_float; stage=:SUBSTEP_AFTER_ADV, substeps=substeps)
            end

            syncField!(
                MD.sync_data[:intm_state],
                MD.jdi,
                :S2M,
                :BND,
            )

            syncField!(
                MD.sync_data[:intm_state],
                MD.jdi,
                :M2S,
                :BND,
            )
        end

        if ! is_master

            EMOM.stepColumn!(MD.mb, Δt_float)
            EMOM.checkBudget!(MD.mb, Δt_float; stage=:AFTER_STEPPING)

            # important: update X
            MD.mb.fi._X_ .= MD.mb.tmpfi._NEWX_
        end

        syncField!(
            MD.sync_data[:output_state],
            MD.jdi,
            :S2M,
            :BLOCK,
        )

        syncField!(
            MD.sync_data[:thermo_state],
            MD.jdi,
            :M2S,
            :BND,
        )

        if write_restart
            writeLog("`wrtie_restart` is true.")
            if is_master
                writeRestart(MD)
            end
        end

    end

    function final(MD::EMOM_DATA)

        writeLog("Finalizing the model.")

    end

    function createRecordFile!(
        MD     :: EMOM_DATA,
        group  :: String,
        recorder :: Recorder,
    )

        t = dt2tuple(MD.clock.time)

        filename = format("{}.EMOM.{}.{:04d}-{:02d}.nc", MD.casename, group, t[1], t[2])

        setNewNCFile!(
            recorder,
            joinpath(MD.config["DRIVER"]["caserun"], filename)
        )

        appendLine(joinpath(MD.config["DRIVER"]["caserun"], MD.config["DRIVER"]["archive_list"]),
            format("mv,{:s},{:s},{:s}",
                filename,
                MD.config["DRIVER"]["caserun"],
                joinpath(MD.config["DRIVER"]["archive_root"], "ocn", "hist"),
            )
        )

    end

    function writeRestart(
        MD :: EMOM_DATA,
    )

        clock_time = MD.clock.time

        timestamp_str = format(
            "{:s}-{:05d}",
            Dates.format(clock_time, "yyyy-mm-dd"),
            floor(Int64, Dates.hour(clock_time)*3600+Dates.minute(clock_time)*60+Dates.second(clock_time)),
        )

        snapshot_filename = format(
            "{:s}.snapshot.{:s}.jld2",
            MD.config["DRIVER"]["casename"],
            timestamp_str,
        )

        EMOM.takeSnapshot(
            MD.clock.time,
            MD.mb,
            joinpath(
                MD.config["DRIVER"]["caserun"],
                snapshot_filename,
            ),
        )

        println("(Over)write restart pointer file: ", MD.config["MODEL_MISC"]["rpointer_file"])
        open(joinpath(MD.config["DRIVER"]["caserun"], MD.config["MODEL_MISC"]["rpointer_file"]), "w") do io
            write(io, snapshot_filename, "\n")
        end

        appendLine(joinpath(MD.config["DRIVER"]["caserun"], MD.config["DRIVER"]["archive_list"]),
            format("cp,{:s},{:s},{:s}",
                snapshot_filename,
                MD.config["DRIVER"]["caserun"],
                joinpath(MD.config["DRIVER"]["archive_root"], "rest", timestamp_str),
            )
        )

        appendLine(joinpath(MD.config["DRIVER"]["caserun"], MD.config["DRIVER"]["archive_list"]),
            format("cp,{:s},{:s},{:s}",
                MD.config["MODEL_MISC"]["rpointer_file"],
                MD.config["DRIVER"]["caserun"],
                joinpath(MD.config["DRIVER"]["archive_root"], "rest", timestamp_str),
            )
        )

    end

    function syncM2S!(
        MD,
    )
        syncField!(
            MD.sync_data[:qflx_direct],
            MD.jdi,
            :M2S,
            :BLOCK,
        )

    end
end

max_stars_repo_name: meteorologytoday/IOM

max_issues_repo_name: meteorologytoday/IOM

max_forks_repo_name: meteorologytoday/IOM


text:
## same for all soundingss
nsamples = 100001
nchainspersounding = 7
ppn = 48
@assert mod(ppn, nchainspersounding+1) == 0
zfixed   = [-1e5]
ρfixed   = [1e12]
zstart = 0.0
extendfrac, dz = 1.06, 1.25
nlayers = 52
ρbg = 5.
ntimesperdecade = 10
nfreqsperdecade = 5
## make transD options
nmin, nmax = 2, 40
K = transD_GP.GP.OrstUhn()
demean = false
sampledc = true
fbounds = [-1.0 3.0]
sddc = 0.01 # fraction
sdpos = 0.05 # fraction
sdprop = 0.05 # fraction
λ, δ = [2.], 0.1
save_freq = 50
Tmax = 2.50
nchainsatone = 1
vectorsum = true
useML = false
## make nuisance options
# sdev are fractions of width
nuisance_sdev = [
        0.0 #zTx
        0.005 #zRx
        0.005 #x_rx
        0.0 #y_rx
        0.0 #rx_roll
        0.005 #rx_pitch
        0.0 #rx_yaw
        0.0 #tx_roll
        0.0 #tx_pitch
        0.0 #tx_yaw
           ]
 nuisance_bounds = [
      0.   0.
     -3.0  3.0
     -3.0  3.0
      0.   0.
      0.   0.
     -1.5  1.5
      0.   0.
      0.   0.
      0.   0.
      0.   0.
          ]
updatenuisances = true
## see if options set can make a reasonable looking forward
transD_GP.TEMPEST1DInversion.makeoperatorandoptions(
        soundings;
        zfixed             = zfixed,
        ρfixed             = ρfixed,
        zstart             = zstart,
        extendfrac         = extendfrac,
        dz                 = dz,
        ρbg                = ρbg,
        nlayers            = nlayers,
        ntimesperdecade    = ntimesperdecade,
        nfreqsperdecade    = nfreqsperdecade,
        nmin               = nmin,
        nmax               = nmax,
        K                  = K,
        demean             = demean,
        sampledc           = sampledc,
        sddc               = sddc,
        sdpos              = sdpos,
        sdprop             = sdprop,
        fbounds            = fbounds,
        save_freq          = save_freq,
        λ                  = λ,
        δ                  = δ,
        nuisance_bounds    = nuisance_bounds,
        nuisance_sdev      = nuisance_sdev,
        updatenuisances    = updatenuisances,
        useML              = useML,
        vectorsum          = vectorsum)

ext: jl

max_stars_repo_name: GeoscienceAustralia/HiQGA

max_issues_repo_name: GeoscienceAustralia/HiQGA

max_forks_repo_name: GeoscienceAustralia/HiQGA


text:
#=
Find nodes for triangular elements.
Use the book.
=#
include("triangle_quadrature_table.jl");

# Build global nodes for a triangular element with refel and vertices v
function triangle_element_nodes(refel, v)
    return  triangle_refel_to_xy(refel.r[:,1], refel.r[:,2], v);
end

# Set up refel nodal array
function triangle_refel_nodes!(refel)
    (eqx, eqy) = triangle_equilateral_nodes(refel.N);
    (r, s) = triangle_equilateral_to_rs(eqx,eqy);

    refel.r = zeros(refel.Np,2);
    refel.wr = zeros(refel.Np);
    #refel.g = zeros(refel.Np,2);
    #refel.wg = zeros(refel.Np);

    refel.r[:,1] = r;
    refel.r[:,2] = s;

    # quadrature nodes/weights from a table
    xyw = triangle_quadrature_nodes_weights(refel.N+1);
    refel.g = xyw[:,1:2];
    refel.wg = xyw[:,3];
    tol = 1e-8;
    tf1(x) = abs(x[1] + 1) < tol;
    tf2(x) = abs(x[2] + 1) < tol;
    tf3(x) = abs(x[1] + x[2]) < tol;
    refel.face2local = [get_face2local_map(refel.r, tf1),
                        get_face2local_map(refel.r, tf2),
                        get_face2local_map(refel.r, tf3)];

    ### Unfinished: find correct surface quadrature weights.
    ### For now just take average(correct for order=1).
    refel.surf_r = [refel.r[refel.face2local[1],:], refel.r[refel.face2local[2],:], refel.r[refel.face2local[3],:]];
    refel.surf_g = [refel.r[refel.face2local[1],:], refel.r[refel.face2local[2],:], refel.r[refel.face2local[3],:]];

    tmp = [length(refel.face2local[1]), length(refel.face2local[2]), length(refel.face2local[3])];
    refel.surf_wr = [ones(tmp[1])./tmp[1], ones(tmp[2])./tmp[2], ones(tmp[3])./tmp[3]];
    refel.surf_wg = refel.surf_wr;
end

# Purpose  : Compute (x,y) nodes in equilateral triangle for polynomial of order N
function triangle_equilateral_nodes(N)
    # optimal alpha values for N up to 16 (from book)
    alpopt = [0.0000 0.0000 1.4152 0.1001 0.2751 0.9800 1.0999 1.2832 1.3648 1.4773 1.4959 1.5743 1.5770 1.6223 1.6258];

    # Set optimized parameter, alpha, depending on order N
    if (N<16)
        alpha = alpopt[N];
    else
        alpha = 5/3;
    end

    # total number of nodes
    Np = Int64((N+1)*(N+2)/2);

    # Create equidistributed nodes on equilateral triangle
    L1 = zeros(Np);
    L2 = zeros(Np);
    L3 = zeros(Np);
    sk = 1;
    for n=1:(N+1)
        for m=1:(N+2-n)
            L1[sk] = (n-1)/N;
            L3[sk] = (m-1)/N;
            sk = sk+1;
        end
    end
    L2 = 1 .- L1 .- L3;
    x = -L2+L3;
    y = (-L2 - L3 + 2 .* L1) .* 0.5773502691896258; # 1/sqrt(3)=0.5773502691896258

    # Compute blending function at each node for each edge
    blend1 = 4 .* L2 .* L3;
    blend2 = 4 .* L1 .* L3;
    blend3 = 4 .* L1 .* L2;

    # Amount of warp for each node, for each edge
    warpf1 = triangle_warpfactor(N,L3-L2);
    warpf2 = triangle_warpfactor(N,L1-L3);
    warpf3 = triangle_warpfactor(N,L2-L1);

    # Combine blend & warp
    warp1 = blend1 .* warpf1 .* (1 .+ (alpha*L1).^2);
    warp2 = blend2 .* warpf2 .* (1 .+ (alpha*L2).^2);
    warp3 = blend3 .* warpf3 .* (1 .+ (alpha*L3).^2);

    # Accumulate deformations associated with each edge
    # x = x + 1*warp1 + cos(2*pi/3)*warp2 + cos(4*pi/3)*warp3;
    # y = y + 0*warp1 + sin(2*pi/3)*warp2 + sin(4*pi/3)*warp3;
    x = x + warp1 - 0.5*warp2 - 0.5*warp3;
    y = y + 0.8660254037844387*warp2 - 0.8660254037844387*warp3;

    return (x, y);
end

# Compute scaled warp function at order N based on rout interpolation nodes
function triangle_warpfactor(N, rout)
    # Compute LGL and equidistant node distribution
    #(LGLr, w) = jacobi_gauss_quad(0,0,N);
    (LGLr, w) = jacobi_LGL_quad(N);

    h = 2/N;
    req = -1:h:1;

    # Compute Vandermonde based on req
    Veq = zeros(N+1, N+1);
    for i=1:N+1
        Veq[:,i] = jacobi_polynomial(req, 0, 0, i-1);
    end

    # Evaluate Lagrange polynomial at rout
    Nr = length(rout);
    Pmat = zeros(N+1,Nr);
    for i=1:N+1
      Pmat[i,:] = jacobi_polynomial(rout, 0, 0, i-1);
    end
    Lmat = Veq'\Pmat;

    # Compute warp factor
    warp = Lmat'*(LGLr - req);

    # Apply Scale factor
    for i=1:length(rout)
        if abs(rout[i]) < 0.9999999999
            warp[i] = warp[i] / (1-rout[i]*rout[i]);
        else
            warp[i] = 0;
        end
    end

    return warp;
end

# From (x,y) in equilateral triangle to (r,s) coordinates in standard triangle
function triangle_equilateral_to_rs(x,y)
    # sqrt(3)=1.7320508075688772
    L1 = (1.7320508075688772 .* y .+ 1.0) ./ 3.0;
    L2 = (-3.0 .* x .- 1.7320508075688772 .* y .+ 2.0) ./ 6.0;
    L3 = ( 3.0 .* x .- 1.7320508075688772 .* y .+ 2.0) ./ 6.0;

    r = -L2 + L3 - L1;
    s = -L2 - L3 + L1;

    return (r, s);
end

# # From (r,s) coordinates in reference triangle to (x,y) in triangle with vertices v
# # v is a 2x3 array [x1 x2 x3; y1 y2 y3]
# function triangle_refel_to_xy(r, s, v)
#     x = 0.5 .* (-(r+s) .* v[1,1] .+ (1+r) .* v[1,2] .+ (1+s) .* v[1,3]);
#     y = 0.5 .* (-(r+s) .* v[2,1] .+ (1+r) .* v[2,2] .+ (1+s) .* v[2,3]);

#     return (x, y);
# end

# function get_face2local_map(r, compare)
#     n = size(r,1);
#     map = zeros(Int,n);
#     nf = 0;
#     for i=1:n
#         if compare(r[i,:])
#             nf = nf+1;
#             map[nf] = i;
#         end
#     end

#     return map[1:nf];
# end


text:
using BenchmarkTools, Profile, Traceur, InteractiveUtils, Lasso

include("helpers.jl")

BenchmarkTools.DEFAULT_PARAMETERS.seconds = 1

kwargs = Dict(:verbose=>false)

n = 300
p = 3
covars, counts = bdata(n, p, 1000)
@assert sum(all(iszero.(counts), dims=1)) == 0 "some columns are always zero"

@info "dmr:"
m1 = @btime dmr(covars, counts; local_cluster=true, kwargs...)
m2 = @btime dmr(covars, counts; local_cluster=false, kwargs...)
m3 = @btime dmr(covars, counts; parallel=false, kwargs...)

@assert m1.coefs == m2.coefs
@assert m1.coefs == m3.coefs

@info "hdmr:"
m1 = @btime hdmr(covars, counts; local_cluster=true, kwargs...)
m2 = @btime hdmr(covars, counts; local_cluster=false, kwargs...)
m3 = @btime hdmr(covars, counts; parallel=false, kwargs...)

@which(hcat(v1, v2))

@assert m1.coefspos == m2.coefspos
@assert m1.coefszero == m2.coefszero
@assert m1.coefspos == m3.coefspos
@assert m1.coefszero == m3.coefszero
@assert m1.coefszero == m2.coefszero

@info "fit(InclusionRepetition,...)"
function minifit(counts, covars)
    n,p = size(covars)
    inpos = 1:p-1
    inzero = 1:p
    j = 1
    M = InclusionRepetition
    offsetpos = zeros(n)
    offsetzero = zeros(n)

    cj = Vector(counts[:,j])
    covarspos, covarszero = HurdleDMR.incovars(covars,inpos,inzero)
    # we use the same offsets for pos and zeros
    @code_warntype HurdleDMR.fit(M,GammaLassoPath,covarszero,cj;
        Xpos=covarspos, offsetpos=offsetpos, offsetzero=offsetzero)
    # coef(path)
end

minifit(counts, covars)

function minifit2(counts, covars)
    n,p = size(covars)
    inpos = 1:p-1
    inzero = 1:p
    j = 1
    M = InclusionRepetition
    offsetpos = zeros(n)
    offsetzero = zeros(n)

    cj = Vector(counts[:,j])
    covarspos, covarszero = HurdleDMR.incovars(covars,inpos,inzero)
    # we use the same offsets for pos and zeros
    @code_warntype HurdleDMR.fit(M,GammaLassoPath,covarszero,cj;
        Xpos=covarspos, offsetpos=offsetpos, offsetzero=offsetzero)
    # coef(path)
end
minifit2(counts, covars)

function minifitzero(counts, covars)
    n,p = size(covars)
    inpos = 1:p-1
    inzero = 1:p
    j = 1
    offsetpos = zeros(n)
    offsetzero = zeros(n)

    covarspos, covarszero = HurdleDMR.incovars(covars,inpos,inzero)

    X = covarszero
    y = Vector(counts[:,j])
    M = GammaLassoPath
    Xpos = covarspos
    dzero = Binomial()
    dpos = Poisson()
    lzero = canonicallink(dzero)
    lpos = canonicallink(dpos)
    dofit = true
    wts = fill(one(eltype(y)),size(y))
    offset = similar(y, 0)
    verbose = false
    showwarnings = false

    ixpos, Iy = HurdleDMR.getIy(y)

    offsetzero, offsetpos = HurdleDMR.setoffsets(y, ixpos, offset, offsetzero, offsetpos)

    # @code_warntype HurdleDMR.finiteoffsetobs(:zeros, X, Iy, offsetzero, wts, showwarnings)
    HurdleDMR.fitzero(M, X, Iy, dzero, lzero, dofit, wts, offsetzero, verbose, showwarnings)
end
@btime minifitzero(counts, covars)

function minifitpos(counts, covars, j)
    n,p = size(covars)
    inpos = 1:p-1
    inzero = 1:p
    offsetpos = zeros(n)
    offsetzero = zeros(n)

    covarspos, covarszero = HurdleDMR.incovars(covars,inpos,inzero)

    X = covarszero
    y = Vector(counts[:,j])
    M = GammaLassoPath
    Xpos = covarspos
    dzero = Binomial()
    dpos = Poisson()
    lzero = canonicallink(dzero)
    lpos = canonicallink(dpos)
    dofit = true
    wts = fill(one(eltype(y)),size(y))
    offset = similar(y, 0)
    verbose = false
    showwarnings = true

    ixpos, Iy = HurdleDMR.getIy(y)

    offsetzero, offsetpos = HurdleDMR.setoffsets(y, ixpos, offset, offsetzero, offsetpos)

    # Xpos optional argument allows to specify a data matrix only for positive counts
    if Xpos == nothing
      # use X for Xpos too
      Xpos = X[ixpos,:]
    elseif size(Xpos,1) == length(y)
      # Xpos has same dimensions as X, take only positive y ones
      Xpos = Xpos[ixpos,:]
    end

    ypos = y[ixpos]
    wtspos = wts[ixpos]

    HurdleDMR.excessy!(ypos, InclusionRepetition)
    # @code_warntype HurdleDMR.finiteoffsetobs(:positives, Xpos, ypos, offsetpos, wtspos, showwarnings)
    # @btime HurdleDMR.finiteoffsetobs(:positives, $Xpos, $ypos, $offsetpos, $wtspos, $showwarnings)
    X, y, offset, wts = HurdleDMR.finiteoffsetobs(:positives, Xpos, ypos, offsetpos, wtspos, showwarnings)
    if length(y) < length(ypos)
        @warn("j=$j has inf offsets")
    end
    X, y, offset, wts
end
function minifitpos2(counts, covars, j)
    n,p = size(covars)
    inpos = 1:p-1
    inzero = 1:p
    offsetpos = zeros(n)
    offsetzero = zeros(n)

    covarspos, covarszero = HurdleDMR.incovars(covars,inpos,inzero)

    X = covarszero
    y = Vector(counts[:,j])
    M = GammaLassoPath
    Xpos = covarspos
    dzero = Binomial()
    dpos = Poisson()
    lzero = canonicallink(dzero)
    lpos = canonicallink(dpos)
    dofit = true
    wts = fill(one(eltype(y)),size(y))
    offset = similar(y, 0)
    verbose = false
    showwarnings = true

    ixpos, Iy = HurdleDMR.getIy(y)

    offsetzero, offsetpos = HurdleDMR.setoffsets(y, ixpos, offset, offsetzero, offsetpos)

    # Xpos optional argument allows to specify a data matrix only for positive counts
    if Xpos == nothing
      # use X for Xpos too
      Xpos = X[ixpos,:]
    elseif size(Xpos,1) == length(y)
      # Xpos has same dimensions as X, take only positive y ones
      Xpos = Xpos[ixpos,:]
    end

    ypos = y[ixpos]
    wtspos = wts[ixpos]

    HurdleDMR.excessy!(ypos, InclusionRepetition)
    # @code_warntype HurdleDMR.finiteoffsetobs!(Xpos, ypos, offsetpos, wtspos, showwarnings, :positives)
    @btime HurdleDMR.finiteoffsetobs!($Xpos, $ypos, $offsetpos, $wtspos, $showwarnings, :positives)
    HurdleDMR.finiteoffsetobs!(Xpos, ypos, offsetpos, wtspos, showwarnings, :positives)
end
for j=1:1000
    Xpos1, ypos1, offsetpos1, wtspos1 = minifitpos(counts, covars, j)
end
Xpos2, ypos2, offsetpos2, wtspos2 = minifitpos2(counts, covars, 3)
minifitpos2(counts, covars)

@code_warntype HurdleDMR.getlogger(true)
@code_warntype HurdleDMR.getlogger(Logging.Error)
using Logging
with_logger(HurdleDMR.getlogger(true)) do
    @info "infoing!"
    @warn "warning!"
    @error "erroring!"
end

cpos1, czero1 = @btime minifit($counts, $covars)
cpos2, czero2 = @btime minifit2($counts, $covars)
@assert cpos1 == cpos2
@assert czero1 == czero2

Profile.init(delay=0.0005)
Juno.@profiler dmr(covars, counts; local_cluster=true, kwargs...)
Juno.@profiler dmr(covars, counts; local_cluster=false, kwargs...)

@trace hdmr(covars, counts; local_cluster=true, kwargs...) maxdepth=3 modules=["HurdleDMR", "Lasso"]

@trace hdmr(covars, counts; local_cluster=true, kwargs...) maxdepth=3 modules=["HurdleDMR", "Lasso"]

##############################
# memory allocation
##############################

# # first start julia with:
# julia --track-allocation=user
local_cluster = true

# Run whatever commands you wish to test. This first run is to ensure that everything is compiled (because compilation allocates memory).
hdmr(covars, counts; local_cluster=local_cluster, kwargs...)

# clear stuff contaminated by compilation
Profile.clear_malloc_data()

# Run your commands again
hdmr(covars, counts; local_cluster=local_cluster, kwargs...)

# Quit julia!
exit()

# Finally, navigate to the directory holding your source code.
# Start julia (without command-line flags), and analyze the results using

using Coverage
ma = analyze_malloc(".")  # could be "." for the current directory, or "src", etc.
ma

local_cluster = true
f = open("benchmark/malloc.$local_cluster.log", "w")
    for m in ma
        write(f, string(m), "\n")
    end
close(f)

# This will return a vector of MallocInfo objects, specifying the number of bytes allocated
# , the file name, and the line number. These are sorted in increasing order of allocation size.

# now cleanup
memfiles = Coverage.find_malloc_files(["."])
rm.(memfiles)

######################################
# focus
######################################
y = Vector(counts[:,1])
y = counts[:,1]
ixpos1, Iy1 = @btime HurdleDMR.getIy(y)
ixpos2, Iy2 = @btime HurdleDMR.getIy2(y)
@assert ixpos2 == collect(1:n)[ixpos1]
@assert Iy1 == Iy2

@code_warntype HurdleDMR.getIy(y)
@code_warntype HurdleDMR.getIy2(y)


text:
#=
Utility functions used in the QuantEcon library

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-08-14

=#
meshgrid(x::AbstractVector, y::AbstractVector) = (repmat(x, 1, length(y))',
                                                  repmat(y, 1, length(x)))

fix(x::Real) = x >= 0 ? floor(Int, x) : ceil(Int, x)
fix!{T<:Real}(x::AbstractArray{T}, out::Array{Int}) = map!(fix, out, x)
fix{T<:Real}(x::AbstractArray{T}) = fix!(x, similar(x, Int))

"""
`fix(x)`

Round `x` towards zero. For arrays there is a mutating version `fix!`
"""
fix

ckron(A::AbstractArray, B::AbstractArray) = kron(A, B)
ckron(arrays::AbstractArray...) = reduce(kron, arrays)

"""
`ckron(arrays::AbstractArray...)`

Repeatedly apply kronecker products to the arrays. Equilvalent to
`reduce(kron, arrays)`
"""
ckron

"""
`gridmake!(out::AbstractMatrix, arrays::AbstractVector...)`

Like `gridmake`, but fills a pre-populated array. `out` must have size
`prod(map(length, arrays), length(arrays))`
"""
function gridmake!(out, arrays::AbstractVector...)
    lens = Int[length(e) for e in arrays]

    n = length(arrays)
    l = prod(lens)
    @assert size(out) == (l, n)

    reverse!(lens)
    repititions = cumprod(vcat(1, lens[1:end-1]))
    reverse!(repititions)
    reverse!(lens)  # put lens back in correct order

    @inbounds for col=1:n
        row = 1
        arr = arrays[col]
        outer = repititions[col]
        inner = floor(Int, l / (outer * lens[col]))
        for _ in 1:outer
            for ix in 1:lens[col]
                v = arr[ix]
                for _ in 1:inner
                    out[row, col] = v
                    row += 1
                end
            end
        end
    end
    return out
end

function gridmake(arrays::AbstractVector...)
    l = prod([length(a) for a in  arrays])
    T = reduce(promote_type, [eltype(a) for a in arrays])
    out = Array(T, l, length(arrays))
    gridmake!(out, arrays...)
    out
end

# type stable version if all arrays have the same eltype
function gridmake{T}(arrays::AbstractVector{T}...)
    out = Array(T, prod([length(a) for a in  arrays]), length(arrays))
    gridmake!(out, arrays...)
end

function gridmake(t::Tuple)
    all(map(x -> isa(x, Integer), t)) ||
        error("gridmake(::Tuple) only valid when all elements are integers")
    gridmake(map(x->1:x, t)...)::Matrix{Int}
end

"""
`gridmake(arrays::AbstractVector...)`

Expand one or more vectors into a matrix where rows span the cartesian product
of combinations of the input vectors. Each input array will correspond to one
column of the output matrix. The first array varies the fastest (see example)

##### Example

```jlcon
julia> x = [1, 2, 3]; y = [10, 20]; z = [100, 200];

julia> gridmake(x, y, z)
12x3 Array{Int64,2}:
 1  10  100
 2  10  100
 3  10  100
 1  20  100
 2  20  100
 3  20  100
 1  10  200
 2  10  200
 3  10  200
 1  20  200
 2  20  200
 3  20  200
```
"""
gridmake


text:
using DataDeps

register(DataDep(
  "year-prediction-msd",
  "http://archive.ics.uci.edu/ml/datasets/YearPredictionMSD",
  "https://archive.ics.uci.edu/ml/machine-learning-databases/00203/YearPredictionMSD.txt.zip",
  "06f801af323bb7798e800583acce4ea1ed2697ac12c23f4424aea0a7a3d09e11",
  post_fetch_method=(path -> begin
    unpack(path)
    new_path = joinpath(dirname(path), "YearPredictionMSD.txt")
    mv(new_path, path)

    UCIData.process_dataset(path,
        target_index=1,
        feature_indices=2:91,
    )
  end),
))

max_stars_repo_name: JackDunnNZ/uci-data

max_issues_repo_name: JackDunnNZ/uci-data

max_forks_repo_name: JackDunnNZ/uci-data


text:
module ZeroProblems

using ..Lenny: AbstractContinuationProblem

#--- Forward definitions

"""
	name(x)

Return the name of a variable or zero problem as a Symbol.
"""
function name end

"""
	dependencies(x)

Return the variable dependencies of a zero problem as a Tuple.
"""
function dependencies end

#--- Variables that zero problems depend on

"""
	Var

A placeholder for a continuation variable. As a minimum it comprises a name and
length (in units of the underlying numerical type). Optionally it can include a
parent variable and offset into the parent variable (negative offsets represent
the offset from the end).

# Example

```
coll = Var(:collocation, 20)  # an array of 20 Float64 (or other numerical type)
x0 = Var(:x0, 1, coll, 0)  # a single Float64 at the start of the collocation array
x1 = Var(:x1, 1, coll, -1)  # a single Float64 at the end of the collocation array
```
"""
mutable struct Var
	name::Symbol
	len::Int
	parent::Union{Var, Nothing}
	offset::Int
end

Var(name::Symbol, len::Integer) = Var(name, len, nothing, 0)

name(u::Var) = u.name
Base.length(u::Var) = u.len

#--- ZeroSubproblem

# For simple subproblems the ZeroSubproblem type can be used to define a
# subproblem and its associated variable dependencies.

# For more complicated subproblems (or subproblems that require access to the
# full problem structure) you should inherit from AbstractZeroSubproblem.

abstract type AbstractZeroSubproblem end

name(ϕ::AbstractZeroSubproblem) = ϕ.name
dependencies(ϕ::AbstractZeroSubproblem) = ϕ.deps

struct ZeroSubproblem{F} <: AbstractZeroSubproblem
	name::Symbol
	deps::Vector{Var}
	f!::F
end

@inline (z::ZeroSubproblem)(res, prob::AbstractContinuationProblem, u...) = z.f!(res, u...)
@inline (z::ZeroSubproblem)(res, u...) = z.f!(res, u...)

#--- ZeroProblem - the full problem structure

struct ZeroProblem{D, U, Φ}
	u::U
	ui::Vector{UnitRange{Int}}
	ϕ::Φ
	ϕi::Vector{UnitRange{Int}}
	ϕdeps::Vector{Tuple{Vararg{Int, N} where N}}
end

ZeroProblem() = ZeroProblem{Nothing, Vector{Var}, Vector{Any}}(Vector{Var}(), Vector{UnitRange{Int}}(), Vector{Any}(), Vector{UnitRange{Int}}())

function Base.push!(zp::ZeroProblem{Nothing}, u::Var)
	if !(u in zp.u)
		if u.parent === nothing
			last = maximum(maximum.(zp.ui))
			push!(zp.u, u)
			push!(zp.ui, last + 1:last + length(u))
		else
			idx = findfirst(isequal(u.parent), zp.u)
			start = (u.offset < 0) ? (zp.ui[idx][end] + u.offset + 1) : (zp.ui[idx][1] + u.offset)
			push!(zp.u, u)
			push!(zp.ui, start:start + length(u) - 1)
		end
	end
	return zp
end

function Base.push!(zp::ZeroProblem{Nothing}, subprob::AbstractZeroSubproblem)
	if subprob in zp.ϕ
		throw(ArgumentError("Subproblem is already part of the zero problem"))
	end
	depidx = Vector{Int}()
	for dep in dependencies(subprob)
		push!(zp, dep)
		push!(depidx, findfirst(isequal(dep), zp.u))
	end
	last = maximum(maximum.(zp.ϕi))
	push!(zp.ϕ, ϕ)
	push!(zp.ϕi, last + 1:last + length(ϕ))
	push!(zp.ϕdeps, (depidx...,))
	return zp
end

function specialize(zp::ZeroProblem)
	u = (zp.u...,)
	ui = zp.ui
	ϕ = (zp.ϕ...,)
	ϕi = zp.ϕi
	ϕdeps = zp.ϕdeps
	return ZeroProblem{(ϕdeps...,), typeof(u), typeof(ϕ)}(u, ui, ϕ, ϕi, ϕdeps)
end

function residual!(res, zp::ZeroProblem{Nothing}, prob::AbstractContinuationProblem, u)
	# TODO: implement anyway?
	throw(ArgumentError("Specialize the zero problem before calling residual!"))
end

@generated function residual!(res, zp::ZeroProblem{D, U, Φ}, prob::AbstractContinuationProblem, u) where {D, U <: Tuple, Φ <: Tuple}
	body = quote
		# Construct views into u for each variable
		uv = ($((:(view(u, zp.ui[$i])) for i in eachindex(U.parameters))...),)
		# Construct views into res for each subproblem
		resv = ($((:(view(res, zp.ϕi[$i])) for i in eachindex(Φ.parameters))...),)
	end
	# Call each of the subproblems
	for i in eachindex(D)
        if length(D[i]) == 0
            # No dependencies means pass everything
            push!(body.args, :(zp.ϕ[$i](resv[$i], prob, u)))
        else
            push!(body.args, :(zp.ϕ[$i](resv[$i], prob, $((:(uv[$(D[i][j])]) for j in eachindex(D[i]))...))))
        end
	end
	# Return res
    push!(body.args, :res)
    @show body
    body
end

end


text:
"""
    SphereCapsuleCollision

    collision between sphere and capsule

    origin_sphere:    position of sphere contact relative to body center of mass
    origin_capsule_a: position of capsule contact a relative to body center of mass
    origin_capsule_b: position of capsule contact b relative to body center of mass
    radius_sphere:    radius of sphere contact
    radius_capsule:   radius of capsule contact
"""
mutable struct SphereCapsuleCollision{T,O,I,OI} <: Collision{T,O,I,OI}
    origin_sphere::SVector{I,T}
    origin_capsule_a::SVector{I,T}
    origin_capsule_b::SVector{I,T}
    radius_sphere::T
    radius_capsule::T
end

# distance
function distance(collision::SphereCapsuleCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_sphere)
    coc = contact_point_segment(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)

    # distance between contact origins
    d = norm(cop - coc, 2)

    # minimum distance between spheres
    return d - (collision.radius_sphere + collision.radius_capsule)
end

function ∂distance∂x(gradient::Symbol, collision::SphereCapsuleCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_sphere)
    coc = contact_point_segment(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)

    # distance between contact origins
    d = norm(cop - coc, 2)
    ∂norm∂d = ∂norm∂x(cop - coc)

    if gradient == :parent
        D = ∂norm∂d * (1.0 * ∂contact_point_origin∂x(xp, qp, collision.origin_sphere) - ∂contact_point_segment∂p(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b) * ∂contact_point_origin∂x(xp, qp, collision.origin_sphere))
    elseif gradient == :child
        D = ∂norm∂d * -1.0 * ∂contact_point_segment∂x(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)
    end

    # if gradient == :parent
    #     FD = FiniteDiff.finite_difference_jacobian(x -> distance(collision, x, qp, xc, qc), xp)
    # elseif gradient == :child
    #     FD = FiniteDiff.finite_difference_jacobian(x -> distance(collision, xp, qp, x, qc), xc)
    # end

    # return FD

    # @assert norm(D - FD, Inf) < 1.0e-5

    return D
end

function ∂distance∂q(gradient::Symbol, collision::SphereCapsuleCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_sphere)
    coc = contact_point_segment(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)

    # distance between contact origins
    d = norm(cop - coc, 2)
    ∂norm∂d = ∂norm∂x(cop - coc)

    if gradient == :parent
        D = ∂norm∂d * (1.0 * ∂contact_point_origin∂q(xp, qp, collision.origin_sphere) - ∂contact_point_segment∂p(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b) * ∂contact_point_origin∂q(xp, qp, collision.origin_sphere))
    elseif gradient == :child
        D = ∂norm∂d * -1.0 * ∂contact_point_segment∂q(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)
    end

    # if gradient == :parent
    #     FD = FiniteDiff.finite_difference_jacobian(q -> distance(collision, xp, Quaternion(q..., false), xc, qc), vector(qp))
    # elseif gradient == :child
    #     FD = FiniteDiff.finite_difference_jacobian(q -> distance(collision, xp, qp, xc, Quaternion(q..., false)), vector(qc))
    # end

    # return FD

    # @assert norm(D - FD, Inf) < 1.0e-5

    return D
end

# contact point in world frame
function contact_point(relative::Symbol, collision::SphereCapsuleCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_sphere)
    coc = contact_point_segment(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)

    # direction of minimum distance (child to parent)
    d = cop - coc
    dir = normalize(d)

    # contact point
    if relative == :parent
        return cop - collision.radius_sphere * dir
    elseif relative == :child
        return coc + collision.radius_capsule * dir
    end
end

function ∂contact_point∂x(relative::Symbol, jacobian::Symbol, collision::SphereCapsuleCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_sphere)
    coc = contact_point_segment(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)

    # direction of minimum distance (child to parent)
    d = cop - coc
    dir = normalize(d)

    if relative == :parent
        # cop - collision.radius_sphere * dir
        if jacobian == :parent
            ∂c∂x = ∂contact_point_origin∂x(xp, qp, collision.origin_sphere)
            X = ∂c∂x
            X -= collision.radius_sphere * ∂normalize∂x(d) * (∂c∂x - ∂contact_point_segment∂p(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b) * ∂c∂x)
        elseif jacobian == :child
            X = -1.0 * collision.radius_sphere * ∂normalize∂x(d) * -1.0 * ∂contact_point_segment∂x(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)
        end
    elseif relative == :child
        # coc + collision.radius_capsule * dir
        if jacobian == :parent
            X = ∂contact_point_segment∂p(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b) * ∂contact_point_origin∂x(xp, qp, collision.origin_sphere)
            X += 1.0 * collision.radius_capsule * ∂normalize∂x(d) * (∂contact_point_origin∂x(xp, qp, collision.origin_sphere) - ∂contact_point_segment∂p(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b) * ∂contact_point_origin∂x(xp, qp, collision.origin_sphere))
        elseif jacobian == :child
            X = ∂contact_point_segment∂x(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)
            X += collision.radius_capsule * ∂normalize∂x(d) * (-∂contact_point_segment∂x(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b))
        end
    end

    # if jacobian == :parent
    #     FD =  FiniteDiff.finite_difference_jacobian(x -> contact_point(relative, collision, x, qp, xc, qc), xp)
    # elseif jacobian == :child
    #     FD = FiniteDiff.finite_difference_jacobian(x -> contact_point(relative, collision, xp, qp, x, qc), xc)
    # end

    # return FD

    # @assert norm(X - FD, Inf) < 1.0e-5

    return X
end

function ∂contact_point∂q(relative::Symbol, jacobian::Symbol, collision::SphereCapsuleCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_sphere)
    coc = contact_point_segment(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)

    # direction of minimum distance (child to parent)
    d = cop - coc
    dir = normalize(d)

    if relative == :parent
        # cop - collision.radius_sphere * dir
        if jacobian == :parent
            ∂c∂q = ∂contact_point_origin∂q(xp, qp, collision.origin_sphere)
            Q = ∂c∂q
            Q -= collision.radius_sphere * ∂normalize∂x(d) * (∂c∂q - ∂contact_point_segment∂p(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b) * ∂c∂q)
        elseif jacobian == :child
            Q = -1.0 * collision.radius_sphere * ∂normalize∂x(d) * -1.0 * ∂contact_point_segment∂q(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)
        end
    elseif relative == :child
        # coc + collision.radius_capsule * dir
        if jacobian == :parent
            Q = ∂contact_point_segment∂p(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b) * ∂contact_point_origin∂q(xp, qp, collision.origin_sphere)
            Q += 1.0 * collision.radius_capsule * ∂normalize∂x(d) * (∂contact_point_origin∂q(xp, qp, collision.origin_sphere) - ∂contact_point_segment∂p(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b) * ∂contact_point_origin∂q(xp, qp, collision.origin_sphere))
        elseif jacobian == :child
            Q = ∂contact_point_segment∂q(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b)
            Q += collision.radius_capsule * ∂normalize∂x(d) * (-∂contact_point_segment∂q(cop, xc, qc, collision.origin_capsule_a, collision.origin_capsule_b))
        end
    end

    # if jacobian == :parent
    #     FD = FiniteDiff.finite_difference_jacobian(q -> contact_point(relative, collision, xp, Quaternion(q..., false), xc, qc), vector(qp))
    # elseif jacobian == :child
    #     FD = FiniteDiff.finite_difference_jacobian(q -> contact_point(relative, collision, xp, qp, xc, Quaternion(q..., false)), vector(qc))
    # end

    # return FD

    # @assert norm(Q - FD, Inf) < 1.0e-5

    return Q
end


text:
struct ArcIter{T} <: AbstractArcIter{T}
    list::Vector{Arc{T}}
end

@forward ArcIter.list iterate, eltype, length, rand

function arcs(g::Graph{T}) where {T}
    list = Arc{T}[]
    sizehint!(list, g.na)
    @inbounds for (u,nbr) in g.fadj, (v,keylist) in nbr, (key,fin) in keylist
        Arc(v, u, key) in list || push!(list, Arc(u, v, key))
    end
    return ArcIter(list)
end

function arcs(g::Graph{T}, n::Union{T,Vector{T}}) where {T}
    # all arc adjacent to n in UNDIRECTED Graph
    list = Arc{T}[]
    sizehint!(list, g.na)
    @inbounds for u in n, (v,keylist) in g.fadj[u], (key,fin) in keylist
        Arc(v, u, key) in list || push!(list, Arc(u, v, key))
    end
    return ArcIter(list)
end

# arcs dependent on directedness. why? karena klo Graph pas ditarik arc list-nya,
# cuma ada satu representasi dari dua arc yg bolak balik.
# which means (u -- v dan v -- u direpresentasikan oleh u -- v)

function arcs(g::Digraph{T}) where {T}
    list = Arc{T}[]
    sizehint!(list, g.na)
    @inbounds for (u,nbr) in g.fadj, (v,keylist) in nbr, (key,fin) in keylist
        push!(list, Arc(u, v, key))
    end
    return ArcIter(list)
end

function arcs(g::Digraph{T}, n::Union{T,Vector{T}}, ::Colon) where {T}
    list = Arc{T}[]
    sizehint!(list, g.na)
    @inbounds for u in n, (v,keylist) in g.fadj[u], (key,fin) in keylist
        push!(list, Arc(u, v, key))
    end
    return ArcIter(list)
end

function arcs(g::Digraph{T}, ::Colon, n::Union{T,Vector{T}}) where {T}
    list = Arc{T}[]
    sizehint!(list, g.na)
    @inbounds for v in n, (u,keylist) in g.badj[v], (key,fin) in keylist
        push!(list, Arc(u, v, key))
    end
    return ArcIter(list)
end


text:
"return the corresponding `DataType` in julia"
_jltype(dtype::GDALDataType) = _jltype(Val{dtype})
_jltype(::Type{Val{GDT_Unknown}}) = Any
_jltype(::Type{Val{GDT_Byte}}) = UInt8
_jltype(::Type{Val{GDT_UInt16}}) = UInt16
_jltype(::Type{Val{GDT_Int16}}) = Int16
_jltype(::Type{Val{GDT_UInt32}}) = UInt32
_jltype(::Type{Val{GDT_Int32}}) = Int32
_jltype(::Type{Val{GDT_Float32}}) = Float32
_jltype(::Type{Val{GDT_Float64}}) = Float64

"return the corresponding `GDALDataType`"
_gdaltype(dtype::DataType) = _gdaltype(Val{dtype})
_gdaltype(::Type{Val{Any}}) = GDT_Unknown
_gdaltype(::Type{Val{UInt8}}) = GDT_Byte
_gdaltype(::Type{Val{UInt16}}) = GDT_UInt16
_gdaltype(::Type{Val{Int16}}) = GDT_Int16
_gdaltype(::Type{Val{UInt32}}) = GDT_UInt32
_gdaltype(::Type{Val{Int32}}) = GDT_Int32
_gdaltype(::Type{Val{Float32}}) = GDT_Float32
_gdaltype(::Type{Val{Float64}}) = GDT_Float64

"""
`TRUE` if the passed type is complex (one of `GDT_CInt16`, `GDT_CInt32`,
`GDT_CFloat32` or `GDT_CFloat64`)
"""
_datatypeiscomplex(dtype::GDALDataType) = GDALDataTypeIsComplex(dtype)::Cint

datatypeiscomplex(dtype::GDALDataType) = Bool(_datatypeiscomplex(dtype))


text:
module TestExamples
using Test
dir = joinpath("..", "examples")
@testset "example $filename" for filename in readdir(dir)
    path = joinpath(dir, filename)
    include(path)
end
end#module


text:
struct Note <: Object
    variables::NamedTuple
    metadata::Metadata
end

Base.NamedTuple(note::Note) = getfield(note, :variables)
Metadata(note::Note) = getfield(note, :metadata)

Base.propertynames(note::Note) = propertynames(NamedTuple(note))
Base.getproperty(note::Note, name::Symbol) = NamedTuple(note)[name]

"""
    Recalls.NOTES

This is a vector of notes created by `@note`.
"""
const NOTES = GlobalRecord{Note}()

function _record_note(metadata; variables...)
    @nospecialize variables
    put!(NOTES, Note((; variables...), metadata))
end

raw"""
    @note var₁ var₂ … varₙ

Record variables `varᵢ`.

Each `varᵢ` can take the following form:

1. a symbol prefixed by `$` (e.g., `$x`)
2. an assignment with the right hand side prefixed by `$` (e.g., `lhs = $rhs`)
3. a symbol (e.g., `x`)
4. an assignment (e.g., `lhs = rhs`)

Expressions prefixed by `$` will be `deepcopy`ed.
"""
macro note(variables...)
    metadata = metadata_expr(__source__, __module__)
    kwargs = map(variables) do v
        if v isa Symbol
            Expr(:kw, v, v)
        elseif Meta.isexpr(v, :$, 1) && v.args[1] isa Symbol
            Expr(:kw, v.args[1], :($deepcopy($(v.args[1]))))
        elseif (
                (Meta.isexpr(v, :(=), 2) || Meta.isexpr(v, :kw, 2)) &&
                Meta.isexpr(v.args[2], :$, 1)
            )
            Expr(:kw, v.args[1], :($deepcopy($(v.args[2].args[1]))))
        else
            v
        end
    end
    quote
        $_record_note($metadata; $(kwargs...))
        nothing
    end |> esc
end

function Base.summary(io::IO, note::Note)
    print(io, "@note with ", length(NamedTuple(note)), " variable(s) for ")
    _summary(io, Metadata(note))
end

function Base.show(io::IO, ::MIME"text/plain", note::Note)
    summary(io, note)
    get(io, :compact, false) && return
    for (k, v) in pairs(NamedTuple(note))
        println(io)
        print(io, k, " = ")
        show(io, v)
    end
end

function sortbytime!(notes = NOTES)
    sort!(notes.buffer, by = n -> Metadata(n).timestamp)
    return notes
end

function note_table(notes = NOTES)
    table = (
        variables = NamedTuple[],
        line = Int[],
        file = Union{Nothing,Symbol}[],
        _module = Module[],
        location = UUID[],
        threadid = Int[],
        timestamp = TimeStamp[],
    )
    for i in eachindex(notes)
        n = notes[i]
        variables = NamedTuple(n)
        metadata = Metadata(n)
        push!(table.variables, variables)
        push!(table.line, metadata.source.line)
        push!(table.file, metadata.source.file)
        push!(table._module, metadata._module)
        push!(table.location, metadata.location)
        push!(table.threadid, metadata.threadid)
        push!(table.timestamp, metadata.timestamp)
    end
    return table
end


text:
"""
    MPO

A finite size matrix product operator type.
Keeps track of the orthogonality center.
"""
mutable struct MPO <: AbstractMPS
  length::Int
  data::Vector{ITensor}
  llim::Int
  rlim::Int
  function MPO(N::Int,
               A::Vector{<:ITensor},
               llim::Int = 0,
               rlim::Int = N+1)
    new(N, A, llim, rlim)
  end
end

MPO() = MPO(0, Vector{ITensor}(), 0, 0)

function MPO(sites::Vector{<:Index})
  N = length(sites)
  v = Vector{ITensor}(undef, N)
  l = [Index(1, "Link,l=$ii") for ii in 1:N-1]
  for ii ∈ eachindex(sites)
    s = sites[ii]
    sp = prime(s)
    if ii == 1
      v[ii] = ITensor(s, sp, l[ii])
    elseif ii == N
      v[ii] = ITensor(l[ii-1], s, sp)
    else
      v[ii] = ITensor(l[ii-1], s, sp, l[ii])
    end
  end
  return MPO(N, v)
end

MPO(A::Vector{<:ITensor}) = MPO(length(A), A)

"""
    MPO(N::Int)

Make an MPO of length `N` filled with default ITensors.
"""
MPO(N::Int) = MPO(N, Vector{ITensor}(undef, N))

"""
    MPO(sites, ops::Vector{String})

Make an MPO with pairs of sites `s[i]` and `s[i]'`
and operators `ops` on each site.
"""
function MPO(sites,
             ops::Vector{String})
  N = length(sites)
  its = Vector{ITensor}(undef, N)
  links = Vector{Index}(undef, N)
  for ii ∈ eachindex(sites)
    si = sites[ii]
    d = dim(si)
    spin_op = op(sites, ops[ii], ii)
    links[ii] = Index(1, "Link,n=$ii")
    local this_it
    if ii == 1
      this_it = ITensor(links[ii], si, si')
      for jj in 1:d, jjp in 1:d
        this_it[links[ii](1), si[jj], si'[jjp]] = spin_op[si[jj], si'[jjp]]
      end
    elseif ii == N
      this_it = ITensor(links[ii-1], si, si')
      for jj in 1:d, jjp in 1:d
        this_it[links[ii-1](1), si[jj], si'[jjp]] = spin_op[si[jj], si'[jjp]]
      end
    else
      this_it = ITensor(links[ii-1], links[ii], si, si')
      for jj in 1:d, jjp in 1:d
        this_it[links[ii-1](1),
                links[ii](1),
                si[jj],
                si'[jjp]] = spin_op[si[jj], si'[jjp]]
      end
    end
    its[ii] = this_it
  end
  MPO(N, its)
end

"""
    MPO(sites, ops::Vector{String})

Make an MPO with pairs of sites `s[i]` and `s[i]'`
and operators `ops` on each site.
"""
MPO(sites, ops::String) = MPO(sites, fill(ops, length(sites)))

function randomMPO(sites, m::Int=1)
  M = MPO(sites)
  for i ∈ eachindex(sites)
    randn!(M[i])
    normalize!(M[i])
  end
  m > 1 && throw(ArgumentError("randomMPO: currently only m==1 supported"))
  return M
end

"""
    siteind(A::MPO, x::MPS, j::Int)

Get the site index of MPO `A` that is unique to
`A` (not shared with MPS `x`).
"""
function siteind(A::MPO, x::MPS, j::Int)
  N = length(A)
  if j == 1
    si = uniqueind(A[j], A[j+1], x[j])
  elseif j == N
    si = uniqueind(A[j], A[j-1], x[j])
  else
    si = uniqueind(A[j], A[j-1], A[j+1], x[j])
  end
  return si
end

"""
    siteinds(A::MPO, x::MPS)

Get the site indices of MPO `A` that are unique to
`A` (not shared with MPS `x`), as a `Vector{<:Index}`.

This is the same as getting the `siteinds` of `A|x>`, i.e.
`siteinds(A * x)`, without doing the contraction.
"""
siteinds(A::MPO, x::MPS) = [siteind(A, x, j) for j in eachindex(A)]

"""
    dot(y::MPS, A::MPO, x::MPS; make_inds_match::Bool = true)
    inner(y::MPS, A::MPO, x::MPS; make_inds_match::Bool = true)

Compute `<y|A|x> = <y|Ax>.

If `make_inds_match = true`, the function attempts to make
the site indices of `A*x` match with the site indices of `y`
before contracting (so for example, the inputs `y` and `A*x`
can have different site indices, as long as they have the same
dimensions or QN blocks).

`A` and `x` must have common site indices.
"""
function LinearAlgebra.dot(y::MPS, A::MPO, x::MPS;
                           make_inds_match::Bool = true)::Number
  N = length(A)
  if length(y) != N || length(x) != N
      throw(DimensionMismatch("inner: mismatched lengths $N and $(length(x)) or $(length(y))"))
  end
  ydag = dag(y)
  simlinkinds!(ydag)
  if make_inds_match
    sAx = siteinds(A, x)
    replacesiteinds!(ydag, sAx)
  end
  O = ydag[1]*A[1]*x[1]
  for j in 2:N
    O = O*ydag[j]*A[j]*x[j]
  end
  return O[]
end

inner(y::MPS, A::MPO, x::MPS) = dot(y, A, x)

"""
    dot(B::MPO, y::MPS, A::MPO, x::MPS; make_inds_match::Bool = true)
    inner(B::MPO, y::MPS, A::MPO, x::MPS; make_inds_match::Bool = true)

Compute `<By|A|x> = <By|Ax>`.

If `make_inds_match = true`, the function attempts to make
the site indices of `A*x` match with the site indices of `B*y`
before contracting (so for example, the inputs `B*y` and `A*x`
can have different site indices, as long as they have the same
dimensions or QN blocks).

`A` and `x` must have common site indices, and `B` and `y`
must have common site indices.
"""
function LinearAlgebra.dot(B::MPO, y::MPS,
                           A::MPO, x::MPS;
                           make_inds_match::Bool = true)::Number
  !make_inds_match && error("make_inds_match = false not currently supported in dot(::MPO, ::MPS, ::MPO, ::MPS)")
  N = length(B)
  if length(y) != N || length(x) != N || length(A) != N
      throw(DimensionMismatch("inner: mismatched lengths $N and $(length(x)) or $(length(y)) or $(length(A))"))
  end
  ydag = dag(y)
  prime!(ydag, 2)
  Bdag = dag(B)
  prime!(Bdag)
  # Swap prime levels 1 -> 2 and 2 -> 1.
  for j in eachindex(Bdag)
    Axcommon = commonind(A[j], x[j])
    ABcommon = uniqueind(inds(A[j], "Site"), IndexSet(Axcommon))
    swapprime!(Bdag[j],2,3)
    swapprime!(Bdag[j],1,2)
    swapprime!(Bdag[j],3,1)
    noprime!(Bdag[j],prime(ABcommon,2))
  end
  yB = ydag[1] * Bdag[1]
  Ax = A[1] * x[1]
  O = yB*Ax
  for j in 2:N
    yB = ydag[j] * Bdag[j]
    Ax = A[j] * x[j]
    yB *= O
    O = yB * Ax
  end
  return O[]
end

inner(B::MPO, y::MPS, A::MPO, x::MPS) = dot(B, y, A, x)

"""
    error_contract(y::MPS, A::MPO, x::MPS;
                   make_inds_match::Bool = true)
    error_contract(y::MPS, x::MPS, x::MPO;
                   make_inds_match::Bool = true)

Compute the distance between A|x> and an approximation MPS y:
`| |y> - A|x> |/| A|x> | = √(1 + (<y|y> - 2*real(<y|A|x>))/<Ax|A|x>)`.

If `make_inds_match = true`, the function attempts match the site
indices of `y` with the site indices of `A` that are not common
with `x`.
"""
function error_contract(y::MPS, A::MPO, x::MPS; kwargs...)
  N = length(A)
  if length(y) != N || length(x) != N
    throw(DimensionMismatch("inner: mismatched lengths $N and $(length(x)) or $(length(y))"))
  end
  iyy   = dot(y, y; kwargs...)
  iyax  = dot(y, A, x; kwargs...)
  iaxax = dot(A, x, A, x; kwargs...)
  return sqrt(abs(1. + (iyy - 2*real(iyax))/iaxax))
end

error_contract(y::MPS, x::MPS, A::MPO) = error_contract(y, A, x)

"""
    contract(::MPS, ::MPO; kwargs...)
    *(::MPS, ::MPO; kwargs...)

    contract(::MPO, ::MPS; kwargs...)
    *(::MPO, ::MPS; kwargs...)

Contract the MPO with the MPS, returning an MPS with the unique
site indices of the MPO.

Choose the method with the `method` keyword, for example
"densitymatrix" and "naive".
"""
function Base.:*(A::MPO, psi::MPS; kwargs...)::MPS
  method = get(kwargs, :method, "densitymatrix")

  # Keyword argument deprecations
  if method == "DensityMatrix"
    @warn "In contract, method DensityMatrix is deprecated in favor of densitymatrix"
    method = "densitymatrix"
  end

  if method == "Naive"
    @warn "In contract, method Naive is deprecated in favor of naive"
    method = "naive"
  end

  if method == "densitymatrix"
    return _contract_densitymatrix(A, psi; kwargs...)
  elseif method == "naive"
    return _contract_naive(A, psi; kwargs...)
  end
  throw(ArgumentError("Method $method not supported"))
end

Base.:*(psi::MPS, A::MPO; kwargs...) = *(A, psi; kwargs...)

function _contract_densitymatrix(A::MPO, psi::MPS; kwargs...)::MPS
  n = length(A)
  n != length(psi) && throw(DimensionMismatch("lengths of MPO ($n) and MPS ($(length(psi))) do not match"))
  psi_out         = similar(psi)
  cutoff::Float64 = get(kwargs, :cutoff, 1e-13)
  maxdim::Int     = get(kwargs,:maxdim,maxlinkdim(psi))
  mindim::Int     = max(get(kwargs,:mindim,1), 1)
  normalize::Bool = get(kwargs, :normalize, false)
  all(x -> !isnothing(x),
      [siteind(A, psi, j) for j in 1:n]) ||
  throw(ErrorException("MPS and MPO have different site indices in contract method 'densitymatrix'"))

  rand_plev = 14741
  psi_c     = dag(copy(psi))
  A_c       = dag(copy(A))
  prime!(psi_c, rand_plev)
  prime!(A_c, rand_plev)
  for j in 1:n
    s = siteind(A, psi, j)
    s_dag = siteind(A_c, psi_c, j)
    replaceind!(A_c[j], s_dag, s)
  end
  E = Vector{ITensor}(undef, n-1)
  E[1] = psi[1]*A[1]*A_c[1]*psi_c[1]
  for j in 2:n-1
    E[j] = E[j-1]*psi[j]*A[j]*A_c[j]*psi_c[j]
  end
  O     = psi[n] * A[n]
  ρ     = E[n-1] * O * dag(prime(O, rand_plev))
  ts    = tags(commonind(psi[n], psi[n-1]))
  Lis   = commonind(ρ, A[n])
  Ris   = prime(Lis, rand_plev)
  FU, D = eigen(ρ, Lis, Ris; ishermitian=true,
                             tags=ts,
                             kwargs...)
  psi_out[n] = dag(FU)
  O     = O * FU * psi[n-1] * A[n-1]
  for j in reverse(2:n-1)
    dO  = prime(dag(O), rand_plev)
    ρ   = E[j-1] * O * dO
    ts  = tags(commonind(psi[j], psi[j-1]))
    Lis = IndexSet(commonind(ρ, A[j]),
                   commonind(ρ, psi_out[j+1]))
    Ris = prime(Lis, rand_plev)
    FU, D = eigen(ρ, Lis, Ris; ishermitian=true,
                               tags=ts,
                               kwargs...)
    psi_out[j] = dag(FU)
    O = O * FU * psi[j-1] * A[j-1]
  end
  if normalize
    O /= norm(O)
  end
  psi_out[1]    = copy(O)
  setleftlim!(psi_out, 0)
  setrightlim!(psi_out, 2)
  return psi_out
end

function _contract_naive(A::MPO, psi::MPS; kwargs...)::MPS
  truncate = get(kwargs,:truncate,true)

  N = length(A)
  if N != length(psi)
    throw(DimensionMismatch("lengths of MPO ($N) and MPS ($(length(psi))) do not match"))
  end

  psi_out = MPS(N)
  for j=1:N
    psi_out[j] = A[j]*psi[j]
  end

  for b=1:(N-1)
    Al = commonind(A[b],A[b+1])
    pl = commonind(psi[b],psi[b+1])
    C = combiner(Al,pl)
    psi_out[b] *= C
    psi_out[b+1] *= dag(C)
  end

  if truncate
    truncate!(psi_out;kwargs...)
  end

  return psi_out
end

function Base.:*(A::MPO, B::MPO; kwargs...)
  cutoff::Float64 = get(kwargs, :cutoff, 1e-14)
  resp_degen::Bool = get(kwargs, :respect_degenerate, true)
  maxdim::Int = get(kwargs,:maxdim,maxlinkdim(A)*maxlinkdim(B))
  mindim::Int = max(get(kwargs,:mindim,1), 1)
  N = length(A)
  N != length(B) && throw(DimensionMismatch("lengths of MPOs A ($N) and B ($(length(B))) do not match"))
  A_ = copy(A)
  orthogonalize!(A_, 1)
  B_ = copy(B)
  orthogonalize!(B_, 1)

  links_A = inds.(A.data, "Link")
  links_B = inds.(B.data, "Link")

  res = deepcopy(A_)
  for i in 1:N-1
    ci = commonind(res[i], res[i+1])
    new_ci = Index(dim(ci), tags(ci))
    replaceind!(res[i], ci, new_ci)
    replaceind!(res[i+1], ci, new_ci)
    @assert commonind(res[i], res[i+1]) != commonind(A[i], A[i+1])
  end
  sites_A = Index[]
  sites_B = Index[]
  for (AA, BB) in zip(data(A_), data(B_))
    sda = setdiff(inds(AA, "Site"), inds(BB, "Site"))
    sdb = setdiff(inds(BB, "Site"), inds(AA, "Site"))
    length(sda) != 1 && error("In contract(::MPO, ::MPO), MPOs must have exactly one shared site index")
    length(sdb) != 1 && error("In contract(::MPO, ::MPO), MPOs must have exactly one shared site index")
    push!(sites_A, sda[1])
    push!(sites_B, sdb[1])
  end
  res[1] = ITensor(sites_A[1], sites_B[1], commonind(res[1], res[2]))
  for i in 1:N-2
    if i == 1
      clust = A_[i] * B_[i]
    else
      clust = nfork * A_[i] * B_[i]
    end
    lA = commonind(A_[i], A_[i+1])
    lB = commonind(B_[i], B_[i+1])
    nfork = ITensor(lA, lB, commonind(res[i], res[i+1]))
    res[i], nfork = factorize(clust,
                              inds(res[i]),
                              ortho="left",
                              tags=tags(lA),
                              cutoff=cutoff,
                              maxdim=maxdim,
                              mindim=mindim)
    mid = dag(commonind(res[i], nfork))
    res[i+1] = ITensor(mid,
                       sites_A[i+1],
                       sites_B[i+1],
                       commonind(res[i+1], res[i+2]))
  end
  clust = nfork * A_[N-1] * B_[N-1]
  nfork = clust * A_[N] * B_[N]

  # in case we primed A
  A_ind = uniqueind(inds(A_[N-1], "Site"), inds(B_[N-1], "Site"))
  Lis = IndexSet(A_ind, sites_B[N-1], commonind(res[N-2], res[N-1]))
  U, V = factorize(nfork, Lis;
                   ortho="right",
                   cutoff=cutoff,
                   tags="Link,n=$(N-1)",
                   maxdim=maxdim,
                   mindim=mindim)
  res[N-1] = U
  res[N] = V
  truncate!(res;kwargs...)
  return res
end


text:
# brent-tests.jl -
#
# Tests for Brent's algorithms.
#

module BrentTests

using Test, Printf
using OptimPackNextGen.Brent

# Counter of function evaluations.
const cnt = Ref{Int}(0)

# Yields `y` converted to a floating-point type according to the type of `x`.
fp(x::Number, y::Real) = convert(float(typeof(x)), y)

fzero_test0(x::Real) = (cnt[] += 1; 1/(x - 3) - 6)
params(::typeof(fzero_test0)) = (3, 4, 19/6)

# Test functions for root finding from Vakkalagadda Satya Sai Prakash,
# "Implementation of Brent-Dekker and A Better Root Finding Method and
# Brent-Dekker Method's Parallelization".

fzero_test1(x::Real) = (cnt[] += 1; exp(x)/2 - 5x + 2)
params(::typeof(fzero_test1)) = (1, 4, NaN)

fzero_test2(x::Real) = (cnt[] += 1;
                                 -2x^4 + 2x^3 - 16x^2 - 60x + 100)
params(::typeof(fzero_test2)) = (0, 4, NaN)

fzero_test3(x::Real) = (cnt[] += 1; exp(x)*cos(x) - x*sin(x))
params(::typeof(fzero_test3)) = (0, 3, NaN)

fzero_test4(x::Real) = (cnt[] += 1; x^5 - 5x + 3)
params(::typeof(fzero_test4)) = (0, 1, NaN)

fzero_test5(x::Real) = (cnt[] += 1;
                        (x^3 - fp(x,0.926)*x^2 +
                         fp(x,0.0371)*x + fp(x,0.043)))
params(::typeof(fzero_test5)) = (0, 0.8, NaN)

fzero_test6(x::Real) = (cnt[] += 1;
                        (-9 + sqrt(99 + 2x - x^2) + cos(2x)))
params(::typeof(fzero_test6)) = (-2, 6, NaN)

fzero_test7(x::Real) = (cnt[] += 1; sin(cosh(x)))
params(::typeof(fzero_test7)) = (0, 2, NaN)

fzero_test8(x::Real) = (cnt[] += 1; exp(-exp(-x)) - x)
params(::typeof(fzero_test8)) = (0, 1, NaN)

# Test functions from the GNU Scientific Library (GSL).

fzero_gsl1(x::Real) = (cnt[] += 1; x^20 - 1)
params(::typeof(fzero_gsl1)) = (0.1, 2.0, 1.0)

fzero_gsl2(x::Real) = (cnt[] += 1; sqrt(abs(x))*sign(x))
params(::typeof(fzero_gsl2)) = (-1.0/3.0, 1.0, 0.0)

fzero_gsl3(x::Real) = (cnt[] += 1; x^2 - fp(x, 1e-8))
params(::typeof(fzero_gsl3)) = (0, 1, sqrt(1e-8))

fzero_gsl4(x::Real) = (cnt[] += 1; x*exp(-x))
params(::typeof(fzero_gsl4)) = (-1.0/3.0, 2.0, 0.0)

fzero_gsl5(x::Real) = (cnt[] += 1; (x - 1)^7)
params(::typeof(fzero_gsl5)) = (0, 3, 1) # (0.9995, 1.0002, 1)

# Problem 4 from AMPGO, xm = 2.868034, fm = -3.85045:
problem04(x) = (cnt[] += 1; -(16x^2 - 24x + 5)*exp(-x))
bounds(::typeof(problem04)) = [1.9, 3.9]
solution(::typeof(problem04)) = [2.868034, -3.85045]

# Problem 13 from AMPGO, xm = 1/sqrt(2), fm = -1.5874:
problem13(x) = (cnt[] += 1; -x^(2/3) - (1 - x^2)^(1/3))
bounds(::typeof(problem13)) = [0.001, 0.99]
solution(::typeof(problem13)) = [1/sqrt(2), -1.5874]

# Problem 18 from AMPGO, xm = 2, fm = 0:
problem18(x) = (cnt[] += 1; x ≤ 3 ? (x - 2)^2 : 2log(x - 2) + 1)
bounds(::typeof(problem18)) = [0, 6]
solution(::typeof(problem18)) = [2, 0]

verb = true
@testset "Brent fzero" begin
    let f = fzero_test0
        (a, b, x0) = params(f)
        @test_throws ArgumentError fzero(f, a, b; rtol=0)
        @test_throws ArgumentError fzero(f, a, b; rtol=1)
        @test_throws ArgumentError fzero(f, a, b; atol=NaN)
    end
    k = 0
    for f in (fzero_test0,
              fzero_test1, fzero_test2, fzero_test3, fzero_test4,
              fzero_test5, fzero_test6, fzero_test7, fzero_test8,
              fzero_gsl1, fzero_gsl2, fzero_gsl3, fzero_gsl4, fzero_gsl5)
        for T in (Float32, Float64)
            (a, b, x0) = params(f)
            rtol = Brent.fzero_rtol(T)
            atol = (x0 == 0 ? eps(T) : Brent.fzero_atol(T))
            prec = sqrt(eps(T))
            cnt[] = 0
            (x, fx) = if T === Float64
                fzero(f, a, b; rtol = rtol, atol = atol)
            else
                fzero(T, f, a, b; rtol = rtol, atol = atol)
            end
            if verb
                if k ≤ 8
                    name = "fzero_test$k"
                else
                    name = "fzero_gsl$(k - 8)"
                end
                @printf("%-12s T = %s, n = %3d, x = % .15f, f(x) = % .15f\n",
                        name, repr(T), cnt[], x, fx)
            end
            @test abs(fx) ≈ zero(T) rtol=zero(T) atol=prec
            if !isnan(x0)
                @test x ≈ T(x0) rtol=rtol atol=atol
            end
            fa = f(a)
            (x1, f1) =  if T === Float64
                fzero(f, a, fa, b; rtol = rtol, atol = atol)
            else
                fzero(T, f, a, fa, b; rtol = rtol, atol = atol)
            end
            @test x1 == x
            fb = f(b)
            (x2, f2) =  if T === Float64
                fzero(f, a, fa, b, fb; rtol = rtol, atol = atol)
            else
                fzero(T, f, a, fa, b, fb; rtol = rtol, atol = atol)
            end
            @test x2 == x
        end
        k += 1
    end
end

@testset "Brent fmin" begin
    for (f, id) in ((problem04,  4),
                    (problem13, 13),
                    (problem18, 18),)
        for (T, prec) in ((Float16, 3e-2),
                          (Float32, 3e-6),
                          (Float64, 6e-15))
            atol = Brent.fmin_atol(T)
            rtol = Brent.fmin_rtol(T)

            cnt[] = 0
            (xm, fm, lo, hi) = if T === Float64
                fmin(f, bounds(f)...)
            else
                fmin(T, f, bounds(f)...)
            end
            (xm_true, fm_true) = solution(f)
            if verb
                println("Problem $id (T=$T):")
                @printf(" ├─ nevals = %d\n", cnt[])
                @printf(" ├─ x_approx = %.15f ± %.3e\n", xm, (hi - lo)/2)
                @printf(" ├─ x_true   = %.15f\n", xm_true)
                @printf(" ├─ x_error  = %.3e\n", abs(xm - xm_true))
                @printf(" ├─ f(x_approx) = %.15f\n", fm)
                @printf(" ├─ f(x_true)   = %.15f\n", fm_true)
                @printf(" └─ f_error     = %.3e\n", abs(fm - fm_true))
            end
            @test abs(xm - xm_true) ≤ atol + rtol*abs(xm_true)
            #@test abs(fm - solution(f)[2]) ≤ prec
        end
        verb && println()
    end
end

end # module


text:
immutable Estimate
    point  :: Float64
    lbound :: Float64
    ubound :: Float64
    confidence_level :: Float64

    function Estimate {T<:Real} (pt::T, lb::T, ub::T, cl::T)
        @assert lb <= ub
        @assert cl > 0 && cl < 1
        return new(float64(pt), float64(lb),
	           float64(ub), float64(cl))
    end
end

function uniform_dist(max_val::Real)
    rand() * max_val
end

function sample_uniform{T<:Integer}(::Type{T}, n::Integer, max_val::T)
    [1 + int(uniform_dist(max_val - 1)) for _ in 1:n]
end

function sample_uniform{T<:FloatingPoint}(::Type{T}, n::Integer, max_val::T)
    [1 + uniform_dist(max_val - 1) for _ in 1:n]
end

function sample{T<:Real}(x::Vector{T})
    n = length(x)
    @assert n > 0
    samples = Array(T, n)
    if n == 1
        samples[1] = x[1]
    elseif n == 2
        # pick pair with replacement
        ridx1 = (rand() < 0.5) ? 1 : 2
        ridx2 = (rand() < 0.5) ? 1 : 2
        samples[1] = x[ridx1]
        samples[2] = x[ridx2]
    else
        # sample with replacement
        for s in 1:n
            ridx = int(round(1 + rand() * (n - 1)))
            samples[s] = x[ridx]
        end
    end
    return samples
end

function bootstrap_sample{T<:Real}(data::Vector{T},
				   stat::Function,
				   nsamples::Integer)
    samples = Array(T, nsamples)
    for idx in 1:nsamples
        samples[idx] = stat(sample(data))
    end
    return samples
end

function confidence_interval(mean::Float64,
			     variance::Float64)
    delta = 1.96 * sqrt(variance)
    (mean - delta, mean + delta)
end

function normal_cdf(x::Float64)
    0.5 * (1.0 + erf(x / sqrt(2.0)))
end

function normal_quantile(x::Float64)

    a = [2509.0809287301226727,
	 33430.575583588128105,
	 67265.770927008700853,
	 45921.953931549871457,
	 13731.693765509461125,
   	 1971.5909503065514427,
	 133.14166789178437745,
	 3.3871328727963666080]
    b = [5226.4952788528545610,
	 28729.085735721942674,
	 39307.895800092710610,
	 21213.794301586595867,
	 5394.1960214247511077,
	 687.18700749205790830,
	 42.313330701600911252,
	 1.0]
    c = [0.000774545014278341407640,
	 0.0227238449892691845833,
	 0.241780725177450611770,
	 1.27045825245236838258,
	 3.64784832476320460504,
	 5.76949722146069140550,
	 4.63033784615654529590,
	 1.42343711074968357734]
    d = [1.05075007164441684324e-9,
	 0.000547593808499534494600,
	 0.0151986665636164571966,
         0.148103976427480074590,
	 0.689767334985100004550,
	 1.67638483018380384940,
	 2.05319162663775882187,
	 1.0]
    e = [2.01033439929228813265e-7,
	 0.0000271155556874348757815,
	 0.00124266094738807843860,
	 0.0265321895265761230930,
	 0.296560571828504891230,
	 1.78482653991729133580,
	 5.46378491116411436990,
	 6.6579046435011037772]
    f = [2.04426310338993978564e-15,
	 1.42151175831644588870e-7,
	 1.84631831751005468180e-5,
	 0.000786869131145613259100,
	 0.0148753612908506148525,
	 0.136929880922735805310,
	 0.599832206555887937690,
	 1.0]

    function polynomial_value(x, coeff)
        fcoeff = coeff[1]
        lcoeff = coeff[2:end]
        reduce((x1, x2) -> (x * x1 + x2), coeff[1], coeff[2:end])
    end

    if 0.075 <= x <= 0.925
        v = x - 0.5
        r = 180625e-6 - v^2
  	pa = polynomial_value(r, a)
        pb = polynomial_value(r, b)
        return v * (pa / pb)
    else
        r = (x < 0.5) ? x : 1.0 - x
        r = sqrt(-log(r))
        if r <= 0.5
            r = r - (16.0 / 10.0)
            pc = polynomial_value(r, c)
            pd = polynomial_value(r, d)
            return sign(x - 0.5) * (pc / pd)
        else
            r = r - 5.0
            pe = polynomial_value(r, e)
            pf = polynomial_value(r, f)
            return sign(x - 0.5) * (pe / pf)
        end
    end
end

function jacknife{T<:Real}(data::Vector{T},
			   stat::Function)
    n = length(data)
    samples = Array(Float64, n)
    sample_data = Array(T, n-1)
    for i in 1:n
        idx = 0
        for j in 1:n
            if j != i
                idx += 1
                sample_data[idx] = data[i]
            end
        end
        samples[i] = stat(sample_data)
    end
    return samples
end

function bootstrap_percentile{T<:Real}(data::Vector{T},
				       statistic::Function,
				       alpha::Float64,
				       nsamples::Integer)
    @assert nsamples > 0 && alpha > 0.0
    n = length(data)
    boot_samples = bootstrap_sample(data, statistic, nsamples)
    sort!(boot_samples)
    idx_L = int(round(nsamples * alpha/2.0))
    idx_U = int(round(nsamples * (1.0 - alpha/2.0)))
    ci_L = boot_samples[idx_L]
    ci_U = boot_samples[idx_U]
    return (ci_L, ci_U)
end

function bootstrap_bca{T<:Float64}(data::Vector{T},
			           statistic::Function,
				   alpha::Float64,
				   nsamples::Integer)
    @assert nsamples > 0 && alpha > 0.0  && alpha < 1.0

    n = length(data)
    ci = 1.0 - (2.0 * alpha)

    boot_samples = bootstrap_sample(data, statistic, nsamples)
    est = statistic(data)
    z0 = normal_quantile(sum(boot_samples .< est) / nsamples)

    jack_samples = jacknife(data, statistic)
    jack_mean = mean(jack_samples)

    jack_dev = jack_samples - jack_mean
    jack_dev = Array(Float64, n)

    neps = 0
    for i in 1:length(jack_dev)
        d = jack_samples[i] - jack_mean
        if isapprox(d, 0.0)
            neps += 1
        end
        jack_dev[i] = d
    end

    # special case when all numbers are exactly the same or within fp tolerance
    if neps == length(jack_dev)
        return Estimate(jack_mean, jack_mean, jack_mean, ci)
    end

    a = sum((jack_dev .^ 3) / (6.0 * sum((jack_dev) .^ 2) ^ 1.5))
    #zL = z0 - normal_quantile((1.0 - alpha) / 2.0)
    #zU = z0 + normal_quantile((1.0 - alpha) / 2.0)

    #aL = normal_cdf(z0 + zL / (1.0 - a * zL))
    #aU = normal_cdf(z0 + zU / (1.0 - a * zU))

    #a_cdf = (zalpha) -> normal_cdf(z0 + (z0 + zalpha) / (1.0 - a * (z0 + zalpha)))

    aL = normal_cdf(z0 + (z0 + alpha) / (1.0 - a * (z0 + alpha)))
    aM = normal_cdf(z0 + (z0 + 0.5) / (1.0 - a * (z0 + 0.5)))
    aU = normal_cdf(z0 + (z0 + (1.0 - alpha)) / (1.0 - a * (z0 + (1.0 - alpha))))

    sort!(boot_samples)
    idxL = int(round(nsamples * aL))
    idxM = int(round(nsamples * aM))
    idxU = int(round(nsamples * aU))

    ciL = boot_samples[idxL]
    med = boot_samples[idxM]
    ciU = boot_samples[idxU]

    return Estimate(med, ciL, ciU, ci)
end

function scale_estimate(est::Estimate, factor::Real)
    return Estimate(factor * est.point,
	            factor * est.lbound,
		    factor * est.ubound,
		    est.confidence_level)
end

function scale_bootstrap_estimate(est::(Float64, Float64),
				  scale::Float64)
    (est[1] * scale, est[2] * scale)
end

function model_estimation_const(h_k::Float64,
				samp_var::Float64)
    @assert samp_var != 0.0
    sqrt(1.0 + h_k^2 / samp_var)
end

function smoothed_sample{T<:Float64}(c_k::T, h_k::T,
			             data::Vector{T},
			             deviates::Vector{T})
    [c_k * (d + h_k * dev) for d in data, dev in deviates]
end

function gaussian_weight(t::Float64)
    k = sqrt(2.0 * pi)
    k * exp(t^2  / -2.0)
end

function kernel_density_est(h::Float64, K::Function,
		        n::Float64, X::Float64, x::Float64)
    reduce((a,b) -> (a + K((x - b) / h)), 0, X) / (n * h)
end


text:
using SimpleSMC
using RNGPool
import SequentialMonteCarlo.SMCModel
using SMCExamples.MVLinearGaussian
using StaticArrays
import SMCExamples.Particles.MVFloat64Particle
using Test
using Random

setRNGs(0)

function mvlgtest(d::Int64)
  model, theta, ys, ko = MVLinearGaussian.defaultMVLGModel(5, d)

  numParticles = 2^16

  smcio = SMCIO{model.particle, model.pScratch}(numParticles, model.maxn)
  smc!(model, smcio)

  @test smcio.logZhats ≈ ko.logZhats atol=0.1*d
end

function testmvlgcsmc()
  model, theta, ys, ko = MVLinearGaussian.defaultMVLGModel(2, 10)
  nsamples = 2^12

  smcio = SMCIO{model.particle, model.pScratch}(256, model.maxn)

  v = Vector{MVFloat64Particle{2}}(undef, 10)
  for p = 1:10
    v[p] = MVFloat64Particle{2}()
    v[p].x .= zeros(MVector{2, Float64})
  end

  meanEstimates = Vector{MVector{2,Float64}}(undef, 10)
  for p = 1:10
    meanEstimates[p] = zeros(MVector{2, Float64})
  end
  for i = 1:nsamples
    csmc!(model, smcio, v, v)
    for p = 1:10
      meanEstimates[p] .+= v[p].x
    end
  end
  meanEstimates ./= nsamples

  testapproxequal((x -> x[1]).(meanEstimates), (x -> x[1]).(ko.smoothingMeans),
    0.1, false)
  testapproxequal((x -> x[2]).(meanEstimates), (x -> x[2]).(ko.smoothingMeans),
    0.1, false)
end

mvlgtest(1)
mvlgtest(2)

testmvlgcsmc()


text:
writing_output = false
if VERSION < v"1.5"
    ver = "111"
else
    ver = "150"
end
@everywhere Random.seed!(42)

####################################################################
# Testing Adaptive Φ Solution
####################################################################
file   = JLD2.jldopen(joinpath(dirname(@__FILE__), "reference/solve_adaptive_phi.jld2"), "r")
cloud  = read(file, "cloud")
i_smc  = read(file, "i")
j_smc  = read(file, "j")
ϕ_prop = read(file, "phi_prop")
ϕ_n1   = read(file, "phi_n1")
proposed_fixed_schedule = read(file, "proposed_fixed_schedule")
tempering_target        = read(file, "tempering_target")
resampled_last_period   = read(file, "resampled_last_period")
close(file)

test_ϕ_n, test_resampled_last_period, test_j, test_ϕ_prop = SMC.solve_adaptive_ϕ(cloud,
                                                                proposed_fixed_schedule,
                                                                i_smc, j_smc, ϕ_prop, ϕ_n1,
                                                                tempering_target,
                                                                resampled_last_period)
if writing_output
    jldopen(joinpath(dirname(@__FILE__), "reference/helpers_output_version=" * ver * ".jld2"),
            true, true, true, IOStream) do file
        write(file, "phi_n", test_ϕ_n)
        write(file, "resampled_last_period", test_resampled_last_period)
        write(file, "j", test_j)
        write(file, "phi_prop", test_ϕ_prop)
    end
end

file = JLD2.jldopen(joinpath(dirname(@__FILE__), "reference/helpers_output_version="
                             * ver * ".jld2"), "r")
saved_ϕ_n = read(file, "phi_n")
saved_resampled_last_period = read(file, "resampled_last_period")
saved_j = read(file, "j")
saved_ϕ_prop = read(file, "phi_prop")
close(file)

####################################################################
@testset "Solve Adaptive Φ" begin
    @test test_ϕ_n ≈ saved_ϕ_n
    @test test_resampled_last_period == saved_resampled_last_period
    @test test_j == saved_j
    @test test_ϕ_prop == saved_ϕ_prop
end

####################################################################
# Testing MvNormal Mixture Draw Function
####################################################################
file = JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/mvnormal_inputs.jld2"))
    para_subset = read(file, "para_subset")
    d_subset    = read(file, "d_subset")
    α           = read(file, "α")
    c           = read(file, "c")
close(file)

test_θ_new = SMC.mvnormal_mixture_draw(para_subset, d_subset; c=c, α=α)

if writing_output
    JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/mvnormal_output_version=" * ver * ".jld2"), true, true, true, IOStream) do file
        write(file, "θ_new", test_θ_new)
    end
end

file = JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/mvnormal_output_version=" * ver * ".jld2"))
    saved_θ_new = read(file, "θ_new")
close(file)

####################################################################
@testset "MvNormal Mixture Draw" begin
    @test test_θ_new ≈ saved_θ_new
end

####################################################################
# Test: get_cov()
####################################################################
d = JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/mutation_inputs.jld2"), "r") do file
    file["d"]
end
d_deg = DegenerateMvNormal(d.μ, d.Σ.mat)

####################################################################
@testset "Test: get_cov(d)" begin
    @test SMC.get_cov(d)     == d.Σ.mat
    @test SMC.get_cov(d_deg) == d_deg.σ
end

####################################################################
# Test: compute_proposal_densities()
####################################################################
file = JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/proposal_densities_in.jld2"))
    para_draw   = read(file, "para_draw")
    para_subset = read(file, "para_subset")
    d_subset    = read(file, "d_subset")
    α           = read(file, "α")
    c           = read(file, "c")
close(file)

q0, q1 = SMC.compute_proposal_densities(para_draw, para_subset, d_subset; α = α,
                                        c = c)
if writing_output
    JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/proposal_densities_output_version=" * ver * ".jld2"), true, true, true, IOStream) do file
        file["q0"] = q0
        file["q1"] = q1
    end
end

file = JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/proposal_densities_output_version=" * ver * ".jld2"))
    saved_q0 = read(file, "q0")
    saved_q1 = read(file, "q1")
close(file)

####################################################################
@testset "Proposal densities" begin
    @test q0 ≈ saved_q0
    @test q1 ≈ saved_q1
end

####################################################################
# Testing ESS Computation
####################################################################
file = JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/ess_inputs.jld2"))
    loglh           = read(file, "loglh")
    current_weights = read(file, "current_weights")
    ϕ_n             = read(file, "ϕ_n")
    ϕ_n1            = read(file, "ϕ_n1")
close(file)

test_ESS = SMC.compute_ESS(loglh, current_weights, ϕ_n, ϕ_n1)

if writing_output
    file = JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/smc_sw_cloud_fix=true_blocks=3.jld2"), "r")
    saved_cloud = read(file, "cloud")
    current_weights = read(file, "w")[:,3]
    close(file)

    n_part    = length(saved_cloud.particles)
    println(typeof(saved_cloud))
    loglh     = [saved_cloud.particles[i].loglh for i=1:n_part]
    old_loglh = [saved_cloud.particles[i].old_loglh for i=1:n_part]
    ϕ_n       = 9.25022e-6
    ϕ_n1      = 2.15769e-6

    JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/ess_inputs.jld2"), true, true, true, IOStream) do file
        write(file, "loglh", loglh)
        write(file, "current_weights", current_weights)
        write(file, "ϕ_n", ϕ_n)
        write(file, "ϕ_n1", ϕ_n1)
        write(file, "old_loglh", old_loglh)
    end

    JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/ess_output_version=" * ver * ".jld2"), true, true, true, IOStream) do file
        write(file, "ess", test_ESS)
    end
end

file = JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/ess_output_version=" * ver * ".jld2"))
    saved_ESS = read(file, "ess")
close(file)

####################################################################
@testset "Compute ESS" begin
    @test test_ESS ≈ saved_ESS
end

####################################################################
# Testing Block Creation
####################################################################
m = AnSchorfheide()

free_para_inds = findall(x -> !x.fixed, m.parameters)
n_free_para    = length(free_para_inds)
n_blocks       = 3

test_blocks_free = SMC.generate_free_blocks(n_free_para, n_blocks)
test_blocks_all  = SMC.generate_all_blocks(test_blocks_free, free_para_inds)
test_blocks      = SMC.generate_param_blocks(length(m.parameters), n_blocks)

if writing_output
    JLD2.jldopen(joinpath(dirname(@__FILE__),"reference/helpers_blocking_version=" * ver * ".jld2"), true, true, true, IOStream) do file
        file["blocks_free"] = test_blocks_free
        file["blocks_all"]  = test_blocks_all
        file["blocks"]      = test_blocks
    end
end

saved_blocks_free = load(joinpath(dirname(@__FILE__),"reference/helpers_blocking_version=" * ver * ".jld2"), "blocks_free")
saved_blocks_all  = load(joinpath(dirname(@__FILE__),"reference/helpers_blocking_version=" * ver * ".jld2"), "blocks_all")
saved_blocks      = load(joinpath(dirname(@__FILE__),"reference/helpers_blocking_version=" * ver * ".jld2"), "blocks")

####################################################################
@testset "Mutation block generation" begin
    @test test_blocks_free == saved_blocks_free
    @test test_blocks_all  == saved_blocks_all
    @test test_blocks      == saved_blocks
end


text:
using Pkg
using LinearAlgebra
using DataFrames

println("test corr-pf")

# load functions
include(pwd()*"/src/SDEMEM OU process/ou_sdemem.jl")
include(pwd()*"/src/SDEMEM OU process/mcmc.jl")

N_time = 200
M_subjects = 10
nbr_particles = 500

ρ = 0.9

θ_1 = 0.5
σ_ϵ = 0.2
θ_3 = 0.1

y,x,t_vec,dt,η,κ,σ_ϵ,ϕ,prior_η,prior_κ,prior_σ_ϵ = set_up(σ_ϵ = σ_ϵ, θ_3 = θ_3,θ_1=θ_1, M=M_subjects,N=N_time)

κ = [θ_1,θ_3]

function time_pf()

    for i = 1:1000
        ll = @time sum(pf(y, κ, σ_ϵ, ϕ, dt, nbr_particles))
    end

end

function time_corr_pf()

    for i = 1:1000
        ll = sum(cpf(y, κ, σ_ϵ, ϕ, dt, randn(nbr_particles,N_time + 1,M_subjects), randn(N_time,2,M_subjects), nbr_particles))
    end

end

runtime_pf = @elapsed time_pf()

runtime_pf = runtime_pf/1000

runtime_cpf = @elapsed time_corr_pf()

runtime_cpf = runtime_cpf/1000

runtime_pf/runtime_cpf


text:
#!/usr/bin/env julia
#=
   Author: Tim Sterne-Weiler, 9/13/2015
   e-mail: tim.sterne.weiler@utoronto.ca
=#

include("dictext.jl")

using ArgParse
using DataFrames
using Distributions

function parse_cmd()
  s = ArgParseSettings()

  @add_arg_table s begin
    "--bed"
      help = "toggle output as bed format (default on)"
      action = :store_true
    "--biot"
      help = "only output this specific biotype (default off)"
      arg_type = ASCIIString
    "--background"
      help = "gene expression levels background file {geneid first column, then data replicates cols}"
      arg_type = ASCIIString
    "--backheader"
      help = "is there a header in the background file? (default true)"
      action = :store_false
  end
  return parse_args(s)
end

###################################################################
function splitgb( coord )
  spl = split( coord, ':' )
  @assert( length(spl) > 1, "$coord does not split : into two!" )
  nums = map( x->parse(Int, x), split( spl[2], '-' ) )
  length(spl) >= 3 ? (spl[1], nums..., spl[3]) : (spl[1], nums...)
end

function read_lig_parse_reg( io; biot=nothing )
   const geneind = 25
   const biotind = 24
   const seqind = 11
   const posind = 17
   const refind = 14

   expfile = []
   bygenesym = Dict{ASCIIString,Int}()
   genesum = 0

   for l::ASCIIString in eachline( io )
      l[1] == 'S' || continue
      s = split(strip(l), '\t')
      g1,g2 = split( s[geneind], ':' )
      b1,b2 = split( s[biotind], ':' )
      #println("$g1 -- $g2, $b1 -- $b2")
      g1 == g2 || continue
      (biot != nothing && (b1 != biot || b2 != biot)) && continue
      s[posind][1:2] == "NA" && continue
      c1,c2 = map( x->splitgb(x), split( s[posind], ',' ) )
      ci1,ci2 = c1[2],c2[2]
      c1[1] == c2[1] || continue
      @assert( c1[end] == c2[end], "$(c1[end]) does not equal $(c2[end]) !!" )
      s1,s2 = replace(s[seqind], r"[a-z]", "") |> x->split( x, '_' )
      #println("$(s[refind]) is the ref ind")
      r1,r2 = split( s[refind], ',' ) |> x->tuple(map(y->parse(Int,y), x)...)
#=      if c1[end] == "+" # if positive
         ci1 += length(s1) - 1
      else # negative
         ci1 -= length(s1) + 1
      end=#
      first,second = tuple(sort([ci1,ci2])...)
      dist = abs( r1 - r2 )
      push!(expfile, (c1[1], first, second, g1, dist, c1[end]) )
      dinc!( bygenesym, ASCIIString(g1) )
      genesum += 1
   end
   expfile, bygenesym, genesum
end

function print_bed( io, reg )
   for bedtup in reg
      bedstr = join(bedtup, '\t')
      println( io , bedstr )
   end
end

function print_pvals( io, pvals::Dict{ASCIIString,Float64} )
   for k in keys(pvals)
      println( io, "$k\t$(pvals[k])" )
   end
end

function load_gene_exp( filename, keyset; header=true, ind=2 )
   df = readtable( filename, separator='\t', header=header)
   rename!( df, Dict{Symbol,Symbol}( names(df)[ind] => :GENESYM ) )
   nf = DataFrame( GENESYM=keyset )
   df = join( df, nf, on = :GENESYM )
   num_inds = get_numeric_cols( df )
   divsum( arr ) = arr / sum(arr)
   for i in num_inds
      df[i] = divsum( df[i] )
   end
   da = Vector{Float64}()
   for r in eachrow( df[num_inds] )
      push!(da, mean(convert(Array, r)))
   end
   df[:mean] = divsum( da )
   geneDic = Dict{ASCIIString,Float64}()
   for i in 1:nrow(df)
      geneDic[ df[i,ind] ] = df[i,:mean]
   end
   geneDic
end

function get_numeric_cols( df::DataFrame )
   num_inds = Int[]
   for i in 1:length(df)
      if typeof(df[i][1]) <: Real
         push!(num_inds, i)
      end
   end
   num_inds
end

function main()
   pargs = parse_cmd()

   reg,genecnt,genesum = read_lig_parse_reg( STDIN, biot=pargs["biot"] )
   if pargs["bed"]
      #print_bed( STDOUT, reg )
   end

   if pargs["background"] != nothing
      back = load_gene_exp( pargs["background"], collect(keys(genecnt)), header=pargs["backheader"] )
      pvals = Dict{ASCIIString,Float64}()
      for k in keys(genecnt)
         haskey(back, k) || continue
         bin = Binomial(Int(genesum), back[k])
         pvals[ k ] = ccdf( bin, genecnt[k]-1 )
      end
      n = length(keys(pvals))
      for k in keys(pvals)
         pvals[k] = min( pvals[k] * n, 1 )
      end
      open( "n_vs_ge-back.pval", "w+" ) do fh
         print_pvals( fh, pvals )
      end
   end

end

main()

max_forks_repo_name: timbitz/aligater