using Statistics
using Plots
using BenchmarkTools
using FFTW
using Statistics
using Optim
using Images, FileIO, ImageIO
using Printf
using Revise
using Profile
using LinearAlgebra
using Distributions
using FITSIO

push!(LOAD_PATH, pwd()*"/main")
using DHC_2DUtils
#include("Deriv_Utils_New.jl")
#import Deriv_Utils_New
push!(LOAD_PATH, pwd()*"/scratch_NM")
using Deriv_Utils_New
using Data_Utils
using Visualization
using LossFuncs
#Remove later

#Jacobian Test Funcs
function derivtestS1(Nx)
   eps = 1e-4
   fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
   im = rand(Float64, Nx, Nx).*0.1
   im[6,6]=1.0
   im1 = copy(im)
   im0 = copy(im)
   im1[2,3] += eps/2
   im0[2,3] -= eps/2
   blarg = Deriv_Utils_New.wst_S1_deriv(im, fhash)
   der0=DHC_compute(im0,fhash,doS2=false,norm=false)
   der1=DHC_compute(im1,fhash,doS2=false,norm=false)
   dS = (der1-der0) ./ eps

   diff_old = dS[3:end]-blarg[2,3,:]

   println(dS[3:end])
   println("and")
   println(blarg[2,3,:])
   println("Mean abs | Mean abs frac", mean(abs.(diff_old)), mean(abs.(diff_old./dS[3:end])))
   println("stdev: ",std(diff_old))
   return blarg[2, 3, :], dS[3:end]
end

