"""

gradcheck(f, w, x...; gcheck=10, o...)

Numerically checks the gradient of `f(w,x...;o...)` with respect to
its first argument `w`, which could be a Number, Array, Tuple or Dict.
`gcheck` largest entries from each numeric array in gradient
`dw=(grad(f))(w,x...;o...)` is compared to its numerical estimate and
a Warning is issued for comparisons that fail.

"""
function gradcheck(f, w, x...; gcheck=10, o...)
    g = grad(f)
    d = g(w, x...; o...)
    if isa(w, Number)
        gc_number(d, f, w, x...; o...)
    elseif isa(w, KnetArray) || (isa(w, Array) && isbits(eltype(w)))
        gc_array(w, d, f, w, x...; gcheck=gcheck, o...)
    else
        k = gc_indices(w)
        for i in k
            gc_index(w, d, i, f, w, x...; gcheck=gcheck, o...)
        end
    end
end

function gc_index(w, d, i, f, w0, x...; gcheck=10, o...)
    if isa(w[i], Number)
        gc_array(w, d, f, w0, x...; gcheck=1, icheck=i, o...)
    elseif isa(w[i],KnetArray) || (isa(w[i], Array) && isbits(eltype(w[i])))
        gc_array(w[i], d[i], f, w0, x...; gcheck=gcheck, o...)
    else
        k = gc_indices(w[i])
        for j in k
            gc_index(w[i], d[i], j, f, w0, x...; gcheck=gcheck, o...)
        end
    end
end

function gc_array(w, d, f, worig, x...; gcheck=10, icheck=0, o...)
    irange = (icheck > 0 ? (icheck:icheck) :
              length(w) <= gcheck ? (1:length(w)) :
              d == nothing ? rand(1:length(w), gcheck) :
              sortperm(abs(vec(Array(d))),rev=true)[1:gcheck])
    (delta, atol, rtol) = gc_params(typeof(w[first(irange)]))
    for i in irange
        w0 = w[i]
        (w1, w2) = gc_interval(w0, delta)
        w[i] = w1
        f1 = f(worig, x...; o...)
        w[i] = w2
        f2 = f(worig, x...; o...)
        w[i] = w0
        nd = (f2-f1) / (w2-w1)
        di = (d==nothing ? 0 : d[i])
        if !isapprox(di, nd; rtol=rtol, atol=atol)
            warn("d=$di nd=$nd")
        else
            println("gcheck: d=$di nd=$nd")
        end
    end
end

function gc_number(d, f, w, x...; o...)
    (delta, atol, rtol) = gc_params(typeof(w))
    (w1, w2) = gc_interval(w, delta)
    (f1, f2) = (f(w1,x...;o...), f(w2,x...;o...))
    nd = (f2-f1) / (w2-w1)
    if !isapprox(d, nd; rtol=rtol, atol=atol)
        warn("d=$d nd=$nd")
    else
        println("gcheck: d=$d nd=$nd")
    end
end

gc_params(t)=(a=cbrt(eps(t)); (a,a,a))
gc_indices(w)=eachindex(w)
gc_indices(w::Tuple)=(1:length(w))

function gc_interval(w,d)
    w1=w-d/2
    w2=w+d/2
    (w1 < 0 < w) && (w1=zero(w))
    (w2 > 0 > w) && (w2=zero(w))
    return (w1,w2)
end


import ForwardDiff
using ForwardDiff.MAX_CHUNK_SIZE
using Base.Test

const XLEN = MAX_CHUNK_SIZE
const YLEN = div(MAX_CHUNK_SIZE, 2) + 1
const X, Y = rand(XLEN), rand(YLEN)
const CHUNK_SIZES = (1, div(MAX_CHUNK_SIZE, 3), div(MAX_CHUNK_SIZE, 2), MAX_CHUNK_SIZE)
const FINITEDIFF_ERROR = 1e-5

# used to test against results calculated via finite difference
test_approx_eps(a::Array, b::Array) = @test_approx_eq_eps a b EPS

##################
# Test Functions #
##################

# f: Number -> Number #
#---------------------#

deriv_test_1(x) = sin(x)^2 / cos(x)^2
deriv_test_2(x) = 2*x + sqrt(x*x*x)
deriv_test_3(x) = 10.31^(x + x) - x
deriv_test_4(x) = x
deriv_test_5(x) = 1

const NUMBER_TO_NUMBER_FUNCS = tuple(deriv_test_1, deriv_test_2,
                                     deriv_test_3, deriv_test_4,
                                     deriv_test_5)

# f: Number -> Array #
#--------------------#

function deriv_test_6(x)
    return reshape([deriv_test_1(x),
                    deriv_test_2(x),
                    deriv_test_3(x),
                    deriv_test_1(x) - deriv_test_2(x),
                    deriv_test_2(x),
                    deriv_test_3(x),
                    deriv_test_2(x),
                    deriv_test_3(x)], 2, 2, 2)
end

const NUMBER_TO_ARRAY_FUNCS = tuple(deriv_test_6)

# f: Vector -> Number #
#---------------------#

function rosenbrock(x::AbstractVector)
    a = one(eltype(x))
    b = 100 * a
    result = zero(eltype(x))
    for i in 1:length(x)-1
        result += (a - x[i])^2 + b*(x[i+1] - x[i]^2)^2
    end
    return result
end

function ackley(x::AbstractVector)
    a, b, c = 20.0, -0.2, 2.0*π
    len_recip = inv(length(x))
    sum_sqrs = zero(eltype(x))
    sum_cos = sum_sqrs
    for i in x
        sum_cos += cos(c*i)
        sum_sqrs += i^2
    end
    return (-a * exp(b * sqrt(len_recip*sum_sqrs)) -
            exp(len_recip*sum_cos) + a + e)
end

self_weighted_logit(x::AbstractVector) = inv(1.0 + exp(-dot(x, x)))

const VECTOR_TO_NUMBER_FUNCS = tuple(rosenbrock, ackley, self_weighted_logit)

# f: Vector -> Vector #
#---------------------#

# Credit for `chebyquad!`, `brown_almost_linear!`, and `trigonometric!` goes to
# Kristoffer Carlsson (@KristofferC) - I (@jrevels) just ported them over from PR #104.
# I've modified these functions to allow for tunable input/output sizes. These changes
# might make these functions incorrect in terms of their original design, but shouldn't
# be too different computationally (which is what we care about for tests/benchmarks).

function chebyquad!(y::Vector, x::Vector)
    tk = 1/length(x)
    for j = 1:length(x)
        temp1 = 1.0
        temp2 = 2x[j]-1
        temp = 2temp2
        for i = 1:length(y)
            y[i] += temp2
            ti = temp*temp2 - temp1
            temp1 = temp2
            temp2 = ti
        end
    end
    iev = -1.0
    for k = 1:length(y)
        y[k] *= tk
        if iev > 0
            y[k] += 1/(k^2-1)
        end
        iev = -iev
    end
    return y
end

chebyquad(x) = (y = zeros(eltype(x), length(Y)); chebyquad!(y, x); return y)

function brown_almost_linear!(y::Vector, x::Vector)
    c = sum(x) - (length(x) + 1)
    for i = 1:(length(x)-1), j = 1:(length(y)-1)
        y[j] += x[i] + c
    end
    y[length(y)] = prod(x) - 1
    return nothing
end

brown_almost_linear(x) = (y = zeros(eltype(x), length(Y)); brown_almost_linear!(y, x); return y)

function trigonometric!(y::Vector, x::Vector)
    for i in x, j in eachindex(y)
        y[j] = cos(i)
    end
    c = sum(y)
    n = length(x)
    for i in x, j in eachindex(y)
        y[j] = sin(i) * y[j] + n - c
    end
    return y
end

trigonometric(x) = (y = zeros(eltype(x), length(Y)); trigonometric!(y, x); return y)

const VECTOR_TO_VECTOR_FUNCS = tuple(tuple(chebyquad!, chebyquad),
                                     tuple(brown_almost_linear!, brown_almost_linear),
                                     tuple(trigonometric!, trigonometric))


# %% Chaotic scattering
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using PyPlot, DynamicalBilliards
fig, axs = subplots(1,3)
axs[3].axis("off")
ax3 = fig.add_subplot(1, 6, 5)
ax4 = fig.add_subplot(1, 6, 6)
DynamicalBilliards.obcolor(::Obstacle) = matplotlib.colors.to_rgb(COLORS[3])

# First plot: illustration of scattering function
# I GIVE UP, I'll Make this plot in PowerPoint
off = 0.3
r = 0.2
offset = [0.0, off]
center = 2.5
R(φ) = [cos(φ) -sin(φ);
        sin(φ)  cos(φ)]
R3 = R(2π/3)
enclosing = billiard_sinai(r, 2center, 2center)
enclosing = enclosing[2:5]

# disk0 = Disk([center,0], 1.5r, "0")
# sca(axs[1])
# plot(disk0)
# axs[1].set_xlim(-1.5 + center, 0.5 + center)
# axs[1].set_ylim(-1+0.3, 1+0.3)
axs[1].axis("off")

# p0 = Particle(center-1, 0.22, 0)
# bd0 = Billiard(enclosing..., disk0)
# x,y,vx,vy = timeseries(p0, bd0, 2; dt = 0.01)
# axs[1].plot(x[1:130],y[1:130]; color = "C1")
# x0, y0 = p0.pos
# vx0, vy0 = p0.vel
# axs[1].quiver(x0, y0, 0.08vx0, 0.08vy0; angles = "xy", scale = 0.5, width = 0.018, color="C1", zorder = 99)
# axs[1].text(x0, y0-0.2, "\$\\phi\$")
# axs[1].text(x0, y0-0.2, "\$a\$")

# Okay now clarify the scattering function
disk1 = Disk([center, center] .+ offset, r, "green")
disk2 = Disk([center, center] .+ R3*offset, r, "red")
disk3 = Disk([center, center] .+ R3*R3*offset, r, "purple")

disks = Billiard(disk1, disk2, disk3)
plot(disks; ax = axs[2])
axs[2].set_xlim(1.9,3.1)
axs[2].set_ylim(2.0,3.2)
scattering = Billiard(enclosing..., disk1, disk2, disk3)

axs[2].axis("off")

# make some particles

terminate(n, τ, i, p) = i ∈ 1:4
φs = 5π/6 .+ range(0; step = 0.001, length = 3) .+ 0.17

ps = [Particle((R(φ)*[0.5, 0] .+ [center, center])..., φ+π) for φ in φs]

partcol = ( "C3", "C1", "C0",)
for (i, p) in enumerate(ps);
    x0, y0 = p.pos
    vx0, vy0 = p.vel
    axs[2].quiver(x0, y0, 0.04vx0, 0.04vy0; angles = "xy", scale = 0.5, width = 0.01,
    color = partcol[i], zorder = 99)
    x,y = timeseries(p, scattering, terminate)
    axs[2].plot(x, y; color = partcol[i], lw = 2, alpha = 0.75)
end

# Detailed computation of input-output function
φs = range(0, 2π/3; length = 100_000)
θs = zero(φs)

for (i, φ) in enumerate(φs);
    p = Particle((R(φ)*[0.5, 0] .+ [center, center])..., φ+π)
    timeseries!(p, scattering, terminate)
    θs[i] = atan(p.vel[2], p.vel[1])
end

ax3.plot(φs, θs .+ π; ls = "None", marker = "o", ms = 0.5, alpha = 0.5)
ax3.set_xticks([0, 2π/3])
ax3.set_xticklabels(["0", "2π/3"])
ax3.set_yticks([0, 2π])
ax3.set_yticklabels(["0", "2π"])
ax4.set_yticks([])
ax3.set_xlim(0, 2π/3)
ax4.set_ylim(0, 2π)
ax3.set_ylim(0, 2π)

φs = range(0.1555, 0.1567; length = 100_000)
θs = zero(φs)
for (i, φ) in enumerate(φs);
    p = Particle((R(φ)*[0.5, 0] .+ [center, center])..., φ+π)
    timeseries!(p, scattering, terminate)
    θs[i] = atan(p.vel[2], p.vel[1])
end
ax4.plot(φs, θs .+ π; ls = "None", marker = "o", ms = 0.5, alpha = 0.5)
ax4.set_xlim(φs[1], φs[end])
ax4.set_xticks([])
ax3.set_xlabel("\$\\phi_i\$", labelpad = -25)
ax3.set_ylabel("\$\\phi_o = \\mathbf{S}(\\phi_i)\$", labelpad = -25)
# ax3.axvspan(φs[1], φs[end]; color = "C1") # not even visible lol

fig.tight_layout(;pad = 0.25)
fig.subplots_adjust(wspace = 0.1)

# wsave(plotsdir("8", "chaoticscattering"), fig)


module ModelDerivatives

using ForwardDiff
using Morbit
using Test

using Combinatorics

const scalar_test_funcs = [
	x -> sum(x),
	x -> sum( x.^2 ),
	x -> sum( x.^(2 : (length(x) + 1)) ),
	x -> exp( sum( x.^2 ) / 100 ),
	x -> ( 1 / (10 + abs(sum(x))) ) * sin( sum(x) )
]

const configs = [
	ExactConfig(;gradients = :autodiff),
	ExactConfig(;gradients = :fdm),
	LagrangeConfig(; degree = 1 ),
	LagrangeConfig(; degree = 2 ),
	RbfConfig(;kernel = :gaussian),
	RbfConfig(;kernel = :multiquadric),
	RbfConfig(;kernel = :inv_multiquadric),
	RbfConfig(;kernel = :cubic),
	RbfConfig(;kernel = :gaussian, polynomial_degree = -1),
	RbfConfig(;kernel = :inv_multiquadric, polynomial_degree = -1),
	RbfConfig(;kernel = :cubic, polynomial_degree = -1),
	RbfConfig(;kernel = :gaussian, polynomial_degree = 0),
	RbfConfig(;kernel = :inv_multiquadric, polynomial_degree = 0),
	RbfConfig(;kernel = :cubic, polynomial_degree = 0),
	TaylorConfig(;degree=1),
	TaylorConfig(;degree=2),
	TaylorApproximateConfig(mode=:autodiff),
	TaylorApproximateConfig(mode=:fdm)
]

