# pre-defined factors

@deffactor EntropyFactor [x] H(x)

@deffactor LogNormalFactor [x, μ, τ] begin
    -(1/2) * ((E(τ) * ( V(x) + V(μ) + (E(x) - E(μ))^2 ) + log(2π) - Elog(τ)))
end

@deffactor LogGammaCanonFactor [x, α, β] begin
    (E(α) - 1) * Elog(x) - E(x) * E(β) + E(α) * Elog(β) - Eloggamma(α)
end

@deffactor LogGammaFactor [x, α, θ] begin
    (E(α) - 1) * Elog(x) - E(x)/E(θ) - E(α) * Elog(θ) - Eloggamma(α)
end

@deffactor LogMvNormalCanonFactor [x, μ, Λ] begin
    δ = E(x) - E(μ)
    EΛ = E(Λ)
    -(1/2) * (trace(EΛ * (C(x) .+ C(μ) .+ δ * δ')) + length(x) * log(2π) - Elogdet(Λ))
end

@deffactor LogMvNormalDiagCanonFactor [x, μ, τ] begin
    δ = E(x) - E(μ)
    Eτ = E(τ)
    -(1/2) * (sum(Eτ .* (V(x) .+ V(μ) .+ δ.^2 )) + length(x) * log(2π) - Elogdet(τ))
end

@deffactor LogDirichletFactor [x, α] begin
    dot(E(α) - 1, Elog(x)) - ElogB(α)
end

@deffactor LogWishartFactor [X ν V] begin
    EX = E(X)
    p = size(EX, 1)
    n = E(ν)
    0.5 * (n - p - 1) * Elogdet(X) - 0.5 * n * Elogdet(V) - 0.5 * trace(Einv(V) * EX) - 0.5 * n * p * log(2) - Elogmvgamma(p, ν/2)
end

################### Dealing with HMMs in factors #####################
@deffactor LogMarkovChainFactor [z, π0, A] begin
    dot(E(z)[:, 1], Elog(π0)) + sum(C(z) .* Elog(A))
end

@deffactor LogMarkovMatrixFactor [x, A] begin
    sum((E(A) - 1) .* Elog(x) .- ElogB(A))
end


module Moora

import GreyDecision.GreyNumbers: GreyNumber
import GreyDecision.Utility: unitize, normalize, weightize, apply_columns, reverseminmax, rowmaxs

function moora(decisionMat::Array{GreyNumber{T}, 2}, weights::Array{Float64,1}, fns::Array{Function,1};  v::Float64=0.5) where {T <: Real}

    w = unitize(weights)

    nalternatives, ncriteria = size(decisionMat)

    normalizedMat = normalize(decisionMat)
    weightednormalizedMat = weightize(normalizedMat, w)

    # cmaxs = colmaxs(weightednormalizedMat)
    cmaxs = apply_columns(fns, weightednormalizedMat)
    cmins = apply_columns(reverseminmax(fns), weightednormalizedMat)

    refmat = similar(weightednormalizedMat)

    for rowind in 1:nalternatives
        if fns[rowind] == maximum
            refmat[rowind, :] .= cmaxs - weightednormalizedMat[rowind, :]
        elseif fns[rowind] == minimum
            refmat[rowind, :] .= weightednormalizedMat[rowind, :] - cmins
        else
            @warn fns[rowind]
            error("Function must be either maximize or minimize")
        end
    end

    scores = rmaxs = rowmaxs(refmat)

    bestIndex = sortperm(rmaxs) |> first

    return Dict(
        :scores => scores,
        :bestindex => bestIndex
    )
end

end # end of module Moora


include("dkk17.jl")

function rcov_quantum_filter(reps, eps, k, α=4, τ=0.1; limit1=2, limit2=1.5)
    d, n = size(reps)
    reps_pca, U = pca(reps, k)
    if k == 1
        reps_estimated_white = reps_pca
        Σ′ = ones(1, 1)
    else
        selected = cov_estimation_iterate(reps_pca, eps/n, τ, nothing, limit=round(Int, limit1*eps))
        reps_pca_selected = reps_pca[:, selected]
        Σ = cov(reps_pca_selected', corrected=false)
        reps_estimated_white = Σ^(-1/2)*reps_pca
        Σ′ = cov(reps_estimated_white')
    end
    M = k > 1 ? exp(α*(Σ′- I)/(opnorm(Σ′) - 1)) : ones(1, 1)
    M /= tr(M)
    estimated_poison_ind = k_lowest_ind(
        -[x'M*x for x in eachcol(reps_estimated_white)],
        round(Int, limit2*eps)
    )
    return .! estimated_poison_ind
end

function rcov_auto_quantum_filter(reps, eps, α=4, τ=0.1; limit1=2, limit2=1.5)
    reps_pca, U = pca(reps, 100)
    best_opnorm, best_selected, best_k = -Inf, nothing, nothing
    for k in round.(Int, range(1, sqrt(100), length=10) .^ 2)
        selected = rcov_quantum_filter(reps, eps, k, α, τ; limit1=limit1, limit2=limit2)
        Σ = cov(reps_pca[:, selected]')
        Σ′ = cov((Σ^(-1/2)*reps_pca)')
        M = exp(α*(Σ′- I)/(opnorm(Σ′) - 1))
        M /= tr(M)
        op = tr(Σ′ * M) / tr(M)
        poison_removed = sum((.! selected)[end-eps+1:end])
        @show k, op, poison_removed
        if op > best_opnorm
            best_opnorm, best_selected, best_k = op, selected, k
        end
    end
    @show best_k, best_opnorm
    return best_selected
end


function integrant_volume_to_surface_area(f)
    f_x(u) = ForwardDiff.derivative(x -> f(x, u[2]), u[1])
    f_z(u) = ForwardDiff.derivative(z -> f(u[1], z), u[2])
    u -> sqrt(f_x(u)^2 + f_z(u)^2 + 1)
end

function integrate_xonly(f, xmin, xmax)
    res = hquadrature(x -> f(SA[x, bd.zmax-bd.hb]), xmin, xmax)
    res[1]
end
function integrate_limit(f, f_zlimit; xmin, xmax, zmin = bd.zmax - bd.hb, zmax = bd.zmax)
    function hcfun(u)
        _x, _z = u
        _z > f_zlimit(_x) ? 0.0 : f(u)
    end
    a = SA[xmin, zmin]
    b = SA[xmax, zmax]
    res = @test_nowarn hcubature(hcfun, a, b, maxevals = 1_000_000)
    # println(res[2])
    res[1]
end
# Start with the middle part
integrate_middle(f = bd.fi2m; kwargs...) =
    integrate_limit(f, (x) -> bd.zmax; xmin = bd.xmin0, xmax = bd.xmax0, kwargs...)
# dziob - psi2, fi2f
integrate_front(f = bd.fi2f; kwargs...) =
    integrate_limit(f, bd.psi2; xmin = bd.xmax0, xmax = bd.xmax, kwargs...)
# rufa - psi1, fi2r
integrate_rear(f = bd.fi2r; kwargs...) =
    integrate_limit(f, bd.psi1; xmin = bd.xmin, xmax = bd.xmin0, kwargs...)

@testset "Use adaptive cubature to approximate boat mass" begin
    # kadlub
    m_m_b = integrate_middle(integrant_volume_to_surface_area(bd.fi2m))
    m_f_b = integrate_front(integrant_volume_to_surface_area(bd.fi2f))
    m_r_b = integrate_rear(integrant_volume_to_surface_area(bd.fi2r))
    snb = 2 * (m_m_b + m_f_b + m_r_b)
    m_hull = snb * bd.rhob * bd.deltab
    # @test snb ≈ bd.snb
    @test 2m_m_b ≈ bd.snbmv
    @test m_f_b ≈ bd.iinbfvv rtol = 1e-4
    @test m_r_b ≈ bd.iinbrvv rtol = 1e-4

    # poklad
    m_m_d = integrate_xonly(u -> bd.fi2m(u...), bd.xmin0, bd.xmax0)
    m_f_d = integrate_xonly(u -> bd.fi2f(u...), bd.xmax0, bd.xmax)
    m_r_d = integrate_xonly(u -> bd.fi2r(u...), bd.xmin, bd.xmin0)
    sbd = 2 * (m_m_d + m_f_d + m_r_d)
    sbde = bd.bbe * (bd.xmaxe - bd.xmine)
    m_deck_full = sbd * bd.rhod * bd.deltad
    m_deck_hole = bd.rhod * bd.deltad * sbde
    @test sbd ≈ bd.sbd
    @test sbde ≈ bd.sbde
    # podłoga
    sbdg = bd.bbg * (bd.xmaxg - bd.xming)
    m_ground = bd.rhog * bd.deltag * sbdg
    @test sbdg ≈ bd.sbdg

    m = m_hull + m_deck_full - m_deck_hole + m_ground
    @test m ≈ bd.m_B rtol = 1e-5
end

@testset "Use adaptive cubature to approximate boat static inertia sxmb" begin
    # kadlub
    sxfun(f) = u -> u[1] * f(u)
    sxfun_area(f) = sxfun(integrant_volume_to_surface_area(f))
    sx_m_b = integrate_middle(sxfun_area(bd.fi2m))
    sx_f_b = integrate_front(sxfun_area(bd.fi2f))
    sx_r_b = integrate_rear(sxfun_area(bd.fi2r))
    x1snb = 2 * (sx_m_b + sx_f_b + sx_r_b)
    sx_hull = x1snb * bd.rhob * bd.deltab
    # @test x1snb ≈ bd.x1snb
    @test 2sx_m_b ≈ bd.x1snbmv
    @test sx_f_b ≈ bd.ix1inbfvv rtol = 1e-3
    @test_broken sx_r_b ≈ bd.ix1inbrvv rtol = 1e-3

    # poklad
    sxfun2(f) = sxfun(u -> f(u[1], u[2]))
    sx_m_d = integrate_xonly(sxfun2(bd.fi2m), bd.xmin0, bd.xmax0)
    sx_f_d = integrate_xonly(sxfun2(bd.fi2f), bd.xmax0, bd.xmax)
    sx_r_d = integrate_xonly(sxfun2(bd.fi2r), bd.xmin, bd.xmin0)
    x1sbd = 2 * (sx_m_d + sx_f_d + sx_r_d)
    x1sbde = 0.5bd.bbe * (bd.xmaxe - bd.xmine)^2
    sx_deck_full = x1sbd * bd.rhod * bd.deltad
    sx_deck_hole = bd.rhod * bd.deltad * x1sbde
    @test x1sbd ≈ bd.x1sbd
    @test_broken x1sbde ≈ bd.x1sbde
    # podłoga
    x1sbdg = 0.5bd.bbg * (bd.xmaxg - bd.xming)^2
    sx_ground = bd.rhog * bd.deltag * x1sbdg
    @test_broken x1sbdg ≈ bd.x1sbdg

    sx = sx_hull + sx_deck_full - sx_deck_hole + sx_ground
    @test_broken sx ≈ bd.sxmb rtol = 1e-5
    # println((sx / bd.m_B, bd.xo))
end

@testset "Use adaptive cubature to approximate boat static inertia szmb" begin
    # kadlub
    szfun(f) = u -> u[2] * f(u)
    szfun_area(f) = szfun(integrant_volume_to_surface_area(f))
    sz_m_b = integrate_middle(szfun_area(bd.fi2m))
    sz_f_b = integrate_front(szfun_area(bd.fi2f))
    sz_r_b = integrate_rear(szfun_area(bd.fi2r))
    z1snb = 2 * (sz_m_b + sz_f_b + sz_r_b)
    sz_hull = z1snb * bd.rhob * bd.deltab
    # @test x1snb ≈ bd.x1snb
    @test 2sz_m_b ≈ bd.z1snbmnv
    @test sz_f_b ≈ bd.iz1inbfvv rtol = 8e-3
    @test sz_r_b ≈ bd.iz1inbrvv rtol = 1e-2

    # poklad
    m_m_d = integrate_xonly(u -> bd.fi2m(u...), bd.xmin0, bd.xmax0)
    m_f_d = integrate_xonly(u -> bd.fi2f(u...), bd.xmax0, bd.xmax)
    m_r_d = integrate_xonly(u -> bd.fi2r(u...), bd.xmin, bd.xmin0)
    sbd = 2 * (m_m_d + m_f_d + m_r_d)
    z1sbd = sbd * (bd.zmax - bd.hb)
    @test z1sbd ≈ bd.z1sbd
    z1sbde = bd.bbe * (bd.xmaxe - bd.xmine) * (bd.zmax - bd.hb)
    @test z1sbde ≈ bd.z1sbde
    sz_deck_full = z1sbd * bd.rhod * bd.deltad
    sz_deck_hole = bd.rhod * bd.deltad * z1sbde
    # podłoga
    z1sbdg = bd.bbg * (bd.xmaxg - bd.xming) * (bd.zmax - bd.hb + bd.zg)
    @test z1sbdg ≈ bd.z1sbdg
    sz_ground = bd.rhog * bd.deltag * z1sbdg

    sz = sz_hull + sz_deck_full - sz_deck_hole + sz_ground
    @test sz ≈ bd.szmb rtol = 5e-3
    # println((sz / bd.m_B, bd.zo))
end

@testset "Use adaptive cubature to approximate boat static inertia j_By" begin
    # kadlub
    x2fun(f) = u -> u[1]^2 * f(u)
    z2fun(f) = u -> u[2]^2 * f(u)
    x2fun_area(f) = x2fun(integrant_volume_to_surface_area(f))
    z2fun_area(f) = z2fun(integrant_volume_to_surface_area(f))
    x2_m_b = integrate_middle(x2fun_area(bd.fi2m))
    x2_f_b = integrate_front(x2fun_area(bd.fi2f))
    x2_r_b = integrate_rear(x2fun_area(bd.fi2r))
    z2_m_b = integrate_middle(z2fun_area(bd.fi2m))
    z2_f_b = integrate_front(z2fun_area(bd.fi2f))
    z2_r_b = integrate_rear(z2fun_area(bd.fi2r))
    x2snb = 2 * (x2_m_b + x2_f_b + x2_r_b)
    z2snb = 2 * (z2_m_b + z2_f_b + z2_r_b)
    J_hull = (x2snb + z2snb) * bd.rhob * bd.deltab
    @test 2x2_m_b ≈ bd.x2snbmv
    @test x2_f_b ≈ bd.ix2inbfvv rtol = 4e-4
    @test_broken x2_r_b ≈ bd.ix2inbrvv rtol = 1e-2

    @test 2z2_m_b ≈ bd.z2snbmnv
    @test z2_f_b ≈ bd.iz2inbfvv rtol = 1e-2
    @test z2_r_b ≈ bd.iz2inbrvv rtol = 1e-2

    # poklad
    m_m_d = integrate_xonly(u -> bd.fi2m(u...), bd.xmin0, bd.xmax0)
    m_f_d = integrate_xonly(u -> bd.fi2f(u...), bd.xmax0, bd.xmax)
    m_r_d = integrate_xonly(u -> bd.fi2r(u...), bd.xmin, bd.xmin0)
    sbd = 2 * (m_m_d + m_f_d + m_r_d)
    z2sbd = sbd * (bd.zmax - bd.hb)^2
    @test z2sbd ≈ bd.z2sbd
    x2_m_d = integrate_xonly(u -> u[1]^2 * bd.fi2m(u...), bd.xmin0, bd.xmax0)
    x2_f_d = integrate_xonly(u -> u[1]^2 * bd.fi2f(u...), bd.xmax0, bd.xmax)
    x2_r_d = integrate_xonly(u -> u[1]^2 * bd.fi2r(u...), bd.xmin, bd.xmin0)
    x2sbd = 2 * (x2_m_d + x2_f_d + x2_r_d)
    @test x2sbd ≈ bd.x2sbd
    z2sbde = bd.bbe * (bd.xmaxe - bd.xmine)^2 * (bd.zmax - bd.hb)
    @test z2sbde ≈ bd.z2sbde
    x2sbde = bd.bbe * (bd.xmaxe^3 - bd.xmine^3) / 3
    @test x2sbde ≈ bd.x2sbde
    J_deck_full = (x2sbd + z2sbd) * bd.rhod * bd.deltad
    J_deck_hole = bd.rhod * bd.deltad * (x2sbde + z2sbde)
    # podłoga
    x2sbdg = bd.bbg * (bd.xmaxg^3 - bd.xming^3) / 3
    z2sbdg = bd.bbg * (bd.xmaxg - bd.xming) * (bd.zmax - bd.hb + bd.zg)^2
    @test x2sbdg ≈ bd.x2sbdg
    @test z2sbdg ≈ bd.z2sbdg
    J_ground = bd.rhog * bd.deltag * (x2sbdg + z2sbdg)

    J_st = bd.m_B * (bd.xo^2 + bd.zo^2)
    J = J_hull + J_deck_full - J_deck_hole + J_ground - J_st
    @test_broken J ≈ bd.J_By rtol = 5e-3
    println((J, bd.J_By))
end

@testset "Check if approximation of vh2o is correct" begin
    tol = eps()
    @test abs(bd.vh2o(0.0, bd.zmax - tol)) < tol
    # vh2o(0.0, 0.0) - should be the full boat volume!
    V_m_b = integrate_middle(u -> bd.fi2m(u...))
    V_f_b = integrate_front(u -> bd.fi2f(u...))
    V_r_b = integrate_rear(u -> bd.fi2r(u...))
    @test bd.vh2o(0.0, 0.0) ≈ 2 * (V_m_b + V_f_b + V_r_b) rtol = 1e-6
    # compute for some bh2o value
    bh2o = 0.111245
    V_m_b = integrate_middle(u -> bd.fi2m(u...), zmin = bh2o)
    V_f_b = integrate_front(u -> bd.fi2f(u...), zmin = bh2o)
    V_r_b = integrate_rear(u -> bd.fi2r(u...), zmin = bh2o)
    @test bd.vh2o(0.0, bh2o) ≈ 2 * (V_m_b + V_f_b + V_r_b) rtol = 1e-6
    # the = 0.003
    # R = SA[cos(the) -sin(the); sin(the) cos(the)]
    # V_m_b = integrate_middle(u -> bd.fi2m((R'*u)...), zmin = bh2o)
    # @test bd.vh2om(the, bh2o) ≈ 2 * (V_m_b) rtol = 1e-6
    # V_f_b = integrate_front(u -> bd.fi2f(u...), zmin = bh2o)
    # V_r_b = integrate_rear(u -> bd.fi2r(u...), zmin = bh2o)
    # @test bd.vh2o(the, bh2o) ≈ 2 * (V_m_b + V_f_b + V_r_b) rtol = 1e-6
end


using ThreeBodyDecay
using Test
using Parameters

@testset "BasicParityRecoupling" begin
    reaction = jp"1/2+"=>(jp"1/2+",jp"0-")

    H_pc = ParityRecoupling(1,0, jp"1/2+"=>(jp"1/2+",jp"0-"))
    @test H_pc == ParityRecoupling(1,0,'-')

    H_ls = RecouplingLS(possible_ls(reaction)[1] |> x2, reaction)

    factor = -1 / sqrt(2)
    @test amplitude(H_pc, 1, 0) * factor ≈ amplitude(H_ls, 1, 0)
    @test amplitude(H_pc,-1, 0) * factor ≈ amplitude(H_ls,-1, 0)
end

@testset "DecayChainCouplings" begin
    mΛb = 5.61960
    mπ = 0.14
    # intermediate resonances
    mΣb = 5.81065; ΓΣb = 0.005
    mΣb_x = 5.83032; ΓΣb_x = 0.0094

    mΛb2S = 6.3; # just a peak of the plot

    tbs = ThreeBodySystem(mπ,mΛb,mπ; m0=mΛb2S,
        two_js = ThreeBodySpins(0, 1, 0; two_h0=1))  # (0- 1/2+ 0- | 1/2+)

    dc_pc = DecayChain(
        k=1,
        Xlineshape=σ->BW(σ,mΣb,ΓΣb),
        tbs=tbs, two_s = 1,
        Hij=ParityRecoupling(1,0,'-'),
        HRk=ParityRecoupling(1,0,'-'))
    #
    σs = randomPoint(tbs.ms)
    #
    @test amplitude(σs, [0, 1,0, 1], dc_pc) != 0
    @test amplitude(σs, [0, 1,0,-1], dc_pc) != 0
    @test amplitude(σs, [0,-1,0, 1], dc_pc) != 0
    @test amplitude(σs, [0,-1,0,-1], dc_pc) != 0
    #
    dc_pv = DecayChain(
        k=1,
        Xlineshape=σ->BW(σ,mΣb,ΓΣb),
        tbs=tbs, two_s = 1,
        Hij=ParityRecoupling(1,0,'-'),
        HRk=NoRecoupling(1,0))
    #
    @test amplitude(σs, [0, 1,0, 1], dc_pv) != 0
    @test amplitude(σs, [0,-1,0, 1], dc_pv) != 0
    @test amplitude(σs, [0,-1,0,-1], dc_pv) == 0im
    @test amplitude(σs, [0,-1,0,-1], dc_pv) == 0im
    #
    dc_pv = DecayChain(
        k=1,
        Xlineshape=σ->BW(σ,mΣb,ΓΣb),
        tbs=tbs, two_s = 1,
        Hij=ParityRecoupling(1,0,'-'),
        HRk=NoRecoupling(1,0))
    #
    @test amplitude(σs, [0, 1,0, 1], dc_pv) != 0
    @test amplitude(σs, [0,-1,0, 1], dc_pv) != 0
    @test amplitude(σs, [0,-1,0,-1], dc_pv) == 0im
    @test amplitude(σs, [0,-1,0,-1], dc_pv) == 0im
end


#
# This module represents an explicit, direct sparse 2D finite-difference Poisson solver for heterogeneous media,
# i.e. media having spatially varying (space-dependent) medium parameters.
# Current implementation assumes Neumann boundary conditions at all the boundaries.
#
# Consider the following Poisson experiment:
# ```math
# ∇⋅(σ(x,z)∇) ψ(t) = ∇⋅(Q(x,z)∇) p(t),
# ```
# ```math
# Q = k * Q_v / η.
# ```


include("CalcNumCTEC.jl")

# f(x,y) = 2x*y + 2x - x^2 - 2y^2 + 72
# # f(x,y) = 1 + x + y

# a = 0.
# b = 8.

# c = 0.
# d = 6.

# t0 = time()
# I_ap = IntegralDupla_Trapezio(f, (a,b), (c,d))
# t1 = time()
# I_ex = 2816.

# println("Integral aproximada: ", I_ap)
# println("Integral exata.....: ", I_ex)
# println("Erro absoluto......: ", abs(I_ap - I_ex))
# println("Tempo de execução..: ", round(t1-t0, digits=4), "s")

v1 = Vertex2d(1,2)
println(v)


const SMALL_K = 0.075

abstract type AbstractStructureFactor <: Function end

struct StructureFactor{L, S} <: AbstractStructureFactor
    liquid::L
    scheme::S
end

function StructureFactor(liquid::L, T::Type; kw...) where {L<:AbstractLiquid}
    scheme = T(liquid; kw...)
    return StructureFactor{L, typeof(scheme)}(liquid, scheme)
end

(S::StructureFactor)(k) = structure_factor(S.liquid, S.scheme, k)

"""
    structure_factor(liquid, scheme, k)

Returns the static structure factor of a `liquid` using the approximation
defined by `scheme` at the wavenumber `k`.
"""
function structure_factor(liquid, scheme, k)
    Ck = Ĉ(liquid, scheme, k)
    return 1 / (1 - Ck)
end

function structure_factor(liquid::DipolarHardSpheres, scheme::MSA, k)
    C₀₀k, C₁₀k, C₁₁k = Ĉ(liquid, scheme, k)

    S₀₀k = 1 / (1 - C₀₀k)
    S₁₀k = 1 / (1 - C₁₀k)
    S₁₁k = 1 / (1 - C₁₁k)

    return (S₀₀k, S₁₀k, S₁₁k)
end

function Ĉ(liquid::HardDisks, scheme::RosenfeldFMT, k′)
    η = volume_fraction(liquid)
    σ = length_scale(potential(liquid))

    k  = σ * k′
    k² = k * k
    J₀ = besselj0(k / 2)
    J₁ = besselj1(k / 2)
    J₁′ = besselj1(k)

    A, B, G = scheme.A, scheme.B, scheme.G

    smallk = k < SMALL_K # for small wavenumber use a Taylor series approximation

    C = smallk ? (A * (1 -  k² / 16) / 4 + B * (1 - 3k² / 32) / 2 +
                  G * (1 -  k² / 8 )) :
                 (A * (2J₁ / k)^2 + B * 2J₀ * J₁ / k + G * 2J₁′ / k)

    return -4η * C
end

"""
    Ĉ(liquid::HardSpheres, scheme::PercusYevick, k)

Returns the product of the bulk density and the Fourier transform of the direct
correlation function for a hard-spheres liquid using the Percus-Yevick
approximation.

`k` is the wavenumber.
"""
function Ĉ(liquid::HardSpheres, scheme::PercusYevick, k′)
    η = volume_fraction(liquid)
    σ = length_scale(potential(liquid))

    k  = σ * k′
    k² = k * k
    k³ = k² * k
    k⁴ = k² * k²
    k⁶ = k³ * k³
    sink, cosk = sin(k), cos(k)

    α, β, δ = scheme.α, scheme.β, scheme.δ

    smallk = k < SMALL_K # for small wavenumber use a Taylor series approximation

    C₀ = smallk ? (1 - k² / 10) / 3 :
                  (sink - k * cosk) / k³
    C₁ = smallk ? (1 - k² / 9 ) / 4 :
                  (2k * sink - (k² - 2) * cosk - 2) / k⁴
    C₃ = smallk ? (1 - k² / 8 ) / 6 :
                  ((4k³ - 24k) * sink - (k⁴ - 12k² + 24) * cosk + 24) / k⁶

    return 24η * (α * C₀ + β * C₁ + δ * C₃)
end

"""
    Ĉ(::HardSpheres, ::VerletWeis, k)

Returns the product of the bulk density and the Fourier transform of the direct
correlation function for a hard-spheres liquid using the Percus-Yevick
approximation with the Verlet-Weis correction.

`k` is the wavenumber.
"""
Ĉ(liquid::HardSpheres, scheme::VerletWeis, k) =
    Ĉ(scheme.coreliquid, scheme.subscheme, scheme.α * k)

"""
    Ĉ(::AttractiveHardSpheres{<:YHSPotential}, ::SharmaSharma, k)

Returns the product of the bulk density and the Fourier transform (FT) of the
direct correlation function for an attractive Yukawa hard-spheres liquid using
the Sharma and Sharma approximation, that is, the sum of the FT of the direct
correlation function for the hard spheres core and ``\\beta\\cdot\\hat{u}(k)``,
where ``\\hat{u}(k)`` is the FT of the Yukawa potential:

```math
    u(r) = -\\epsilon\\exp(-zr/\\sigma) / (r/\\sigma)
```

`k` is the wavenumber.
"""
function Ĉ(liquid::AttractiveHardSpheres{P, U}, scheme::SharmaSharma, k′) where
    {P, U<:YHSPotential}

    η = volume_fraction(liquid)
    T = temperature(liquid)
    σ, z = length_scale(potential(liquid))

    k  = σ * k′
    k² = k * k
    k⁴ = k² * k²
    sink, cosk = sin(k), cos(k)

    smallk = k < SMALL_K # for small wavenumber use a Taylor series approximation

    C = smallk ? ((1 + z) - (3 + z) / 6 * k² + (5 + z) / 120 * k⁴) :
                  (k * cosk + z * sink) / k
    C = C / (k² + z * z)

    C₀ = Ĉ(scheme.coreliquid, scheme.subscheme, k′)

    return C₀ + 24η / T * C
end

"""
    Ĉ(::AttractiveHardSpheres{<:SWHSPotential}, ::SharmaSharma, k)

Returns the product of the bulk density and the Fourier transform (FT) of the
direct correlation function for an attractive square-well hard-spheres liquid
using the Sharma and Sharma approximation, that is, the sum of the FT of the
direct correlation function for the hard spheres core and
``\\beta\\cdot\\hat{u}(k)``, where ``\\hat{u}(k)`` is the FT of the square well
potential:

```math
    u(r) = -\\epsilon\\sigma, \\qquad \\sigma < r \\le \\lambda\\sigma
```

`k` is the wavenumber.
"""
function Ĉ(liquid::AttractiveHardSpheres{P, U}, scheme::SharmaSharma, k′) where
    {P, U<:SWHSPotential}

    η = volume_fraction(liquid)
    T = temperature(liquid)
    σ, λ = length_scale(potential(liquid))

    λ³ = λ * λ * λ
    λ⁵ = λ³ * λ * λ

    k  = σ * k′
    k² = k * k
    k³ = k² * k
    sink, cosk = sin(k), cos(k)
    sinλk, cosλk = sin(λ * k), cos(λ * k)

    smallk = k < SMALL_K # for small wavenumber use a Taylor series approximation

    C = smallk ? (λ³ - 1) / 3 - (λ⁵ - 1) / 30 * k² :
                 (cosk - λ * cosλk) / k² + (sinλk - sink) / k³

    C₀ = Ĉ(scheme.coreliquid, scheme.subscheme, k′)

    return C₀ + 24η / T * C
end

"""
    Ĉ(::AttractiveHardSpheres{<:SWHSPotential}, ::NonLinearSharmaSharma, k)

Returns the product of the bulk density and the Fourier transform (FT) of the
direct correlation function for an attractive square-well hard-spheres liquid
using the Non-Linear Sharma and Sharma approximation, that is, the sum of the FT of the
direct correlation function for the hard spheres core and
``exp[\\beta\\cdot\\u(r)]-1``, where ``u(r)`` is the potential:

```math
    u(r) = -\\epsilon\\sigma, \\qquad \\sigma < r \\le \\lambda\\sigma
```

`k` is the wavenumber.
"""
function Ĉ(liquid::AttractiveHardSpheres{P, U}, scheme::NonLinearSharmaSharma, k′) where
    {P, U<:SWHSPotential}

    η = volume_fraction(liquid)
    T = temperature(liquid)
    σ, λ = length_scale(potential(liquid))
    x = exp(1 / T) - 1

    λ³ = λ * λ * λ
    λ⁵ = λ³ * λ * λ

    k = σ * k′
    k² = k * k
    k³ = k² * k
    sink, cosk = sin(k), cos(k)
    sinλk, cosλk = sin(λ * k), cos(λ * k)

    smallk = k < 0.075

    C = smallk ? (λ³ - 1) / 3 - (λ⁵ - 1) / 30 * k² :
                 (cosk - λ * cosλk) / k² + (sinλk - sink) / k³

    C₀ = Ĉ(scheme.coreliquid, scheme.subscheme, k′)

    return C₀ + 24η * x * C
end

"""
    Ĉ(liquid::DipolarHardSpheres, scheme::MSA, k)

Returns the product of the bulk density and the Fourier transform of
projections of the direct correlation function `(C₀₀, C₁₀, C₁₁)` for a
dipolar hard-spheres liquid using the MSA approximation.

`k` is the wavenumber.
"""
function Ĉ(liquid::DipolarHardSpheres, scheme::MSA, k′)
    η = volume_fraction(liquid)
    σ = length_scale(potential(liquid))

    k  = σ * k′
    k² = k * k
    k⁶ = k² * k² * k²
    sink, cosk = sin(k), cos(k)

    α₀, α₁, α₂, α₃ = scheme.α₀, scheme.α₁, scheme.α₂, scheme.α₃
    β₀, β₁, β₂, β₃ = scheme.β₀, scheme.β₁, scheme.β₂, scheme.β₃
    a₁, a₃, b₁, b₃ = scheme.a₁, scheme.a₃, scheme.b₁, scheme.b₃

    smallk = k < SMALL_K # for small wavenumber use a Taylor series approximation

    C₁₀ = smallk ? (α₀ + 2β₀ - (α₁ + 2β₁) * k²) :
                   (  a₃ - a₁ * k² + (-a₃ + (α₃ + 2β₃) * k² / 3) * k * sink +
                    (-a₃ + ((a₁ + a₃ / 2) - (α₂ + 2β₂) * k² / 3) * k²) * cosk
                   ) / k⁶
    C₁₁ = smallk ? (α₀ -  β₀ - (α₁ -  β₁) * k²) :
                   (  b₃ - b₁ * k² + (-b₃ + (α₃ -  β₃) * k² / 3) * k * sink +
                    (-b₃ + ((b₁ + b₃ / 2) - (α₂ -  β₂) * k² / 3) * k²) * cosk
                   ) / k⁶

    C₀₀ = Ĉ(scheme.coreliquid, scheme.subscheme, k)
    C₁₀ = 24η * C₁₀
    C₁₁ = 24η * C₁₁

    return (C₀₀, C₁₀, C₁₁)
end


# this is the 2nd attempt; in this attempt I will try use MCTS and gradients
using StatsBase, DataFrames, Flux, JLD2, Plots
include("2048_fn.jl")

model = Flux.Chain(
  Dense(14*16+4, 128, relu),
  Dense(128, 64, relu),
  Dense(64, 32, relu),
  Dense(32, 1, relu));

loss(x,y) = sum((model(x) .- y).^2)
opt = ADAM(params(model))

# play one game to see how it goes
function assess_model(model)
  init_grid = init_game()
  @time (init_grid, fnl, seq, cart, tf, reward) = simulate_game!(init_grid, model, 0.0)

  # obtain the discounted rewards
  y = discounted_rewards(reward, 0.99)[1:end-1];

  # score every move using  model
  updated_model = [Tracker.data(model(grid_move_to_feature(sim_seq(init_grid, seq[1:i+1], cart[1:i+1], tf[1:i+1]), seq[i])))[1] for i=1:length(seq)-1]
  display(sim_seq(init_grid, seq, cart, tf))
  g = plot(updated_model, label="model")
  plot!(y, label="discounted rewards")
  (init_grid, fnl, seq, cart, tf, reward, g)
end

ϵ = 0.5

model = nothing
@load "no_neg_model.jld2"
model

@time for i=1:10000
  global ϵ
  ϵ -= 0.0001
  simulate_one_game_and_train!(model, opt, ϵ)
  if i % 100 == 0
    @save "no_neg_model.jld2" model;
    println("epsilon: $ϵ")
    init_grid = init_game()
    res = simulate_game!(init_grid, model, 0)
    display(res[2])
    @save "no_neg_model.jld2" model;
  end
end

@time (init, fnl, seq, cart, tf, reward, g) = assess_model(model);
g

last_n_move = length(seq)-20
grid_new = sim_seq(init, seq[1:end-last_n_move], cart[1:end-last_n_move], tf[1:end - last_n_move])
grid_try = copy(grid_new)
display(grid_new)
vm = valid_moves(grid_new)
x = vec(Tracker.data(model(hcat(grid_move_to_feature.([grid_new], vm)...))))
y = x .- minimum(x)

using DataFrames
DataFrame(vm=vm, chosen=vm.== seq[end-last_n_move], x=vec(x),  prob_exp = round.(exp.(y)/sum(exp.(y)),digits=2), prob_square =((y.+1).^2)./sum((y.+1).^2))


using BayesianDataFusion
using Base.Test

using DataFrames

A = randn(30,2);
B = randn(40,2);

X = Float64[ sum(A[i,:].*B[j,:]) for i in 1:size(A,1), j in 1:size(B,1) ]

df = DataFrame(A=Int64[], B=Int64[], v=Float64[])
for i=1:size(A,1), j=1:size(B,1)
  push!(df, Any[i, j, X[i,j]])
end

## adding features
feat = randn(size(A,1), 2)

rd = RelationData(df)
rd.entities[1].F = feat
rd.entities[1].lambda_beta_sample = true

assignToTest!(rd.relations[1], 10)

result = macau(rd, burnin=50, psamples=10, num_latent=2, verbose=false)


using Gcvspl, Test

# Dummy data
x = collect(0.:1:100)
y = 2*x

spline_coeff = 1.0
SplOrder = 3

c, WK, IER = Gcvspl.gcv_spl(x,y,ese,spline_coeff;SplineOrder = Int32(SplOrder-1)) # with cubic spline as the default
y_calc = Gcvspl.spl_der(x,x,c,SplineOrder= Int32(SplOrder-1))

@test isapprox(y,y_calc,atol=1e-1)


module Deriv_Utils_New
    using Statistics
    using Plots
    using BenchmarkTools
    using Profile
    using FFTW
    using Statistics
    using Optim
    using Images, FileIO, ImageIO
    using Printf

    #using Profile
    using LinearAlgebra

    # put the cwd on the path so Julia can find the module
    push!(LOAD_PATH, pwd()*"/main")
    using eqws

    #Jacobians of coeffs wrt image #######################################################

    function wst_S1_deriv(image::Array{Float64,2}, filter_hash::Dict; FFTthreads::Int=1)
        #=
        output: Nx, Nx, Nf
        =#
        # Use nthread threads for FFT -- but for Nx<512 nthread=1 is fastest.  Overhead?
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate output array
        dS1dα  = zeros(Float64, Nx, Nx, Nf)

        ## 1st Order
        im_fd = fft(image)

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]
        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals

        # make a FFTW "plan" for an array of the given size and type
        P = plan_ifft(im_fd)  # P is an operator, P*im is ifft(im)

        # Loop over filters
        for f = 1:Nf
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i

            zarr[f_i] = (f_v.*f_v) .* im_fd[f_i]
            dS1dα[:,:,f] = 2 .* real.(P*zarr)

            zarr[f_i] .= 0   # reset zarr for next loop
        end
        return dS1dα
    end

    function wst_S20_deriv(image::Array{Float64,2}, filter_hash::Dict; FFTthreads::Int=1)
        #=
        output: Nx, Nx, Nf, Nf
        =#
        # Use nthread threads for FFT -- but for Nx<512 nthread=1 is fastest.  Overhead?
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate output array
        dS20dα  = zeros(Float64, Nx, Nx, Nf, Nf)

        # allocate image arrays for internal use
        im_rdc = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_rd  = Array{Float64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_fd  = fft(image)

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]
        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals

        # make a FFTW "plan" a complex array, both forward and inverse transform
        P_fft  = plan_fft(im_fd)   # P_fft is an operator,  P_fft*im is fft(im)
        P_ifft = plan_ifft(im_fd)  # P_ifft is an operator, P_ifft*im is ifft(im)

        # Loop over filters
        for f = 1:Nf
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i

            zarr[f_i] = f_v .* im_fd[f_i]
            Z_λ = P_ifft*zarr  # complex valued ifft of zarr
            zarr[f_i] .= 0   # reset zarr for next loop
            im_rdc[:,:,f] = Z_λ
            im_rd[:,:,f]  = abs.(Z_λ)
        end

        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals
        for f2 = 1:Nf
            f_i = f_ind[f2]  # CartesianIndex list for filter
            f_v = f_val[f2]  # Values for f_i
            uvec = im_rdc[:,:,f2] ./ im_rd[:,:,f2]
            for f1 = 1:Nf
                temp = P_fft*(im_rd[:,:,f1].*uvec)
                zarr[f_i] = f_v .* temp[f_i]

                Z1dZ2 = real.(P_ifft*zarr)
                #  It is possible to do this with rifft, but it is not much faster...
                #   Z1dZ2 = myrealifft(zarr)
                dS20dα[:,:,f1,f2] += Z1dZ2
                dS20dα[:,:,f2,f1] += Z1dZ2
                zarr[f_i] .= 0   # reset zarr for next loop
            end
        end
        return dS20dα
    end

    function wst_S12_deriv(image::Array{Float64,2}, filter_hash::Dict; FFTthreads::Int=1)
        #=
        output: Nx, Nx, Nf, Nf
        =#
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]

        zarr = zeros(ComplexF64, Nf, Nx, Nx)  # temporary array to fill with zvals
        ## 1st Order
        im_fd_0 = fft(image)  # total power=1.0
        # make a FFTW "plan" for an array of the given size and type
        P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)
        for f1 = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1
            zarr[f1, f_i1] = im_fd_0[f_i1] .* f_v1 #check if this notation is fine
        end
        zarrc =conj.(zarr)
        absz = abs.(zarr)

        #Implemented in the more memory intensive / parallelized ops way
        #fhash_indsarr = [hcat((x->[x[1], x[2]]).(fh)...)' for fh in fhash["filt_index"]] #Each elem gives you fhash["filt_index"] in non-CartesianIndices
        #dabsz = #Mem Req: Nf*Nx^2*sp*Nx^2 and sp~2%

        #Less Mem Route
        dabsz_absz = zeros(Float64, Nf, Nf, Nx, Nx)  # Lookup arrays for dS12 terms
        xigrid = reshape(CartesianIndices((1:Nx, 1:Nx)), (1, Nx^2))
        xigrid = hcat((x->[x[1], x[2]]).(xigrid)...) #2*Nx^2
        #Calculate pnz's here and pre-allocate memory

        for f1 = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1

            for f2 = 1:Nf
                f_i2 = f_ind[f2]  # CartesianIndex list for filter1
                f_v2 = f_val[f2]  # Values for f_i1
                pnz = intersect(f_i1, f_i2) #Fd indices where both filters are nonzero
                if length(pnz)!=0
                    pnz_arr = vcat((x->[x[1], x[2]]').(pnz)...) #pnz*2
                    Φmat =  exp.((-2π*im)/Nx .* ((pnz_arr .- 1) * (xigrid .- 1))) #pnz*Nx^2, BN1
                    f_v1pnz = f_v1[findall(in(pnz), f_i1)]
                    t2 = real.(zarrc[f1, pnz] .* Φmat) #Check, BN3
                    #println(size(f_v1pnz), size(t2), size(absz[f2, pnz]))
                    term = sum(((absz[f2, pnz] ./ absz[f1, pnz]) .* f_v1pnz) .* t2, dims=1) #p*Nx^2 -> 1*Nx^2, BN2
                    dabsz_absz[f1, f2, :, :] = reshape(term, (Nx, Nx))
                end
            end
        end
        dS12 = dabsz_absz + permutedims(dabsz_absz, [2, 1, 3, 4])
        return permutedims(dS12, [3, 4, 1, 2]) #MODIFIED
    end
    #=
    function wst_S12_derivfast(image::Array{Float64,2}, filter_hash::Dict; FFTthreads::Int=1)
        #=
        output: Nx, Nx, Nf, Nf
        =#
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]

        zarr = zeros(ComplexF64, Nf, Nx, Nx)  # temporary array to fill with zvals
        ## 1st Order
        im_fd_0 = fft(image)  # total power=1.0
        # make a FFTW "plan" for an array of the given size and type
        P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)
        for f1 = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1
            zarr[f1, f_i1] = im_fd_0[f_i1] .* f_v1 #check if this notation is fine
        end
        zarrc =conj.(zarr)
        absz = abs.(zarr)

        #Implemented in the more memory intensive / parallelized ops way
        #fhash_indsarr = [hcat((x->[x[1], x[2]]).(fh)...)' for fh in fhash["filt_index"]] #Each elem gives you fhash["filt_index"] in non-CartesianIndices
        #dabsz = #Mem Req: Nf*Nx^2*sp*Nx^2 and sp~2%

        #Less Mem Route
        dabsz_absz = zeros(Float64, Nf, Nf, Nx, Nx)  # Lookup arrays for dS12 terms
        xigrid = reshape(CartesianIndices((1:Nx, 1:Nx)), (1, Nx^2))
        xigrid = hcat((x->[x[1], x[2]]).(xigrid)...) #2*Nx^2
        #Calculate pnz's here and pre-allocate memory
        for f1  = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1
            for f2 = 1:Nf
                f_i2 = f_ind[f2]  # CartesianIndex list for filter1
                f_v2 = f_val[f2]  # Values for f_i1

        for f1 = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1

            for f2 = 1:Nf
                f_i2 = f_ind[f2]  # CartesianIndex list for filter1
                f_v2 = f_val[f2]  # Values for f_i1
                pnz = intersect(f_i1, f_i2) #Fd indices where both filters are nonzero
                if length(pnz)!=0
                    pnz_arr = vcat((x->[x[1], x[2]]').(pnz)...) #pnz*2
                    Φmat =  exp.((-2π*im)/Nx .* ((pnz_arr .- 1) * (xigrid .- 1))) #pnz*Nx^2, BN1
                    f_v1pnz = f_v1[findall(in(pnz), f_i1)]
                    t2 = real.(zarrc[f1, pnz] .* Φmat) #Check, BN3
                    #println(size(f_v1pnz), size(t2), size(absz[f2, pnz]))
                    term = sum(((absz[f2, pnz] ./ absz[f1, pnz]) .* f_v1pnz) .* t2, dims=1) #p*Nx^2 -> 1*Nx^2, BN2
                    dabsz_absz[f1, f2, :, :] = reshape(term, (Nx, Nx))
                end
            end
        end
        dS12 = dabsz_absz + permutedims(dabsz_absz, [2, 1, 3, 4])
        return permutedims(dS12, [3, 4, 1, 2]) #MODIFIED
    end
    =#
    function wst_S1S2_deriv(image::Array{Float64,2}, filter_hash::Dict; FFTthreads::Int=1)
        #=
        output: Nx, Nx, Nf+Nf^2
        =#
        # Use 2 threads for FFT
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate output array
        dS1dp  = zeros(Float64, Nf, Nx, Nx) #Mod zeros(Float64, Nx, Nx, Nf)
        dS2dp  = zeros(Float64, Nf, Nf, Nx, Nx) #Mod

        # allocate image arrays for internal use
        im_rdc_0_1 = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_rd_0_1  = Array{Float64, 3}(undef, Nx, Ny, Nf)  # real domain, complex

        # Not sure what to do here -- for gradients I don't think we want these
        ## 0th Order
        #S0[1]   = mean(image)
        #norm_im = image.-S0[1]
        #S0[2]   = sum(norm_im .* norm_im)/(Nx*Ny)
        #norm_im ./= sqrt(Nx*Ny*S0[2])

        ## 1st Order
        im_fd_0 = fft(image)  # total power=1.0

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]

        f_ind_rev = [[CartesianIndex(mod1(Nx+2 - ci[1],Nx), mod1(Nx+2 - ci[2],Nx)) for ci in f_Nf] for f_Nf in f_ind]

        #CartesianIndex(17,17) .- f_ind
        # tmp arrays for all tmp[f_i] = arr[f_i] .* f_v opns
        zarr1 = zeros(ComplexF64, Nx, Nx)
        fz_fψ1 = zeros(ComplexF64, Nx, Nx)
        fterm_a = zeros(ComplexF64, Nx, Nx)
        fterm_ct1 = zeros(ComplexF64, Nx, Nx)
        fterm_ct2 = zeros(ComplexF64, Nx, Nx)
        #rterm_bt2 = zeros(ComplexF64, Nx, Nx)
        # make a FFTW "plan" for an array of the given size and type
        P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)

        # Loop over filters
        for f1 = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1
            f_i1rev = f_ind_rev[f1]

            zarr1[f_i1] = f_v1 .* im_fd_0[f_i1] #F[z] #BUG: here????
            zarr1_rd = P*zarr1 #for frzi
            fz_fψ1[f_i1] = f_v1 .* zarr1[f_i1]
            fz_fψ1_rd = P*fz_fψ1
            dS1dp[f1, :,:] = 2 .* real.(fz_fψ1_rd) #real.(conv(I_λ, ψ_λ)) #Mod
            #CHECK; that this equals derivS1fast code

            #dS2 loop prep
            #ψ_λ1  = realspace_filter(Nx, f_i1, f_v1) #Slow
            #fcψ_λ1 = fft(conj.(ψ_λ1)) #Slow
            I_λ1 = sqrt.(abs2.(zarr1_rd))
            fI_λ1 = fft(I_λ1)
            rterm_bt1 = zarr1_rd./I_λ1 #Z_λ1/I_λ1_bar.
            rterm_bt2 = conj.(zarr1_rd)./I_λ1
            for f2 = 1:Nf
                f_i2 = f_ind[f2]  # CartesianIndex list for filter2
                f_v2 = f_val[f2]  # Values for f_i2

                fterm_a[f_i2] = fI_λ1[f_i2] .* (f_v2).^2 #fterm_a = F[I_λ1].F[ψ_λ]^2
                rterm_a = P*fterm_a                #Finv[fterm_a]
                fterm_bt1 = fft(rterm_a .* rterm_bt1) #F[Finv[fterm_a].*Z_λ1/I_λ1_bar] for T1
                fterm_bt2 = fft(rterm_a .* rterm_bt2) #F[Finv[fterm_a].*Z_λ1_bar/I_λ1_bar] for T2
                fterm_ct1[f_i1] = fterm_bt1[f_i1] .* f_v1    #fterm_b*F[ψ_λ]
                #fterm_ct2 = fterm_bt2 .* fcψ_λ1             #Slow
                #println(size(fterm_ct2[f_i1rev]),size(fterm_bt2[f_i1rev]),size(f_v1),size(conj.(f_v1)))
                fterm_ct2[f_i1rev] = fterm_bt2[f_i1rev] .* f_v1 #f_v1 is real
                #fterm_ct2slow = fterm_bt2 .* fcψ_λ1
                dS2dp[f1, f2, :, :] = real.(P*(fterm_ct1 + fterm_ct2)) #Mod THIS SHOULD be Nx*f2-1 + f1 or do the reshaping in loss func.

                #Reset f2 specific tmp arrays to 0
                fterm_a[f_i2] .= 0
                fterm_ct1[f_i1] .=0
                fterm_ct2[f_i1rev] .=0

            end
            # reset all reused variables to 0
            zarr1[f_i1] .= 0
            fz_fψ1[f_i1] .= 0

        end
        return permutedims(vcat(dS1dp, reshape(dS2dp, (Nf^2, Nx, Nx))), [2, 3, 1]) #Mod #MODIFIED
    end

    #Chi-Square Derivatives###############################################################
    function wst_S2_deriv_sum(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64}; FFTthreads::Int=1)
        #=
        output:
        =#
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        if size(wt)!=(Nf, Nf) error("Wt has wrong shape") end

        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]
        f_ind_rev = [[CartesianIndex(mod1(Nx+2 - ci[1],Nx), mod1(Nx+2 - ci[2],Nx)) for ci in f_Nf] for f_Nf in f_ind]

        im_fd_0 = fft(image)
        P = plan_fft(im_fd_0)
        Pi = plan_ifft(im_fd_0)
        zarr1 = zeros(ComplexF64, Nx, Nx)
        Ilam1 = zeros(ComplexF64, Nx, Nx)
        fterm_bt1 = zeros(ComplexF64, Nx, Nx)
        fterm_bt2 = zeros(ComplexF64, Nx, Nx)
        ifterm = zeros(ComplexF64, Nx, Nx)
        dLossdα = zeros(Float64, Nx, Nx)
        for f1=1:Nf
            vlam1 = zeros(ComplexF64, Nx, Nx) #defined for each lambda1, computed by summing over lambda2
            f_i1 = f_ind[f1]
            f_v1 = f_val[f1]
            f_i1rev = f_ind_rev[f1]

            zarr1[f_i1] = im_fd_0[f_i1] .* f_v1
            zarr1_rd = Pi*(zarr1)
            Ilam1 = abs.(zarr1_rd)
            fIlam1 = P*Ilam1

            for f2=1:Nf
                f_i2 = f_ind[f2]
                f_v2 = f_val[f2]
                vlam1[f_i2] += wt[f1, f2] .* f_v2.^2
            end
            #vlam1= sum_λ2 w.ψl1,l2^2
            vlam1 = (Pi*(fIlam1 .* vlam1))./Ilam1 #sum over f2 -> (Nx, Nx) -> real space Vλ1/Iλ1
            fterm_bt1 = P*(vlam1 .* zarr1_rd)
            fterm_bt2 = P*(vlam1 .* conj.(zarr1_rd))
            ifterm[f_i1] = fterm_bt1[f_i1] .* f_v1
            ifterm[f_i1rev] += fterm_bt2[f_i1rev] .* f_v1
            dLossdα += real.(Pi*ifterm)

            #zero out alloc
            zarr1[f_i1] .= 0
            ifterm[f_i1] .=0
            ifterm[f_i1rev] .=0

        end
        return reshape(dLossdα, (Nx^2, 1))

    end

    function wst_S1S2_derivsum_comb(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64}; FFTthreads::Int=1)
        FFTW.set_num_threads(FFTthreads)
        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        if size(wt)!=(Nf + (Nf*Nf),) error("Wt has wrong shape: the argument should contain both S1 and S2 weights") end
        wts1 = reshape(wt[1:Nf], (Nf, 1))
        wts2 = reshape(wt[Nf+1:end], (Nf, Nf))

        dwS1 = reshape(wst_S1_deriv(image, filter_hash), (Nx^2, Nf)) * wts1 #NOTE: Replace this with eqws_2DUtils.wst_S1_deriv
        dwS2 = reshape(wst_S2_deriv_sum(image, filter_hash, wts2), (Nx^2, 1))
        #println(size(dwS1), size(dwS2))
        return dwS1 + dwS2 #Why was another reshape needed here?

    end
    #=
    function wst_S20_deriv_sum_old(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64}; FFTthreads::Int=1)
        # Sum over (f1,f2) filter pairs for S20 derivative.  This is much faster
        #   than calling wst_S20_deriv() because the sum can be moved inside the FFT.
        # Use FFTthreads threads for FFT -- but for Nx<512 FFTthreads=1 is fastest.  Overhead?
        # On Cascade Lake box, 4 is good for 2D, 8 or 16 for 3D FFTs
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate image arrays for internal use
        Uvec   = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_rd  = Array{Float64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_fd  = fft(image) #Should this be the apodized image?

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]
        zarr    = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals

        # make a FFTW "plan" a complex array, both forward and inverse transform
        P_fft  = plan_fft(im_fd)   # P_fft is an operator,  P_fft*im is fft(im)
        P_ifft = plan_ifft(im_fd)  # P_ifft is an operator, P_ifft*im is ifft(im)

        # Loop over filters
        for f = 1:Nf
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i

            zarr[f_i] = f_v .* im_fd[f_i]
            Z_λ = P_ifft*zarr  # complex valued ifft of zarr
            zarr[f_i] .= 0     # reset zarr for next loop
            im_rd[:,:,f] = abs.(Z_λ)
            Uvec[:,:,f]  = Z_λ ./ im_rd[:,:,f]
        end

        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals
        for f2 = 1:Nf
            f_i = f_ind[f2]  # CartesianIndex list for filter
            f_v = f_val[f2]  # Values for f_i

            Wtot = reshape( reshape(im_rd,Nx*Nx,Nf)*wt[:,f2], Nx, Nx) #SPEED: Preallocate mem for Wtot?
            temp = P_fft*(Wtot.*Uvec[:,:,f2])
            zarr[f_i] .+= f_v .* temp[f_i]
        end
        ΣdS20dα = 2* real.(P_ifft*zarr) #added here

        return reshape(ΣdS20dα, (Nx^2, 1))
    end
    =#

    function wst_S20_deriv_sum(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64}; FFTthreads::Int=1)
        # Sum over (f1,f2) filter pairs for S20 derivative.  This is much faster
        #   than calling wst_S20_deriv() because the sum can be moved inside the FFT.
        # Use FFTthreads threads for FFT -- but for Nx<512 FFTthreads=1 is fastest.  Overhead?
        # On Cascade Lake box, 4 is good for 2D, 8 or 16 for 3D FFTs
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate image arrays for internal use
        Uvec   = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_rd  = Array{Float64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_fd  = fft(image) #Should this be the apodized image?

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]
        zarr    = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals

        # make a FFTW "plan" a complex array, both forward and inverse transform
        P_fft  = plan_fft(im_fd)   # P_fft is an operator,  P_fft*im is fft(im)
        P_ifft = plan_ifft(im_fd)  # P_ifft is an operator, P_ifft*im is ifft(im)

        # Loop over filters
        for f = 1:Nf
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i

            zarr[f_i] = f_v .* im_fd[f_i]
            Z_λ = P_ifft*zarr  # complex valued ifft of zarr
            zarr[f_i] .= 0     # reset zarr for next loop
            im_rd[:,:,f] = abs.(Z_λ)
            Uvec[:,:,f]  = Z_λ ./ im_rd[:,:,f]
        end

        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals
        for f2 = 1:Nf
            f_i = f_ind[f2]  # CartesianIndex list for filter
            f_v = f_val[f2]  # Values for f_i

            Wtot = reshape( reshape(im_rd,Nx*Nx,Nf)* (wt[:,f2] + wt'[:, f2]), Nx, Nx) #SPEED: Preallocate mem for Wtot?
            temp = P_fft*(Wtot.*Uvec[:,:,f2])
            zarr[f_i] .+= f_v .* temp[f_i]
        end
        ΣdS20dα = real.(P_ifft*zarr) #added here

        return reshape(ΣdS20dα, (Nx^2, 1))
    end

    function wst_S20_deriv_sum_norm(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64}, std::Float64; FFTthreads::Int=1)
        # Sum over (f1,f2) filter pairs for S20 derivative.  This is much faster
        #   than calling wst_S20_deriv() because the sum can be moved inside the FFT.
        # Use FFTthreads threads for FFT -- but for Nx<512 FFTthreads=1 is fastest.  Overhead?
        # On Cascade Lake box, 4 is good for 2D, 8 or 16 for 3D FFTs
        FFTW.set_num_threads(FFTthreads)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate image arrays for internal use
        Uvec   = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_rd  = Array{Float64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_fd  = fft(image) #Should this be the apodized image?

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]
        zarr    = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals

        # make a FFTW "plan" a complex array, both forward and inverse transform
        P_fft  = plan_fft(im_fd)   # P_fft is an operator,  P_fft*im is fft(im)
        P_ifft = plan_ifft(im_fd)  # P_ifft is an operator, P_ifft*im is ifft(im)

        # Loop over filters
        for f = 1:Nf
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i

            zarr[f_i] = f_v .* im_fd[f_i]
            Z_λ = P_ifft*zarr  # complex valued ifft of zarr
            zarr[f_i] .= 0     # reset zarr for next loop
            im_rd[:,:,f] = abs.(Z_λ)
            Uvec[:,:,f]  = Z_λ ./ im_rd[:,:,f]
        end

        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals
        for f2 = 1:Nf
            f_i = f_ind[f2]  # CartesianIndex list for filter
            f_v = f_val[f2]  # Values for f_i

            Wtot = reshape(reshape(im_rd,Nx*Nx,Nf)* (wt[:,f2] + wt'[:, f2]), Nx, Nx) #SPEED: Preallocate mem for Wtot?
            temp = P_fft*(Wtot.*Uvec[:,:,f2])
            zarr[f_i] .+= f_v .* temp[f_i]
        end
        ΣdS20dα = real.(P_ifft*zarr).*(1.0/(Nx*std)) #added here

        return reshape(ΣdS20dα, (Nx^2, 1))
    end

    #Image Reconstruction Funcs##########################################################
    #=
    function image_recon_derivsum(input::Array{Float64, 2}, filter_hash::Dict, s_targ_mean::Array{Float64, 1}, s_targ_invcov::Array{Float64, 2}, pixmask::BitArray{2}, coeff_choice;
        FFTthreads::Int=1, optim_settings=Dict([("iterations", 10)]), coeff_mask=nothing) #add iso here and a check that returns error if both coeffmask is not nothing and iso is present.
        #=
        Cases to be handled by this function:
        S2 | S20 | S12
        Select pixels masked or all floating
        Select coefficients optimized wrt

        Cases NOT handled by this function:
        Iso
        Only S1
        No Regularization
        Non-Gaussian loss function
        Using the S0 params (mean and variance) in the loss function
        eqws_compute with norm=true
        Log
        Apodizing
        =#
        (Nf,) = size(filter_hash["filt_index"])
        println("Coeff mask:", (coeff_mask!=nothing))

        (Nx, Ny)  = size(input)
        if Nx != Ny error("Input image must be square") end

        if Nf == 0 error("filter hash corrupted") end
        if coeff_choice=="S2"
            println("S2")
            pixmask = pixmask[:] #flattened: Nx^2s

            if coeff_mask!=nothing
                if length(coeff_mask)!= (2+Nf + Nf^2) error("Wrong dim mask") end
                if size(s_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
                if size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_invcov should only have coeffs to be optimized") end
            else #No mask: all coeffs (default: All S1 and S2 coeffs will be optim wrt. Mean, var not incl)
                #Currently assuming inputs have mean, var params in case that's something we wanna optimize at some point
                if size(s_targ_mean)[1]!=(Nf+Nf^2) error("s_targ_mean should only contain coeffs to be optimized") end
                if (size(s_targ_invcov)!=(Nf+Nf^2, Nf+Nf^2))  error("s_targ_invcov of wrong size") end #(s_targ_invcov!=I) &
                #At this point both have dims |S1+S2|
                #Create coeff_mask subsetting 3:end
                coeff_mask = fill(true, 2+Nf+Nf^2)
                coeff_mask[1] = false
                coeff_mask[2] = false
            end

            num_freecoeff = count((i->(i==true)), coeff_mask) #Nf+Nf^2 if argument coeff_mask was empty
            num_freecoeffS1 = count((i->(i==true)), coeff_mask[3:Nf+2])
            num_freecoeffS2 = count((i->(i==true)), coeff_mask[Nf+3:end])

            #After this all cases have a coeffmask, and s_targ_mean and s_targ_invcov have the shapes of the coefficients that we want to select.

            #Optimization Hyperparameters
            norm=false
            numitns_dict = get(optim_settings, "iterations", 100)

            function loss_func(img_curr)
                #size(img_curr) must be (Nx^2, 1)
                s_curr = eqws_compute(reshape(img_curr, (Nx, Nx)),  filter_hash, doS2=true, doS12=false, doS20=false, norm=tonorm)[coeff_mask]
                neglogloss = 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_targ_mean)
                return neglogloss[1]
            end

            function dloss(storage_grad, img_curr)
                #storage_grad, img_curr must be (Nx^2, 1)
                s_curr = eqws_compute(reshape(img_curr, (Nx, Nx)), filter_hash, doS2=true, doS12=false, doS20=false, norm=tonorm)[coeff_mask]
                diff = s_curr - s_targ_mean
                wts1s2 = zeros(Nf+Nf^2)
                wtall = reshape(diff' * s_targ_invcov, num_freecoeff)
                wts1s2[coeff_mask[3:end]] .= wtall
                dterm = wst_S1S2_derivsum_comb(reshape(img_curr, (Nx, Nx)), filter_hash, wts1s2, FFTthreads=FFTthreads)
                storage_grad .= reshape(dterm, (Nx^2, 1))
                storage_grad[pixmask, 1] .= 0 # better way to do this by taking pixmask as an argument wst_s2_deriv_sum?
                return
            end

            #Debugging stuff
            println("Diff check")
            eps = zeros(size(input))
            eps[1, 2] = 1e-4
            chisq1 = loss_func(input+eps./2)
            chisq0 = loss_func(input-eps./2)
            brute  = (chisq1-chisq0)/1e-4
            #df_brute = eqws_compute(reshape(input, (Nx, Nx)), filter_hash, doS2=true, doS12=false, doS20=false, norm=false)[coeff_mask] - s_targ_mean
            clever = reshape(zeros(size(input)), (Nx*Nx, 1))
            _bar = dloss(clever, reshape(input, (Nx^2, 1)))
            #println("dS1S2comb check")
            #wts1s2 = zeros(Nf+Nf^2)
            #wtall = reshape(df' * s_targ_invcov, num_freecoeff)
            #wts1s2[coeff_mask[3:end]] .= wtall
            #dterm = Deriv_Utils.wst_S1S2_derivsum_comb(reshape(input, (Nx, Nx)), filter_hash, wts1s2)
            #dS1S2comb = Deriv_Utils.wst_S1S2_derivfast(reshape(input, (Nx, Nx)), filter_hash)
            #dS1dp = permutedims(dS1S2comb[1:Nf, :, :], [2, 3, 1])
            #dS2dp = permutedims(dS1S2comb[Nf+1:end, :, :], [2, 3, 1])
            #dS2dp = reshape(dS2dp, Nx, Nx, Nf^2)
            #sum2 = zeros(Float64, Nx, Nx)
            #for i=1:Nf*Nf sum2 += (dS2dp[:,:,i].*wts1s2[Nf+i]) end
            #term1 = reshape(dS1dp, (Nx^2, Nf)) * reshape(wts1s2[1:Nf], (Nf, 1))
            #dterm_brute = term1 + reshape(sum2, (Nx^2, 1))
            println("Chisq Derve Check")
            println("Brute:  ",brute)
            println("Clever: ",clever[Nx*(1)+1])

            res = optimize(loss_func, dloss, reshape(input, (Nx^2, 1)), ConjugateGradient(), Optim.Options(iterations = numitns_dict, store_trace = true, show_trace = true))
            result_img = Optim.minimizer(res)
        elseif coeff_choice=="S20"
            println("S20")
            pixmask = pixmask[:] #flattened: Nx^2s

            if coeff_mask!=nothing
                if count((i->i==true), coeff_mask[1:Nf+2])!=0 println("Warning: S1 coeffs being double counted / S0 aren't 0") end #wrap
                if length(coeff_mask)!= (2+Nf + Nf^2) error("Wrong dim mask") end #wrap
                if size(s_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
                if size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_invcov should only have coeffs to be optimized") end
            else #No mask: all coeffs (default: Only S20 coeffs will be optim wrt.)
                #Currently assuming inputs have mean, var params in case that's something we wanna optimize at some point
                if size(s_targ_mean)[1]!=(Nf^2) error("s_targ_mean should only contain coeffs to be optimized") end
                if (size(s_targ_invcov)!=(Nf^2, Nf^2))  error("s_targ_invcov of wrong size") end #(s_targ_invcov!=I) &
                #At this point both have dims |S1+S2|
                #Create coeff_mask subsetting 3:end
                coeff_mask = fill(true, 2+Nf+Nf^2)
                coeff_mask[1:Nf+2] .= false #Mask out S0 and S1
            end

            num_freecoeff = count((i->(i==true)), coeff_mask) #Nf^2 if argument coeff_mask was empty

            #After this all cases have a coeffmask, and s_targ_mean and s_targ_invcov have the shapes of the coefficients that we want to select.

            #Optimization Hyperparameters
            tonorm = get(optim_settings, "norm", false)
            numitns_dict = get(optim_settings, "iterations", 100)

            function loss_func20(img_curr)
                #size(img_curr) must be (Nx^2, 1)
                s_curr = eqws_compute(reshape(img_curr, (Nx, Nx)),  filter_hash, doS2=false, doS12=false, doS20=true, norm=tonorm)[coeff_mask]
                neglogloss = 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_targ_mean)
                return neglogloss[1]
            end

            function dloss20(storage_grad, img_curr)
                #storage_grad, img_curr must be (Nx^2, 1)
                s_curr = eqws_compute(reshape(img_curr, (Nx, Nx)), filter_hash, doS2=false, doS12=false, doS20=true, norm=tonorm)[coeff_mask]
                diff = s_curr - s_targ_mean
                wt = reshape(convert(Array{Float64, 2}, transpose(diff) * s_targ_invcov), (Nf, Nf))
                dterm = wst_S20_deriv_sum(reshape(img_curr, (Nx, Nx)), filter_hash, wt, FFTthreads=FFTthreads)
                storage_grad .= reshape(dterm, (Nx^2, 1))
                storage_grad[pixmask, 1] .= 0 # better way to do this by taking pixmask as an argument wst_s2_deriv_sum?
                return
            end

            #Debugging stuff
            println("Diff check")
            eps = zeros(size(input))
            eps[1, 2] = 1e-4
            chisq1 = loss_func20(input+eps./2)
            chisq0 = loss_func20(input-eps./2)
            brute  = (chisq1-chisq0)/1e-4
            #df_brute = eqws_compute(reshape(input, (Nx, Nx)), filter_hash, doS2=true, doS12=false, doS20=false, norm=false)[coeff_mask] - s_targ_mean
            clever = reshape(zeros(size(input)), (Nx*Nx, 1))
            _bar = dloss20(clever, reshape(input, (Nx^2, 1)))
            println("Chisq Derve Check")
            println("Brute:  ",brute)
            println("Clever: ",clever[Nx*(1)+1])

            res = optimize(loss_func20, dloss20, reshape(input, (Nx^2, 1)), ConjugateGradient(), Optim.Options(iterations = numitns_dict, store_trace = true, show_trace = true))
            result_img = Optim.minimizer(res)
        else
             if coeff_choice!="S12" error("Invalid coeff_choice") end
             println("S12")
             pixmask = pixmask[:] #flattened: Nx^2s

             if coeff_mask!=nothing
                 num_freecoeff = count((i->(i==true)), coeff_mask)
                 println("Number of coefficients being optimized= ", num_freecoeff)
                 if count((i->i==true), coeff_mask[1:Nf+2])!=0 println("Warning: S1 coeffs being double counted / S0 aren't 0") end #wrap
                 if length(coeff_mask)!= (2+Nf + Nf^2) error("Wrong dim mask") end #wrap
                 if size(s_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
                 if size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_invcov should only have coeffs to be optimized") end
             else #No mask: all coeffs (default: Only S20 coeffs will be optim wrt.)
                 error("S12 must be run with the coeff_mask argument")
                 #=
                 if size(s_targ_mean)[1]!=(Nf^2) error("s_targ_mean should only contain coeffs to be optimized") end
                 if (size(s_targ_invcov)!=(Nf^2, Nf^2))  error("s_targ_invcov of wrong size") end #(s_targ_invcov!=I) &
                 #At this point both have dims |S1+S2|
                 #Create coeff_mask subsetting 3:end
                 coeff_mask = fill(true, 2+Nf+Nf^2)
                 coeff_mask[1:Nf+2] .= false #Mask out S0 and S1 =#
             end

             #After this all cases have a coeffmask, and s_targ_mean and s_targ_invcov have the shapes of the coefficients that we want to select.

             #Optimization Hyperparameters
             tonorm = get(optim_settings, "norm", false)
             numitns_dict = get(optim_settings, "iterations", 100)
             coeffmask_S12 = coeff_mask[Nf+3:end]

             function loss_func12(img_curr)
                 #size(img_curr) must be (Nx^2, 1)
                 s_curr = eqws_compute(reshape(img_curr, (Nx, Nx)),  filter_hash, doS2=false, doS12=true, doS20=false, norm=tonorm)[coeff_mask]
                 neglogloss = 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_targ_mean)
                 return neglogloss[1]
             end

             function dloss12(storage_grad, img_curr)
                 #storage_grad, img_curr must be (Nx^2, 1)
                 s_curr = eqws_compute(reshape(img_curr, (Nx, Nx)), filter_hash, doS2=false, doS12=true, doS20=false, norm=tonorm)[coeff_mask]
                 diff = s_curr - s_targ_mean
                 wt = transpose(diff) * s_targ_invcov
                 storage_grad .= reshape(wst_S12_deriv(reshape(img_curr, (Nx, Nx)), filter_hash, FFTthreads=FFTthreads), (Nx^2, Nf^2))[:, coeffmask_S12] * transpose(wt) #->(nx^2, 1)
                 storage_grad[pixmask, 1] .= 0 # better way to do this by taking pixmask as an argument wst_s2_deriv_sum?
                 return
             end

             #Debugging stuff
             println("Diff check")
             eps = zeros(size(input))
             eps[1, 2] = 1e-4
             chisq1 = loss_func12(input+eps./2)
             chisq0 = loss_func12(input-eps./2)
             brute  = (chisq1-chisq0)/1e-4
             #df_brute = eqws_compute(reshape(input, (Nx, Nx)), filter_hash, doS2=true, doS12=false, doS20=false, norm=false)[coeff_mask] - s_targ_mean
             clever = reshape(zeros(size(input)), (Nx*Nx, 1))
             _bar = dloss12(clever, reshape(input, (Nx^2, 1)))
             println("Chisq Derve Check")
             println("Brute:  ",brute)
             println("Clever: ",clever[Nx*(1)+1])

             res = optimize(loss_func12, dloss12, reshape(input, (Nx^2, 1)), ConjugateGradient(), Optim.Options(iterations = numitns_dict, store_trace = true, show_trace = true))
             result_img = Optim.minimizer(res)

        end
        return reshape(result_img, (Nx, Nx))
    end
    =#
    function image_recon_derivsum_regularized(input::Array{Float64, 2}, filter_hash::Dict, s_targ_mean::Array{Float64, 1}, s_targ_invcov, pixmask::BitArray{2}, dhc_args;
        FFTthreads::Int=1, optim_settings=Dict([("iterations", 10)]), coeff_mask::BitArray{1}=nothing, lambda=0.001,show_trace0::Bool=true) #add iso here and a check that returns error if both coeffmask is not nothing and iso is present.
        #=
        Cases to be handled by this function:
        S20
        Apodized or not
        Iso
        Log or Regular: If log then input, targ and cov should have been constructed using log
        Select pixels masked or all floating--not using for masked currently
        Select coefficients optimized wrt
        Lambda: regularization

        Cases NOT handled by this function:
        Only S1
        Non-Gaussian loss function
        Using the S0 params (mean and variance) in the loss function
        eqws_compute with norm=true
        =#
        (N1iso, Nf) = size(filter_hash["S1_iso_mat"])
        println("Coeff mask:", (coeff_mask!=nothing))

        (Nx, Ny)  = size(input)
        if Nx != Ny error("Input image must be square") end
        if !dhc_args[:doS20] error("Not implemented") end
        if Nf == 0 error("filter hash corrupted") end

        println("S20")
        #pixmask = pixmask[:] #flattened: Nx^2s #DEB
        ori_input = reshape(copy(input), (Nx^2, 1))

        if coeff_mask!=nothing
            if (!dhc_args[:iso]) & (count((i->i==true), coeff_mask[1:Nf+2])!=0) error("Code assumes S1 coeffs and S0 are masked out") end #wrap
            if (dhc_args[:iso]) & (count((i->i==true), coeff_mask[1:N1iso+2])!=0) error("Code assumes S1iso coeffs and S0 are masked out") end
            if size(s_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
            if size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_invcov should only have coeffs to be optimized") end
        else #No mask: all coeffs (default: Only S20 coeffs will be optim wrt.)
            error("You must supply a coeff_mask")
        end

        #Apodization adjustments
        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
        #For the Jacobian dA*P/dP
        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

        num_freecoeff = count((i->(i==true)), coeff_mask)

        #After this all cases have a coeffmask, and s_targ_mean and s_targ_invcov have the shapes of the coefficients that we want to select.
        #Optimization Hyperparameters
        tonorm = false
        numitns_dict = get(optim_settings, "iterations", 100)
        minmethod = get(optim_settings, "minmethod", ConjugateGradient())

        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})
            s_curr = eqws_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]
            regterm =  0.5*lambda*sum((adaptive_apodizer(img_curr) - adaptive_apodizer(input)).^2)
            lnlik = ( 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_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)

        function dloss20(storage_grad::Array{Float64, 2}, img_curr::Array{Float64, 2})
            s_curr = eqws_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]
            diff = s_curr - s_targ_mean
            #Add branches here
            wt = augment_weights(convert(Array{Float64,1}, reshape(transpose(diff) * s_targ_invcov, (length(diff),))))
            apdimg_curr = adaptive_apodizer(img_curr)
            storage_grad .= reshape((wst_S20_deriv_sum(apdimg_curr, filter_hash, wt, FFTthreads=FFTthreads)' + 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
        epsmag = 1e-4
        eps[row, col] = epsmag
        chisq1 = loss_func20(input+eps./2) #DEB
        chisq0 = loss_func20(input-eps./2) #DEB
        brute  = (chisq1-chisq0)/epsmag
        #df_brute = eqws_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], " Mean ", meanval) #DEB
        println("Initial Loss ", loss_func20(input))
        res = optimize(loss_func20, dloss20, input, minmethod, Optim.Options(iterations = numitns_dict, store_trace = show_trace0, show_trace = show_trace0))
        result_img = Optim.minimizer(res)
        println("Final Loss ", loss_func20(result_img))
        return res, reshape(result_img, (Nx, Nx))
    end

    function image_recon_derivsum_regularized_S2(input::Array{Float64, 2}, filter_hash::Dict, s_targ_mean::Array{Float64, 1}, s_targ_invcov, pixmask::BitArray{2}, dhc_args;
        FFTthreads::Int=1, optim_settings=Dict([("iterations", 10)]), coeff_mask=nothing, lambda=0.001,show_trace0::Bool=true) #add iso here and a check that returns error if both coeffmask is not nothing and iso is present.
        #=
        Cases to be handled by this function:
        S20
        Apodized or not
        Iso
        Log or Regular: If log then input, targ and cov should have been constructed using log
        Select pixels masked or all floating--not using for masked currently
        Select coefficients optimized wrt
        Lambda: regularization

        Cases NOT handled by this function:
        Only S1
        Non-Gaussian loss function
        Using the S0 params (mean and variance) in the loss function
        eqws_compute with norm=true
        =#
        (N1iso, Nf) = size(filter_hash["S1_iso_mat"])
        println("Coeff mask:", (coeff_mask!=nothing))

        (Nx, Ny)  = size(input)
        if Nx != Ny error("Input image must be square") end
        if !dhc_args[:doS2] error("Not implemented") end
        if Nf == 0 error("filter hash corrupted") end

        println("S2")
        #pixmask = pixmask[:] #flattened: Nx^2s #DEB
        ori_input = reshape(copy(input), (Nx^2, 1))

        if coeff_mask!=nothing
            if (!dhc_args[:iso]) & (count((i->i==true), coeff_mask[1:Nf+2])!=0) error("Code assumes S1 coeffs and S0 are masked out") end #wrap
            if (dhc_args[:iso]) & (count((i->i==true), coeff_mask[1:N1iso+2])!=0) error("Code assumes S1iso coeffs and S0 are masked out") end
            if size(s_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
            if size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_invcov should only have coeffs to be optimized") end
        else #No mask: all coeffs (default: Only S20 coeffs will be optim wrt.)
            error("You must supply a coeff_mask")
        end

        #Apodization adjustments
        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
        #For the Jacobian dA*P/dP
        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

        num_freecoeff = count((i->(i==true)), coeff_mask)

        #After this all cases have a coeffmask, and s_targ_mean and s_targ_invcov have the shapes of the coefficients that we want to select.
        #Optimization Hyperparameters
        tonorm = false
        numitns_dict = get(optim_settings, "iterations", 100)
        minmethod = get(optim_settings, "minmethod", ConjugateGradient())

        if (dhc_args[:doS2]) & (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[:doS2]) & (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_func2(img_curr::Array{Float64, 2})
            s_curr = eqws_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]
            regterm =  0.5*lambda*sum((adaptive_apodizer(img_curr) - adaptive_apodizer(input)).^2)
            lnlik = ( 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_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)

        function dloss2(storage_grad::Array{Float64, 2}, img_curr::Array{Float64, 2})
            s_curr = eqws_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]
            diff = s_curr - s_targ_mean
            #Add branches here
            wt = augment_weights(convert(Array{Float64,1}, reshape(transpose(diff) * s_targ_invcov, (length(diff),))))
            apdimg_curr = adaptive_apodizer(img_curr)
            storage_grad .= reshape((wst_S2_deriv_sum(apdimg_curr, filter_hash, wt, FFTthreads=FFTthreads)' + 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
        epsmag = 1e-4
        eps[row, col] = epsmag
        chisq1 = loss_func2(input+eps./2) #DEB
        chisq0 = loss_func2(input-eps./2) #DEB
        brute  = (chisq1-chisq0)/epsmag
        #df_brute = eqws_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 = dloss2(clever, input) #DEB
        println("Chisq Derve Check")
        println("Brute:  ",brute)
        println("Clever: ",clever[row, col], " Difference: ", brute - clever[row, col], " Mean ", meanval) #DEB
        println("Initial Loss ", loss_func2(input))
        res = optimize(loss_func2, dloss2, input, minmethod, Optim.Options(iterations = numitns_dict, store_trace = show_trace0, show_trace = show_trace0))
        result_img = Optim.minimizer(res)
        println("Final Loss ", loss_func2(result_img))
        return res, reshape(result_img, (Nx, Nx))
    end

    function image_recon_derivsum_regularized_transformed_gaussian(input::Array{Float64, 2}, filter_hash::Dict, fs_targ_mean::Array{Float64, 1}, fs_targ_invcov::Array{Float64, 2}, pixmask::BitArray{2}, dhc_args, func, dfunc;
        FFTthreads::Int=1, optim_settings=Dict([("iterations", 10)]), coeff_mask=nothing, lambda=0.001) #add iso here and a check that returns error if both coeffmask is not nothing and iso is present.
        #=
        f(S) ~ G(fs_targ_mean, fs_targ_invcov)
        =#
        (N1iso, Nf) = size(filter_hash["S1_iso_mat"])
        println("Coeff mask:", (coeff_mask!=nothing))

        (Nx, Ny)  = size(input)
        if Nx != Ny error("Input image must be square") end
        if !dhc_args[:doS20] error("Not implemented") end
        if Nf == 0 error("filter hash corrupted") end

        println("S20")
        #pixmask = pixmask[:] #flattened: Nx^2s #DEB
        ori_input = reshape(copy(input), (Nx^2, 1))

        if coeff_mask!=nothing
            if (!dhc_args[:iso]) & (count((i->i==true), coeff_mask[1:Nf+2])!=0) error("Code assumes S1 coeffs and S0 are masked out") end #wrap
            if (dhc_args[:iso]) & (count((i->i==true), coeff_mask[1:N1iso+2])!=0) error("Code assumes S1iso coeffs and S0 are masked out") end
            #if length(coeff_mask)!= (2+Nf + Nf^2) error("Wrong dim mask") end #wrap
            if size(fs_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
            if size(fs_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_invcov should only have coeffs to be optimized") end
        else #No mask: all coeffs (default: Only S20 coeffs will be optim wrt.)
            error("You must supply a coeff_mask")
        end

        #Apodization adjustments
        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
        #For the Jacobian dA*P/dP
        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

        num_freecoeff = count((i->(i==true)), coeff_mask)

        #After this all cases have a coeffmask, and s_targ_mean and s_targ_invcov have the shapes of the coefficients that we want to select.
        #Optimization Hyperparameters
        tonorm = false
        numitns_dict = get(optim_settings, "iterations", 100)
        minmethod = get(optim_settings, "minmethod", ConjugateGradient())

        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(eqws_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]))
            s_curr = func(eqws_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)
            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)

        function dloss20(storage_grad::Array{Float64, 2}, img_curr::Array{Float64, 2})
            s_curr = eqws_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]
            diff = func(s_curr) - fs_targ_mean
            #Add branches here
            wt = convert(Array{Float64,1}, reshape(transpose(diff) * fs_targ_invcov, (length(diff),))) .* dfunc(s_curr)
            wtaug = augment_weights(wt)
            #println(size(wtaug))
            apdimg_curr = adaptive_apodizer(img_curr)
            storage_grad .= reshape((wst_S20_deriv_sum(apdimg_curr, filter_hash, wtaug, FFTthreads=FFTthreads)' + 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
        epsmag = 1e-6
        eps[row, col] = epsmag
        chisq1 = loss_func20(input+eps./2) #DEB
        chisq0 = loss_func20(input-eps./2) #DEB
        brute  = (chisq1-chisq0)/epsmag
        #df_brute = eqws_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], " Mean ", meanval) #DEB
        println("Initial Loss ", loss_func20(input))
        res = optimize(loss_func20, dloss20, input, minmethod, Optim.Options(iterations = numitns_dict, store_trace = true, show_trace = true))
        result_img = Optim.minimizer(res)
        println("Final Loss ", loss_func20(result_img))
        return res, reshape(result_img, (Nx, Nx))
    end

    function image_recon_derivsum_separate(input::Array{Float64, 2}, iteration_input::Array{Float64, 2}, filter_hash::Dict, s_targ_mean::Array{Float64, 1}, s_targ_invcov, pixmask::BitArray{2}, dhc_args;
        FFTthreads::Int=1, optim_settings=Dict([("iterations", 10)]), coeff_mask=nothing, lambda=0.001) #add iso here and a check that returns error if both coeffmask is not nothing and iso is present.
        #=
        Cases to be handled by this function:
        S20
        Apodized or not
        Iso
        Log or Regular: If log then input, targ and cov should have been constructed using log
        Select pixels masked or all floating--not using for masked currently
        Select coefficients optimized wrt
        Lambda: regularization

        Cases NOT handled by this function:
        Only S1
        Non-Gaussian loss function
        Using the S0 params (mean and variance) in the loss function
        eqws_compute with norm=true
        =#
        (N1iso, Nf) = size(filter_hash["S1_iso_mat"])
        println("Coeff mask:", (coeff_mask!=nothing))

        (Nx, Ny)  = size(input)
        if Nx != Ny error("Input image must be square") end
        if !dhc_args[:doS20] error("Not implemented") end
        if Nf == 0 error("filter hash corrupted") end

        println("S20")
        #pixmask = pixmask[:] #flattened: Nx^2s #DEB
        ori_input = reshape(copy(input), (Nx^2, 1))

        if coeff_mask!=nothing
            if (!dhc_args[:iso]) & (count((i->i==true), coeff_mask[1:Nf+2])!=0) error("Code assumes S1 coeffs and S0 are masked out") end #wrap
            if (dhc_args[:iso]) & (count((i->i==true), coeff_mask[1:N1iso+2])!=0) error("Code assumes S1iso coeffs and S0 are masked out") end
            if size(s_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
            if size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_invcov should only have coeffs to be optimized") end
        else #No mask: all coeffs (default: Only S20 coeffs will be optim wrt.)
            error("You must supply a coeff_mask")
        end

        #Apodization adjustments
        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
        #For the Jacobian dA*P/dP
        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

        num_freecoeff = count((i->(i==true)), coeff_mask)

        #After this all cases have a coeffmask, and s_targ_mean and s_targ_invcov have the shapes of the coefficients that we want to select.
        #Optimization Hyperparameters
        tonorm = false
        numitns_dict = get(optim_settings, "iterations", 100)
        minmethod = get(optim_settings, "minmethod", ConjugateGradient())

        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})
            s_curr = eqws_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]
            regterm =  0.5*lambda*sum((adaptive_apodizer(img_curr) - adaptive_apodizer(input)).^2)
            lnlik = ( 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_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)

        function dloss20(storage_grad::Array{Float64, 2}, img_curr::Array{Float64, 2})
            s_curr = eqws_compute_wrapper(img_curr, filter_hash, norm=tonorm; dhc_args...)[coeff_mask]
            diff = s_curr - s_targ_mean
            #Add branches here
            wt = augment_weights(convert(Array{Float64,1}, reshape(transpose(diff) * s_targ_invcov, (length(diff),))))
            apdimg_curr = adaptive_apodizer(img_curr)
            storage_grad .= reshape((wst_S20_deriv_sum(apdimg_curr, filter_hash, wt, FFTthreads=FFTthreads)' + 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
        epsmag = 1e-4
        eps[row, col] = epsmag
        chisq1 = loss_func20(input+eps./2) #DEB
        chisq0 = loss_func20(input-eps./2) #DEB
        brute  = (chisq1-chisq0)/epsmag
        #df_brute = eqws_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], " Mean ", meanval) #DEB
        println("Initial Loss ", loss_func20(input))
        res = optimize(loss_func20, dloss20, iteration_input, minmethod, Optim.Options(iterations = numitns_dict, store_trace = true, show_trace = true))
        result_img = Optim.minimizer(res)
        println("Final Loss ", loss_func20(result_img))
        return res, reshape(result_img, (Nx, Nx))
    end
end


using Graphs

function create_graph(start_node, end_node)
    @assert length(start_node)==length(end_node)

    no_node = max(maximum(start_node), maximum(end_node))
    no_arc = length(start_node)

    graph = simple_inclist(no_node)
    for i=1:no_arc
        add_edge!(graph, start_node[i], end_node[i])
    end
    return graph
end

function get_vector(state, origin, destination, link_dic)
    current = destination
    parent = -1
    x = zeros(Int, maximum(link_dic))

    while parent != origin
        parent = state.parents[current]

        link_idx = link_dic[parent,current]

        if link_idx != 0
            x[link_idx] = 1
        end

        current = parent
    end

    return x
end

# preparing a graph
graph = create_graph(start_node, end_node);
link_dic = sparse(start_node, end_node, 1:numLinks);

function BPR(flowVec, fcoeffs)
    bpr = similar(flowVec)
    for a = 1:length(bpr)
        bpr[a] = free_flow_time[a] * sum([fcoeffs[i] * (flowVec[a]/capacity[a])^(i-1) for i = 1:length(fcoeffs)])
        if bpr[a] < free_flow_time[a]
            bpr[a] = free_flow_time[a]
        end
    end
    return bpr
end

function all_or_nothing(travel_time, travel_demand_dict)
    state = []
    path = []
    x = zeros(size(start_node))

    for r=1:numNodes
        # for each origin node r, find shortest paths to all destination nodes
        state = dijkstra_shortest_paths(graph, travel_time, r)

        for s=1:numNodes
            # for each destination node s, find the shortest-path vector
            # load travel demand
            x = x + travel_demand_dict[(r,s)] * get_vector(state, r, s, link_dic)
        end
    end

    return x
end

function tapMSA_Multi(travel_demand_car, travel_demand_truck, fcoeffs, numIter=1000, tol=1e-6)

    # Finding a starting feasible solution
    travel_time = BPR(zeros(numLinks), fcoeffs);

    travel_time_car = 1.0 * BPR(zeros(numLinks), fcoeffs);
    travel_time_truck = 1.1 * BPR(zeros(numLinks), fcoeffs);

    xl_car = all_or_nothing(travel_time_car, travel_demand_car);
    xl_truck = all_or_nothing(travel_time_truck, travel_demand_truck);

    l = 1

    while l < numIter
        l += 1

        # Finding yl_car
        xl_car_old = xl_car
        travel_time_car = 1.0 * BPR(xl_car + 2.0 * xl_truck, fcoeffs)
        yl_car = all_or_nothing(travel_time_car, travel_demand_car);
        xl_car = xl_car + (yl_car - xl_car)/l
        xl_car_new = xl_car

        # Finding yl_truck
        xl_truck_old = xl_truck
        travel_time_truck = 1.1 * BPR(xl_car + 2.0 * xl_truck, fcoeffs)
        yl_truck = all_or_nothing(travel_time_truck, travel_demand_truck);
        xl_truck = xl_truck + (yl_truck - xl_truck)/l
        xl_truck_new = xl_truck

        relative_gap_car = norm(xl_car_new - xl_car_old, 1) / norm(xl_car_new, 1)
        relative_gap_truck = norm(xl_truck_new - xl_truck_old, 1) / norm(xl_truck_new, 1)

        if relative_gap_car < tol && relative_gap_truck < tol
            break
        end
    end

    tapFlows = Dict{}()
    tapFlowsCar = Dict{}()
    tapFlowsTruck = Dict{}()
    tapFlowVect = zeros(2,length(xl_car))

    for i = 1:length(ta_data.start_node)
        key = (ta_data.start_node[i], ta_data.end_node[i])
        tapFlowsCar[key] = xl_car[i]
        tapFlowsTruck[key] = xl_truck[i]
    end
    tapFlows["car"] = tapFlowsCar
    tapFlows["truck"] = tapFlowsTruck

    for j = 1:length(xl_car)
        tapFlowVect[1,j] = xl_car[j]
        tapFlowVect[2,j] = xl_truck[j]
    end

    return tapFlows, tapFlowVect

end


# DP formulation for additively separable constraints
# State is Array{Real}: right-hand sides of constraints

include("mathprog_eval.jl")

"""^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Transition function for additively separable constraints

Input:
evals: An array of decomposed terms of a constraint
doms: Array of range representing domain of variables
ords: Function for ordering of the variables

Output: Transition function
"""
function create_constraint_transition_function(evals::SeparableFunctionEvaluation, doms::Array{Range}, ords::Ordering)
	@assert(length(doms) == length(ords)) # number of variables are consistent
	n = length(doms)	# number of variables
	min_ahead = Array{Real}(undef, n)		# For each variable, stores the minimum value that can be obtained starting from the next layer to the terminal
	max_ahead = Array{Real}(undef, n)		# For each variable, stores the maximum value that can be obtained starting from the next layer to the terminal

	#initialize values for the last layer
	min_ahead[get_var(ords,n)] = 0
	max_ahead[get_var(ords,n)] = 0
	for i in n-1:-1:1
		var = get_var(ords,i)		# gets the variable of current layer i
		var_next = get_var(ords,i+1)		# gets the variable in layer i+1
		if isassigned(evals.separate_functions, var_next)		#checks if the function is defined for the variable
			min_v, max_v = min_max_calculator(evals.separate_functions[var_next], doms[var_next])		#computes the min and max of the function over varibale domain
			min_ahead[var] = min_v + min_ahead[var_next]
			max_ahead[var] = max_v + max_ahead[var_next]
		else
			min_ahead[var] = min_ahead[var_next]	#if the function is not defined for the next layer, it does not affect the min max values
			max_ahead[var] = max_ahead[var_next]
		end
	end
	return function (state::Real, variable::Int, value::Int)
		@assert(value in doms[variable])		#ensures that the value belongs to the variable domain

		if isassigned(evals.separate_functions, variable)
			new_state = state + evals.separate_functions[variable](value)		#computes the new state value
		else
			new_state = state 		#if the function is not defined for the variable, the state value remains unchanged
		end

		#To determine feasibility, it is assumed that the constraint form is f(x) <= 0
		if new_state <= -max_ahead[variable]			#new state is feasible no matter what variable assignments are chosen from now on
			new_state = -max_ahead[variable]
		elseif new_state > -min_ahead[variable]			#new state is infeasible no matter what variable assignments are chosen
			new_state = nothing
		end

		return new_state
	end
end

"""^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Initial state for additively separable constraints

Input:
evals: An array of decomposed terms of a constraint

"""
function create_constraint_initial_state(evals::SeparableFunctionEvaluation)::Real
    return eval.constant  	# constant value of the constraint as moved to lhs of <= constraints
end

"""^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Domain function for additively separable constraints

Input:
m: math prog model

Output:
domain function that takes a variable index and returns the domain as range
"""
function create_constraint_domain_function(m::Model)
	nvars = MathProgBase.numvar(m)
	ranges = Array{Range}(undef, nvars)
	for v in 1:nvars
		v_obj = Variable(m, v)
		ranges[v] = ceil(Int, getlowerbound(v_obj)):floor(Int, getupperbound(v_obj))
	end
	return ranges
end


struct SLk0k1lMixingRule <: SLMixingRule
    components::Vector{String}
    k0::PairParam{Float64}
    k1::PairParam{Float64}
    l::PairParam{Float64}
end

@registermodel SLk0k1lMixingRule

function SLk0k1lMixingRule(components; userlocations=String[], verbose=false)
    params = getparams(components, ["LatticeFluid/SanchezLacombe/mixing/k0k1l_unlike.csv"]; userlocations=userlocations, verbose=verbose)
    k0 = params["k0"]
    k1 = params["k1"]
    l = params["l"]
    model = SLk0k1lMixingRule(components,k0,k1,l)
    return model
end

function sl_mix(unmixed_vol,unmixed_epsilon,mixmodel::SLk0k1lMixingRule)
    #dont mind the function names, it performs the correct mixing
    premixed_vol= epsilon_LorentzBerthelot(unmixed_vol,mixmodel.l)
    premixed_epsilon = sigma_LorentzBerthelot(unmixed_epsilon)
    return premixed_vol,premixed_epsilon
end

function mix_vε(model::SanchezLacombe,V,T,z,mix::SLk0k1lMixingRule,r̄,Σz)
    r =  model.params.segment.values
    v = model.params.vol.values
    k0 = mix.k0.values
    k1 = mix.k1.values
    ε = model.params.epsilon.values
    r̄inv = one(r̄)/r̄
     ϕ = @. r* z* r̄inv/Σz
    v_r = zero(V+T+first(z))
    ε_r = v_r
    Σz2 = 1/(Σz*Σz)
    for i in @comps
        for j in @comps
            ϕi = ϕ[i]
            ϕj = ϕ[j]
            ϕiϕj = ϕi*ϕj
            v_r += ϕiϕj*v[i,j]
            δij = (i===j)
            kmi = view(k1,:,i)
            kmj = view(k1,:,j)
            kij = k0[i,j] + (1-δij)* (dot(ϕ,kmi) + dot(ϕ,kmj))

            εij = ε[i,j]*(1-kij)
            ε_r += ϕiϕj*εij
        end
    end
    return v_r,ε_r
end