function derivtestS1S2(Nx) #MUST RUN DHC_COMPUTE WITH NORM=FALSE.
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    dS1dp_existing = Deriv_Utils_New.wst_S1_deriv(im, fhash)
    dS1S2dp = Deriv_Utils_New.wst_S1S2_deriv(im, fhash) #Deriv_Utils_New.wst_S1S2_deriv(im, fhash)

    der0=DHC_compute(im0,fhash,doS2=true,doS20=false, norm=false)
    der1=DHC_compute(im1,fhash,doS2=true,doS20=false, norm=false)
    dSlim = (der1-der0) ./ eps
    Nf = length(fhash["filt_index"])
    lasts1ind = 2+Nf

    dS1lim23 = dSlim[3:lasts1ind]
    dS2lim23 = dSlim[lasts1ind+1:end]
    dS1dp_23 = dS1S2dp[2, 3, 1:Nf]
    dS2dp_23 = dS1S2dp[2, 3, Nf+1:end]
    #=
    println("Checking dS1dp using existing deriv")
    Nf = length(fhash["filt_index"])
    lasts1ind = 2+Nf
    diff = dS1dp_existing[2, 3, :]-dSlim[3:lasts1ind]
    println(dSlim[3:lasts1ind])
    println("and")
    println(dS1dp_existing[2, 3, :])
    println("stdev: ",std(diff))
    println("--------------------------")
    =#
    println("Checking dS1dp using dS1S2")
    derdiff = dS1dp_23 - dS1lim23
    #println(dS1lim23)
    #println("and")
    #println(dS1dp_23)
    txt= @sprintf("Range of S1 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E ",maximum(der0[3:lasts1ind]),maximum(der1[3:lasts1ind]),minimum(der0[3:lasts1ind]),minimum(der1[3:lasts1ind]))
    print(txt)
    txt = @sprintf("Checking dS1dp using S1S2 deriv, mean(abs(diff/dS1lim)): %.3E", mean(abs.(derdiff./dS1lim23)))
    print(txt)
    println("stdev: ",std(derdiff))
    txt= @sprintf("Range of dS1lim: Max= %.3E Min= %.3E",maximum(abs.(dS1lim23)),minimum(abs.(dS1lim23)))
    print(txt)
    println("--------------------------")
    println("Checking dS2dp using S1S2 deriv")
    Nf = length(fhash["filt_index"])
    println("Shape check",size(dS2dp_23),size(dS2lim23))
    derdiff = dS2dp_23 - dS2lim23 #Column major issue here?
    println("Difference",derdiff)
    txt = @sprintf("Range of S2 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0[lasts1ind+1:end]),maximum(der1[lasts1ind+1:end]),minimum(der0[lasts1ind+1:end]),minimum(der1[lasts1ind+1:end]))
    print(txt)
    txt = @sprintf("Range of dS2lim: Max=%.3E  Min=%.3E \n",maximum(abs.(dS2lim23)),minimum(abs.(dS2lim23)))
    print(txt)
    txt = @sprintf("Difference between dS2dp and dSlim Mean: %.3E stdev: %.3E mean(abs(derdiff/dS2lim)): %.3E \n",mean(derdiff),std(derdiff), mean(abs.(derdiff./dS2lim23)))
    print(txt)

    println("Examining only those S2 coeffs which are greater than eps=1e-3 for der0")
    eps_mask = findall(der0[lasts1ind+1:end] .> 1E-3)
    der0_s2good = der0[lasts1ind+1:end][eps_mask]
    der1_s2good = der1[lasts1ind+1:end][eps_mask]
    dSlim_s2good = dS2lim23[eps_mask]
    derdiff_good = derdiff[eps_mask]
    print(derdiff_good)
    txt = @sprintf("\nRange of S2 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0_s2good),maximum(der1_s2good),minimum(der0_s2good),minimum(der1_s2good))
    print(txt)
    txt = @sprintf("Range of dS2lim: Max=%.3E  Min=%.3E \n",maximum(abs.(dSlim_s2good)),minimum(abs.(dSlim_s2good)))
    print(txt)
    txt = @sprintf("Mean: %.3E stdev: %.3E mean(abs(derdiff/dS2lim)): %.3E \n",mean(derdiff_good),std(derdiff_good), mean(abs.(derdiff_good ./dSlim_s2good)))
    print(txt)
    return dS1dp_23, dS1lim23, reshape(dS2dp_23, (Nf, Nf)), reshape(dS2lim23, (Nf, Nf))
end

function derivtestS12(Nx)
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    dS12dp = Deriv_Utils_New.wst_S12_deriv(im, fhash)

    der0=DHC_compute(im0, fhash, doS2=false, doS20=false, doS12=true, norm=false)
    der1=DHC_compute(im1, fhash, doS2=false, doS20=false, doS12=true, norm=false)
    dS = (der1-der0) ./ eps

    Nf = length(fhash["filt_index"])
    i0 = 3+Nf
    blarg = dS12dp[2, 3, :, :]
    diff = dS[i0:end]-reshape(blarg,Nf*Nf)
    println(dS[i0:end])
    println("and")
    println(blarg)
    println("stdev: ",std(diff))
    println(diff)
    println(size(diff), size(dS), size(dS12dp))
    txt = @sprintf("Mean: %.3E stdev: %.3E mean(abs(diff/dS12dp)): %.3E, mean(abs(dS12dp/dS)) : %.3E \n",mean(diff),std(diff), mean(abs.(diff ./dS12dp[2, 3, :, :][:])), mean(abs.(dS12dp[2, 3, :, :][:]/dS[i0:end])))
    print(txt)

    #plot(dS[3:end])
    #plot!(blarg[2,3,:])
    #plot(diff)
    return reshape(dS12dp[2, 3, :, :], (Nf, Nf)), reshape(dS[i0:end], (Nf, Nf))
end

#Deriv sum test funcs
function dS1S2sum_combtest(fhash) #Tested for non-iso
    (Nf, )    = size(fhash["filt_index"])
    Nx        = fhash["npix"]
    im = rand(Nx,Nx)
    mywts = rand(Nf + (Nf*Nf))
    println(size(mywts))

    # Use new faster code
    sum1 = Deriv_Utils_New.wst_S1S2_derivsum_comb(im, fhash, mywts)
    wts1 = reshape(mywts[1:Nf], (Nf, 1))
    dwS1 = reshape(Deriv_Utils_New.wst_S1_deriv(im, fhash), (Nx^2, Nf)) * wts1

    # Compare to established code
    dS1S2dp = Deriv_Utils_New.wst_S1S2_deriv(im, fhash)
    dS1dp = dS1S2dp[:, :, 1:Nf]
    dS2dp = dS1S2dp[:, :, Nf+1:end]
    term1 = reshape(dS1dp, (Nx^2, Nf)) * reshape(mywts[1:Nf], (Nf, 1))
    sum2 = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (dS2dp[:,:,i].*mywts[i+Nf]) end
    dScomb_brute = term1 + reshape(sum2, (Nx^2, 1))

    println("tot abs mean", mean(abs.(dScomb_brute -sum1)))
    println("tot Stdev: ",std(dScomb_brute-sum1))
    println("Ratio: ", mean(dScomb_brute./sum1))
    println("S1")
    println("Abs mean", mean(abs.(term1 - dwS1)))
    println("Abs mean", std(term1 - dwS1))
    println("Ratio: ", mean(term1./dwS1))
    println("S2 separate")
    sum1 = Deriv_Utils_New.wst_S2_deriv_sum(im, fhash, reshape(mywts[Nf+1:end], (Nf, Nf)))
    sum2 = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (dS2dp[:,:,i].*mywts[Nf+i]) end
    sum2 = reshape(sum2, (Nx^2, 1))
    println("abs mean", mean(abs.(sum2 -sum1)))
    println("Stdev: ",std(sum1-sum2))
    println("Ratio: ", mean(sum2./sum1))
    return
end

function dS20sum_test(fhash) #MODIFIED from sandbox.jl
    (Nf, )    = size(fhash["filt_index"])
    Nx        = fhash["npix"]
    im = rand(Nx,Nx)
    mywts = rand(Nf*Nf)

    # this is symmetrical, but not because S20 is symmetrical!
    wtgrid = reshape(mywts, Nf, Nf)
    wtvec  = reshape(wtgrid, Nf*Nf)

    # Use new faster code
    sum1 = Deriv_Utils_New.wst_S20_deriv_sum(im, fhash, wtgrid, FFTthreads=1)

    # Compare to established code
    dS20 = reshape(Deriv_Utils_New.wst_S20_deriv(im, fhash, FFTthreads=1),Nx,Nx,Nf*Nf)
    sum2 = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (dS20[:,:,i].*wtvec[i]) end
    sum2 = reshape(sum2, (Nx^2, 1))
    println("Abs mean", mean(abs.(sum2 - sum1)))
    println("Abs mean", std(sum2 - sum1))
    println("Ratio: ", mean(sum2./sum1))
    return
end

function imgreconS2test(Nx, pixmask; norm=norm, bias=0.0, std_added=nothing, std_sim=nothing, sim_smoothed=false, sim_smoothedval=0.8, input_smoothed=false, coeff_choice="S2", invcov="Diagonal", mask=nothing, lambda=nothing)
    #=
    Experiments handled by this wrapper:
    S2 | S12 | S20
    Create covariance using simulated noise + img true (Choice of smoothing or not smoothing)
    Init that is smoothed or not smoothed
    Choice of Diagonal | Diagonal+Eps | Full | Full+Eps:
        +Eps adds an epsilon to the diagonal of the covariance matrix to reduce Cond No
        Diagonal considers just the standard deviation of the simulated weights and Full adds an epsilon
    Pixmask: Keeps certain pixels fixed to true values in the image (Not tested using this wrapper yet). False=> Floating pixels.
    Coeffmask: Of length 2+Nf+Nf^2. Choice of coefficients to be optimized. True=>Optimized
    =#

    img = readdust(Nx)
    if std_added==nothing
        std_added=std(img)
    else

    end
    if std_sim==nothing
        std_sim =std(img)
    else

    end
    noise = reshape(rand(Normal(bias, std_added), Nx^2), (Nx, Nx))
    init = img+noise
    if input_smoothed init = imfilter(init, Kernel.gaussian(0.8)) end

    if coeff_choice=="S2"
        #=
        Default assumes coefficients to be optimized are: All S1 and S2 coefficients larger than 1e-5
        =#
        epsilon=1e-5
        fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
        (Nf,) = size(fhash["filt_index"])

        #Uniform Added Noise
        #noise = rand(Nx, Nx).*(2*high) .- high
        #init = copy(img)+noise
        #init[1, 2] += 10.0
        #init[23, 5] -= 25.0
        s2targ = DHC_compute(img, fhash, doS2=true, norm=norm, iso=false)
        #Mask
        if mask==nothing
            mask = (s2targ .>= epsilon)
            mask[1]=false
            mask[2]=false
        else
            if size(mask)[1]!=(2+Nf+Nf^2) error("Wrong length mask") end
        end
        sig2, s2cov = Data_Utils.S2_whitenoiseweights(img, fhash, Nsam=10, loc=bias, sig=std_sim, smooth=sim_smoothed, smoothval = sim_smoothedval, coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
        sig2 = sig2[:]
        println("NF=", size(fhash["filt_index"]), "Sel Coeffs", count((i->(i==true)), mask), size(mask), " Size s2targ, type ", typeof(s2targ[mask]), " Size s2w ", typeof(sig2))

        if invcov=="Diagonal"
            s2icov = invert_covmat(sig2)
        elseif invcov=="Diagonal+Eps"
            s2icov = invert_covmat(sig2, 1e-5)
        elseif invcov=="Full+Eps"
            s2icov = invert_covmat(s2cov, 1e-5)
        else#s2icov
            s2icov = invert_covmat(s2cov)
        end
        recon_img = Deriv_Utils_New.image_recon_derivsum(init, fhash, Float64.(s2targ[mask]), s2icov, pixmask, "S2", coeff_mask=mask, optim_settings=Dict([("iterations", 1000), ("norm", norm)]))
        #recon_img = Deriv_Utils_New.img_reconfunc_coeffmask(init, fhash, s2targ[mask], s2icov[mask, mask], pixmask, Dict([("iterations", 500)]), coeff_mask=mask)
    elseif coeff_choice=="S20"
        #=
        Default assumes coefficients to be optimized are: All S20 coefficients
        =#
        fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
        (Nf,) = size(fhash["filt_index"])
        #img = (img .- mean(img))./std(img) #normalized

        #Uniform Added Noise
        #noise = rand(Nx, Nx).*(2*high) .- high
        #init = copy(img)+noise
        #init[1, 2] += 10.0
        #init[23, 5] -= 25.0
        #Mask
        if mask==nothing
            mask = trues(2+Nf+Nf^2)
            mask[1:2+Nf] .= false
        else
            if size(mask)[1]!=(2+Nf+Nf^2) error("Wrong length mask") end
        end

        s2targ = DHC_compute(img, fhash, doS2=false, doS20=true, norm=norm, iso=false)
        println("NF=", size(fhash["filt_index"]))

        sig2, s2cov = Data_Utils.S20_whitenoiseweights(img, fhash, Nsam=10, loc=bias, sig=std_sim, smooth=sim_smoothed, smoothval = sim_smoothedval, coeff_mask=mask, coeff_choice="S20") #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
        sig2 = sig2[:]
        println("NF=", size(fhash["filt_index"]), "Sel Coeffs", count((i->(i==true)), mask), size(mask), " Size s2targ, type ", typeof(s2targ[mask]), " Size s2w ", typeof(sig2))

        if invcov=="Diagonal"
            s2icov = invert_covmat(sig2)
        elseif invcov=="Diagonal+Eps"
            s2icov = invert_covmat(sig2, 1e-10)
        elseif invcov=="Full+Eps"
            s2icov = invert_covmat(s2cov, 1e-10)
        else#s2icov
            s2icov = invert_covmat(s2cov)
        end
        if lambda==nothing
            recon_img = Deriv_Utils_New.image_recon_derivsum(init, fhash, Float64.(s2targ[mask]), s2icov, pixmask, "S20", coeff_mask=mask, optim_settings=Dict([("iterations", 1000), ("norm", norm)]))
        else
            recon_img = Deriv_Utils_New.image_recon_derivsum_regularized(init, fhash, Float64.(s2targ[mask]), s2icov, pixmask, "S20", coeff_mask=mask, optim_settings=Dict([("iterations", 1000), ("norm", norm)]), lambda=0.001)
        end

    else
        if coeff_choice!="S12" error("Not implemented") end
        img = readdust(Nx)
        fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
        (Nf,) = size(fhash["filt_index"])

        #Uniform Added Noise
        #noise = rand(Nx, Nx).*(2*high) .- high
        #init = copy(img)+noise
        #init[1, 2] += 10.0
        #init[23, 5] -= 25.0
        s2targ = DHC_compute(img, fhash, doS2=false, doS20=false, doS12=true, norm=norm, iso=false)
        #Mask
        if mask==nothing
            mask = (s2targ .>= 1e-5)
            mask[1:Nf+2] .= false
        else

        end
        sig2, s2cov = Data_Utils.S20_whitenoiseweights(img, fhash, Nsam=10, loc=bias, sig=std_sim, smooth=sim_smoothed, smoothval = sim_smoothedval, coeff_choice="S12", coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
        sig2 = sig2[:]
        println("NF=", size(fhash["filt_index"]), "Sel Coeffs", count((i->(i==true)), mask), size(mask), " Size s2targ, type ", typeof(s2targ[mask]), " Size s2w ", typeof(sig2))

        if invcov=="Diagonal"
            s2icov = invert_covmat(sig2)
        elseif invcov=="Diagonal+Eps"
            s2icov = invert_covmat(sig2, 1e-10)
        elseif invcov=="Full+Eps"
            s2icov = invert_covmat(s2cov, 1e-10)
        else#s2icov
            s2icov = invert_covmat(s2cov)
        end
        recon_img = Deriv_Utils_New.image_recon_derivsum(init, fhash, Float64.(s2targ[mask]), s2icov, pixmask, "S12", coeff_mask=mask, optim_settings=Dict([("iterations", 1000), ("norm", norm)]))

    end
    return img, init, recon_img
end

function logimgreconS2test(Nx, pixmask; norm=norm, std_added=nothing, std_sim=nothing, sim_smoothed=false, sim_smoothedval=0.8, input_smoothed=false, coeff_choice="S2", invcov="Diagonal", mask=nothing)
    #=
    Uses the coefficients of the LOG of the image
    Experiments handled by this wrapper:
    S2 | S12 | S20
    Create covariance using simulated noise + img true (Choice of smoothing or not smoothing)
    Init that is smoothed or not smoothed
    Choice of Diagonal | Diagonal+Eps | Full | Full+Eps:
        +Eps adds an epsilon to the diagonal of the covariance matrix to reduce Cond No
        Diagonal considers just the standard deviation of the simulated weights and Full adds an epsilon
    Pixmask: Keeps certain pixels fixed to true values in the image (Not tested using this wrapper yet). False=> Floating pixels.
    Coeffmask: Of length 2+Nf+Nf^2. Choice of coefficients to be optimized. True=>Optimized
    =#
    img = readdust(Nx)
    oriimg = copy(img)
    logimg = log.(img)
    if std_added==nothing
        std_added = std(oriimg)
    else

    end

    epsilon=1e-5
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
    (Nf,) = size(fhash["filt_index"])

    #Std Normal
    noise = reshape(rand(Normal(0.0, std_added), Nx^2), (Nx, Nx))
    init = oriimg+noise

    if std_sim==nothing
        std_sim = std(oriimg)
    else

    end

    if input_smoothed init = imfilter(init, Kernel.gaussian(0.8)) end

    if coeff_choice=="S2"
        s2targ = DHC_compute(logimg, fhash, doS2=true, norm=norm, iso=false)
        #Mask: Default assumes coefficients to be optimized are: All S1 and S2 coefficients larger than 1e-5
        if mask==nothing
            mask = (s2targ .>= epsilon)
            mask[1]=false
            mask[2]=false
        else
            if size(mask)[1]!=(2+Nf+Nf^2) error("Wrong length mask") end
        end
    elseif coeff_choice=="S20"
        s2targ = DHC_compute(logimg, fhash, doS20=true, doS2=false, norm=norm, iso=false)
        #Mask: Default assumes only S20
        if mask==nothing
            mask = trues(2+Nf+Nf^2)
            mask[1:2+Nf] .= false
        else
            if size(mask)[1]!=(2+Nf+Nf^2) error("Wrong length mask") end
        end
    else
        if coeff_choice=="S12"
            s2targ = DHC_compute(logimg, fhash, doS20=false, doS12=true, doS2=false, norm=norm, iso=false)
            #Mask
            if mask==nothing
                mask = (s2targ .>= 1e-5)
                mask[1:Nf+2] .= false
            else
                if size(mask)[1]!=(2+Nf+Nf^2) error("Wrong length mask") end
            end
        end
    end

    sig2, s2cov = whitenoiseweights_forlog(oriimg, fhash, coeff_choice, Nsam=10, loc=0.0, sig=std_sim, smooth=sim_smoothed, smoothval = sim_smoothedval, coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
    sig2 = sig2[:]
    println("NF=", size(fhash["filt_index"]), "Sel Coeffs", count((i->(i==true)), mask), size(mask), " Size s2targ, type ", typeof(s2targ[mask]), " Size s2w ", typeof(sig2))

    if invcov=="Diagonal"
        s2icov = invert_covmat(sig2)
    elseif invcov=="Diagonal+Eps"
        s2icov = invert_covmat(sig2, 1e-5)
    elseif invcov=="Full+Eps"
        s2icov = invert_covmat(s2cov, 1e-5)
    else#s2icov
        s2icov = invert_covmat(s2cov)
    end
    logrecon_img = Deriv_Utils_New.image_recon_derivsum(log.(init), fhash, Float64.(s2targ[mask]), s2icov, pixmask, coeff_choice, coeff_mask=mask, optim_settings=Dict([("iterations", 1000), ("norm", norm)]))
    recon_img = exp.(logrecon_img)
    return oriimg, init, recon_img
end#redundant?

function dbn_experiment(true_img, noisy_init, dbn_covar, pixmask; norm=norm, bias=0.0, std_sim=nothing, sim_smoothed=false, sim_smoothedval=0.8, input_smoothed=false, coeff_choice="S2", invcov="Diagonal", mask=nothing, lambda=nothing)
    #=
    Experiments handled by this wrapper:
    S20
    Regularized vs Non regularized
    Use externally calculated covariance
    Choice of Diagonal | Diagonal+Eps | Full | Full+Eps:
        +Eps adds an epsilon to the diagonal of the covariance matrix to reduce Cond No
        Diagonal considers just the standard deviation of the simulated weights and Full adds an epsilon
    Pixmask: Keeps certain pixels fixed to true values in the image (Not tested using this wrapper yet). False=> Floating pixels.
    Coeffmask: Of length 2+Nf+Nf^2. Choice of coefficients to be optimized. True=>Optimized
    =#

    (Nx, ) = size(true_img)

    #=
    Default assumes coefficients to be optimized are: All S20 coefficients
    =#
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
    (Nf,) = size(fhash["filt_index"])
    #img = (img .- mean(img))./std(img) #normalized

    #Uniform Added Noise
    #noise = rand(Nx, Nx).*(2*high) .- high
    #init = copy(img)+noise
    #init[1, 2] += 10.0
    #init[23, 5] -= 25.0
    #Mask
    if mask==nothing
        mask = trues(2+Nf+Nf^2)
        mask[1:2+Nf] .= false
    else
        if size(mask)[1]!=(2+Nf+Nf^2) error("Wrong length mask") end
    end

    s2targ = DHC_compute(img, fhash, doS2=false, doS20=true, norm=false, iso=false)
    println("NF=", size(fhash["filt_index"]))
    println("NF=", size(fhash["filt_index"]), "Sel Coeffs", count((i->(i==true)), mask), size(mask), " Size s2targ, type ", typeof(s2targ[mask]), " Size s2w ", typeof(sig2))

    if invcov=="Diagonal"
        s2icov = invert_covmat(Diagonal(dbn_covar))
    elseif invcov=="Diagonal+Eps"
        s2icov = invert_covmat(Diagonal(dbn_covar), 1e-10)
    elseif invcov=="Full+Eps"
        s2icov = invert_covmat(dbn_covar, 1e-10)
    else#s2icov
        s2icov = invert_covmat(dbn_covar)
    end
    if lambda==nothing
        recon_img = Deriv_Utils_New.image_recon_derivsum(init, fhash, Float64.(s2targ[mask]), s2icov, pixmask, "S20", coeff_mask=mask, optim_settings=Dict([("iterations", 1000), ("norm", norm)]))
    else
        recon_img = Deriv_Utils_New.image_recon_derivsum_regularized(init, fhash, Float64.(s2targ[mask]), s2icov, pixmask, "S20", coeff_mask=mask, optim_settings=Dict([("iterations", 1000), ("norm", norm)]), lambda=0.001)
    end
    return img, init, recon_img
end#redundant?

#Calling functions############################################
ds1code, ds1lim = derivtestS1(16)
ds1code, ds1lim, ds2code, ds2lim = derivtestS1S2(16)
ds12code, ds12lim = derivtestS12(16)
#All above work correctly for norm=false

Nx=16
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
dS1S2sum_combtest(fhash)

dS20sum_test(fhash)

#Image Recon tests############################################
#S2 with no pixmask, coeff_mask for large selection
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S2", invcov="Diagonal")
Visualization.plot_synth_QA(img, init, recon, fink_filter_hash(1, 8, nx=32, pc=1, wd=1, Omega=true), fname="scratch_NM/TestPlots/WhiteNoiseS2tests/S20_passed_chisq.png")
mean(abs.(init - img)), mean(abs.(recon - img))
#Not working
#Compare all quantities with the equivalentimgrecon code that didnt use derivsum and only uses coeff_mask
#That fails with DHC_compute too. Try DHC_compute_old

#Examining S2w with new dhc-compute
Nx=32
img = readdust(Nx)
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf,) = size(fhash["filt_index"])
s2w, s2cov = S2_weights(img, fhash, 10, iso=false, norm=true)
mask = s2w .>= 1e-5
mask[1]=false
mask[2]=false
println("Max/Min coeffmasked S2w", maximum(s2w[mask]), minimum(s2w[mask]))
println("Max/Min coeffmasked S2icov", maximum(s2icov[mask, mask]), minimum(s2icov[mask, mask]))

Nx=64
img = readdust(Nx)
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf,) = size(fhash["filt_index"])
s2targ = DHC_compute(img, fhash, doS2=false, doS20=false, doS12=true, norm=false, iso=false)
std_sim=std(img)
s2w, s2icov = Data_Utils.S20_whitenoiseweights(img, fhash, Nsam=10, loc=0.0, sig=std_sim, smooth=false, smoothval =0.8, coeff_choice="S12") #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
mask = s2targ .>= 1e-5
mask[1:Nf+2] .= false
println("Nf=", size(fhash["filt_index"]))
println("NumSelCoeffs=", count((i->(i==1)), mask))

#NOTE: #1 Choosing what the covariance matrix should be for different coefficients
#S2
Nx=64
img = readdust(Nx)
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf,) = size(fhash["filt_index"])
s2targ = DHC_compute(img, fhash, doS2=true, doS20=false, doS12=false, norm=false, iso=false)
#Incl only S1, S2 that are larger than 1e-5
mask = s2targ .>= 1e-5
mask[1:2] .= false
std_sim=std(img)
s2w, s2cov = Data_Utils.S2_whitenoiseweights(img, fhash, Nsam=10, loc=0.0, sig=std_sim, smooth=false, smoothval =0.8, coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
Data_Utils.invert_covmat(s2w) #1.41E18, #1.26E-20
Data_Utils.invert_covmat(s2w, 1e-10) #39k, #1.26E-20
Data_Utils.invert_covmat(s2w, 1e-5) #4.92, #1.26E-20 Try this!

Data_Utils.invert_covmat(s2cov) #2e33, #1e6 CAN'T DO THIS!
Data_Utils.invert_covmat(s2cov, 1e-10) #800k
Data_Utils.invert_covmat(s2cov, 1e-5) #9.24, 1e-19 Try this!

#Incl all S1, S2
mask = trues(2+Nf+Nf^2)
mask[1:2] .= false
std_sim=std(img)
s2w, s2cov = Data_Utils.S2_whitenoiseweights(img, fhash, Nsam=10, loc=0.0, sig=std_sim, smooth=false, smoothval =0.8, coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
Data_Utils.invert_covmat(s2w) #1.41E28, #1.26E-20
Data_Utils.invert_covmat(s2w, 1e-5) #6.2, 1-20
Data_Utils.invert_covmat(s2cov) #2e43, 1e11 CAN'T DO THIS!
Data_Utils.invert_covmat(s2cov, 1e-5) #12, 1e-19 Try this!

#S20
Nx=64
img = readdust(Nx)
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf,) = size(fhash["filt_index"])
s2targ = DHC_compute(img, fhash, doS2=false, doS20=true, doS12=false, norm=false, iso=false)
#Including all S20
mask=trues(2+Nf+Nf^2)
mask[1:Nf+2] .= false
std_sim=std(img)
s2w, s2cov = Data_Utils.S20_whitenoiseweights(img, fhash, Nsam=10, loc=0.0, sig=std_sim, smooth=false, smoothval =0.8, coeff_choice="S20", coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
Data_Utils.invert_covmat(s2w) #513, #1.26E-20
Data_Utils.invert_covmat(s2w, 1e-5) #6.7, #1.26E-20 Try this!

Data_Utils.invert_covmat(s2cov) #Breaks
Data_Utils.invert_covmat(s2cov, 1e-10) #5e6, 1e-10
Data_Utils.invert_covmat(s2cov, 1e-5) #55, 1e-17 Choosing this!

#Incl only S20 that are larger than 1e-5
mask = s2targ .>= 1e-5
mask[1:Nf+2] .= false
std_sim=std(img)
s2w, s2cov = Data_Utils.S20_whitenoiseweights(img, fhash, Nsam=10, loc=0.0, sig=std_sim, smooth=false, smoothval =0.8, coeff_choice="S20", coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
Data_Utils.invert_covmat(s2w) #361
Data_Utils.invert_covmat(s2w, 1e-5) #6
Data_Utils.invert_covmat(s2cov) #e7, Breaks
Data_Utils.invert_covmat(s2cov, 1e-5) #51

#S12
Nx=64
img = readdust(Nx)
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf,) = size(fhash["filt_index"])
s2targ = DHC_compute(img, fhash, doS2=false, doS20=false, doS12=true, norm=false, iso=false)
#Including all S12:
mask=trues(2+Nf+Nf^2)
mask[1:Nf+2] .= false
std_sim=std(img)
s2w, s2cov = Data_Utils.S20_whitenoiseweights(img, fhash, Nsam=10, loc=0.0, sig=std_sim, smooth=false, smoothval =0.8, coeff_choice="S12", coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
Data_Utils.invert_covmat(s2w, 1e-5) #e7
Data_Utils.invert_covmat(s2cov, 1e-5) #e8

#Including only S12 larger than some eps
mask = s2targ .>=1e-5
mask[1:Nf+2] .= false
std_sim=std(img)
s2w, s2cov = Data_Utils.S20_whitenoiseweights(img, fhash, Nsam=10, loc=0.0, sig=std_sim, smooth=false, smoothval =0.8, coeff_choice="S12", coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
Data_Utils.invert_covmat(s2w, 1e-5) #e7
Data_Utils.invert_covmat(s2cov, 1e-5) #e8
#appear to be some coefficients which aren't too small O(0.6) but still have 0 std
findall(s2w .<1e-5)
s2targ[mask][findall(s2w .<1e-5)]
#Exclude all below 1e0
mask = s2targ .>=1 #This does not exclude any of the diagonal elements (S1)
mask[1:Nf+2] .= false
std_sim=std(img)
s2w, s2cov = Data_Utils.S20_whitenoiseweights(img, fhash, Nsam=10, loc=0.0, sig=std_sim, smooth=false, smoothval =0.8, coeff_choice="S12", coeff_mask=mask) #S2_uniweights(img, fhash, Nsam=10, high=high, iso=false, norm=norm)
Data_Utils.invert_covmat(s2w, 1e-5) #100k Try this!
Data_Utils.invert_covmat(s2cov, 1e-5) #e8, Num Err=1e-8

#NOTE: Trying all the variants with reasonable condition numbers
#S2: Diagonal, All S1, S2 coeffs greater than eps=1e-5: 22->7 Works well
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S2", invcov="Diagonal+Eps")
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S2_Diag_allgreaterthaneps.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#S2: FullCov, All S1, S2 coeffs greater than eps=1e-5: Breaks
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S2", invcov="Full+Eps")
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S2_Cov_allgreaterthaneps.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#S2: Incl All S1+S2, adding an eps to the super small coeffs
#Diagonal mat,  22->8 Works well
Nx=64
img = readdust(Nx)
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf,) = size(fhash["filt_index"])
s2targ = DHC_compute(img, fhash, doS2=true, doS20=false, doS12=false, norm=false, iso=false)
mask = trues(2+Nf+Nf^2)
mask[1:2] .= false
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S2", invcov="Diagonal+Eps",mask=mask)
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S2_DiagEps_allincl.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#Cov, 22->8 Works well
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S2", invcov="Full+Eps",mask=mask)
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S2_CovEps_allincl.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#S20: All S20 only incl: Best so far.
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S20", invcov="Diagonal+Eps")
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S20_DiagEps_allincl.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#S20: All S20 only incl: Breaks
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S20", invcov="Full+Eps")
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S20_FullEps_allincl.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#S12:

#NOTE: With Smoothing
#S2
#Cov, 22->8 Works well
mask = trues(2+Nf+Nf^2)
mask[1:2] .= false
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=true, input_smoothed=true, coeff_choice="S2", invcov="Full+Eps",mask=mask)
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Smoothed/S2_CovEps_allincl.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#S20: All S20 only incl: Best so far. All S20 incl.
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=true, input_smoothed=true, coeff_choice="S20", invcov="Diagonal+Eps")
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Smoothed/S20_DiagEps_allincl.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#S12: All S12 only incl
Nx=64
img = readdust(Nx)
mean(img)
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf,) = size(fhash["filt_index"])
s12targ = DHC_compute(img, fhash, doS2=false, doS20=false, doS12=true, norm=false, iso=false)
mask = s2targ .>=1 #This does not exclude any of the diagonal elements (S1)
mask[1:Nf+2] .= false
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=true, input_smoothed=true, coeff_choice="S12", invcov="Diagonal+Eps", mask=mask)
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Smoothed/S12_DiagEps_allincl.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#NOTE: With the log of the image's coefficients
img = readdust(64)
img, init, recon = logimgreconS2test(64, falses((64, 64)), norm=false, std_added=std(img)/5.0, std_sim=std(img)/5.0, sim_smoothed=false, input_smoothed=false, coeff_choice="S20", invcov="Diagonal+Eps")
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S20_LogCoeff_DiagEps_allincl.png")
mean(abs.(init - img)), mean(abs.(recon - img))

img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, std_added=std(img)/5.0, std_sim=std(img)/5.0, input_smoothed=false, coeff_choice="S20", invcov="Diagonal+Eps")
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S20_DiagEps_allincl_comparison-with-log.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#NOTE: Can it correct a bias?
img = readdust(64)
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, bias=mean(img)/10.0, sim_smoothed=false, input_smoothed=false, coeff_choice="S20", invcov="Diagonal+Eps")
mean(abs.(init - img)), mean(abs.(recon - img))
Visualization.plot_diffscales([img, init, recon, recon - img], ["Ground Truth", "Init (N(0.1<I_gt>, std(I_gt))", "Reconstruction", "Reconstruction - Ground Truth"], fname="scratch_NM/TestPlots/CoeffCombinations_WhiteNoise_Nosmooth/S20_DiagEps_allincl_BiasedWhiteNoise")

