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

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 ReconFuncs

codefile = "scratch_NM/Experiments_3-24.jl"
Nx=64
im = readdust(Nx)
true_img = im[:, :, 1]

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 20), ("norm", false)])
recon_settings = Dict([("target_type", "ground truth"), ("covar_type", "white_noise"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("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)
#################################

Nx=64
im = readdust(Nx)
true_img = im[:, :, 1]

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 20), ("norm", false)])
recon_settings = Dict([("target_type", "ground truth"), ("covar_type", "white_noise"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.1), ("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)
#################################

Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 20), ("norm", false)])
recon_settings = Dict([("target_type", "ground truth"), ("covar_type", "white_noise"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.1), ("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)

########################
#Weird behavior here: just memory issues -- works afetr closing window and reopening

Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false)])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(true_img).^(-2)), ("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
fname_save = "scratch_NM/NewWrapper/3-24/Exp1.jld2"
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save
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/PCA/sfd_dbntargwhitenoisecov_lam_imglevel_smoothinit_apd_logspace.png")

println("Mean Abs Res", mean(abs.(init - true_img)), " ", mean(abs.(recon_img - true_img)))

aptrue = apodizer(true_img)
apinit = apodizer(init)
aprecon = apodizer(recon_img)
Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname="scratch_NM/PCA/sfd_dbntargwhitenoisecov_lam_imglevel_smoothinit_apd_logspacekern1.png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname="scratch_NM/PCA/sfd_dbntargwhitenoisecov_lam_imglevel_smoothinit_apd_logspace_6p.png")
println("Mean Abs Res: Init = ", mean(abs.(apinit - aptrue)), " Recon = ", mean(abs.(aprecon - aptrue)))

##3-35###################################
#Testing JLD
using JLD2
@save "scratch_NM/NewWrapper/3-24/test.jld2" true_img init recon_img recon_settings

@load "scratch_NM/NewWrapper/3-24/test.jld2" load_gt=true_img load_init=init load_recon=recon_img load_dict=recon_settings

using FileIO

save("scratch_NM/NewWrapper/3-24/test_fiolarge.jld2", Dict("true_img"=>true_img, "init"=>init, "recon"=>recon_img, "dict"=>recon_settings))

load("scratch_NM/")

##3-26###############################################################################################################
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false)])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(true_img).^(-2)), ("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
fname_save = "scratch_NM/NewWrapper/3-24/Exp1.jld2"
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save
recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)
#Stops after 5 iterations

#Without apodization | SFD | SFDCovar
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false)])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(true_img).^(-2)), ("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
fname_save = "scratch_NM/NewWrapper/3-24/Exp1.jld2"
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save
recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

#With apd | SFD | SFDCovar
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false)])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(true_img).^(-2)), ("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
fname_save = "scratch_NM/NewWrapper/3-24/SfdExp1.jld2"
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save
recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)
Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname="scratch_NM/NewWrapper/3-24/Exp1.png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname="scratch_NM/NewWrapper/3-24/Exp1_6p.png")
println("Mean abs Recon-True", round(mean(abs.(recon_img - true_img)), digits=6))
println("Mean abs Recon-Init", round(mean(abs.(recon_img - init)), digits=6))
println("Mean abs Recon-True", round(mean(abs.(init- true_img)), digits=6))

#With apd | GT Target | SFDCovar
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :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", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(true_img).^(-2)), ("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
fname_save = "scratch_NM/NewWrapper/3-24/Sfd_gttargExp2.jld2"
recon_settings["safemode"] = false
recon_settings["fname_save"] = fname_save
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = SFD Dbn Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, "scratch_NM/NewWrapper/3-24/Trace_Sfd_gttargExp2.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname="scratch_NM/NewWrapper/3-24/Sfd_gttargExp2.png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname="scratch_NM/NewWrapper/3-24/Sfd_gttargExp2_6p.png")
convert(Array, Optim.trace(resobj))

#With apd | GT Target | GTCovar
fname_save = "scratch_NM/NewWrapper/3-24/WNtargWNcov"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "white_noise"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(true_img).^(-2)), ("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_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = White Noise Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname="scratch_NM/NewWrapper/3-24/Sfd_gttargExp2.png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname="scratch_NM/NewWrapper/3-24/Sfd_gttargExp2_6p.png")
convert(Array, Optim.trace(resobj))

#Same as above with LBFGS(): Does the descent not stall?
#With apd | GT Target | WNCovar
fname_save = "scratch_NM/NewWrapper/3-24/GTtargWNcov_LBFGS"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", LBFGS())])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "white_noise"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(true_img).^(-2)), ("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_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = White Noise Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname="scratch_NM/NewWrapper/3-24/Sfd_gttargExp2.png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname="scratch_NM/NewWrapper/3-24/Sfd_gttargExp2_6p.png")
convert(Array, Optim.trace(resobj))

#GT Targ | SFD cov | Lambda=0: DECENT-ISH
fname_save = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_noreg_CG"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("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...
s_true = DHC_compute_apd(log.(true_img), filter_hash;norm=false, iso=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=s_true)
#Check
regs_true = DHC_compute_apd(true_img, filter_hash; dhc_args...)[coeff_mask]
s2mean - s_true

recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov
recon_settings["safemode"] = false
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

#GT Targ | SFD cov
fname_save = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_lamcorr_LBFGS"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", LBFGS())])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(apodizer(log.(true_img))).^(-2)), ("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_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

#GT Targ | WN Cov
#Same as above with LBFGS(): Does the descent not stall?
#With apd | GT Target | WNCovar
fname_save = "scratch_NM/NewWrapper/3-24/GTtargWNcov_CGExp3"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "white_noise"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 1.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_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = SFD Cov Lam", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

##3-26 What's wrong with the sfd dbn loss?
#Are you already closer than 1 std to the true image?
fname = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_lamcorr_LBFGS"
gttarget = load(fname*".jld2")
#order: A*log*Im
ltrue = log.(gttarget["true_img"])
linit = log.(gttarget["init"])
lrecon =log.(gttarget["recon"])

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
#init
l2init = 0.5*lam_old*sum((alinit - altrue).^2)
s_init = DHC_compute_apd(linit, filter_hash; dhc_args...)[coeff_mask]
s_recon = DHC_compute_apd(lrecon, filter_hash; dhc_args...)[coeff_mask]
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
#819, 322

#recon
l2init = 0.5*lam_old*sum((alrecon - altrue).^2)
l1init = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
#812, 439

#Are you already closer than 1 std to the true image?
fname = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_noreg_LBFGS"
gttarget = load(fname*".jld2")
#order: A*log*Im
ltrue = log.(gttarget["true_img"])
linit = log.(gttarget["init"])
lrecon =log.(gttarget["recon"])

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
#init
l2init = 0.5*lam_old*sum((alinit - altrue).^2)
s_init = DHC_compute_apd(linit, filter_hash; dhc_args...)[coeff_mask]
s_recon = DHC_compute_apd(lrecon, filter_hash; dhc_args...)[coeff_mask]
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
#819

#recon
l2init = 0.5*lam_old*sum((alrecon - altrue).^2)
l1init = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
#821

##
fname = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_noreg_CG"
gttarget = load(fname*".jld2")
#order: A*log*Im
ltrue = log.(gttarget["true_img"])
linit = log.(gttarget["init"])
lrecon =log.(gttarget["recon"])

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
#init
l2init = 0.5*lam_old*sum((alinit - altrue).^2)
s_init = DHC_compute_apd(linit, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_recon = DHC_compute_apd(lrecon, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_true = DHC_compute_apd(ltrue, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
#16.1

#recon
l2init = 0.5*lam_old*sum((alrecon - altrue).^2)
l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
#1e-5

##

#SFD Targ | SFD cov | Lambda=Corr
fname_save = "scratch_NM/NewWrapper/3-24/SFDtargSFDcov_lamcorr_CG"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = imfilter(true_img + noise, Kernel.gaussian(1.0))
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(apodizer(log.(true_img))).^(-2)), ("white_noise_args", white_noise_args)]) #Add constraints
#"white_noise_args": args for white_noise sims if GaussianLoss and Covar_type='white_noise'. :loc...
#s_true = DHC_compute_apd(log.(true_img), filter_hash;norm=false, iso=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
#regs_true = DHC_compute_apd(true_img, filter_hash; dhc_args...)[coeff_mask]
#s2mean - s_true

recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov
recon_settings["safemode"] = false
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

#Testing

fname = "scratch_NM/NewWrapper/3-24/SFDtargSFDcov_lamcorr_CG"
gttarget = load(fname*".jld2")
#order: A*log*Im
ltrue = log.(gttarget["true_img"])
linit = log.(gttarget["init"])
lrecon =log.(gttarget["recon"])

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
#init
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
s_init = DHC_compute_apd(linit, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_recon = DHC_compute_apd(lrecon, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_true = DHC_compute_apd(ltrue, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
#True: 277,

#recon
l2init = 0.5*lam_old*sum((alrecon - altrue).^2)
l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
#1e-5

##No smooth
#GT Targ | SFD cov | Lambda=0: DECENT-ISH
fname_save = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_nosmooth_CG"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(apodizer(log.(true_img))).^(-2)), ("white_noise_args", white_noise_args)]) #Add constraints
#"white_noise_args": args for white_noise sims if GaussianLoss and Covar_type='white_noise'. :loc...
s_true = DHC_compute_apd(log.(true_img), filter_hash;norm=false, iso=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=s_true)
#Check
regs_true = DHC_compute_apd(true_img, filter_hash; dhc_args...)[coeff_mask]
#s2mean - s_true

recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

##SFD Targ | SFD Cov
fname_save = "scratch_NM/NewWrapper/3-24/SFDtargSFDcov_lamcorr_nosmooth_CG_recheck"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(apodizer(log.(true_img))).^(-2)), ("white_noise_args", white_noise_args)]) #Add constraints
#"white_noise_args": args for white_noise sims if GaussianLoss and Covar_type='white_noise'. :loc...
#s_true = DHC_compute_apd(log.(true_img), filter_hash;norm=false, iso=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
regs_true = DHC_compute_apd(true_img, filter_hash; dhc_args...)[coeff_mask]
#s2mean - s_true

recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(SFD) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
maximum(im)
mean(im)
minimum(im)

##SFD Targ | SFD Cov | No Log
fname_save = "scratch_NM/NewWrapper/3-24/SFDtargSFDcov_lamcorr_nosmooth_reg_CG"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(apodizer(true_img)).^(-2)), ("white_noise_args", white_noise_args)]) #Add constraints

