text:
## This file contains direct Julia implementations of 2d and 3d convolutions

# Helper functions for restricting x/w overreach
function clamp_lo(x, w)
    idx = 1
    while idx <= length(x) && x[idx] <= 0
        idx += 1
    end
    return (x[idx:end], w[idx:end])
end
function clamp_hi(x, w, L)
    idx = length(x)
    while idx >= 1 && x[idx] > L
        idx -= 1
    end
    return (x[1:idx], w[1:idx])
end

"""
    conv_direct!(y, x, w, cdims; alpha=1, beta=0)

Direct convolution implementation; used for debugging, tests, and mixing/matching of
strange datatypes within a single convolution.  Uses naive nested for loop implementation
and does not attempt to optimize performance.  Rather, this implementation is intended to
be maximally understandable and debuggable, to aid in testing other, more performant
implementations.  We also explicitly support mixing and matching of strange datatypes,
so that if the user really wants to convolve an image of `UInt8`'s with a `Float16`
kernel, storing the result in a `Float32` output, there is at least a function call
for that madness.

The keyword arguments `alpha` and `beta` control accumulation behavior; this function
calculates `y = alpha * x * w + beta * y`, therefore by setting `beta` to a nonzero
value, the user is able to accumulate values into a preallocated `y` buffer, or by
setting `alpha` to a nonunitary value, an arbitrary gain factor can be applied.

By defaulting `beta` to `false`, we make use of the Bradbury promotion trick to override
`NaN`'s that may pre-exist within our output buffer, as `false*NaN == 0.0`, whereas
`0.0*NaN == NaN`.  Only set `beta` if you are certain that none of the elements within
`y` are `NaN`.

The basic implementation performs 3-dimensional convolution; 1-dimensional and 2-
dimensional casesa are supported by simply reshaping `y`, `x` and `w`, for which
wrapper methods are available.
"""
conv_direct!

function conv_direct!(y::AbstractArray{yT,5}, x::AbstractArray{xT,5},
                      w::AbstractArray{wT,5}, cdims::DenseConvDims;
                      alpha::yT = yT(1), beta = false) where {yT, xT, wT}
    check_dims(size(x), size(w), size(y), cdims)

    width, height, depth = input_size(cdims)
    kernel_w, kernel_h, kernel_d = kernel_size(cdims)
    out_c = channels_out(cdims)
    pad_w_lo, pad_w_hi, pad_h_lo, pad_h_hi, pad_d_lo, pad_d_hi = padding(cdims)
    dil_w, dil_h, dil_d = dilation(cdims)
    stride_w, stride_h, stride_d = stride(cdims)
    out_width, out_height, out_depth = output_size(cdims)

    # If we're doing crosscorr instead of conv, then don't bother to flip `w`
    if !flipkernel(cdims)
        w = w[end:-1:1, end:-1:1, end:-1:1, :, :]
    end

    # A helper function to project from output (w, h) to input (input_w, input_h)
    @inline project(idx, stride, pad) = (idx - 1)*stride - pad + 1

    # explicit formulation of convolution.  Oh hoisting gods, hear my plea.
    @inbounds for batch in 1:size(x)[end],
        c_out in 1:out_c,
        d_idx in 1:out_depth,
        h_idx in 1:out_height,
        w_idx in 1:out_width

        # Starting points of the window of x we're going to grab
        x_w = project(w_idx, stride_w, pad_w_lo)
        x_h = project(h_idx, stride_h, pad_h_lo)
        x_d = project(d_idx, stride_d, pad_d_lo)

        # Grow that starting point into ranges
        x_widxs = x_w .+ (0:dil_w:(dil_w*kernel_w-1))
        x_hidxs = x_h .+ (0:dil_h:(dil_h*kernel_h-1))
        x_didxs = x_d .+ (0:dil_d:(dil_d*kernel_d-1))
        w_widxs = 1:kernel_w
        w_hidxs = 1:kernel_h
        w_didxs = 1:kernel_d

        # Clamp the ranges to simulate padding
        x_widxs, w_widxs = clamp_lo(x_widxs, w_widxs)
        x_widxs, w_widxs = clamp_hi(x_widxs, w_widxs, width)
        x_hidxs, w_hidxs = clamp_lo(x_hidxs, w_hidxs)
        x_hidxs, w_hidxs = clamp_hi(x_hidxs, w_hidxs, height)
        x_didxs, w_didxs = clamp_lo(x_didxs, w_didxs)
        x_didxs, w_didxs = clamp_hi(x_didxs, w_didxs, depth)

        # Grab our slices
        x_slice = view(x, x_widxs, x_hidxs, x_didxs, :, batch)
        w_slice = view(w, w_widxs, w_hidxs, w_didxs, :, c_out)

        # Do the dotproduct dance, then weight by alpha/beta and git 'er done
        dotprod = sum(x_slice .* w_slice)
        y[w_idx, h_idx, d_idx, c_out, batch] = alpha*convert(yT, dotprod) +
                                               beta*y[w_idx, h_idx, d_idx, c_out, batch]
    end

    return y
end

## Gradient definitions
"""
    ∇conv_data_direct!(dx, dy, w, cdims; alpha=1, beta=0)

Calculate the gradient imposed upon `x` in the convolution `y = x * w`.
"""
∇conv_data_direct!

function ∇conv_data_direct!(dx::AbstractArray{xT,5}, dy::AbstractArray{yT,5},
                            w::AbstractArray{wT,5}, cdims::DenseConvDims;
                            alpha::xT=xT(1), beta=false) where {xT, yT, wT}
    w = transpose_swapbatch(w[end:-1:1, end:-1:1, end:-1:1, :, :])
    dy = predilate(dy, stride(cdims))
    ctdims = DenseConvDims(dy, w; padding=transpose_pad(cdims),
                                  dilation=dilation(cdims),
                                  flipkernel=flipkernel(cdims))
    dx = conv_direct!(dx, dy, w, ctdims; alpha=alpha, beta=beta)
    return transpose_swapbatch(dx)
end

"""
    ∇conv_filter_direct!(dw, x, dy, cdims; alpha=1, beta=0)

Calculate the gradient imposed upon `w` in the convolution `y = x * w`.
"""
∇conv_filter_direct!

function ∇conv_filter_direct!(dw::AbstractArray{wT,5}, x::AbstractArray{xT,5},
                              dy::AbstractArray{yT,5}, cdims::DenseConvDims;
                              alpha::wT=wT(1), beta=false) where {xT, yT, wT}
    x = transpose_swapbatch(x[end:-1:1, end:-1:1, end:-1:1, :, :])
    dy = transpose_swapbatch(predilate(dy, stride(cdims)))
    ctdims = DenseConvDims(dy, x; padding=transpose_pad(cdims),
                                    stride=dilation(cdims))
    conv_direct!(dw, dy, x, ctdims; alpha=alpha, beta=beta)
    if flipkernel(cdims)
        dw .= dw[end:-1:1, end:-1:1, end:-1:1, :, :]
    end
    return dw
end


text:
@testset "Horizontal diffusion of random" begin

    mx,nx,nlev = 31,32,8

    for T in (Float16,Float32,Float64)

        A           = rand(Complex{T},mx,nx,nlev)
        tendency    = rand(Complex{T},mx,nx,nlev)
        D1          = rand(T,mx,nx)
        D2          = rand(T,mx,nx)

        A1 = A[:,:,1]
        tendency1 = copy(tendency[:,:,1])

        SpeedyWeather.horizontal_diffusion!(A,tendency,D1,D2)
        SpeedyWeather.horizontal_diffusion!(A1,tendency1,D1,D2)

        @test tendency1 == tendency[:,:,1]
    end
end

ext: jl


text:
function Δxtest(nx)
	(nx<6) && error("test requires nx>6")

	p=Random.randn(nx)
	p[1:2] .= 0.0  # it is necessary that starting models have these boundary conditions
	p[nx-1:nx] .= 0.0

	pout=Δxsample(p)

	pout2=randn(nx)
	pout2[1:2] .= 0.0  # it is necessary that starting models have these boundary conditions
	pout2[nx-1:nx] .= 0.0

	p2=Δxsample(pout2)

	@test dot(pout2,pout) ≈ dot(p,p2)
end

function Δttest(nx)
	(nx<6) && error("test requires nx>6")

	p=Random.randn(nx)

	pout=Δtsample(p)
	println(pout)

	pout2=randn(nx)

	p2=Δtsample(pout2)

	@test dot(pout2,pout) ≈ dot(p,p2)
end

function testing2(nx,nz)
	ntimes=1

#	p=zeros(nz,nx,3)
#	p[div(nz,2),div(nx,2),1]=1.0
	p=Random.randn(nz,nx,3)
	modrrvz=Random.randn(nz,nx)
	modrrvx=Random.randn(nz,nx)
	modttI=Random.randn(nz,nx)
	a_x=Random.randn(nz,nx)

	pout=advance_sample(ntimes, p, modrrvx, modrrvz, modttI)
#	println(pout[:,:,1])

	pout2=randn(nz,nx,3)

	p2=advance_sample(ntimes, pout2, modrrvx, modrrvz, modttI)

	i=3
	@test dot(pout2[:,:,i],pout[:,:,i]) ≈ dot(p[:,:,i],p2[:,:,i])
	#@test dot(pout2,pout-p) ≈ dot(p,p2-pout2)
end


text:
using MPI_Tools
using Test

@testset "MPI_Tools.jl" begin
    # Write your tests here.
end


text:
import Base.convert
using CUDA: CuArray

convert(::Type{Array{T}}, x::CuArray{S,N}) where {T,N,S} = convert(Array{T,N}, x)
convert(::Type{Array{T,N}}, x::CuArray{S,N}) where {T,N,S} = convert(Array{T,N},Array(x))


text:
using ROC
using DataFrames
using CSV
using Test
using Random

@testset "Input checks" begin
    scores = [1, 2, 3]
    labels = ['a', 'b', 'b']

    @test_throws ArgumentError roc(scores, ['a', 'b'])

    curve = roc(scores, labels)
    @test_throws ErrorException AUC(curve, 0.0)
    @test_throws ErrorException AUC(curve, 1.0)
end