# TODO `TaylorCallbackConfig` cannot be tested this way

NVAR_RANGE = 1:3
NOUT_RANGE = 1:2
#%%

function perform_tests( cfg, n_vars, n_out, k, unconstrained, other_func_indices, x0 = Float64[] )
	f1 = scalar_test_funcs[k]

	lb = fill(-5.0, n_vars )
	ub = fill(6.0, n_vars )

	if unconstrained
		mop = Morbit.MixedMOP(n_vars)
	else
		mop = Morbit.MixedMOP(lb,ub)
	end

	test_func = isempty(other_func_indices) ? f1 : x -> [ f1(x); [ φ(x) for φ in scalar_test_funcs[other_func_indices] ] ]
	Morbit.add_vector_objective!( mop, test_func, cfg; n_out )

	if isempty(x0) x0 = Morbit._rand_box_point(lb, ub) end
	@show x0

	smop, iter_data, data_base, sc, ac = Morbit.initialize_data(mop, x0, Float32[] )

	_mod = sc.surrogates[1].model
	objf = Morbit.list_of_objectives(smop)[1]
	meta = sc.surrogates[1].meta

	# `Morbit.eval_objf` evaluates at a site from the original domain
	_eval_objf_original = x -> Morbit.eval_objf(objf, x)
	_eval_objf_scaled = x -> Morbit.eval_objf( objf, Morbit.unscale(x, smop) )

	# `eval_models`, `get_gradient` etc., on the other hand,
	# take input from [0,1]^n:
	_eval_model = x -> Morbit.eval_models( _mod, x )

	x0_s = Morbit.scale( x0, smop )

	@test x0_s ≈ Morbit.get_x( iter_data )
	@test _eval_objf_original( x0 ) ≈ _eval_objf_scaled( x0_s )

	@test isapprox(Morbit.get_gradient( _mod, x0_s, 1), ForwardDiff.gradient( x -> _eval_model(x)[1], x0_s ); rtol = 0.001 )
	@test isapprox(Morbit.get_jacobian( _mod, x0_s ), ForwardDiff.jacobian( _eval_model, x0_s ); rtol = 0.001 )

	for i = 1 : 10
		ξ0= Morbit._rand_box_point( lb, ub )
		ξ0_s = Morbit.scale( ξ0, smop )
		@test isapprox(Morbit.get_gradient( _mod, ξ0_s, 1), ForwardDiff.gradient( x -> _eval_model(x)[1], ξ0_s ); rtol = 0.001 )
		@test isapprox(Morbit.get_jacobian( _mod, ξ0_s ), ForwardDiff.jacobian( _eval_model, ξ0_s ); rtol = 0.001 )
	end
end#function

for cfg in configs
	for n_vars = NVAR_RANGE
		for n_out = NOUT_RANGE
			for k in eachindex(scalar_test_funcs)
				for unconstrained in [false,true]
					@info n_vars, n_out, k, unconstrained
					@info cfg

					other_func_indices = rand( collect(eachindex(scalar_test_funcs)), n_out - 1 )

#					for other_func_indices in combinations( eachindex(scalar_test_funcs), n_out - 1 )
						perform_tests(cfg, n_vars, n_out, k, unconstrained, other_func_indices )
#					end
				end#for
			end#for
		end#for
	end#for
end#for

end#module

using .ModelDerivatives


@testset "categorical sampler" begin

	cc = dummyChoiceContext()

	@testset "getparams, numparams, paramranges, sample" begin

		@testset "using default construction" begin

			s = DataGenerators.CategoricalSampler(4)

			@testset "numparams and paramranges" begin
			    @test DataGenerators.numparams(s) == 4
			    prs = DataGenerators.paramranges(s)
			    @test typeof(prs) <: Vector{Tuple{Float64,Float64}}
			    @test prs == fill((0.0,1.0), DataGenerators.numparams(s))
			end

			@testset "default params" begin
			    @test DataGenerators.getparams(s) == [0.25, 0.25, 0.25, 0.25]
			end

			@mtestset "default sampling" reps=Main.REPS alpha=Main.ALPHA begin
			    x, trace = DataGenerators.sample(s, (0,1), cc)
			    @test typeof(x) <: Int
			    @mtest_values_are [1,2,3,4] x
		        @mtest_distributed_as Categorical([0.25,0.25,0.25,0.25]) x
			end

		end

		@testset "non-default construction" begin

			s = DataGenerators.CategoricalSampler(5, [0.3,0.2,0.1,0.2,0.2])

			@testset "constructor with params" begin

			    @test getparams(s) == [0.3,0.2,0.1,0.2,0.2]

			    @mtestset "consistent with categorical" reps=Main.REPS alpha=Main.ALPHA begin
	            	x, trace = DataGenerators.sample(s, (0,1), cc)
			        @mtest_distributed_as Categorical([0.3,0.2,0.1,0.2,0.2]) x
			    end

			end

		end

	end

	@testset "setparams" begin

		s = DataGenerators.CategoricalSampler(4)
		prs = DataGenerators.paramranges(s)
		midparams = map(pr->robustmidpoint(pr[1],pr[2]), prs)

		@testset "valid parameters $params" for params in [[0.1,0.2,0.3,0.4], [0.4,0.1,0.1,0.4], [0.3,0.2,0.3,0.2],]

		    DataGenerators.setparams!(s, params)

		    @mtestset "consistent with categorical" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s, (0,1), cc)
		        @mtest_distributed_as Categorical(params) x
			end

		end

		@testset "boundary values for parameter index $pidx" for pidx in 1:length(prs)

		    pr = prs[pidx]
		    params = copy(midparams)

			@testset "bound index $bidx" for bidx in 1:2

		        params[pidx] = pr[bidx]
		        DataGenerators.setparams!(s, params)

			    @mtestset "consistent with categorical" reps=Main.REPS alpha=Main.ALPHA begin
		        	x, trace = DataGenerators.sample(s, (0,1), cc)
			        @mtest_distributed_as  Categorical(params./sum(params)) x
				end

				@testset "range check exception" begin
		            params[pidx] = bidx == 1 ? prevfloat(pr[bidx]) : nextfloat(pr[bidx])
		            @test_throws ErrorException DataGenerators.setparams!(s, params)
				end

			end

		end

		@testset "setparams normalises weights" begin

		    DataGenerators.setparams!(s, [0.4, 0.6, 0.7, 0.3])
		    @test getparams(s) == [0.2, 0.3, 0.35, 0.15]

		    @mtestset "consistent with categorical" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s, (0,1), cc)
		        @mtest_distributed_as Categorical([0.2,0.3,0.35,0.15]) x
			end

		end

		@testset "setparams adjusts when all weights are zero" begin

		    DataGenerators.setparams!(s, [0.0, 0.0, 0.0, 0.0])
		    @test getparams(s) == [0.25, 0.25, 0.25, 0.25]

		    @mtestset "consistent with categorical" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s, (0,1), cc)
		        @mtest_distributed_as Categorical([0.25,0.25,0.25,0.25]) x
			end

		end

		@testset "setparams with wrong number of parameters" begin
		    @test_throws ErrorException DataGenerators.setparams!(s, midparams[1:end-1])
		    @test_throws ErrorException DataGenerators.setparams!(s, [midparams; 0.5])
		end

	end

	@testset "estimateparams" begin

		otherparams = [0.35, 0.15, 0.1, 0.25, 0.15]

		@testset "equal bounds" begin

		    params = [0.2, 0.2, 0.2, 0.2, 0.2]
		    s1 = DataGenerators.CategoricalSampler(5, params)
		    s2 = DataGenerators.CategoricalSampler(5, otherparams)
		    traces = map(1:100) do i
		        x, trace = DataGenerators.sample(s1, (0,1), cc)
		        trace
		    end
		    estimateparams!(s2, traces)

		    @mtestset "consistent with categorical" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s2, (0,1), cc)
		        @mtest_distributed_as Categorical(params) x
			end

		end

		@testset "last category never sampled" begin

		    params = [0.3, 0.4, 0.2, 0.1, 0.0]
		    s1 = DataGenerators.CategoricalSampler(5, params)
		    s2 = DataGenerators.CategoricalSampler(5, otherparams)
		    traces = map(1:100) do i
		        x, trace = DataGenerators.sample(s1, (0,1), cc)
		        trace
		    end
		    estimateparams!(s2, traces)

		    @mtestset "consistent with categorical" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s2, (0,1), cc)
		        @mtest_distributed_as Categorical(params) x
			end

		end

		@testset "too few traces" begin

		    params = [0.1, 0.2, 0.4, 0.2, 0.1]
		    s1 = DataGenerators.CategoricalSampler(5, params)
		    s2 = DataGenerators.CategoricalSampler(5, otherparams)
		    traces = map(1:0) do i
		        x, trace = DataGenerators.sample(s1, (0,1), cc)
		        trace
		    end
		    estimateparams!(s2, traces)

		    @mtestset "consistent with categorical" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s2, (0,1), cc)
		        @mtest_distributed_as Categorical(otherparams) x
			end

		end

	end

end


using DiffEqBase, ParameterizedFunctions, DiffEqBiological, Markdown
#SDE Example Problems
export prob_sde_wave, prob_sde_linear, prob_sde_cubic, prob_sde_2Dlinear,
      prob_sde_lorenz, prob_sde_2Dlinear, prob_sde_additive,
      prob_sde_additivesystem, oval2ModelExample, prob_sde_bistable,
      prob_sde_bruss, prob_sde_oscilreact
#SDE Stratonovich Example Problems
export prob_sde_linear_stratonovich, prob_sde_2Dlinear_stratonovich

### SDE Examples

f_linear(u,p,t) = 1.01u
σ_linear(u,p,t) = 0.87u
linear_analytic(u0,p,t,W) = @.(u0*exp(0.63155t+0.87W))

@doc doc"""
```math
du_t = βudt + αudW_t
```
where β=1.01, α=0.87, and initial condtion u0=1/2, with solution

```math
u(u0,p,t,W_t)=u0\\exp((α-\\frac{β^2}{2})t+βW_t)
```

"""
prob_sde_linear = SDEProblem(SDEFunction(f_linear,σ_linear,
                             analytic=linear_analytic),σ_linear,1/2,(0.0,1.0))

linear_analytic_strat(u0,p,t,W) = @.(u0*exp(1.01t+0.87W))
prob_sde_linear_stratonovich = SDEProblem(SDEFunction(f_linear,σ_linear,
                               analytic=linear_analytic_strat),
                               σ_linear,1/2,(0.0,1.0))
f_linear_iip(du,u,p,t) = @.(du = 1.01*u)
σ_linear_iip(du,u,p,t) = @.(du = 0.87*u)
@doc doc"""
8 linear SDEs (as a 4x2 matrix):

```math
du_t = βudt + αudW_t
```
where β=1.01, α=0.87, and initial condtion u0=1/2 with solution

```math
u(u0,p,t,W_t)=u0\\exp((α-\\frac{β^2}{2})t+βW_t)
```
"""
prob_sde_2Dlinear = SDEProblem(SDEFunction(f_linear_iip,σ_linear_iip,
                               analytic=linear_analytic),
                               σ_linear_iip,ones(4,2)/2,(0.0,1.0))
prob_sde_2Dlinear_stratonovich = SDEProblem(SDEFunction(f_linear_iip,σ_linear_iip,
                               analytic=linear_analytic_strat),
                               σ_linear_iip,ones(4,2)/2,(0.0,1.0))

f_cubic(u,p,t) = -.25*u*(1-u^2)
σ_cubic(u,p,t) = .5*(1-u^2)
cubic_analytic(u0,p,t,W) = @. ((1+u0)*exp(W)+u0-1)/((1+u0)*exp(W)+1-u0)
ff_cubic = SDEFunction(f_cubic,σ_cubic,analytic = cubic_analytic)
@doc doc"""
```math
du_t = \\frac{1}{4}u(1-u^2)dt + \\frac{1}{2}(1-u^2)dW_t
```

and initial condtion u0=1/2, with solution

```math
u(u0,p,t,W_t)=\\frac{(1+u0)\\exp(W_t)+u0-1}{(1+u0)\\exp(W_t)+1-u0}
```
"""
prob_sde_cubic = SDEProblem(ff_cubic,σ_cubic,1/2,(0.0,1.0))

f_wave(u,p,t) = @. -0.01*sin(u)*cos(u)^3
σ_wave(u,p,t) = @. 0.1*cos(u)^2
wave_analytic(u0,p,t,W) = @. atan(0.1*W + tan(u0))
ff_wave = SDEFunction(f_wave,σ_wave,analytic=wave_analytic)
@doc doc"""
```math
du_t = -\\frac{1}{100}\sin(u)\cos^3(u)dt + \\frac{1}{10}\cos^{2}(u_t) dW_t
```

and initial condition `u0=1.0` with solution

```math
u(u0,p,t,W_t)=\\arctan(\\frac{W_t}{10} + \\tan(u0))
```
"""
prob_sde_wave = SDEProblem(ff_wave,σ_wave,1.0,(0.0,1.0))