s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
regs_true = DHC_compute_apd(true_img, filter_hash; dhc_args...)[coeff_mask]
#s2mean - s_true

recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(SFD) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
maximum(im)
mean(im)
minimum(im)

##GT Targ | SFD Cov | Reg
fname_save = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_lam0_nosmooth_reg_CG"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
heatmap(init)
#heatmap(true_img)

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) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "sfd_dbn"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.0), ("white_noise_args", white_noise_args)]) #Add constraints

s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
regs_true = DHC_compute_apd(true_img, filter_hash; dhc_args...)[coeff_mask]
#s2mean - s_true

recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

##GT Targ | GT Cov | Reg
fname_save = "scratch_NM/NewWrapper/3-24/GTtargWNcov_lamcorr_nosmooth_reg_CG"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
heatmap(init)
#heatmap(true_img)

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, :coeff_mask=>coeff_mask) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :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), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "white_noise"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(apodizer(true_img)).^(-2)), ("white_noise_args", white_noise_args)]) #Add constraints

s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
regs_true = DHC_compute_apd(true_img, filter_hash; dhc_args...)[coeff_mask]
#s2mean - s_true

recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov
recon_settings["safemode"] = true
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(GT) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

#Testing: Where it DID work
fname = "scratch_NM/NewWrapper/3-24/GTtargWNcov_lamcorr_nosmooth_reg_CG"
gttarget = load(fname * ".jld2")
#order: A*log*Im
true_img = gttarget["true_img"]
init = gttarget["init"]
recon =gttarget["recon"]

altrue = apodizer(gttarget["true_img"])
alinit = apodizer(gttarget["init"])
alrecon = apodizer(gttarget["recon"])

