module CommonToAll
using PyPlot, StatsBase, Statistics, Distances, LinearAlgebra,
      DelimitedFiles, ..AbstractOperator, NearestNeighbors, CSV, DataFrames

import ..Options, ..OptionsStat, ..OptionsNonstat, ..OptionsNuisance,
       ..history, ..GP.κ, ..calcfstar!, ..AbstractOperator.Sounding

export trimxft, assembleTat1, gettargtemps, checkns, getchi2forall, nicenup,
        plot_posterior, make1Dhist, make1Dhists, setupz, zcontinue, makezρ, plotdepthtransforms,
        unwrap, getn, geomprogdepth, assemblemodelsatT, getstats, gethimage,
        assemblenuisancesatT, makenuisancehists, stretchexists,
        makegrid, whichislast, makesummarygrid, makearray, plotNEWSlabels,
        plotprofile, gridpoints, splitsoundingsbyline, dfn2hdr

function trimxft(opt::Options, burninfrac::Float64, temperaturenum::Int)
    x_ft = assembleTat1(opt, :x_ftrain, burninfrac=burninfrac, temperaturenum=temperaturenum)
    n = assembleTat1(opt, :nodes, burninfrac=burninfrac, temperaturenum=temperaturenum)
    x, ft = zeros(size(opt.xall, 1), sum(n)), zeros(size(opt.fbounds, 1), sum(n))
    nlast = 0
    for (i, xft) in enumerate(x_ft)
        x[:,nlast+1:nlast+n[i]]  = xft[1:size(opt.xall, 1), 1:n[i]]
        ft[:,nlast+1:nlast+n[i]] = xft[size(opt.xall, 1)+1:end, 1:n[i]]
        nlast += n[i]
    end
    x, ft, n
end

function assembleTat1(optin::Options, stat::Symbol; burninfrac=0.5, temperaturenum=-1)
    @assert temperaturenum!=-1 "Please explicitly specify which temperature number"
    if stat == :fstar && temperaturenum!= 1
        return assemblemodelsatT(optin, burninfrac=burninfrac, temperaturenum=temperaturenum)
    end
    isns = checkns(optin)
    @assert 0.0<=burninfrac<=1.0
    Tacrosschains = gettargtemps(optin)
    temps = sort(unique(Tacrosschains))
    niters = size(Tacrosschains, 1)
    start = round(Int, niters*burninfrac)
    start == 0 && (start = 1)
    ttarget = temps[temperaturenum]
    nmodels = sum((Tacrosschains[start:end,:] .== ttarget))
    if stat == :fstar || stat == :x_ftrain
        mat1 = Array{Array{Float64}, 1}(undef, nmodels)
    else
        mat1 = Array{Real, 1}(undef, nmodels)
    end
    opt = deepcopy(optin)
    imodel = 0
    for ichain in 1:size(Tacrosschains, 2)
        opt.fstar_filename = "models_"*opt.fdataname*isns*"_$ichain.bin"
        opt.x_ftrain_filename = "points_"*opt.fdataname*isns*"_$ichain.bin"
        opt.costs_filename = "misfits_"*opt.fdataname*isns*"_$ichain.bin"
        if stat == :fstar
            at1idx = findall(Tacrosschains[:,ichain].==ttarget) .>= start
        else
            at1idx = Tacrosschains[:,ichain].==ttarget
            at1idx[1:start-1] .= false
        end
        ninchain = sum(at1idx)
        @info "chain $ichain has $ninchain models"
        ninchain == 0 && continue
        mat1[imodel+1:imodel+ninchain] .= history(opt, stat=stat)[at1idx]
        imodel += ninchain
    end
    iters = history(opt, stat=:iter)
    @info "obtained models $(iters[start]) to $(iters[end]) at T=$ttarget"
    mat1
end

function assemblemodelsatT(opt::OptionsStat; burninfrac=0.9, temperaturenum=-1)
    @assert temperaturenum!=-1 "Please explicitly specify which temperature number"
    x_ft = assembleTat1(opt, :x_ftrain, burninfrac=burninfrac, temperaturenum=temperaturenum)
    n = assembleTat1(opt, :nodes, burninfrac=burninfrac, temperaturenum=temperaturenum)
    nmodels = length(n)
    matT = Array{Array{Float64}, 1}(undef, nmodels)
    K_y = zeros(opt.nmax, opt.nmax)
    Kstar = zeros(Float64, size(opt.xall,2), opt.nmax)
    xtest = opt.xall
    for imodel = 1:nmodels
        xtrain = @view x_ft[imodel][1:size(opt.xall, 1), 1:n[imodel]]
        ftrain = x_ft[imodel][size(opt.xall, 1)+1:end, :]
        ky = view(K_y, 1:n[imodel], 1:n[imodel])
        map!(x->κ(opt.K, x),ky,pairwise(WeightedEuclidean(1 ./opt.λ² ), xtrain, dims=2))
        K_y[diagind(K_y)] .+= opt.δ^2
        ks = view(Kstar, :, 1:n[imodel])
        map!(x->κ(opt.K, x),Kstar,pairwise(WeightedEuclidean(1 ./opt.λ² ), xtest, xtrain, dims=2))
        matT[imodel] = zeros(size(opt.fbounds, 1), size(opt.xall, 2))
        fstar = matT[imodel]
        calcfstar!(fstar, ftrain, opt, K_y, Kstar, n[imodel])
    end
    matT
end

function assemblenuisancesatT(optn::OptionsNuisance;
    burninfrac = 0.5, temperaturenum = -1)
    @assert temperaturenum != -1 "Please specify temperature idx explicitly"
    @assert 0.0 <= burninfrac < 1.0
    Tacrosschains = gettargtemps(optn)
    #this is probably insanely inefficient
    #Θ(niters*nchains) to run the unique() call.
    #however, we need to iterate over the entire array
    #at some point to build the ensemble so it's
    #fine. Possible improvement: store temperatures
    #as the "temperature number" to start with and keep
    #a separate mapping from temperature index to the float
    #value. This means our sortedTs is just 1:nchains and
    #also avoids any issues with equality testing of floats.
    sortedTs = sort(unique(Tacrosschains))
    niters = size(Tacrosschains,1)
    #this will never give a bounds error
    #because of the assert above
    firsti = round(Int, niters*burninfrac)
    firsti == 0 && (start = 1)
    # firsti = 1 + floor(Int, niters*burninfrac)
    ttarg = sortedTs[temperaturenum]
    nmodels = sum(Tacrosschains[firsti:end,:] .== ttarg)
    fdataname = optn.fdataname
    #drop iteration number
    mvals = zeros(nmodels,optn.nnu)
    imodel = 0
    for ichain = 1:size(Tacrosschains,2)
        at1idx = Tacrosschains[:,ichain].==ttarg
        at1idx[1:firsti-1] .= false
        ninchain = sum(at1idx)
        @info "chain $ichain has $ninchain models"
        ninchain == 0 && continue
        vals_filename = "values_nuisance_"*fdataname*"$ichain.bin"
        ndat = readdlm(vals_filename, ' ', Float64)[:,2:end]
        mvals[imodel+1:imodel+ninchain,:] .= ndat[at1idx,:]
        imodel += ninchain
    end
    mvals
end

function getnchains(costs_filename)
c = 0
    for fname in readdir(pwd())
           r = Regex(costs_filename)
           c += !isa(match(r, fname), Nothing)
    end
    c
end

function gettargtemps(opt_in::Options)
    isns = checkns(opt_in)
    opt = deepcopy(opt_in)
    costs_filename = "misfits_"*opt.fdataname*isns
    nchains = getnchains(costs_filename)
    @info "Number of chains is $nchains"
    # now look at any chain to get how many iterations
    opt.costs_filename    = costs_filename*"_1.bin"
    iters          = history(opt, stat=:iter)
    niters         = length(iters)
    @info "McMC has run for $(iters[end]) iterations"
    # then create arrays of unsorted by temperature T
    Tacrosschains  = zeros(Float64, niters, nchains)
    # get the values into the arrays
    for ichain in 1:nchains
        opt.costs_filename = costs_filename*"_$ichain.bin"
        Tacrosschains[:,ichain] = history(opt, stat=:T)
    end
    Tacrosschains
end

function gettargtemps(optn_in::OptionsNuisance)
    optn = deepcopy(optn_in)
    costs_filename = "misfits_"*optn.fdataname*"nuisance"
    nchains = getnchains(costs_filename)
    @info "Number of chains is $nchains"
    optn.costs_filename = costs_filename*"_1.bin"
    iters = history(optn, stat=:iter)
    niters = length(iters)
    @info "MCMC has run for $(iters[end]) iterations"
    Tacrosschains = zeros(Float64, niters, nchains)
    for ichain in 1:nchains
        optn.costs_filename = costs_filename*"_$ichain.bin"
        Tacrosschains[:,ichain] = history(optn, stat=:T)
    end
    Tacrosschains
end

function getstats(optin::Options;
                  figsize=(5,6), fontsize=12,
                  nxticks=5, nyticks=5, alpha=0.6)
    isns = checkns(optin)
    opt = deepcopy(optin)
    costs_filename = "misfits_"*opt.fdataname*isns
    nchains = getnchains(costs_filename)
    chains = 1:nchains
    @info "Number of chains is $nchains"
    opt.costs_filename    = costs_filename*"_1.bin"
    iters          = history(opt, stat=:iter)
    statnames = [:acceptanceRateBirth, :acceptanceRateDeath,
                 :acceptanceRatePosition, :acceptanceRateProperty, :acceptanceRateDC]
    f,ax = plt.subplots(length(statnames), 1,
                        sharex=true, sharey=true, figsize=figsize)
    maxar = 0
    for (ichain, chain) in enumerate(chains)
        opt.costs_filename = costs_filename*"_$chain.bin"
        for (istat, statname) in enumerate(statnames)
            ar = history(opt, stat=statname)
            mx = maximum(ar[.!isnan.(ar)])
            mx > maxar && (maxar = mx)
            ax[istat].plot(iters, ar, alpha=alpha)
            if ichain == nchains
                ax[istat].set_title("$statname "*isns)
                ax[istat].set_ylabel("acc. %")
                if istat == length(statnames)
                    ax[istat].set_xlabel("mcmc step no.")
                    ax[istat].set_xticks(LinRange(iters[1], iters[end], nxticks))
                    ax[istat].set_xlim(extrema(iters))
                    ax[istat].set_yticks(LinRange(0, maxar, nyticks))
                    ax[istat].set_ylim(0, maxar)
                end
            end
        end
    end
    nicenup(f, fsize=fontsize)
end

function getstats(optin::OptionsNuisance;
                  figsize=(5,6), fontsize=12,
                  nxticks=5, nyticks=5, alpha=0.6)
    opt = deepcopy(optin)
    costs_filename = "misfits_"*opt.fdataname*"nuisance"
    nchains = getnchains(costs_filename)
    chains = 1:nchains
    @info "Number of chains is $nchains"
    opt.costs_filename = costs_filename*"_1.bin"
    iters          = history(opt, stat=:iter)
    statname = :acceptanceRate
    f,ax = plt.subplots(length(optin.idxnotzero), 1, sharex=true, sharey=true, figsize=figsize)
    maxar = 0
    for (ichain, chain) in enumerate(chains)
        opt.costs_filename = costs_filename*"_$chain.bin"
        ar = history(opt, stat=statname)[:,optin.idxnotzero]
        for (i, idx) in enumerate(optin.idxnotzero)
            mx = maximum(ar[.!isnan.(ar)])
            mx > maxar && (maxar = mx)
            ax[i].plot(iters, ar[:,i], alpha=alpha)
            if ichain == nchains
                ax[i].set_title("Nuisance #$idx")
                ax[i].set_ylabel("acc. %")
                if i == length(optin.idxnotzero)
                    ax[i].set_xlabel("mcmc step no.")
                    ax[i].set_xticks(LinRange(iters[1], iters[end], nxticks))
                    ax[i].set_xlim(extrema(iters))
                end
            end
        end
    end
    nicenup(f, fsize=fontsize)
end

function checkns(optin::Options)
    isns = typeof(optin) == OptionsNonstat
    @info "ns is $isns"
    ns = "ns"
    isns || (ns="s")
    return ns
end

function getchi2forall(optn_in::OptionsNuisance;
                       nchains = 1,
                       figsize = (6,4),
                       fsize = 8,
                       alpha = 0.25,
                       nxticks = 0,
                       gridon = false)
    optn = deepcopy(optn_in)
    fdataname = optn.fdataname
    costs_filename = "misfits_"*fdataname*"nuisance"
    if nchains == 1
        nchains = getnchains(costs_filename)
    end
    optn.costs_filename = costs_filename*"_1.bin"
    iters = history(optn, stat=:iter)
    niters = length(iters)
    Tacrosschains = zeros(Float64, niters, nchains)
    χ2acrosschains = zeros(Float64, niters, nchains)
    for chain in 1:nchains
        optn.costs_filename = costs_filename*"_$(chain).bin"
        Tacrosschains[:,chain] = history(optn, stat=:T)
        χ2acrosschains[:,chain] = history(optn, stat=:misfit)
    end
    Torder = sort([(i,j) for i=1:niters, j=1:nchains],
                    by = ix->Tacrosschains[ix...],
                    dims = 2)
    χ2sorted = [χ2acrosschains[Torder[i,j]...] for i=1:niters, j=1:nchains]
    fig, ax = subplots(3,1, sharex=true, figsize=figsize)
    ax[1].plot(iters, χ2acrosschains, alpha=alpha)
    ax[1].set_xlim(extrema(iters)...)
    ax[1].set_ylim(0,100)
    ax[1].set_title("unsorted χ^2 misfit")
    ax[2].plot(iters, Tacrosschains, alpha=alpha)
    ax[2].set_title("temperature")
    ax[3].plot(iters, χ2sorted, alpha=alpha)
    ax[3].set_ylim(0,100)
    ax[3].set_title("sorted χ^2 misfit")
    ax[3].set_xlabel("iterations")

    nxticks == 0 || ax[3].set_xticks(iters[1]:div(iters[end],nxticks):iters[end])
    nicenup(fig, fsize = fsize)