f_additive(u,p,t) = @. p[2]/sqrt(1+t) - u/(2*(1+t))
σ_additive(u,p,t) = @. p[1]*p[2]/sqrt(1+t)
p = (0.1,0.05)
additive_analytic(u0,p,t,W) = @. u0/sqrt(1+t) + p[2]*(t+p[1]*W)/sqrt(1+t)
ff_additive = SDEFunction(f_additive,σ_additive,analytic=additive_analytic)
@doc doc"""
Additive noise problem

```math
u_t = (\\frac{β}{\\sqrt{1+t}}-\\frac{1}{2(1+t)}u_t)dt + \\frac{αβ}{\\sqrt{1+t}}dW_t
```

and initial condition u0=1.0 with α=0.1 and β=0.05, with solution

```math
u(u0,p,t,W_t)=\\frac{u0}{\\sqrt{1+t}} + \\frac{β(t+αW_t)}{\\sqrt{1+t}}
```
"""
prob_sde_additive = SDEProblem(ff_additive,σ_additive,1.0,(0.0,1.0),p)

f_additive_iip(du,u,p,t) = @.(du = p[2]/sqrt(1+t) - u/(2*(1+t)))
σ_additive_iip(du,u,p,t) = @.(du = p[1]*p[2]/sqrt(1+t))
ff_additive_iip = SDEFunction(f_additive_iip,σ_additive_iip,analytic=additive_analytic)
p = ([0.1;0.1;0.1;0.1],[0.5;0.25;0.125;0.1115])
@doc doc"""
A multiple dimension extension of `additiveSDEExample`

"""
prob_sde_additivesystem = SDEProblem(ff_additive_iip,σ_additive_iip,[1.;1.;1.;1.],(0.0,1.0),p)

f_lorenz = @ode_def_bare LorenzSDE begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ ρ β
σ_lorenz(du,u,p,t) = @.(du = 3.0)
@doc doc"""
Lorenz Attractor with additive noise

```math
\\begin{align}
dx &= σ*(y-x)dt + αdW_t \\\\
dy &= (x*(ρ-z) - y)dt + αdW_t \\\\
dz &= (x*y - β*z)dt + αdW_t \\\\
\\end{align}
```

with ``σ=10``, ``ρ=28``, ``β=8/3``, ``α=3.0`` and inital condition ``u0=[1;1;1]``.
"""
prob_sde_lorenz = SDEProblem(f_lorenz,σ_lorenz,ones(3),(0.0,10.0),(10.0,28.0,2.66))

f_nltest(u,p,t) = (1/3)*u^(1/3) + 6*u^(2/3)
σ_nltest(u,p,t) = u^(2/3)
analytic_nltest(u0,p,t,W) = (2t + 1 + W/3)^3
ff_nltest = SDEFunction(f_nltest,σ_nltest,analytic=analytic_nltest)
@doc doc"""
Runge–Kutta methods for numerical solution of stochastic differential equations
Tocino and Ardanuy
"""
prob_sde_nltest = SDEProblem(ff_nltest,σ_nltest,1.0,(0.0,10.0))

function oval2ModelExample(;largeFluctuations=false,useBigs=false,noiseLevel=1)
  #Parameters
  J1_200=3.
  J1_34=0.15
  J2_200=0.2
  J2_34=0.35
  J_2z=0.9
  J_O=0.918
  J_SO=0.5
  # J_ZEB=0.06
  J_ecad1=0.1
  J_ecad2=0.3
  J_ncad1=0.4
  J_ncad2=0.4
  J_ncad3 = 2
  J_snail0=0.6
  J_snail1=1.8
  J_zeb=3.0
  K1=1.0
  K2=1.0
  K3=1.0
  K4=1.0
  K5=1.0
  KTGF=20.
  Ks=100.
  # TGF0=0
  TGF_flg=0.
  Timescale=1000.
  dk_ZR1=0.5
  dk_ZR2=0.5
  dk_ZR3=0.5
  dk_ZR4=0.5
  dk_ZR5=0.5
  k0O=0.35
  k0_200=0.0002
  k0_34=0.001
  k0_snail=0.0005
  k0_zeb=0.003
  kO=1.2
  kOp=10.
  k_200=0.02
  k_34=0.019
  k_OT=1.1
  k_SNAIL=16.
  k_TGF=1.5
  k_ZEB=16.
  k_ecad0=5.
  k_ecad1=15.
  k_ecad2=5.
  k_ncad0=5.
  k_ncad1=2.
  k_ncad2=5.
  k_snail=0.05
  k_tgf=0.05
  k_zeb=0.06
  kdO=1.
  kd_200=0.035
  kd_34=0.035
  kd_SNAIL=1.6
  kd_SR1=0.9
  kd_TGF=0.9
  kd_ZEB=1.66
  kd_ecad=0.05
  kd_ncad=0.05
  kd_snail=0.09
  kd_tgf=0.1
  kd_tgfR=1.0
  kd_zeb=0.1
  kd_Op = 10.
  lamda1=0.5
  lamda2=0.5
  lamda3=0.5
  lamda4=0.5
  lamda5=0.5
  lamdas=0.5
  lamdatgfR=0.8
  nO=6.
  nSO=2.
  nzo=2.
  GE = 1.
  f = function (dy,y,p,t)
    # y(1) = snailt
    # y(2) = SNAIL
    # y(3) = miR34t
    # y(4) = SR1 # abundance of SNAIL/miR34 complex
    # y(5) = zebt
    # y(6) = ZEB
    # y(7) = miR200t
    # y(8) = ZR1 # abundance of ZEB/miR200 complex with i copies of miR200 bound on the sequence of ZEB1
    # y(9) = ZR2
    # y(10) = ZR3
    # y(11) = ZR4
    # y(12) = ZR5
    # y(13) = tgft
    # y(14) = TGF
    # y(15) = tgfR # abundance of TGF/miR200 complex
    # y(16) = Ecad
    # y(17) = Ncad
    # y(18) = Ovol2
    TGF0=.5(t>100)
    #ODEs
    dy[1]=k0_snail+k_snail*(((y[14]+TGF0)/J_snail0))^2/(1+(((y[14]+TGF0)/J_snail0))^2+(y[19]/J_SO)^nSO)/(1+y[2]/J_snail1)-kd_snail*(y[1]-y[4])-kd_SR1*y[4]
    dy[2]=k_SNAIL*(y[1]-y[4])-kd_SNAIL*y[2]
    dy[3]=k0_34+k_34/(1+((y[2]/J1_34))^2+((y[6]/J2_34))^2)-kd_34*(y[3]-y[4])-kd_SR1*y[4]+lamdas*kd_SR1*y[4]
    dy[4]=Timescale*(Ks*(y[1]-y[4])*(y[3]-y[4])-y[4])
    dy[5]=k0_zeb+k_zeb*((y[2]/J_zeb))^2/(1+((y[2]/J_zeb))^2+((y[19]/J_2z))^nO)-kd_zeb*(y[5]-(5*y[8]+10*y[9]+10*y[10]+5*y[11]+y[12]))-dk_ZR1*5*y[8]-dk_ZR2*10*y[9]-dk_ZR3*10*y[10]-dk_ZR4*5*y[11]-dk_ZR5*y[12]
    dy[6]=k_ZEB*(y[5]-(5*y[8]+10*y[9]+10*y[10]+5*y[11]+y[12]))-kd_ZEB*y[6]
    dy[7]=k0_200+k_200/(1+((y[2]/J1_200))^3+((y[6]/J2_200))^2)-kd_200*(y[7]-(5*y[8]+2*10*y[9]+3*10*y[10]+4*5*y[11]+5*y[12])-y[15])-dk_ZR1*5*y[8]-dk_ZR2*2*10*y[9]-dk_ZR3*3*10*y[10]-dk_ZR4*4*5*y[11]-dk_ZR5*5*y[12]+lamda1*dk_ZR1*5*y[8]+lamda2*dk_ZR2*2*10*y[9]+lamda3*dk_ZR3*3*10*y[10]+lamda4*dk_ZR4*4*5*y[11]+lamda5*dk_ZR5*5*y[12]-kd_tgfR*y[15]+lamdatgfR*kd_tgfR*y[15]
    dy[8]=Timescale*(K1*(y[7]-(5*y[8]+2*10*y[9]+3*10*y[10]+4*5*y[11]+5*y[12])-y[15])*(y[5]-(5*y[8]+10*y[9]+10*y[10]+5*y[11]+y[12]))-y[8])
    dy[9]=Timescale*(K2*(y[7]-(5*y[8]+2*10*y[9]+3*10*y[10]+4*5*y[11]+5*y[12])-y[15])*y[8]-y[9])
    dy[10]=Timescale*(K3*(y[7]-(5*y[8]+2*10*y[9]+3*10*y[10]+4*5*y[11]+5*y[12])-y[15])*y[9]-y[10])
    dy[11]=Timescale*(K4*(y[7]-(5*y[8]+2*10*y[9]+3*10*y[10]+4*5*y[11]+5*y[12])-y[15])*y[10]-y[11])
    dy[12]=Timescale*(K5*(y[7]-(5*y[8]+2*10*y[9]+3*10*y[10]+4*5*y[11]+5*y[12])-y[15])*y[11]-y[12])
    dy[13]=k_tgf-kd_tgf*(y[13]-y[15])-kd_tgfR*y[15]
    dy[14]=k_OT+k_TGF*(y[13]-y[15])-kd_TGF*y[14]
    dy[15]=Timescale*(TGF_flg+KTGF*(y[7]-(5*y[8]+2*10*y[9]+3*10*y[10]+4*5*y[11]+5*y[12])-y[15])*(y[13]-y[15])-y[15])
    dy[16]=GE*(k_ecad0+k_ecad1/(((y[2]/J_ecad1))^2+1)+k_ecad2/(((y[6]/J_ecad2))^2+1)-kd_ecad*y[16])
    dy[17]=k_ncad0+k_ncad1*(((y[2]/J_ncad1))^2)/(((y[2]/J_ncad1))^2+1)+k_ncad2*(((y[6]/J_ncad2))^2)/(((y[6]/J_ncad2)^2+1)*(1+y[19]/J_ncad3))-kd_ncad*y[17]
    dy[18]=k0O+kO/(1+((y[6]/J_O))^nzo)-kdO*y[18]
    dy[19]=kOp*y[18]-kd_Op*y[19]
  end

  σ1 = function (dσ,y,p,t)
    dσ[1] = noiseLevel*1.5y[1]
    dσ[18]= noiseLevel*6y[18]
  end

  σ2 = function (dσ,y,p,t)
    dσ[1] = 0.02y[1]
    dσ[16]= 0.02y[16]
    dσ[18]= 0.2y[18]
    dσ[17]= 0.02y[17]
  end

  if largeFluctuations
    σ = σ1
  else
    σ = σ2
  end

  u0 = [0.128483;1.256853;0.0030203;0.0027977;0.0101511;0.0422942;0.2391346;0.0008014;0.0001464;2.67e-05;4.8e-6;9e-7;0.0619917;1.2444292;0.0486676;199.9383546;137.4267984;1.5180203;1.5180203] #Fig 9B
  if useBigs
    u0 = big.(u0)
  end
  #u0 =  [0.1701;1.6758;0.0027;0.0025;0.0141;0.0811;0.1642;0.0009;0.0001;0.0000;0.0000;0.0000;0.0697;1.2586;0.0478;194.2496;140.0758;1.5407;1.5407] #Fig 9A
  SDEProblem(f,σ,u0,(0.0,500.0))
end

stiff_quad_f_ito(u,p,t) = -(p[1]+(p[2]^2)*u)*(1-u^2)
stiff_quad_f_strat(u,p,t) = -p[1]*(1-u^2)
stiff_quad_g(u,p,t) = p[2]*(1-u^2)

function stiff_quad_f_ito_analytic(u0,p,t,W)
    α = p[1]
    β = p[2]
    exp_tmp = exp(-2*α*t+2*β*W)
    tmp = 1 + u0
    (tmp*exp_tmp + u0 - 1)/(tmp*exp_tmp - u0 + 1)
end

ff_stiff_quad_ito = SDEFunction(stiff_quad_f_ito,stiff_quad_g,
                                analytic=stiff_quad_f_ito_analytic)

function stiff_quad_f_strat_analytic(u0,p,t,W)
    α = p[1]
    β = p[2]
    exp_tmp = exp(-2*α*t+2*β*W)
    tmp = 1 + u0
    (tmp*exp_tmp + u0 - 1)/(tmp*exp_tmp - u0 + 1)
end

ff_stiff_quad_strat = SDEFunction(stiff_quad_f_strat,stiff_quad_g,
                                analytic=stiff_quad_f_strat_analytic)

@doc doc"""
The composite Euler method for stiff stochastic
differential equations

Kevin Burrage, Tianhai Tian

And

S-ROCK: CHEBYSHEV METHODS FOR STIFF STOCHASTIC
DIFFERENTIAL EQUATIONS

ASSYR ABDULLE AND STEPHANE CIRILLI

Stiffness of Euler is determined by α+β²<1
Higher α or β is stiff, with α being deterministic stiffness and
β being noise stiffness (and grows by square).
"""
prob_sde_stiffquadito = SDEProblem(ff_stiff_quad_ito,stiff_quad_g,0.5,(0.0,3.0),(1.0,1.0))

@doc doc"""
The composite Euler method for stiff stochastic
differential equations

Kevin Burrage, Tianhai Tian

And

S-ROCK: CHEBYSHEV METHODS FOR STIFF STOCHASTIC
DIFFERENTIAL EQUATIONS

ASSYR ABDULLE AND STEPHANE CIRILLI

Stiffness of Euler is determined by α+β²<1
Higher α or β is stiff, with α being deterministic stiffness and
β being noise stiffness (and grows by square).
"""
prob_sde_stiffquadstrat = SDEProblem(ff_stiff_quad_strat,stiff_quad_g,0.5,(0.0,3.0),(1.0,1.0))