println("Mean Abs Res: Init-True = ", mean(abs.(alinit - altrue)), " Recon-True = ", mean(abs.(alrecon - altrue)))
println("Mean Abs Frac Res", mean(abs.((alinit - altrue)./altrue)), " Recon-True=", mean(abs.((alrecon - altrue)./altrue)))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
#init
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
s_init = DHC_compute_apd(init, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_recon = DHC_compute_apd(recon, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_true = DHC_compute_apd(true_img, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*lam_old*sum((alinit - altrue).^2)
#True: 0, 2280
#Init: 21321, 0

#recon
l2init = 0.5*lam_old*sum((alrecon - alinit).^2)
l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
#Recon: 173, 959

#Testing: Where it DID NOT work
fname = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_lamcorr_nosmooth_reg_CG"
gttarget = load(fname * ".jld2")
#order: A*log*Im
true_img = gttarget["true_img"]
init = gttarget["init"]
recon =gttarget["recon"]

altrue = apodizer(gttarget["true_img"])
alinit = apodizer(gttarget["init"])
alrecon = apodizer(gttarget["recon"])

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
#init
s_init = DHC_compute_apd(init, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_recon = DHC_compute_apd(recon, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_true = DHC_compute_apd(true_img, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
l1truecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*lam_old*sum((alinit - altrue).^2)
#True: 0, 2327
#Init: 0.0009, 0

#recon
l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*lam_old*sum((alrecon - alinit).^2)
#Recon: 0.0009, 1e-9

#Testing where it did work but in log space.
fname = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_nosmooth_CG"
gttarget = load(fname*".jld2")
#order: A*log*Im
ltrue = log.(gttarget["true_img"])
linit = log.(gttarget["init"])
lrecon =log.(gttarget["recon"])

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
#init
s_init = DHC_compute_apd(linit, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_recon = DHC_compute_apd(lrecon, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
s_true = DHC_compute_apd(ltrue, filter_hash; norm=false, iso=false, dhc_args...)[coeff_mask]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*lam_old*sum((alinit - altrue).^2)
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
#True: 0, 2716
#Init: 3717, 0

#recon
l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*lam_old*sum((alrecon - alinit).^2)
#306, 549

true_img = gttarget["true_img"]
init = gttarget["init"]
recon_img =gttarget["recon"]

fname = "scratch_NM/NewWrapper/3-24/GTtargSFDcov_noreg_nosmooth_CG"
gttarget = load(fname*".jld2")
println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

##4-1: Generating standardized init images and experiments
Nx=64
imall = readsfd(Nx)
Random.seed!(41)

imlist = [1, 5, 10, 50, 100, 500, 1000, 5000, 10000]
for imid in imlist
  fname_save = "scratch_NM/StandardizedExp/Nx64/" * "Data_" * string(imid) * ".jld2"
  true_img = imall[:, :, imid]
  stdval = std(true_img)
  noise = rand(Normal(0.0, stdval), Nx, Nx)
  init = true_img + noise
  zind = findall(init.<0)
  println("Percentage set to 1e-6", 100*length(zind)/length(init))
  init[zind] .= 1e-6
  save(fname_save, Dict("true_img"=>true_img, "seed"=>41, "init"=>init, "noise model"=>"White noise, No smoothing. sigma=std(true_img)", "std"=>stdval))
end

savexp = load("scratch_NM/StandardizedExp/Nx64/Data_5000.jld2")
heatmap(savexp["true_img"])
heatmap(savexp["init"])

##4-2 Comparing SFD-SFD runs with and without apd
#Log-Apd-Iso
fname = "scratch_NM/NewWrapper/3-29/SFDTargSFDCov/log_apd_iso"
gttarget = load(fname*".jld2")
#order: A*log*Im
ltrue = log.(gttarget["true_img"])
linit = log.(gttarget["init"])
lrecon =log.(gttarget["recon"])

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
filter_hash = fink_filter_hash(1, 8, nx=64, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])
coeff_mask = falses(2+S1iso+S2iso)
coeff_mask[S1iso+3:end] .= true
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>true)
length(coeff_mask)
#init
s_init = DHC_compute_wrapper(linit, filter_hash; norm=false,  dhc_args...)[coeff_mask]
s_recon = DHC_compute_wrapper(lrecon, filter_hash; norm=false, dhc_args...)[coeff_mask]
s_true = DHC_compute_wrapper(ltrue, filter_hash; norm=false, dhc_args...)[coeff_mask]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*lam_old*sum((alinit - altrue).^2)
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
#True:43.1, 2505
#Init: 903, 0

#recon
l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*lam_old*sum((alrecon - alinit).^2)
#228, 184

true_img = gttarget["true_img"]
init = gttarget["init"]
recon_img =gttarget["recon"]

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("RMSE Init-True= ", (mean((init - true_img).^2)).^0.5, " Recon-True=", (mean((recon_img - true_img).^2)).^0.5)

##Log-Apd-NoIso
fname = "scratch_NM/NewWrapper/3-29/SFDTargSFDCov/log_apd_iso"
gttarget = load(fname*".jld2")
#order: A*log*Im
ltrue = log.(gttarget["true_img"])
linit = log.(gttarget["init"])
lrecon =log.(gttarget["recon"])

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
filter_hash = fink_filter_hash(1, 8, nx=64, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])
coeff_mask = falses(2+S1iso+S2iso)
coeff_mask[S1iso+3:end] .= true
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>true)
length(coeff_mask)
#init
s_init = DHC_compute_wrapper(linit, filter_hash; norm=false,  dhc_args...)[coeff_mask]
s_recon = DHC_compute_wrapper(lrecon, filter_hash; norm=false, dhc_args...)[coeff_mask]
s_true = DHC_compute_wrapper(ltrue, filter_hash; norm=false, dhc_args...)[coeff_mask]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*lam_old*sum((alinit - altrue).^2)
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
#True:43.1, 2505
#Init: 903, 0

#recon
l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*lam_old*sum((alrecon - alinit).^2)
#228, 184

true_img = gttarget["true_img"]
init = gttarget["init"]
recon_img =gttarget["recon"]

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("RMSE Init-True= ", (mean((init - true_img).^2)).^0.5, " Recon-True=", (mean((recon_img - true_img).^2)).^0.5)

##Log-Apd-Noiso
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/log_apd_noiso/nophi_lamcorr_recon_1000"
gttarget = load(fname*".jld2")
(Nf,) = size(filter_hash["filt_index"])
reshape(gttarget["coeff_mask"][37:end], (Nf, Nf))
#order: A*log*Im
ltrue = log.(gttarget["true_img"])
linit = log.(gttarget["init"])
lrecon =log.(gttarget["recon"])

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

lam_corr = std(altrue).^(-2)
lam_old = gttarget["dict"]["lambda"]

#Loss contrib
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
filter_hash = fink_filter_hash(1, 8, nx=64, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])
coeff_mask = falses(2+S1iso+S2iso)
coeff_mask[S1iso+3:end] .= true
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>true)
length(coeff_mask)
#init
s_init = DHC_compute_wrapper(linit, filter_hash; norm=false,  dhc_args...)[coeff_mask]
s_recon = DHC_compute_wrapper(lrecon, filter_hash; norm=false, dhc_args...)[coeff_mask]
s_true = DHC_compute_wrapper(ltrue, filter_hash; norm=false, dhc_args...)[coeff_mask]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*lam_old*sum((alinit - altrue).^2)
l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
#True:43.1, 2505
#Init: 903, 0

#recon
l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*lam_old*sum((alrecon - alinit).^2)
#228, 184

true_img = gttarget["true_img"]
init = gttarget["init"]
recon_img =gttarget["recon"]

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("RMSE Init-True= ", (mean((init - true_img).^2)).^0.5, " Recon-True=", (mean((recon_img - true_img).^2)).^0.5)

##4-2
#Log Apd Iso
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/log_apd_iso/lamcorr_recon_1000"
gttarget = load(fname*".jld2")
#Calc 1dps log
Nx = size(gttarget["true_img"])[1]
kbins = collect(1:Nx/2.0)

true_img = gttarget["true_img"]
init = gttarget["init"]
recon_img =gttarget["recon"]

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("RMSE Init-True= ", (mean((init - true_img).^2)).^0.5, " Recon-True=", (mean((recon_img - true_img).^2)).^0.5)

true_lps = calc_1dps(apodizer(log.(gttarget["true_img"])), kbins)
recon_lps = calc_1dps(apodizer(log.(gttarget["recon"])), kbins)
init_lps = calc_1dps(apodizer(log.(gttarget["init"])), kbins)
p = plot(log.(kbins), log.(true_lps), label="True")
plot!(log.(kbins), log.(recon_lps), label="Recon")
plot!(log.(kbins), log.(init_lps), label="Init")
plot!(title="P(k) of Log Image")

true_ps = calc_1dps(apodizer(gttarget["true_img"]), kbins)
recon_ps = calc_1dps(apodizer(gttarget["recon"]), kbins)
init_ps = calc_1dps(apodizer(gttarget["init"]), kbins)
p = plot(log.(kbins), log.(true_ps), label="True")
plot!(log.(kbins), log.(recon_ps), label="Recon")
plot!(log.(kbins), log.(init_ps), label="Init")
plot!(title="P(k) of Image")

s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>true)

coeff_mask = gttarget["coeff_mask"]
s_true = DHC_compute_wrapper(log.(gttarget["true_img"]), filter_hash, norm=false; dhc_args...)[coeff_mask]
s_init = DHC_compute_wrapper(log.(gttarget["init"]), filter_hash, norm=false; dhc_args...)[coeff_mask]
s_recon = DHC_compute_wrapper(log.(gttarget["recon"]), filter_hash, norm=false; dhc_args...)[coeff_mask]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*lam_old*sum((alinit - altrue).^2)
#78.2, 2505

l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
#2.68e7, 0

l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*lam_old*sum((alrecon - alinit).^2)
#4168, 183.61

#Log, Apd, Noiso

fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/log_apd_noiso/lamcorr_recon_1000"
gttarget = load(fname*".jld2")
#Calc 1dps log
Nx = size(gttarget["true_img"])[1]
kbins = collect(1:Nx/2.0)

true_img = gttarget["true_img"]
init = gttarget["init"]
recon_img =gttarget["recon"]

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("RMSE Init-True= ", (mean((init - true_img).^2)).^0.5, " Recon-True=", (mean((recon_img - true_img).^2)).^0.5)

true_lps = calc_1dps(apodizer(log.(gttarget["true_img"])), kbins)
recon_lps = calc_1dps(apodizer(log.(gttarget["recon"])), kbins)
init_lps = calc_1dps(apodizer(log.(gttarget["init"])), kbins)
p = plot(log.(kbins), log.(true_lps), label="True")
plot!(log.(kbins), log.(recon_lps), label="Recon")
plot!(log.(kbins), log.(init_lps), label="Init")
plot!(title="P(k) of Log Image: Noiso")

true_ps = calc_1dps(apodizer(gttarget["true_img"]), kbins)
recon_ps = calc_1dps(apodizer(gttarget["recon"]), kbins)
init_ps = calc_1dps(apodizer(gttarget["init"]), kbins)
p = plot(log.(kbins), log.(true_ps), label="True")
plot!(log.(kbins), log.(recon_ps), label="Recon")
plot!(log.(kbins), log.(init_ps), label="Init")
plot!(title="P(k) of Image: NoIso")

s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false)
gttarget["coeff_mask"]

s_true = DHC_compute_wrapper(log.(gttarget["true_img"]), filter_hash, norm=false; dhc_args...)[gttarget["coeff_mask"]]
s_init = DHC_compute_wrapper(log.(gttarget["init"]), filter_hash, norm=false; dhc_args...)[gttarget["coeff_mask"]]
s_recon = DHC_compute_wrapper(log.(gttarget["recon"]), filter_hash, norm=false; dhc_args...)[gttarget["coeff_mask"]]

ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*lam_old*sum((alinit - altrue).^2)
#648, 2505