@testset "ROCR data" begin
    data = CSV.read(joinpath(@__DIR__, "data", "ROCRdata.csv"))
    scores = data[!, :predictions]
    labels = data[!, :labels]

    curve = roc(scores, labels)
    @test abs( AUC(curve) - 0.834187 ) < 0.000001 # ROCR 0.8341875
    @test abs( AUC(curve, 0.01) - 0.000329615 ) < 0.000000001 # ROCR 0.0003296151
    @test abs( AUC(curve, 0.1) - 0.0278062 ) < 0.0000001 # ROCR 0.02780625

    @testset "Missings" begin
        na_scores = [missing, scores[2:end]...]
        na_labels = [missing, labels[2:end]...]

        @test abs( AUC(roc(na_scores, labels)) - 0.834187 ) < 0.001
        @test abs( AUC(roc(scores,  na_labels)) - 0.834187 ) < 0.001
        @test abs( AUC(roc(na_scores, na_labels)) - 0.834187 ) < 0.001
    end

    @testset "BitVector" begin
        curve = roc(scores, BitVector(labels))

        @test abs( AUC(curve) - 0.834187 ) < 0.000001 # ROCR 0.8341875
    end

    @testset "Distances" begin
        curve = roc(1.0 .- scores, labels, distances=true)

        @test abs( AUC(curve) - 0.834187 ) < 0.000001 # ROCR 0.8341875
        @test abs( AUC(curve, 0.01) - 0.000329615 ) < 0.000000001 # ROCR 0.0003296151
        @test abs( AUC(curve, 0.1) - 0.0278062 ) < 0.0000001 # ROCR 0.02780625
    end
end

# Are AUC and ROC consistent after permutation of scores and labels?
@testset "Permutations" begin
    perm = randperm(100)
    scores = rand(100)
    labels = [-ones(50); ones(50)]

    @test AUC(roc(scores, labels)) == AUC(roc(scores[perm], labels[perm]))
end

@testset "ROC analysis: web-based calculator for ROC curves' example" begin
    scores = [1 , 2 , 3 , 4 , 6 , 5 , 7 , 8 , 9 , 10]
    labels = [0 , 0 , 0 , 0 , 0 , 1 , 1 , 1 , 1 , 1]
    res = roc(scores, labels, 1)

    @test AUC(res) ≈ 0.96
    @test res.TPR == [0.0, 0.2, 0.4, 0.6, 0.8, 0.8, 1.0, 1.0, 1.0, 1.0, 1.0]
    @test res.FPR == [0.0, 0.0, 0.0, 0.0, 0.0, 0.2, 0.2, 0.4, 0.6, 0.8, 1.0]
end

@testset "Ties" begin
    # Example from https://www.epeter-stats.de/roc-curves-and-ties/
    # This package uses the first strategy, i.e.  AUC is equivalent to the
    # Mann-Whitney U statistic

    label = vcat([true for _ in 1:1000], [false for _ in 1:1000])
    pred = vcat(rand(7:14, 1000), rand(1:8, 1000))
    roc_data = roc(pred, label)
    @test abs( round(AUC(roc_data), digits=2) - 0.97 ) ≤ 0.015

    # Example from https://github.com/brian-lau/MatlabAUC/issues/1
    data = [-1 1; 1 2; -1 3; -1 4; 1 5; -1 6; 1 7; -1 8; 1 9; 1 10;
            1 11; -1 13; 1 13; 1 14; 1 14]
    roc_data = roc(data[:,2], data[:,1], 1)
    spss = [1.000 1.000
            1.000 .833
             .889 .833
             .889 .667
             .889 .500
             .778 .500
             .778 .333
             .667 .333
             .667 .167
             .556 .167
             .444 .167
             .333 .167
             .222 .000
             .000 .000]
    @test size(spss, 1) == length(roc_data.TPR)
    @test size(spss, 1) == length(roc_data.FPR)
    for i in 1:size(spss, 1)
        @test isapprox(roc_data.TPR[end - i + 1], spss[i, 1], atol=1e-3)
        @test isapprox(roc_data.FPR[end - i + 1], spss[i, 2], atol=1e-3)
    end
end

@testset "Distances and Ties" begin

    labels = Bool[false, true, false, false, true, false, true, false, true,
                  true, true, false, true, true, true]
    normal = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 13, 14, 14] # 14 is better
    distances = 14 .- normal # 0 is better
    roc_normal = roc(normal, labels)
    roc_distances = roc(distances, labels, distances=true)

    # Same ROC curve:
    @test roc_normal.TPR ≈ roc_distances.TPR
    @test roc_normal.FPR ≈ roc_distances.FPR
    @test AUC(roc_normal) ≈ AUC(roc_distances)
end


text:
#!/usr/bin/env julia
using Documenter
using HypertextLiteral

# Setup for doctests embedded in docstrings.
DocMeta.setdocmeta!(HypertextLiteral, :DocTestSetup, :(using HypertextLiteral))

# Highlight indented code blocks as Julia code.
using Documenter.Expanders: ExpanderPipeline, Selectors, Markdown, iscode
abstract type DefaultLanguage <: ExpanderPipeline end
Selectors.order(::Type{DefaultLanguage}) = 99.0
Selectors.matcher(::Type{DefaultLanguage}, node, page, doc) =
    iscode(node, "")
Selectors.runner(::Type{DefaultLanguage}, node, page, doc) =
    page.mapping[node] = Markdown.Code("julia", node.code)

custom_footer = """
Powered by [Documenter.jl](https://github.com/JuliaDocs/Documenter.jl),
[NarrativeTest.jl](https://github.com/MechanicalRabbit/NarrativeTest.jl),
and the [Julia Programming Language](https://julialang.org/).
"""

makedocs(
    sitename = "HypertextLiteral.jl",
    format = Documenter.HTML(prettyurls=(get(ENV, "CI", nothing) == "true"),
                             footer=custom_footer),
    pages = [
        "Overview" => "index.md",
        "content.md",
        "attribute.md",
        "script.md",
        "notation.md",
        "design.md",
        "primitives.md",
        "reference.md",
    ],
    modules = [HypertextLiteral])

deploydocs(
    repo = "github.com/MechanicalRabbit/HypertextLiteral.jl.git",
)


text:
# This file was generated, do not modify it. # hide
ŷ = predict(mtm, Xtrain)
cross_entropy(ŷ, ytrain) |> mean


text:
struct Graph
    vertices::Set{String}
    edges::Set{Tuple}
end

function generateGraph(filename)
    values = readlines(filename*".txt")
    edges = Set()
    vert = Set()
    for line in values
        st = split(line, '-')
        v1 = string(st[1])
        v2 = string(st[2])
        vert = push!(vert, v1)
        vert = push!(vert, v2)
        if !((v2,v1) in edges)
            edges = push!(edges, (v1,v2))
        end
    end
    return Graph(vert, edges)
end

function getNeighbors(vertex, graph)
    ngb = Set()
    for e in graph.edges
        if e[1] == vertex
            ngb = push!(ngb, e[2])
        elseif e[2] == vertex
            ngb = push!(ngb, e[1])
        end
    end
    return ngb
end

function dfs(graph::Graph, node::String, stop::String, visited::Dict, currentpath::Vector{String}, paths)
    visited[node] = true
    if (node == stop)
        println(currentpath)
        paths = push!(paths, currentpath)
        currentpath = String[]
    end
    neighbors = getNeighbors(node, graph)
    for n in neighbors
        if isuppercase(n[1]) || (islowercase(n[1]) && visited[n] == false)
            currentpath = push!(currentpath, n)
            dfs(graph, n, stop, visited, currentpath, paths)
            currentpath = deleteat!(currentpath, findlast(x -> x == n, currentpath))
        end
    end
    visited[node] = false
end

function dfsparttwo(graph::Graph, node::String, stop::String, visited::Dict, currentpath::Vector{String}, paths, visitedtwice::Bool)

    visited[node] = true
    if (node == stop)
        #println("path ended")
        #println(currentpath)
        paths = push!(paths, currentpath)
        currentpath = String[]
        visitedtwice = false
    end
    neighbors = getNeighbors(node, graph)

    for n in neighbors
        if n in currentpath
            visited[n] = true
        end

        if visitedtwice
            check = true
            for el in currentpath
                if islowercase(el[1]) && count(x -> x == el, currentpath) > 1
                    check = false
                    break
                end
            end
            if check
                visitedtwice = false
            end
        end

        if isuppercase(n[1]) || (islowercase(n[1]) && visited[n] == false) || (islowercase(n[1]) && n in currentpath && visitedtwice == false && n != "start" && n != "end")

            if islowercase(n[1]) && n in currentpath
                visitedtwice = true
            end
            currentpath = push!(currentpath, n)
            dfsparttwo(graph, n, stop, visited, currentpath, paths, visitedtwice)
            currentpath = deleteat!(currentpath, findlast(x -> x == n, currentpath))
        end
    end
    visited[node] = false
end

function findPaths(graph::Graph, start::String, stop::String)
    paths = []
    visited = Dict(x => false for x in graph.vertices)
    currentpath = String[]
    currentpath = push!(currentpath, start)
    dfs(graph, start, stop, visited, currentpath, paths)
    return paths

end

function findPathsparttwo(graph::Graph, start::String, stop::String)
    paths = []
    visited = Dict(x => false for x in graph.vertices)
    currentpath = String[]
    currentpath = push!(currentpath, start)
    visitedtwice = false
    dfsparttwo(graph, start, stop, visited, currentpath, paths, visitedtwice)
    return paths
end

function solve(filename)
    graph = generateGraph(filename)
    println(graph.edges)
    paths = findPaths(graph, "start", "end")
    println(length(paths))
end

function solve2(filename)
    graph = generateGraph(filename)
    paths = findPathsparttwo(graph, "start", "end")
    println(length(paths))
end

#solve("input")
#solve("smallexample")
solve2("smallexample")
solve2("midexample")
solve2("largeexample")
solve2("input")


text:
import Tk

function tkwindow(name, w, h, closecb=nothing)
    win = Tk.Window(name, w, h)
    c = Tk.Canvas(win, w, h)
    Tk.pack(c, expand = true, fill = "both")
    if !is(closecb,nothing)
        Tk.bind(win, "<Destroy>", closecb)
    end
    c
end

function tkdestroy(c::Tk.Canvas)
    w = Tk.toplevel(c)
    Tk.destroy(w)
    nothing
end