@doc doc"""
Stochastic Heat Equation with scalar multiplicative noise

S-ROCK: CHEBYSHEV METHODS FOR STIFF STOCHASTIC
DIFFERENTIAL EQUATIONS

ASSYR ABDULLE AND STEPHANE CIRILLI

Raising D or k increases stiffness
"""
function generate_stiff_stoch_heat(D=1,k=1;N = 100)
    A = full(Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1]))
    dx = 1/N
    A = D/(dx^2) * A
    function f(du,u,p,t)
        mul!(du,A,u)
    end
    #=
    function f(::Type{Val{:analytic}},u0,p,t,W)
        expm(A*t+W*I)*u0
    end
    =#
    function g(du,u,p,t)
        @. du = k*u
    end
    SDEProblem(f,g,ones(N),(0.0,3.0),noise=WienerProcess(0.0,0.0,0.0,rswm=RSWM(adaptivealg=:RSwM3)))
end

bistable_f(du,u,p,t) = du[1] = p[1] + p[2]*u[1]^4/(u[1]^4+11.9^4) - p[3]*u[1]
function bistable_g(du,u,p,t)
  du[1,1] = .1*sqrt(p[1] + p[2]*u[1]^4/(u[1]^4+11.9^4))
  du[1,2] = -.1*sqrt(p[3]*u[1])
end
p = (5.0,18.0,1.0)
"""
Bistable chemical reaction network with a semi-stable lower state.
"""
prob_sde_bistable = SDEProblem(bistable_f,bistable_g,[3.],(0.,300.),p,
                               noise_rate_prototype=zeros(1,2))

function bruss_f(du,u,p,t)
    du[1] = p[1] + p[2]*u[1]*u[1]*u[2] - p[3]*u[1]-p[4]*u[1]
    du[2] = p[3]*u[1] - p[2]*u[1]*u[1]*u[2]
end
function bruss_g(du,u,p,t)
  du[1,1] = 0.15*sqrt(p[1])
  du[1,2] = 0.15*sqrt(p[2]*u[1]*u[1]*u[2])
  du[1,3] = -0.15*sqrt(p[3]*u[1])
  du[1,4] = -0.15*sqrt(p[4]*u[1])
  du[2,1] = 0
  du[2,2] = -0.15*sqrt(p[2]*u[1]*u[1]*u[2])
  du[2,3] = 0.15*sqrt(p[3]*2.5*u[1])
  du[2,4] = 0
end
p = (1.0,1.0,2.5,1.0)
"""
Stochastic Brusselator
"""
prob_sde_bruss = SDEProblem(bruss_f,bruss_g,[3.,2.],(0.,100.),p,
                            noise_rate_prototype=zeros(2,4))

network = @reaction_network rnType  begin
    p1, (X,Y,Z) --> 0
    hill(X,p2,100.,-4), 0 --> Y
    hill(Y,p3,100.,-4), 0 --> Z
    hill(Z,p4,100.,-4), 0 --> X
    hill(X,p5,100.,6), 0 --> R
    hill(Y,p6,100.,4)*0.002, R --> 0
    p7, 0 --> S
    R*p8, S --> SP
    p9, SP + SP --> SP2
    p10, SP2 --> 0
end p1 p2 p3 p4 p5 p6 p7 p8 p9 p10
p = (0.01,3.0,3.0,4.5,2.,15.,20.0,0.005,0.01,0.05)

"""
An oscillatory chemical reaction system
"""
prob_sde_oscilreact = SDEProblem(network, [200.,60.,120.,100.,50.,50.,50.], (0.,4000.),p)


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.7.0
#     language: julia
#     name: julia-1.7
# ---

# %% [markdown]
# # ベイズハッキング
#
# * 黒木玄
# * 2021-12-04
# * [nbviewer](https://nbviewer.org/github/genkuroki/public/blob/main/0025/Bayes%20hacking.ipynb)
#
# 関連ノートブック
#
# * [止め方で結果が変わる？](https://nbviewer.org/github/genkuroki/public/blob/main/0025/%E6%AD%A2%E3%82%81%E6%96%B9%E3%81%A7%E7%B5%90%E6%9E%9C%E3%81%8C%E5%A4%89%E3%82%8F%E3%82%8B%EF%BC%9F.ipynb)
# * [Clopper-Pearsonの信頼区間](https://nbviewer.org/github/genkuroki/public/blob/main/0025/Clopper-Pearson%20confidence%20interval.ipynb)

# %%
using Distributions
using StatsPlots
using StatsFuns: logit, logistic
using Roots
using StatsBase: ecdf
using Memoization
using SpecialFunctions

using Base64
showimg(mime, fn; tag="img") = open(fn) do f
    base64 = base64encode(f)
    display("text/html", """<$tag src="data:$mime;base64,$base64" />""")
end

safemul(x, y) = iszero(x) ? x : x*y
x ⪅ y = x < y || x ≈ y

# %%
showimg("image/jpeg", "IMG_2864.jpg"; tag="img width=50%")

# %%
showimg("image/jpeg", "IMG_2865.jpg"; tag="img width=60%")

# %%
"""
事前分布 Beta(a, b) のもとでn回中ｋ回成功の場合の事後分布は Beta(k+a, n-k+b) になる。
デフォルトでは a = b = 1 の一様事前分布を採用することにする。
"""
posterior_dist(n, k; a=1, b=a) = Beta(k + a, n - k + b)

"""pvalue_bayes(n, p, k; a, b) < α ⇔ p ∉ (ベイズ版100(1-α)%信用区間)"""
@memoize function pvalue_bayes(n, p, k; a = 1, b = a)
    posterior = posterior_dist(n, k; a, b)
    v0 = logpdf(posterior, p)
    f(t) = logpdf(posterior, logistic(t)) - v0
    m = params(posterior) |> ((α, β),) -> (α - 1)/(α + β - 2)
    if m ≤ 0
        s = ccdf(posterior, p)
    elseif m ≥ 1
        s = cdf(posterior, p)
    elseif p ≤ m
        q = logistic(find_zero(f, min(50, logit(m) + 1)))
        s = cdf(posterior, p) + ccdf(posterior, q)
    else
        q = logistic(find_zero(f, max(-50, logit(m) - 1)))
        s = cdf(posterior, q) + ccdf(posterior, p)
    end
    min(1, s)
end

"""
Bernoulli(p)分布の独立試行をNまで行うとき、
pがベイズ版信用区間から外れるまでの試行回数をL個計算
"""
function bayeshacking(N, p; α = 0.05, L = 10^4, a = 1, b = a)
    numtrials = fill(N + 1, L) # N＋１で埋め尽くされた長さLのベクトル
    Threads.@threads for i in 1:L # 並列のでL回のシミュレーションを実行
        k = 0
        for n in 1:N # ベルヌイ試行回数nの最大値はN
            k += rand(Bernoulli(p)) # kに1または0を足す。1を足す確率はp.
            if pvalue_bayes(n, p, k; a, b) < α # もしもpの値がベイズ信用区間を外れたら
                numtrials[i] = n  # n回で目的を達成したことを記録に残す
                break
            end
        end
    end
    numtrials
end

# %%
numtrials_bhack = bayeshacking(10^3, 0.5)
[(n, ecdf(numtrials_bhack)(n)) for n in (10, 100, 1000)]

# %%
plot(n -> ecdf(numtrials_bhack)(n), 0, 10^3; label="", xtick=0:100:1000)
title!("ecdf(number of trials for Bayes hacking)"; titlefontsize=12)

# %%
"""通常のP値函数(exact版)"""
@memoize function pvalue_exact(n, p, k)
    bin = Binomial(n, p)
    p0 = pdf(bin, k)
    s = sum(pdf(bin, j) for j in support(bin) if pdf(bin, j) ⪅ p0)
    min(1, s)
end

"""通常のP値函数(正規分布近似)"""
@memoize function pvalue_normal(n, p, k)
    bin = Binomial(n, p)
    normal = Normal(mean(bin), std(bin))
    min(1, 2cdf(normal, k), 2ccdf(normal, k))
end

"""
Bernoulli(p)分布の独立試行をNまで行うとき、
P値がα未満になるまでの試行回数をL個計算
"""
function phacking(pvalue_func, N, p; α = 0.05, L = 10^4)
    numtrials = fill(N + 1, L) # N＋１で埋め尽くされた長さLのベクトル
    Threads.@threads for i in 1:L # 並列のでL回のシミュレーションを実行
        k = 0
        for n in 1:N # ベルヌイ試行回数nの最大値はN
            k += rand(Bernoulli(p)) # kに1または0を足す。1を足す確率はp.
            if pvalue_func(n, p, k) < α # もしもpのP値が有意水準を下回ったら
                numtrials[i] = n # n回で目的を達成したことを記録に残す
                break
            end
        end
    end
    numtrials
end

# %%
numtrials_phack_exact = phacking(pvalue_exact, 10^3, 0.5)
[(n, ecdf(numtrials_phack_exact)(n)) for n in (10, 100, 1000)]

# %%
numtrials_phack_normal = phacking(pvalue_normal, 10^3, 0.5)
[(n, ecdf(numtrials_phack_normal)(n)) for n in (10, 100, 1000)]

# %%
"""ベイズ信頼区間ハックをプロット"""
function plot_bhack(; N=10^3, p=0.5, α = 0.05, L = 10^4, a = 1, b = a, kwargs...)
    numtrials_bhack = bayeshacking(N, p; α, L, a, b)
    numtrials_phack_exact = phacking(pvalue_exact, N, p; α, L)
    numtrials_phack_normal = phacking(pvalue_normal, N, p; α, L)

    P = plot(; legend=:bottomright)
    plot!(n -> ecdf(numtrials_bhack)(n), 0, N; label="Bayes hacking")
    plot!(n -> ecdf(numtrials_phack_exact)(n), 0, N; label="p-hacking (exact p-value)", ls=:dash)
    plot!(n -> ecdf(numtrials_phack_normal)(n), 0,N; label="p-hacking (normal dist. approx.)", ls=:dashdot)
    title!("ecdf(number of trials),  α = $α"; titlefontsize=12)
    plot!(; xlabel="n", ylabel="P(number of trials when stopped ≤ n)", guidefontsize=10)
    plot!(; kwargs...)

    Q = plot(Beta(a, b); ylim=(0, maximum(x -> pdf(Beta(a, b), x), 0.001:0.001:0.999) + 1), label="")
    plot!(; xtick=0:0.1:1, xlabel="parameter", ylabel="probability density")
    title!("prior = Beta($a, $b)"; titlefontsize=12)

    plot(P, Q; size=(640, 800), layout=(2, 1), leftmargin=2Plots.mm)
end

# %%
plot_bhack(; N=10^3, p=0.5, α = 0.05, a = 1, b = 1, xtick=0:100:1000)

# %%
plot_bhack(; N=10^3, p=0.5, α = 0.05, a = 0.5, b = 0.5, xtick=0:100:1000)

# %%
plot_bhack(; N=10^3, p=0.5, α = 0.05, a = 0.1, b = 0.1, xtick=0:100:1000)

# %%
plot_bhack(; N=10^3, p=0.5, α = 0.05, a = 10, b = 10, xtick=0:100:1000)

# %%
plot_bhack(; N=10^3, p=0.5, α = 0.05, a = 5, b = 15, xtick=0:100:1000)

# %% [markdown]
# 以下では通常のP値とベイズ版でのP値の類似物がほぼ等しくなることを確認しよう.

# %%
function pvalue_functions(; n = 20, k = 6, a = 1, b = a)
    m = k/n
    s = √(m*(1-m)/n)
    pmin, pmax = max(0.001, m - 4s), min(0.999, m + 4s)
    p = range(pmin, pmax; length=1000)

    pval_bayes = pvalue_bayes.(n, p, k; a, b)
    pval_exact = pvalue_exact.(n, p, k)
    pval_normal = pvalue_normal.(n, p, k)

    plot(;)
    plot!(p, pval_bayes; label="pvalue_bayes")
    plot!(p, pval_exact; label="pvalue_exact", ls=:dash)
    plot!(p, pval_normal; label="pvalue_normal", ls=:dashdot)
    title!("n = $n,  k = $k,  prior = Beta($a, $b)"; titlefontsize=12)
end

# %%
pvalue_functions(n = 10, k = 3)

# %%
pvalue_functions(n = 20, k = 6)

# %%
pvalue_functions(n = 50, k = 15)

# %%
pvalue_functions(n = 100, k = 30)

# %% [markdown]
# Stirlingの公式を使った計算によって, $k = np$ のとき, $n\to\infty$ で
#
# $$
# \begin{aligned}
# \text{log(marginal likelihood ratio)} &=
# \log B(k+a, n-k+b) - \log B(a, b)) - \log(p^k(1-p)^{n-p}
# \\ &=
# a\log p + b\log(1-p) - \frac{1}{2}\log\frac{np(1-p)}{2\pi} - \log B(a, b) + o(1)
# \end{aligned}
# $$
#
# 右辺は対数周辺尤度比の漸近挙動であり, $-(1/2)\log n + \mathrm{const.}$ の形をしている. 実際にはχ²分布のスケールに合わせるために２倍してあったことに注意せよ. そのおかげで, `k` が `Binomial(n, p)` 分布に従う確率変数のとき `logmarginallikrat(n, p, k; a, b)` はその最小値(上の２倍)と自由度１のχ²分布に従う確率変数の和に近似的に等しい.