l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*lam_old*sum((alinit - alinit).^2)
#7.25e7, 0

l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*lam_old*sum((alrecon - alinit).^2)
#2948, 184

sig = (diag(s_targ_invcov)).^(-0.5)
z_true_targ = (s_true - s_targ_mean)./sig
z_recon_targ = (s_recon - s_targ_mean)./sig
z_recon_true = (s_recon - s_true)./sig
z_init_targ = (s_init - s_targ_mean)./sig

zscores = hcat(z_true_targ, z_recon_targ, z_recon_true)
Nf = size(filter_hash["filt_index"])[1]
images = [reshape(z_true_targ, (Nf, Nf)), reshape(z_recon_targ, (Nf, Nf)), reshape(z_recon_true, (Nf, Nf)), reshape(z_init_targ, (Nf, Nf))]
titles= ["True-Target", "Recon-Target", "Recon-True", "Init-Targ"]
pl12 = plot(
        heatmap(images[1], title=titles[1]),
        heatmap(images[2], title=titles[2]),
        heatmap(images[3],title= titles[3]),
        heatmap(images[4],title= titles[4]))
plot!(suptitle="S20 Z Scores")

clim = (minimum(z_true_targ), maximum(z_true_targ))
pl12 = plot(
        heatmap(images[1], title=titles[1], clim=clim),
        heatmap(images[2], title=titles[2], clim=clim),
        heatmap(images[3],title= titles[3], clim=clim),
        heatmap(images[4],title= titles[4], clim=clim))
plot!(suptitle="S20 Z Scores")
z_recon_targ = reshape(z_recon_targ, (Nf, Nf))

##Covariance
Nx=64
dbnimg = readsfd(Nx, logbool=true)
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/log_apd_noiso/lamcorr_recon_1000"
gttarget = load(fname*".jld2")
fhash = gttarget["filter_hash"]
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false)
coeff_mask = gttarget["coeff_mask"]
smean1, sdiag1, scov1 = dbn_coeffs_calc(dbnimg[:, :, 1:5000], fhash, dhc_args, coeff_mask)
smean2, sdiag2, scov2 = dbn_coeffs_calc(dbnimg[:, :, 5000:end], fhash, dhc_args, coeff_mask)

dbn1 = get_dbn_coeffs(dbnimg[:, :, 1:5000], fhash, dhc_args, coeff_mask = coeff_mask)
dbn2 = get_dbn_coeffs(dbnimg[:, :, 5001:end], fhash, dhc_args, coeff_mask= coeff_mask)
scov1inv = invert_covmat(scov1, 1e-6)
scov2inv = invert_covmat(scov2, 1e-6)

diff = dbn1 .- reshape(smean1, (1, 1156))
chisq1 = sum((diff * scov1inv) .* diff, dims=2)

diff = dbn2 .- reshape(smean2, (1, 1156))
chisq2 = sum((diff * scov2inv) .* diff, dims=2)

p = histogram(chisq1[:])
p = histogram(chisq2[:])
chisqsamp = rand(Distributions.Chisq(1156), 5000)
histogram!(chisqsamp)
########################

sd1inv = invert_covmat(sdiag1, 1e-10)
sd2inv = invert_covmat(sdiag2, 1e-10)

diff = dbn1 .- reshape(smean1, (1, 1156))
chisq1 = sum((diff * sd1inv) .* diff, dims=2)

diff = dbn2 .- reshape(smean2, (1, 1156))
chisq2 = sum((diff * sd2inv) .* diff, dims=2)

p = histogram(chisq1[:], label="Subset1")
p= histogram(chisq2[:], label="Subset2")
chisqsamp = rand(Distributions.Chisq(1156), 5000)
p = histogram(chisqsamp, label="True Chisq")

####################################
Nx=64
dbnimg = readsfd(Nx, logbool=true)
meanimg = mean(dbnimg, dims=[1, 2])
dbnimgms = dbnimg .- meanimg
dbn1 = get_dbn_coeffs(dbnimgms[:, :, 1:5000], fhash, dhc_args, coeff_mask = coeff_mask)
dbn2 = get_dbn_coeffs(dbnimgms[:, :, 5001:end], fhash, dhc_args, coeff_mask= coeff_mask)
smean1, sdiag1, scov1 = dbn_coeffs_calc(dbnimgms[:, :, 1:5000], fhash, dhc_args, coeff_mask)
smean2, sdiag2, scov2 = dbn_coeffs_calc(dbnimgms[:, :, 5001:end], fhash, dhc_args, coeff_mask)

scov1inv = invert_covmat(scov1, 1e-6)
scov2inv = invert_covmat(scov2, 1e-6)

diff = dbn1 .- reshape(smean1, (1, 1156))
chisq1 = sum((diff * scov1inv) .* diff, dims=2)

diff = dbn2 .- reshape(smean2, (1, 1156))
chisq2 = sum((diff * scov2inv) .* diff, dims=2)

p = histogram(chisq1[:])
p = histogram(chisq2[:])
chisqsamp = rand(Distributions.Chisq(1156), 5000)
histogram!(chisqsamp)

#########################################
diff = dbn1 .- reshape(smean1, (1, 1156))
chisq12 = sum((diff * scov2inv) .* diff, dims=2)

diff = dbn2 .- reshape(smean2, (1, 1156))
chisq21 = sum((diff * scov1inv) .* diff, dims=2)

p = histogram(chisq1[:], xlims=(0, 5000), label="ChiSq1-1")
histogram!(chisq12[:], xlims=(0, 5000), label="ChiSq1-2")

p = histogram(chisq2[:], xlims=(0, 5000), label="ChiSq2")
histogram!(chisq21[:], xlims=(0, 5000), label="ChiSq2-1")

chisqsamp = rand(Distributions.Chisq(1156), 5000)
histogram!(chisqsamp)
#########################################
#Log coeff, log im
Nx=64
dbnimg = readsfd(Nx, logbool=true)
#dbnimgms = dbnimg .- meanimg
dbn1 = log.(get_dbn_coeffs(dbnimg[:, :, 1:5000], fhash, dhc_args, coeff_mask = coeff_mask))
dbn2 = log.(get_dbn_coeffs(dbnimg[:, :, 5001:end], fhash, dhc_args, coeff_mask= coeff_mask))

#Mean dia, cov
Ncovsamp = 5000
s_targ_mean = mean(dbn1, dims=1)
scov  = (dbn1 .- s_targ_mean)' * (dbn1 .- s_targ_mean) ./(Ncovsamp-1)
smean1, sdiag1, scov1 = s_targ_mean, diag(scov), scov
s_targ_mean = mean(dbn2, dims=1)
scov  = (dbn2 .- s_targ_mean)' * (dbn2 .- s_targ_mean) ./(Ncovsamp-1)
smean2, sdiag2, scov2 = s_targ_mean, diag(scov), scov

scov1inv = invert_covmat(scov1, 1e-8)
scov2inv = invert_covmat(scov2, 1e-8)

diff = dbn1 .- reshape(smean1, (1, 1156))
chisq1 = sum((diff * scov1inv) .* diff, dims=2)

diff = dbn2 .- reshape(smean2, (1, 1156))
chisq2 = sum((diff * scov2inv) .* diff, dims=2)

p = histogram(chisq1[:])
p = histogram(chisq2[:])
chisqsamp = rand(Distributions.Chisq(1156/2), 5000)
histogram!(chisqsamp)

diff = dbn1 .- reshape(smean1, (1, 1156))
chisq12 = sum((diff * scov2inv) .* diff, dims=2)

diff = dbn2 .- reshape(smean2, (1, 1156))
chisq21 = sum((diff * scov1inv) .* diff, dims=2)

p = histogram(chisq1[:], label="ChiSq1-1")
histogram!(chisq12[:], label="ChiSq1-2")
histogram!(chisqsamp, label="ChiSqExpected")