end

function getchi2forall(opt_in::Options;
                        nchains         = 1,
                        figsize         = (6,4),
                        fsize           = 8,
                        alpha           = 0.25,
                        nxticks         = 0,
                        gridon          = false,
                        omittemp        = false,
                      )
    # now look at any chain to get how many iterations
    isns = checkns(opt_in)
    opt = deepcopy(opt_in)
    costs_filename = "misfits_"*opt.fdataname*isns
    if nchains == 1 # then actually find out how many chains there are saved
        nchains = getnchains(costs_filename)
    end
    opt.costs_filename    = costs_filename*"_1.bin"
    iters          = history(opt, stat=:iter)
    niters         = length(iters)
    # then create arrays of unsorted by temperature T, k, and chi2
    Tacrosschains  = zeros(Float64, niters, nchains)
    kacrosschains  = zeros(Int, niters, nchains)
    X2by2inchains  = zeros(Float64, niters, nchains)
    # get the values into the arrays
    for ichain in 1:nchains
        opt.costs_filename = costs_filename*"_$ichain.bin"
        Tacrosschains[:,ichain] = history(opt, stat=:T)
        kacrosschains[:,ichain] = history(opt, stat=:nodes)
        X2by2inchains[:,ichain] = history(opt, stat=:U)
    end

    f, ax = plt.subplots(3,1, sharex=true, figsize=figsize)
    ax[1].plot(iters, kacrosschains, alpha=alpha)
    ax[1].set_xlim(extrema(iters)...)
    ax[1].set_title(isns*" unsorted by temperature")
    gridon && ax[1].grid()
    ax[1].set_ylabel("# nodes")
    ax[2].plot(iters, X2by2inchains, alpha=alpha)
    gridon && ax[2].grid()
    ax[2].set_ylabel("-Log L")
    gridon && ax[3].grid()
    ax[3].plot(iters, Tacrosschains, alpha=alpha)
    ax[3].set_ylabel("Temperature")
    ax[3].set_xlabel("iterations")
    nxticks == 0 || ax[3].set_xticks(iters[1]:div(iters[end],nxticks):iters[end])
    nicenup(f, fsize=fsize)

    Tunsorted = copy(Tacrosschains)
    for jstep = 1:niters
        sortidx = sortperm(vec(Tacrosschains[jstep,:]))
        X2by2inchains[jstep,:] = X2by2inchains[jstep,sortidx]
        kacrosschains[jstep,:] = kacrosschains[jstep,sortidx]
        Tacrosschains[jstep,:] = Tacrosschains[jstep,sortidx]
    end
    nrows = omittemp ? 2 : 3
    f, ax = plt.subplots(nrows, 1, sharex=true, figsize=figsize)
    nchainsatone = sum(Tacrosschains[1,:] .== 1)
    ax[1].plot(iters, kacrosschains, alpha=alpha)
    ax[1].set_xlim(extrema(iters)...)
    ax[1].set_ylabel("# nuclei")
    ax[1].set_title(isns*" sorted by temperature")
    ax[1].plot(iters, kacrosschains[:,1:nchainsatone], "k", alpha=alpha)
    gridon && ax[1].grid()
    ax[2].plot(iters, X2by2inchains, alpha=alpha)
    ax[2].plot(iters, X2by2inchains[:,1:nchainsatone], "k", alpha=alpha)
    ax[2].set_ylabel("-Log L")
    gridon && ax[2].grid()
    if !omittemp
        ax[3].plot(iters, Tunsorted, alpha=alpha, color="gray")
        ax[3].set_ylabel("Temperature")
        gridon && ax[3].grid()
    end
    nxticks == 0 || ax[3].set_xticks(iters[1]:div(iters[end],nxticks):iters[end])
    ax[nrows].set_xlabel("iterations")
    nicenup(f, fsize=fsize)

end

function geomprogdepth(n, dy, c)
    dy*(1.0-c^n)/(1-c)
end

function getn(z, dy, c)
    log(1 - z/dy*(1-c))/log(c)
end

function unwrap(v, inplace=false)
  # currently assuming an array
  unwrapped = inplace ? v : copy(v)
  for i in 2:length(v)
    while unwrapped[i] - unwrapped[i-1] >= pi
      unwrapped[i] -= 2pi
    end
    while unwrapped[i] - unwrapped[i-1] <= -pi
      unwrapped[i] += 2pi
    end
  end
  return unwrapped
end

function setupz(zstart, extendfrac;n=100, dz=10, showplot=false, forextension=false)
    @assert extendfrac>1
    znrange        = 1.0:n
    zboundaries    = [zstart, zstart .+ geomprogdepth.(znrange, dz, extendfrac)...]
    thickness      = diff(zboundaries)
    zall           = zboundaries[1:end-1] + thickness/2
    znall          = getn.(zall .- zstart, dz, extendfrac)
    showplot && plotdepthtransforms(zall, znall, zboundaries)
    zbreturn = forextension ? zboundaries : zboundaries[1:end-1]
    return zall, znall, zbreturn
end

function zcontinue(;zall=nothing, znall=nothing, zboundaries=nothing, n=nothing,
                extendfrac=nothing, dz=nothing, showplot=false)
    isa(dz, Nothing) && (dz=diff(zboundaries)[end])
    zall_, znall_, zboundaries_ = setupz(zboundaries[end], extendfrac, n=n, dz=dz, showplot=false)
    zall, znall, zboundaries = [zall;zall_], [znall; znall[end] .+ znall_], [zboundaries[1:end-1]; zboundaries_]
    zlast = zboundaries[end] + diff(zboundaries)[end]
    showplot && plotdepthtransforms(zall, znall, [zboundaries;zlast])
    @assert all(diff(zall).>0)
    @assert all(diff(znall).>0)
    @assert all(diff(zboundaries).>0)
    zall, znall, zboundaries
end

function setupz(zstart;n=100, dz=10)
    zboundaries    = range(zstart, step=dz, length=n)
    zall           = zboundaries[1:end-1] .+ dz/2
    znall          = 1.5:1:n-0.5
    plotdepthtransforms(zall, znall, zboundaries)
    return zall, znall, zboundaries[1:end-1]
end

function plotdepthtransforms(zall, znall, zboundaries)
    figure()
    plot(znall, zall)
    xlabel("depth index")
    ylabel("depth associated m")
    grid()
    nicenup(gcf())

    f, ax = plt.subplots(1, 1, figsize=(10,5))
    ax.stem(zboundaries[1:end-1], zboundaries[1:end-1], markerfmt="")
    ax.stem(zall, zall, "k--", markerfmt=" ")
    ax.set_xlabel("depth m")
    ax.set_ylabel("depth m")
    nicenup(gcf())

    thickness = diff(zboundaries)
    f, ax = plt.subplots(1, 2, sharey=true, figsize=(11,5))
    ax[1].stem(zall,thickness, "k--", markerfmt=" ")
    ax[1].set_ylabel("thickness m")
    ax[1].set_xlabel("depth m")
    ax[1].yaxis.grid(which="major")
    ax[2].stem(znall,thickness, "k--", markerfmt=" ")
    ax[2].set_xlabel("depth index")
    ax[2].yaxis.grid(which="major")
    nicenup(gcf())
end

function makezρ(zboundaries::Array{Float64, 1};
          zfixed      = [-1e6,    0.],
          ρfixed      = [1e12     0.3])
    @assert length(zfixed) == length(ρfixed)
    z = [zfixed..., zboundaries...]
    @assert all(diff(z).>0)
    ρ = zeros(length(z))
    nfixed = length(ρfixed)
    ρ[1:nfixed] .= ρfixed
    z, ρ, nfixed
end

function nicenup(g::PyPlot.Figure;fsize=16)
    for ax in gcf().axes
        ax.tick_params("both",labelsize=fsize)
        ax.xaxis.label.set_fontsize(fsize)
        ax.yaxis.label.set_fontsize(fsize)
        any(keys(ax) .== :zaxis) && ax.zaxis.label.set_fontsize(fsize)
        ax.title.set_fontsize(fsize)

        if typeof(ax.get_legend_handles_labels()[1]) != Array{Any,1}
            ax.legend(loc="best", fontsize=fsize)
        end
    end
    g.tight_layout()
end

function plot_posterior(F::Operator1D,
                        optns::OptionsNonstat,
                        opts::OptionsStat;
                        temperaturenum = 1,
                        nbins = 50,
                        burninfrac=0.5,
                        qp1=0.05,
                        qp2=0.95,
                        vmaxpc=1.0,
                        cmappdf = "inferno",
                        figsize=(12,5),
                        pdfnormalize=false,
                        fsize=14,
                        showlscale1vd=false,
                        istothepow=false,
                        doplot = true,
                        a = 0.25,
                        lw = 1)
    @assert 0<vmaxpc<=1
    M = assembleTat1(optns, :fstar, burninfrac=burninfrac, temperaturenum=temperaturenum)
    himage_ns, edges_ns, CI_ns, meanimage_ns,
    meandiffimage_ns, sdslope_ns, = gethimage(F, M, optns, temperaturenum=temperaturenum,
                            nbins=nbins, qp1=qp1, qp2=qp2, istothepow=istothepow,
                            islscale=false, pdfnormalize=pdfnormalize)
    M = assembleTat1(opts, :fstar, burninfrac=burninfrac, temperaturenum=temperaturenum)
    himage, edges, CI, meanimage,
    meandiffimage, sdslope, = gethimage(F, M, opts, temperaturenum=temperaturenum,
                            nbins=nbins, qp1=qp1, qp2=qp2, istothepow=false,
                            islscale=true, pdfnormalize=pdfnormalize)
    f,ax = plt.subplots(1, 3+convert(Int, showlscale1vd), sharey=true, figsize=figsize)
    xall = opts.xall
    diffs = diff(xall[:])
    xmesh = vcat(xall[1:end-1] - diffs/2, xall[end]-diffs[end]/2, xall[end])
    vmin, vmax = extrema(himage_ns)
    vmax = vmin+vmaxpc*(vmax-vmin)
    im1 = ax[1].pcolormesh(edges_ns[:], xmesh, himage_ns, cmap=cmappdf, vmax=vmax)
    cb1 = colorbar(im1, ax=ax[1])
    cb1.ax.set_xlabel("pdf \nns")
    @show propmin, propmax = getbounds(CI_ns, optns.fbounds)
    ax[1].set_xlim(propmin, propmax)
    ax[2].fill_betweenx(xall[:],meandiffimage_ns[:]-sdslope_ns[:],meandiffimage_ns[:]+sdslope_ns[:],alpha=a, color="gray")
    ax[2].plot(meandiffimage_ns, xall[:], linewidth=lw, color="k")
    ax[2].set_xlabel("slope")

    vmin, vmax = extrema(himage)
    vmax = vmin+vmaxpc*(vmax-vmin)
    im3 = ax[3].pcolormesh(edges[:], xmesh, himage, cmap=cmappdf, vmax=vmax)
    ax[3].set_ylim(extrema(xall)...)
    propmin, propmax = getbounds(CI, opts.fbounds)
    ax[3].set_xlim(propmin, propmax)
    ax[3].invert_yaxis()
    cb3 = colorbar(im3, ax=ax[3])
    cb3.ax.set_xlabel("pdf \nstationary")

    if showlscale1vd
        ax[4].plot(meandiffimage[:], xall[:], linewidth=2, color="g", alpha=0.5)
        ax[4].fill_betweenx(xall[:], meandiffimage[:]-sdslope[:],meandiffimage[:]+sdslope[:], alpha=.5)
        ax[4].set_xlabel("slope")
    end

    ax[1].plot(CI_ns, xall[:], linewidth=lw, color="w", alpha=a)
    # ax[1].plot(CI_ns, xall[:], linewidth=1, color="c", linestyle="--", alpha=0.5)
    # ax[1].plot(meanimage_ns[:], xall[:], linewidth=2, color="w", alpha=0.5)
    # ax[1].plot(meanimage_ns[:], xall[:], linewidth=2, color="k", linestyle="--", alpha=0.5)
    ax[3].plot(CI, xall[:], linewidth=lw, color="w", alpha=a)
    # ax[3].plot(CI, xall[:], linewidth=1, color="c", linestyle="--", alpha=0.5)
    # ax[3].plot(meanimage[:], xall[:], linewidth=2, color="w", alpha=0.5)
    # ax[3].plot(meanimage[:], xall[:], linewidth=2, color="k", linestyle="--", alpha=0.5)
    ax[1].set_xlabel(L"\log_{10} \rho")
    ax[1].set_ylabel("depth (m)")
    ax[3].set_xlabel(L"\log_{10} λ")
    nicenup(f, fsize=fsize)
end