# %%
"""
対数周辺尤度比 (χ²分布のスケールに合わせるために2倍しておく)
Bayes因子 = exp(logmarginallikrat(n, p, k; a, b)/2)
"""
@memoize function logmarginallikrat(n, p, k; a=1, b=a)
    logmarginallik = 2(logbeta(k + a, n - k + b) - logbeta(a, b))
    logmarginallik0 = 2(safemul(k, log(p)) + safemul(n - k, log(1 - p)))
    logmarginallik - logmarginallik0
end

"""`logmarginallikrat(n, p, k; a, b)` の最小値の定数項"""
lmlr_min_const(n, p; a=1, b=a) = 2(a*log(p) + b*log(1-p) - 1/2*log(p*(1-p)/(2π)) - logbeta(a, b))

"""`logmarginallikrat(n, p, k; a, b)` の最小値の対数項＋定数項"""
lmlr_min_approx(n, p; a=1, b=a) = -log(n) + lmlr_min_const(n, p; a, b)

"""`logmarginallikrat(n, p, k; a, b)` の分布のプロット"""
function plot_ecdflmlr(n, p=0.5; L=10^4, a=1, b=a)
    k = rand(Binomial(n, p), L)
    lmlr = @. logmarginallikrat(n, p, k; a, b)
    @show lmlr_min = logmarginallikrat.(n, p, n*p; a, b)
    @show lmlr_min_approx(n, p; a, b)
    @show -log(n)
    @show lmlr_min_const(n, p; a, b)
    lmlrmmin = lmlr .- lmlr_min
    plot(; legend=:bottomright)
    plot!(x -> ecdf(lmlrmmin)(x), 0, 8; label="ecdf(log marginal likelihood ratios - min)")
    plot!(x -> cdf(Chisq(1), x); label="cdf(Chisq(1), x)")
    title!("n = $n,  p = $p,  prior = Beta($a, $b)"; titlefontsize=12)
end

# %%
plot_ecdflmlr(10)

# %%
plot_ecdflmlr(100)

# %%
plot_ecdflmlr(1000)

# %%
"""
Bayes因子ハッキングに挑戦
Bayes因子 > threshold となるまでデータを取得し続ける。
"""
function try_bfhacking(N, p; threshold=10, L=10^4, a=1, b=a)
    numtrials = fill(N + 1, L) # N＋１で埋め尽くされた長さLのベクトル
    Threads.@threads for i in 1:L # 並列のでL回のシミュレーションを実行
        k = 0
        for n in 1:N # ベルヌイ試行回数nの最大値はN
            k += rand(Bernoulli(p)) # kに1または0を足す。1を足す確率はp.
            if logmarginallikrat(n, p, k; a, b) > 2log(threshold) # もしもベイズ因子が閾値を超えたら
                numtrials[i] = n # n回で目的を達成したこ"""logmarginallikrat(n, p, k; a, b)の最小値の定数項
                break
            end
        end
    end
    numtrials
end

"""Bayes因子ハッキングに類似のpハッキングに挑戦"""
function try_phacking_like_bfhacking(pvalue_func, N, p; threshold=10, L=10^4, a=1, b=a)
    numtrials = fill(N + 1, L) # N＋１で埋め尽くされた長さLのベクトル
    Threads.@threads for i in 1:L # 並列のでL回のシミュレーションを実行
        k = 0
        for n in 1:N # ベルヌイ試行回数nの最大値はN
            k += rand(Bernoulli(p)) # kに1または0を足す。1を足す確率はp.
            chisqthreshold_n = 2log(threshold) + log(n)
            α_n = ccdf(Chisq(1), chisqthreshold_n) # nごとに有意水準を設定する
            if pvalue_func(n, p, k) < α_n # もしもpのP値が有意水準を下回ったら
                numtrials[i] = n # n回で目的を達成したことを記録に残す
                break
            end
        end
    end
    numtrials
end

"""ベイズ因子ハッキングへの挑戦の様子をプロット"""
function plot_bfhacking(; N=5000, p=0.5, threshold=10, L=10^4, a=1, b=a)
    numtrials_bfhack = try_bfhacking(N, 0.5; threshold, L, a, b)
    numtrials_phack_like_bfhack_exact = try_phacking_like_bfhacking(pvalue_exact, N, 0.5; threshold, L, a, b)
    numtrials_phack_like_bfhack_normal = try_phacking_like_bfhacking(pvalue_normal, N, 0.5; threshold, L, a, b)

    P = plot(; legend=:bottomright)
    plot!(n -> ecdf(numtrials_bfhack)(n), 0, N; label="Bayes factor hacking")
    plot!(n -> ecdf(numtrials_phack_like_bfhack_exact)(n), 0, N; label="p-hacking like the above (exact)", ls=:dash)
    plot!(n -> ecdf(numtrials_phack_like_bfhack_normal)(n), 0, N; label="p-hacking like the above (normal dist. approx.)", ls=:dashdot)
    title!("ecdf(number of trials),  threshold = $threshold"; titlefontsize=12)
    plot!(; xlabel="n", ylabel="P(number of trials when stopped ≤ n)", guidefontsize=10)

    Q = plot(Beta(a, b); ylim=(0, maximum(x -> pdf(Beta(a, b), x), 0.001:0.001:0.999) + 1), label="")
    plot!(; xtick=0:0.1:1, xlabel="parameter", ylabel="probability density")
    title!("prior = Beta($a, $b)"; titlefontsize=12)

    plot(P, Q; size=(640, 800), layout=(2, 1), leftmargin=2Plots.mm)
end

# %% [markdown]
# threshold = 10 のとき、 Bayes因子 > 10 という条件でデータ取得を止めることに成功する確率は数パーセントのオーダーになっている。これの類似をP値を使ってやっても概ね結果は同じになる。有意水準を
#
# $$
# \alpha_n = \operatorname{ccdf}(\operatorname{Chisq}(1), 2\log(\mathrm{threshold}) + \log n)
# $$
#
# によって $n$ について単調減少するように決めている。

# %%
plot_bfhacking(; N=5000, p=0.5, threshold=10, a=1, b=1)

# %%
plot_bfhacking(; N=5000, p=0.5, threshold=10, a=0.5, b=0.5)

# %%
plot_bfhacking(; N=5000, p=0.5, threshold=10, a=0.1, b=0.1)

# %%
plot_bfhacking(; N=5000, p=0.5, threshold=10, a=10, b=10)

# %%
plot_bfhacking(; N=5000, p=0.5, threshold=10, a=5, b=15)

# %%


# ---
# title: 740. Delete and Earn
# id: problem740
# author: Indigo
# date: 2021-05-14
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/delete-and-earn/description/>
# hidden: true
# ---
#
# Given an array `nums` of integers, you can perform operations on the array.
#
# In each operation, you pick any `nums[i]` and delete it to earn `nums[i]`
# points. After, you must delete **every** element equal to `nums[i] - 1` or
# `nums[i] + 1`.
#
# You start with 0 points. Return the maximum number of points you can earn by
# applying such operations.
#
# **Example 1:**
#
#
#
#     Input: nums = [3, 4, 2]
#     Output: 6
#     Explanation:
#     Delete 4 to earn 4 points, consequently 3 is also deleted.
#     Then, delete 2 to earn 2 points. 6 total points are earned.
#
#
#
#
# **Example 2:**
#
#
#
#     Input: nums = [2, 2, 3, 3, 3, 4]
#     Output: 9
#     Explanation:
#     Delete 3 to earn 3 points, deleting both 2's and the 4.
#     Then, delete 3 again to earn 3 points, and 3 again to earn 3 points.
#     9 total points are earned.
#
#
#
#
# **Note:**
#
#   * The length of `nums` is at most `20000`.
#   * Each element `nums[i]` is an integer in the range `[1, 10000]`.
#
#
#
#
## @lc code=start
using LeetCode

function delete_and_earn(nums::Vector{Int})
    if length(nums) == 0
        return 0
    elseif length(nums) == 1
        return nums[1]
    end
    cntr = fill(0, maximum(nums))
    for num in nums
        cntr[num] += num
    end
    fst, scd = cntr[1], max(cntr[2], cntr[1])
    for i in 3:length(cntr)
        fst, scd = scd, max(fst + cntr[i], scd)
    end
    return scd
end
## @lc code=end


################################################################################
#  Copyright 2020, Tom Van Acker                                               #
################################################################################
# MultiStateSystems.jl                                                         #
# A Julia package to solve multi-state system models.                          #
# See http://github.com/timmyfaraday/MultiStateSystems.jl                      #
################################################################################
"""
Example considering a flow transmission system with three pipes taken from:

> Multi-State System Reliability Analysis and Optimization for Engineers and
  Industrial Managers by A. Lisnianski, I. Frenkel, and Y. Ding, pg. 11/209
  (2010)

The oil flow is transmitted from a source S to a user U (Fig. 1).

        n₁  pipe 1  n₂
        |----///----|   pipe 3  n₃
    S⋅--|           |----///----|--⋅ U
        |----///----|
            pipe 2
Fig. 1: Flow transmission system.

The pipes' performance is measured by their transmission capacity (m³/hr)
(Table 1).

Table 1: Transmission capacity of the pipes in the flow transmission system.
| pipe 1               | pipe 2               | pipe 3               |
|----------------------|----------------------|----------------------|
| g¹₁ = 0    m³/hr     | g²₁ = 0    m³/hr     | g³₁ = 0    m³/hr     |
| g¹₂ = 1500 m³/hr     | g²₂ = 2000 m³/hr     | g³₂ = 1800 m³/hr     |
|                      |                      | g³₃ = 4000 m³/hr     |
"""

# load pkgs
using Unitful
using MultiStateSystems

# initialize the state-transition diagrams corresponding with the pipes
stdᵖ¹, stdᵖ², stdᵖ³ = STD(), STD(), STD()

# add the states and transitions to the respective state-transition diagrams
## pipe 1
add_states!(stdᵖ¹, flow = [0u"m^3/hr", 1500u"m^3/hr"],
                   init = [0.0, 1.0])
add_transitions!(stdᵖ¹, states = [(1,2), (2,1)],
                        rate   = [μ¹, λ¹])
## pipe 2
add_states!(stdᵖ², flow = [0u"m^3/hr", 2000u"m^3/hr"],
                   init = [0.0, 1.0])
add_transitions!(stdᵖ², states = [(1,2), (2,1)],
                        rate   = [μ², λ²])
## pipe 3
add_states!(stdᵖ³, flow = [0u"m^3/hr", 1800u"m^3/hr", 4000u"m^3/hr"],
                   init = [0.0, 0.0, 1.0])
add_transitions!(stdᵖ³, states = [(1,2), (2,1), (2,3), (3,2)],
                        rate   = [μ³ᵃ, λ³ᵃ, μ³ᵇ, λ³ᵇ])

# solve the state-transition diagrams
solve!(stdᵖ¹, SteadyStateProcess())
solve!(stdᵖ², SteadyStateProcess())
solve!(stdᵖ³, SteadyStateProcess())

# initialize the network corresponding to the flow transmission system
ntwᶠᵗˢ = Network()

# add the user, source and components to the network ntwᶠᵗˢ
add_user!(ntwᶠᵗˢ, node = 3)
add_source!(ntwᶠᵗˢ, node = 1)
add_components!(ntwᶠᵗˢ, edge = [(1,2),(1,2),(2,3)],
                        name = ["pipe 1", "pipe 2", "pipe 3"],
                        std  = [stdᵖ¹, stdᵖ², stdᵖ³])

# return the network
return ntwᶠᵗˢ


"""
    UniformMesh{T}
    UniformMesh(start::T, stop::T, length::Int) where {T}

1D uniform mesh data.

# Arguments
- `start::T` : beginning of the mesh
- `stop::T` : end of the mesh
- `length::Int` : number of cells of the mesh

# Implementation
- `step::T`  : size step
- `points::Vector{T}`: Array with node positions
- `width::T` : Distance between left and right edges.

"""
struct UniformMesh{T}
    step::T
    points::Vector{T}
    width::T

    function UniformMesh(start::T, stop::T, length::Int) where {T}
        pdeb = range(start, stop = stop, length = length + 1)
        points = pdeb[1:end-1]
        step_loc = step(pdeb)
        width = stop - start
        new{T}(step_loc, points, width)
    end
end
"""
    Base.step(mesh::UniformMesh)

Get the step of the mesh

# Argument
- `mesh::UniformMesh` : the mesh

# Return
- `step` : the step of the mesh that is the difference between two contiguous points
"""
Base.step(mesh::UniformMesh) = mesh.step

start(mesh::UniformMesh) = mesh.points[1]
stop(mesh::UniformMesh) = mesh.points[end] + mesh.step
"""
    Base.length(mesh::UniformMesh)

Get the length of the mesh

# Argument
- `mesh::UniformMesh` : the mesh

# Return
- `length` : the length of the mesh that is the number of points and cells
"""
Base.length(mesh::UniformMesh) = length(mesh.points)
"""
    points(mesh::UniformMesh)

Get the points of the mesh

# Argument
- `mesh::UniformMesh` : the mesh

# Return
- `points` : the points of the mesh that is the vector of all points of the mesh except the last
"""
points(mesh::UniformMesh) = mesh.points
"""
    width(mesh::UniformMesh)

Get the width of the mesh

# Argument
- `mesh::UniformMesh` : the mesh

# Return
- `width` : the width that is step*length or distance between left and right edges.
"""
width(mesh::UniformMesh) = mesh.width
"""
    vec_k_fft(mesh::UniformMesh{T}) where{T}

Get the fft coefficients of the mesh

# Argument
- `mesh::UniformMesh{T}` : the mesh

# Return
- fft coefficients
"""
function vec_k_fft(mesh::UniformMesh{T}) where {T}
    midx = div(length(mesh), 2)
    k = 2T(pi) / (width(mesh))
    res = k * vcat(0:midx-1, -midx:-1)
    return res