p = histogram(chisq2[:],label="ChiSq2")
histogram!(chisq21[:],  label="ChiSq2-1")
histogram!(chisqsamp, label="ChiSqExpected")

##Save for MAF
Nx=64
dbnimg = readsfd(Nx, logbool=true)
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/log_apd_noiso/lamcorr_recon_1000"
gttarget = load(fname*".jld2")
fhash = gttarget["filter_hash"]
Nf = size(fhash["filt_index"])[1]
coeff_mask = falses(2+Nf+Nf^2)
coeff_maskS20 = trues((Nf, Nf))
coeff_mask[3+Nf:end] .= triu(coeff_maskS20)[:]
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false)
dbn = get_dbn_coeffs(dbnimg, fhash, dhc_args, coeff_mask= coeff_mask)

f = FITS("scratch_NM/dbncoeffs/S20sym_noiso_log.fits", "w")
write(f, dbn)
close(f)

##Implementing log S20 derivatives
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
#heatmap(init)
#heatmap(true_img)

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false) #Iso #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
if dhc_args[:iso]
    coeff_mask = falses(2+S1iso+S2iso)
    coeff_mask[S1iso+3:end] .= true
else #Not iso
    coeff_mask = falses(2+Nf+Nf^2)
    coeff_maskS20 = trues((Nf, Nf))
    coeff_mask[3+Nf:end] .= true
end

white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :norm=>false, :smooth=>false, :smoothval=>0.8) #Iso #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", false),("TransformedGaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.00456)]) #Add constraints

dbnocffs = log.(get_dbn_coeffs(log.(im), filter_hash, dhc_args, coeff_mask = coeff_mask))
fmean = mean(dbnocffs, dims=1)
fcov = (dbnocffs .- fmean)' * (dbnocffs .- fmean) ./(size(dbnocffs)[1] -1)
fcovinv = invert_covmat(fcov, 1e-8)

recon_settings["fs_targ_mean"] = fmean[:]
recon_settings["fs_invcov"] = fcovinv
recon_settings["safemode"] = true
fname_save = "scratch_NM/NewWrapper/TransfGaussian/test"
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
recon_settings["transform_func"] = Data_Utils.fnlog
recon_settings["transform_dfunc"] = Data_Utils.fndlog
resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)
ar1 = convert(Array{Any, 1}, exp.(rand(2)))
Data_Utils.fnlog(ar1)
p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(SFD) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

##########################################################################
##DEBUG
tonorm = false
input = init
fs_targ_mean = fmean[:]
fs_targ_invcov = fcovinv
lambda = 0.004
Nx=64
coeff_mask = falses(2+Nf+Nf^2)
coeff_mask[3+Nf:end] .= true
if dhc_args[:apodize]
    Ap = wind_2d(Nx)
    cA = sum(Ap)
    Apflat = reshape(Ap, Nx^2)
    od_nz_idx = findall(!iszero, Apflat) #findall((x->((x!=0) & (x!=1))), Apflat)
    #od_zidx = findall(iszero, Apflat)
    avals = Apflat[od_nz_idx]
    dA = zeros(Nx^2, Nx^2)
    dA[:, od_nz_idx] .= ((1.0 .- Apflat) * avals')./cA
    dA[diagind(dA)] += Apflat
else
    dA = I
end

function adaptive_apodizer(input_image::Array{Float64, 2})
    if dhc_args[:apodize]
        apd_image = apodizer(input_image)
    else
        apd_image = input_image
    end
    return apd_image
end

if (dhc_args[:doS20]) & (dhc_args[:iso])
    iso2nf2mask = zeros(Int64, Nf^2)
    M20 = filter_hash["S2_iso_mat"]
    for id=1:Nf^2 iso2nf2mask[M20.colptr[id]] = M20.rowval[id] end
    coeff_masks20 = coeff_mask[3+N1iso:end]
else# (dhc_args[:doS20]) & (!dhc_args[:iso])
    coeff_masks20 = coeff_mask[3+Nf:end]
end

function augment_weights(inpwt::Array{Float64})
    if (dhc_args[:doS20]) & (dhc_args[:iso]) #wt (Nc) -> |S2_iso| -> |Nf^2|
        w_s2iso = zeros(size(filter_hash["S2_iso_mat"])[1])
        w_s2iso[coeff_masks20] .= inpwt
        w_nf2 = zeros(Nf^2)
        w_nf2 .= w_s2iso[iso2nf2mask]
        return reshape(w_nf2, (Nf, Nf))
    else#if (dhc_args[:doS20]) & (!dhc_args[:iso])
        w_nf2 = zeros(Nf^2)
        w_nf2[coeff_masks20] .= inpwt
        return reshape(w_nf2, (Nf, Nf))
    end
end

function loss_func20(img_curr::Array{Float64, 2})
    #println(typeof(DHC_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]))
    s_curr = DHC_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask] #Edit
    regterm =  0.5*lambda*sum((adaptive_apodizer(img_curr) - adaptive_apodizer(input)).^2)
    println(size(s_curr), size(fs_targ_mean), size(fs_targ_invcov))
    lnlik = ( 0.5 .* (s_curr - fs_targ_mean)' * fs_targ_invcov * (s_curr - fs_targ_mean))
    neglogloss = lnlik[1] + regterm
    return neglogloss
end
#TODO: Need to replace the deriv_sums20 with a deriv_sum wrapper that handles all cases (S12, S20, S2)

pixmask=falses(Nx, Nx)
function dloss20(storage_grad::Array{Float64, 2}, img_curr::Array{Float64, 2})
    s_curr = DHC_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]
    diff = s_curr - fs_targ_mean
    #Add branches here
    wt1 = convert(Array{Float64,1}, reshape(transpose(diff) * fs_targ_invcov, (length(diff),)))
    println(size(wt1), typeof(wt1))
    wt = wt1# .* fill(1.0, length(diff))#(1.0./s_curr)
    wtaug = augment_weights(wt)
    #println(wt - wt1)
    apdimg_curr = adaptive_apodizer(img_curr)
    storage_grad .= reshape((Deriv_Utils_New.wst_S20_deriv_sum(apdimg_curr, filter_hash, wtaug, FFTthreads=1)' + reshape(lambda.*(apdimg_curr - adaptive_apodizer(input)), (1, Nx^2))) * dA, (Nx, Nx))
    storage_grad[pixmask] .= 0 # better way to do this by taking pixmask as an argument wst_s2_deriv_sum?
    return
end

println("Diff check")
eps = zeros(size(input))
row, col = 24, 18 #convert(Int8, Nx/2), convert(Int8, Nx/2)+3
eps[row, col] = 1e-6
chisq1 = loss_func20(input+eps./2) #DEB
chisq0 = loss_func20(input-eps./2) #DEB
brute  = (chisq1-chisq0)/1e-6
#df_brute = DHC_compute(reshape(input, (Nx, Nx)), filter_hash, doS2=true, doS12=false, doS20=false, norm=false)[coeff_mask] - s_targ_mean
clever = zeros(size(input)) #DEB
meanval = mean(adaptive_apodizer(input)) ./ mean(wind_2d(Nx))
_bar = dloss20(clever, input) #DEB
println("Chisq Derve Check")
println("Brute:  ",brute)
println("Clever: ",clever[row, col], " Difference: ", brute - clever[row, col]) #DEB

#####################################
##Debug: Check old code: Works with triangular mask
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
#heatmap(init)
#heatmap(true_img)

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false) #Iso #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
if dhc_args[:iso]
    coeff_mask = falses(2+S1iso+S2iso)
    coeff_mask[S1iso+3:end] .= true
else #Not iso
    coeff_mask = falses(2+Nf+Nf^2)
    coeff_maskS20 = trues((Nf, Nf))
    coeff_mask[3+Nf:end] .= triu(coeff_maskS20)[:]
end

white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :norm=>false, :smooth=>false, :smoothval=>0.8) #Iso #only if you're using noise based covar
optim_settings = Dict([("iterations", 10), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true),("TransformedGaussianLoss", false), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.00456)]) #Add constraints