function display(c::Tk.Canvas, pc::PlotContainer)
    c.draw = let bad=false
        function (_)
            bad && return
            ctx = getgc(c)
            set_source_rgb(ctx, 1, 1, 1)
            paint(ctx)
            try
                Winston.page_compose(pc, Tk.cairo_surface(c))
            catch e
                bad = true
                isa(e, WinstonException) || rethrow(e)
                println("Winston: ", e.msg)
            end
        end
    end
    Tk.draw(c)
    Tk.update()
end

function get_context(c::Tk.Canvas, pc::PlotContainer)
    device = CairoRenderer(Tk.cairo_surface(c))
    ext_bbox = BoundingBox(0,width(c),0,height(c))
    _get_context(device, ext_bbox, pc)
end

get_context(pc::PlotContainer) = get_context(curfig(_display), pc)


text:
#Demo 14: Ck2Q vs Delay tests
#-------------------------------------------------------------------------------

using Colors
using CMDimCircuits
CMDimCircuits.@using_CData()

#Get a demo display:
include(CMDimCircuits.demoplotcfgscript)
#Normally use something like:
#CMDimData.@includepkg EasyPlotInspect

#==Constants
===============================================================================#
LBLAX_AMPV = "Amplitude (V)"
LBLAX_TIME = "Time (s)"
LBLAX_DELPS = "Delay (ps)"
LBLAX_ΔDELPS = "𝛥Delay (ps)"
delayattr = cons(:a,
	line = set(style=:solid, width=3, color=1),
	glyph = set(shape=:x, size=2),
)
color_ck2q = cons(:a, line = set(color=2))
color_ck2qI = cons(:a, line = set(color=3))

#==Input data
===============================================================================#
tbit = 1e-9 #Bit period
osr = 20 #samples per bit
nsamples = 127

#==Computations
===============================================================================#
seq = 1.0*prbs(reglen=7, seed=1, nsamples=nsamples)
t = DataF1(0:(tbit/osr):(nsamples*tbit))
τref = tbit/5
τ = tbit/1.5
ck = sin(2pi*t/tbit)

#Generate data:
Πref = pulse(t, Pole(1/τref,:rad), tpw=tbit)
Π = pulse(t, Pole(1/τ,:rad), tpw=tbit)

patref = (pattern(seq, Πref, tbit=tbit)-0.5)*2 #Center data pattern
pat = (pattern(seq, Π, tbit=tbit)-0.5)*2 #Center data pattern

Δ = measdelay(patref, pat, xing_ref=CrossType(:risefall), xing_main=CrossType(:risefall))
ck2q = measck2q(ck, pat, xing_ck=CrossType(:rise), xing_q=CrossType(:risefall))
@show tstart_ck = xcross1(patref)-50e-12
#Relative to ideal reference clock:
ck2q_ideal = measck2q(pat, tbit, tstart_ck=tstart_ck, xing_q=CrossType(:risefall))

#==Generate plot
===============================================================================#
LBLAX_AMPV = "Amplitude (V)"
LBLAX_TIME = "Time (s)"
LBLAX_DELPS = "Delay (ps)"

plot = cons(:plot, nstrips=3,
	ystrip1 = set(axislabel=LBLAX_AMPV, striplabel="Patterns"),
	ystrip2 = set(axislabel=LBLAX_DELPS, striplabel="Delays (measdelay, measck2q)"),
	ystrip3 = set(axislabel=LBLAX_ΔDELPS, striplabel="Delay Differences (wrt ck2q & ideal ck2q)"),
	xaxis = set(label=LBLAX_TIME, max=maximum(t)+3*tbit),
)
push!(plot,
	cons(:wfrm, ck, line=set(width=1, color=RGB24(.5, .5, .5)), label="clock", strip=1),
	cons(:wfrm, patref, line=set(width=2), label="ref", strip=1),
	cons(:wfrm, pat, label="pat", strip=1),

	cons(:wfrm, Δ/1e-12, delayattr, label="delay", strip=2),
	cons(:wfrm, ck2q/1e-12, delayattr, color_ck2q, label="ck2q", strip=2),
	cons(:wfrm, ck2q_ideal/1e-12, delayattr, color_ck2qI, label="ck2qI (ideal)", strip=2),

	cons(:wfrm, (Δ-ck2q)/1e-12, delayattr, color_ck2q, label="del-ck2q", strip=3),
	cons(:wfrm, (Δ-ck2q_ideal)/1e-12, delayattr, color_ck2qI, label="del-ck2qI", strip=3),
)

pcoll = push!(cons(:plotcoll, title="Compare measdelay & measck2q"), plot)
	pcoll.ncolumns = 1

#==Display results in pcoll
===============================================================================#
EasyPlot.displaygui(pcoll)

#==Return pcoll to user (call evalfile(...))
===============================================================================#
pcoll


text:
################################################################################
#  Copyright 2020, Tom Van Acker                                               #
################################################################################
# MultiStateSystems.jl                                                         #
# A Julia package to solve multi-state system models.                          #
# See http://github.com/timmyfaraday/MultiStateSystems.jl                      #
################################################################################

"""
    EENS(usr)

Expected Energy Not Served (EENS) [MWh]

This is the amount of electricity demand of a user, measured in MWh, that is
expected not to be met in a given year.
    - EENS(usr) gives the EENS when an user's demand equals the maximal output
      of the system for that user.
"""
EENS(usr::PropDict) =
    8760u"hr"*sum((maximum(usr[:ugf].val).-usr[:ugf].val).*usr[:ugf].prb) |> u"MWh"

"""
    GRA(usr)

Generation Ratio Availability (GRA) [-]

This is the probability of at least transferring a specific percentage of the
generation, Generation Ration (GR) [-], to a user through the network.
    - GRA(usr,GR) gives the probability of transferring at least the GR to a
      user through the network.
    - GRA(usr) gives a user's GRA for a GR ranging from 0.0 to 1.0.
    - GRO(usr,GR) gives the output towards a user for a given GR
"""
GRA(usr::PropDict) = [GRA(usr,GR) for GR in 0.0:0.01:1.0]
GRA(usr::PropDict,GR::Float64) =
    sum(usr[:ugf].prb[findfirst(x -> x .>= GRO(usr,GR),usr[:ugf].val):end])
GRO(usr::PropDict,GR::Float64) = GR*maximum(usr[:ugf].val)


text:
# collection of high level basic monte carlo tasks

# high level data structues

abstract type AbstractChain end
abstract type AbstractState end
abstract type AbstractPolicy end
abstract type AbstractMove end
abstract type AbstractObservable end
abstract type ChainCollection end

mutable struct MarkovChain <: AbstractChain
    model::Dict # lattice, temperature, hamiltonian, etc...
    state::AbstractState # spins, fermions, whatever
    policy::AbstractPolicy # random choice, neighbors, ...
    # move::AbstractMove # flip a spin, exchange neighbors, ...
    mc_spec::Dict # MC steps, MC warmup, measurement interval
    observable::Any # some function
    data::Dict # outcomes of observable
    diagnostics::Dict # information on the MC run
    MarkovChain() = new()
end

# functions to access parameters of a generic markov chain
get_MarkovChain(c::MarkovChain) = c
get_Model(c::AbstractChain) = get_MarkovChain(c).model
get_State(c::AbstractChain) = get_MarkovChain(c).state
get_Policy(c::AbstractChain) = get_MarkovChain(c).policy
get_Observable(c::AbstractChain) = get_MarkovChain(c).observable
get_Diagnostics(c::AbstractChain) = get_MarkovChain(c).diagnostics
get_Mc_Spec(c::AbstractChain) = get_MarkovChain(c).mc_spec
get_Data(c::AbstractChain) = get_MarkovChain(c).data

# extension to chain collection
get_Chains(c::ChainCollection) = get_MarkovChain.(c.replicas)
get_GlobalChain(c::ChainCollection) = c.basechain
get_Model(c::ChainCollection) = get_Model(get_GlobalChain(c))
get_State(c::ChainCollection) = get_State.(get_Chains(c))
get_Policy(c::ChainCollection) = get_Policy(get_GlobalChain(c))
get_Observable(c::ChainCollection) = get_Observable(get_GlobalChain(c))
get_Diagnostics(c::ChainCollection) = get_Diagnostics(get_GlobalChain(c))
get_Mc_Spec(c::ChainCollection) = get_Mc_Spec(get_GlobalChain(c))
get_Data(c::ChainCollection) = get_Data(get_GlobalChain(c))

""" Run Monte Carlo"""
function runMC!(chain::AbstractChain)
    mc_specs = get_Mc_Spec(chain)
    NUM_MC_STEPS = mc_specs["mc_steps"]
    SAMPLE_INTERVAL = mc_specs["sample_interval"]

    # run burn in steps
    do_warmup!(chain)

    # carry out monte carlo sampling of observable
    for i = 1:NUM_MC_STEPS
        do_move!(chain)
        if mod(i,SAMPLE_INTERVAL) == 0
            measure_observable!(chain)
        end
    end

    # condense measurement statistics and diagnostics from MC sampling
    aggregate_measurements!(chain)
    aggregate_diagnostics!(chain)
end

function runMC!(chain::ChainCollection)
    mc_specs = get_Mc_Spec(chain)
    NUM_MC_STEPS = mc_specs["mc_steps"]
    SAMPLE_INTERVAL = mc_specs["sample_interval"]

    replicas = get_Replicas(chain)
    # run burn in steps
    do_warmup!.(replicas)

    # take first measurement
    observable = get_Observable(chain)

    if observable != Nothing
        chain.basechain.data = observable(chain)
    end

    # carry out monte carlo sampling of observable
    for i = 1:NUM_MC_STEPS
        do_move!.(replicas)
        if mod(i,SAMPLE_INTERVAL) == 0
            measure_observable!(chain)
            println("On measurerement ", i/SAMPLE_INTERVAL," out of ", NUM_MC_STEPS/SAMPLE_INTERVAL)
        end
    end

    # condense measurement statistics and diagnostics from MC sampling
    aggregate_measurements!(chain)
    # aggregate_diagnostics!(chain)
end

function aggregate_measurements!(chain::AbstractChain)
    data = get_Data(chain)
    diagnostics = get_Diagnostics(chain)

    for (key,val) in data
        data[key] = val/diagnostics["num_measurements"]
    end
end

function aggregate_measurements!(chain::ChainCollection)
    data = get_Data(chain)
    diagnostics = get_Diagnostics(chain)

    for (key,val) in data
        data[key] = val/diagnostics["num_measurements"]
    end