#NOTE: What effect does adding a regularizing term have?
img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S20", invcov="Diagonal+Eps")
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/3-10_Plots/NoReg_1ksamps.png")
mean(abs.(init - img)), mean(abs.(recon - img))

img, init, recon = imgreconS2test(64, falses((64, 64)), norm=false, sim_smoothed=false, input_smoothed=false, coeff_choice="S20", invcov="Diagonal+Eps", lambda=0.1)
Visualization.plot_synth_QA(img, init, recon, fname="scratch_NM/3-10_Plots/Reg_1ksamps_0-1.png")
mean(abs.(init - img)), mean(abs.(recon - img))

#NOTE: Using the Gaussian covariance of fits
im, covar = mldust(64)
im, covar = mldust(64, logbool=False)
true_img = im[:, :, 1]
Nx=64
noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
Ncovsamp = 10000
(Nf, ) = size(filter_hash["filt_index"])
s20_dbn = zeros(Float64, Ncovsamp, 2+Nf+Nf^2)
for idx=1:Ncovsamp
    s20_dbn[idx, :] = DHC_compute(im[:, :, idx], filter_hash, doS2=false, doS20=true, norm=false, iso=false)
end

s_targ_mean = mean(s20_dbn, dims=1)
scov  = (s20_dbn .- s_targ_mean)' * (s20_dbn .- s_targ_mean) ./(Ncovsamp-1)

sfddbn_experiment(im[:, :, 1], init, scov, falses(Nx, Nx), norm=false,) #complete

mean(im[:, :, 1])
#####################################################################################################
#Using NEW Reconstruction wrapper and SFD dbn

img, covar = mldust(64, logbool=false)

true_img = im[:, :, 1]
Nx=64
noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf, ) = size(filter_hash["filt_index"])
coeff_mask = falses(2+Nf+Nf^2)
coeff_mask[Nf+3:end] .= true

dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>false)
white_noise_args = Dict(:loc=>0.0, :sig=>1.0, :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>false, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false)])
recon_settings = Dict([("target_type", "ground truth"),("covar_type", "white_noise"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
    ("Covar_type", "sfd_dbn"), ("apd", "non_apd"), ("optim_settings", optim_settings), ("lambda", (std(true_img).^(-2))./100), ("white_noise_args", white_nois)]) #Add constraints
#"white_noise_args": args for white_noise sims if GaussianLoss and Covar_type='white_noise'. :loc...

s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings)
recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov

recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)
#Visualization.plot_synth_QA(true_img, init, recon_img, fname="scratch_NM/NewWrapper/sfd_dbntest.png")
Visualization.plot_diffscales([true_img, init, recon_img, recon_img - true_img], ["GT", "Init", "Reconstruction", "Residual"], fname="scratch_NM/NewWrapper/sfd_dbntest_std-by100.png")

@benchmark dhc = DHC_compute_apd(true_img, filter_hash; Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>false)...)
@benchmark dhc_apd = DHC_compute_apd(true_img, filter_hash; Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true)...)

fname = "scratch_NM/data/dust10000.fits"
f = FITS(fname, "r")
big = read(f[1])
big64 = imresize(big, 64, 64, 10000)
img_old, covar = mldust(64, logbool=false)
mean(big64[:, :, 1])
bigf64 = imresize(Float64.(big), 64, 64, 10000)

(_,__,Nslice) = size(big)
println(Nslice, " slices")
size(big)
minimum(big64)
minimum(im)
heatmap(big[:, :, 1])
heatmap(im[:, :, 1])
##############################################################################################################################

#Using NEW Reconstruction wrapper and SFD dbn

im, covar = mldust(, logbool=false)
true_img = im[:, :, 1]
Nx=64
noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(0.8))
filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf, ) = size(filter_hash["filt_index"])
coeff_mask = falses(2+Nf+Nf^2)
coeff_mask[Nf+3:end] .= true

dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true)
white_noise_args = Dict(:loc=>0.0, :sig=>1.0, :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>false, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 100), ("norm", false)])
recon_settings = Dict([("target_type", "ground truth"),("covar_type", "white_noise"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
 ("apd", "non_apd"), ("optim_settings", optim_settings), ("lambda", 0.0), ("white_noise_args", white_noise_args)]) #Add constraints
#"white_noise_args": args for white_noise sims if GaussianLoss and Covar_type='white_noise'. :loc...

s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings)
recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov

recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)
Visualization.plot_diffscales([true_img, init, recon_img, recon_img - true_img], ["GT", "Init", "Reconstruction", "Residual"], fname="scratch_NM/NewWrapper/sfd_gtwhitenoisecov_lam0_smoothinit.png")

#Old image##############################################################################
Nx=64
im = readdust(Nx)
true_img = im

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(0.8))
filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf, ) = size(filter_hash["filt_index"])
coeff_mask = falses(2+Nf+Nf^2)
coeff_mask[Nf+3:end] .= true

dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>false)
white_noise_args = Dict(:loc=>0.0, :sig=>1.0, :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>false, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false)])
recon_settings = Dict([("target_type", "ground truth"),("covar_type", "white_noise"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
 ("apd", "non_apd"), ("optim_settings", optim_settings), ("lambda", 0.0), ("white_noise_args", white_noise_args)]) #Add constraints
#"white_noise_args": args for white_noise sims if GaussianLoss and Covar_type='white_noise'. :loc...

s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings)
recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov

recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)
Visualization.plot_diffscales([true_img, init, recon_img, recon_img - true_img], ["GT", "Init", "Reconstruction", "Residual"], fname="scratch_NM/NewWrapper/sfd_gtwhitenoisecov_lam0_smoothinit_apd.png")

##############################################################################################################################
using StatsBase

image = readdust(256)
apdimg_fixed = apodizer(image)
apdimg = apodizer_ori(image, 1, 1, 255)
mean(image)
mean(apdimg_fixed)

heatmap(image)
heatmap(apdimg)

image = readdust(64)
apdimg = apodizer(image, 1, 1, 63)

heatmap(image)
heatmap(apdimg)

big64

apdimg_new = apodizer_new(image)
heatmap(apdimg_new)
mean(image)

mean(apdimg_new)
mean(apdimg)

#Apodization func check
image = readdust(64)
Nx=64
Amat = wind_2d(Nx)
datad_w = fweights(Amat);#why semi-col?#BUG: Replace 256 with im_size? #<A>
meanVal = mean(image,datad_w) #<AF>
temp2d_a = (image.-meanVal).*wind_2d(Nx).+meanVal #<AF>
mean(Amat.* (image .- meanVal)) #E-14

##After modifying dS20 derivsum

filter_hash = fink_filter_hash(1, 8, nx=64, Omega=true)

dS20sum_test(filter_hash)