end

# @inline function newval(valref, val, borne, width)
#     diff = val - valref
#     if abs(diff) > borne
#         s = sign(diff)
#         divdiff, moddiff = divrem(diff, width)
#         if abs(moddiff) > borne
#             val -= (divdiff + s) * width
#         else
#             val -= divdiff * width
#         end
#     end
#     return val
# end

@inline function newval(valref, val, borne, width)
    diff = val - valref
    return abs(diff) >= borne ? valref + mod(diff +borne, width) - borne : val
end
# @inline function newval(valref::OpTuple{N,T}, val::OpTuple{N,T}, borne::OpTuple{N,T}, width::OpTuple{N,T}) where{N,T}
#     return OpTuple(ntuple(x -> newval(valref[x],val[x],borne[x],width[x]),N))
# end
function traitmodend!(lg::T2, f::Array{T2,N}) where {N, T2<:Union{OpTuple{N,<:Number},Number}} # T2 must OpTuple or Number
    f .= mod.(f, lg)
end
function traitmodbegin!(lg::T2, f::Array{T2,N}) where {N, T2<:Union{OpTuple{N,<:Number},Number}} # T2 must OpTuple or Number
    sz = size(f)
    borne = lg / 2
    flagmod = false
    list = [CartesianIndex(ntuple(x->1,N))]
    flagfait = zeros(Bool, sz)
    incr = map(x -> CartesianIndex(ntuple(y -> y == x, N)), 1:N)
    while !isempty(list)
        indice = popfirst!(list)
#        @show indice
        valref = f[indice]
        for i = 1:N
            newindice = indice + incr[i]
            if newindice[i] <= sz[i]
                if !flagfait[newindice]
                    val = f[newindice]
                    nval = newval(valref, val, borne, lg)
                    if nval != val
                        f[newindice] = nval
                        flagmod = true
                    end
                    flagfait[newindice] = true
                    push!(list, newindice)
                end
            end
        end
    end
    return flagmod
end

traitmodbegin!(mesh::UniformMesh{T}, f::Array{T,N}) where {T,N} = traitmodbegin!(width(mesh),f)
# traitmodbegin!(mesh::NTuple{N,UniformMesh{T}}, f::Array{OpTuple{T,N},N}) where {T,N} = traitmodbegin!(OpTuple(width.(mesh)),f)

# function traitmodbegin!(mesh::Union{NTuple{N,UniformMesh{T}},UniformMesh{T}}, f::Array{T,N}) where {T,N}
#     sz = size(f)
#     wdth = width.(mesh)
#     borne = wdth ./ 2
#     flagmod = false
#     list = [ones(Int, N)]
#     flagfait = zeros(Bool, sz)
#     incr = map(x -> map(y -> y == x, 1:N), 1:N)
#     while !isempty(list)
#         indice = popfirst!(list)
# #        @show indice
#         valref = f[indice...]
#         for i = 1:N
#             newindice = indice + incr[i]
#             if newindice[i] <= sz[i]
#                 if !flagfait[newindice...]
#                     val = f[newindice...]
#                     nval = newval(valref, val, borne, wdth)
#                     if nval != val
#                         f[newindice...] = nval
#                         flagmod = true
#                     end
#                     flagfait[newindice...] = true
#                     push!(list, newindice)
#                 end
#             end
#         end
#     end
#     return flagmod
# end
function traitmodend!(mesh::UniformMesh{T}, res::Array{T,N}) where {T,N}
    strt = start(mesh)
    res .-= strt
    traitmodend!(width(mesh), res)
    res .+= strt
end
# function traitmodend!(t_mesh::NTuple{N, UniformMesh{T}}, res::Array{OpTuple{N,T},N}) where {T,N}
#     strt = OpTuple(start.(t_mesh))
#     res .-= strt
#     traitmodend!(OpTuple(width.(t_mesh)), res)
#     res .+= strt
# end

stdtomesh(mesh::UniformMesh{T}, v) where{T}=mesh.step*v .+ mesh.points[1]
meshtostd(mesh::UniformMesh{T}, v) where{T}=(v .- mesh.points[1]) / mesh.step

# function gettuple_x(t_mesh::NTuple{N,UniformMesh{T}}) where {T,N}
#     sz = length.(t_mesh)
#     ret = ntuple(x -> zeros(T, sz), N)
#     for ind in CartesianIndices(sz)
#         for i = 1:N
#             ret[i][ind] = t_mesh[i].points[ind.I[i]]
#         end
#     end
#     return ret
# end


module day11

include("inputs.jl")

ex = inputs.ex_d11()
ex1 = inputs.ex1_d11()
ex2 = inputs.ex2_d11()
input = inputs.input_d11()

build_array(input) = collect.(input) |> x->hcat(x...) |> permutedims

function king(idx, grid, dirs)
    """ Neighbors as the King could reach."""
    return map(dirs) do dir
        checkbounds(Bool, grid, dir + idx) ? grid[idx+dir] : '.'
    end
end

function queen(idx, grid, dirs)
    """ Neighbors as the Queen could reach."""
    seat = idx # needed to reset the index after each inner loop
    neighbors = map(dirs) do dir
        idx = seat
        while checkbounds(Bool, grid, dir+idx)
            grid[dir+idx] ∈ ('L', '#') && return grid[dir+idx]
            idx += dir
        end
        return '.'
    end
end

function step(grid, get_neighbors, seen)
    next_grid = copy(grid)
    dirs = [CartesianIndex(i,j) for i ∈ -1:1, j ∈ -1:1 if !(i==j==0)]
    for idx in CartesianIndices(grid)
        neighbors = get_neighbors(idx, grid, dirs)
        if grid[idx] == 'L' && ('#' ∉ neighbors)
            next_grid[idx] = '#'
        elseif grid[idx] == '#' && count(x->x=='#', neighbors) >= seen
            next_grid[idx] = 'L'
        end
    end
    return next_grid
end

function count_occupied(input, get_neighbors, seen)
    grid = build_array(input)
    grid_ = similar(grid)

    while !(grid_ == grid)
        grid_ = copy(grid)
        grid = step(grid_, get_neighbors, seen)
    end
    return count(x->x=='#', grid)
end

part1(input) = count_occupied(input, king, 4)
part2(input) = count_occupied(input, queen, 5)

@assert part1(ex) == 37
@assert part1(input) == 2273
@info "Day 11, Part 1 answer: $(part1(input))"

@assert part2(ex) == 26
@assert part2(input) == 2064
@info "Day 11, Part 2 answer: $(part2(input))"

end


# ==================================
# International Space Station (ISS)
#
# system type: continuous LTI system
# state dimension: 270
# input dimension: 3
# ==================================
using ReachabilityBenchmarks, MathematicalSystems, LazySets, MAT