function plot_posterior(F::Operator1D,
                        opt::OptionsStat;
                    temperaturenum = 1,
                    nbins = 50,
                    burninfrac=0.5,
                    qp1=0.05,
                    qp2=0.95,
                    vmaxpc=1.0,
                    cmappdf = "inferno",
                    figsize=(5,5),
                    pdfnormalize=false,
                    istothepow=false,
                    fsize=14,
                    plotCI = true,
                    plotmean = true,
                    alpha = 1.0,
                    CIcolor = ["w", "k"],
                    meancolor = ["m", "r"],
                    lwidth = 2,
                    doplot = true)
    @assert 0<vmaxpc<=1

    M = assembleTat1(opt, :fstar, burninfrac=burninfrac, temperaturenum=temperaturenum)
    himage, edges, CI, meanimage, meandiffimage, sdslope, = gethimage(F, M, opt, burninfrac=burninfrac, temperaturenum=temperaturenum,
                nbins=nbins, qp1=qp1, qp2=qp2, istothepow=istothepow,
                islscale=false, pdfnormalize=pdfnormalize)

    if doplot
        f,ax = plt.subplots(1,2, sharey=true, figsize=figsize)
        xall = opt.xall
        diffs = diff(xall[:])
        xmesh = vcat(xall[1:end-1] - diffs/2, xall[end]-diffs[end]/2, xall[end])
        vmin, vmax = extrema(himage)
        vmax = vmin+vmaxpc*(vmax-vmin)
        im1 = ax[1].pcolormesh(edges[:], xmesh, himage, cmap=cmappdf, vmax=vmax)
        ax[1].set_ylim(extrema(xall)...)
        ax[1].invert_yaxis()
        cb1 = colorbar(im1, ax=ax[1])
        cb1.ax.set_xlabel("pdf \nstationary")
        plotCI && ax[1].plot(CI, xall[:], linewidth=lwidth, color=CIcolor[1], alpha=alpha)
        plotCI && ax[1].plot(CI, xall[:], linewidth=lwidth, color=CIcolor[2], linestyle="--", alpha=alpha)
        plotmean && ax[1].plot(meanimage[:], xall[:], linewidth=lwidth, color=meancolor[1], alpha=alpha)
        plotmean && ax[1].plot(meanimage[:], xall[:], linewidth=lwidth, color=meancolor[2], linestyle="--", alpha=alpha)
        ax[1].set_xlabel(L"\log_{10} \rho")
        ax[1].set_ylabel("depth (m)")
        bounds = copy(opt.fbounds)
        if stretchexists(F)
            bounds = [minimum(F.low) maximum(F.low + F.Δ)]
        end
        istothepow && (bounds .= 10 .^ bounds)
        propmin, propmax = getbounds(CI, bounds)
        ax[1].set_xlim(propmin, propmax)
        ax[2].plot(meandiffimage[:], xall[:], linewidth=2, color="k", linestyle="-")
        zeroside = meandiffimage[:]-sdslope[:]
        zeroside[zeroside .< 0] .= 0
        ax[2].fill_betweenx(xall[:],zeroside,meandiffimage[:]+sdslope[:],alpha=.25)
        ax[2].set_xlabel("mean slope")
        nicenup(f, fsize=fsize)
    end
    CI[:,1], CI[:,2], CI[:,3], meanimage, meandiffimage, sdslope
end

function getbounds(CI, bounds)
    propmin = min(minimum(CI), minimum(bounds))
    propmax = max(maximum(CI), maximum(bounds))
    propmin, propmax
end

function plot_posterior(operator::Operator1D,
                        optn::OptionsNuisance;
                        temperaturenum = 1,
                        nbins = 50,
                        burninfrac=0.5,
                        figsize=(8,16),
                        pdfnormalize=false,
                        fsize=14,
                        qp1 = 0.05,
                        qp2 = 0.95,
                        doplot=true)
    hists, CI = makenuisancehists(optn, qp1, qp2, burninfrac = burninfrac,
                                     nbins = nbins, temperaturenum = temperaturenum,
                                     pdfnormalize = pdfnormalize)
    if doplot
        fig,ax = subplots(length(hists), 1, figsize=figsize)
        length(hists) == 1 && (ax = [ax])
        for (i, h) = enumerate(hists)
            bwidth = diff(h.edges[1])[1]
            bx = h.edges[1][1:end-1] .+ bwidth/2
            ax[i].bar(bx, h.weights, width=bwidth, edgecolor="black")
        end
        nicenup(fig, fsize=fsize)
    end
    hists, CI
end

function makenuisancehists(optn::OptionsNuisance, qp1, qp2; burninfrac = 0.5, nbins = 50,
    temperaturenum = -1, pdfnormalize = false)
    @assert temperaturenum != -1 "Please explicitly specify the temperature index"
    nuisanceatT = assemblenuisancesatT(optn, burninfrac = burninfrac,
                                    temperaturenum = temperaturenum)
    nnu = length(optn.idxnotzero)
    CI = zeros(Float64, nnu, 3)
    h = Vector{Histogram}(undef, nnu)
    for (i, islice) in enumerate(optn.idxnotzero)
        numin, numax = extrema(optn.bounds[islice,:])
        edges = LinRange(numin, numax, nbins+1)
        h[i] = fit(Histogram, nuisanceatT[:,islice], edges)
        CI[i,:] = quantile(nuisanceatT[:,islice], [qp1, 0.5, qp2])
    end
    h, CI
end

function firstderiv(x)
    fd = similar(x)
    fd[1]        =  x[2]     - x[1]
    fd[end]      =  x[end]   - x[end-1]
    fd[2:end-1] .= (x[3:end] - x[1:end-2])/2
    abs.(fd)
    # fd
end

function secondderiv(x)
    sd = similar(x)
    sd[1]        = 2x[1] - 5x[2] + 4x[3] - x[4]
    sd[2:end-1] .= (x[3:end] - x[2:end-1] + x[1:end-2])
    sd[end]      = 2x[end] - 5x[end-1] + 4x[end-2] - x[end-3]
    abs.(sd)
end

function gethimage(F::Operator, M::AbstractArray, opt::Options;
                burninfrac = 0.5,
                nbins = 50,
                rhomin=Inf,
                rhomax=-Inf,
                qp1=0.05,
                qp2=0.95,
                islscale = false,
                pdfnormalize=false,
                istothepow = false,
                temperaturenum=1)
    T = x->x
    if (rhomin == Inf) && (rhomax == -Inf)
        if !stretchexists(F) # if no stretch
            for (i,mm) in enumerate(M)
                rhomin_mm = minimum(mm)
                rhomax_mm = maximum(mm)
                rhomin_mm < rhomin && (rhomin = rhomin_mm)
                rhomax_mm > rhomax && (rhomax = rhomax_mm)
            end
            if (typeof(opt) == OptionsStat && opt.needλ²fromlog) && islscale
                T = x->0.5log10(x)
            end
        else # there is a stretch
            rhomin = minimum(F.low)
            rhomax = maximum(F.low + F.Δ)
        end
    else
        @assert rhomin < rhomax
    end
    if istothepow && !islscale
        T = x->10. ^x
    end
    rhomin, rhomax = map(x->T(x), (rhomin, rhomax))
    edges = LinRange(rhomin, rhomax, nbins+1)
    himage = zeros(Float64, length(M[1]), nbins)
    CI = zeros(Float64, length(M[1]), 3)
    meanimage = zeros(length(M[1]))
    for ilayer=1:length(M[1])
        if !stretchexists(F) # if no stretch
            mthislayer = [T(m[ilayer]) for m in M]
            himage[ilayer,:] = fit(Histogram, mthislayer, edges).weights
            CI[ilayer,:] = [quantile(mthislayer,(qp1, 0.5, qp2))...]
            meanimage[ilayer] = mean(vec(mthislayer))
        else # there is an affine stretch with depth
            mthislayer = [m[ilayer] for m in M]
            expandedthislayer = T.(F.low[ilayer] .+ mthislayer*F.Δ[ilayer])
            himage[ilayer,:] = fit(Histogram, expandedthislayer, edges).weights
            CI[ilayer,:] = [quantile(expandedthislayer,(qp1, 0.5, qp2))...]
            meanimage[ilayer] = mean(vec(expandedthislayer))
        end
        himage[ilayer,:] = himage[ilayer,:]/sum(himage[ilayer,:])/(diff(edges)[1])
        pdfnormalize && (himage[ilayer,:] = himage[ilayer,:]/maximum(himage[ilayer,:]))
    end
    if !stretchexists(F)
        Tm = [T.(m) for m in M]
    else
        Tm = [T.(F.low + vec(m).*F.Δ) for m in M]
    end
    Mslope = mean(firstderiv.(Tm))
    sdevslope = std(firstderiv.(Tm))
    himage, edges, CI, meanimage, Mslope, sdevslope
end

function stretchexists(F::Operator)
    in(:stretch, fieldnames(typeof(F))) && F.stretch
end

# plotting codes for 2D sections in AEM

function splitsoundingsbyline(soundings::Array{S, 1}) where S<:Sounding
    alllines = [s.linenum for s in soundings]
    linenum  = unique(alllines)
    nlines   = length(linenum)
    linestartidx = zeros(Int, nlines)
    for i = 1:nlines
        linestartidx[i] = findfirst(alllines .== linenum[i])
    end
    linestartidx
end