#=

   Dice game.

  From
  http://www.informs.org/ORMS-Today/Public-Articles/December-Volume-38-Number-6/THE-PUZZLOR
  """
  A         B           C           D
     7          4          6            5
   1 1 1 1    4 4 4 4    2 2 2 2      3 5 3 3
     7          4          6            5
  Figure 1 shows four dice with varying numbers on each face. You and
  three friends will play a simple game where you will each roll one
  of the dice and the highest number wins. You get first pick from
  the dice.

  Which die should you choose in order to maximize your chance of winning?
  """

  Analysis
    A vs B:  A wins 2 times, B wins 4 times
    A vs C:  A wins 2 times, C wins 4 times
    A vs D:  A wins 2 times, D wins 4 times
    B vs C:  B wins 4 times, C wins 2 times
    B vs D:  B wins 3 times, D wins 3 times
    C vs D:  C wins 2 times, D wins 4 times

    Wins (of row vs col)
      A  B  C  D    Sum (winning)
   A  -  2  2  2      6    0.166667
   B  4  -  4  3     11    0.305556   <---
   C  4  2  -  2      8    0.222222
   D  4  4  3  -     11    0.305556   <---
                    ---
                     36

  This model give simular result:

  Summary Statistics
  parameters      mean       std   naive_se      mcse        ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64    Float64   Float64       Float64

           A    2.0490    1.9556     0.0618    0.0520   638.1785    0.9994      115.5283
           B    3.5920    1.8480     0.0584    0.0572   600.4959    0.9990      108.7067
           C    2.6830    0.6519     0.0206    0.0173   638.1785    0.9994      115.5283
           D    3.6760    0.9726     0.0308    0.0263   677.3357    1.0016      122.6169
          Ap    0.1708    0.1630     0.0052    0.0043   638.1785    0.9994      115.5283
          Bp    0.2993    0.1540     0.0049    0.0048   600.4959    0.9990      108.7067
          Cp    0.2236    0.0543     0.0017    0.0014   638.1785    0.9994      115.5283
          Dp    0.3063    0.0811     0.0026    0.0022   677.3357    1.0016      122.6169

   Matrix of throws[i,j]:
   0.0000 1.6550 2.3100 2.9650
   1.6620 0.0000 2.3380 3.0700
   2.3240 2.3450 0.0000 3.6550
   2.9860 3.0140 3.6620 0.0000

   So one should select either die D (or perhaps die B).

=#

using Turing, StatsPlots, DataFrames
using Printf
include("jl_utils.jl")

@model function dice_game()
    S = [1 1 1 1 7 7
         4 4 4 4 4 4
         2 2 2 2 6 6
         3 5 3 3 5 5]

    throws = tzeros(4,4)
    for I in 1:4
        for J in 1:4
            if I != J
                # Roll the two dice
                Iix ~ DiscreteUniform(1,6)
                Jix ~ DiscreteUniform(1,6)
                # Save the winner
                throws[I,J] ~ Dirac(S[I,Iix] > S[J,Jix] ? I : J)
            else
                throws[I,J] ~ Dirac(0)
            end
        end
    end

    # Number of wins
    A ~ Dirac(sum([throws[i,j] == 1 for i in 1:4 for j in 1:4]))
    B ~ Dirac(sum([throws[i,j] == 2 for i in 1:4 for j in 1:4]))
    C ~ Dirac(sum([throws[i,j] == 3 for i in 1:4 for j in 1:4]))
    D ~ Dirac(sum([throws[i,j] == 4 for i in 1:4 for j in 1:4]))

    # Percentages
    total = A + B + C + D
    Ap ~ Dirac(A / total)
    Bp ~ Dirac(B / total)
    Cp ~ Dirac(C / total)
    Dp ~ Dirac(D / total)

end

model = dice_game()

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 1_000)
chns = sample(model, PG(5), 1_000)
# chns = sample(model, SMC(10), 1_000)
# chns = sample(model, IS(), 1_000)

display(chns[[:A,:B,:C,:D,:Ap,:Bp,:Cp,:Dp]])
# display(chns)
# display(plot(chns))

# show_var_dist_pct(chns,:A)
# show_var_dist_pct(chns,:B)
# show_var_dist_pct(chns,:C)
# show_var_dist_pct(chns,:D)

println("Matrix of throws[i,j]:")
for i in 1:4
    for j in 1:4
        @printf("%2.4f ", mean(chns[Symbol("throws[$i,$j]")]))
    end
    println()
end


using StochasticDiffEq, DiffEqDevTools, ParameterizedFunctions
using Plots; gr()
const N = 100

f = @ode_def LotkaVolterraTest begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d

p = [1.5,1.0,3.0,1.0]

function g(du,u,p,t)
  du .= 0.1u
end
u0 = [1.0;1.0]
tspan = (0.0,10.0)
prob = SDEProblem(f,g,u0,tspan,p);

sol = solve(prob,SRIW1(),abstol=1e-4,reltol=1e-4)
plot(sol)

reltols = 1.0 ./ 4.0 .^ (2:4)
abstols = reltols#[0.0 for i in eachindex(reltols)]
setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./12.0.^((1:length(reltols)) .+ 1.5))
          Dict(:alg=>RKMil(),:dts=>1.0./12.0.^((1:length(reltols)) .+ 1.5),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./4.0.^((1:length(reltols)) .+ 5),:adaptive=>false)
          Dict(:alg=>SRIW2())
          Dict(:alg=>SOSRI())
          Dict(:alg=>SOSRI2())
          ]
test_dt = 1/10^2
appxsol_setup = Dict(:alg=>SRIW1(),:abstol=>1e-4,:reltol=>1e-4)
wp = WorkPrecisionSet(prob,abstols,reltols,setups,test_dt;
                                     maxiters = 1e7,
                                     verbose=false,save_everystep=false,
                                     parallel_type = :threads,
                                     appxsol_setup = appxsol_setup,
                                     numruns_error=N,error_estimate=:final)
plot(wp)

reltols = 1.0 ./ 4.0 .^ (2:4)
abstols = reltols#[0.0 for i in eachindex(reltols)]
setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./12.0.^((1:length(reltols)) .+ 1.5))
          Dict(:alg=>RKMil(),:dts=>1.0./12.0.^((1:length(reltols)) .+ 1.5),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./4.0.^((1:length(reltols)) .+ 5),:adaptive=>false)
          Dict(:alg=>SRIW2())
          Dict(:alg=>SOSRI())
          Dict(:alg=>SOSRI2())
          ]
test_dt = 1e-2
appxsol_setup = Dict(:alg=>SRIW1(),:abstol=>1e-4,:reltol=>1e-4)
wp = WorkPrecisionSet(prob,abstols,reltols,setups,test_dt;
                                     maxiters = 1e7,
                                     verbose=false,save_everystep=false,
                                     parallel_type = :none,
                                     appxsol_setup = appxsol_setup,
                                     numruns_error=N,error_estimate=:weak_final)
plot(wp;legend=:topleft)

sample_size = Int[10;1e2;1e3;1e4]
se = get_sample_errors(prob,setups[6],test_dt,numruns=sample_size,
                                      sample_error_runs = 100_000,solution_runs=100)

plot(wp;legend=:topleft)
times = [wp[i].times for i in 1:length(wp)]
times = [minimum(minimum(t) for t in times),maximum(maximum(t) for t in times)]
plot!([se[end];se[end]],times,color=:orange,linestyle=:dash,label="Sample Error: 1000",lw=3)

using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


using DataFrames
using CategoricalArrays
using Dates

using RCall
let dir = datadir("pest_data.rds")
    R"""
    library(dplyr)
    library(lubridate)

    pest_data <- readRDS($dir)
        N_buildings <- length(unique(pest_data$building_id))
    """
end

R"""
N_months <- length(unique(pest_data$date))

# Add some IDs for building and month
pest_data <- pest_data %>%
  mutate(
    building_fac = factor(building_id, levels = unique(building_id)),
    building_idx = as.integer(building_fac),
    ids = rep(1:N_months, N_buildings),
    mo_idx = lubridate::month(date)
  )
"""

R"""
# Center and rescale the building specific data
building_data <- pest_data %>%
    select(
      building_idx,
      live_in_super,
      age_of_building,
      total_sq_foot,
      average_tenant_age,
      monthly_average_rent
    ) %>%
    unique() %>%
    arrange(building_idx) %>%
    select(-building_idx) %>%
    scale(scale=FALSE) %>%
    as.data.frame() %>%
    mutate( # scale by constants
      age_of_building = age_of_building / 10,
      total_sq_foot = total_sq_foot / 10000,
      average_tenant_age = average_tenant_age / 10,
      monthly_average_rent = monthly_average_rent / 1000
    ) %>%
    as.matrix()
"""

# TODO: Fix the sorting in the Julia version
building_data_r = rcopy(R"building_data")

# Data maniplation in Julia

pest_data.building_fac = categorical(pest_data.building_id)
pest_data.building_idx = levelcode.(pest_data.building_fac)

pest_data.mo_idx = month.(pest_data.date)
n_months = length(unique(pest_data.date))
n_buildings = length(unique(pest_data.building_id))
pest_data.ids = repeat(1:n_months, n_buildings)

using MLDataUtils

building_data = pest_data[:, [:building_idx, :live_in_super, :age_of_building, :total_sq_foot,
        :average_tenant_age, :monthly_average_rent]]
building_data = select(sort(unique(building_data), :building_idx), Not(:building_idx))
μ = center!(building_data)

building_data.age_of_building = building_data.age_of_building ./ 10
building_data.total_sq_foot = building_data.total_sq_foot ./ 10000
building_data.average_tenant_age = building_data.average_tenant_age ./ 10
building_data.monthly_average_rent = building_data.monthly_average_rent ./ 1000;

building_data_julia = Matrix(building_data)

building_data = building_data_r


using SingleFloats
using Test

include("frombase.jl")

@test Float64(sqrt(Single32(2))) == sqrt(2.0)
@test Float32(sqrt(Single32(2.0))) == Float32(sqrt(Float64(2.0)))

#@test sum(cot.(Single32.(collect(1.0:20)))) == sum(cot.(reverse(Single32.(collect(1.0:20)))))


using Base.Test
using Laplacians

# generate chimeric graphs, and test run many routines on them

function isTree(gr::SparseMatrixCSC)
    isConnected(gr) && (nnz(gr) == 2*(gr.n-1))
end

function testSolvers(a)

    n = a.n
    excess = zeros(n); excess[1] = excess[n] = 1e-5;
    la = lap(a);
    sdd = la + spdiagm(excess);
    b = rand(n); b = b - mean(b);

    for solver in SDDSolvers
        f = solver(sdd, tol=1e-6, maxtime=0.001);
        x = f(b);
    end

    for solver in LapSolvers
        f = solver(a, tol=1e-6, maxtime=0.001);
        x = f(b);
    end

end

n = 1000
tol = 1e-11

for i in 1:100
    println("wtedChimera($n, $i)")
    gr = wtedChimera(n,i)
    @test gr.n == n
    @test isConnected(gr)
    @test isTree(kruskal(gr))
    @test isTree(randishKruskal(gr))
    @test isTree(randishPrim(gr))
    @test abs(sum(kruskal(gr)) - sum(prim(gr))) < tol
    @test sum(kruskal(gr,kind=:min)) <= sum(kruskal(gr,kind=:max))

    testSolvers(gr)
end

# every single function/method should have at least one test!

# assume each vertex pair has at most ONE unique edge
function isSubsetOfEdges(subset::SparseMatrixCSC, superset::SparseMatrixCSC, marginOfError)
  nSubsetVertices = subset.n

  for vInd in 1:nSubsetVertices
    for eSubInd in subset.colptr[vInd]:(subset.colptr[vInd+1]-1)
      for eSupInd in superset.colptr[vInd]:(superset.colptr[vInd+1]-1)
        if superset.rowval[eSupInd] == subset.rowval[eSubInd]
          if abs(superset.nzval[eSupInd] - subset.nzval[eSubInd]) > marginOfError
            println("edge: v1: ", vInd, " v2: ", subset.rowval[eSubInd], " edgeWeight: ", subset.nzval[eSubInd])
            println("\t=/= edge: v1: ", vInd, " v2: ", superset.rowval[eSupInd], " edgeWeight: ", superset.nzval[eSupInd])
            return false
          end

          break
        end
      end
    end
  end #for

  return true

end #isSubsetOfEdges

# Test Function