function iss_model()
    file = matopen(@relpath "iss.mat")

    # system matrix
    A = read(file, "A")

    # input matrix
    B = read(file, "B")

    # state domain
    X = Universe(270)

    # input domain
    U = Hyperrectangle([0.05, 0.9, 0.95], [0.05, 0.1, 0.05])

    # continuous LTI system
    S = @system(x' = Ax + Bu, x ∈ X, u ∈ U)

    return S
end


module Day10

using AdventOfCode2019
using LinearAlgebra

function day10(input::String = readInput(joinpath(@__DIR__, "..", "data", "day10.txt")))
    asteroidMap = _generate_map(input)
    station, maxAsteroids = _evaluate_best_asteroid(asteroidMap)
    vapo = _vaporize(asteroidMap, station, 200)
    vapo_puzzle_coords = _to_puzzle_coordinates(vapo)
    solution_part2 = 100vapo_puzzle_coords[1] + vapo_puzzle_coords[2]
    return [maxAsteroids, solution_part2]
end

function _to_puzzle_coordinates(c::CartesianIndex)
    return CartesianIndex(c[2]-1, c[1]-1)
end

function _evaluate_best_asteroid(asteroidMap::Array{Bool,2})
    maxCoord = CartesianIndex(-1, -1)
    maxCount = 0
    asteroidCoords = findall(asteroidMap)
    for a in asteroidCoords
        count = 0
        cp = copy(asteroidCoords)
        filter!(x -> x != a, cp)
        while length(cp) != 0
            b = pop!(cp)
            list = Array{CartesianIndex,1}()
            for ast in cp
                if _on_line(ast, a, b)
                    push!(list, ast)
                end
            end
            filter!(x -> !(x in list), cp)
            if _in_between(a, b, list)
                count += 2
            else
                count += 1
            end
        end
        if count > maxCount
            maxCount = count
            maxCoord = a
        end
    end
    return maxCoord, maxCount
end

function _generate_map(input::String)
    sinp = split(input)
    nrows = length(sinp)
    ncols = length(sinp[1])
    asteroidMap = zeros(Bool, (nrows, ncols))
    for (i, row) in enumerate(sinp)
        for (j, col) in enumerate(row)
            if col == '#'
                asteroidMap[i, j] = 1
            end
        end
    end
    return asteroidMap
end

function _on_line(p::CartesianIndex, a::CartesianIndex, b::CartesianIndex)
    return dot(b.I .- a.I, p.I .- a.I)^2 == sum((b.I .- a.I).^2) * sum((p.I .- a.I).^2)
end

function _in_between(p::CartesianIndex, b::CartesianIndex, list::Array{CartesianIndex,1})
    m = min(p, b, list...)
    M = max(p, b, list...)
    if p[1] == m[1]
        (p[2] == m[2] || p[2] == M[2]) && return false
    elseif p[1] == M[1]
        (p[2] == m[2] || p[2] == M[2]) && return false
    end
    return true
end

function _vaporize(asteroidMap::Array{Bool,2}, station::CartesianIndex, n::Int)
    asteroidCoords = findall(asteroidMap)
    filter!(x->x!=station, asteroidCoords)

    asteroidCoordsRight = copy(asteroidCoords)
    filter!(x -> (x[2] > station[2] || (x[2] == station[2] && x[1] < station[1])), asteroidCoordsRight)
    asteroidCoordsLeft = copy(asteroidCoords)
    filter!(x->!(x in asteroidCoordsRight), asteroidCoordsLeft)

    dirsRight = Array{CartesianIndex,1}()
    while length(asteroidCoordsRight) != 0
        ast = pop!(asteroidCoordsRight)
        filter!(x->!_on_line(x, station, ast), asteroidCoordsRight)
        push!(dirsRight, ast)
    end

    up = CartesianIndex(-1, 0)
    sort!(dirsRight, by=x->dot(up.I, (x-station).I)/(norm(up.I)*norm((x-station).I)), rev=true)

    dirsLeft = Array{CartesianIndex,1}()
    while length(asteroidCoordsLeft) != 0
        ast = pop!(asteroidCoordsLeft)
        filter!(x->!_on_line(x, station, ast), asteroidCoordsLeft)
        push!(dirsLeft, ast)
    end

    down = CartesianIndex(1, 0)
    sort!(dirsLeft, by=x->dot(down.I, (x-station).I)/(norm(down.I)*norm((x-station).I)), rev=true)

    dirs = union(dirsRight, dirsLeft)

    i = 1
    c = 0
    filter!(x->x!=station, asteroidCoords)
    length(asteroidCoords) < n && throw(AssertionError("not enough asteroids"))
    while true
        dir = dirs[i]
        candidates = []
        for ast in asteroidCoords
            if _on_line(ast, station, dir) && dot((ast-station).I, (dir-station).I) > 0
                push!(candidates, ast)
            end
        end
        if length(candidates) > 0
            ind = argmin([norm((x-station).I) for x in candidates])
            vapo = candidates[ind]
            filter!(x->x!=vapo, asteroidCoords)
            c += 1
            if c == n
                return vapo
            end
        end
        i = mod1(i + 1, length(dirs))
    end
end

end # module


let
  # testing constructor and pdf
  d = BoolDistribution(0.3)
  @test pdf(d, true) == 0.3
  @test pdf(d, false) == 0.7

  # testing support
  @test support(d) == [true, false]

  # testing ==
  d2 = BoolDistribution(0.3)
  @test d == d2
  @test BoolDistribution(0.4) != BoolDistribution(0.1)

  # testing hash
  @test hash(d) == hash(d2)

  @test sprint((io,d)->show(io,MIME("text/plain"),d), d) == sprint((io,d)->showdistribution(io,d,title="BoolDistribution"), d)
end


module UtilsTest

using Test
using ProximalBase
using Random
using LinearAlgebra
using SparseArrays
using Statistics

@testset "shrink" begin

  @test shrink(3.0, 1.0) == 2.0
  @test shrink(-2.0, 1.0) == -1.0
  @test shrink(0.5, 1.0) == 0.

  Random.seed!(1)
  p = 10
  x = randn(p)
  c = norm(x) + 0.1
  @test shrinkL2!(similar(x), x, c) == zeros(p)
  c = c - 0.15
  @test shrinkL2!(similar(x), x, c) == max(1. - c / norm(x), 0,) * x

  x = randn(p, 2*p)
  c = norm(x) + 0.1
  @test shrinkL2!(similar(x), x, c) == zeros(p, 2*p)
  c = c - 0.15
  @test shrinkL2!(similar(x), x, c) == max(1. - c / norm(x), 0,) * x
end

@testset "multiply" begin

  Random.seed!(1)
  n, m = 100, 100
  A = randn(n, m)
  b = randn(m)
  x = A * b
  xt = A' * b
  for i=1:n
    @test A_mul_B_row(A, b, i) ≈ x[i]
    @test At_mul_B_row(A, b, i) ≈ xt[i]
  end

  bs = sprand(m, 0.3)
  x = A * bs
  xt = A' * bs
  bi = SparseIterate(bs)
  xi = A * bi
  xti = A' * bi
  for i=1:n
    @test A_mul_B_row(A, bs, i) ≈ x[i]
    @test At_mul_B_row(A, bs, i) ≈ xt[i]

    @test A_mul_B_row(A, bi, i) ≈ x[i]
    @test At_mul_B_row(A, bi, i) ≈ xt[i]

    @test xi[i] ≈ x[i]
    @test xti[i] ≈ xt[i]
  end

  # symmetric A * X * B
  X = randn(100, 50)
  Σx = Symmetric(cov(X))
  Y = randn(100, 50)
  Σy = Symmetric(cov(Y))
  Δ = sprand(50, 50, 0.1)
  Δ = (Δ + Δ') / 2.

  A = Σx * Δ * Σy
  Δs = SymmetricSparseIterate(Δ)
  @test A_mul_X_mul_B(Σx, Δs, Σy) ≈ A

  # symmetric A * U * U' * B
  X = randn(100, 50)
  Σx = Symmetric(cov(X))
  Y = randn(100, 50)
  Σy = Symmetric(cov(Y))
  U = randn(50, 5)

  A = (Σx * U)*(U' * Σy)
  @test A_mul_UUt_mul_B(Σx, U, Σy) ≈ A

  # asymmetric A * X * B
  X = randn(100, 50)
  Σx = Symmetric(cov(X))
  Y = randn(100, 50)
  Σy = Symmetric(cov(Y))
  Δ = sprand(50, 50, 0.1)
  A = Σx * Δ * Σy

  Δs = SparseIterate( Δ[:] )
  R = zeros(50, 50)
  for col=1:50
      for row=1:50
          R[row, col] = A_mul_X_mul_B_rc(Σx, Δs, Σy, row, col)
      end
  end

  @test R ≈ A
end

end


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

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

Random.seed!(9)

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

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

dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true) #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
white_noise_args = Dict(:loc=>0.0, :sig=>1.0, :Nsam=>1000, :iso=>false, :norm=>false, :smooth=>true, :smoothval=>0.8) #only if you're using noise based covar
optim_settings = Dict([("iterations", 20), ("norm", false)])
recon_settings = Dict([("target_type", "ground truth"), ("covar_type", "white_noise"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.0), ("white_noise_args", white_noise_args)]) #Add constraints
#"white_noise_args": args for white_noise sims if GaussianLoss and Covar_type='white_noise'. :loc...

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

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

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

dhc_args_nonapd = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>false)
dhc_args_apd = dhc_args
@benchmark s_apd = DHC_compute_apd(true_img, filter_hash; dhc_args_apd...)
@benchmark s_nonapd = DHC_compute_apd(true_img, filter_hash; dhc_args_nonapd...)

heatmap(true_img)
heatmap(init)

minimum(true_img)

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

function adaptive_apodizer_fast(input_image::Array{Float64, 2}, dhc_args)#SPEED
        if dhc_args[:apodize]
            input_image .= apodizer(input_image)
        else

        end
        return input_image
end

@benchmark apdval = adaptive_apodizer(true_img, dhc_args_nonapd)
@benchmark apdvalfast = adaptive_apodizer_fast(true_img, dhc_args_nonapd)
#No difference

@benchmark apdimg = apodizer(true_img) #39 μs
@benchmark img = DHC_compute_apd(true_img, filter_hash; dhc_args_nonapd...) #3.475 ms
@benchmark img = DHC_compute_apd(true_img, filter_hash; dhc_args_apd...) #3.595 ms

Profile.clear()
@profile img = DHC_compute_apd(true_img, filter_hash; dhc_args_apd...)
Juno.profiler()

using StatsBase
arr1 = rand(16, 16)
wts = fweights(arr1)

#3/21###########################################################################

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

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

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

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

function dapodizer_test(img)
    eps = 1e-4
    im1 = copy(img)
    im0 = copy(img)
    im1[6, 30] += (eps/2)
    im0[6, 30] -= (eps/2)
    brute = ((apodizer(im1) - apodizer(im0))./ eps)[6, 30]
    Apmat = wind_2d(size(img)[1])
    c = sum(Apmat)
    ct = 1.0 + (1.0/c)
    term2 = (Apmat./c)
    Apmat .*= (ones(size(Apmat)).*ct .- term2)
    clever = Apmat[6, 30]
    println("Brute ", brute)
    println("Clever ", clever)
    println("diff= ", clever - brute)
end

function dwtapdtest(img, filter_hash)
    (Nf, ) = size(filter_hash["filt_index"])
    eps = 1e-4
    im1 = copy(img)
    im0 = copy(img)
    im1[6, 30] += (eps/2)
    im0[6, 30] -= (eps/2)
    brute = ((apodizer(im1) - apodizer(im0))./ eps)[6, 30]
    Apmat = wind_2d(size(img)[1])
    c = sum(Apmat)
    ct = 1.0 + (1.0/c)
    term2 = (Apmat./c)
    Apmat .*= (ones(size(Apmat)).*ct .- term2)
    clever = Apmat[6, 30]
    println("Brute ", brute)
    println("Clever ", clever)
    println("diff= ", clever - brute)

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

    sum1  = Deriv_Utils_New.wst_S20_deriv_sum(apodizer(img), filter_hash, wtgrid, FFTthreads=1)
    Sjac = reshape(Deriv_Utils_New.wst_S20_deriv(apodizer(img), filter_hash, FFTthreads=1), Nx,Nx,Nf*Nf)
    sum2  = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (Sjac[:,:,i].*wtvec[i]) end #THIS NEEDS TO BE WTGRID
    sum2 = reshape(sum2, (Nx^2, 1))
    println("Abs mean", mean(abs.(sum2 - sum1)))
    println("Abs mean", std(sum2 - sum1))
    println("Ratio: ", mean(sum2./sum1))
end

dapodizer_test(true_img)
dS20sum_test(adaptive_apodizer(true_img), filter_hash)
dwtapdtest(true_img, filter_hash)

#######DEBUG ZONE##############
tonorm = false
s_targ_mean, s_targ_invcov = s2mean, s2icov
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 derivative term dA*P/dP
if dhc_args[:apodize]
    #=
    Apmat = wind_2d(Nx)
    c = sum(Apmat)
    ct = 1.0 + (1.0/c)
    term2 = (Apmat./c)
    Apmat = Apmat .* (ct .- term2)
    =#
    Ap = wind_2d(Nx)
    cA = sum(Ap)
    Apflat = reshape(Ap, Nx^2)
    od_nz_idx = findall((x->((x!=0) & (x!=1))), Apflat)
    #tol = 1e-3
    #od_nz_idx2 = findall((x->((abs(x)>tol) & (abs(x-1)>tol))), Apflat)
    avals = Apflat[od_nz_idx]
    dA = zeros(Nx^2, Nx^2)
    dA[od_nz_idx, od_nz_idx] .= ((1.0 .- avals) * avals')./cA
    dA[diagind(dA)] += Apflat

else
    dA = I
end

lambda=0.0
FFTthreads=1
pixmask = falses(size(init))

s_targ_invcov = I
#DEB: input->init
function loss_func20(img_curr::Array{Float64, 2})
    s_curr = DHC_compute_apd(img_curr,  filter_hash, iso=false, norm=tonorm; dhc_args...)[coeff_mask]
    #regterm =  0.5*lambda*sum((adaptive_apodizer(img_curr) - adaptive_apodizer(init)).^2)
    lnlik = ( 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_targ_mean))
    #println("Lnlik size | Reg Size", size(lnlik), size(regterm))
    neglogloss = lnlik[1] #+ regterm #BUG:??
    return s_curr - s_targ_mean, lnlik, 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 = DHC_compute_apd(img_curr, filter_hash, norm=tonorm, iso=false; dhc_args...)[coeff_mask]
    diff = s_curr - s_targ_mean
    wt = reshape(convert(Array{Float64, 2}, transpose(diff) * s_targ_invcov), (Nf, Nf))
    apdimg_curr = adaptive_apodizer(img_curr)
    storage_grad .= reshape(Deriv_Utils_New.wst_S20_deriv_sum(apdimg_curr, filter_hash, wt, FFTthreads=FFTthreads)' * dA, (Nx, Nx)) #DEBUG
    #dnll_wrong = reshape(Deriv_Utils_New.wst_S20_deriv_sum(wind_2d(Nx).*img_curr, filter_hash, wt, FFTthreads=FFTthreads), (Nx, Nx)) #DEBUG
    #storage_grad .= Apmat .* (dnll + lambda.*(apdimg_curr - adaptive_apodizer(init)))
    storage_grad[pixmask] .= 0 # better way to do this by taking pixmask as an argument wst_s2_deriv_sum?
    return diff, wt, dA, dnll#, dnll_wrong
end

println("Diff check")
eps = zeros(size(init))
row, col = 24, 18 #convert(Int8, Nx/2), convert(Int8, Nx/2)+3
eps[row, col] = 1e-4
d1, ll1, chisq1 = loss_func20(init+eps./2) #DEB
d0, ll0, chisq0 = loss_func20(init-eps./2) #DEB
d, ll, chisq = loss_func20(init)
brute  = (chisq1-chisq0)/1e-4
#df_brute = DHC_compute(reshape(input, (Nx, Nx)), filter_hash, doS2=true, doS12=false, doS20=false, norm=false)[coeff_mask] - s_targ_mean
clever = zeros(size(init)) #DEB
meanval = mean(adaptive_apodizer(init)) ./ mean(wind_2d(Nx))
ddiff, dwt, dApmat, dnll = dloss20(clever, init) #DEB
println("Chisq Derve Check")
println("Brute:  ",brute)
println("Clever: ",clever[row, col], " Difference: ", brute - clever[row, col], " Mean ", meanval) #DEB

#diffs: same
brute_dnll = ((ll1 - ll0)/1e-4)[1]
dl_dnll = (Apmat .* dnll)[row, col]
brute_dnll - dl_dnll

dl_dnll_wrong = (Apmat .* dnll_wrong)[row, col]
brute_dnll - dl_dnll_wrong

#heatmap(Apmat)
heatmap(Apmat - wind_2d(Nx))

#3-22-23#############################
#Examining dA
Nx=64
Ap = wind_2d(Nx)
cA = sum(Ap)
true_img = readdust(Nx)
filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
Apflat = reshape(Ap, Nx^2)
od_nz_idx = findall((x->((x!=0) & (x!=1))), Apflat)
#tol = 1e-3
#od_nz_idx2 = findall((x->((abs(x)>tol) & (abs(x-1)>tol))), Apflat)
avals = Apflat[od_nz_idx]
dA = zeros(Nx^2, Nx^2)
dA[od_nz_idx, od_nz_idx] .= ((1.0 .- avals) * avals')./cA
dA[diagind(dA)] += Apflat
#dA[od_nz_idx, od_nz_idx][diagind(dA[od_nz_idx, od_nz_idx])] .+= avals
count((x->x!=0), dA) / length(dA)

spdA = sparse(dA)

mywts = reshape(rand(Nf*Nf), (Nf, Nf))
dls20 = Deriv_Utils_New.wst_S20_deriv_sum(true_img, filter_hash, mywts, FFTthreads=FFTthreads)'
@benchmark prod = dls20 * dA #21ms mem: 32KiB, allocs=2
@benchmark prod = dls20 * spdA #4ms mem / allocs same
#Conclusion: Sparse is definitely faster (factor of 5)
dA[od_nz_idx, od_nz_idx]
#check
xi = 12
xj = 500
dA[od_nz_idx, od_nz_idx][xj, xi]
dA[od_nz_idx, od_nz_idx][xi, xj]
dA[od_nz_idx, od_nz_idx][xj, xi]
axi = Apflat[od_nz_idx][xi]
axj = Apflat[od_nz_idx][xj]
dA[od_nz_idx, od_nz_idx][xi, xi] - axi*(1 + (1/cA) -(axi/cA))
dA[od_nz_idx, od_nz_idx][xj, xi] - axi*(1-axj)/cA
dA[od_nz_idx, od_nz_idx][xi, xj] - axj*(1-axi)/cA

#Difference STILL EXISTS
Nx=64
Ap = wind_2d(Nx)
cA = sum(Ap)
true_img = readdust(Nx)
filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
function dwtapdtest_single(img, filter_hash)
    (Nf, ) = size(filter_hash["filt_index"])
    eps = 1e-4
    im1 = copy(img)
    im0 = copy(img)
    row, col = 20, 39
    flat_ind = Nx*(col-1) + row
    im1[row, col] += (eps/2)
    im0[row, col] -= (eps/2)
    brute = ((apodizer(im1) - apodizer(im0))./ eps)#[row, col]
    #=Apmat = wind_2d(size(img)[1])
    c = sum(Apmat)
    ct = 1.0 + (1.0/c)
    term2 = (Apmat./c)
    Apmat .*= (ones(size(Apmat)).*ct .- term2)
    =#
    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
    #tol = 1e-3
    #od_nz_idx2 = findall((x->((abs(x)>tol) & (abs(x-1)>tol))), Apflat)
    #avals = Apflat[od_nz_idx]
    #
    #dA[od_nz_idx, od_nz_idx] .= ((1.0 .- avals) * avals')./cA
    println("Check dA[:, xi]")
    clever = reshape(dA[:, flat_ind], (Nx, Nx))
    #println("Brute ", brute)
    #println("Clever ", clever)
    println("Ratio= ", mean(abs.(clever./brute)))
    println("Mean Abs diff= ", mean(abs.(clever - brute)))

    # this is symmetrical, but not because S20 is symmetrical!
    mywts = rand(Nf, Nf)
    wtgrid = reshape(mywts, Nf, Nf) + reshape(mywts, Nf, Nf)'
    #only one coeff
    #wtgrid = zeros(Nf, Nf)
    #wtgrid[1, 1] = 1.0
    wtvec  = reshape(wtgrid, Nf*Nf)

    sum1  = Deriv_Utils_New.wst_S20_deriv_sum(apodizer(img), filter_hash, wtgrid, FFTthreads=1) #BUG: Is apodizer image doing something weird??
    Sjac = reshape(Deriv_Utils_New.wst_S20_deriv(apodizer(img), filter_hash, FFTthreads=1), Nx,Nx,Nf*Nf)
    sum2  = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (Sjac[:,:,i].*wtvec[i]) end #THIS NEEDS TO BE WTGRID
    sum2 = reshape(sum2, (Nx^2, 1))
    println("Check wst_deriv_sum vs longer route")
    println("Abs mean", mean(abs.(sum2 - sum1)))
    println("Abs mean", std(sum2 - sum1))
    println("Ratio: ", mean(sum2./sum1))

    println("Check deriv of Weighted Sum{S(A*P)}")
    dll = reshape(sum1' * dA, (Nx, Nx))
    coeffmask = falses(2+Nf+Nf^2)
    coeffmask[Nf+3:end] .= true
    Sh = DHC_compute_apd(im1, filter_hash; doS2=false, doS20=true, doS12=false, apodize=true, norm=false, iso=false)[coeffmask]
    Su =  DHC_compute_apd(im0, filter_hash; doS2=false, doS20=true, doS12=false, apodize=true, norm=false, iso=false)[coeffmask]

    dll_brute = ((wtvec' * Sh) - (wtvec' * Su))./eps
    println("Diff", dll[row, col] - dll_brute)
    println("Ratio", dll[row, col]./dll_brute)
    #dll_brute = (Deriv_Utils_New.wst_S20_deriv_sum(apodizer(im1), filter_hash, wtgrid, FFTthreads=1) - Deriv_Utils_New.wst_S20_deriv_sum(apodizer(im0), filter_hash, wtgrid, FFTthreads=1))./eps
    return clever, brute, dll, dll_brute, dA
end

dAclev, dAbrute, dll_cl, dll_brute, dAjac = dwtapdtest_single(true_img, filter_hash)
heatmap(reshape(dll_cl, (Nx, Nx)))
heatmap(reshape(dll_brute, (Nx, Nx)))
heatmap(apodizer(true_img))

row, col=6, 30
flat_ind = Nx*(col-1) + row
println(dA[1, flat_ind])

maximum(dll_cl)
maximum(dll_brute)

#TODO:
#Replace with fixed mu and look at comparison
#Why is the brute force test and clever SO DIFFERENT?
#shifting>????
#not using coeffmask in one place but using it elsewhere??/
#Revisit window func

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

count((x->x!=0), dA) / length(dA)

spdA = sparse(dA)

mywts = reshape(rand(Nf*Nf), (Nf, Nf))
dls20 = Deriv_Utils_New.wst_S20_deriv_sum(true_img, filter_hash, mywts, FFTthreads=1)'
@benchmark prod = dls20 * dA #22ms mem: 32KiB, allocs=2
@benchmark prod = dls20 * spdA #12ms mem / allocs same

#3/24###########################################################################
#test wrapper
Nx=16
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf, ) = size(fhash["filt_index"])
coeff_mask = falses(2+Nf+Nf^2)
coeff_mask[Nf+3:end] .= true
function test_wrapper(Nx, filter_hash, coeff_mask)
    true_img = readdust(Nx)
    swrap = DHC_compute_wrapper(true_img, filter_hash, coeff_mask=coeff_mask)
    sapd = DHC_compute_apd(true_img, filter_hash)[coeff_mask]
    print("Mean Abs Diff ", mean(abs.(swrap - sapd)))
end

test_wrapper(16, fhash, coeff_mask)

#=
filter_hash = fink_filter_hash(1, 8, nx=64, pc=1, wd=1, Omega=true)
Nx=64
true_img = readdust(Nx)
filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
DHC_compute(true_img, filter_hash, filter_hash2=filter_hash)

function test(a, b, c=b; d=3)
    return (a/b) -c + d
end

test(3, 4, 2, d=2)
=#

##3-28: Adding Iso and propagating changes everywhere
Nx=16
im = readsfd(Nx)
true_img = im[:, :, 1]
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
siso = DHC_compute(true_img, fhash, iso=true, norm=false, doS2=true)

#S20+NoIsoChecking that you didnt break old code
fname_save = "scratch_NM/NewWrapper/IsoTests/s20_noiso"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

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

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(Nf, ) = size(filter_hash["filt_index"])
coeff_mask = falses(2+Nf+Nf^2)
coeff_mask[Nf+3:end] .= true

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

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

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

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

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

#S20 + Iso check
#S20+NoIsoChecking that you didnt break old code
fname_save = "scratch_NM/NewWrapper/3-29/GTTargWNCov/log_apd_iso"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

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

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nf) = size(filter_hash["S2_iso_mat"])
coeff_mask = falses(2+S1iso+S2iso)
coeff_mask[S1iso+3:end] .= true
#coeff_mask = falses(2+Nf+Nf^2)
#coeff_mask[Nf+3:end] .= true

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

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

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

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

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

#S20 + Iso check
#S20+NoIsoChecking that you didnt break old code
fname_save = "scratch_NM/NewWrapper/3-29/SFDTargSFDCov/log_apd_iso"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

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

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

white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :norm=>false, :smooth=>false, :smoothval=>0.8) #Iso #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", std(apodizer(log.(true_img))).^(-2)), ("white_noise_args", white_noise_args)]) #Add constraints
regs_true = DHC_compute_wrapper(log.(true_img), filter_hash; norm=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
#s2mean - s_true

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

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

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

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

fname_save = "scratch_NM/NewWrapper/3-29/GTTargWNCov/log_apd_iso"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

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

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nf) = size(filter_hash["S2_iso_mat"])
coeff_mask = falses(2+S1iso+S2iso)
coeff_mask[S1iso+3:end] .= true
#coeff_mask = falses(2+Nf+Nf^2)
#coeff_mask[Nf+3:end] .= true

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

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

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

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

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

#S20 + Iso check
#S20+NoIsoChecking that you didnt break old code
fname_save = "scratch_NM/NewWrapper/3-29/SFDTargSFDCov/log_apd_iso_exp2"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

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

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

white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :norm=>false, :smooth=>false, :smoothval=>0.8) #Iso #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.00456), ("white_noise_args", white_noise_args)]) #Add constraints
regs_true = DHC_compute_wrapper(log.(true_img), filter_hash; norm=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
#s2mean - s_true

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

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

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

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

##fname_save = "scratch_NM/NewWrapper/3-29/SFDTargSFDCov/log_apd_iso_exp2"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]

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

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

white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :norm=>false, :smooth=>false, :smoothval=>0.8) #Iso #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", true), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.00456), ("white_noise_args", white_noise_args)]) #Add constraints
regs_true = DHC_compute_wrapper(log.(true_img), filter_hash; norm=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
#s2mean - s_true

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

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

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

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


############################################
#
#    dolfin.cpp.geometry
# https://fenicsproject.org/docs/dolfin/2018.1.0/python/_autogenerated/dolfin.cpp.geometry.html
#
#    ufl.geometry
#
############################################

@pyclass dolfin Point
@pyclass dolfin BoundingBoxTree
@pyclass dolfin MeshPointIntersection

@pyclass ufl FacetNormal Expression
@pyclass ufl CellDiameter Expression

@pyclass ufl CellNormal
@pyclass ufl CellVolume
@pyclass ufl Circumradius
@pyclass ufl MaxCellEdgeLength
@pyclass ufl MaxFacetEdgeLength
@pyclass ufl MinCellEdgeLength
@pyclass ufl MinFacetEdgeLength
@pyclass ufl SpatialCoordinate

export Point, BoundingBoxTree, MeshPointIntersection, FacetNormal, CellDiameter,
       CellNormal, CellVolume, Circumradius, MaxCellEdgeLength, MaxFacetEdgeLength,
       MinCellEdgeLength, MinFacetEdgeLength, SpatialCoordinate


# Code to asses the effect of moment-compensation on the diffusion signal

## Read
#Waveform
using MRIsim, PlotlyJS
b = [0, 50, 200, 400, 600, 790]
E = Array{Any, 2}(undef,6,4)
FID = false
path_file = "/home/ccp/Documents/MRIsim.jl/"
#Timings
f = readlines(path_file*"qte_timings.txt")
timings = tryparse.(Float64, split(f[2]," "))
δ1, rf180, δ2 = timings[timings .≠ nothing]*1e-3
#Normal DTI
G, n1, n2 = MRIsim.read_diff_fwf(path_file*"QTE_Waveforms/MX_MC2_b790.txt")
Gmax = 21e-3
g1, g2 = G[1:20:n1,1], G[1:20:n2,4]
seq_dti = Sequence(Gmax*[Grad(g1,δ1) delay(rf180) -Grad(g2,δ2)])
#ACQ
if FID
    seq_ACQ = Sequence([Grad(0,0)], [RF(0,0)], DAC(1,0))
else
    FOV = 4e-2 #10cm
    Nxy = 20
    ΔT = 20e-6 #100um
    Gmax_acq = 60e-3
    seq_ACQ,_,_,_ = PulseDesigner.EPI_base(FOV, Nxy, ΔT, Gmax_acq)
end
#Phantom
Ns = 10_000; Nsx = floor(Int64, sqrt(Ns) ); Ns = Nsx*Nsx
ADC = 2e-9 #Normal diffusion coefficient
v = 6e-2 # 2mm/s Capilar flow
Δx = .5e-2 #3mm
x = range(-Δx/2,Δx/2,length=Nsx) #spin coordinates
y = range(-Δx/2,Δx/2,length=Nsx)  #spin coordinates
xx, yy = x .+ y'*0, x*0 .+ y' #grid points
T = dur(seq_dti)+dur(seq_ACQ)/2
phantom  = Phantom(x=xx[:], y=yy[:], Dλ1=ADC*ones(Ns), Dλ2=ADC*ones(Ns))
phantom2 = Phantom(x=xx[:], y=yy[:], ux=(x,y,z,t)->v*t, Dλ1=ADC*ones(Ns), Dλ2=ADC*ones(Ns))
#Iter
for idx = 1:5
#Seq
G, n1, n2 = MRIsim.read_diff_fwf(path_file*"QTE_Waveforms/MX_MC2_b$(b[idx+1]).txt")
g1, g2 = G[1:20:n1,1], G[1:20:n2,4]
seq = Sequence(Gmax*[Grad(g1,δ1) delay(rf180) -Grad(g2,δ2)])
seq += seq_ACQ
#Simulate
f = sqrt(b[idx+1]/790)
simParams = Dict(:Δt => 10e-6, :Nblocks=> 1)
if FID
    recParams = Dict(:skip_rec => true)
else
    recParams = Dict(:epi => true, :recon => :fft, :Nx => Nxy)
end
S_MM = MRIsim.simulate(phantom, seq, simParams, recParams)
S_DTI = MRIsim.simulate(phantom, f*seq_dti+seq_ACQ, simParams, recParams)
S_MM2 = MRIsim.simulate(phantom2, seq, simParams, recParams)
S_DTI2 = MRIsim.simulate(phantom2, f*seq_dti+seq_ACQ, simParams, recParams)
S0 = MRIsim.simulate(phantom, 0*seq, simParams, recParams)
S02 = MRIsim.simulate(phantom2, 0*seq, simParams, recParams)
#Saving results
E[idx+1,1] = S_MM
E[idx+1,2] = S_DTI
E[idx+1,3] = S_MM2
E[idx+1,4] = S_DTI2
E[1,1] = S0
E[1,2] = S0
E[1,3] = S02
E[1,4] = S02
end

##Plot
i = 4
zmax = 12
plot(heatmap(z=abs.(E[i,1]), transpose = false, zmin=0, zmax=zmax), Layout(;title="No mov. MM-DTI b$(b[i])"))
plot(heatmap(z=abs.(E[i,2]), transpose = false, zmin=0, zmax=zmax), Layout(;title="No mov. DTI b$(b[i])"))
plot(heatmap(z=abs.(E[i,3]), transpose = false, zmin=0, zmax=zmax), Layout(;title="Mov. MM-DTI b$(b[i])"))
plot(heatmap(z=abs.(E[i,4]), transpose = false, zmin=0, zmax=zmax), Layout(;title="Mov. DTI b$(b[i])"))

p = MRIsim.plot_grads_moments(seq_dti+seq_ACQ)


acceptance(v::AbstractArray{Bool}) = mean(v)

function acceptance(v::AbstractArray)
  accepted = 1
  n = length(v)

  for i in 2:n
    if v[i] != v[i-1]
      accepted += 1
    end
  end

  return accepted/n
end

acceptance(v::AbstractArray, region) = mapslices(acceptance, v, region)

function acceptance(s::UnivariateParameterNState; key::Symbol=:accept, diagnostics::Bool=true)
  if diagnostics
    return acceptance(convert(Array{Bool, 2}, s.diagnosticvalues[findfirst(s.diagnostickeys, key), :]))
  else
    return acceptance(s.value)
  end
end

acceptance(s::VariableNState{Univariate}) = acceptance(s.value)

function acceptance(s::MultivariateParameterNState; key::Symbol=:accept, diagnostics::Bool=true)
  if diagnostics
    return acceptance(convert(Array{Bool, 2}, s.diagnosticvalues[findfirst(s.diagnostickeys, key), :]))
  else
    return acceptance(Any[s.value[:, i] for i in 1:s.n])
  end
end

acceptance(s::VariableNState{Multivariate}) = acceptance(Any[s.value[:, i] for i in 1:s.n])