function makegrid(vals::AbstractArray, soundings::Array{S, 1};
    dr=10, zall=[NaN], dz=-1) where S<:Sounding
    @assert all(.!isnan.(zall))
    @assert dz>0
    X = [s.X for s in soundings]
    Y = [s.Y for s in soundings]
    x0, y0 = X[1], Y[1]
    R  = sqrt.((X .- x0).^2 + (Y .- y0).^2)
    rr, zz = [r for z in zall, r in R], [z for z in zall, r in R]
    topo = [s.Z for s in soundings]
    zz = topo' .- zz # mAHD
    kdtree = KDTree([rr[:]'; zz[:]'])
    gridr = range(R[1], R[end], step=dr)
    gridz = reverse(range(extrema(zz)..., step=dz))
    rr, zz = [r for z in gridz, r in gridr], [z for z in gridz, r in gridr]
    idxs, = nn(kdtree, [rr[:]'; zz[:]'])
    img = zeros(size(rr))
    for i = 1:length(img)
        img[i] = vals[idxs[i]]
    end
    topofine = gridpoints(R, gridr, topo)
    img[zz .>topofine'] .= NaN
    img, gridr, gridz, topofine, R
end

function makegrid(vals::AbstractArray, soundings::Array{S, 1}, XYZrho;
    dr=10, zall=[NaN], dz=-1) where S<:Sounding
    @assert all(.!isnan.(zall))
    @assert dz>0
    X = [s.X for s in soundings]
    Y = [s.Y for s in soundings]
    x0, y0 = X[1], Y[1]
    R  = sqrt.((X .- x0).^2 + (Y .- y0).^2)
    rr, zz = [r for z in zall, r in R], [z for z in zall, r in R]
    topo = [s.Z for s in soundings]
    zz = topo' .- zz # mAHD
    kdtree = KDTree([rr[:]'; zz[:]'])
    gridr = range(R[1], R[end], step=dr)
    gridz = reverse(range(extrema(zz)..., step=dz))
    rr, zz = [r for z in gridz, r in gridr], [z for z in gridz, r in gridr]
    idxs, = nn(kdtree, [rr[:]'; zz[:]'])
    # comparison with ZYZrho
    x0, y0 = XYZrho[1,1], XYZrho[1,2]
    rrcompare = sqrt.((XYZrho[:,1] .- x0).^2 + (XYZrho[:,2] .- y0).^2)
    kdtreecompare = KDTree([rrcompare[:]'; XYZrho[:,3]'])
    idxscompare, = nn(kdtreecompare, [rr[:]'; zz[:]'])
    img, imgcompare = zeros(size(rr)), zeros(size(rr))
    for i = 1:length(img)
        img[i] = vals[idxs[i]]
        imgcompare[i] = XYZrho[idxscompare[i], 4]
    end
    topofine = gridpoints(R, gridr, topo)
    img[zz .>topofine'] .= NaN
    imgcompare[zz .> topofine'] .= NaN
    img, gridr, gridz, topofine, R, imgcompare
end

function gridpoints(R, gridr, points::Array{Float64, 1})
    @assert length(R) == length(points)
    kdtree = KDTree(R[:]')
    idxs, = nn(kdtree, gridr[:]')
    pointsgrid = [points[idxs[i]] for i = 1:length(gridr)]
end

function gridpoints(R, gridr, points::Array{Float64, 2})
    nsound = length(R)
    nprops = size(points, 1)
    @assert nsound == size(points, 2)
    pointsgrid = zeros(length(gridr), nprops)
    for i = 1:nprops
        pointsgrid[:,i] = gridpoints(R, gridr, vec(points[i,:]))
    end
    pointsgrid
end

function whichislast(soundings::AbstractArray)
    X = [s.X for s in soundings]
    Y = [s.Y for s in soundings]
    Eislast, Nislast = true, true
    X[1]>X[end] && (Eislast = false)
    Y[1]>Y[end] && (Nislast = false)
    Eislast, Nislast
end

function makesummarygrid(soundings, pl, pm, ph, ρmean, vdmean, vddev, zall, dz; dr=10)
    # first flip ρ to σ and so pl and ph are interchanged
    phgrid, gridx, gridz, topofine, R = makegrid(-pl, soundings, zall=zall, dz=dz, dr=dr)
    plgrid,                           = makegrid(-ph, soundings, zall=zall, dz=dz, dr=dr)
    pmgrid,                           = makegrid(-pm, soundings, zall=zall, dz=dz, dr=dr)
    σmeangrid,                        = makegrid(-ρmean, soundings, zall=zall, dz=dz, dr=dr)
    ∇zmeangrid,                       = makegrid(vdmean, soundings, zall=zall, dz=dz, dr=dr)
    ∇zsdgrid,                         = makegrid(vddev, soundings, zall=zall, dz=dz, dr=dr)
    Z = [s.Z for s in soundings]
    phgrid, plgrid, pmgrid, σmeangrid, ∇zmeangrid, ∇zsdgrid, gridx, gridz, topofine, R, Z
end

function makexyzrho(soundings, zall)
    linename = "_line_$(soundings[1].linenum)_summary.txt"
    fnames = ["rho_low", "rho_mid", "rho_hi", "rho_avg"].*linename
    dlow, dmid, dhigh, davg = map(x->readdlm(x), fnames)
    ndepths = length(zall)
    for (i, d) in enumerate([dlow, dmid, dhigh, davg])
        @assert ndims(d) == 2
        @assert size(d, 2) == length(soundings)
        @assert size(d, 1) == ndepths
        xyzrho = makearray(soundings, d, zall)
        writedlm(fnames[i][1:end-4]*"_xyzrho.txt", xyzrho)
    end
end

function makearray(soundings, d, zall)
    outarray = zeros(length(d), 4)
    ndepths = length(zall)
        for (i, s) in enumerate(soundings)
            z = s.Z .- zall # convert depths to mAHD
            x, y = s.X, s.Y
            idx = (i-1)*ndepths+1:i*ndepths
            outarray[idx,:] = [[x y].*ones(ndepths) z d[:,i]]
        end
    outarray
end

function plotNEWSlabels(Eislast, Nislast, gridx, gridz, axarray)
    for s in axarray
        Eislast ? s.text(gridx[1], minimum(s.get_ylim()), "W", backgroundcolor="w") : s.text(gridx[1], minimum(s.get_ylim()), "E", backgroundcolor="w")
        Nislast ? s.text(gridx[end], minimum(s.get_ylim()), "N", backgroundcolor="w") : s.text(gridx[end], minimum(s.get_ylim()), "S", backgroundcolor="w")
    end
end

function plotprofile(ax, idxs, Z, R)
    for idx in idxs
        ax.plot(R[idx]*[1,1], [ax.get_ylim()[1], Z[idx]], "-w")
        ax.plot(R[idx]*[1,1], [ax.get_ylim()[1], Z[idx]], "--k")
    end
end

function selectwithin1Dinterval(M::AbstractVector, z,
                                zbounds, vbounds,
                                cond = :mean)
    @assert(any(cond .== [:all, :any, :median, :mean]))
    @assert size(vbounds, 1) == size(zbounds, 1)
    @assert length(z) == length(M[1])
    @assert all(vbounds[:,1] .< vbounds[:,2])
    @assert all(zbounds[:,1] .< zbounds[:,2])
    nconditions = size(vbounds, 1)
    idx = zeros(Bool, length(M), nconditions)
    for (i, m) in enumerate(M)
       for j in 1:nconditions
            idxdepth = zbounds[j,1] .< z .< zbounds[j,2]
            if any(cond .== [:median, :mean])
                if vbounds[j,1] < eval(cond)(m[idxdepth]) < vbounds[j,2]
                    idx[i,j] = true
                end
            else
                if eval(cond)(vbounds[j,1] .< (m[idxdepth]) .< vbounds[j,2])
                    idx[i,j] = true
                end
            end
       end
    end
    vec(reduce(&, idx, dims=2))
end

function block1Dvalues(M::AbstractVector, z, zbounds, cond = :median)
    @assert(any(cond .== [:median, :mean]))
    @assert length(z) == length(M[1])
    @assert all(zbounds[:,1] .< zbounds[:,2])
    nconditions = size(zbounds, 1)
    Mblock = zeros(length(M), nconditions)
    for (i, m) in enumerate(M)
        for j in 1:nconditions
            idxdepth = zbounds[j,1] .< z .<= zbounds[j,2]
            Mblock[i,j] = eval(cond)(m[idxdepth])
        end
     end
    Mblock
end

function correlationplot(M::Array{Float64, 2}; figsize=(5,5), nbins=25)
    f, ax = plt.subplots(size(M, 2), size(M, 2), figsize=figsize, sharex=true, sharey=true)
    nrows = size(M, 2)
    for j = 1:nrows
        for i = j:nrows
            h = normalize(fit(Histogram, (M[:,i], M[:,j]), nbins=nbins))
            ax[nrows*(j-1)+i].pcolormesh(h.edges[2], h.edges[1], h.weights)
        end
    end
    ax[1].set_aspect(1)
end

function plot1Dpatches(ax, zlist, xlist; alpha=0.25, fc="red", ec="blue", lw=2)
    @assert length(zlist) == length(xlist)
    for i in 1:size(zlist, 1)
        x0, z0 = xlist[i,1], zlist[i,1]
        delx = xlist[i,2] - x0
        delz = zlist[i,2] - z0
        ax.add_patch(matplotlib.patches.Rectangle((x0,z0),delx,delz, alpha=alpha, fc=fc, ec=fc))
        ax.add_patch(matplotlib.patches.Rectangle((x0,z0),delx,delz, fill=false, ec=ec, lw=lw, linestyle="--"))
    end
end

#filtering the projection added part from oasis
function filter_proj!(df)
    filter!(row -> !contains(row.combined, "RT=PROJ"), df)
end

#filtering the transformation added part from oasis
function filter_trans!(df)
    filter!(row -> !contains(row.combined,"RT=TRNS"), df)
end

#function to read the *dfn file and extract the column number and column names as a *.txt file
function dfn2hdr(dfnfile::String)
    df = CSV.read(dfnfile, DataFrame; header=false);
    df = coalesce.(df,"")

    colname = []
    for col in 1:length(names(df))
        push!(colname,names(df)[col])
    end
    t = df[!, colname[1]]
    for i in colname[2:end]
		t = t .* df[!, i]
	end
    df[!,"combined"] = t
    filter_proj!(df)
    filter_trans!(df)
    row_number = nrow(df)
    insertcols!(df, 1, :inc => 0)
    insertcols!(df, 1, :first => 0)
    insertcols!(df,1, :second => 0)

    cumulative_columns = 0  #this will set up a cumulative variable

    for i in 2:size(df,1)   #this will do operations in row level
        regex_literal1 = r"NAME=" #few parameters have their names after this keyword
        regex_literal2 = r"RT="  #four parameters have their names after this keyword
        if !contains(string(df[i,:].combined),regex_literal1)
            continue
        end

        if contains(df[i,:].combined,"RT:A4F")
            df[i,:].combined = replace(df[i,:].combined, ":A4F" => "AA")
        end

        inc_regex = r":([0-9]+)F"
        inc_match = match(inc_regex, df[i,:].combined)  #matching the keyword with the string
        if isnothing(inc_match)
            inc = 0
            idx2 =  idx2 = split(string(df[i,:].combined), regex_literal1)
            idx2_r = first.(split.(idx2[2], ":"))

        else
            inc = parse(Int64, inc_match.captures[1]) #the outcome is int64
            idx2 =  idx2 = split(string(df[i,:].combined), regex_literal2)
            idx2_r = first.(split.(idx2[2], ":"))

        end

        df[i,:].inc = inc
        df[i,:].first = cumulative_columns + 1
        df[i,:].second = df[i,:].first + inc
        cumulative_columns += inc + 1

        if occursin(";END DEFN", idx2_r)  #type is string
            idx2_r = first.(split.(idx2_r,";"))
        else
            idx2_r = first.(split.(idx2_r,"END"))
        end

        #start writing to a file here
        io = open("header.txt","a")
        if (df[i,:].first != df[i,:].second)
            writedlm(io,[df[i,:].first df[i,:].second idx2_r])
        else
            writedlm(io, [df[i,:].first idx2_r])
        end
        close(io)
    end
end

end # module CommonToAll


function TabularRegression(
        blocks::Tuple{<:TableRow, <:Continuous},
        data)
    tabledata, targetdata = data
    tabledata isa TableDataset || error("`data` needs to be a tuple of a `TableDataset` and targets")
    return BlockMethod(
        blocks,
        (setup(TabularPreprocessing, blocks[1], tabledata),),
        outputblock=blocks[2]
    )
end

"""
    TabularRegression(blocks, data)

Learning method for tabular regression. Continuous columns are
normalized and missing values are filled, categorical columns are label encoded
taking into account any missing values which might be present.
 `blocks` should be an input and target block `(TableRow(...), Continuous(...))`.

    TabularRegression(n, tabledata [; catcols, contcols])

Construct learning method with `classes` to classify into and a `TableDataset`
`tabledata`. The column names can be passed in or guessed from the data. The
regression target is a vector of `n` values.
"""
function TabularRegression(
        n::Int,
        tabledata::TableDataset;
        catcols = nothing,
        contcols = nothing)
    blocks = (
        setup(TableRow, tabledata; catcols=catcols, contcols=contcols),
        Continuous(n)
    )
    return TabularRegression(blocks, (tabledata, nothing))
end

# ## Tests

@testset "TabularRegression [method]" begin
    df = DataFrame(A = 1:4, B = ["M", "F", "F", "M"], C = 10:13)
    td = TableDataset(df)
    targets = [rand(2) for _ in 1:4]
    method = TabularRegression(2, td; catcols=(:B,), contcols=(:A,))
    testencoding(method.encodings, method.blocks)
    DLPipelines.checkmethod_core(method)
    @test_nowarn methodlossfn(method)
    @test_nowarn methodmodel(method)

    @testset "`encodeinput`" begin
        row = mockblock(method.blocks[1])

        xtrain = encodeinput(method, Training(), row)
        @test length(xtrain[1]) == length(method.blocks[1].catcols)
        @test length(xtrain[2]) == length(method.blocks[1].contcols)

        @test eltype(xtrain[1]) <: Number
    end

    @testset "`encodetarget`" begin
        target = 11
        y = encodetarget(method, Training(), target)
        @test target == y
    end

    @test_nowarn method = TabularRegression(2, td)

    FastAI.test_method_show(method, ShowText(Base.DevNull()))
end


module SatelliteImages
using Dates, ArchGDAL, GeoArrays, DataFrames, Shapefile, GeoInterface, ProgressMeter, SparseArrays, StatsBase, Statistics, RecursiveArrayTools, JLD, CubicSplines
# Write your package code here.

export load_img, load_datacube, save_img, save_datacube,
lat_to_row, row_to_lat, lat_to_row, long_to_column, column_to_long, translate_geometry, Coordinate, CoordinateSystem, image_to_coordinate,
polygon_mask, load_shapefile, mask_area, long_apply, cross_apply, apply_mask, view_img, aggregate, aggregate_timeseries, aggregate_dataframe, bias_correction,  sparse_cube, make_datacube, load_example, radiance_datacube, clouds_datacube, mumbai_map, MUMBAI_COORDINATE_SYSTEM, INDIA_COORDINATE_SYSTEM, coordinate_to_image, aggregate_per_area_dataframe

include("data_io.jl")
include("view_image.jl")
include("coordinate_system.jl")
include("polygons.jl")
include("area.jl")
include("f_apply.jl")
include("sparse_datacube.jl")
include("aggregate.jl")
include("example.jl")
include("masks.jl")
include("luxor_functions.jl")
end


char(x) = convert(Char, x)
char(x::FloatingPoint) = char(iround(x))

integer(x::Char) = int(x)

## char promotions ##

promote_rule(::Type{Char}, ::Type{Int8})    = Int32
promote_rule(::Type{Char}, ::Type{Uint8})   = Int32
promote_rule(::Type{Char}, ::Type{Int16})   = Int32
promote_rule(::Type{Char}, ::Type{Uint16})  = Int32
promote_rule(::Type{Char}, ::Type{Int32})   = Int32
promote_rule(::Type{Char}, ::Type{Uint32})  = Uint32
promote_rule(::Type{Char}, ::Type{Int64})   = Int64
promote_rule(::Type{Char}, ::Type{Uint64})  = Uint64
promote_rule(::Type{Char}, ::Type{Int128})  = Int128
promote_rule(::Type{Char}, ::Type{Uint128}) = Uint128

## character operations & comparisons ##

# numeric operations
# TODO: this should be removed, but needs to be here as long as Char <: Integer
+(x::Char   , y::Char   ) = int(x)+int(y)

# ordinal operations
+(x::Char   , y::Integer) = char(int(x)+int(y))
+(x::Integer, y::Char   ) = y+x
-(x::Char   , y::Char   ) = int(x)-int(y)
-(x::Char   , y::Integer) = char(int(x)-int(y))

# bitwise operations
(~)(x::Char) = char(~uint32(x))
(&)(x::Char, y::Char) = char(uint32(x) & uint32(y))
(|)(x::Char, y::Char) = char(uint32(x) | uint32(y))
($)(x::Char, y::Char) = char(uint32(x) $ uint32(y))

bswap(x::Char) = char(bswap(uint32(x)))

<<(x::Char, y::Int32)  = uint32(x) << y
>>(x::Char, y::Int32)  = uint32(x) >>> y
>>>(x::Char, y::Int32) = uint32(x) >>> y

< (x::Char, y::Char) = uint32(x) <  uint32(y)
<=(x::Char, y::Char) = uint32(x) <= uint32(y)

## traits ##

sizeof(::Type{Char}) = 4

## printing & showing characters ##

print(io::IO, c::Char) = (write(io,c); nothing)
show(io::IO, c::Char) = (print(io,'\''); print_escaped(io,UTF32String(c),"'"); print(io,'\''))


#!/usr/bin/env julia

function packagelt(a::String, b::String)
    a == "julia" && b != "julia" && return true
    a != "julia" && b == "julia" && return false
    return lowercase(a) < lowercase(b)
end

≲(v::VersionNumber, t::NTuple{0,Int}) = true
≲(v::VersionNumber, t::NTuple{1,Int}) = v.major ≤ t[1]
≲(v::VersionNumber, t::NTuple{2,Int}) = v.major < t[1] ||
                                        v.major ≤ t[1] && v.minor ≤ t[2]
≲(v::VersionNumber, t::NTuple{3,Int}) = v.major < t[1] ||
                                        v.major ≤ t[1] && v.minor < t[2] ||
                                        v.major ≤ t[1] && v.minor ≤ t[2] && v.patch ≤ t[3]

≲(t::NTuple{0,Int}, v::VersionNumber) = true
≲(t::NTuple{1,Int}, v::VersionNumber) = t[1] ≤ v.major
≲(t::NTuple{2,Int}, v::VersionNumber) = t[1] < v.major ||
                                        t[1] ≤ v.major && t[2] ≤ v.minor
≲(t::NTuple{3,Int}, v::VersionNumber) = t[1] < v.major ||
                                        t[1] ≤ v.major && t[2] < v.minor ||
                                        t[1] ≤ v.major && t[2] ≤ v.minor && t[3] ≤ v.patch

function compress_versions(inc::Vector{VersionNumber}, from::Vector{VersionNumber})
    issorted(inc) || (inc = sort(inc))
    exc = sort!(setdiff(from, inc))
    pairs = []
    if isempty(exc)
        lo, hi = first(inc), last(inc)
        push!(pairs, (lo.major, lo.minor) => (hi.major, hi.minor))
    else
        for v in inc
            t = (v.major, v.minor)
            if any(t ≲ w ≲ t for w in exc)
                t = (v.major, v.minor, v.patch)
            end
            if isempty(pairs) || any(pairs[end][1] ≲ w ≲ t for w in exc)
                push!(pairs, t => t) # need a new interval
            else
                pairs[end] = pairs[end][1] => t # can be merged with last
            end
        end
    end
    @assert all(any(p[1] ≲ v ≲ p[2] for p ∈ pairs) for v ∈ inc)
    @assert all(!any(p[1] ≲ v ≲ p[2] for p ∈ pairs) for v ∈ exc)
    return pairs
end
compress_versions(f::Function, from::Vector{VersionNumber}) =
    compress_versions(filter(f, from), from)
compress_versions(vi::VersionInterval, from::Vector{VersionNumber}) =
    compress_versions(v->v in vi, from)
compress_versions(inc, from) = compress_versions(inc, collect(from))

versions_string(p::Pair) = versions_string(p...)
versions_string(a::Tuple{}, b::Tuple{}) = "*"
versions_string(a::NTuple{m,Int}, b::NTuple{n,Int}) where {m,n} =
    a == b ? join(a, '.') : "$(join(a, '.'))-$(join(b, '.'))"

versions_repr(x) = repr(versions_string(x))
versions_repr(v::Vector) = length(v) == 1 ? repr(versions_string(v[1])) :
    "[" * join(map(repr∘versions_string, v), ", ") * "]"


function optimize(obj::BaseOptimizer, arg0::Vector{OptimizationData})
    return jcall(obj, "optimize", Object, (Vector{OptimizationData},), arg0)
end


function allunitsharvest!(game)
	#groundlevel=false
	#if unit.pl=[2]
	#	groundlevel=true
	#end
	#bools=(unit.passover,unit.passoverself,unit.inclusive)
	#influencemap=allinfluence(game,unit.groundlevel,bools)
	allunitslive!(game)
	points=[0.0,0,0,0,0]
	for (loc,unit) in game.map
		if unit!=0
			points.+=unit.harvest!(game,unit)
#			println(points)
		end
	end
	#game.points.+=points
	game.season+=1
	push!(game.sequence,:harvest)
	GAccessor.text(game.g[1,2],pointslabel(game))
	return points
end

function lifluence(game,unit::Unit)
	lifemap=unitslive(game,unit.color)
	connectedunits=Unit[unit]
	cwhite=Unit[]
	group=[unit.loc]
	if unit.color==(1,1,1)
		push!(cwhite,unit)
		stuff=Dict()
		#stuff[:loced]=[unit.loc]
		reach=makegrid(7)
		for loc in reach
			u=game.map[loc]
			if isa(u,Unit) && u.color!=(1,1,1) && !haskey(stuff,u.color)	#!in(loc,stuff[:loced])
				lmu=lifluence(game,u)
				if in(unit,lmu[3])
					stuff[u.color]=lmu
				end
			end
		end
		for (col,lm) in stuff
			for l in lm[1]
				if !in(l,group)
					push!(group,l)
				end
			end
			for cu in lm[2]
				if !in(cu,connectedunits)
					push!(connectedunits,cu)
				end
			end
			for cw in lm[3]
				if !in(cw,cwhite)
					push!(cwhite,cw)
				end
			end
		end
		return (group,connectedunits,cwhite)
	end
	temp=[unit.loc]
	while !isempty(temp)
		temp2=Tuple[]
		for t in temp
			for h in adjacent(t)
#				lmt=lifemap[t]
#				lmh=lifemap[h]
#				canspread=true
				if !in(h,group) && !in(h,temp) && !in(h,temp2) && in(h,game.board.grid) && sum(lifemap[h])>0
					push!(temp2,h)
					u=game.map[h]
					if isa(u,Unit)
						push!(connectedunits,u)
						if u.color==(1,1,1)
							push!(cwhite,u)
						end
					end
				end
			end
		end
		for t2 in temp2
			push!(group,t2)
		end
		temp=temp2
	end
	return (group,connectedunits,cwhite)
end
function connectedunits(game,unit)
	(lif,cu,cw)=lifluence(game,unit)
	cellconnected=getcellgroup(game,unit)
	return (cw,cellconnected,cu) #connected white units, connected by cells, connected by lifluence
end

function influence(game,hex,groundlevel=false,passover=false,passoverself=true,inclusive=true)
	unit=game.map[hex]
	white=(1,1,1)
	group=Dict(hex=>6.0)
	temp=Dict(hex=>6.0)
	for rad in 1:unit.ir
		temp2=Dict()
		for t in temp
			for h in adjacent(t[1],1,groundlevel)
				if !in(h,keys(group)) && !in(h,keys(temp)) && !in(h,keys(temp2)) && in(h,keys(game.map))
					inf=1/rad
					if game.map[h]==0 || passover
						temp2[h]=inf
					elseif passoverself && (game.map[h].color==unit.color || game.map[h].color==white)
						temp2[h]=inf
					end
					if inclusive && game.map[h]!=0 && !in(h,keys(temp2))
						group[h]=inf
					end
				end
			end
		end
		for (h2,i2) in temp2
			group[h2]=i2
		end
		temp=temp2
	end
	return group
end
function allinfluence(game,groundlevel=false,bools=(true,false,true,true))
	influencemap=Dict()
	for (loc,player) in game.map
		if !groundlevel || loc[3]==2
			influencemap[loc]=[0.0,0,0]
		end
	end
	for (loc,unit) in game.map
		if unit!=0
			col=unit.color
			infl=influence(game,loc,bools...)
			for inf in infl
				influencemap[inf[1]].+=inf[2].*col
			end
		end
	end
	return influencemap
end

function peekharvest(game,groundlevel=false,bools=(true,false,true,true))
	influencemap=allinfluence(game,groundlevel,bools)
	brgbw=[0.0,0,0,0,0]
	for (iloc,inf) in influencemap
		ninf=numcolors(inf)
		if ninf==3
			brgbw[5]+=min(inf...)
		end
		if ninf==1
			brgbw[1]+=min(sum(inf),1)
		else
			for c in 1:3
				brgbw[c+1]+=min(inf[c],inf[c%3+1])
			end
		end
	end
	for c in 2:4
		brgbw[c]-=brgbw[5]
	end
	return brgbw
end

function getgroup_dep(game,hex)
	player=game.map[hex]
	white=(1,1,1)
	if player==0
		return []
	end
	group=Tuple[hex]
	temp=[hex]
	while !isempty(temp)
		temp2=Tuple[]
		for t in temp
			for h in adjacent(t)
				if !in(h,group) && !in(h,temp) && !in(h,temp2) && in(h,keys(game.map)) && (game.map[h]==player || game.map[h]==white)
					push!(temp2,h)
				end
			end
		end
		for t2 in temp2
			push!(group,t2)
		end
		temp=temp2
	end
	return group
end

function unitcost(game,loc,unitparams)
	distance=nearestwhite(game,loc,loc[3]==2)
	cost=distance*unitparams[2]
	cost*=sum(unitparams[1]) #only if sum>1?
	return cost
end
function unitcost(game,unit::Unit)
	return unit.costfun(game)
end
function subtractcost(game,cost,color)
	rgb=cost.*color
	srgb=sum(rgb)
	if rgb[1]<=game.points[2] && rgb[2]<=game.points[3] && rgb[3]<=game.points[4]
		game.points[2]-=rgb[1]
		game.points[3]-=rgb[2]
		game.points[4]-=rgb[3]
	elseif srgb/2<=game.points[5] && srgb/2<=game.points[1]
		game.points[1]-=srgb/2
		game.points[5]-=srgb/2
	else
		return false
	end
	return true
end
function updategroups_dep!(game::Game) #sometimes doesn't find most recent unit..? Or is there an issue when deleting units? Maybe fixed. Now it sometimes adds the most recent unit twice...
	groups=Group[]
	for spawn in game.spawns
		push!(groups,getgroup(game,spawn))
	end
	unique=Group[groups[1]]
	for group in groups
		for uniq in unique
			if samegroup(group,uniq)
				break
			end
			push!(unique,group)
		end
	end
	game.groups=unique
end

function getgroup_dep(game,unit::Unit,color=-1,connectedunits=Unit[]) #why don't white units get added to spawns? Maybe they shouldn't, bug in our favor. They should be available as partial spawns, divide white into 3 spawn units. Fixed some stuff, still good bugs? Well, this should be rewritten rather than bugged down, first a initgroup that gets the body.
	if color==-1
		color=unit.color
	end
	lifemap=unitslive(game,unit.color)
	if !in(unit,connectedunits)
		push!(connectedunits,unit)
	end
	cwhite=Unit[] #connected spawns
	if color==(1,1,1) #this needs to be rewritten to allow colored spawns, sorta works now but... nonbody groups automerge with white group
		push!(cwhite,unit)
		stuff=Dict()
		reach=makegrid(7) #should be minimum twice the maximum ir of all units +1, should be a better way
		for lo in reach
			loc=lo.+(unit.loc.-(0,0,2))
			if !in(loc,keys(game.map))
				continue
			end
			u=game.map[loc]
			if isa(u,Unit) && distance(u.loc,unit.loc)<(u.ir+unit.ir+2) && !haskey(stuff,u.color) && !in(u,connectedunits) && u.color!=(1,1,1)
				subgroup=getgroup(game,u,u.color,connectedunits)
				if in(unit,subgroup.spawns)
					stuff[u.color]=subgroup
				end
			end
		end
		for (col,sg) in stuff
			for cu in sg.units
				if !in(cu,connectedunits)
					push!(connectedunits,cu)
				end
			end
			for cw in sg.spawns
				if !in(cw,cwhite)
					push!(cwhite,cw)
				end
			end
		end
		body=getcellgroup(game,cwhite[1])
		return newgroup(cwhite,body,connectedunits)
	end
	ulocs=[unit.loc]
	temp=[unit.loc]
	while !isempty(temp)
		temp2=Tuple[]
		for t in temp
			for h in adjacent(t)
				if !in(h,ulocs) && !in(h,temp) && !in(h,temp2) && in(h,game.board.grid) && sum(lifemap[h])>0
					push!(temp2,h)
					u=game.map[h]
					if isa(u,Unit) && !in(u,connectedunits)
						push!(connectedunits,u)
						if u.color==(1,1,1)
							push!(cwhite,u)
							for adjwu in adjacent(u.loc)
								adju=game.map[adjwu]
								if isa(adju,Unit) && adju.color==color && !in(adju,connectedunits)
									nsg=getgroup(game,adju,color,connectedunits)
									for cu in nsg.units
										if !in(cu,connectedunits)
											push!(connectedunits,cu)
										end
									end
									for cw in nsg.spawns
										if !in(cw,cwhite)
											push!(cwhite,cw)
										end
									end
								end
							end
						end
					end
				end
			end
		end
		for t2 in temp2
			push!(ulocs,t2)
		end
		temp=temp2
	end
	body=getcellgroup(game,connectedunits[1]) #tempfix
	return newgroup(cwhite,body,connectedunits)
end

function allunitslive!(game)
	lifemap=Dict()
	for loc in game.board.grid
		lifemap[loc]=[0.0,0,0]
	end
	for (loc,unit) in game.map
		if isa(unit,Unit)
			unit.live!(game,unit,lifemap)
		end
	end
	game.lifemap=lifemap
	return lifemap
end
#lifemap=allunitslive! #maybe not store lifemap in game since it gets modified in so many places. Rewrite everything to work with a temp lifemap. Or rather have a solid lifemap and a temp harvest ledger

function printpoints(game)
	harv=round.(peekharvest(game),1,10)
	println("Black: ",harv[1]," Red: ",harv[2]," Green: ",harv[3]," Blue: ",harv[4]," White: ",harv[5])#," Total: ",sum(harv))
end
function nearestwhite(game,hex,layer=false) #rewrite to check for spawn? Or spawning is always done from groups? This isn't even needed now that costs are gone
	white=(1,1,1)
	group=[hex]
	temp=[hex]
	for rad in 1:1000000
		temp2=[]
		for t in temp
			for h in adjacent(t,1,layer)
				if in(h,keys(game.map))
					unit=game.map[h]
					if unit!=0 && unit.color==white
						return rad
					end
					if !in(h,group) && !in(h,temp) && !in(h,temp2)
						push!(temp2,h)
					end
				end
			end
		end
		for h2 in temp2
			push!(group,h2)
		end
		temp=temp2
	end
	return Inf
end
function getexpandcost(shells::Integer=6,initlocs=[(6,6,2)],basecost=50) #deprecated
	patch=makegrid(shells,initlocs)
	cost=length(patch)+basecost*length(initlocs)
	return cost
end
function harvest!(game,group::Group)
	if !group.harvested
		harv=checkharvest(game,group) #some stuff needs to be rewritten to avoid the deepcopy above
		group.points+=harv
		group.harvested=true
		for unit in group.units
			unit.harvested=true
		end
	end
end
function allgroupsharvest!(game)
	for group in game.groups
		harvest!(game,group)
	end
end

function collectharvest!(game::Game)
	for group in game.groups
		game.points+=group.points
		group.points-=group.points
	end
end
function newseason!(game)
	for group in game.groups
		group.harvested=false
		for unit in group.units
			unit.harvested=false
		end
	end
	game.season+=1
end
function harvest!(game::Game)
	allunitslive!(game) #maybe call only when new units are placed. But then placing units takes longer...
	updategroups!(game)
	allgroupsharvest!(game)
	collectharvest!(game)
	newseason!(game)
	push!(game.sequence,:harvest)
	GAccessor.text(game.g[1,2],pointslabel(game))
	return game.points
end


# SPLITTING DATA SETS

# Helper function for partitioning in the non-stratified case
function _partition(rows, fractions, ::Nothing)
    # container for the row selections (head:tail)
    n_splits = length(fractions) + 1
    heads    = zeros(Int, n_splits)
    tails    = zeros(Int, n_splits)
    n_rows   = length(rows)
    head     = 1
    for (i, p) in enumerate(fractions)
        n = round(Int, p * n_rows)
        iszero(n) && (@warn "A split has only one element."; n = 1)
        # update tail
        tail = head + n - 1
        # store
        heads[i] = head
        tails[i] = tail
        # update head
        head = tail + 1
    end
    if head > n_rows
        @warn "Last vector in the split has only one element."
        head = n_rows
    end
    heads[end] = head
    tails[end] = n_rows
    return tuple((rows[h:t] for (h, t) in zip(heads, tails))...)
end

_make_numerical(v::AbstractVector) =
    throw(ArgumentError("`stratify` must have `Count`, `Continuous` "*
                        "or `Finite` element scitpye. Consider "*
                        "`coerce(stratify, Finite)`. "))
_make_numerical(v::AbstractVector{<:Union{Missing,Real}}) = v
_make_numerical(v::AbstractVector{<:Union{Missing,CategoricalValue}}) =
                                int.(v)

# Helper function for partitioning in the stratified case
function _partition(rows, fractions, raw_stratify::AbstractVector)
    stratify = _make_numerical(raw_stratify)
    length(stratify) == length(rows) ||
        throw(ArgumentError("The stratification vector must "*
                            "have as many entries as " *
                            "the rows to partition."))
    uv    = unique(stratify)
    # construct table (n_classes * idx_of_that_class)
    # NOTE use of '===' is important to handle missing.
    idxs  = [[i for i in rows if stratify[rows[i]] === v] for v in uv]

    # number of occurences of each class and proportions
    nidxs = length.(idxs)
    props = length.(idxs) ./ length(rows)

    n_splits = length(fractions) + 1
    n_rows   = length(rows)

    ns_props = round.(Int, n_rows * fractions * props')
    ns_props = vcat(ns_props, nidxs' .- sum(ns_props, dims=1))

    # warn if anything is >= 1
    if !all(e -> e > 1, ns_props)
        @warn "Some splits have a single or no representative of some class."
    end
    # container for the rows
    split_rows = []
    heads      = ones(Int, length(uv))
    for r in 1:size(ns_props, 1)
        tails   = heads .+ ns_props[r, :] .- 1
        # take chunks of the indices corresponding to the current fraction
        indices = vcat((idxs[i][heads[i]:tails[i]] for i in eachindex(uv))...)
        # rearrange by order of appearance
        indices = sort(indices)
        push!(split_rows, rows[indices])
        heads .= tails .+ 1
    end
    if !all(sl -> sl > 1, length.(split_rows))
        @warn "Some splits have a single or no representative of some class."
    end
    return tuple(split_rows...)
end

const ERR_PARTITION_UNSUPPORTED = ArgumentError(
    "Function `partition` only supports "*
    "AbstractVector, AbstractMatrix or containers implementing the "*
    "Tables interface.")
const ERR_PARTITION_DIMENSION_MISMATCH = DimensionMismatch(
    "Expected a tuple of objects with a common length. ")

__nrows(X) = Tables.istable(X) ? nrows(X) : throw(ERR_PARTITION_UNSUPPORTED)
__nrows(X::Union{AbstractMatrix,AbstractVector}) = nrows(X)

"""
    partition(X, fractions...;
              shuffle=nothing,
              rng=Random.GLOBAL_RNG,
              stratify=nothing,
              multi=false)

Splits the vector, matrix or table `X` into a tuple of objects of the
same type, whose vertical concatenation is `X`. The number of rows in
each component of the return value is determined by the
corresponding `fractions` of `length(nrows(X))`, where valid fractions
are floats between 0 and 1 whose sum is less than one. The last
fraction is not provided, as it is inferred from the preceding ones.

For "synchronized" partitioning of multiple objects, use the
`multi=true` option described below.

    julia> partition(1:1000, 0.8)
    ([1,...,800], [801,...,1000])

    julia> partition(1:1000, 0.2, 0.7)
    ([1,...,200], [201,...,900], [901,...,1000])

    julia> partition(reshape(1:10, 5, 2), 0.2, 0.4)
    ([1 6], [2 7; 3 8], [4 9; 5 10])

    X, y = make_blobs() # a table and vector
    Xtrain, Xtest = partition(X, 0.8, stratify=y)

    (Xtrain, Xtest), (ytrain, ytest) = partition((X, y), 0.8, rng=123, multi=true)

## Keywords

* `shuffle=nothing`: if set to `true`, shuffles the rows before taking
  fractions.

* `rng=Random.GLOBAL_RNG`: specifies the random number generator to be
  used, can be an integer seed. If specified, and `shuffle ===
  nothing` is interpreted as true.

* `stratify=nothing`: if a vector is specified, the partition will
  match the stratification of the given vector. In that case,
  `shuffle` cannot be `false`.

* `multi=false`: if `true` then `X` is expected to be a `tuple` of
  objects sharing a common length, which are each partitioned
  separately using the same specified `fractions` *and* the same row
  shuffling. Returns a tuple of partitions (a tuple of tuples).

"""
function partition(X, fractions::Real...;
                   shuffle::Union{Nothing,Bool}=nothing,
                   rng=Random.GLOBAL_RNG,
                   stratify::Union{Nothing,AbstractVector}=nothing,
                   multi=false)

    # check the fractions
    if !all(e -> 0 < e < 1, fractions) || sum(fractions) >= 1
        throw(DomainError(fractions,
                          "Fractions must be in (0, 1) with sum < 1."))
    end

    # determinen `n_rows`:
    if X isa Tuple && multi
        isempty(X) && return tuple(fill((), length(fractions) + 1)...)
        x = first(X)
        n_rows = __nrows(x)
        all(X[2:end]) do x
            nrows(x) === n_rows
        end || throw(ERR_PARTITION_DIMENSION_MISMATCH)
    else
        n_rows = __nrows(X)
    end

    # check the rng & adjust shuffling
    if rng isa Integer
        rng = MersenneTwister(rng)
    end
    if rng != Random.GLOBAL_RNG && shuffle === nothing
        shuffle = true
    end
    rows = collect(1:n_rows)
    shuffle !== nothing && shuffle && shuffle!(rng, rows)

    # determine the partition of `rows`:
    row_partition = _partition(rows, collect(fractions), stratify)
    return _partition(X, row_partition)
end

function _partition(X, row_partition)
#    _X = Tables.istable(X) ? X : collect(X)
    return tuple((selectrows(X, p) for p in row_partition)...)
end

_partition(X::Tuple, row_partition) =
    map(x->_partition(x, row_partition), X)

# # UNPACK

"""
    unpack(table, f1, f2, ... fk;
           wrap_singles=false,
           shuffle=false,
           rng::Union{AbstractRNG,Int,Nothing}=nothing,
           coerce_options...)

Horizontally split any Tables.jl compatible `table` into smaller
tables or vectors by making column selections determined by the
predicates `f1`, `f2`, ..., `fk`. Selection from the column names is
without replacement. A *predicate* is any object `f` such that
`f(name)` is `true` or `false` for each column `name::Symbol` of
`table`.

Returns a tuple of tables/vectors with length one greater than the
number of supplied predicates, with the last component including all
previously unselected columns.

```
julia> table = DataFrame(x=[1,2], y=['a', 'b'], z=[10.0, 20.0], w=["A", "B"])
2×4 DataFrame
 Row │ x      y     z        w
     │ Int64  Char  Float64  String
─────┼──────────────────────────────
   1 │     1  a        10.0  A
   2 │     2  b        20.0  B

Z, XY, W = unpack(table, ==(:z), !=(:w))
julia> Z
2-element Vector{Float64}:
 10.0
 20.0

julia> XY
2×2 DataFrame
 Row │ x      y
     │ Int64  Char
─────┼─────────────
   1 │     1  a
   2 │     2  b

julia> W  # the column(s) left over
2-element Vector{String}:
 "A"
 "B"
```

Whenever a returned table contains a single column, it is converted to
a vector unless `wrap_singles=true`.

If `coerce_options` are specified then `table` is first replaced
with `coerce(table, coerce_options)`. See
[`ScientificTypes.coerce`](@ref) for details.

If `shuffle=true` then the rows of `table` are first shuffled, using
the global RNG, unless `rng` is specified; if `rng` is an integer, it
specifies the seed of an automatically generated Mersenne twister. If
`rng` is specified then `shuffle=true` is implicit.

"""
function unpack(X, predicates...;
                wrap_singles=false,
                shuffle=nothing,
                rng=nothing, pairs...)

    # add a final predicate to unpack all remaining columns into to
    # the last return value:
    predicates = (predicates..., _ -> true)

    shuffle, rng = shuffle_and_rng(shuffle, rng)

    shuffle && (X = selectrows(X, Random.shuffle(rng, 1:nrows(X))))

    if isempty(pairs)
        Xfixed = X
    else
        Xfixed = coerce(X, pairs...)
    end

    unpacked = Any[]
    names_left = schema(Xfixed).names |> collect

    for c in predicates
        names = filter(c, names_left)
        filter!(!in(names), names_left)
        length(names) == 1 && !wrap_singles && (names = names[1])
        push!(unpacked, selectcols(Xfixed, names))
    end
    return Tuple(unpacked)
end

## RESTRICTING TO A FOLD

struct FoldRestrictor{i,N}
    f::NTuple{N,Vector{Int}}
end
(r::FoldRestrictor{i})(X) where i = selectrows(X, (r.f)[i])

"""
    restrict(X, folds, i)

The restriction of `X`, a vector, matrix or table, to the `i`th fold
of `folds`, where `folds` is a tuple of vectors of row indices.

The method is curried, so that `restrict(folds, i)` is the operator
on data defined by `restrict(folds, i)(X) = restrict(X, folds, i)`.

### Example

    folds = ([1, 2], [3, 4, 5],  [6,])
    restrict([:x1, :x2, :x3, :x4, :x5, :x6], folds, 2) # [:x3, :x4, :x5]

See also [`corestrict`](@ref)

"""
restrict(f::NTuple{N}, i) where N = FoldRestrictor{i,N}(f)
restrict(X, f, i) = restrict(f, i)(X)

## RESTRICTING TO A FOLD COMPLEMENT

"""
    complement(folds, i)

The complement of the `i`th fold of `folds` in the concatenation of
all elements of `folds`. Here `folds` is a vector or tuple of integer
vectors, typically representing row indices or a vector, matrix or
table.

    complement(([1,2], [3,], [4, 5]), 2) # [1 ,2, 4, 5]

"""
complement(f, i) = reduce(vcat, collect(f)[Not(i)])

struct FoldComplementRestrictor{i,N}
    f::NTuple{N,Vector{Int}}
end
(r::FoldComplementRestrictor{i})(X) where i =
    selectrows(X, complement(r.f, i))

"""
    corestrict(X, folds, i)

The restriction of `X`, a vector, matrix or table, to the *complement*
of the `i`th fold of `folds`, where `folds` is a tuple of vectors of
row indices.

The method is curried, so that `corestrict(folds, i)` is the operator
on data defined by `corestrict(folds, i)(X) = corestrict(X, folds, i)`.

### Example

    folds = ([1, 2], [3, 4, 5],  [6,])
    corestrict([:x1, :x2, :x3, :x4, :x5, :x6], folds, 2) # [:x1, :x2, :x6]

"""
corestrict(f::NTuple{N}, i) where N = FoldComplementRestrictor{i,N}(f)
corestrict(X, f, i) = corestrict(f, i)(X)

## to be replaced (not used anywhere):
## ACCESSORS FOR JULIA NDSPARSE ARRAYS (N=2)

# nrows(::Val{:sparse}, X) = maximum([r[1] for r in keys(X)])
# function select(::Val{:sparse}, X, r::Integer, c::Symbol)
#     try
#         X[r,c][1]
#     catch exception
#         exception isa KeyError || throw(exception)
#         missing
#     end
# end
# select(::Val{:sparse}, X, r::AbstractVector{<:Integer}, c::Symbol) = [select(X, s, c) for s in r]
# select(::Val{:sparse}, X, ::Colon, c::Symbol) = [select(X, s, c) for s in 1:nrows(X)]
# selectrows(::Val{:sparse}, X, r::Integer) = X[r:r,:]
# selectrows(::Val{:sparse}, X, r) = X[r,:]
# selectcols(::Val{:sparse}, X, c::Symbol) = select(X, :, c)
# selectcols(::Val{:sparse}, X, c::AbstractVector{Symbol}) = X[:,sort(c)]
# selectcols(::Val{:sparse}, X, ::Colon) = X
# select(::Val{:sparse}, X, r::Integer, c::AbstractVector{Symbol}) = X[r,sort(c)]
# select(::Val{:sparse}, X, r::Integer, ::Colon) = X[r,:]
# select(::Val{:sparse}, X, r, c) = X[r,sort(c)]

## TRANSFORMING BETWEEN CATEGORICAL ELEMENTS AND RAW VALUES

MLJModelInterface.transform(
    e::Union{CategoricalArray,CategoricalValue,CategoricalPool},
    arg) = CategoricalDistributions.transform(e, arg)

## SKIPPING MISSING AND NAN: skipinvalid

_isnan(x) = false
_isnan(x::Number) = isnan(x)

skipnan(x) = Iterators.filter(!_isnan, x)

"""
    skipinvalid(itr)

Return an iterator over the elements in `itr` skipping `missing` and
`NaN` values. Behaviour is similar to [`skipmissing`](@ref).

    skipinvalid(A, B)

For vectors `A` and `B` of the same length, return a tuple of vectors
`(A[mask], B[mask])` where `mask[i]` is `true` if and only if `A[i]`
and `B[i]` are both valid (non-`missing` and non-`NaN`). Can also
called on other iterators of matching length, such as arrays, but
always returns a vector. Does not remove `Missing` from the element
types if present in the original iterators.

"""
skipinvalid(v) = v |> skipmissing |> skipnan

isinvalid(x) = ismissing(x) || _isnan(x)

function skipinvalid(yhat, y)
    mask = .!(isinvalid.(yhat) .| isinvalid.(y))
    return yhat[mask], y[mask]
end

# TODO: refactor balanced accuracy to get rid of these:

function _skipinvalid(yhat, y, w::Arr)
    mask = .!(isinvalid.(yhat) .| isinvalid.(y))
    return yhat[mask], y[mask], w[mask]
end

function _skipinvalid(yhat, y, w::Union{Nothing,AbstractDict})
    mask = .!(isinvalid.(yhat) .| isinvalid.(y))
    return yhat[mask], y[mask], w
end


@rule typeof(-)(:in2, Marginalisation) (m_out::Any, m_in1::Any, meta::Any) = begin
    @call_rule typeof(+)(:in2, Marginalisation) (m_out = m_in1, m_in1 = m_out, meta = meta)
end


#===============================================================================
                    New definitions of introduced updates
===============================================================================#

#===============================================================================
                            Imputation of paths
===============================================================================#
"""
    struct PathImputation{T} <: MCMCDiffusionImputation
        ρs::Vector{Vector{T}}
        aux_laws
    end

Update type. An indicator for imputation of unobserved path segments. `ρs` are
the memory parameters for the preconditioned Crank–Nicolson scheme and
`aux_laws` are the laws for the auxiliary diffusions.

    PathImputation(ρ::T, P) where T<:Number

Base constructor v1. Initialize all `ρs` with the same value `ρ` and set
auxiliary laws to `P`.

    PathImputation(ρ::T, P) where T<:Vector
Base constructor v2. Initialize each recording with its own value of `ρ` and set
auxiliary laws to `P`.
"""
struct PathImputation{T,A} <: MCMCDiffusionImputation
    ρs::Vector{Vector{T}}
    aux_laws
    adpt::A

    function PathImputation(ρ::T, P; adpt=NoAdaptation()) where T<:Number
        _adpt, A = init_path_imp_adpt(adpt)
        new{T,A}([[ρ]], P, _adpt)
    end

    function PathImputation(ρ::T, P; adpt=NoAdaptation()) where T<:Vector
        _adpt, A = init_path_imp_adpt(adpt)
        new{eltype(ρ),A}([ρ], P, _adpt)
    end

    function PathImputation(
            ρs::T, P; adpt=NoAdaptation()
        ) where T<:Vector{<:Vector}
        _adpt, A = init_path_imp_adpt(adpt)
        new{eltype(ρs[1]),A}(ρs, P, _adpt)
    end
end

init_update!(updt::eMCMC.MCMCUpdate, block_layout) = nothing

function init_update!(updt::PathImputation, block_layout)
    num_recordings = length(block_layout)
    Δ = num_recordings - length(updt.ρs)
    if Δ > 0
        for i in 1:Δ
            push!(updt.ρs, [updt.ρs[1][1]])
        end
    end
    for i in 1:length(updt.ρs)
        update_ρ!(updt.ρs[i], block_layout[i])
    end
    resize_adpt!(updt.adpt, num_recordings)
end

function update_ρ!(ρs::Vector, block_layout)
    N = length(block_layout[2])
    if length(ρs) < N
        for i in 1:(N-length(ρs))
            push!(ρs, ρs[1])
        end
    end
end

struct StartingPointsImputation <: MCMCDiffusionImputation end

struct StartingPointsLangevinImputation <: MCMCDiffusionImputation end

#===============================================================================
                        Conjugate Gaussian updates
===============================================================================#

struct DiffusionConjugGsnUpdate{K,TP} <: eMCMC.MCMCConjugateParamUpdate
    coords::K
    prior::TP
    adpt::NoAdaptation

    function DiffusionConjugGsnUpdate(coords, prior)
        if typeof(prior) <: Normal
            prior = Gaussian( [prior.μ], reshape(prior.σ, (1,1)) )
        elseif typeof(prior) <: MvNormal
            prior = Gaussian(prior.μ, prior.Σ.mat)
        else
            @assert typeof(prior) <: Gaussian
        end

        if typeof(coords) <: Number
            coords = [coords]
        end
        new{typeof(coords), typeof(prior)}(coords, prior, NoAdaptation())
    end
end

struct ConjugGsnHelper{T,N,K}
    relevant_recordings::Vector{Int64}
    sub_μ_idx::Vector{K}
    θ_names::Vector{Val}
    θᶜ_names::Vector{Val}

    function ConjugGsnHelper(
            updt::DiffusionConjugGsnUpdate,
            global_ws,
            PPP,
            lpn::LocalUpdtParamNames;
            computation_type=:outofplace
        )
        _glob_pnames = global_ws.pnames[updt.coords]
        relevant_recordings = Int64[]
        sub_μ_idx = Vector{Vector{Int64}}(undef, 0)
        θᶜ_names = Val[]
        θ_names = Val[]
        all_θ_names = Symbol[]

        for (i, PP) in enumerate(PPP)
            P = first(PP).P_target
            if length(lpn.θ_local_names) > 0
                push!(relevant_recordings, i)

                # sort by the local index of θ
                θ_loc_names = sort(collect(lpn.θ_local_names[i]), by=x->x[1])
                sub_idx = getindex.(θ_loc_names, 1)
                push!(sub_μ_idx, sub_idx)

                sub_name = getindex.(θ_loc_names, 2)
                push!(
                    θᶜ_names,
                    Val(
                        Tuple(
                            filter(
                                pn->(
                                    !(pn in sub_name) &&
                                    !DD.ignore_for_cu(Val(pn), P)
                                ),
                                DD.parameter_names(P)
                            )
                        )
                    )
                )
                push!( θ_names, Val(Tuple(sub_name)) )
                append!(all_θ_names, sub_name)
            else
                push!(sub_μ_idx, Int64[])
                push!(θᶜ_names, Val[])
                push!(θ_names, Val[])
            end
        end
        u_θ_names = unique(all_θ_names)
        sub_μ_idx = transform_μ_idx(Val(computation_type), sub_μ_idx)
        new{computation_type, length(u_θ_names), eltype(sub_μ_idx)}(
            relevant_recordings,
            sub_μ_idx,
            θ_names,
            θᶜ_names,
        )
    end

    ConjugGsnHelper() = new{:none, 0, Nothing}()
end

transform_μ_idx(::Val{:inplace}, sub_μ_idx) = sub_μ_idx

function transform_μ_idx(v::Val{:outofplace}, sub_μ_idx)
    map(i->transform_μ_idx(v, i, Val(length(i))), sub_μ_idx)
end

function transform_μ_idx(::Val{:outofplace}, sub_μ_idx, ::Val{N}) where N
    SVector{N,Int64}(sub_μ_idx)
end

function μ_init(cu_helper::ConjugGsnHelper{:outofplace,N}) where N
    zero(MVector{N,Float64})
end

function W_init(cu_helper::ConjugGsnHelper{:outofplace,N}) where N
    zero(MMatrix{N,N,Float64})
end

function μ_init(cu_helper::ConjugGsnHelper{:outofplace}, ::SVector{N}) where N
    zero(SVector{N,Float64})
end

μ_init(cu_helper::ConjugGsnHelper{:inplace,N}) where N = zeros(Float64, N)
W_init(cu_helper::ConjugGsnHelper{:inplace,N}) where N = zeros(Float64, (N,N))

function μ_init(cu_helper::ConjugGsnHelper{:inplace,N}, v) where N
    zeros(Float64, length(v))
end

@generated function conjug_θᶜ(
        P_target, pnames::Val{T}, ::ConjugGsnHelper{:outofplace}
    ) where T
    N = length(T)
    data = Expr(
        :tuple,
        1.0,
        (:(getfield(P_target, T[$i])) for i in 1:N)...,
    )
    vec = Expr(:call, MVector{N+1}, data)
end

@generated function conjug_θᶜ(
        P_target, pnames::Val{T}, ::ConjugGsnHelper{:inplace}
    ) where T
    N = length(T)
    data = Expr(
        :tuple,
        1.0,
        (:(getfield(P_target, T[$i])) for i in 1:N)...,
    )
    vec = Expr(:call, :collect, data)
end

function eMCMC.proposal!(
        updt::DiffusionConjugGsnUpdate,
        global_ws,
        local_ws::eMCMC.LocalWorkspace,
        step,
    )
    state°(local_ws) .= conjugate_draw(
        local_ws.cu_helper,
        local_ws.sub_ws_diff.P,
        updt.prior,
        local_ws.sub_ws_diff.XX,
    )
end

function conjugate_draw(cu_helper, PP, prior, XX)
    μ = μ_init(cu_helper)
    W = W_init(cu_helper)
    for i in cu_helper.relevant_recordings
        μi = μ_init(cu_helper, cu_helper.sub_μ_idx[i])
        Wi = μi*μi'
        P = first(PP[i]).P_target
        θᶜ = conjug_θᶜ(P, cu_helper.θᶜ_names[i], cu_helper)

        μi, Wi = compute_μ_and_W(
            μi, Wi, cu_helper.θ_names[i], cu_helper.θᶜ_names[i], θᶜ, P, XX[i],
            cu_helper
        )

        sidx = cu_helper.sub_μ_idx[i]
        view(μ, sidx) .+= μi
        view(W, sidx, sidx) .+= Wi
    end
    μ_post, Σ_post = compute_posterior_μ_and_Σ(μ, W, prior)
    rand(Gaussian(μ_post, Σ_post))
end

function compute_μ_and_W(
        μ, W, θ_names, θᶜ_names, θᶜ, P_target::S, XX, ::ConjugGsnHelper{:outofplace};
        nnh = DD.num_non_hypo(S)
    ) where S
    for X in XX
        num_pts = length(X)
        for i in 1:num_pts-1
            φₜ = φ(θ_names, X.t[i], X.x[i], P_target, nnh)
            φᶜₜ = φᶜ(θᶜ_names, θᶜ, X.t[i], X.x[i], P_target, nnh)
            Γ⁻¹ = DD.hypo_a_inv(X.t[i], X.x[i], P_target)
            dt = X.t[i+1] - X.t[i]
            dy = DD.nonhypo(X.x[i+1], P_target) - DD.nonhypo(X.x[i], P_target)
            μ = μ + φₜ'*Γ⁻¹*dy - φₜ'*Γ⁻¹*φᶜₜ*dt
            W = W + φₜ'*Γ⁻¹*φₜ*dt
        end
    end
    μ, W
end

function compute_μ_and_W(
        μ, W, θ_names, P_target, XX, θᶜ, ::ConjugGsnHelper{:inplace}
    )
    error("not implemented, need buffers...")
end

function compute_posterior_μ_and_Σ(μ, W, prior)
    Σ = Symmetric(inv(W + inv(prior.Σ)))
    Σ * (μ .+ prior.Σ\prior.μ), Σ
end

const _GENERATED = Val(:gen)

function φ(v::Val, t, x, P::S, num_non_smooth_coords=DD.num_non_hypo(S)) where S
    φ(_GENERATED, v, t, x, P, num_non_smooth_coords)
end

@generated function φ(::Val{:gen}, ::Val{Ts}, t, x, P::S, ::Val{N}) where {Ts,S,N}
	args = [Expr(:call, :Val, :(Ts[$i])) for i in 1:length(Ts)]
    data = Expr(
        :call,
        :tuplejoin,#TODO try to do tuple join not in the expression
        ( :(DD.phi($a, t, x, P)) for a in args )...
    )
    mat = Expr(:call, SMatrix{N,length(Ts)}, data)
    return mat
end

"""
    φᶜ(v, θᶜ, t, x, P::S, n=DD.num_non_hypo(S)) where S

Remainder term in the drift after removing φ'θ from it.
"""
function φᶜ(v, θᶜ, t, x, P::S, n=DD.num_non_hypo(S)) where S
    φᶜlinear(v, t, x, P, n) * θᶜ
end

function φᶜlinear(v, t, x, P::S, nnh=DD.num_non_hypo(S)) where S
    φᶜlinear(_GENERATED, v, t, x, P, nnh)
end

@generated function φᶜlinear(::Val{:gen}, ::Val{Ts}, t, x, P::S, ::Val{N}) where {Ts,S,N}
    args = [Expr(:call, :Val, :(Ts[$i])) for i in 1:length(Ts)]
    data = Expr(
        :call,
        :tuplejoin,
        :( DD.phi(Val(:intercept), t, x, P) ),
        ( :(DD.phi($a, t, x, P)) for a in args )...
    )
    mat = Expr(:call, SMatrix{N,length(Ts)+1}, data)
    return mat
end

@inline tuplejoin(x) = x
@inline tuplejoin(x, y) = (x..., y...)
@inline tuplejoin(x, y, z...) = tuplejoin(tuplejoin(x, y), z...)
#TODO at the accept reject step it needs to accept a boolean not only a vec of booleans


function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{R, D},
    ::Type{S},
) where {
    R <: PSY.RenewableGen,
    D <: AbstractRenewableDispatchFormulation,
    S <: PM.AbstractPowerModel,
}
    devices = get_available_components(R, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    # Expression
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    if has_service_model(model)
        add_to_expression!(
            container,
            ActivePowerRangeExpressionLB,
            ActivePowerVariable,
            devices,
            model,
            S,
        )
        add_to_expression!(
            container,
            ActivePowerRangeExpressionUB,
            ActivePowerVariable,
            devices,
            model,
            S,
        )
    end
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{R, D},
    ::Type{S},
) where {
    R <: PSY.RenewableGen,
    D <: AbstractRenewableDispatchFormulation,
    S <: PM.AbstractPowerModel,
}
    devices = get_available_components(R, sys)

    if has_service_model(model)
        add_constraints!(
            container,
            ActivePowerVariableLimitsConstraint,
            ActivePowerRangeExpressionUB,
            devices,
            model,
            S,
        )
    else
        add_constraints!(
            container,
            ActivePowerVariableLimitsConstraint,
            ActivePowerVariable,
            devices,
            model,
            S,
        )
    end

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{R, D},
    ::Type{S},
) where {
    R <: PSY.RenewableGen,
    D <: AbstractRenewableDispatchFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(R, sys)

    add_variables!(container, ActivePowerVariable, devices, D())

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    # Expression
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    if has_service_model(model)
        add_to_expression!(
            container,
            ActivePowerRangeExpressionLB,
            ActivePowerVariable,
            devices,
            model,
            S,
        )
        add_to_expression!(
            container,
            ActivePowerRangeExpressionUB,
            ActivePowerVariable,
            devices,
            model,
            S,
        )
    end
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{R, D},
    ::Type{S},
) where {
    R <: PSY.RenewableGen,
    D <: AbstractRenewableDispatchFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(R, sys)

    if has_service_model(model)
        add_constraints!(
            container,
            ActivePowerVariableLimitsConstraint,
            ActivePowerRangeExpressionUB,
            devices,
            model,
            S,
        )
    else
        add_constraints!(
            container,
            ActivePowerVariableLimitsConstraint,
            ActivePowerVariable,
            devices,
            model,
            S,
        )
    end
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{R, FixedOutput},
    ::Type{S},
) where {R <: PSY.RenewableGen, S <: PM.AbstractPowerModel}
    devices = get_available_components(R, sys)

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)
    add_parameters!(container, ReactivePowerTimeSeriesParameter, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{R, FixedOutput},
    ::Type{S},
) where {R <: PSY.RenewableGen, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(R, sys)

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{R, FixedOutput},
    ::Type{S},
) where {R <: PSY.RenewableGen, S <: PM.AbstractPowerModel}
    # FixedOutput doesn't add any constraints to the model. This function covers
    # AbstractPowerModel and AbtractActivePowerModel
    return
end


using ExpmV, Expokit
export AQCprocessor

abstract AQCprocessor{N}

include("hamiltonian.jl")
include("aqc.jl")
include("adiaprocessing.jl")
include("demoncooling.jl")


select(pars::AbstractArray{T,N}, elements_per_dim...; newfun = randoutzeroin) where {T, N} = NaiveNASlib.parselect(pars, elements_per_dim...; newfun)
select(::Missing, args...;kwargs...) = missing
select(::Flux.Zeros, args...;kwargs...) = Flux.Zeros()

struct WeightParam end
struct BiasParam end
struct RecurrentWeightParam end
struct RecurrentState end

"""
    neuroninsert(lt, partype)

Return a function which creates new parameters for layers of type `lt` to use for [`select`](@Ref).
"""
neuroninsert(t, partype) = randoutzeroin
neuroninsert(t, parname::Symbol) = neuroninsert(t, Val(parname))
neuroninsert(t, name::Val) = randoutzeroin

neuroninsert(lt::FluxParNorm, t::Val) = norminsert(lt, t)
norminsert(::FluxParNorm, ::Union{Val{:β},Val{:μ}}) = (args...) -> 0
norminsert(::FluxParNorm, ::Union{Val{:γ},Val{:σ²}}) = (args...) -> 1

# Coupling between input and output weights make it difficult to do anything else?
neuroninsert(::FluxDepthwiseConv, partype) = (args...) -> 0

randoutzeroin(T, d, s...) = _randoutzeroin(T,d,s)
_randoutzeroin(T, d, s) = 0
_randoutzeroin(T, d, s::NTuple{2, Int}) = d == indim(FluxDense()) ? 0 : randn(T, s) ./ prod(s)
_randoutzeroin(T, d, s::NTuple{N, Int}) where N = d == indim(FluxConv{N-2}()) ? 0 : randn(T, s) ./ prod(s)

"""
    KernelSizeAligned(Δsize; pad)
    KernelSizeAligned(Δs::Integer...;pad)

Strategy for changing kernel size of convolutional layers where filters remain phase aligned. In other words, the same
element indices are removed/added for all filters and only 'outer' elements are dropped or added.

Call with vertex as input to change weights.

### Examples

```jldoctest
julia> using NaiveNASflux, Flux

julia> cv = fluxvertex(Conv((3,3), 1=>1;pad=SamePad()), conv2dinputvertex("in", 1));

julia> cv(ones(Float32, 4,4,1,1)) |> size
(4, 4, 1, 1)

julia> layer(cv).weight |> size
(3, 3, 1, 1)

julia> cv |> KernelSizeAligned(-1, 1; pad=SamePad());

julia> cv(ones(Float32, 4,4,1,1)) |> size
(4, 4, 1, 1)

julia> layer(cv).weight |> size
(2, 4, 1, 1)
```
"""
struct KernelSizeAligned{T, P}
    Δsize::T
    pad::P
end
KernelSizeAligned(Δs::Integer...;pad = ntuple(i -> 0, length(Δs))) = KernelSizeAligned(Δs, pad)

(s::KernelSizeAligned)(l) = selectfilters(layertype(l), l, s)
(s::KernelSizeAligned)(v::AbstractVertex) = mutate_weights(v, s)

otherpars(s::KernelSizeAligned, l) = paddingfor(layertype(l), l, s)
paddingfor(lt, l, s) = ()
paddingfor(::FluxConvolutional{N}, l, s) where N = (;pad = Flux.calc_padding(typeof(l), s.pad, size(weights(l))[1:N] .+ s.Δsize, l.dilation, l.stride))

selectfilters(t, l, s) = ()
selectfilters(::FluxConvolutional, l, s) = selectfilters(s, weights(l))

function selectfilters(s::KernelSizeAligned, pars)
    csize = size(pars)
    mids = csize ./ 2 .+ 0.5
    byfun(dim) = i -> abs(mids[dim] - i) * 1 / (1 + mean(abs.(selectdim(pars, dim, i))))
    ps = Pair{Int, Vector{Int}}[]
    for (dim, Δsize) in enumerate(s.Δsize)
        if Δsize < 0
            els = sort(partialsort(1:csize[dim], 1:csize[dim]+Δsize, by=byfun(dim)))
            push!(ps, dim => els)
        elseif Δsize > 0
            els = -ones(Int, csize[dim] + Δsize)
            offs = Δsize ÷ 2
            els[offs+1:csize[dim]+offs] = 1:csize[dim]
            push!(ps, dim => els)
        end
    end
    return ps
end


@testset "generalwhittle" begin
    @testset "Gradient" begin
        @testset "univariate" begin
            @test_throws ArgumentError WhittleLikelihoodInference.grad_generalwhittle!(ones(2), ones(10), ones(2,10), ones(11))
            @test_throws ArgumentError WhittleLikelihoodInference.grad_generalwhittle!(ones(2), ones(10), ones(2,11), ones(10))
        end
        @testset "multivariate" begin
            S,I = fill([@SMatrix ones(2,2) for i in 1:10],2)
            ∇S = [@SMatrix ones(2,2) for j in 1:2, i in 1:10]
            @test_throws ArgumentError WhittleLikelihoodInference.grad_generalwhittle!(ones(2), S[1:9], ∇S[1:9], I)
            @test_throws ArgumentError WhittleLikelihoodInference.grad_generalwhittle!(ones(2), S[1:9], ∇S, I[1:9])
        end
    end
    @testset "Hessian" begin
        @testset "univariate" begin
            @test_throws ArgumentError WhittleLikelihoodInference.hess_generalwhittle!(ones(2), ones(10), ones(2,10), ones(3,10), ones(11))
            @test_throws ArgumentError WhittleLikelihoodInference.hess_generalwhittle!(ones(2), ones(10), ones(2,11), ones(3,10), ones(10))
            @test_throws ArgumentError WhittleLikelihoodInference.hess_generalwhittle!(ones(2), ones(10), ones(2,10), ones(3,11), ones(10))
        end
        @testset "multivariate" begin
            S,I = fill([@SMatrix ones(2,2) for i in 1:10],2)
            ∇S = [@SMatrix ones(2,2) for j in 1:2, i in 1:10]
            ∇²S = [@SMatrix ones(2,2) for j in 1:3, i in 1:10]
            @test_throws ArgumentError WhittleLikelihoodInference.hess_generalwhittle!(ones(2), S[1:9], ∇S[1:9], ∇²S[1:9], I)
            @test_throws ArgumentError WhittleLikelihoodInference.hess_generalwhittle!(ones(2), S[1:9], ∇S, ∇²S[1:9], I[1:9])
            @test_throws ArgumentError WhittleLikelihoodInference.hess_generalwhittle!(ones(2), S[1:9], ∇S[1:9], ∇²S, I[1:9])
        end
    end
    @testset "Hessian offdiag" begin
        @testset "univariate" begin
            @test_throws ArgumentError WhittleLikelihoodInference.offdiag_add_hess_generalwhittle!(ones(2), ones(10), ones(2,10), ones(2,10), ones(11))
            @test_throws ArgumentError WhittleLikelihoodInference.offdiag_add_hess_generalwhittle!(ones(2), ones(10), ones(2,11), ones(2,10), ones(10))
            @test_throws ArgumentError WhittleLikelihoodInference.offdiag_add_hess_generalwhittle!(ones(2), ones(10), ones(2,10), ones(2,11), ones(10))
        end
        @testset "multivariate" begin
            S,I = fill([@SMatrix ones(2,2) for i in 1:10],2)
            ∇S1,∇S2 = fill([@SMatrix ones(2,2) for j in 1:2, i in 1:10],2)
            @test_throws ArgumentError WhittleLikelihoodInference.offdiag_add_hess_generalwhittle!(ones(2), S[1:9], ∇S1[1:9], ∇S2[1:9], I)
            @test_throws ArgumentError WhittleLikelihoodInference.offdiag_add_hess_generalwhittle!(ones(2), S[1:9], ∇S1, ∇S2[1:9], I[1:9])
            @test_throws ArgumentError WhittleLikelihoodInference.offdiag_add_hess_generalwhittle!(ones(2), S[1:9], ∇S1[1:9], ∇S2, I[1:9])
        end
    end
end