end

function aggregate_diagnostics!(chain::AbstractChain)
    diagnostics = get_Diagnostics(chain)
    diagnostics["acceptance_ratio"] =
    diagnostics["accepted_moves"]/diagnostics["mc_steps"]
end

# measures a generic observable
function measure_observable!(chain::AbstractChain)
    observable = get_Observable(chain)
    result = observable(chain) # expect a dict
    data = get_Data(chain)
    diagnostics = get_Diagnostics(chain)

    for (key, val) in result
        data[key] += val
    end

    diagnostics["num_measurements"] +=1
end

function measure_observable!(chain::ChainCollection)
    observable = get_Observable(chain)
    result = observable(chain) # expect a dict
    data = get_Data(chain)
    diagnostics = get_Diagnostics(chain)

    for (key, val) in result
        data[key] += val
    end

    diagnostics["num_measurements"] +=1
end

""" Initialize a markov chain """
function init_Chain!(chain::AbstractChain;
    model=model, observable=observable, policy=policy, mc_spec=mc_spec)
    chain.basechain = MarkovChain()
    chain.basechain.model = model
    chain.basechain.state = get_init_state(chain)
    chain.basechain.policy = policy
    chain.basechain.mc_spec = mc_spec
    chain.basechain.diagnostics = get_init_diagnostics(chain)
    chain.basechain.observable = observable
end

function init_Chain!(chain::ChainCollection;
    model=model, observable=observable, policy=policy, mc_spec=mc_spec)
        """ function overload for chain collection """
    # chain.basechain = MarkovChain() # for global data
    init_Chain!.(get_Replicas(chain),
    model=model, observable=Nothing, policy=policy, mc_spec=mc_spec) # initializes the replicas
    chain.basechain.policy = policy
    chain.basechain.model = model
    chain.basechain.mc_spec = mc_spec

    chain.basechain.diagnostics = get_init_diagnostics(chain)
    chain.basechain.observable = observable

end

function do_move!(chain)
    move = get_Move(chain)
    process_move!(chain,move)
end

# recursion for warmup

function do_warmup!(chain)
    mc_specs = get_Mc_Spec(chain)
    num_mc_warmup_steps = mc_specs["mc_warmup_steps"]
    for i = 1:num_mc_warmup_steps
        do_move!(chain)
    end

    observable = get_Observable(chain)

    if observable != Nothing
        chain.basechain.data = observable(chain)
    end
end

function process_move!(chain,move; forceaccept = false)
    """ Get the ratio for the proposed move
    and decide whether to accept or reject. Forceaccept
    is an optional keyword that can be useful for debugging"""
    ratio, extras = compute_ratio(chain,move) # extras is an optional
    # model based extra return
    forceaccept ? ACCEPT = true : ACCEPT =  ratio > rand()
    update_chain!(ACCEPT,chain,move,extras=extras)
end


text:
@testset "factorial_ij" begin
    for i in 0:40, j in i:40
        @test SqState.factorial_ij(i, j) == factorial(big(i)) / factorial(big(j))
        @test SqState.factorial_ij(j, i) == factorial(big(i)) / factorial(big(j))
    end
end