dbnocffs = get_dbn_coeffs(log.(im), filter_hash, dhc_args, coeff_mask = coeff_mask) #removed log from here
fmean = mean(dbnocffs, dims=1)
fcov = (dbnocffs .- fmean)' * (dbnocffs .- fmean) ./(size(dbnocffs)[1] -1)
fcovinv = invert_covmat(fcov, 1e-5)

recon_settings["s_targ_mean"] = fmean[:]
recon_settings["s_invcov"] = fcovinv
recon_settings["safemode"] = true
fname_save = "scratch_NM/NewWrapper/TransfGaussian/nof_test"
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings

resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(SFD) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

##Log-Log check: Works

Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

noise = rand(Normal(0.0, std(true_img)), Nx, Nx)
init = true_img + noise
#heatmap(init)
#heatmap(true_img)

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false) #Iso #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
if dhc_args[:iso]
    coeff_mask = falses(2+S1iso+S2iso)
    coeff_mask[S1iso+3:end] .= true
else #Not iso
    coeff_mask = falses(2+Nf+Nf^2)
    coeff_maskS20 = trues((Nf, Nf))
    coeff_mask[3+Nf:end] .= triu(coeff_maskS20)[:]
end

white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :norm=>false, :smooth=>false, :smoothval=>0.8) #Iso #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", false),("TransformedGaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(apodizer(log.(init)))^(-2) )]) #Add constraints

dbnocffs = log.(get_dbn_coeffs(log.(im), filter_hash, dhc_args, coeff_mask = coeff_mask)) #removed log from here
fmean = mean(dbnocffs, dims=1)
fcov = (dbnocffs .- fmean)' * (dbnocffs .- fmean) ./(size(dbnocffs)[1] -1)
fcovinv = invert_covmat(fcov, 1e-6)

recon_settings["transform_func"] = Data_Utils.fnlog
recon_settings["transform_dfunc"] = Data_Utils.fndlog
recon_settings["fs_targ_mean"] = fmean[:]
recon_settings["fs_invcov"] = fcovinv
recon_settings["safemode"] = false
fname_save = "scratch_NM/NewWrapper/TransfGaussian/nof_test"
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings

resobj, recon_img = reconstruction_wrapper(true_img, init, filter_hash, dhc_args, coeff_mask, recon_settings)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(SFD) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
using Plots

using simulator.materials.spesifics
using simulator.targetfigures
using simulator.analyticalmaterials
using simulator.ema

# Simulation parameters start

save = true

n1 = SiO2_core_Sellmeier
n2 = Pd000S5
n3 = Pd053S6
n4 = Air
n4_ema(λ) = n_Bruggeman(ϵ(n3(λ)), ϵ(Air(λ)), 0.5)
n5 = Air

d1 = 10e-9
d2 = 30e-9
d3 = 10e-9
d4 = 0e-9
d5 = 10e-9

d2smax = 40e-9
d3smax = 40e-9
d4smax = 5e-9
d2ss = 1e-9
d3ss = 1e-10
d4ss = 1e-10
d3smrd2sm = 0.4
d4smrd3sm = 0.1

λ = 1200e-9
λs = [x for x in 800e-9:1e-10:1400e-9]

Hcs = [x for x in 1e-2:1e-2:1e-1]
Pd_c = 0.52

n3sd = [Pd000S5, Pd041S21, Pd051S22, Pd053S6, Pd061S19, Pd064S7and18, Pd069S24, Pd070S9to17, Pd074S8drude, Pd100S20]
n3sh = [Pd034_unloaded, Pd042_unloaded, Pd052_unloaded, Pd073_unloaded, Pd_unloaded]
Pd_cs_d = [0.00, 0.41, 0.51, 0.53, 0.61, 0.64, 0.69, 0.70, 0.74, 1.00]
Pd_cs_h = [0.34, 0.42, 0.52, 0.73, 1.00]
Pd_d3s = [10e-9, 10e-9, 10e-9, 10e-9, 10e-9]

θd = 45
θds = [x for x in 40:0.001:60]

fig_path = "../rapport/figures/thesisfigures_staged/"

# Simulation parameters end

if save
    gr()
else
    plotly()
end
θ1 = θd * π / 180
θs = θds .* (π / 180)

# Run thickness simulations
println("\nScan d3:")
@time begin
    system_d3 = make_d3_system(n1, n2, n3, n4, n5, d1, d2, d4, d5, θ1)
    d3_xs, d3_ysr, d3_ysλ, d3_ysw, d3_ysΔ = scan_plasmon_singleparameter(system_d3, [x for x in 0:d3ss:min(d3smax, d2*d3smrd2sm)], λs, 2)
end
d3_xsnm = d3_xs *= 1e9
fig_d3r = plot(d3_xsnm, d3_ysr, title="Plasmon reflection minimum", xaxis="Palladium thickness [nm]", yaxis="Power reflection coefficient minimum", legend=false, tiks=:native)
fig_d3λ = plot(d3_xsnm, d3_ysλ .* 1e9, title="Plasmon minimum wavelength", xaxis="Palladium thickness [nm]", yaxis="Dip wavelength [nm]", legend=false, tiks=:native)
fig_d3w = plot(d3_xsnm, d3_ysw .* 1e9, title="Plasmon dip full width half maximum", xaxis="Palladium thickness [nm]", yaxis="Dip full width half maximum [nm]", legend=false, tiks=:native)
fig_d3Δ = plot(d3_xsnm, d3_ysΔ, title="Plasmon dip peak to peak", xaxis="Palladium thickness [nm]", yaxis="Dip peak to peak", legend=false, tiks=:native)
if save
    savefig(fig_d3r, fig_path * "d3r.pdf")
    savefig(fig_d3λ, fig_path * "d3l.pdf")
    savefig(fig_d3w, fig_path * "d3w.pdf")
    savefig(fig_d3Δ, fig_path * "d3d.pdf")
else
    display(fig_d3r)
    display(fig_d3λ)
    display(fig_d3w)
    display(fig_d3Δ)
end

println("\nScan d2 and d3:")
@time begin
    system_d23 = make_d2_dualparameter_system(n1, n2, n3, n4, n5, d1, d4, d5, θ1)
    d23_xs, d23_ysr, d23_ysλ, d23_ysw, d23_ysΔ = scan_plasmon_dualthicknesses(system_d23, 0, d2smax, d2ss, 0, d3smax, d2ss, λs, 1, d3smrd2sm)
end
d23_xsnm = d23_xs .* 1e9
fig_d23r = plot(d23_xsnm, d23_ysr .* 1e9, title="Minimum plasmon dip", xaxis="Gold thickness [nm]", yaxis="Palladium thickness [nm]", legend=false, tiks=:native)
fig_d23w = plot(d23_xsnm, d23_ysw .* 1e9, title="Minimum dip width", xaxis="Gold thickness [nm]", yaxis="Palladium thickness [nm]", legend=false, tiks=:native)
fig_d23Δ = plot(d23_xsnm, d23_ysΔ .* 1e9, title="Maximum dip peak to peak", xaxis="Gold thickness [nm]", yaxis="Palladium thickness [nm]", legend=false, tiks=:native)
if save
    savefig(fig_d23r, fig_path * "d23r.pdf")
    savefig(fig_d23w, fig_path * "d23w.pdf")
    savefig(fig_d23Δ, fig_path * "d23d.pdf")
else
    display(fig_d23r)
    display(fig_d23w)
    display(fig_d23Δ)
end

# Run demo of structure with and without gold
println("\nSingle plasmon with without gold")
@time begin
    system_with = make_layered_tm_system([n1, n2, n3, n4, n5], [d1, 30e-9, 10e-9, d4, d5], θ1)
    system_without = make_layered_tm_system([n1, n2, n3, n4, n5], [d1, 0e-9, 10e-9, d4, d5], θ1)
    with_xs, with_ys = scann_singleparameter(system_with, λs)
    without_xs, without_ys = scann_singleparameter(system_without, λs)

    system_with_θ = make_layered_θ_tm_system([n1, n2, n3, n4, n5], [d1, 30e-9, 10e-9, d4, d5], λ)
    system_without_θ = make_layered_θ_tm_system([n1, n2, n3, n4, n5], [d1, 0e-9, 10e-9, d4, d5], λ)
    with_xsθ, with_ysθ = scann_singleparameter(system_with_θ, θs)
    without_xsθ, without_ysθ = scann_singleparameter(system_without_θ, θs)