# uses four different types of graphs:
# dim will increase by "iterationsize" each time for each number of iterations
"""
testAKPW runs tests on akpw to make sure it runs correctly. It takes a few parameters:

startingDim is the dimension of the first set of graphs it will test (the number of nodes is dim squared)
iterationsize is the amount that dim will grow each iteration
numIterations is the number of iterations
kindVar, randomClustersVar, metisClusteringVar, exponentialX, and shuffleClustersVar each pertain to the equivalent options for akpw
marginOfError is the allowable margin of error for running edgeSubset tests (if the difference in edge weights between the tree and graph is above this margin, it will not pass the test)

In each iteration, 4 graphs (one grid, one product, one necklace and one chimera) each with
(startingDim+((numIterations-1)*iterationsize))^2 nodes are produced, and akpw is run on each. The resulting tree
is tested using two functions: isTree and isSubsetOfEdges. isTree tests to see if it is, in fact, a tree.
isSubsetOfEdges makes sure all of the edges in the tree are also edges in the original graph (this works with
  akpw! as well, if you want to test that function instead, by using the marginOfError option).

for example:

  testAKPW(4, 2, 3) should 3 sets of four (grid, product, necklace and chimera) graphs: with 16, 36, and 64 nodes.
  None should produce any errors
"""
function testAKPW(startingDim, iterationSize, numIterations; kindVar=:max, randomClustersVar=false, metisClusteringVar=false, shuffleClustersVar=true, exponentialX=true, marginOfError=.0000001)

  if startingDim < 4
    println("please choose a higher startingDim (min: 4)")
    return
  end

  dim = startingDim

  for i in 1:numIterations

    a = grid2(dim)
    n = size(a)[1]
    (ai,aj,av) = findnz(triu(a))
    gridGraph = sparse(ai,aj,rand(size(av)),n,n)
    gridGraph = gridGraph + gridGraph';

    a = productGraph(generalizedRing(dim,[1 5]), grownGraphD(dim,3));
    n = size(a)[1]
    (ai,aj,av) = findnz(triu(a))
    productGraph1 = sparse(ai,aj,rand(size(av)),n,n)
    productGraph1 = productGraph1 + productGraph1'

    a = generalizedNecklace(generalizedRing(dim,[1 5]), grownGraphD(dim,3),5);
    n = size(a)[1]
    (ai,aj,av) = findnz(triu(a))
    necklaceGraph = sparse(ai,aj,rand(size(av)),n,n)
    necklaceGraph = necklaceGraph + necklaceGraph'

    chimGraph = wtedChimera(dim^2)

    rows, columns, edgeWeights = findnz(gridGraph)
    akpwTreeGrid = akpw(gridGraph, kind=kindVar, metisClustering = metisClusteringVar, randomClusters = randomClustersVar, shuffleClusters = shuffleClustersVar, exponentialX = exponentialX)
    oldGrid = sparse(rows, columns, edgeWeights)

    rows, columns, edgeWeights = findnz(productGraph1)
    akpwTreeProduct = akpw(productGraph1, kind=kindVar, metisClustering = metisClusteringVar, randomClusters = randomClustersVar, shuffleClusters = shuffleClustersVar, exponentialX = exponentialX)
    oldProductGraph = sparse(rows, columns, edgeWeights)

    rows, columns, edgeWeights = findnz(necklaceGraph)
    akpwTreeNecklace = akpw(necklaceGraph, kind=kindVar, metisClustering = metisClusteringVar, randomClusters = randomClustersVar, shuffleClusters = shuffleClustersVar, exponentialX = exponentialX)
    oldNecklace = sparse(rows, columns, edgeWeights)

    rows, columns, edgeWeights = findnz(chimGraph)
    akpwTreeChim = akpw(chimGraph, kind=kindVar, metisClustering = metisClusteringVar, randomClusters = randomClustersVar, shuffleClusters = shuffleClustersVar, exponentialX = exponentialX)
    oldChimera = sparse(rows, columns, edgeWeights)

    if !isTree(akpwTreeGrid)
      write(STDERR, "Grid is not a tree\n")
    end

    if !isSubsetOfEdges(akpwTreeGrid, oldGrid, marginOfError)
      write(STDERR, "akpwTreeGrid is not a subset of the edges in gridGraph\n")
    end

    if !isTree(akpwTreeProduct)
      write(STDERR, "productGraph is not a tree\n")
    end

    if !isSubsetOfEdges(akpwTreeProduct, oldProductGraph, marginOfError)
      write(STDERR, "akpwTreeProduct is not a subset of the edges in productGraph\n")
    end

    if !isTree(akpwTreeNecklace)
      write(STDERR, "necklace is not a tree\n")
    end

    if !isSubsetOfEdges(akpwTreeNecklace, oldNecklace, marginOfError)
      write(STDERR, "akpwTreeNecklace is not a subset of the edges in necklaceGraph\n")
    end

    if !isTree(akpwTreeChim)
      write(STDERR, "chimera is not a tree\n")
    end

    if !isSubsetOfEdges(akpwTreeChim, oldChimera, marginOfError)
      write(STDERR, "akpwTreeChim is not a subset of the edges in chimGraph\n")
    end

    dim += iterationSize

  end#for

  println("Tests PASS!")

end #testAKPW

include("testPCG.jl")


@testset "group" begin
    @test group(identity, 11:20) == Dict(Pair.(11:20, (x->[x]).(11:20)))
    @test group(iseven, 1:10) == Dict(true => [2,4,6,8,10], false => [1,3,5,7,9])

    @test group(iseven, x -> x*2, 1:10) == Dict(true => [4,8,12,16,20], false => [2,6,10,14,18])

    @test group((x,y) -> iseven(x+y), (x,y) -> x, 1:10, [1,3,4,2,5,6,4,2,3,9]) == Dict(true => [1,4,5,6,8,9], false => [2,3,7,10])
end

@testset "groupinds" begin
    @test groupinds(identity, 11:20) == Dict(Pair.(11:20, (x->[x]).(1:10)))
    @test groupinds(iseven, 11:20) == Dict(true => [2,4,6,8,10], false => [1,3,5,7,9])
end

@testset "groupview" begin
    @test groupview(identity, 11:20)::Groups == group(identity, 11:20)::Dict
    @test groupview(iseven, 11:20)::Groups == group(iseven, 11:20)::Dict
end

@testset "groupreduce" begin
    @test groupreduce(identity, +, 1:10) == Dict(Pair.(1:10, 1:10))
    @test groupreduce(iseven, +, 1:10) == Dict(true => 30, false => 25)

    @test groupreduce(iseven, x -> x*2, +, 1:10) == Dict(true => 60, false => 50)

    @test groupreduce(iseven, x -> x*2, +, 1:10; init=10) == Dict(true => 70, false => 60)

    @test groupreduce((x,y) -> iseven(x+y), (x,y) -> x+y, +, 1:10, 1:10; init=10) == Dict(true => 120)
    @test groupreduce((x,y) -> iseven(x+y), (x,y) -> x+y, +, 1:10, [1,3,4,2,5,6,4,2,3,9]; init=10) == Dict(true => 62, false => 52)

    @test groupcount(iseven, 1:10) == Dict(true => 5, false => 5)
    @test groupsum(iseven, 1:10) == Dict(true => 30, false => 25)
    @test groupprod(iseven, 1:10) == Dict(true => 2*4*6*8*10, false => 1*3*5*7*9)
end


using TransposedVectors
using Base.Test

include("TransposedVector.jl")
include("hcat.jl")
include("mul.jl")
include("ldiv.jl")
include("rdiv.jl")
include("diagonal.jl")
include("triangular.jl")
include("sparse.jl")
include("qr.jl")
include("norm.jl")


"""
m: minibatch size
N: full data size
"""
function subsample_index(m::Integer, N::Integer)
  @assert m < N
  shuffled_idx = Random.shuffle(1:N)
  return shuffled_idx[1:m], shuffled_idx[(m+1):end]
end

function sample_minibatch(batchsizes::Vector{S},
                          y::Vector{Matrix{T}}) where {T <: AbstractFloat,
                                                       S <: Integer}
  I = length(y)
  @assert length(batchsizes) == I
  N = size.(y, 1)
  J = size(y[1], 2)

  idx_mini, idx_mega = zip([subsample_index(batchsizes[i], N[i]) for i in 1:I]...)
  ymini = [y[i][idx_mini[i], :] for i in 1:I]
  ymega = [y[i][idx_mega[i], :] for i in 1:I]

  return idx_mini, idx_mega, Data(ymini), Data(ymega)
end


using Test, ArgParse
using SignalingDimension

try print_test_results catch err
    include("../utilities.jl")
    using .ScriptUtilities: print_test_results
end

# parsing command line arguments
arg_table = ArgParseSettings()
@add_arg_table! arg_table begin
    "-N"
        help = "max number of inputs/outputs to consider"
        arg_type = Int64
        default = 45
    "--dir"
        help = "directory to print results .txt to"
        arg_type = String
        default = "./script/facet_verifications/results/"
    "--stdout"
        help = "Print output to STDOUT. If not flagged, output is printed to `.txt` file."
        action = :store_true
end
args = parse_args(ARGS, arg_table)

# verification script: verifies all maximum likelihood facets with N ∈ [3, N_max]
function verify_maximum_likelihood_facet(N_max :: Int64)
    @testset "Testing all maximum likelihood facets of size `($N_max, $N_max)` or less." begin
        @testset "scanning over all `d ∈ [2, $N - 1]`" for N in 3:N_max
            @testset "Verifying N = $N and d = $d" for d in 2:N-1
                vertices = aff_ind_maximum_likelihood_vertices(N,d)
                facet = maximum_likelihood_facet(N,d)

                @test verify_facet(d, facet, vertices)
            end
        end
    end
end

# outputting result to file or stdout
print_test_results(
    verify_maximum_likelihood_facet,
    params=[args["N"]],
    stdout=args["stdout"],
    dir=args["dir"]
)


import EnergyExpressions: showcoeff

"""
    LinearCombination(Ts::Vector{T}, coeffs::Vector{N})

Represents a general linear combination of objects of type `T`.
"""
struct LinearCombination{T,N<:Number}
    Ts::Vector{T}
    coeffs::Vector{N}
    function LinearCombination(Ts::Vector{T}, coeffs::Vector{N}) where {T,N}
        sel = .!iszero.(coeffs)
        new{T,N}(Ts[sel], coeffs[sel])
    end
end

Base.length(lc::LinearCombination) = length(lc.Ts)
Base.eltype(lc::LinearCombination{T,N}) where {T,N} = (T,N)

Base.:(==)(a::LinearCombination, b::LinearCombination) =
    a.Ts == b.Ts && a.coeffs == b.coeffs

Base.hash(lc::LinearCombination, h::UInt) =
    hash(lc.Ts, hash(lc.coeffs, h))

Base.zero(::LinearCombination{T,N}) where {T,N} =
    LinearCombination(Vector{T}(), Vector{N}())

Base.getindex(lc::LinearCombination, i) =
    (lc.Ts[i],lc.coeffs[i])

Base.iterate(lc::LinearCombination, i=1) =
    i > length(lc) ? nothing : (lc[i], i+=1)

function Base.show(io::IO, lc::LinearCombination)
    n = length(lc)
    for (i,(t,c)) in enumerate(lc)
        showcoeff(io, c, i > 1)
        write(io, " ")
        show(io, t)
        i < n && write(io, " ")
    end
end

Base.:(+)(A::LinearCombination, B::LinearCombination) =
    LinearCombination(vcat(A.Ts, B.Ts), vcat(A.coeffs, B.coeffs))

Base.:(-)(A::LinearCombination, B::LinearCombination) =
    LinearCombination(vcat(A.Ts, B.Ts), vcat(A.coeffs, -1 * B.coeffs))

Base.:(+)(A::LC, B::T) where {T,N,LC<:LinearCombination{<:T,N}} =
    LinearCombination(vcat(A.Ts, B), vcat(A.coeffs, one(N)))

Base.:(+)(A::T, B::LC) where {T,N,LC<:LinearCombination{<:T,N}} =
    LinearCombination(vcat(A, B.Ts), vcat(one(N), B.coeffs))

Base.:(-)(A::LC) where {LC<:LinearCombination} =
    LinearCombination(A.Ts, -1 * A.coeffs)

Base.:(*)(A::LinearCombination, b::Number) =
    LinearCombination(A.Ts, b * A.coeffs)

Base.:(*)(a::Number, B::LinearCombination) =
    LinearCombination(B.Ts, a * B.coeffs)

Base.:(/)(A::LinearCombination, b::Number) =
    LinearCombination(A.Ts, A.coeffs/b)

"""
    @linearly_combinable TT

Turns the type `TT` into a linearly combinable type, i.e. defines
arithmetic operators.

# Examples

```julia
julia> @linearly_combinable Symbol

julia> 4*(:x) - 5*(:y)
4 :x - 5 :y
```
"""
macro linearly_combinable(TT)
    quote
        Base.:(+)(A::T, B::U) where {T<:$(esc(TT)),U<:$(esc(TT))} =
            LinearCombination(promote_type(T,U)[A,B], [1,1])

        Base.:(-)(A::$(esc(TT))) =
            LinearCombination($(esc(TT))[A], [-1])

        Base.:(*)(a::Number, B::$(esc(TT))) =
            LinearCombination([B], [a])

        Base.:(*)(A::$(esc(TT)), b::Number) =
            LinearCombination([A], [b])
    end
end


@testset "allapproxequal()" begin
    @test nsi.allapproxequal([])                     == true
    @test nsi.allapproxequal([1.0])                  == true
    @test nsi.allapproxequal([1.00000001, 1.0, 1.0]) == true
    # tolerance is 1e-8          1234567
    @test nsi.allapproxequal([ 1.0000001, 1.0, 1.0]) == false
    @test nsi.allapproxequal([ 1, 1, 1])             == true
    @test nsi.allapproxequal([ 2, 1, 1])             == false
    @test_throws MethodError nsi.allapproxequal([ "A", "A"])
end

@testset "verify_same_length()" begin
    @test nsi.verify_same_length([], [])                 === nothing
    @test nsi.verify_same_length(["A"], ["B"])           === nothing
    @test nsi.verify_same_length(["A", "B"], ["C", "D"]) === nothing
    @test nsi.verify_same_length([1], [2])               === nothing
    @test nsi.verify_same_length([1, 2], [3, 4])         === nothing
    @test nsi.verify_same_length([1 2; 3 4], [5 6; 7 8]) === nothing
    @test_throws ArgumentError nsi.verify_same_length([1], [])
    @test_throws ArgumentError nsi.verify_same_length(["A"], [])
    @test_throws ArgumentError nsi.verify_same_length(["A", "B"], ["C"])
    @test_throws ArgumentError nsi.verify_same_length([1], [])
    @test_throws ArgumentError nsi.verify_same_length([1, 2], [3])
    @test_throws ArgumentError nsi.verify_same_length([1 2; 3 4], [5 6;])
end

@testset "left_right_from_peak()" begin
    x = collect(Float64, 1:20)
    y = zeros(Float64, 20)
    y[10] = 1
    @test nsi.left_right_from_peak(x, y, 10, 1) == [9.5, 10.5]
    @test nsi.left_right_from_peak(x, y, 10, 0.7) == [9.65, 10.35]
    y[10] = 0
    y[1]  = 1
    @test nsi.left_right_from_peak(x, y, 1, 1) == [0.5, 1.5]
    y[1] = 0
    y[20]  = 1
    @test nsi.left_right_from_peak(x, y, 20, 1) == [19.5, 20.5]
end