@testset "z" begin
    x_range = -10:0.01:10
    p_range = -6:0.01:6

    for x in x_range, p in p_range
        @test SqState.z(x, p) == sqrt(2.)*(x + p*im)
    end

    xs = collect(x_range)
    ps = collect(p_range)
    @test SqState.z(xs, ps) == sqrt(2.).*(xs .+ ps' .* im)
end

@testset "α" begin
    m_range = n_range = 1:35

    for m in m_range, n in m:n_range.stop
        @test SqState.α(m, n) == max(m, n) - min(m, n)
    end

    ms = collect(m_range)
    ns = collect(n_range)
    @test SqState.α(ms, ns) == max.(ms, ns') .- min.(ms, ns')
end

@testset "gaussian_function" begin
    tol = 1e-14

    x_range = -10:0.01:10
    p_range = -6:0.01:6

    for x in x_range, p in p_range
        @test isapprox(SqState.gaussian_function(x, p), exp(-(x^2 + p^2)) / π, atol=tol)
    end

    xs = collect(x_range)
    ps = collect(p_range)
    @test isapprox(
        SqState.gaussian_function(xs, ps),
        exp.(-(xs.^2 .+ (ps').^2)) ./ π,
        atol=tol
    )
end

@testset "(-1)^i" begin
    for i in 1:50
        @test SqState.neg_one_to_power_of(i) == (-1)^i
    end
end

@testset "coefficient_of_wave_function" begin
    tol = 1e-14

    m_range = n_range = 1:35

    for m in m_range, n in n_range
        m1 = (n ≥ m) ? (-1)^m : (-1)^n
        @test isapprox(
            SqState.coefficient_of_wave_function(m, n),
            m1 * sqrt(factorial(big(min(m, n))) / factorial(big(max(m, n)))),
            atol=tol
        )
    end

    ms = collect(m_range)
    ns = collect(n_range)
    @test SqState.coefficient_of_wave_function(ms, ns) == SqState.coefficient_of_wave_function.(ms, ns')
end

@testset "z_to_power" begin
    tol = 1e-14

    m_range = n_range = 1:35
    x_range = -1:0.1:1
    p_range = -0.6:0.1:0.6

    for m in m_range, n in n_range, x in x_range, p in p_range
        z = (n >= m) ? sqrt(2.)*(x - p*im) : sqrt(2.)*(x + p*im)
        @test SqState.z_to_power(m, n, x, p) == z^(max(m, n) - min(m, n))
    end

    ms = collect(m_range)
    ns = collect(n_range)
    xs = collect(x_range)
    ps = collect(p_range)
    xs_ = reshape(collect(x_range), 1, 1, length(x_range))
    ps_ = reshape(collect(p_range), 1, 1, 1, length(p_range))
    @test SqState.z_to_power(ms, ns, xs, ps) == SqState.z_to_power.(ms, ns', xs_, ps_)
    @test SqState.z_to_power(ms, ns)(xs, ps) == SqState.z_to_power.(ms, ns', xs_, ps_)
    @test SqState.z_to_power(xs, ps)(ms, ns) == SqState.z_to_power.(ms, ns', xs_, ps_)
end


text:
#=function wst_S1S2_deriv_fd(image::Array{Float64,2}, filter_hash::Dict)
    #Works now.
    #Possible Bug: You assumed zero freq in wrong place and N-k is the wrong transformation.
    # Use 2 threads for FFT
    FFTW.set_num_threads(2)

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

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

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

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

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

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

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

    #CartesianIndex(17,17) .- f_ind
    # tmp arrays for all tmp[f_i] = arr[f_i] .* f_v opns
    zarr1 = zeros(ComplexF64, Nx, Nx)
    fz_fψ1 = zeros(ComplexF64, Nx, Nx)
    fterm_a = zeros(ComplexF64, Nx, Nx)
    fterm_ct1 = zeros(ComplexF64, Nx, Nx)
    fterm_ct2 = zeros(ComplexF64, Nx, Nx)
    #rterm_bt2 = zeros(ComplexF64, Nx, Nx)
    # make a FFTW "plan" for an array of the given size and type
    P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)
    #phase_mat = reshape(collect(1:1:Nx), (Nx, 1)) .* reshape(collect(1:1:Nx), (1, Nx)) #Need to inst an Nx^2Nx2 mat. Instead just use it on the fly
    #phase_mat = 2π*
    #phase_mat = exp()
    # Loop over filters
    xImat = CartesianIndices((Nx, Nx))
    for f1 = 1:Nf
        f_i1 = f_ind[f1]  # CartesianIndex list for filter1
        f_v1 = f_val[f1]  # Values for f_i1
        f_i1rev = f_ind_rev[f1]

        zarr1[f_i1] = f_v1 .* im_fd_0[f_i1] #F[z]
        #zarr1_rd = P*zarr1 #for frzi
        fz_fψ1[f_i1] = f_v1 .* zarr1[f_i1]
        fz_fψ1_rd = P*fz_fψ1
        dS1dp[:,:,f1] = 2 .* real.(fz_fψ1_rd) #real.(conv(I_λ, ψ_λ))
        #Step 2: Do this using the fd space trick

        #dS2 loop prep
        #=I_λ1 = sqrt.(abs2.(zarr1_rd))
        fI_λ1 = fft(I_λ1)
        rterm_bt1 = zarr1_rd./I_λ1 #Z_λ1/I_λ1_bar.
        rterm_bt2 = conj.(zarr1_rd)./I_λ1=#

        for f2 = 1:Nf
            f_i2 = f_ind[f2]  # CartesianIndex list for filter2
            f_v2 = f_val[f2]  # Values for f_i2

            a1 = fz_fψ1[f_i2] .* (f_v2).^2
            b1 = exp.((-2π*im*) .* (f_i2 .* xImat))  #Shape: #f_i2*Nx*Nx
            term1 = a1 .* b1
            term2 =
            #fterm_ct2 = fterm_bt2 .* fcψ_λ1             #Slow
            #println(size(fterm_ct2[f_i1rev]),size(fterm_bt2[f_i1rev]),size(f_v1),size(conj.(f_v1)))
            fterm_ct2[f_i1rev] = fterm_bt2[f_i1rev] .* f_v1 #f_v1 is real
            #fterm_ct2slow = fterm_bt2 .* fcψ_λ1
            dS2dp[:, :, f1, f2] = real.(P*(fterm_ct1 + fterm_ct2))

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

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

    end
    return dS1dp, dS2dp
end
=#

#=
function wst_S1S2_deriv_fdsparse(image::Array{Float64,2}, filter_hash::Dict)
    #Works now.
    #Possible Bug: You assumed zero freq in wrong place and N-k is the wrong transformation.
    # Use 2 threads for FFT
    FFTW.set_num_threads(2)

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

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

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

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

    #Less Mem Route
    dS2dp = zeros(Float64, Nf, Nf, Nx, Nx)  # Lookup arrays for dS12 terms
    xigrid = reshape(CartesianIndices((1:Nx, 1:Nx)), (1, Nx^2))
    xigrid = hcat((x->[x[1], x[2]]).(xigrid)...) #2*Nx^2

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

        for f2 = 1:Nf
            f_i2 = f_ind[f2]  # CartesianIndex list for filter1
            f_v2 = f_val[f2]  # Values for f_i1
            pnz = intersect(f_i1, f_i2) #Fd indices where both filters are nonzero
            if length(pnz)!=0
                pnz_arr = vcat((x->[x[1], x[2]]').(pnz)...) #pnz*2: Assumption that the correct p to use in the phase are the indices!!! THIS IS WRONG instead create a kgrid and use pnz to index from that
                Φmat =  exp.((-2π*im)/Nx .* ((pnz_arr .- 1) * (xigrid .- 1))) #pnz*Nx^2
                f_v1pnz = f_v1[findall(in(pnz), f_i1)]
                f_v2pnz = f_v2[findall(in(pnz), f_i2)]
                t2 = real.(zarr[f1, pnz] .* Φmat)
                #println(size(f_v1pnz), size(t2), size(absz[f2, pnz]))
                term = 2*real.(sum((f_v2pnz.^2 .* f_v1pnz .*fim_fac[pnz]) .* t2, dims=1)) #p*Nx^2 -> 1*Nx^2
                dS2dp[f1, f2, :, :] = reshape(term/(Nx^2), (Nx, Nx))
            end
        end
    end
    return dS2dp
end
=#
#=
function realconv_test(Nx, ar1, ar2)
    soln = zeros(Float64, 8)
    for x=1:8:
        soln[x] = ar1 .* ar2[x-collect(1:1:8)]
=#

#=#Inner function code
#Code for img_reconfunc using wst_S2_deriv_sum: doesn't pass chi-square check even though inner machinery (derivsumcombtest, S1 and S2 separate, adding) all verified??

=#
#Old Feature Visualization
#=
fname_save = direc * "SFDTargSFDCov/" * ARGS[1] * "_" * ARGS[2] * "_" * ARGS[3] * "/LambdaVary/" * string(numfile) * ARGS[5]  #Change
Nx=size(true_img)[1]

(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=>apdbool, :iso=>isobool) #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] .= triu(trues(Nf, Nf))[:]
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(linesearch=LineSearches.BackTracking()))])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", logbool), ("GaussianLoss", true), ("Invcov_matrix", ARGS[6]),
  ("optim_settings", optim_settings), ("white_noise_args", white_noise_args)]) #Add constraints

recon_settings["datafile"] = datfile

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

if recon_settings["log"] & dhc_args[:apodize]
    lval = std(apodizer(log.(true_img))).^(-2)
elseif !recon_settings["log"] & dhc_args[:apodize]
    s_init = DHC_compute_wrapper(apodizer(init), filter_hash, norm=false; dhc_args...)[coeff_mask]
    l1init = ( 0.5 .* (s_init - s2mean)' * s2icov * (s_init - s2mean))[1]
    println("L1init", l1init)
    est1 = std(apodizer(true_img)).^(-2)
    wn_exp = mean(wind_2d(Nx).^2)*Nx*Nx*(loaddf["std"])^2
    lval = minimum([est1, 0.01*(l1init - 1.0)/(wn_exp)])
elseif recon_settings["log"] & !dhc_args[:apodize]
    lval = std(log.(true_img)).^(-2)
else
    s_init = DHC_compute_wrapper(init, filter_hash, norm=false; dhc_args...)[coeff_mask]
    l1init = ( 0.5 .* (s_init - s2mean)' * s2icov * (s_init - s2mean))[1]
    println("L1init", l1init)
    est1 = std(true_img).^(-2)
    wn_exp = Nx*Nx*(loaddf["std"])^2
    lval = minimum([est1, 0.01*(l1init - 1.0)/(wn_exp)])
end

recon_settings["lambda"] = lval
println("Regularizer Lambda=", round(lval, sigdigits=3))

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("Input Data", datfile)
println("Output File", fname_save)
println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("Mean Abs Res: Init-True = ", mean((init - true_img).^2).^0.5, " Recon-True = ", mean((recon_img - true_img).^2).^0.5)

#S1
fname_save = "scratch_NM/NewWrapper/FeatureVis/NotIso/New"
Nx=64
im = readsfd(Nx)
garbage_true = im[:, :, 34]
init = fill(mean(im), (Nx, Nx))

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=>false, :iso=>false) #Iso #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
if dhc_args[:iso]
    coeff_mask = falses(2+S1iso+S2iso)
    coeff_mask[S1iso+3:end] .= true
else #Not iso
    coeff_mask = falses(2+Nf+Nf^2)
    coeff_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", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.0)]) #Removed white_noise_args
#regs_true = DHC_compute_wrapper(log.(true_img), filter_hash; norm=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(garbage_true, 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(garbage_true, init, filter_hash, dhc_args, coeff_mask, recon_settings)

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

p=heatmap(recon_img, title="Init: Mean SFD, Optimizing only wrt SFD S20")
savefig(p, fname_save * "_result.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")

maximum(s2mean)
minimum(s2mean)
maximum(s2icov)
minimum(diag(s2icov))

#S1_Target
fname_save = "scratch_NM/NewWrapper/FeatureVis/NotIso/meaninit_allS1_target"
Nx=64
im = readsfd(Nx)
true_img = im[:, :, 1]
init = fill(mean(im), (Nx, Nx))
#heatmap(init)
#heatmap(true_img)

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false) #Iso #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
if dhc_args[:iso]
    coeff_mask = falses(2+S1iso+S2iso)
    coeff_mask[S1iso+3:end] .= true
else #Not iso
    coeff_mask = falses(2+Nf+Nf^2)
    coeffsS20 = falses(Nf, Nf)
    coeffsS20[diagind(coeffsS20)] .= true
    coeff_mask[Nf+3:end] .= coeffsS20[:]
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", "ground_truth"), ("covar_type", "sfd_dbn"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("lambda", 0.0)]) #Removed white_noise_args
#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: S1 only | Targ = S(SFD) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

p=heatmap(recon_img, title="Init: Mean SFD, Optimizing only wrt SFD S1")
savefig(p, fname_save * "_result.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")

heatmap(recon_img)

recon_img[1, 2]
fname = "scratch_NM/NewWrapper/FeatureVis/NotIso/meaninit_allS1"
gttarget = load(fname*".jld2")

##4-2 KLDiv Examination\
Nx=64
regsfd = readsfd(Nx, logbool=false)
dbn_coeffs_calc()
=#

#Old stuff from Deriv_Utils_New
#=
function sfd_derivsum_s20_reg(input::Array{Float64, 2}, filter_hash::Dict, s_targ_mean::Array{Float64, 1}, s_targ_invcov, pixmask::BitArray{2}, coeff_choice;
    FFTthreads::Int=1, optim_settings=Dict([("iterations", 10)]), coeff_mask=nothing, lambda=0.001)
    (Nf,) = size(filter_hash["filt_index"])
    println("Coeff mask:", (coeff_mask!=nothing))
    if coeff_choice!="S20" error("Not implemented") end
    (Nx, Ny)  = size(input)
    if Nx != Ny error("Input image must be square") end

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

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

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

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

    function loss_func20(img_curr)
        #size(img_curr) must be (Nx^2, 1)
        return LossGaussianS20(img_curr, filter_hash, ori_input, coeff_mask, s_targ_mean, s_invcov, lambda)
    end

    function dloss20(storage_grad, img_curr) #CHECK that this does update storage_grad
        #storage_grad, img_curr must be (Nx^2, 1)
        dLossGaussian(storage_grad, img_curr, filter_hash, ori_input, coeff_mask, s_targ_mean, s_invcov, lambda)
        return
    end

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

    res = optimize(loss_func20, dloss20, reshape(input, (Nx^2, 1)), ConjugateGradient(), Optim.Options(iterations = numitns_dict, store_trace = true, show_trace = true))
    result_img = Optim.minimizer(res)
end
=#
#=
function img_reconfunc_coeffmask(input, filter_hash, s_targ_mean, s_targ_invcov, pixmask, optim_settings; coeff_mask=nothing, coeff_type="S2")
    #=
    input: Initial image. When trues in pixmask, assumed to contain the correct value in those pixels.
    coeff_mask: Allows you to only optimize wrt a subset of coefficients (eg: j1<=j2)
    pixmask: Which pixels are floating (false) vs fixed (true)
    =#
    println("Coeff mask:", (coeff_mask!=nothing))
    (Nx, Ny)  = size(input)
    if Nx != Ny error("Input image must be square") end
    #filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)

    (Nf, )    = size(filter_hash["filt_index"])
    if Nf == 0 error("filter hash corrupted") end
    if coeff_type=="S2"
        println("S2")
        pixmask = pixmask[:] #flattened: Nx^2s
        cpvals = copy(input[:])[pixmask] #constrained pix values

        if coeff_mask!=nothing
            if length(coeff_mask)!= 2+Nf + Nf^2 error("Wrong dim mask") end
            if size(s_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
            #if ((s_targ_invcov!=I) & size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask)) error("s_targ_invcov should only contain coeffs to be optimized") end
            #if (s_targ_invcov!=I)  error("E1") end
            if size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("E2") end
        else #No mask: all coeffs (default: All S1 and S2 coeffs will be optim wrt)
            #Currently assuming inputs have mean, var params in case that's something we wanna optimize at some point
            if size(s_targ_mean)[1]== (2+Nf+Nf^2)
                s_targ_mean = s_targ_mean[3:end]
            else
                error("s_targ_mean of wrong size")
            end
            if (s_targ_invcov!=I) & (size(s_targ_invcov)==(2+Nf+Nf^2, 2+Nf+Nf^2))
                s_targ_invcov = s_targ_invcov[3:end, 3:end]
            else
                error("s_targ_invcov of wrong size")
            end
            #At this point both have dims S1+S2
            #3:end
            coeff_mask = fill(true, 2+Nf+Nf^2)
            coeff_mask[1] = false
            coeff_mask[2] = false
        end
        #After this all cases have a coeffmask
        function loss_func(img_curr)
            #println("Img",size(img_curr))
            s_curr = DHC_2DUtils.DHC_compute(reshape(img_curr, (Nx, Nx)),  filter_hash, doS2=true, doS12=false, doS20=false, norm=false)[coeff_mask]
            neglogloss = 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_targ_mean)
            return neglogloss[1]
        end

        function dloss(storage_grad, img_curr)
            s_curr = DHC_2DUtils.DHC_compute(reshape(img_curr, (Nx, Nx)), filter_hash, doS2=true, doS12=false, doS20=false, norm=false)[coeff_mask]
            diff = s_curr - s_targ_mean
            #diff = diff
            dS1S2 = reshape(wst_S1S2_deriv(reshape(img_curr, (Nx, Nx)), filter_hash), (Nx^2, Nf+Nf^2)) #(Nf+Nf^2, Nx, Nx)->(Nx^2, Nf+Nf^2) Cant directly do this without handling S1 coeffs
            dS1S2 = dS1S2[:, coeff_mask[3:end]] #Nx^2 * |SelCoeff|
            dS1S2[pixmask, :] .= 0 #Zeroing out wrt fixed params
            term1 = s_targ_invcov * diff #(Nf+Nf^2) or |SelCoeff| x 1
            #WARNING: Uses the Deriv_Utils version of dS1S2. Need to rewrite if you want to use the DHC_2DUtils one.
            #println("In dLoss:",size(diff), size(term1), size(term2))
            #println(size(term1),size(term2),size(storage_grad))
            mul!(storage_grad, dS1S2, term1) #Nx^2x1
        end
        #Debugging stuff
        println("Chisq derivative check")
        eps = zeros(size(input))
        eps[1, 2] = 1e-4
        chisq1 = loss_func(input+eps./2)
        chisq0 = loss_func(input-eps./2)
        brute  = (chisq1-chisq0)/1e-4

        clever = reshape(zeros(size(input)), (Nx*Nx, 1))
        _bar = dloss(clever, reshape(input, (Nx^2, 1)))
        println("Brute:  ",brute)
        println("Clever: ",clever[Nx*(1)+1])

        res = optimize(loss_func, dloss, reshape(input, (Nx^2, 1)), ConjugateGradient(), Optim.Options(iterations = get(optim_settings, "iterations", 100), store_trace = true, show_trace = true))
        result_img = zeros(Float64, Nx, Nx)
        result_img = Optim.minimizer(res)
    elseif coeff_type=="S12"
        println("S12")
        pixmask = pixmask[:] #flattened: Nx^2s
        cpvals = copy(input[:])[pixmask] #constrained pix values

        if coeff_mask!=nothing
            if length(coeff_mask)!= 2+Nf + Nf^2 error("Wrong dim mask") end
            if size(s_targ_mean)[1]!=count((i->(i==true)), coeff_mask) error("s_targ_mean should only contain coeffs to be optimized") end
            if size(s_targ_invcov)[1]!=count((i->(i==true)), coeff_mask) error("E2") end
            #Add WARNING that flags when the S1 are additionally incl in coeffmask
        else #No mask: all coeffs (default: All S1 and S2 coeffs will be optim wrt)
            #Currently assuming inputs have mean, var params in case that's something we wanna optimize at some point
            if size(s_targ_mean)[1]== (2+Nf+Nf^2)
                s_targ_mean = s_targ_mean[3:end]
            else
                error("s_targ_mean of wrong size")
            end
            if (s_targ_invcov!=I) & (size(s_targ_invcov)==(2+Nf+Nf^2, 2+Nf+Nf^2))
                s_targ_invcov = s_targ_invcov[3:end, 3:end]
            else
                error("s_targ_invcov of wrong size")
            end
            #At this point both have dims S1+S2
            #3:end
            coeff_mask = fill(true, 2+Nf+Nf^2)
            coeff_mask[1] = false
            coeff_mask[2] = false
        end
        #After this all cases have a coeffmask
        function loss_func12(img_curr)
            #println("Img",size(img_curr))
            s_curr = Deriv_Utils.DHC(reshape(img_curr, (Nx, Nx)),  filter_hash, doS2=false, doS12=true, doS20=false)[coeff_mask]
            neglogloss = 0.5 .* (s_curr - s_targ_mean)' * s_targ_invcov * (s_curr - s_targ_mean)
            return neglogloss[1]
        end

        function dloss12(storage_grad, img_curr)
            s_curr = Deriv_Utils.DHC(reshape(img_curr, (Nx, Nx)), filter_hash, doS2=false, doS12=true, doS20=false)[coeff_mask]
            diff = s_curr - s_targ_mean
            dS12 = reshape(permutedims(Deriv_Utils.wst_S12_deriv(reshape(img_curr, (Nx, Nx)), filter_hash), [3, 4, 1, 2]), (Nx^2, Nf^2)) #Nf, Nf, Nx, Nx -> (Nx^2, Nf^2)
            println(size(dS12))
            dS12 = dS12[:, coeff_mask[3+Nf:end]] #Nx^2 * |SelCoeff|
            dS12[pixmask, :] .= 0   #Zeroing out wrt fixed pix
            term1 = s_targ_invcov * diff #CHECK: Same as diff^Cinv? Nf^2 x1
            mul!(storage_grad, dS12, term1) #term1: 1x (Nf+Nf^2) | term2: #(Nf+Nf^2) x Nx^2
        end

        #Debugging stuff
        println("Chisq derivative check")
        eps = zeros(size(input))
        eps[1, 2] = 1e-4
        chisq1 = loss_func12(input+eps./2)
        chisq0 = loss_func12(input-eps./2)
        brute  = (chisq1-chisq0)/1e-4

        clever = reshape(zeros(size(input)), (Nx*Nx, 1))
        _bar = dloss12(clever, reshape(input, (Nx^2, 1)))
        println("Brute:  ",brute)
        println("Clever: ",clever[Nx*(1)+1])

        res = optimize(loss_func12, dloss12, reshape(input, (Nx^2, 1)), ConjugateGradient(), Optim.Options(iterations = get(optim_settings, "iterations", 100), store_trace = true, show_trace = true))
        result_img = zeros(Float64, Nx, Nx)
        result_img = Optim.minimizer(res)
    else error("S20 not implemented")
    end
    return reshape(result_img, (Nx, Nx))
end
=#

##5-31: Biased coefficients: wrong for even S1 (excluded the Second order terms)
#=
function DHC_compute_S20r_noisy(image::Array{Float64,2}, filter_hash::Dict, sigim::Array{Float64,2};
    doS2::Bool=true, doS20::Bool=false, apodize=false, norm=false, iso=false, FFTthreads=1, filter_hash2::Dict=filter_hash, coeff_mask=nothing)
    #Not using coeff_mask here after all. ASSUMES ANY ONE of doS12, doS2 and doS20 are true, when using coeff_mask
    #@assert !iso "Iso not implemented yet"
    Nf = size(filter_hash["filt_index"])[1]
    Nx = size(true_img)[1]

    if apodize
        ap_img = apodizer(image)
        dA = get_dApodizer(Nx, Dict([(:apodize => apodize)]))
    else
        ap_img = image
        dA = get_dApodizer(Nx, Dict([(:apodize => apodize)]))
    end

    function DHC_compute_biased(image::Array{Float64,2}, filter_hash::Dict, filter_hash2::Dict, sigim::Array{Float64,2};
        doS2::Bool=true, doS20::Bool=false, norm=true, iso=false, FFTthreads=2, normS1::Bool=false, normS1iso::Bool=false)
        # image        - input for WST
        # filter_hash  - filter hash from fink_filter_hash
        # filter_hash2 - filters for second order.  Default to same as first order.
        # doS2         - compute S2 coeffs
        # doS20        - compute S2 coeffs
        # norm         - scale to mean zero, unit variance
        # iso          - sum over angles to obtain isotropic coeffs

        # Use 2 threads for FFT
        FFTW.set_num_threads(FFTthreads)

        (Nx, Ny)  = size(image)
        if Nx != Ny error("Input image must be square") end
        (Nf, )    = size(filter_hash["filt_index"])
        if Nf == 0  error("filter hash corrupted") end
        @assert Nx==filter_hash["npix"] "Filter size should match npix"
        @assert Nx==filter_hash2["npix"] "Filter2 size should match npix"
        @assert (normS1 && normS1iso) != 1 "normS1 and normS1iso are mutually exclusive"

        # allocate coeff arrays
        out_coeff = []
        S0  = zeros(Float64, 2)
        S1  = zeros(Float64, Nf)
        if doS2  S2  = zeros(Float64, Nf, Nf) end  # traditional 2nd order
        if doS20 S20 = zeros(Float64, Nf, Nf) end  # real space correlation
        anyM2 = doS2 | doS20
        anyrd = doS2 | doS20             # compute real domain with iFFT

        # allocate image arrays for internal use
        if anyrd
            im_rd_0_1  = Array{Float64, 3}(undef, Nx, Ny, Nf)
            ψsqfac = Array{Float64, 4}(undef, Nf, Nx, Ny, 2)
            Iψfac = Array{ComplexF64, 4}(undef, Nf, Nx, Ny, 2)
            ψfaccross = Array{Float64, 2}(undef, Nf, Nf)
        end
        varim  = sigim.^2
        Pf = plan_fft(varim)
        fvar = Pf*(varim)

        ## 0th Order
        S0[1]   = mean(image)
        norm_im = image.-S0[1]
        S0[2]   = sum(norm_im .* norm_im)/(Nx*Ny)
        if norm
            norm_im ./= sqrt(Nx*Ny*S0[2])
        else
            norm_im = copy(image)
        end

        append!(out_coeff,S0[:])

        ## 1st Order
        im_fd_0 = Pf*(norm_im)  # total power=1.0

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

        zarr = zeros(ComplexF64, Nx, Ny)  # temporary array to fill with zvals

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

        ## Main 1st Order and Precompute 2nd Order
        for f = 1:Nf
            S1tot = 0.0
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i
            # for (ind, val) in zip(f_i, f_v)   # this is slower!
            if length(f_i) > 0
                for i = 1:length(f_i)
                    ind       = f_i[i]
                    zval      = f_v[i] * im_fd_0[ind]
                    S1tot    += abs2(zval)
                    zarr[ind] = zval        # filter*image in Fourier domain
                end
                S1[f] = S1tot/(Nx*Ny)  # image power
                if anyrd
                    psi_pow = sum(f_v.^2)./(Nx*Ny)
                    im_rd_0_1[:,:,f] .= abs2.(P*zarr)
                    rsψ = realspace_filter(Nx, f_i, f_v)
                    frealψ = Pf*(real.(rsψ))
                    fimψ = Pf*(imag.(rsψ))
                    ψsqfac[f, :, :, 1] = real.(rsψ) #check that real
                    ψsqfac[f, :, :, 2] = imag.(rsψ) #check that real
                    Iψfac[f, :, :, 1] = P*(im_fd_0 .* frealψ)
                    Iψfac[f, :, :, 2] = P*(im_fd_0 .* fimψ)
                end

                zarr[f_i] .= 0
            end
        end

        append!(out_coeff, iso ? filter_hash["S1_iso_mat"]*S1 : S1)

        if normS1iso
            S1iso = vec(reshape(filter_hash["S1_iso_mat"]*S1,(1,:))*filter_hash["S1_iso_mat"])
        end

        # we stored the abs()^2, so take sqrt (this is faster to do all at once)
        if anyrd
            im_rd_0_1 .= sqrt.(im_rd_0_1)
            for f1=1:Nf
                #println("f1", f1)
                for f2=1:Nf
                    #println("f2", f2)
                    val1 = Pf*(ψsqfac[f1, :, :, 1] .* ψsqfac[f2, :, :, 1])
                    #println(size(val1), size(fvar))
                    term1 = (P*(fvar .* val1) .* Iψfac[f1, :, :, 1]) .* Iψfac[f2, :, :, 1]
                    term2 = (P*(fvar .* (Pf*(ψsqfac[f1, :, :, 1] .* ψsqfac[f2, :, :, 2]))).* Iψfac[f1, :, :, 1]) .* Iψfac[f2, :, :, 2]
                    term3 = (P*(fvar .* (Pf*(ψsqfac[f1, :, :, 2] .* ψsqfac[f2, :, :, 1]))).* Iψfac[f1, :, :, 2]) .* Iψfac[f2, :, :, 1]
                    term4 = (P*(fvar .* (Pf*(ψsqfac[f1, :, :, 2] .* ψsqfac[f2, :, :, 2]))).* Iψfac[f1, :, :, 2]) .* Iψfac[f2, :, :, 2]
                    combined = sum((term1 + term2 + term3 + term4)./(im_rd_0_1[:, :, f1] .* im_rd_0_1[:, :, f2]))
                    println(imag(combined))
                    ψfaccross[f1, f2] = real(combined)
                end
            end
        end

        Mat2 = filter_hash["S2_iso_mat"]
        if doS2
            f_ind2   = filter_hash2["filt_index"]  # (J, L) array of filters represented as index value pairs
            f_val2   = filter_hash2["filt_value"]

            ## Traditional second order
            for f1 = 1:Nf
                thisim = fft(im_rd_0_1[:,:,f1])  # Could try rfft here
                # Loop over f2 and do second-order convolution
                if normS1
                    normS1pwr = S1[f1]
                elseif normS1iso
                    normS1pwr = S1iso[f1]
                else
                    normS1pwr = 1
                end

                for f2 = 1:Nf
                    f_i = f_ind2[f2]  # CartesianIndex list for filter
                    f_v = f_val2[f2]  # Values for f_i
                    # sum im^2 = sum(|fft|^2/npix)
                    #intfterm = fft((real.(rsψ1) + imag.(rsψ1)).*(real.(rsψ2) + imag.(rsψ2)))
                    S2[f1,f2] = sum(abs2.(f_v .* thisim[f_i]))/(Nx*Ny)/normS1pwr
                end
            end
            append!(out_coeff, iso ? Mat2*S2[:] : S2[:])
        end

        # Real domain 2nd order
        if doS20
            Amat = reshape(im_rd_0_1, Nx*Ny, Nf)
            S20  = Amat' * Amat
            #println(size(S20))
            S20 .+= ψfaccross #assuming the nx*ny factor above was for the parseval correction
            append!(out_coeff, iso ? Mat2*S20[:] : S20[:])
        end

        return out_coeff
    end

    sall = DHC_compute_biased(ap_img, filter_hash, filter_hash2, sigim, doS2=doS2, doS20=doS20, norm=norm, iso=iso, FFTthreads=FFTthreads)
    dS20dp = wst_S20_deriv(ap_img, filter_hash)
    ap_dS20dp = dA * reshape(dS20dp, Nx*Nx, Nf*Nf)
    G =  ap_dS20dp .* reshape(sigim, Nx*Nx, 1)
    cov = G'*G

    if coeff_mask!=nothing
        @assert count(!iszero, coeff_mask[1:Nf+2])==0 "This function only handles S20r"
        @assert length(coeff_mask)==length(sall) "Mask must have the same length as the total number of coefficients"
        s20rmask = coeff_mask[Nf+3:end]
        return sall[coeff_mask], cov[s20rmask, s20rmask]
    else
        return sall, cov
    end

end
=#

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
import TaylorModels

function integrate(f, I; steps = 1024, degree = 6)
    lo = I.lo
    hi = I.hi
    l = 2*radius(I)
    int_center = Interval(0.0)
    int_error = Interval(0.0)
    for i in 1:steps
        left = lo+(i-1)*(l/steps)
        right = lo+i*l/steps
        r = (1/2)*(l/steps)
        J = interval(left, right)
        TM = TaylorModels.TaylorModel1(degree, J)
        FM = f(TM)
        #@info FM
        for i in 0:Int64(floor(degree/2))
            int_center+=2*(FM.pol[2*i]*r^(2*i+1))/(2*i+1)
            int_error +=2*FM.rem*r
        end
    end
    return int_center+int_error
end

function adaptive_integration(f, I::Interval; tol = 2^-10, steps = 8, degree = 6) # tol 2^-10, steps = 8 are default values
    lo = I.lo
    hi = I.hi
    l = 2*radius(I)
    int_value = Interval(0.0)
    for i in 1:steps
        left = lo+(i-1)*(l/steps)
        right = lo+i*l/steps
        Istep = Interval(left, right)
        val = integrate(f, Istep)
        if radius(val)<tol
            int_value += val
        else
            I₁, I₂ = bisect(I)
            val₁ = adaptive_integration(f, I₁; tol = tol/2, steps = steps, degree = degree+2)
            val₂ = adaptive_integration(f, I₁; tol = tol/2, steps = steps, degree = degree+2)
            int_value +=val₁+val₂
        end
    end
    return int_value
end

struct Observable
    B
    v::Vector
    infbound
end

### TODO: Actually some assumptions are made, as the fact that
# the Ulam base is equispaced
function Observable(B::Ulam, ϕ::Function; tol = 2^-10)
    v = zeros(Interval{Float64}, length(B))
    for i in 1:length(B)
        I = Interval(B.p[i], B.p[i+1])
        v[i] = adaptive_integration(ϕ, I; tol = tol, steps =1, degree = 2)*length(B)
    end
    infbound = maximise(x->abs(ϕ(x)), Interval(0,1))[1]
    return Observable(B, v, infbound)
end

import TaylorSeries

function discretizationlogder(B, D::PwMap; degree = 7)
    v = zeros(Interval{Float64}, length(B))
    infbound  = emptyinterval()
    endpoints = D.endpoints
    delicate_indexes = Int64.(floor.(length(B)*[x.lo for x in endpoints])).+1
    for i in 1:(length(delicate_indexes)-1)
        for j in delicate_indexes[i]:delicate_indexes[i+1]-1
            I = Interval(B.p[j], B.p[j+1])
            r = Interval(radius(I))
            Tmid = TaylorSeries.Taylor1([Interval(mid(I)), Interval(1)], degree)
            Tint = TaylorSeries.Taylor1([I, Interval(1)], degree)
            Fmid = log(TaylorModels.derivative(D.Ts[i](Tmid)))
            Fint = log(TaylorModels.derivative(D.Ts[i](Tint)))
            infbound = infbound ∪ abs(Fint[0])
            ϵ = mag(Fint[degree]-Fmid[degree])

            for k in 0:Int64(floor(Float64(degree)/2))
                v[j]+=2*(Fmid[2*k]*r^(2*k+1))/(2*k+1)
            end
            v[j]+=Interval(-ϵ, ϵ)*r^(degree+1)/(degree+1)
        end
    end

    #correction since the endpoints may be wide intervals
    for i in 2:length(endpoints)-1
        x = endpoints[i]
        Tx = TaylorSeries.Taylor1([x, Interval(1)], 1)

        corr = 2*Interval(radius(x))*(abs(log(TaylorModels.derivative(D.Ts[i-1](Tx))))
                +abs(log(TaylorModels.derivative(D.Ts[i](Tx)))))[0]
        v[delicate_indexes[i]]+=corr
    end
    v*=length(B)
    return Observable(B, v, infbound)
end

function integrateobservable(B::Ulam, ϕ::Observable, f::Vector, error)
    val = (ϕ.v)'*f
    return val/length(B)+(ϕ.infbound.hi)*Interval(-error, error)
end


text:
using OMEinsumContractionOrders, Test, Random
using OMEinsumContractionOrders: random_exprtree, ExprTree, ExprInfo, ruleset, update_tree!, tcscrw, optimize_subtree!, optimize_tree_sa!, labels, tree_timespace_complexity, fast_log2sumexp2
using OMEinsum, Graphs

@testset "random expr tree" begin
    function random_regular_eincode(n, k)
        g = Graphs.random_regular_graph(n, k)
        ixs = [minmax(e.src,e.dst) for e in Graphs.edges(g)]
        return EinCode((ixs..., [(i,) for i in Graphs.vertices(g)]...), ())
    end
    Random.seed!(2)
    tree = random_exprtree([[1,2,5], [2,3], [2,4]], [5], 5)
    @test tree isa OMEinsumContractionOrders.ExprTree
    tree2 = random_exprtree(EinCode(((1,2,5), (2,3), (2,4)), (5,)))
    @test tree isa OMEinsumContractionOrders.ExprTree
    code = random_regular_eincode(20, 3)
    optcode = optimize_greedy(code, uniformsize(code, 2))
    tree3 = ExprTree(optcode)
    @test tree isa OMEinsumContractionOrders.ExprTree
    labelmap = Dict([v=>k for (k,v) in OMEinsumContractionOrders._label_dict(code)])
    optcode_reconstruct = NestedEinsum(tree3, labelmap)
    @test optcode == optcode_reconstruct
end

@testset "rules" begin
    LeafNode(id, labels) = ExprTree(ExprInfo(labels, id))
    t1 = ExprTree(LeafNode(3, [1,2]), ExprTree(LeafNode(1,[2,3]), LeafNode(2,[1,4]), ExprInfo([1,2])), ExprInfo([2]))
    t2 = ExprTree(ExprTree(LeafNode(1, [2,3]), LeafNode(2, [1,4]), ExprInfo([1,2,3])), LeafNode(3,[1,2]), ExprInfo([2]))
    t3 = ExprTree(LeafNode(1,[2,3]), LeafNode(2, [1,2]), ExprInfo([2]))
    t4 = ExprTree(ExprTree(LeafNode(1, [2,3]), LeafNode(2, [1,4]), ExprInfo([1,2])), ExprTree(LeafNode(4,[5,1]), LeafNode(3,[1]), ExprInfo([1])), ExprInfo([2]))
    @test ruleset(t1) == 3:4
    @test ruleset(t2) == 1:2
    @test ruleset(t3) == 1:0
    @test ruleset(t4) == 1:4
    log2_sizes = ones(5)
    _tcsc(t, l) = tcscrw(labels(t.left), labels(t.right), labels(t), l, true)
    @test all(_tcsc(t1, log2_sizes) .≈ (2.0, 1.0, log2(10)))
    @test all(_tcsc(t2, log2_sizes) .≈ (2.0, 1.0, log2(14)))
    @test all(_tcsc(t3, log2_sizes) .≈ (1.0, 1.0, log2(10)))
    @test all(_tcsc(t4, log2_sizes) .≈ (2.0, 1.0, log2(8)))
    t11 = update_tree!(copy(t1), 3, [2])
    @test t11 == ExprTree(LeafNode(1,[2,3]), ExprTree(LeafNode(3,[1,2]), LeafNode(2,[1,4]), ExprInfo([2])), ExprInfo([2]))
    t11_ = update_tree!(copy(t1), 4, [1,2])
    @test t11_ == ExprTree(LeafNode(2,[1,4]), ExprTree(LeafNode(1,[2,3]), LeafNode(3,[1,2]), ExprInfo([1,2])), ExprInfo([2]))
    t22 = update_tree!(copy(t2), 1, [1,2])
    @test t22 == ExprTree(ExprTree(LeafNode(1,[2,3]), LeafNode(3,[1,2]), ExprInfo([1,2])), LeafNode(2,[1,4]), ExprInfo([2]))
    t22_ = update_tree!(copy(t2), 2, [2])
    @test t22_ == ExprTree(ExprTree(LeafNode(3, [1,2]), LeafNode(2,[1,4]), ExprInfo([2])), LeafNode(1,[2,3]), ExprInfo([2]))
    t44 = update_tree!(copy(t4), 1, [1,2])
    @test t44 == ExprTree(ExprTree(LeafNode(1,[2,3]), ExprTree(LeafNode(4,[5,1]), LeafNode(3,[1]), ExprInfo([1])), ExprInfo([1,2])), LeafNode(2,[1,4]), ExprInfo([2]))
end

@testset "optimization" begin
    function random_regular_eincode(n, k)
        g = Graphs.random_regular_graph(n, k)
        ixs = [minmax(e.src,e.dst) for e in Graphs.edges(g)]
        return EinCode((ixs..., [(i,) for i in Graphs.vertices(g)]...), ())
    end
    Random.seed!(2)
    n = 40
    log2_sizes = rand(n+n÷2) * 2
    code = random_regular_eincode(n, 3)
    optcode = optimize_greedy(code, uniformsize(code, 2))
    tree = ExprTree(optcode)
    tc0, sc0, rw0 = tree_timespace_complexity(tree, log2_sizes)
    size_dict = Dict([j=>exp2(log2_sizes[j]) for j=1:length(log2_sizes)])
    tc0_, sc0_ = OMEinsum.timespace_complexity(NestedEinsum(tree), size_dict)
    @test tc0 ≈ tc0_ && sc0 ≈ sc0_
    opt_tree = copy(tree)
    optimize_subtree!(opt_tree, 100.0, log2_sizes, 5, 2.0, 1.0)
    tc1, sc1, rw0 = tree_timespace_complexity(opt_tree, log2_sizes)
    @test sc1 < sc0 || (sc1 == sc0 && tc1 < tc0)
end

@testset "optimize tree sa" begin
    function random_regular_eincode(n, k)
        g = Graphs.random_regular_graph(n, k)
        ixs = [minmax(e.src,e.dst) for e in Graphs.edges(g)]
        return EinCode((ixs..., [(i,) for i in Graphs.vertices(g)]...), ())
    end
    Random.seed!(3)
    n = 60
    ne = n + n÷2
    log2_sizes = ones(ne)
    code = random_regular_eincode(n, 3)
    optcode = optimize_greedy(code, uniformsize(code, 2))
    tree = ExprTree(optcode)
    tc0, sc0, rw0 = tree_timespace_complexity(tree, log2_sizes)
    opttree = copy(tree)
    optimize_tree_sa!(opttree, log2_sizes, Slicer(log2_sizes, 0); sc_target=sc0-2.0, βs=0.1:0.1:10.0, niters=100, sc_weight=1.0, rw_weight=1.0)
    tc1, sc1, rw1 = tree_timespace_complexity(opttree, log2_sizes)
    @test sc1 < sc0 || (sc1 == sc0 && tc1 < tc0)

    slicer = Slicer(log2_sizes, 5)
    optimize_tree_sa!(opttree, log2_sizes, slicer; sc_target=sc0-2.0, βs=0.1:0.1:10.0, niters=100, sc_weight=1.0, rw_weight=1.0)
    tc2, sc2, rw2 = tree_timespace_complexity(opttree, slicer.log2_sizes)
    @test tc2 <= tc1 + 3
    @test sc2 <= sc1 + 3
    @test length(slicer) == 5
    @test all(l->(slicer.log2_sizes[l]==1 && !haskey(slicer.legs, l)) || (slicer.log2_sizes[l]==0 && haskey(slicer.legs, l)), 1:length(log2_sizes))
    @test sc1 < sc0 || (sc1 == sc0 && tc1 < tc0)
end

@testset "sa tree" begin
    function random_regular_eincode(n, k)
        g = Graphs.random_regular_graph(n, k)
        ixs = [minmax(e.src,e.dst) for e in Graphs.edges(g)]
        return EinCode((ixs..., [(i,) for i in Graphs.vertices(g)]...), ())
    end
    Random.seed!(2)
    g = random_regular_graph(220, 3)
    code = random_regular_eincode(220, 3)
    res = optimize_greedy(code,uniformsize(code, 2))
    tc, sc = OMEinsum.timespace_complexity(res, uniformsize(code, 2))

    @test optimize_tree(res,uniformsize(code, 2); sc_target=32, βs=0.1:0.05:20.0, ntrials=0, niters=10, sc_weight=1.0, rw_weight=1.0) isa SlicedEinsum
    optcode = optimize_tree(res,uniformsize(code, 2); sc_target=32, βs=0.1:0.05:20.0, ntrials=2, niters=10, sc_weight=1.0, rw_weight=1.0)
    tc, sc = OMEinsum.timespace_complexity(optcode, uniformsize(code, 2))
    @test sc <= 32
    @test length(optcode.slicing) == 0

    # contraction test
    code = random_regular_eincode(50, 3)
    codek = optimize_greedy(code, uniformsize(code, 2))
    codeg = optimize_tree(code, uniformsize(code, 2); initializer=:random)
    tc, sc = OMEinsum.timespace_complexity(codek, uniformsize(code, 2))
    @test sc <= 12
    xs = [[2*randn(2, 2) for i=1:75]..., [randn(2) for i=1:50]...]
    resg = codeg(xs...)
    resk = codek(xs...)
    @test resg ≈ resk

    # contraction test
    code = random_regular_eincode(50, 3)
    codek = optimize_greedy(code, uniformsize(code, 2))
    codeg = optimize_tree(codek, uniformsize(code, 2); initializer=:specified)
    tc, sc = OMEinsum.timespace_complexity(codek, uniformsize(code, 2))
    @test sc <= 12
    xs = [[2*randn(2, 2) for i=1:75]..., [randn(2) for i=1:50]...]
    resg = codeg(xs...)
    resk = codek(xs...)
    @test resg ≈ resk
end

@testset "fast log2sumexp2" begin
    a, b, c = randn(3)
    @test fast_log2sumexp2(a, b) ≈ log2(sum(exp2.([a,b])))
    @test fast_log2sumexp2(a, b, c) ≈ log2(sum(exp2.([a,b,c])))
end

@testset "slicing" begin
    function random_regular_eincode(n, k)
        g = Graphs.random_regular_graph(n, k)
        ixs = [minmax(e.src,e.dst) for e in Graphs.edges(g)]
        return EinCode((ixs..., [(i,) for i in Graphs.vertices(g)]...), ())
    end
    # contraction test
    Random.seed!(2)
    code = random_regular_eincode(100, 3)
    code0 = optimize_greedy(code, uniformsize(code, 2))
    codek = optimize_tree(code0, uniformsize(code, 2); initializer=:specified, nslices=0, niters=5)
    codeg = optimize_tree(code0, uniformsize(code, 2); initializer=:specified, nslices=5, niters=5)
    tc0, sc0 = OMEinsum.timespace_complexity(codek, uniformsize(code, 2))
    tc, sc = OMEinsum.timespace_complexity(codeg, uniformsize(code, 2))
    @show tc, sc, tc0, sc0
    @test sc <= sc0 - 4
    xs = [[2*randn(2, 2) for i=1:150]..., [randn(2) for i=1:100]...]
    resg = codeg(xs...)
    resk = codek(xs...)
    @test resg ≈ resk

    # with open edges
    Random.seed!(2)
    code = EinCode(random_regular_eincode(100, 3).ixs, [3,81,2])
    codek = optimize_tree(code0, uniformsize(codek, 2); initializer=:specified, nslices=0, niters=5)
    codeg = optimize_tree(code0, uniformsize(codeg, 2); initializer=:specified, nslices=5, niters=5)
    tc0, sc0 = OMEinsum.timespace_complexity(codek, uniformsize(code, 2))
    tc, sc = OMEinsum.timespace_complexity(codeg, uniformsize(code, 2))
    fl = OMEinsum.flop(codeg, uniformsize(code, 2))
    @test tc ≈ log2(fl)
    @show tc, sc, tc0, sc0
    @test sc <= sc0 - 4
    @test sc <= 17
    xs = [[2*randn(2, 2) for i=1:150]..., [randn(2) for i=1:100]...]
    resg = codeg(xs...)
    resk = codek(xs...)
    @test resg ≈ resk
end