end
fig_with = plot(with_xs .* 1e9, with_ys, title="Plasmon structure with and without gold", xaxis="Incident wavelength [nm]", yaxis="Reflection coefficient", label="10 nm palladium on 30 nm gold", legend=:bottomright, tiks=:native)
plot!(fig_with, without_xs .* 1e9, without_ys, label="10 nm palladium")
fig_with_θ = plot(θds, with_ysθ, title="Plasmon structure with and without gold", xaxis="Incident angle [degrees]", yaxis="Reflection coefficient", label="10 nm palladium on 30 nm gold", legend=:bottomright, tiks=:native)
plot!(fig_with_θ, θds, without_ysθ, label="10 nm palladium")
if save
    savefig(fig_with, fig_path * "withwithout.pdf")
    savefig(fig_with_θ, fig_path * "withwithouttheta.pdf")
else
    display(fig_with)
    display(fig_with_θ)
end

# Run hydrogen simulation
println("\nScan hydrogen:")
@time begin
    system_Hc = make_hc_system(n1, n2, n3, n4, n5, d1, d2, d3, d4, d5, θ1, Pd_c)
    Hc_xs, Hc_ysr, Hc_ysλ, Hc_ysw, Hc_ysΔ = scan_plasmon_singleparameter(system_Hc, Hcs, λs, 2)
end
Hc_xsp = Hc_xs .* 1e2
fig_Hcr = plot(Hc_xsp, Hc_ysr, title="Plasmon reflection minimum", xaxis="Hydrogen concentration [%]", yaxis="Power reflection coefficient minimum", legend=false, tiks=:native)
fig_Hcλ = plot(Hc_xsp, Hc_ysλ .* 1e9, title="Plasmon minimum wavelength", xaxis="Hydrogen concentration [%]", yaxis="Dip wavelength [nm]", legend=false, tiks=:native)
fig_Hcw = plot(Hc_xsp, Hc_ysw .* 1e9, title="Plasmon dip full width half maximum", xaxis="Hydrogen concentration [%]", yaxis="Dip full width half maximum [nm]", legend=false, tiks=:native)
fig_HcΔ = plot(Hc_xsp, Hc_ysΔ, title="Plasmon dip peak to peak", xaxis="Hydrogen concentration [%]", yaxis="Dip peak to peak", legend=false, tiks=:native)
if save
    savefig(fig_Hcr, fig_path * "Hcr.pdf")
    savefig(fig_Hcλ, fig_path * "Hcl.pdf")
    savefig(fig_Hcw, fig_path * "Hcw.pdf")
    savefig(fig_HcΔ, fig_path * "Hcd.pdf")
else
    display(fig_Hcr)
    display(fig_Hcλ)
    display(fig_Hcw)
    display(fig_HcΔ)
end

# Run alloy comparisons
println("\nScan alloys d3:")
fig_alloys_d3r = plot(title="Plasmon reflection minimum", xaxis="Palladium thickness [nm]", yaxis="Power reflection coefficient minimum", legend=:bottomright, tiks=:native)
fig_alloys_d3λ = plot(title="Plasmon minimum wavelength", xaxis="Palladium thickness [nm]", yaxis="Dip wavelength [nm]", legend=:bottomleft, tiks=:native)
fig_alloys_d3w = plot(title="Plasmon dip full width half maximum", xaxis="Palladium thickness [nm]", yaxis="Dip full width half maximum [nm]", legend=:bottomright, tiks=:native)
fig_alloys_d3Δ = plot(title="Plasmon dip peak to peak", xaxis="Palladium thickness [nm]", yaxis="Dip peak to peak", legend=:topright, tiks=:native)
@time begin
    for (n3_local, Pd_c_local) in zip(n3sd, Pd_cs_d)
        @show Pd_c_local
        system_alloysd3 = make_d3_system(n1, n2, n3_local, n4, n5, d1, d2, d4, d5, θ1)
        alloysd3_xs, alloysd3_ysr, alloysd3_ysλ, alloysd3_ysw, alloysd3_ysΔ = scan_plasmon_singleparameter(system_alloysd3, [x for x in 0:d3ss:min(d3smax, d2*d3smrd2sm)], λs, 2)

        alloysd3_xsnm = alloysd3_xs .* 1e9
        plot!(fig_alloys_d3r, alloysd3_xsnm, alloysd3_ysr, label=(string(Pd_c_local) * "% palladium"))
        plot!(fig_alloys_d3λ, alloysd3_xsnm, alloysd3_ysλ .* 1e9 , label=(string(Pd_c_local) * "% palladium"))
        plot!(fig_alloys_d3w, alloysd3_xsnm, alloysd3_ysw .* 1e9 , label=(string(Pd_c_local) * "% palladium"))
        plot!(fig_alloys_d3Δ, alloysd3_xsnm, alloysd3_ysΔ, label=(string(Pd_c_local) * "% palladium"))
    end
end
if save
    savefig(fig_alloys_d3r, fig_path * "alloys_d3r.pdf")
    savefig(fig_alloys_d3λ, fig_path * "alloys_d3l.pdf")
    savefig(fig_alloys_d3w, fig_path * "alloys_d3w.pdf")
    savefig(fig_alloys_d3Δ, fig_path * "alloys_d3d.pdf")
else
    display(fig_alloys_d3r)
    display(fig_alloys_d3λ)
    display(fig_alloys_d3w)
    display(fig_alloys_d3Δ)
end

println("\nSingle plasmon in d3 scan")
@time begin
    system_d3_dip = make_layered_tm_system([n1, n2, Pd041S21, n4, n5], [d1, d2, 5e-9, d4, d5], θ1)
    system_d3_nodip = make_layered_tm_system([n1, n2, Pd041S21, n4, n5], [d1, d2, 10e-9, d4, d5], θ1)
    d3dip_xs, d3dip_ys = scann_singleparameter(system_d3_dip, λs)
    d3nodip_xs, d3nodip_ys = scann_singleparameter(system_d3_nodip, λs)

    system_d3_dip_θ = make_layered_θ_tm_system([n1, n2, Pd041S21, n4, n5], [d1, d2, 5e-9, d4, d5], λ)
    system_d3_nodip_θ = make_layered_θ_tm_system([n1, n2, Pd041S21, n4, n5], [d1, d2, 10e-9, d4, d5], λ)
    d3dip_xsθ, d3dip_ysθ = scann_singleparameter(system_d3_dip_θ, θs)
    d3nodip_xsθ, d3nodip_ysθ = scann_singleparameter(system_d3_nodip_θ, θs)
end
fig_d3dip = plot(d3dip_xs .* 1e9, d3dip_ys, title="Reflection spectra for 41% palladium gold structure", xaxis="Incident wavelength [nm]", yaxis="Reflection coefficient", label="5 nm palladium alloy", legend=:bottomright, tiks=:native)
plot!(fig_d3dip, d3nodip_xs .* 1e9, d3nodip_ys, label="10 nm palladium alloy")
fig_d3dip_θ = plot(θds, d3dip_ysθ, title="Angular reflection spectra for 41% palladium gold structure", xaxis="Incident angle [degrees]", yaxis="Reflection coefficient", label="5 nm palladium alloy", legend=:bottomright, tiks=:native)
plot!(fig_d3dip_θ, θds, d3nodip_ysθ, label="10 nm palladium alloy")
if save
    savefig(fig_d3dip, fig_path * "d3dip.pdf")
    savefig(fig_d3dip_θ, fig_path * "d3diptheta.pdf")
else
    display(fig_d3dip)
    display(fig_d3dip_θ)
end