@testset "trapz()" begin
    # Test integration function trapz using analytic integrals;
    # integral functions (F) of functions f from
    # https://en.wikipedia.org/wiki/Lists_of_integrals

    bounds = [(-1, 4), (1, 4), (-1.4, 7.9), (11.6, 17.4)]

    test_data = [
        (f = x->  exp(3x), F = x->        1/3*exp(3x)),
        (f = x->  1.2^(x), F = x->   1.2^(x)/log(1.2)),
        (f = x->   sin(x), F = x->            -cos(x)),
        (f = x-> 1/(2x+3), F = x-> 1/2*log(abs(2x+3)))
    ]
    x = collect(-20:0.0001:20)
    for t in test_data
        for (l, r) in bounds
            f, F = t
            ref_1 = F(r) - F(l)
            ref_2 = ref_1 - 1/2*(f(r) + f(l))*(r - l) # 1/2*(f(r) + f(l))*(r - l) term: subtracting baseline
            @test nsi.trapz(x, f.(x), l, r) ≈ ref_1 rtol=1e-7
        end
    end

    begin
        x = collect(Float64, 1:10)
        y = collect(Float64, 1:10)
        @test nsi.trapz(x, y, 1, 3) == 4.0
        @test nsi.trapz(x, y, 1.1, 2.1) ≈ 1.6
        @test nsi.trapz(x, y, 1.1, 1.2) ≈ 1.1*0.1 + 0.1*0.1*0.5
        @test nsi.trapz(x, y, 1.0, 1.1) ≈ 0.1 + 0.1*0.1*0.5
        @test nsi.trapz(x, y, 3.9, 4.0) ≈ 3.9*0.1 + 0.1*0.1*0.5
    end

    begin
        x = [1, 2, 3, 4, 6, 7.0]
        y = [1, 2, 4, 2.5, 3, 2]
        @test nsi.trapz(x, y, 1, 7) == 15.75
        @test nsi.trapz(x, y, 0, 10) == 15.75
        @test nsi.trapz(x, y, 1, 2) == 1.5
        @test nsi.trapz(x, y, 3, 4) == 3.25
        @test nsi.trapz(x, y, 1, 6) == 13.25
        @test nsi.trapz(x, y, 6, 1) == 13.25
        @test nsi.trapz(x, y, 1.5, 7) == 15.125
    end
end

@testset "@samples" begin
    seed!(1)
    @testset "Normal" begin
        ref = Particles(10_000, Normal(1, 1))
        out = @samples 10_000 1 ± 1
        @test mean(ref) ≈ mean(out)
        @test std(ref) ≈ std(out)
        # pass expression to macro
        out = @samples 10_000 (2 / 2) ± (3 - 2)
        @test mean(ref) ≈ mean(out)
        @test std(ref) ≈ std(out)
        # pass variables to macro
        out = let
            a = 1
            b = 1
            @samples 10_000 a ± b
        end
        @test mean(ref) ≈ mean(out)
        @test std(ref) ≈ std(out)
    end
    @testset "Uniform" begin
        ref = Particles(10_000, Uniform(1, 2))
        out = @samples 10_000 1 .. 2
        @test mean(ref) ≈ mean(out)
        @test std(ref) ≈ std(out)
        # pass expression to macro
        out = @samples 10_000 (2 / 2) .. (3 - 1)
        @test mean(ref) ≈ mean(out)
        @test std(ref) ≈ std(out)
        # pass variables to macro
        out = let
            a = 1
            b = 2
            @samples 10_000 a .. b
        end
        @test mean(ref) ≈ mean(out)
        @test std(ref) ≈ std(out)
    end
    @test_throws ArgumentError (@samples 1 1 + 1)
    @test_throws ArgumentError (@samples 1 1 + 1 + 1)
end

@testset "lininterp" begin
    xs = 0:0.5:10 |> collect
    ys = 0:0.5:10 |> collect
    @test NoisySignalIntegration.lininterp(2.7, xs, ys) == 2.7
    ys = ys.^2
    @test NoisySignalIntegration.lininterp(2.7, xs, ys) ≈ 7.35
    @test NoisySignalIntegration.lininterp(2.5, xs, ys) == 6.25
    @test NoisySignalIntegration.lininterp(10.0, xs, ys) == 100.0
    @test_throws ErrorException NoisySignalIntegration.lininterp(0.0, xs, ys)
    @test_throws ErrorException NoisySignalIntegration.lininterp(10.1, xs, ys)
end


struct EPMat{T<:AbstractFloat} <: AbstractEPMat
    αStS::SparseMatrixCSC{T,Int64} # A container of KK = βS'S (Its diag will change)
    αStS_diag::AbstractArray{T} # Store the original diag of KK = βS'S
    αStb::AbstractArray{T}
    Σ::Matrix{T} # ΣQ Enforced to be dense
    v::AbstractArray{T}
    lb::AbstractArray{T}
    ub::AbstractArray{T}
end