println("\nScan alloys hydrogen:")
fig_alloys_Hcr = plot(title="Plasmon reflection minimum", xaxis="Hydrogen concentration [%]", yaxis="Power reflection coefficient minimum", legend=:bottomright, tiks=:native)
fig_alloys_Hcλ = plot(title="Plasmon minimum wavelength", xaxis="Hydrogen concentration [%]", yaxis="Dip wavelength [nm]", legend=:bottomright, tiks=:native)
fig_alloys_Hcw = plot(title="Plasmon dip full width half maximum", xaxis="Hydrogen concentration [%]", yaxis="Dip full width half maximum [nm]", legend=:bottomright, tiks=:native)
fig_alloys_HcΔ = plot(title="Plasmon dip peak to peak", xaxis="Hydrogen concentration [%]", yaxis="Dip peak to peak", legend=:bottomright, tiks=:native)
@time begin
    for (n3_local, Pd_c_local, d3_local) in collect(zip(n3sh, Pd_cs_h, Pd_d3s))[2:length(n3sh)]
        @show Pd_c_local
        system_alloys_Hc = make_hc_system(n1, n2, n3_local, n4, n5, d1, d2, d3_local, d4, d5, θ1, Pd_c_local)
        alloysHc_xs, alloysHc_ysr, alloysHc_ysλ, alloysHc_ysw, alloysHc_ysΔ = scan_plasmon_singleparameter(system_alloys_Hc, Hcs, λs, 2)

        alloysHc_xsp = alloysHc_xs .* 1e2
        plot!(fig_alloys_Hcr, alloysHc_xsp, alloysHc_ysr, label=(string(Pd_c_local) * "% palladium"))
        plot!(fig_alloys_Hcλ, alloysHc_xsp, alloysHc_ysλ .* 1e9 , label=(string(Pd_c_local) * "% palladium"))
        plot!(fig_alloys_Hcw, alloysHc_xsp, alloysHc_ysw .* 1e9 , label=(string(Pd_c_local) * "% palladium"))
        plot!(fig_alloys_HcΔ, alloysHc_xsp, alloysHc_ysΔ, label=(string(Pd_c_local) * "% palladium"))
    end
end
if save
    savefig(fig_alloys_Hcr, fig_path * "alloys_Hcr.pdf")
    savefig(fig_alloys_Hcλ, fig_path * "alloys_Hcl.pdf")
    savefig(fig_alloys_Hcw, fig_path * "alloys_Hcw.pdf")
    savefig(fig_alloys_HcΔ, fig_path * "alloys_Hcd.pdf")
else
    display(fig_alloys_Hcr)
    display(fig_alloys_Hcλ)
    display(fig_alloys_Hcw)
    display(fig_alloys_HcΔ)
end

# Run roughness simulations
println("\nScan d4 roughness:")
@time begin
    system_d4_ema = make_d4_system(n1, n2, n3, n4_ema, n5, d1, d2, d3, d5, θ1)
    d4_ema_xs, d4_ema_ysr, d4_ema_ysλ, d4_ema_ysw, d4_ema_ysΔ = scan_plasmon_singleparameter(system_d4_ema, [x for x in 0:d4ss:min(d4smax, d3*d4smrd3sm)], λs, 2)
end
d4_ema_xsnm = d4_ema_xs *= 1e9
fig_d4_emar = plot(d4_ema_xsnm, d4_ema_ysr, title="Plasmon reflection minimum", xaxis="Roughness ema thickness [nm]", yaxis="Power reflection coefficient minimum", legend=false, tiks=:native)
fig_d4_emaλ = plot(d4_ema_xsnm, d4_ema_ysλ .* 1e9, title="Plasmon minimum wavelength", xaxis="Roughness ema thickness [nm]", yaxis="Dip wavelength [nm]", legend=false, tiks=:native)
fig_d4_emaw = plot(d4_ema_xsnm, d4_ema_ysw .* 1e9, title="Plasmon dip full width half maximum", xaxis="Roughness ema thickness [nm]", yaxis="Dip full width half maximum [nm]", legend=false, tiks=:native)
fig_d4_emaΔ = plot(d4_ema_xsnm, d4_ema_ysΔ, title="Plasmon dip peak to peak", xaxis="Roughness ema thickness [nm]", yaxis="Dip peak to peak", legend=false, tiks=:native)
if save
    savefig(fig_d4_emar, fig_path * "d4_emar.pdf")
    savefig(fig_d4_emaλ, fig_path * "d4_emal.pdf")
    savefig(fig_d4_emaw, fig_path * "d4_emaw.pdf")
    savefig(fig_d4_emaΔ, fig_path * "d4_emad.pdf")
else
    display(fig_d4_emar)
    display(fig_d4_emaλ)
    display(fig_d4_emaw)
    display(fig_d4_emaΔ)
end

ext: jl


text:
function loop_translation(n)

  reps = parse(Int64, ENV["LAMP_REPS"])

  A = randn(n,n)
  B = randn(n,n)
  C = randn(n,n)

  Benchmarker.add_data(csv, "loop_translation_nai", Benchmarker.measure(reps, naive_loop, A, B, C))
  Benchmarker.add_data(csv, "loop_translation_rec", Benchmarker.measure(reps, recommended_loop, A, B, C))

  @test(isapprox(naive_loop(A,B,C)[2], recommended_loop(A,B,C)[2]))
end

function naive_loop( A::DenseArray{Float64, 2},  B::DenseArray{Float64, 2},  C::DenseArray{Float64, 2})
  result = copy(C)
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  for i in range(1, size(result, 2), step=1)
    @inbounds result[:, i] = A * B[:, i]
  end

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function recommended_loop( A::DenseArray{Float64, 2},  B::DenseArray{Float64, 2},  C::DenseArray{Float64, 2})
  result = copy(C)
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = A * B

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end


text:
# 10-layer (5 down, 5 up) autoencoder model (AEM)
function aem5_layers(image_size::Tuple{Int, Int})
    k = (3, 3);  # kernel size
    s = (2, 2);  # stride size
    d = (1, 1);  # dilation size

    # Size and padding for the first layer
    epad1 = pad2d(image_size, k, s);
    # Size and padding for the second layer
    size2 = nstrides2d(image_size, k, epad1, s);
    epad2 = pad2d(size2, k, s);
    # Size and padding for the third layer
    size3 = nstrides2d(size2, k, epad2, s);
    epad3 = pad2d(size3, k, s);
    # Size and padding for the fourth layer
    size4 = nstrides2d(size3, k, epad3, s);
    epad4 = pad2d(size4, k, s);
    # Size and padding for the fifth layer
    size5 = nstrides2d(size4, k, epad4, s);
    epad5 = pad2d(size5, k, s);

    autoencode = Chain(
                   # Encoder section
                   Conv(k, 3=>32, relu, pad=epad1, stride=s, dilation=d),
                   Conv(k, 32=>64, relu, pad=epad2, stride=s, dilation=d),
                   Conv(k, 64=>128, relu, pad=epad3, stride=s, dilation=d),
                   Conv(k, 128=>256, relu, pad=epad4, stride=s, dilation=d),
                   Conv(k, 256=>512, relu, pad=epad5, stride=s, dilation=d),
                   # Decoder section
                   ConvTranspose(k, 512=>256, relu, pad=epad5, stride=s, dilation=d),
                   ConvTranspose(k, 256=>128, relu, pad=epad4, stride=s, dilation=d),
                   ConvTranspose(k, 128=>64, relu, pad=epad3, stride=s, dilation=d),
                   ConvTranspose(k, 64=>32, relu, pad=epad2, stride=s, dilation=d),
                   ConvTranspose(k, 32=>3, relu, pad=epad1, stride=s, dilation=d)
                   );

    return autoencode
end

struct AEM5 <: AutoEncoderModel
    layers::Chain
end

AEM5(image_size::Tuple{Int, Int}) = AEM5(aem5_layers(image_size))

Base.show(io::IO, ::AEM5) = print(io, "AEM5()")

Flux.@functor AEM5

(m::AEM5)(x) = m.layers(x)