function EPMat(S::AbstractArray{T}, b::AbstractArray{T},
        lb::AbstractArray{T}, ub::AbstractArray{T}, alpha::T) where T <: Real
    isinf(alpha) && error("For Inf alpha use 'EPMatT0'")
    M, N = size(S)
    αStS = sparse(alpha * S' * S)
    return EPMat(
        #= αStS      =# αStS,
        #= αStS_diag =# diag(αStS),
        #= αStb      =# alpha * S' * b,
        #= Σ         =# zeros(T,N,N),
        #= v         =# zeros(T,N),
        #= lb        =# lb,
        #= ub        =# ub
    )
end


using Revise
using ADCME
using ADCMEKit
using NNFEM
using PyPlot
using ProgressMeter
using AdFem
using SparseArrays
#----------------------------------------------------------------------------
# step 1: create a computational domain
# fixed on the bottom and push from right

m = 40
n = 20
h = 1/n
NT = 50
Δt = 0.01/NT

bdedge = bcedge("right", m, n, h)
bdnode = bcnode("lower", m, n, h)

λ = constant(8.641975308641973e9)
μ = constant(3.7037037037037034e9)
invη = 1/1e5*constant(ones(4*m*n))

M = 2000*compute_fem_mass_matrix1(m, n, h)
Zero = spzeros((m+1)*(n+1), (m+1)*(n+1))
M = SparseTensor([M Zero;Zero M])

## alpha-scheme
β = 1/4; γ = 1/2

# invη is a 4*m*n array
function make_matrix(invη)
  fn_G = invη->begin
    G = tensor([1/Δt+2/3*μ*invη -μ/3*invη 0.0
      -μ/3*invη 1/Δt+2/3*μ*invη 0.0
      0.0 0.0 1/Δt+μ*invη])
    invG = inv(G)
  end
  invG = map(fn_G, invη)
  S = tensor([2μ/Δt+λ/Δt λ/Δt 0.0
      λ/Δt 2μ/Δt+λ/Δt 0.0
      0.0 0.0 μ/Δt])

  H = invG*S
  K = compute_fem_stiffness_matrix(H, m, n, h)
  L = M + β*Δt^2*K
  # L = pl[1]*L
  L, Lbd = fem_impose_Dirichlet_boundary_condition_experimental(L, bdnode, m, n, h)

  return K, L, S, invG
end

function simulate()

  a = TensorArray(NT+1); a = write(a, 1, zeros(2(m+1)*(n+1))|>constant)
  v = TensorArray(NT+1); v = write(v, 1, zeros(2(m+1)*(n+1))|>constant)
  d = TensorArray(NT+1); d = write(d, 1, zeros(2(m+1)*(n+1))|>constant)
  U = TensorArray(NT+1); U = write(U, 1, zeros(2(m+1)*(n+1))|>constant)
  Sigma = TensorArray(NT+1); Sigma = write(Sigma, 1, zeros(4*m*n, 3)|>constant)
  Varepsilon = TensorArray(NT+1); Varepsilon = write(Varepsilon, 1,zeros(4*m*n, 3)|>constant)

  Forces = zeros(NT, 2(m+1)*(n+1))
  for i = 1:NT
    T = eval_f_on_boundary_edge((x,y)->1e6, bdedge, m, n, h)

    T = [-T zeros(size(T)...)]
    rhs = compute_fem_traction_term(T, bdedge, m, n, h)

    Forces[i, :] = rhs
  end
  Forces = constant(Forces)

  K, L, S, invG = make_matrix(invη)

  L = factorize(L)

  function condition(i, tas...)
    i <= NT
  end

  function body(i, tas...)
    a_, v_, d_, U_, Sigma_, Varepsilon_ = tas
    a = read(a_, i)
    v = read(v_, i)
    d = read(d_, i)
    U = read(U_, i)
    Sigma = read(Sigma_, i)
    Varepsilon = read(Varepsilon_, i)

    res = batch_matmul(invG/Δt, Sigma)
    F = compute_strain_energy_term(res, m, n, h) - K * U
    rhs = Forces[i] -  F

    td = d + Δt*v + Δt^2/2*(1-2β)*a
    tv = v + (1-γ)*Δt*a
    rhs = rhs - K*td
    rhs = scatter_update(rhs, constant([bdnode; bdnode.+(m+1)*(n+1)]), constant(zeros(2*length(bdnode))))

    ## alpha-scheme
    a = L\rhs # bottleneck
    d = td + β*Δt^2*a
    v = tv + γ*Δt*a
    U_new = d

    Varepsilon_new = eval_strain_on_gauss_pts(U_new, m, n, h)

    res2 = batch_matmul(invG * S, Varepsilon_new-Varepsilon)
    Sigma_new = res +  res2

    i+1, write(a_, i+1, a), write(v_, i+1, v), write(d_, i+1, d), write(U_, i+1, U_new),
          write(Sigma_, i+1, Sigma_new), write(Varepsilon_, i+1, Varepsilon_new)
  end

  i = constant(1, dtype=Int32)
  _, _, _, _, u, sigma, varepsilon = while_loop(condition, body,
                    [i, a, v, d, U, Sigma, Varepsilon])

  U = stack(u)
  Sigma = stack(sigma)
  Varepsilon = stack(varepsilon)

  U = set_shape(U, (NT+1, size(U,2)))
  Sigma = set_shape(Sigma, (NT+1, 4*m*n, 3))

  return U, Sigma
end

U, sigma = simulate()
sess = Session(); init(sess)
U, sigma = run(sess, [U, sigma])

##

elements = []
# prop = Dict("name"=> "PlaneStrain", "rho"=> 2000, "E"=> 1e10, "nu"=> 0.35)
prop = Dict("name"=> "ViscoelasticityMaxwell", "rho"=> 2000, "E"=> 1e10, "nu"=> 0.35, "eta"=>1e7)

coords = zeros((m+1)*(n+1), 2)
for j = 1:n
    for i = 1:m
        idx = (m+1)*(j-1)+i
        elnodes = [idx; idx+1; idx+1+m+1; idx+m+1]
        ngp = 2
        nodes = [
        (i-1)*h (j-1)*h
        i*h (j-1)*h
        i*h j*h
        (i-1)*h j*h
        ]
        coords[elnodes, :] = nodes
        push!(elements, SmallStrainContinuum(nodes, elnodes, prop, ngp))
    end
end

# fixed on the bottom, push on the right
EBC = zeros(Int64, (m+1)*(n+1), 2)
FBC = zeros(Int64, (m+1)*(n+1), 2)
g = zeros((m+1)*(n+1), 2)
f = zeros((m+1)*(n+1), 2)

for i = 1:m+1
    EBC[i+n*(m+1),:] .= -1
end

Edge_Traction_Data = Array{Int64}[]
for i = 1:n
  elem = elements[(i-1)*m + m]
  for k = 1:4
    if elem.coords[k,1]>2-0.001 && elem.coords[k+1>4 ? 1 : k+1,1]>2-0.001
      push!(Edge_Traction_Data, [(i-1)*m + m, k, 1])
    end
  end
end
Edge_Traction_Data = hcat(Edge_Traction_Data...)'|>Array

ndims = 2
domain = Domain(coords, elements, ndims, EBC, g, FBC, f, Edge_Traction_Data)
######

domain.history["state"] = []
domain.history["stress"] = []
for i = 1:NT+1
  push!(domain.history["state"], U[i,:])
  push!(domain.history["stress"], sigma[i,:,:])
end
# visualize_von_mises_stress_on_scoped_body(d_, domain, scale_factor=10.0)
p = visualize_total_deformation_on_scoped_body(U, domain, scale_factor=10.0)
saveanim(p, "adfem.gif")
close("all")


using Latte, ArrayViews

@neuron type MLPNeuron
    weights  :: DenseArray{Float32}
    ∇weights :: DenseArray{Float32}

    bias     :: DenseArray{Float32}
    ∇bias    :: DenseArray{Float32}
end

@neuron forward(neuron::MLPNeuron) do
    # dot product of weights and inputs
    for i in 1:length(neuron.inputs[1])
        neuron.value += neuron.weights[i] * neuron.inputs[1][i]
    end
    # shift by the bias value
    neuron.value += neuron.bias[1]
end

@neuron backward(neuron::MLPNeuron) do
    for i in 1:length(neuron.inputs[1])
        neuron.∇inputs[1][i] += neuron.weights[i] * neuron.∇
    end
    for i in 1:length(neuron.inputs[1])
        neuron.∇weights[i] += neuron.inputs[1][i] * neuron.∇
    end
    neuron.∇bias[1] += neuron.∇
end

function FCLayer(name::Symbol, net::Net, input_ensemble::AbstractEnsemble, num_neurons::Int)
    neurons = Array(MLPNeuron, num_neurons)
    num_inputs = length(input_ensemble)
    weights = xavier(Float32, num_inputs, num_neurons)
    ∇weights = zeros(Float32, num_inputs, num_neurons)

    bias = zeros(Float32, 1, num_neurons)
    ∇bias = zeros(Float32, 1, num_neurons)
    for i in 1:num_neurons
        neurons[i] = MLPNeuron(view(weights, :, i), view(∇weights, :, i),
                               view(bias, :, i), view(∇bias, :, i))
    end
    ens = Ensemble(net, name, neurons, [Param(name,:weights, 1.0f0, 1.0f0),
                                        Param(name,:bias, 2.0f0, 0.0f0)])
    add_connections(net, input_ensemble, ens,
                    (i) -> (tuple([Colon() for d in size(input_ensemble)]... )))
    return ens
end

net = Net(100)
data, label = HDF5DataLayer(net, "data/train.txt", "data/test.txt")

fc1 = FCLayer(:fc1, net, data, 100)
fc2 = FCLayer(:fc2, net, fc1, 10)

loss     = SoftmaxLossLayer(:loss, net, fc2, label)
accuracy = AccuracyLayer(:accuracy, net, fc2, label)

params = SolverParameters(
    lr_policy    = LRPolicy.Inv(0.01, 0.0001, 0.75),
    mom_policy   = MomPolicy.Fixed(0.9),
    max_epoch    = 50,
    regu_coef    = .0005,
    snapshot_dir = "mlp")
sgd = SGD(params)
solve(sgd, net)


function feedforward!(net::Network,input::Vector{Float64})::Vector{Float64}
    net.layers[1] = input

    # a^(l+1) = σ(w^l*a^l + b^1)
    for i=1:(length(net.layers)-1)
        net.z[i] = net.weights[i]*net.neurons[i] + net.biases[i]
        net.neurons[i+1] = net.z[i].net.act()
    end

    return net.neurons[end]
end


"""
    MHAttention(nheads::Integer, qkv_layer, attn_drop, projection)

Multi-head self-attention layer.

# Arguments:

  - `nheads`: Number of heads
  - `qkv_layer`: layer to be used for getting the query, key and value
  - `attn_drop`: dropout rate after the self-attention layer
  - `projection`: projection layer to be used after self-attention
"""
struct MHAttention{P, Q, R}
    nheads::Int
    qkv_layer::P
    attn_drop::Q
    projection::R
end

"""
    MHAttention(planes::Integer, nheads::Integer = 8; qkv_bias::Bool = false, attn_drop = 0., proj_drop = 0.)

Multi-head self-attention layer.

# Arguments:

  - `planes`: number of input channels
  - `nheads`: number of heads
  - `qkv_bias`: whether to use bias in the layer to get the query, key and value
  - `attn_drop`: dropout rate after the self-attention layer
  - `proj_drop`: dropout rate after the projection layer
"""
function MHAttention(planes::Integer, nheads::Integer = 8; qkv_bias::Bool = false,
                     attn_drop = 0.0, proj_drop = 0.0)
    @assert planes % nheads==0 "planes should be divisible by nheads"
    qkv_layer = Dense(planes, planes * 3; bias = qkv_bias)
    attn_drop = Dropout(attn_drop)
    proj = Chain(Dense(planes, planes), Dropout(proj_drop))
    return MHAttention(nheads, qkv_layer, attn_drop, proj)
end

@functor MHAttention

function (m::MHAttention)(x::AbstractArray{T, 3}) where {T}
    nfeatures, seq_len, batch_size = size(x)
    x_reshaped = reshape(x, nfeatures, seq_len * batch_size)
    qkv = m.qkv_layer(x_reshaped)
    qkv_reshaped = reshape(qkv, nfeatures ÷ m.nheads, m.nheads, seq_len, 3 * batch_size)
    query, key, value = chunk(qkv_reshaped, 3; dims = 4)
    scale = convert(T, sqrt(size(query, 1) / m.nheads))
    key_reshaped = reshape(permutedims(key, (2, 1, 3, 4)), m.nheads, nfeatures ÷ m.nheads,
                           seq_len * batch_size)
    query_reshaped = reshape(permutedims(query, (1, 2, 3, 4)), nfeatures ÷ m.nheads, m.nheads, seq_len * batch_size)
    attention = m.attn_drop(softmax(batched_mul(query_reshaped, key_reshaped) .* scale))
    value_reshaped = reshape(permutedims(value, (1, 2, 3, 4)), nfeatures ÷ m.nheads, m.nheads, seq_len * batch_size)
    pre_projection = reshape(batched_mul(attention, value_reshaped),
                             (nfeatures, seq_len, batch_size))
    y = m.projection(reshape(pre_projection, size(pre_projection, 1), :))
    return reshape(y, :, seq_len, batch_size)
end


"Number of variables in the circuit"
num_variables(c::Union{ΔNode,Δ}) = length(variable_scope(c))

"Get the probability that a random world satisties the circuit"
function sat_prob(circuit::Union{ΔNode,Δ})::Rational{BigInt}
    sat_prob(circuit, v -> BigInt(1) // BigInt(2))
end

function sat_prob(circuit::Δ, varprob::Function)::Rational{BigInt}
    sat_prob(circuit[end], varprob)
end

function sat_prob(root::ΔNode, varprob::Function)::Rational{BigInt}
    f_con(n) = is_true(n) ? one(Rational{BigInt}) : zero(Rational{BigInt})
    f_lit(n) = positive(n) ? varprob(variable(n)) : one(Rational{BigInt}) - varprob(variable(n))
    f_a(n, callback) = mapreduce(callback, *, children(n))
    f_o(n, callback) = mapreduce(callback, +, children(n))
    foldup(root, f_con, f_lit, f_a, f_o, Rational{BigInt})
end

"Get the model count of the circuit"
function model_count(circuit::Δ, num_vars_in_scope::Int = num_variables(circuit))::BigInt
    # note that num_vars_in_scope can be more than num_variables(circuit)
    BigInt(sat_prob(circuit) * BigInt(2)^num_vars_in_scope)
end

const Signature = Vector{Rational{BigInt}}

"""
Get a signature for each node using probabilistic equivalence checking
"""
function prob_equiv_signature(circuit::Δ, k::Int)::Dict{Union{Var,ΔNode},Signature}
    prob_equiv_signature(circuit[end],k)
end

function prob_equiv_signature(circuit::ΔNode, k::Int)::Dict{Union{Var,ΔNode},Signature}
    # uses probability instead of integers to circumvent smoothing, no mod though
    signs::Dict{Union{Var,ΔNode},Signature} = Dict{Union{Var,ΔNode},Signature}()
    prime::Int = 7919 #TODO set as smallest prime larger than num_variables
    randprob() = BigInt(1) .// rand(1:prime,k)
    do_signs(v::Var) = get!(randprob, signs, v)
    f_con(n) = (signs[n] = (is_true(n) ? ones(Rational{BigInt}, k) : zeros(Rational{BigInt}, k)))
    f_lit(n) = (signs[n] = (positive(n) ? do_signs(variable(n)) : BigInt(1) .- do_signs(variable(n))))
    f_a(n, call) = (signs[n] = (mapreduce(c -> call(c), (x,y) -> (x .* y), children(n))))
    f_o(n, call) = (signs[n] = (mapreduce(c -> call(c), (x,y) -> (x .+ y), children(n))))
    foldup(circuit, f_con, f_lit, f_a, f_o, Signature)
    signs
end

"Get the variable scope of the circuit node"
function variable_scope(circuit::Δ)::BitSet
    variable_scope(circuit[end])
end

"Get the variable scope of the root of the circuit"
function variable_scope(root::ΔNode)::BitSet
    f_con(n) = BitSet()
    f_lit(n) = BitSet(variable(n))
    f_inner(n, call) = mapreduce(call, union, children(n))
    foldup(root, f_con, f_lit, f_inner, f_inner, BitSet)
end

"Get the variable scope of each node in the circuit"
function variable_scopes(circuit::Δ)::Dict{ΔNode,BitSet}
    variable_scopes(circuit[end])
end

function variable_scopes(root::ΔNode)::Dict{ΔNode,BitSet}
    # variable_scopes(node2dag(root))
    scope = Dict{ΔNode,BitSet}()
    f_con(n) = scope[n] = BitSet()
    f_lit(n) = scope[n] = BitSet(variable(n))
    f_inner(n, call) = scope[n] = mapreduce(call, union, children(n))
    foldup(root, f_con, f_lit, f_inner, f_inner, BitSet)
    scope
end

"Is the circuit smooth?"
function is_smooth(circuit::Δ)::Bool
    is_smooth(circuit[end])
end

function is_smooth(root::ΔNode)::Bool
    result::Bool = true
    f_con(n) = BitSet()
    f_lit(n) = BitSet(variable(n))
    f_a(n, cs) = reduce(union, cs)
    f_o(n, cs) = begin
        scope = reduce(union, cs)
        result = result && all(c -> c == scope, cs)
        scope
    end
    foldup_aggregate(root, f_con, f_lit, f_a, f_o, BitSet)
    result
end

"Is the circuit decomposable?"
function is_decomposable(circuit::Δ)::Bool
    is_decomposable(circuit[end])
end

function is_decomposable(root::ΔNode)::Bool
    result::Bool = true
    f_con(n) = BitSet()
    f_lit(n) = BitSet(variable(n))
    f_a(n, cs) = begin
        result = result && disjoint(cs...)
        reduce(union, cs)
    end
    f_o(n, cs) = reduce(union, cs)
    foldup_aggregate(root, f_con, f_lit, f_a, f_o, BitSet)
    result
end

"Construct a mapping from literals to their canonical node representation"
function literal_nodes(circuit::Union{Δ,ΔNode})::Dict{Lit,ΔNode}
    lit_dict = Dict{Lit,ΔNode}()
    foreach(circuit) do n
        if isliteralgate(n)
            if haskey(lit_dict, literal(n))
                error("Circuit has multiple representations of literal $(literal(n))")
            end
            lit_dict[literal(n)] = n
        end
    end
    lit_dict
end

"Construct a mapping from constants to their canonical node representation"
function constant_nodes(circuit::Δ)::Tuple{Union{Nothing, ΔNode},Union{Nothing, ΔNode}}
    true_node = nothing
    false_node = nothing
    visit(n::ΔNode) = visit(GateType(n),n)
    visit(::GateType, n::ΔNode) = ()
    visit(::ConstantGate, n::ΔNode) = begin
        if is_true(n)
            if issomething(true_node)
                error("Circuit has multiple representations of true")
            end
            true_node = n
        else
            @assert is_false(n)
            if issomething(false_node)
                error("Circuit has multiple representations of false")
            end
            false_node = n
        end
    end
    for node in circuit
        visit(node)
    end
    (false_node, true_node)
end

"Construct a mapping from constants to their canonical node representation"
function constant_nodes2(circuit::Union{Δ,ΔNode})::Tuple{Union{Nothing, ΔNode},Union{Nothing, ΔNode}}
    true_node = nothing
    false_node = nothing
    foreach(circuit) do n
        if isconstantgate(n)
            if is_true(n)
                isnothing(true_node) || error("Circuit has multiple representations of true")
                true_node = n
            else
                @assert is_false(n)
                isnothing(false_node) || error("Circuit has multiple representations of false")
                false_node = n
            end
        end
    end
    (false_node, true_node)
end

"Check whether literal nodes are unique"
function has_unique_literal_nodes(circuit::Δ)::Bool
    literals = Set{Lit}()
    result = true
    visit(n::ΔNode) = visit(GateType(n),n)
    visit(::GateType, n::ΔNode) = ()
    visit(::LiteralGate, n::ΔNode) = begin
        if literal(n) ∈ literals
            result = false
        end
        push!(literals, literal(n))
    end
    for node in circuit
        visit(node)
    end
    return result
end

"Check whether literal nodes are unique"
function has_unique_literal_nodes2(circuit::Δ)::Bool
    has_unique_literal_nodes2(circuit[end])
end

function has_unique_literal_nodes2(root::ΔNode)::Bool
    literals = Set{Lit}()
    @inline f_con(n) = true
    @inline f_lit(n) = begin
        lit = literal(n)
        if lit in literals
            false
        else
            push!(literals, lit)
            true
        end
    end
    @inline f_a(n, cs) = all(cs)
    @inline f_o(n, cs) = all(cs)
    foldup_aggregate(root, f_con, f_lit, f_a, f_o, Bool)
end

"Check whether constant nodes are unique"
function has_unique_constant_nodes(circuit::Δ)::Bool
    seen_false = false
    seen_true = false
    result = true
    visit(n::ΔNode) = visit(GateType(n),n)
    visit(::GateType, n::ΔNode) = ()
    visit(::ConstantGate, n::ΔNode) = begin
        if is_true(n)
            if seen_true
                result = false
            end
            seen_true = true
        else
            @assert is_false(n)
            if seen_false
                result = false
            end
            seen_false = true
        end
    end
    for node in circuit
        visit(node)
    end
    return result
end

function (circuit::Δ)(data::XData)
    circuit[end](data)
end

function (root::ΔNode)(data::XData)
    evaluate(root, data)
end

function evaluate(root::ΔNode, data::XData{Bool})::BitVector
    @inline f_lit(n) = if positive(n)
        [feature_matrix(data)[:,variable(n)]]
    else
        [broadcast(!,feature_matrix(data)[:,variable(n)])]
    end
    @inline f_con(n) =
        is_true(n) ? always(Bool, num_examples(data)) : never(Bool, num_examples(data))
    @inline fa(n, call) = begin
        if num_children(n) < 2
            return call(@inbounds children(n)[1])
        else
            c1 = call(@inbounds children(n)[1])
            c2 = call(@inbounds children(n)[2])
            if num_children(n) == 2 && length(c1) == 1 && length(c2) == 1
                return [c1[1], c2[1]] # no need to allocate a new BitVector, just return pair
            end
            x = always(Bool, num_examples(data))
            accumulate_elements(x, c1, &)
            accumulate_elements(x, c2, &)
            for c in children(n)[3:end]
                accumulate_elements(x, call(c), &)
            end
            return [x]
        end
    end
    @inline fo(n, call) = begin
        x = never(Bool, num_examples(data))
        for c in children(n)
            accumulate_elements(x, call(c), |)
        end
        return [x]
    end
    conjoin_elements(foldup(root, f_con, f_lit, fa, fo, Vector{BitVector}))
end

@inline function conjoin_elements(elems::Vector{BitVector})::BitVector
    reduce((x,y) -> x .& y, elems)
end

@inline function disjoin_elements(elems::Vector{BitVector})::BitVector
    reduce((x, y) -> x .| y, elems)
end

@inline function accumulate_elements(x::BitVector, elems::Vector{BitVector}, op)::BitVector
    if length(elems) == 1
        @inbounds @. x = op(x, elems[1])
    else
        @assert length(elems) == 2
        @inbounds @. x = op(x, elems[1] & elems[2])
    end
end

function pass_down2(circuit::Δ, data::XData{Bool})
    num = num_examples(data)
    @inline f_root(n, f_s) = begin
        [always(Bool, num)]
    end
    @inline f_inner(n, f_s, p_s) = begin
        x = never(Bool, num)
        if length(f_s) == 2
            for y in p_s
                @inbounds @. x = x | ( y[1] & f_s[1] & f_s[2])
            end
        else
            @assert length(f_s) == 1
            for y in p_s
                @inbounds @. x = x | ( y[1] & f_s[1])
            end
        end
        [x]
    end

    f_leaf = f_inner

    folddown_aggregate(circuit, f_root, f_leaf, f_inner, Vector{BitVector})
end

function pass_up_down2(circuit::Δ, data::XData{Bool})
    circuit[end](data)
    pass_down2(circuit, data)
end


"""
Compute the dark matter density profile at the point (x, vr, vt).  Where x is the radius divided by the scale radius (r / rs), vr is the radial velocity, and vt is the tangential velocity.
"""
function profile_density(x::G, vr::G, vt::G, p::NFWParameters{G}; N::G = 1000.0, αt::G = 1.0) where G <: AbstractFloat
    if x < 0.0
        return 0.0
    end

    #Compute Energy
    E = 0.5 * (vt * vt + vr * vr) + gravitational_potential(x, p)
    if E < p.Φlim && E >= 0.0
        hE = N * E^p.a * (E^p.q + p.adjEc^p.q)^(p.d/p.q) * (p.Φlim - E)^p.e
    else
        return 0.0
    end

    #Compute angular momentum
    J = abs(x * p.rs * vt) #vt = v * sin(theta)
    if p.b <= 0.0
        #gJ = 1.0 / (1.0 + (J / p.adjJb)^(-p.b))
        gJ = (1.0 + (J / p.adjJb)^(-p.b / αt))^-αt
    else
        #gJ = 1.0 + (J / p.adjJb)^p.b
        gJ = (1.0 + (J / p.adjJb)^(p.b / αt))^αt
    end

    return hE * gJ
end

function log_profile_density(x::G, vr::G, vt::G, p::NFWParameters{G}; N::G = 1000.0, αt::G = 1.0) where G <: AbstractFloat
    if x < 0.0
        return -Inf
    end

    #Compute Energy
    E = 0.5 * (vt * vt + vr * vr) + gravitational_potential(x, p)
    lE = log(E)
    if E < p.Φlim && E >= 0.0
        lhE = log(N) + p.a * lE + (p.d/p.q) * log(E^p.q + p.adjEc^p.q) + p.e * log(p.Φlim - E)
    else
        return -Inf
    end

    #Compute angular momentum
    J = abs(x * p.rs * vt) #vt = v * sin(theta)
    if p.b <= 0.0
        #gJ = 1.0 / (1.0 + (J / p.adjJb)^(-p.b))
        lgJ = -αt * log(1.0 + (J / p.adjJb)^(-p.b / αt))
    else
        #gJ = 1.0 + (J / p.adjJb)^p.b
        lgJ = αt * log(1.0 + (J / p.adjJb)^(p.b / αt))
    end

    return lhE + lgJ
end

function profile_density(x::G, vr::G, vt::G, p::SFWParameters{G}; N::G = 1000.0, αt::G = 1.0) where G <: AbstractFloat
    if x < 0.0
        return 0.0
    end

    #Compute Energy
    E = 0.5 * (vt * vt + vr * vr) + gravitational_potential(x, p)
    if E < p.Φlim && E >= 0.0
        hE = N * E^p.a * (E^p.q + p.adjEc^p.q)^(p.d/p.q) * (p.Φlim - E)^p.e
    else
        return 0.0
    end

    #Compute angular momentum
    J = abs(x * p.rs * vt) #vt = v * sin(theta)
    if p.b <= 0.0
        #gJ = 1.0 / (1.0 + (J / p.adjJb)^(-p.b))
        gJ = (1.0 + (J / p.adjJb)^(-p.b / αt))^-αt
    else
        #gJ = 1.0 + (J / p.adjJb)^p.b
        gJ = (1.0 + (J / p.adjJb)^(p.b / αt))^αt
    end

    return hE * gJ
end

function log_profile_density(x::G, vr::G, vt::G, p::SFWParameters{G}; N::G = 1000.0, αt::G = 1.0) where G <: AbstractFloat
    if x < 0.0
        return -Inf
    end

    #Compute Energy
    E = 0.5 * (vt * vt + vr * vr) + gravitational_potential(x, p)
    lE = log(E)
    if E < p.Φlim && E >= 0.0
        lhE = log(N) + p.a * lE + (p.d/p.q) * log(E^p.q + p.adjEc^p.q) + p.e * log(p.Φlim - E)
    else
        return -Inf
    end

    #Compute angular momentum
    J = abs(x * p.rs * vt) #vt = v * sin(theta)
    if p.b <= 0.0
        lgJ = -αt * log(1.0 + (J / p.adjJb)^(-p.b / αt))
    else
        lgJ = αt * log(1.0 + (J / p.adjJb)^(-p.b / αt))
    end

    return lhE + lgJ
end


@testset "Kronecker product" begin
    A = [1 0; 0 1]
    B = [1 0; 0 -1]
    C = [1 0 0 0; 0 -1 0 0; 0 0 1 0; 0 0 0 -1]
    @test A⊗B == C
    @test ⊗(A,A,B,B) == A⊗A⊗B⊗B

    A, B, C, D = rand(2, 2), rand(2, 2), rand(2, 2), rand(2, 2)
    @test A⊗(B+C) ≈ A⊗B + A⊗C
    @test (B+C)⊗A ≈ B⊗A + C⊗A
    @test (A⊗B)⊗C ≈ A⊗(B⊗C)
    @test (A⊗B)*(C⊗D) ≈ (A*C)⊗(B*D)
end


using MultiFloats
using Printf

function sqrt_precision(x::MultiFloat{T,N}) where {T,N}
    setprecision(BigFloat, 2200) do
        approx = BigFloat(sqrt(x))
        exact = sqrt(BigFloat(x))
        -Float64(log2(abs((approx - exact) / exact)))
    end
end

function exp_precision(x::MultiFloat{T,N}) where {T,N}
    setprecision(BigFloat, 2200) do
        approx = BigFloat(exp(x))
        exact = exp(BigFloat(x))
        -Float64(log2(abs((approx - exact) / exact)))
    end
end

num_trials = 10^6

println("SQRT (CLEAN)")
use_clean_multifloat_arithmetic()
for N = 2 : 8
    correctbits = floor(Int, minimum(
            sqrt_precision(ldexp(rand(Float64x{N}), rand(-300:300)))
            for _ = 1 : num_trials))
    @printf("Float64x%d: %.3d bits (%2d lost, %.3f%% loss)\n", N, correctbits,
        precision(Float64x{N}) - correctbits,
        100 * abs(precision(Float64x{N}) - correctbits) / precision(Float64x{N}))
end
println()

println("SQRT (STANDARD)")
use_standard_multifloat_arithmetic()
for N = 2 : 8
    correctbits = floor(Int, minimum(
            sqrt_precision(ldexp(rand(Float64x{N}), rand(-300:300)))
            for _ = 1 : num_trials))
    @printf("Float64x%d: %.3d bits (%2d lost, %.3f%% loss)\n", N, correctbits,
        precision(Float64x{N}) - correctbits,
        100 * abs(precision(Float64x{N}) - correctbits) / precision(Float64x{N}))
end
println()

println("SQRT (SLOPPY)")
use_sloppy_multifloat_arithmetic()
for N = 2 : 8
    correctbits = floor(Int, minimum(
            sqrt_precision(ldexp(rand(Float64x{N}), rand(-300:300)))
            for _ = 1 : num_trials))
    @printf("Float64x%d: %.3d bits (%2d lost, %.3f%% loss)\n", N, correctbits,
        precision(Float64x{N}) - correctbits,
        100 * abs(precision(Float64x{N}) - correctbits) / precision(Float64x{N}))
end
println()

println("EXP (CLEAN)")
use_clean_multifloat_arithmetic()
for N = 2 : 8
    correctbits = floor(Int, minimum(
            exp_precision(400 * rand(Float64x{N}) - 200)
            for _ = 1 : num_trials))
    @printf("Float64x%d: %.3d bits (%2d lost, %.3f%% loss)\n", N, correctbits,
        precision(Float64x{N}) - correctbits,
        100 * abs(precision(Float64x{N}) - correctbits) / precision(Float64x{N}))
end
println()

println("EXP (STANDARD)")
use_standard_multifloat_arithmetic()
for N = 2 : 8
    correctbits = floor(Int, minimum(
            exp_precision(400 * rand(Float64x{N}) - 200)
            for _ = 1 : num_trials))
    @printf("Float64x%d: %.3d bits (%2d lost, %.3f%% loss)\n", N, correctbits,
        precision(Float64x{N}) - correctbits,
        100 * abs(precision(Float64x{N}) - correctbits) / precision(Float64x{N}))
end
println()

println("EXP (SLOPPY)")
use_sloppy_multifloat_arithmetic()
for N = 2 : 8
    correctbits = floor(Int, minimum(
            exp_precision(400 * rand(Float64x{N}) - 200)
            for _ = 1 : num_trials))
    @printf("Float64x%d: %.3d bits (%2d lost, %.3f%% loss)\n", N, correctbits,
        precision(Float64x{N}) - correctbits,
        100 * abs(precision(Float64x{N}) - correctbits) / precision(Float64x{N}))
end
println()


using Flux, Trebuchet
using Zygote: forwarddiff
using Statistics: mean
using Random

lerp(x, lo, hi) = x*(hi-lo)+lo

function shoot(wind, angle, weight)
  Trebuchet.shoot((wind, Trebuchet.deg2rad(angle), weight))[2]
end

shoot(ps) = forwarddiff(p -> shoot(p...), ps)

Random.seed!(0)

model = Chain(Dense(2, 16, σ),
              Dense(16, 64, σ),
              Dense(64, 16, σ),
              Dense(16, 2)) |> f64

θ = params(model)

function aim(wind, target)
  angle, weight = model([wind, target])
  angle = σ(angle)*90
  weight = weight + 200
  angle, weight
end

distance(wind, target) =
  shoot(Tracker.collect([wind, aim(wind, target)...]))

function loss(wind, target)
  try
    (distance(wind, target) - target)^2
  catch e
    # Roots.jl sometimes give convergence errors, ignore them
    param(0)
  end
end

DIST  = (20, 100)	# Maximum target distance
SPEED =   5 # Maximum wind speed

target() = (randn() * SPEED, lerp(rand(), DIST...))

meanloss() = mean(sqrt(loss(target()...)) for i = 1:100)

opt = ADAM()

dataset = (target() for i = 1:100_000)
cb = Flux.throttle(() -> @show(meanloss()), 10)

Flux.train!(loss, θ, dataset, opt, cb = cb)


struct ExactFormulation <: AbstractRelaxation end

function reformulate_unary_quadratic_term!(
    model::MOI.ModelLike,
    reformulation::Reformulation{ExactFormulation},
    x::VI,
    y::VI,
    l::Float64,
    u::Float64,
)
    MOI.add_constraint(
        model,
        1.0 * SV(y) - 1.0 * SV(x) * SV(x),
        MOI.EqualTo(0.0),
    )
    return nothing
end


export optimal
using Optim
using Statistics

# sum of absolute differences (element-wise)
function Δ⃯(𝒄₁, 𝒄₂)
  return sum(abs.(𝒄₁- 𝒄₂))
end

# sum of distances outside bounding box (element-wise)
function Δ͍(𝒄, 𝒄₋, 𝒄₊)
  return sum(max.(𝒄₋-𝒄, 0.0, 𝒄-𝒄₊))
end

# collection 𝒄 with last element dropped
function drop(𝒄)
  return 𝒄[1:end-1]
end

# fill collection with length of 𝒄 with average of 𝒄
function avg(𝒄)
  return fill(mean(𝒄), length(𝒄))
end

# map frequencies and diameters to locations, but drop the flute-length value
⇴ = drop ∘ mapflute
# push a value to front of vector, while dropping last value
⥆ = drop ∘ vcat

# generate error function with scoped constants
function mkerrfn(𝒇, 𝒅₊, 𝒑₋, 𝒑₊)
  𝒍⃯ = 𝒇⇴ 𝒅₊ # positions of holes with maximum allowed diameters
  function errfn(𝒅)
    # locations
    𝒍 = 𝒇⇴ 𝒅      # proposed diameters
    𝒍̲ = 𝒇⇴ avg(𝒅) # average of proposed
    𝒍⃮ = 0.0⥆ 𝒍    # shift proposed locations left (push 0.0)
    𝒍⃭ = 𝒍⃮+𝒑₋      # closest desired position to previous hole
    𝒍⃬ = 𝒍⃮+𝒑₊      # furthest desired position from previous hole
    # error terms
    𝑒̲ = Δ⃯(𝒍, 𝒍̲)       # sum of distances from average
    𝑒⃯ = Δ⃯(𝒍, 𝒍⃯)       # sum of distances from max positions
    𝑒͍ = Δ͍(𝒍, 𝒍⃭, 𝒍⃬)^2  # squared sum of distances outside desired from previous holes
    # sum and weigh error terms
    𝑒 = 𝑒̲ + 2𝑒⃯ + 3𝑒͍
    return 𝑒
  end
  return errfn
end

function optimal(𝒇, 𝒅₋, 𝒅₊, 𝒑₋, 𝒑₊; trace=false)
  # minimize error function
  errfn = mkerrfn(𝒇, 𝒅₊, 𝒑₋, 𝒑₊)
  # box-constrained, initial parameters (bad guess)
  𝒅₁ = [8.8, 8.8, 6.6, 8.8, 9.9, 5.5]
  # simulated annealing (round 1, fast cooldown)
  algo = SAMIN(rt=0.9)
  options = Optim.Options(iterations=Int(1e5), show_trace=trace, show_every=Int(1e4))
  𝒅₂ = Optim.minimizer(optimize(errfn, 𝒅₋, 𝒅₊, 𝒅₁, algo, options))
  # simulated annealing (round 2, slow cooldown)
  algo = SAMIN(rt=0.99)
  options = Optim.Options(iterations=Int(1e6), show_trace=trace, show_every=Int(1e4))
  result = optimize(errfn, 𝒅₋, 𝒅₊, 𝒅₂, algo, options)
  𝒅ₕ = Optim.minimizer(result)
  𝑒 = Optim.minimum(result)
  # proposed diameters
  return 𝒅ₕ, 𝑒
end