text:
abstract type TensorStorage{ElT} <: AbstractVector{ElT} end

data(S::TensorStorage) = S.data

datatype(S::TensorStorage) = typeof(data(S))

Base.eltype(::TensorStorage{ElT}) where {ElT} = ElT

Base.eltype(::Type{<:TensorStorage{ElT}}) where {ElT} = ElT

Base.iterate(S::TensorStorage, args...) = iterate(data(S), args...)

dense(S::TensorStorage) = S

# This is necessary since for some reason inference doesn't work
# with the more general definition (eltype(Nothing) === Any)
Base.eltype(::TensorStorage{Nothing}) = Nothing

Base.length(S::TensorStorage) = length(data(S))

Base.size(S::TensorStorage) = size(data(S))

Base.@propagate_inbounds Base.getindex(S::TensorStorage, i::Integer) = data(S)[i]
Base.@propagate_inbounds function Base.setindex!(S::TensorStorage, v, i::Integer)
  return (setindex!(data(S), v, i); S)
end

(S::TensorStorage * x::Number) = setdata(S, x * data(S))
(x::Number * S::TensorStorage) = S * x
(S::TensorStorage / x::Number) = setdata(S, data(S) / x)

-(S::TensorStorage) = setdata(S, -data(S))

similar(S::TensorStorage) = setdata(S, similar(data(S)))
similar(S::TensorStorage, x) = setdata(S, similar(data(S), x))

# Define Base.similar in terms of NDTensors.similar
Base.similar(t::TensorStorage, args...) = similar(t, args...)
# To fix method ambiguity
Base.similar(t::TensorStorage) = similar(t)
Base.similar(t::TensorStorage, T::Type) = similar(t, T)
Base.similar(t::TensorStorage, d::Tuple) = similar(t, d)
Base.similar(t::TensorStorage, d::Vararg{Union{Integer,AbstractUnitRange}}) = similar(t, d)

# Needed for passing Tensor{T,2} to BLAS/LAPACK
function Base.unsafe_convert(::Type{Ptr{ElT}}, T::TensorStorage{ElT}) where {ElT}
  return Base.unsafe_convert(Ptr{ElT}, data(T))
end

# This may need to be overloaded, since storage types
# often have other fields besides data

Base.conj!(S::TensorStorage) = (conj!(data(S)); return S)

Base.conj(S::TensorStorage) = conj(AllowAlias(), S)

function Base.conj(::AllowAlias, S::TensorStorage)
  return setdata(S, conj(data(S)))
end

function Base.conj(::NeverAlias, S::TensorStorage)
  return conj!(copy(S))
end

Base.complex(S::TensorStorage) = setdata(S, complex(data(S)))

Base.real(S::TensorStorage) = setdata(S, real(data(S)))

Base.imag(S::TensorStorage) = setdata(S, imag(data(S)))

Base.copyto!(S1::TensorStorage, S2::TensorStorage) = (copyto!(data(S1), data(S2)); S1)

Random.randn!(S::TensorStorage) = (randn!(data(S)); S)

Base.fill!(S::TensorStorage, v) = (fill!(data(S), v); S)

LinearAlgebra.rmul!(S::TensorStorage, v::Number) = (rmul!(data(S), v); S)

scale!(S::TensorStorage, v::Number) = rmul!(S, v)

norm(S::TensorStorage) = norm(data(S))

Base.convert(::Type{T}, S::T) where {T<:TensorStorage} = S

blockoffsets(S::TensorStorage) = S.blockoffsets

"""
nzblocks(T::TensorStorage)

Return a vector of the non-zero blocks of the BlockSparse storage.
"""
nzblocks(T::TensorStorage) = nzblocks(blockoffsets(T))

eachnzblock(T::TensorStorage) = eachnzblock(blockoffsets(T))

nnzblocks(S::TensorStorage) = length(blockoffsets(S))
nnz(S::TensorStorage) = length(S)

offset(S::TensorStorage, block) = offset(blockoffsets(S), block)


text:
include("tree_mesh.jl")
include("structured_mesh.jl")
include("surface_interpolant.jl")
include("unstructured_mesh.jl")
include("p4est_mesh.jl")
include("mesh_io.jl")

ext: jl


text:
## KWArgs

abstract type KWArgs <: AbstractDict{Symbol, String}; end

mutable struct OwnedKWArgs <: KWArgs
    ptr::SoapySDRKwargs
    function OwnedKWArgs(kw::SoapySDRKwargs)
        this = new(kw)
        finalizer(SoapySDRKwargs_clear, this)
        this
    end
end
Base.unsafe_load(o::OwnedKWArgs) = o.ptr
function ptr(kw::OwnedKWArgs)
    return pointer_from_objref(kw)
end
function SoapySDRKwargs_clear(kw::OwnedKWArgs)
    SoapySDRKwargs_clear(ptr(kw))
end

mutable struct KWArgsList <: AbstractVector{KWArgs}
    ptr::Ptr{SoapySDRKwargs}
    length::Csize_t
    function KWArgsList(ptr::Ptr{SoapySDRKwargs}, length::Csize_t)
        this = new(ptr, length)
        finalizer(this) do this
            SoapySDRKwargsList_clear(this.ptr, this.length)
        end
        this
    end
end
Base.size(kwl::KWArgsList) = (kwl.length,)

"""
    KWArgsListRef

This is returned by calls to e.g. `Devices` and is used to present
arguments for device creation. The Julia API for this should be
similar to a `Dict`.
"""
struct KWArgsListRef <: KWArgs
    list::KWArgsList
    idx::Int
    function Base.getindex(kwl::KWArgsList, i::Integer)
        checkbounds(kwl, i)
        new(kwl, i)
    end
end

function ptr(kwl::KWArgsListRef)
    @assert kwl.list.ptr !== C_NULL
    kwl.list.ptr + (kwl.idx-1)*sizeof(SoapySDRKwargs)
end

function Base.setindex!(kwl::KWArgsListRef, val, key)
    SoapySDRKwargs_set(ptr(kwl), key, val)
end

function Base.get(kwl::KWArgsListRef, key, default)
    res = SoapySDRKwargs_get(ptr(kwl), key)
    res == C_NULL ? default : unsafe_string(res)
end

Base.unsafe_load(kw::KWArgs) = unsafe_load(ptr(kw))
Base.length(kw::KWArgs) = unsafe_load(kw).size
function _getindex(kw::KWArgs, i::Integer)
    1 <= i <= length(kw) || throw(BoundsError(kw, i))
    @GC.preserve kw begin
        return Symbol(unsafe_string(unsafe_load(unsafe_load(kw).keys, i))) =>
               unsafe_string(unsafe_load(unsafe_load(kw).vals, i))
    end
end

Base.iterate(kw::KWArgs, i=1) = i > length(kw) ? nothing : (_getindex(kw, i), i+1)

function SoapySDRKwargsList_clear(list::KWArgsList)
    SoapySDRKwargsList_clear(list.ptr, list.length)
    list.ptr = C_NULL
end

##

mutable struct StringList <: AbstractVector{String}
    strs::Ptr{Cstring}
    length::Csize_t
    function StringList(strs::Ptr{Cstring}, length::Integer)
        this = new(strs, Csize_t(length))
        finalizer(SoapySDRStrings_clear, this)
        this
    end
end
Base.size(s::StringList) = (s.length,)
function Base.getindex(s::StringList, i::Integer)
    checkbounds(s, i)
    unsafe_string(unsafe_load(s.strs, i))
end

SoapySDRStrings_clear(s::StringList) = @GC.preserve s SoapySDRStrings_clear(pointer_from_objref(s), s.length)

function Base.show(io::IO, s::SoapySDRArgInfo)
    println(io, "name: ", unsafe_string(s.units))
    println(io, "key: ", unsafe_string(s.key))
    println(io, "value: ", unsafe_string(s.name))
    println(io, "description: ", unsafe_string(s.description))
    println(io, "units: ", unsafe_string(s.units))
    #type
    #range
    println(io, "options: ", StringList(s.options, s.numOptions))
    println(io, "optionNames: ", StringList(s.optionNames, s.numOptions))
end


text:
using Random
using Printf
using StatsBase
using StaticArrays
using Distributions
using POMDPs
using AutoViz
using AutomotiveDrivingModels
using AutomotivePOMDPs
using POMDPSimulators
using POMDPPolicies
using DeepQLearning
using Flux
using RLInterface
using POMDPModelTools
using BeliefUpdaters
using TensorBoardLogger
using ArgParse
using BSON
using CSV
using LinearAlgebra
using ProgressMeter
include("environment.jl")
include("generative_mdp.jl")
include("masking.jl")
include("cooperative_IDM.jl")
include("belief_updater.jl")
include("belief_mdp.jl")

# BLAS.set_num_threads(8)

s = ArgParseSettings()
@add_arg_table s begin
    "--seed"
        help = "specify the random seed"
        arg_type = Int
        default = 1
    "--training_steps"
        help = "specify the number of training steps"
        arg_type = Int
        default = 10_000
    "--eps_fraction"
        help = "Fraction of the training set to use to decay epsilon from 1.0 to eps_end"
        arg_type = Float64
        default = 0.5
    "--eps_end"
        help = "epsilon value at the end of the exploration phase"
        arg_type = Float64
        default = 0.01
    "--learning_rate"
        help = "learning rate for DQN"
        arg_type = Float64
        default = 1e-4
    "--target_update_freq"
        help = "target update frequency for DQN"
        default = 5000
    "--logdir"
        help = "Directory in which to save the model and log training data"
        arg_type = String
        default = "log"
    "--n_eval"
        help = "Number of episodes for evaluation"
        arg_type = Int64
        default = 1000
    "--max_cars"
        help = "Number of cars on the main road"
        arg_type = Int64
        default = 12
    "--min_cars"
        help = "Number of cars on the main road"
        arg_type = Int64
        default = 0
    "--cooperation"
        help = "whether the ego vehicle observe the cooperation level or not"
        action = :store_true
    "--driver_type"
        help = "the distribution of driver types, choose among cooperative, aggressive or random"
        arg_type = String
        default = "random"
    "--traffic_speed"
        arg_type = String
        default = "mixed"
    "--recurrent"
        help = "whether to use an RNN in DQN"
        action = :store_true
    "--load"
        help = "load a pretrain policy"
        arg_type = Union{Nothing, String}
        default = nothing
    "--collision_cost"
        help = "collision cost in the mdp formulation"
        arg_type = Float64
        default = -1.0
    "--belief"
        help = "train with belief updater"
        action = :store_true
end
parsed_args = parse_args(s)

seed = parsed_args["seed"]
Random.seed!(seed)
rng = MersenneTwister(seed)

mdp = GenerativeMergingMDP(random_n_cars = true,
                            collision_cost = parsed_args["collision_cost"],
                            max_cars = parsed_args["max_cars"],
                            min_cars = parsed_args["min_cars"],
                           driver_type = Symbol(parsed_args["driver_type"]),
                           traffic_speed = Symbol(parsed_args["traffic_speed"]),
                           observe_cooperation=parsed_args["cooperation"])

if parsed_args["belief"]
    for i=2:mdp.max_cars+1
        mdp.driver_models[i] = CooperativeIDM()
    end
    pomdp = FullyObservablePOMDP(mdp)
    up = MergingUpdater(mdp)
    bmdp = GenerativeBeliefMDP{typeof(pomdp), MergingUpdater, MergingBelief,Int64}(pomdp, up)
    s0 = initialstate(bmdp, rng)
    mdp = bmdp
else
    s0 = initialstate(mdp, rng)
end

svec = convert_s(Vector{Float64}, s0, mdp)

input_dims = length(svec)

if parsed_args["recurrent"]
    model = Chain(LSTM(input_dims, 64), Dense(64, 32, relu), Dense(32, n_actions(mdp)))
else
    model = Chain(Dense(input_dims, 64, relu), Dense(64, 32, relu), Dense(32, n_actions(mdp)))
end

if parsed_args["load"] != nothing
    BSON.@load parsed_args["load"] policy
    Flux.loadparams!(model, Flux.params(policy.qnetwork))
end

solver = DeepQLearningSolver(qnetwork = model,
                      max_steps = parsed_args["training_steps"],
                      eps_fraction = parsed_args["eps_fraction"],
                      eps_end = parsed_args["eps_end"],
                      eval_freq = 10_000,
                      save_freq = 10_000,
                      target_update_freq = parsed_args["target_update_freq"],
                      buffer_size = 10_000,
                      batch_size = 32,
                      learning_rate = parsed_args["learning_rate"],
                      train_start = 1000,
                      buffer_size = 400_000,
                      log_freq = 1000,
                      num_ep_eval = 1000,
                      double_q = true,
                      dueling = false,
                      prioritized_replay = true,
                      verbose = true,
                      recurrence = parsed_args["recurrent"],
                      rng = rng,
                      logdir = parsed_args["logdir"])

@time policy = solve(solver, mdp)

env = MDPEnvironment(mdp)
DeepQLearning.evaluation(solver.evaluation_policy,
                policy, env,
                solver.num_ep_eval,
                solver.max_episode_length,
                solver.verbose)

function quick_evaluation(mdp, policy::Policy, rng::AbstractRNG, n_eval=1000)
    avg_r, avg_dr, c_rate, avg_steps, t_out = 0.0, 0.0, 0.0, 0.0, 0.0
    @showprogress for i=1:n_eval
        s0 = initialstate(mdp, rng)
        hr = HistoryRecorder(rng = rng, max_steps=100)
        hist = simulate(hr, mdp, policy, s0)
        avg_r += undiscounted_reward(hist)
        avg_dr += discounted_reward(hist)
        c_rate += undiscounted_reward(hist) <= -1.0
        t_out += n_steps(hist) >= hr.max_steps ? 1.0 : 0.0
        avg_steps += n_steps(hist)
    end
    avg_r /= n_eval
    avg_dr /= n_eval
    c_rate /= n_eval
    avg_steps /= n_eval
    t_out /= n_eval
    return avg_r, avg_dr, c_rate, avg_steps, t_out
end

avg_r, avg_dr, c_rate, avg_steps, t_out = quick_evaluation(mdp, policy, rng, 1000)

println("Collisions ", c_rate*100)
println("Avg steps ", avg_steps)
println("Time outs ", t_out*100)
println("avg disc reward ", avg_dr)
println("avg reward ", avg_r)

BSON.@save joinpath(parsed_args["logdir"], "policy_$(parsed_args["cooperation"]).bson") policy

BSON.@save "results.bson" c_rate avg_steps t_out avg_dr avg_r


text:
#=
""" Give the boundary nodes """
function bdy_nodes(dims::Tuple)
    D = length(dims)
    nbdy = prod(dims) - prod(dims .- 2)
    bdy = CartesianIndex{D}[]
    for (i, d) in enumerate(dims)
        # This only works with julia v 1.6
        push!(bdy, CartesianIndices((dims[1:(i-1)]..., 1:(d-1):d, dims[(i+1):D]...))[:]...)
    end
    bdy = unique(bdy)
    neighbors = zeros(length(bdy), D)
    for (j, b) in enumerate(bdy)
        for i in 1:D
          neighbors[j, i] = ((Tuple(b)[i] == 1)||(Tuple(b)[i] == dims[i])) ? 1 : 2
        end
    end
    return bdy, neighbors
end
=#

"""
  bdy_nodes(dims)

Boundary node computation, for arbitrary dimension

...
# Arguments

  - `dims::Tuple` number of points in each direction
...

...
# Outputs
  - `Vector{Int64}`: vector containing the indices of coordinates
  on the boundary of the hyperrectangle volume
...

"""
function bdy_nodes(dims::Tuple)
    D = length(dims)
    bdy = CartesianIndex[]
    for d in CartesianIndices(dims)
        if (sum(Tuple(d) .== 1) > 0) || (sum(Tuple(d) .== dims) > 0)
            push!(bdy, d)
        end
    end
    neighbors = zeros(Int64, length(bdy), D)
    for (j, b) in enumerate(bdy)
        for i in 1:D
          neighbors[j, i] = ((Tuple(b)[i] == 1)||(Tuple(b)[i] == dims[i])) ? 1 : 2
        end
    end
    return bdy, neighbors
end

""" Write the matrix """
function nablasq_arb(dims, aspect_ratios)
  D = length(dims)
  oh = map(i -> ones(dims[i])/aspect_ratios[i], 1:length(dims))
  del = Vector{SparseArrays.SparseMatrixCSC{Float64, Int64}}(undef, D)
  for (i, d) in enumerate(dims)
    del[i] = spdiagm_nonsquare(d + 1, d, -1 => -oh[i], 0 => oh[i])
  end
  A = zeros(prod(dims), prod(dims))
  for i in 1:D
      A += kron([sparse(I, dims[j], dims[j]) for j in 1:(i-1)]..., del[i]'*del[i],
                [sparse(I, dims[j], dims[j]) for j in (i+1):D]...)
  end
  bdy, neighbors = bdy_nodes(dims)
  for (count, b) in enumerate(bdy)
    i = LinearIndices(zeros(dims...))[b]
    A[i, i] = 0.0
    for j in 1:D
      A[i, i] += neighbors[count, j] / aspect_ratios[j] ^ 2
    end
  end
  return A
end

""" Helper function to give the Matern matrix in arbitrary dimensions """
function _Matern_matrix(dims, m, eps, aspect_ratios)
    A = nablasq_arb(dims, aspect_ratios)
    sizeA = size(A, 1)
    for i = 1:sizeA
        A[i, i] = A[i, i] + eps^2
    end
    return A^m
end

""" Compute the interpolation """
function interp(data, ind, m, eps, aspect_ratios)
    dims = size(data)
    A = (eps == 0.0)&&(m == 1) ?
                nablasq_arb(dims, aspect_ratios) :
                _Matern_matrix(dims, m, eps, aspect_ratios)
    totalsize = prod(dims)
    C = sparse(I, totalsize, totalsize)
    rhs_a = data[:]
    for i in ind
        j = (typeof(i) <: CartesianIndex) ? LinearIndices(data)[i] : i
        C[j, j] = 0.0
        rhs_a[j] = 0.0
    end
    Id = sparse(I, totalsize, totalsize)
    u = ((C - (Id - C) * A)) \ rhs_a
    return reshape(u, dims...)
end

max_stars_repo_name: rayosborn/nxpeaks

max_issues_repo_name: rayosborn/nxpeaks

max_forks_repo_name: axmas-anl/nxrefine


text:
x = readlines("inbuild_funs.txt");
xmls = filter(e->endswith(e, ".xml\n"), x)
# TODO actually parse XML
map!(xmls) do x
    replace(x, ".xml\n", "")
end
final = filter(xmls) do x
    isdefined(Symbol(x)) # only usefull if also defined in julia
end
n_ary_map = Dict{Int, Set{Symbol}}()
for elem in final
    f = Symbol(elem)
    for m in methods(eval(f))
        i = length(m.sig.parameters) - 1
        set = get!(n_ary_map, i, Set{Symbol}())
        push!(set, f)
    end
end

open("gl_inbuild_funs.jl") do io
    println(io, "Dict(")
    for (k, v) in n_ary_map
        println(io, "    ", k, " => ", collect(v))
    end
    println(io, ")")
end


text:
"""
    FlexibilityVariable <: JuMP.AbstractJuMPScalar
An abstract type to define new variable types.
"""
abstract type FlexibilityVariable <: JuMP.AbstractJuMPScalar end

"""
    RandomVariable <: FlexibilityVariable
A DataType for random variables.

**Fields**
- `m::Model` Flexibility model.
- `idx::Int` Index of variable in model.
"""
struct RandomVariable <: FlexibilityVariable
    m::Model
    idx::Int
end

"""
    RecourseVariable <: FlexibilityVariable
A DataType for recourse variables.

**Fields**
- `m::Model` Flexibility model.
- `idx::Int` Index of variable in model.
"""
struct RecourseVariable <: FlexibilityVariable
    m::Model
    idx::Int
end

"""
    FlexibilityExpr <: JuMP.GenericAffExpr
A `GenericAffExpr` that contains random and/or recourse variables.
"""
const FlexibilityExpr = JuMP.GenericAffExpr{AffExpr, Union{RandomVariable, RecourseVariable}}
FlexibilityExpr() = FlexibilityExpr(RandomVariable[], AffExpr[], AffExpr())

"""
    FlexibilityConstraint <: JuMP.AbstractConstraint
A constraint that contains random and/or recourse variables.

**Fields**
- `flex_expr::FlexibilityExpr` Constraint expression.
- `sense::Symbol` The cosntraint sense symbol `:(<=)` or `:(>=)` or `:(==)`.
"""
mutable struct FlexibilityConstraint <: JuMP.AbstractConstraint
    flex_expr::FlexibilityExpr
    sense::Symbol # :(<=) or :(>=) or :(==), right-hand side assumed to be zero
end

"""
    AbstractUncertaintySet
An abstract type to define new uncertainty set types.
"""
abstract type AbstractUncertaintySet end

"""
    EllipsoidalSet <: AbstractUncertaintySet
An ellipsoidal uncertainty set that will use the covariance matrix stored in [`FlexibilityData`](@ref).

**Fields**
- `name::Symbol` Name of the set which will be `:Ellipsoid`.
- `only_positive::Bool` An option to indicate if the set should be intersected with the set all positive numbers.
"""
mutable struct EllipsoidalSet <: AbstractUncertaintySet
    name::Symbol
    only_positive::Bool
end

"""
    HyperboxSet <: AbstractUncertaintySet
A hyperbox uncertainty set whose nomimal dimensions are determined by `neg_dev` and `pos_dev`.

**Fields**
- `name::Symbol` Name of the set which will be `:Hyperbox`.
- `neg_dev::Vector{Number}` A vector of the expected negative deviation of the random variables.
- `pos_dev::Vector{Number}` A vector of the expected positive deviation of the random variables.
- `only_positive::Bool` An option to indicate if the set should be intersected with the set all positive numbers.
"""
mutable struct HyperboxSet <: AbstractUncertaintySet
    name::Symbol
    neg_dev::Vector{Number}
    pos_dev::Vector{Number}
    only_positive::Bool
end

"""
    PNormSet <: AbstractUncertaintySet
A p-norm based uncertainty set based on a bounded p-norm.

**Fields**
- `name::Symbol` Name of the set which will be `:PNorm`.
- `p::Number` The value of p which can be 1, 2, or Inf.
- `only_positive::Bool` An option to indicate if the set should be intersected with the set all positive numbers.
"""
mutable struct PNormSet <: AbstractUncertaintySet
    name::Symbol
    p::Number
    only_positive::Bool
end

"""
    FlexibilityData

A DataType for storing the data necessary to manage the bookkeeping of the flexibility variables (`RandomVariable`
and `RecourseVariable`), the uncertainty set, and solution results.

**Fields**

- `flexibility_constraints::Vector{JuMP.AbstractConstraint}` Constraints that involve flexibility variables.
- `numRVs::Int` The number of `RandomVariable` that have been added to the model.
- `RVmeans::Vector{Number}` The means corresponding to each `RandomVariable`.
- `RVnames::Vector{AbstractString}` The symbolic name of each `RandomVariable`.
- `RVcols::Vector{Int}` The index of each `RandomVariable`.
- `num_recourse_vars::Int` The number of `RecourseVariable` that have been added to the model.
- `recourse_names::Vector{AbstractString}` The symbolic name of each `RecourseVariable`.
- `recourse_cols::Vector{Int}` The index of each `RecourseVariable`.
- `uncertainty_set::AbstractUncertaintySet` The uncertainty set DataType with all of the set specfic attributes.
- `covariance::Matrix{Number}` The covariance matrix.
- `flexibility_index::Union{Nothing, Number}` The flexibility index result obtained from solving the flexibility model.
- `active_constraints::Vector{Int}` The indexes of the active inequality constraints at the solution of the flexibility model.
- 'solution_time::Number' The solution time in seconds.
"""
mutable struct FlexibilityData
    flexibility_constraints::Vector{JuMP.AbstractConstraint}

    # Random variable data
    numRVs::Int
    RVmeans::Vector{Number}
    RVnames::Vector{AbstractString}
    RVcols::Vector{Int}

    # Recourse variable data
    num_recourse_vars::Int
    recourse_names::Vector{AbstractString}
    recourse_cols::Vector{Int}

    # Various formulation data/results
    uncertainty_set::AbstractUncertaintySet  #e.g. ellipsoidal, 1-norm, 2-norm, etc...
    covariance::Matrix{Number}
    flexibility_index::Union{Nothing, Number}
    active_constraints::Vector{Int}
    solution_time::Union{Nothing, Number}
end

# Set methods to extract covariance information if appropriate
EllipsoidalSet() = EllipsoidalSet(:Ellipsoid, false)


text:
using Statistics
using Plots
using BenchmarkTools
using Profile
using FFTW
using Statistics
using Optim
using Images, FileIO, ImageIO
using Printf
using SparseArrays

#using Profile
using LinearAlgebra

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

@time fhash = fink_filter_hash(1, 8, nx=16, pc=1, wd=1)

sum(a2, dims=1)

function realspace_filter(Nx, f_i, f_v)

    zarr = zeros(ComplexF64, Nx, Nx)
    for i = 1:length(f_i)
        zarr[f_i[i]] = f_v[i] # filter*image in Fourier domain
    end
    filt = ifft(zarr)  # real space, complex
    return filt
end

function DHC(image::Array{Float64,2}, filter_hash::Dict, filter_hash2::Dict=filter_hash;
    doS2::Bool=true, doS12::Bool=false, doS20::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
    # doS12        - compute S2 coeffs
    # doS20        - compute S2 coeffs

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

    # array sizes
    (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

    # 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 doS12 S12 = zeros(Float64, Nf, Nf) end  # Fourier correlation
    if doS20 S20 = zeros(Float64, Nf, Nf) end  # real space correlation
    anyM2 = doS2 | doS12 | doS20
    anyrd = doS2 | doS20             # compute real domain with iFFT

    # allocate image arrays for internal use
    if doS12 im_fdf_0_1 = zeros(Float64,           Nx, Ny, Nf) end   # this must be zeroed!
    if anyrd im_rd_0_1  = Array{Float64, 3}(undef, Nx, Ny, Nf) end

    ## 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])
    # println("norm off")
    norm_im = copy(image)

    append!(out_coeff,S0[:])

    ## 1st Order
    im_fd_0 = fft(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
                if doS12 im_fdf_0_1[ind,f] = abs(zval) end
            end
            S1[f] = S1tot/(Nx*Ny)  # image power
            if anyrd
                im_rd_0_1[:,:,f] .= abs2.(P*zarr) end
            zarr[f_i] .= 0
        end
    end
    append!(out_coeff, S1[:]) #Why need the :?

    # 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) end

    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
            # println("  f1",f1,"  sum(fft):",sum(abs2.(thisim))/Nx^2, "  sum(im): ",sum(abs2.(im_rd_0_1[:,:,f1])))
            # Loop over f2 and do second-order convolution
            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)
                S2[f1,f2] = sum(abs2.(f_v .* thisim[f_i]))/(Nx*Ny)
            end
        end
        append!(out_coeff, S2)
    end

    # Fourier domain 2nd order
    if doS12
        Amat = reshape(im_fdf_0_1, Nx*Nx, Nf)
        S12  = Amat' * Amat
        append!(out_coeff, S12)
    end

    # Real domain 2nd order
    if doS20
        Amat = reshape(im_rd_0_1, Nx*Nx, Nf)
        S20  = Amat' * Amat
        append!(out_coeff, S20)
    end

    return out_coeff
end

##Derivative Functions

function wst_S1_deriv(image::Array{Float64,2}, filter_hash::Dict)
    function conv(a,b)
        ifft(fft(a) .* fft(b))
    end

    # 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
    dSdp  = zeros(Float64, Nx, Nx, Nf)

    # allocate image arrays for internal use
    #im_rdc_0_1 = Array{ComplexF64, 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"]

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

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

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

        zarr[f_i] = f_v .* im_fd_0[f_i] #Why do we need .* here??
        #Commented: I_λ = P*zarr  # complex valued ifft of zarr
        #im_rdc_0_1[:,:,f] = I_lambda

        #xfac = 2 .* real(I_λ)
        #yfac = 2 .* imag(I_λ)
        # it is clearly stupid to transform back and forth, but we will need these steps for the S20 part
        ψ_λ  = realspace_filter(Nx, f_i, f_v)
        # convolution requires flipping wavelet direction, equivalent to flipping sign of imaginary part.
        # dSdp[:,:,f] = real.(conv(xfac,real(ψ_λ))) - real.(conv(yfac,imag(ψ_λ)))
        zarr[f_i] = f_v .* zarr[f_i] #Reusing zarr for F[conv(z_λ,ψ_λ)] = F[z_λ].*F[ψ_λ]
        dSdp[:,:,f] = 2 .* real.(P*(zarr)) #real.(conv(I_λ, ψ_λ))
        zarr[f_i] .= 0   # reset zarr for next loop

    end
    return dSdp

end

function wst_S1_deriv_fast(image::Array{Float64,2}, filter_hash::Dict)

    # 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
    dSdp  = zeros(Float64, Nx, Nx, Nf)

    # allocate image arrays for internal use
    #im_rdc_0_1 = Array{ComplexF64, 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"]

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

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

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

        zarr[f_i] = f_v .* im_fd_0[f_i] #Why do we need .* here??
        #Old: I_λ = P*zarr  # complex valued ifft of zarr, Should have been z_lambda in the latex conv
        #im_rdc_0_1[:,:,f] = I_lambda

        #xfac = 2 .* real(I_λ)
        #yfac = 2 .* imag(I_λ)
        # it is clearly stupid to transform back and forth, but we will need these steps for the S20 part
        #Old: ψ_λ  = realspace_filter(Nx, f_i, f_v)
        # convolution requires flipping wavelet direction, equivalent to flipping sign of imaginary part.
        # dSdp[:,:,f] = real.(conv(xfac,real(ψ_λ))) - real.(conv(yfac,imag(ψ_λ)))
        zarr[f_i] = f_v .* zarr[f_i] #Reusing zarr for F[conv(z_λ,ψ_λ)] = F[z_λ].*F[ψ_λ] = fz_fψ
        dSdp[:,:,f] = 2 .* real.(P*(zarr)) #real.(conv(I_λ, ψ_λ))
        zarr[f_i] .= 0   # reset zarr for next loop

    end
    return dSdp

end

function wst_S20_deriv(image::Array{Float64,2}, filter_hash::Dict) #make this faster
    function conv(a,b)
        ifft(fft(a) .* fft(b))
    end

    # 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
    dS20dp  = 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"]

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

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

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

        zarr[f_i] = f_v .* im_fd_0[f_i]
        I_λ = P*zarr  # complex valued ifft of zarr
        zarr[f_i] .= 0   # reset zarr for next loop
        im_rdc_0_1[:,:,f] = I_λ
        im_rd_0_1[:,:,f]  = abs.(I_λ)
    end

    for f2 = 1:Nf
        ψ_λ  = realspace_filter(Nx, f_ind[f2], f_val[f2])
        uvec = im_rdc_0_1[:,:,f2] ./ im_rd_0_1[:,:,f2]
        for f1 = 1:Nf
            cfac = im_rd_0_1[:,:,f1].*uvec
            I1dI2 = real.(conv(cfac,ψ_λ))  #NM: Make the conv to fft change here too
            dS20dp[:,:,f1,f2] += I1dI2
            dS20dp[:,:,f2,f1] += I1dI2
        end
    end

    return dS20dp

end

function wst_S1S2_deriv(image::Array{Float64,2}, filter_hash::Dict)
    #Works as is: some superfluous convolution and FFT opns. Working on fixing this.
    # 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)
    #debug

    # 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"]

    # 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)
    #rterm_bt2 = zeros(ComplexF64, Nx, Nx)
    # make a FFTW "plan" for an array of the given size and type
    P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)

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

        zarr1[f_i1] = f_v1 .* im_fd_0[f_i1] #F[z]
        zarr1_rd = P*zarr1
        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_λ, ψ_λ))
        #CHECK; that this equals derivS1fast code

        #dS2 loop prep
        ψ_λ1  = realspace_filter(Nx, f_i1, f_v1)
        fcψ_λ1 = fft(conj.(ψ_λ1))
        I_λ1 = sqrt.(abs2.(zarr1_rd)) #NM: Alloc+find faster way
        fI_λ1 = fft(I_λ1)
        rterm_bt1 = zarr1_rd./I_λ1 #Z_λ1/I_λ1_bar. Check that I_lambda_bar equals I_lambda. MUST be true.
        rterm_bt2 = conj.(zarr1_rd)./I_λ1
        for f2 = 1:Nf
            f_i2 = f_ind[f2]  # CartesianIndex list for filter2
            f_v2 = f_val[f2]  # Values for f_i2

            fterm_a[f_i2] = fI_λ1[f_i2] .* (f_v2).^2 #fterm_a = F[I_λ1].F[ψ_λ]^2
            rterm_a = P*fterm_a                #Finv[fterm_a]
            fterm_bt1 = fft(rterm_a .* rterm_bt1) #F[Finv[fterm_a].*Z_λ1/I_λ1_bar] for T1
            fterm_bt2 = fft(rterm_a .* rterm_bt2) #F[Finv[fterm_a].*Z_λ1_bar/I_λ1_bar] for T2
            fterm_ct1[f_i1] = fterm_bt1[f_i1] .* f_v1    #fterm_b*F[ψ_λ]
            fterm_ct2 = fterm_bt2 .* fcψ_λ1    #Slow version
            dS2dp[:, :, f1, f2] = real.((P*fterm_ct1) + (P*fterm_ct2)) #CHECK: Should this 2 not be there?? Why does it give an answer thats larger by 2x

            #Reset f2 specific tmp arrays to 0
            fterm_a[f_i2] .= 0
            fterm_ct1[f_i1] .=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_derivfast(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)

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

        zarr1[f_i1] = f_v1 .* im_fd_0[f_i1] #F[z]
        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_λ, ψ_λ))
        #CHECK; that this equals derivS1fast code

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

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

            #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_S12_deriv(image::Array{Float64,2}, filter_hash::Dict)
    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"]

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

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

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

    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)]
                t2 = real.(zarrc[f1, pnz] .* Φmat) #Check
                #println(size(f_v1pnz), size(t2), size(absz[f2, pnz]))
                term = sum(((absz[f2, pnz] ./ absz[f1, pnz]) .* f_v1pnz) .* t2, dims=1) #p*Nx^2 -> 1*Nx^2
                dabsz_absz[f1, f2, :, :] = reshape(term, (Nx, Nx))
            end
        end
    end
    dS12 = dabsz_absz + permutedims(dabsz_absz, [2, 1, 3, 4])

end

function wst_S2_deriv_sum(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64})
    #Currently assumes a weight matrix for only the S2 -- wt[S1] must be handled outside
    FFTW.set_num_threads(2)

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

    if size(wt)!=(Nf,Nf) error("Wt has wrong shape: the argument should contain only S2 weights") end

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

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

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

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

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

    end
    return dLossdα

end

function wst_S1S2_derivsum_comb(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64})
    # array sizes
    (Nx, Ny)  = size(image)
    (Nf, )    = size(filter_hash["filt_index"])

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

    dwS1 = reshape(DHC_2DUtils.wst_S1_deriv(image, filter_hash), (Nx^2, Nf)) * wts1
    dwS2 = reshape(wst_S2_deriv_sum(image, filter_hash, wts2), (Nx^2, 1))
    println(size(dwS1), size(dwS2))
    return dwS1 + dwS2

end
##Derivative test functions
# wst_S1_deriv agrees with brute force at 1e-10 level.
function derivtestS12(Nx)
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    dS12dp = wst_S12_deriv(im, fhash)

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

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

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

function dS1S2sum_combtest(fhash) #Derivsum S2works here too
    (Nf, )    = size(fhash["filt_index"])
    Nx        = fhash["npix"]
    im = rand(Nx,Nx)
    mywts = rand(Nf + (Nf*Nf))
    println(size(mywts))

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

    # Compare to established code
    dS1dp, dS2dp = wst_S1S2_derivfast(im, fhash)
    term1 = reshape(dS1dp, (Nx^2, Nf)) * reshape(mywts[1:Nf], (Nf, 1))
    dS2dp = reshape(dS2dp, (Nx, Nx, Nf^2))
    sum2 = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (dS2dp[:,:,i].*mywts[i+Nf]) end
    sum2 += reshape(term1, (Nx, Nx))
    sum2 = sum2[:]
    println("tot abs mean", mean(abs.(sum2 -sum1)))
    println("tot Stdev: ",std(sum2-sum1))
    println("S1")
    println("Abs mean", mean(abs.(term1 - dwS1)))
    println("Abs mean", std(term1 - dwS1))
    println("S2 separate")
    sum1 = wst_S2_deriv_sum(im, fhash, reshape(mywts[Nf+1:end], (Nf, Nf)))
    # Compare to established code
    dS1dp, dS2dp = wst_S1S2_derivfast(im, fhash)
    dS2dp = reshape(dS2dp, (Nx, Nx, Nf^2))
    sum2 = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (dS2dp[:,:,i].*mywts[Nf+i]) end
    println("abs mean", mean(abs.(sum2 -sum1)))
    println("Stdev: ",std(sum1-sum2))

    return
end

function dS2sum_test(fhash) #DerivsumS2 works here
    (Nf, )    = size(fhash["filt_index"])
    Nx        = fhash["npix"]
    im = rand(Nx,Nx)
    mywts = rand(Nf*Nf)

    wtvec  = reshape(mywts, (Nf, Nf))

    # Use new faster code
    sum1 = wst_S2_deriv_sum(im, fhash, wtvec)

    # Compare to established code
    dS1dp, dS2dp = wst_S1S2_derivfast(im, fhash)
    dS2dp = reshape(dS2dp, (Nx, Nx, Nf^2))
    sum2 = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (dS2dp[:,:,i].*mywts[i]) end
    println("abs mean", mean(abs.(sum2 -sum1)))
    println("Stdev: ",std(sum1-sum2))

    return
end

dS12code, dS12lim = derivtestS12(16)

 function derivtest(Nx)
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    blarg = wst_S1_deriv(im, fhash)
    println("Bl size",size(blarg))

    der0=DHC(im0,fhash,doS2=false)
    der1=DHC(im1,fhash,doS2=false)
    dS = (der1-der0) ./ eps #since eps in only (2, 3) this should equal blarg[2, 3, :]

    diff = dS[3:end]-blarg[2,3,:]
    println(dS[3:end])
    println("and")
    println(blarg[2,3,:])
    println("stdev: ",std(diff))

    #=#Delete: blarg shouldnt be 0 for all other pixels since the coefficient isnt constant wrt other pixels!!
    Nf = length(fhash["filt_index"])
    mask = trues((Nx,Nx, Nf))
    println("Nf",Nf,"Check",(size(mask)==size(blarg)))
    mask[2, 3, :] .= false
    println("Mean deriv wrt to other pixels (should be 0 for all)",mean(blarg[mask]), "Maximum",maximum(blarg[mask]))
    #For Nx=128, Mean is actually 0.002 Seems concerning since eps is 1e-4!? the max is 0.10
    #For Nx=8, Mean is 0.01, max is 1.0 -- some kind of normalization issue?=#
    return
end

 function derivtestfast(Nx, mode="fourier")
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    blarg = wst_S1_deriv(im, fhash)
    if mode=="fourier"
        blargfast = wst_S1S2_deriv_fdsparse(im, fhash)
    else
        blargfast = wst_S1_deriv_fast(im, fhash)
    end
    der0=DHC(im0,fhash,doS2=false)
    der1=DHC(im1,fhash,doS2=false)
    dS = (der1-der0) ./ eps

    diff_old = dS[3:end]-blarg[2,3,:]
    diff = dS[3:end]-blargfast[2,3,:]
    println(dS[3:end])
    println("and")
    println(blarg[2,3,:])
    println("stdev: ",std(diff_old))
    println("New fewer conv ds1")
    println(dS[3:end])
    println("and")
    println(blargfast[2,3,:])
    println("stdev: ",std(diff))
    #=#Delete: blarg shouldnt be 0 for all other pixels since the coefficient isnt constant wrt other pixels!!
    #Check: that the derivative is 0 for all pixels?? THink about this... need to check against autodiff?
    Nf = length(fhash["filt_index"])
    mask = trues((Nx,Nx, Nf))
    println("Nf",Nf,"Check",(size(mask)==size(blargfast)))
    mask[2, 3, :] .= false
    println("Mean deriv wrt to other pixels (should be 0 for all)",mean(blargfast[mask]), "Maximum",maximum(blargfast[mask]))
    #For Nx=128, Mean is actually 0.002, the max is 0.10. Seems concerning since eps is 1e-4!?
    #For Nx=8, Mean is 0.01, max is 0.90 -- some kind of normalization issue?
    #plot(dS[3:end])
    #plot!(blarg[2,3,:])
    #plot(diff)=#
    return
end

dS2code, dS2lim = derivtestS1S2(16, mode="fourier")
fhash= fink_filter_hash(1, 8, nx=16, pc=1, wd=1)

dS2sum_test(fhash)
dS1S2sum_combtest(fhash)

function derivtestS1S2(Nx; mode="slow")
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    dS1dp_existing = wst_S1_deriv(im, fhash)
    if mode=="slow"
        dS1dp, dS2dp = wst_S1S2_deriv(im, fhash)
    elseif mode=="fourier"
        dS2dp = wst_S1S2_deriv_fdsparse(im, fhash)
    elseif mode=="faster"
        dS2dp = wst_S1S2_deriv_fdsparse(im, fhash)
    else
        dS1dp, dS2dp = wst_S1S2_derivfast(im, fhash)
    end

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

    dS1lim23 = dSlim[3:lasts1ind]
    dS2lim23 = dSlim[lasts1ind+1:end]
    if mode!="fourier"
        dS1dp_23 = dS1dp[2, 3, :]
    end
    dS2dp_23 = dS2dp[2, 3, :, :]
    #=
    println("Checking dS1dp using existing deriv")
    Nf = length(fhash["filt_index"])
    lasts1ind = 2+Nf
    diff = dS1dp_existing[2, 3, :]-dSlim[3:lasts1ind]
    println(dSlim[3:lasts1ind])
    println("and")
    println(dS1dp_existing[2, 3, :])
    println("stdev: ",std(diff))
    println("--------------------------")=#
    if mode!="fourier"
        println("Checking dS1dp using dS1S2")
        derdiff = dS1dp_23 - dS1lim23
        #println(dS1lim23)
        #println("and")
        #println(dS1dp_23)
        txt= @sprintf("Range of S1 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E ",maximum(der0[3:lasts1ind]),maximum(der1[3:lasts1ind]),minimum(der0[3:lasts1ind]),minimum(der1[3:lasts1ind]))
        print(txt)
        txt = @sprintf("Checking dS1dp using S1S2 deriv, mean(abs(diff/dS1lim)): %.3E", mean(abs.(derdiff./dS1lim23)))
        print(txt)
        println("stdev: ",std(derdiff))
        txt= @sprintf("Range of dS1lim: Max= %.3E Min= %.3E",maximum(abs.(dS1lim23)),minimum(abs.(dS1lim23)))
        print(txt)
    end
    println("--------------------------")
    println("Checking dS2dp using S1S2 deriv")
    Nf = length(fhash["filt_index"])
    println("Shape check",size(dS2dp_23),size(dS2lim23))
    derdiff = reshape(dS2dp_23, Nf*Nf) - dS2lim23 #Column major issue here?
    println("Difference",derdiff)
    txt = @sprintf("Range of S2 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0[lasts1ind+1:end]),maximum(der1[lasts1ind+1:end]),minimum(der0[lasts1ind+1:end]),minimum(der1[lasts1ind+1:end]))
    print(txt)
    txt = @sprintf("Range of dS2lim: Max=%.3E  Min=%.3E \n",maximum(abs.(dS2lim23)),minimum(abs.(dS2lim23)))
    print(txt)
    txt = @sprintf("Difference between dS2dp and dSlim Mean: %.3E stdev: %.3E mean(abs(derdiff/dS2lim)): %.3E \n",mean(derdiff),std(derdiff), mean(abs.(derdiff./dS2lim23)))
    print(txt)
    println("Examining only those S2 coeffs which are greater than eps=1e-3 for der0")
    eps_mask = findall(der0[lasts1ind+1:end] .> 1E-3)
    der0_s2good = der0[lasts1ind+1:end][eps_mask]
    der1_s2good = der1[lasts1ind+1:end][eps_mask]
    dSlim_s2good = dS2lim23[eps_mask]
    derdiff_good = derdiff[eps_mask]
    print(derdiff_good)
    txt = @sprintf("\nRange of S2 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0_s2good),maximum(der1_s2good),minimum(der0_s2good),minimum(der1_s2good))
    print(txt)
    txt = @sprintf("Range of dS2lim: Max=%.3E  Min=%.3E \n",maximum(abs.(dSlim_s2good)),minimum(abs.(dSlim_s2good)))
    print(txt)
    txt = @sprintf("Mean: %.3E stdev: %.3E mean(abs(derdiff/dS2lim)): %.3E \n",mean(derdiff_good),std(derdiff_good), mean(abs.(derdiff_good ./dSlim_s2good)))
    print(txt)
    println("Derivative Values")
    println("Using Newton lim")
    println(dS2lim23[:])
    println("Using dS2")
    println(dS2dp_23[:])
    println(size(dS2dp_23), Nf^2)
    return reshape(dS2dp_23, (Nf, Nf)), reshape(dS2lim23, (Nf, Nf))
end

dS2code, dS2lim = derivtestS1S2(16, mode="fast") #THIS WORKS with dhc AND NOT WITH dhc compute; FIGURE OUT WHY DERIV UTILS DIDNT

#=OLD
function derivtestS1S2(Nx)
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    dS1dp_existing = wst_S1_deriv(im, fhash)
    dS1dp, dS2dp = wst_S1S2_deriv(im, fhash)

    der0=DHC(im0,fhash,doS2=true,doS20=false)
    der1=DHC(im1,fhash,doS2=true,doS20=false)
    dSlim = (der1-der0) ./ eps
    #=
    println("Checking dS1dp using existing deriv")
    Nf = length(fhash["filt_index"])
    lasts1ind = 2+Nf
    diff = dS1dp_existing[2, 3, :]-dSlim[3:lasts1ind]
    println(dSlim[3:lasts1ind])
    println("and")
    println(dS1dp_existing[2, 3, :])
    println("stdev: ",std(diff)) =#
    println("--------------------------")
    Nf = length(fhash["filt_index"])
    lasts1ind = 2+Nf
    diff = dS1dp[2, 3, :]-dSlim[3:lasts1ind]
    println(dSlim[3:lasts1ind])
    println("and")
    println(dS1dp[2,3,:])
    txt= @sprintf("Range of S1 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E ",maximum(der0[3:lasts1ind]),maximum(der1[3:lasts1ind]),minimum(der0[3:lasts1ind]),minimum(der1[3:lasts1ind]))
    print(txt)
    txt = @sprintf("Checking dS1dp using S1S2 deriv"," mean(abs(diff/dS1lim)): %.3E", mean(abs.(diff./dSlim[3:lasts1ind])))
    print(txt)
    println("stdev: ",std(diff)," mean(abs(diff/dS1lim)): ", mean(abs.(diff./dSlim[3:lasts1ind])))
    txt= @sprintf("Range of dS1lim: Max= ",maximum(abs.(dSlim[3:lasts1ind]))," Min= ",minimum(abs.(dSlim[3:lasts1ind])))
    print(txt)
    println("--------------------------")
    println("Checking dS2dp using S1S2 deriv")
    Nf = length(fhash["filt_index"])

    println("Shape check",size(dS2dp[2,3,:,:]),size(dSlim[lasts1ind+1:end]))
    diff = dSlim[lasts1ind+1:end] - reshape(dS2dp[2, 3, :, :], Nf*Nf) #Column major issue here?
    println("Difference",diff)
    println("Range of S2 der0, der1: Max",)
    print("Range of dS2lim: Max= ",maximum(abs.(dSlim[lasts1ind+1:end]))," Min= ",minimum(abs.(dSlim[lasts1ind+1:end])))
    println("Mean: ",mean(diff)," stdev: ",std(diff)," mean(abs(diff/dS2lim)): ", mean(abs.(diff./dSlim[lasts1ind+1:end])))

    #plot(dS[3:end])
    #plot!(blarg[2,3,:])
    #plot(diff)
    return
end
=#
function derivtestS20(Nx)
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    dS20dp = wst_S20_deriv(im, fhash)

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

    Nf = length(fhash["filt_index"])
    i0 = 3+Nf
    blarg = dS20dp[2,3,:,:]
    diff = dS[i0:end]-reshape(blarg,Nf*Nf)
    println(dS[i0:end])
    println("and")
    println(blarg)
    println("stdev: ",std(diff))
    println()
    println(diff)

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

function dS2sum_test(fhash) #DerivsumS2 works here
    (Nf, )    = size(fhash["filt_index"])
    Nx        = fhash["npix"]
    im = rand(Nx,Nx)
    mywts = rand(Nf*Nf)

    wtvec  = reshape(mywts, (Nf, Nf))

    # Use new faster code
    sum1 = wst_S2_deriv_sum(im, fhash, wtvec)

    # Compare to established code
    dS1dp, dS2dp = wst_S1S2_derivfast(im, fhash)
    dS2dp = reshape(dS2dp, (Nx, Nx, Nf^2))
    sum2 = zeros(Float64, Nx, Nx)
    for i=1:Nf*Nf sum2 += (dS2dp[:,:,i].*mywts[i]) end
    println("abs mean", mean(abs.(sum2 -sum1)))
    println("Stdev: ",std(sum1-sum2))

    return
end

derivtestS1S2(16, mode="fast")

Nx=16
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
img = zeros(Float64, Nx, Nx)
img[6,6]=1.0
@benchmark ds1ds2slow = wst_S1S2_deriv(img, fhash)
@benchmark ds1s2 = wst_S1S2_derivfast(img, fhash)
@benchmark blarg = wst_S1_deriv(img, fhash)
@benchmark blarg = wst_S1_deriv_fast(img, fhash)
@benchmark
# S1 deriv mean time, NM Feb10??
# 32
# 64    3|6|23 ms
# 128   15|21|22 ms
# 256
# 512   1|2|2 s

# S1 derivFAST mean time, NM Feb10??
# 32
# 64    1.5|2|12 ms
# 128   6|10|29 ms
# 256
# 512   582|721|856

# S1 deriv time, Jan 30: Doug / Andrew??
# 32    17 ms
# 64    34 ms
# 128   115 ms
# 256   520 ms
# 512   2500 ms

Nx = 64
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im = zeros(Float64, Nx, Nx)
im[6,6]=1.0
@benchmark blarg = wst_S20_deriv(im, fhash)

Profile.clear()
@profile blarg = wst_S20_deriv(im, fhash)
Juno.profiler()

# S2 deriv time, Jan 30
# 8     28 ms
# 16    112
# 32    320
# 64    1000
# 128   5 sec
# 256   ---
# 512   ---

print(1)

derivtest(128)

Nx=32
im = rand(Nx,Nx)
@time fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
@benchmark Sarr = DHC(im, fhash, doS2=false)  # 0.36 ms
@benchmark Sarr = DHC(im, fhash, doS2=true)   # 6.4 ms
@benchmark Sarr = DHC(im, fhash, doS2=false, doS12=true)  # 0.7 ms
@benchmark Sarr = DHC(im, fhash, doS2=false, doS20=true)  # 2.1 ms

Nx=64
im = rand(Nx,Nx)
@time fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
@benchmark Sarr = DHC(im, fhash, doS2=false)  # 1.3 ms
@benchmark Sarr = DHC(im, fhash, doS2=true)   # 19 ms
@benchmark Sarr = DHC(im, fhash, doS2=false, doS12=true)  # 2.8 ms
@benchmark Sarr = DHC(im, fhash, doS2=false, doS20=true)  # 6.6 ms

Nx=256
im = rand(Nx,Nx)
@time fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
@benchmark Sarr = DHC(im, fhash, doS2=false)  # 25 ms
@benchmark Sarr = DHC(im, fhash, doS2=true)   # 375 ms
@benchmark Sarr = DHC(im, fhash, doS2=false, doS12=true)  # 55 ms
@benchmark Sarr = DHC(im, fhash, doS2=false, doS20=true)  # 130 ms

Nx = 64
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im = rand(Nx,Nx)
S_targ = DHC(im, fhash, doS2=false)

im = rand(Nx,Nx)

derivtestS1S2(128)

function wst_synth(im_init, fixmask)
    # fixmask -  0=float in fit   1=fixed

    function wst_synth_chisq(vec_in)

        thisim = copy(im_init)
        thisim[indfloat] = vec_in

        Sarr  = DHC(thisim, fhash, doS2=false)
        diff  = (Sarr .- S_targ)[3:end]

        # should have some kind of weight here
        chisq = diff'*diff
        return chisq

    end

    function wst_synth_dchisq(storage, vec_in)

        thisim = copy(im_init)
        thisim[indfloat] = vec_in
        dSdp   = wst_S1_deriv(thisim, fhash)
        S1arr  = DHC(thisim, fhash, doS2=false)
        diff  = (S1arr - S_targ)[3:end]

        # dSdp matrix * S1-S_targ is dchisq
        dchisq_im = (reshape(dSdp, Nx*Nx, Nf) * diff).*2
        dchisq = reshape(dchisq_im, Nx, Nx)[indfloat]

        storage .= dchisq
    end

    (Nx, Ny)  = size(im_init)
    if Nx != Ny error("Input image must be square") end
    (Nf, )    = size(fhash["filt_index"])

    # index list of pixels to float in fit
    indfloat = findall(fixmask .== 0)

    # initial values for floating pixels
    vec_init = im_init[indfloat]
    println(length(vec_init), " pixels floating")

    eps = zeros(size(vec_init))
    eps[1] = 1e-4
    chisq1 = wst_synth_chisq(vec_init+eps./2)
    chisq0 = wst_synth_chisq(vec_init-eps./2)
    brute  = (chisq1-chisq0)/1e-4

    clever = zeros(size(vec_init))
    _bar = wst_synth_dchisq(clever, vec_init)
    println("Brute:  ",brute)
    println("Clever: ",clever[1])

    # call optimizer
    res = optimize(wst_synth_chisq, wst_synth_dchisq, vec_init, BFGS())

    # copy results into pixels of output image
    im_synth = copy(im_init)
    im_synth[indfloat] = Optim.minimizer(res)
    println(res)

    return im_synth
end

Nx    = 128
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im    = rand(Nx,Nx)
fixmask = im .> 0.5
S_targ = DHC(im, fhash, doS2=false)

init = copy(im)
init[findall(fixmask .==0)] .= 0

foo = wst_synth(init, fixmask)

img = readdust()
img = (x->convert(Float64, x)).(img[1:256, 1:256])
fimg = fft(img)
rpow = sum(abs2.(img))
fpow = sum(abs2.(fimg))

#Optim.minimizer(optimize(FOM, im, LBFGS()))

# using dchisq function
# size   t(BFGS) t(LBFGS) [sec]
# 16x16    1
# 32x32    4        9
# 64x64    20      52
# 128x128         153 (fitting 50% of pix)

# without dchisq function
# size   t(BFGS) t(LBFGS) [sec]
# 16x16   27
# 32x32  355      110
# 64x64          1645
# 128x128        est 28,000  (8 hrs)

function readdust()

    RGBA_img = load(pwd()*"/scratch_DF/t115_clean_small.png")
    img = reinterpret(UInt8,rawview(channelview(RGBA_img)))

    return img[1,:,:]
end

function wst_synthS20(im_init, fixmask, S_targ, S20sig)
    # fixmask -  0=float in fit   1=fixed

    function wst_synth_chisq(vec_in)

        thisim = copy(im_init)
        thisim[indfloat] = vec_in

        S20arr  = DHC(thisim, fhash, doS2=false, doS20=true)
        i0 = 3+Nf
        diff  = ((S20arr - S_targ)./S20sig)[i0:end]

        # should have some kind of weight here
        chisq = diff'*diff
        println(chisq)
        return chisq

    end

    function wst_synth_dchisq(storage, vec_in)

        thisim = copy(im_init)
        thisim[indfloat] = vec_in
        dS20dp = wst_S20_deriv(thisim, fhash)
        S20arr = DHC(thisim, fhash, doS2=false, doS20=true)
        i0 = 3+Nf
        # put both factors of S20sig in this array to weight
        diff   = ((S20arr - S_targ)./(S20sig.^2))[i0:end]

        # dSdp matrix * S1-S_targ is dchisq
        dchisq_im = (reshape(dS20dp, Nx*Nx, Nf*Nf) * diff).*2
        dchisq = reshape(dchisq_im, Nx, Nx)[indfloat]

        storage .= dchisq
    end

    (Nx, Ny)  = size(im_init)
    if Nx != Ny error("Input image must be square") end
    (Nf, )    = size(fhash["filt_index"])

    # index list of pixels to float in fit
    indfloat = findall(fixmask .== 0)

    # initial values for floating pixels
    vec_init = im_init[indfloat]
    println(length(vec_init), " pixels floating")

    eps = zeros(size(vec_init))
    eps[1] = 1e-4
    chisq1 = wst_synth_chisq(vec_init+eps./2)
    chisq0 = wst_synth_chisq(vec_init-eps./2)
    brute  = (chisq1-chisq0)/1e-4

    clever = zeros(size(vec_init))
    _bar = wst_synth_dchisq(clever, vec_init)
    println("Brute:  ",brute)
    println("Clever: ",clever[1])

    # call optimizer
    res = optimize(wst_synth_chisq, wst_synth_dchisq, vec_init, BFGS())

    # copy results into pixels of output image
    im_synth = copy(im_init)
    im_synth[indfloat] = Optim.minimizer(res)
    println(res)

    return im_synth
end

function S20_weights(im, fhash, Nsam=10)

    (Nx, Ny)  = size(im)
    if Nx != Ny error("Input image must be square") end
    (Nf, )    = size(fhash["filt_index"])

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

    S20   = DHC(im, fhash, doS2=false, doS20=true)
    Ns    = length(S20)
    S20arr = zeros(Float64, Ns, Nsam)
    for j=1:Nsam
        noise = rand(Nx,Nx)
        S20arr[:,j] = DHC(im+noise, fhash, doS2=false, doS20=true)
    end

    wt = zeros(Float64, Ns)
    for i=1:Ns
        wt[i] = std(S20arr[i,:])
    end

    return wt
end

# read dust map
dust = Float64.(readdust())
dust = dust[1:16,1:16]
println("Real=",sum(abs2.(dust)))
println("Real=",sum(abs2.(fft(dust))))

Nx    = 64
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
im    = imresize(dust,(Nx,Nx))
fixmask = rand(Nx,Nx) .< 0.1

S_targ = DHC(im, fhash, doS2=false, doS20=true)
S_targ[end] = 0 #regularizing start
init = copy(im)
floatind = findall(fixmask .==0)
init[floatind] .+= rand(length(floatind)).*5 .-25

S20sig = S20_weights(im, fhash, 100)
foo = wst_synthS20(init, fixmask, S_targ, S20sig)
S_foo = DHC(foo, fhash, doS2=false, doS20=true)

# using dchisq function with S20
# size   t(BFGS) t(LBFGS) [sec]
# 8x8      33
# 16x16    -
# 32x32

#DOING dS1S2 and derivtest code in MAIN TO EXAMINE VARS in workspace
#=
dust = Float64.(readdust())
Nx=16
dust = dust[1:Nx,1:Nx]
eps = 1e-4
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im = rand(Float64, Nx, Nx).*0.1 #dust[1:Nx, 1:Nx]
im[6,6]=1.0
im1 = copy(im)
im0 = copy(im)
im1[2,3] += eps/2
im0[2,3] -= eps/2
dS1dp_existing = wst_S1_deriv(im, fhash)
dS1dp, dS2dp = wst_S1S2_deriv(im, fhash)

der0=DHC(im0,fhash,doS2=true,doS20=false)
der1=DHC(im1,fhash,doS2=true,doS20=false)
dSlim = (der1-der0) ./ eps

println("Checking dS1dp using existing deriv")
Nf = length(fhash["filt_index"])
lasts1ind = 2+Nf
derdiff = dS1dp_existing[2, 3, :]-dSlim[3:lasts1ind]
println(dSlim[3:lasts1ind])
println("and")
println(dS1dp_existing[2, 3, :])
println("stdev: ",std(derdiff))
println("--------------------------")
Nf = length(fhash["filt_index"])
lasts1ind = 2+Nf
derdiff = dS1dp[2, 3, :]-dSlim[3:lasts1ind]
println(dSlim[3:lasts1ind])
println("and")
println(dS1dp[2,3,:])
txt= @sprintf("Range of S1 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0[3:lasts1ind]),maximum(der1[3:lasts1ind]),minimum(der0[3:lasts1ind]),minimum(der1[3:lasts1ind]))
print(txt)
txt = @sprintf("Checking dS1dp using S1S2 deriv, mean(abs(derdiff/dS1lim)): %.3E \n", mean(abs.(derdiff./dSlim[3:lasts1ind])))
print(txt)
println("stdev: ",std(derdiff))
txt= @sprintf("Range of dS1lim: Max= %.3E, Min= %.3E \n",maximum(abs.(dSlim[3:lasts1ind])),minimum(abs.(dSlim[3:lasts1ind])))
print(txt)
txt = @sprintf("Mean: %.3E stdev: %.3E mean(abs(derdiff/dS1lim)): %.3E ",mean(derdiff),std(derdiff), mean(abs.(derdiff./dSlim[3:lasts1ind])))
print(txt)
println("--------------------------")
println("Checking dS2dp using S1S2 deriv")
Nf = length(fhash["filt_index"])

println("Shape check",size(dS2dp[2,3,:,:]),size(dSlim[lasts1ind+1:end]))
derdiff = dSlim[lasts1ind+1:end] - reshape(dS2dp[2, 3, :, :], Nf*Nf) #Column major issue here?
println("derdifference",derdiff)
txt = @sprintf("Range of S2 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0[lasts1ind+1:end]),maximum(der1[lasts1ind+1:end]),minimum(der0[lasts1ind+1:end]),minimum(der1[lasts1ind+1:end]))
print(txt)
txt = @sprintf("Range of dS2lim: Max=%.3E  Min=%.3E \n",maximum(abs.(dSlim[lasts1ind+1:end])),minimum(abs.(dSlim[lasts1ind+1:end])))
print(txt)
txt = @sprintf("Mean: %.3E stdev: %.3E mean(abs(derdiff/dS2lim)): %.3E \n",mean(derdiff),std(derdiff), mean(abs.(derdiff./dSlim[lasts1ind+1:end])))
print(txt)
println("Examining only those S2 coeffs which are greater than eps=1e-5 for der0")
eps_mask = findall(der0[lasts1ind+1:end] .> 1E-3)
der0_s2good = der0[lasts1ind+1:end][eps_mask]
der1_s2good = der1[lasts1ind+1:end][eps_mask]
dSlim_s2good = dSlim[lasts1ind+1:end][eps_mask]
derdiff_good = derdiff[eps_mask]
print(derdiff_good)
txt = @sprintf("\nRange of S2 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0_s2good),maximum(der1_s2good),minimum(der0_s2good),minimum(der1_s2good))
print(txt)
txt = @sprintf("Range of dS2lim: Max=%.3E  Min=%.3E \n",maximum(abs.(dSlim_s2good)),minimum(abs.(dSlim_s2good)))
print(txt)
txt = @sprintf("Mean: %.3E stdev: %.3E mean(abs(derdiff/dS2lim)): %.3E \n",mean(derdiff_good),std(derdiff_good), mean(abs.(derdiff_good ./dSlim_s2good)))
print(txt)
=#

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
"""
A type that stores an array of structures as a structure of arrays.
# Fields:
- `fieldarrays`: a named tuple of arrays. Also `fieldarrays(x)`
"""
struct StructArray{T, N, C<:NamedTuple} <: AbstractArray{T, N}
    fieldarrays::C

    function StructArray{T, N, C}(c) where {T, N, C<:NamedTuple}
        length(c) > 0 || error("must have at least one column")
        ax = axes(c[1])
        length(ax) == N || error("wrong number of dimensions")
        for i = 2:length(c)
            axes(c[i]) == ax || error("all field arrays must have same shape")
        end
        new{T, N, C}(c)
    end
end

StructArray{T}(c::C) where {T, C<:Tuple} = StructArray{T}(NamedTuple{fields(T)}(c))
StructArray{T}(c::C) where {T, C<:NamedTuple} = StructArray{T, length(size(c[1])), C}(c)
StructArray{T}(c::C) where {T, C<:Pair} = StructArray{T}(Tuple(c))
StructArray(c::C) where {C<:NamedTuple} = StructArray{eltypes(C)}(c)
StructArray(c::C) where {C<:Tuple} = StructArray{eltypes(C)}(c)
StructArray(c::Pair{P, Q}) where {P, Q} = StructArray{Pair{eltype(P), eltype(Q)}}(c)

StructArray{T}(; kwargs...) where {T} = StructArray{T}(values(kwargs))
StructArray(; kwargs...) = StructArray(values(kwargs))

StructArray{T}(args...) where {T} = StructArray{T}(NamedTuple{fields(T)}(args))

const StructVector{T, C<:NamedTuple} = StructArray{T, 1, C}
StructVector{T}(args...; kwargs...) where {T} = StructArray{T}(args...; kwargs...)
StructVector(args...; kwargs...) = StructArray(args...; kwargs...)

Base.IndexStyle(::Type{StructArray{T, N, C}}) where {T, N, C} = Base.IndexStyle(gettypes(C).parameters[1])

_undef_array(::Type{T}, sz; unwrap = t -> false) where {T} = unwrap(T) ? StructArray{T}(undef, sz; unwrap = unwrap) : Array{T}(undef, sz)

_similar(v::AbstractArray, ::Type{Z}; unwrap = t -> false) where {Z} =
    unwrap(Z) ? buildfromschema(typ -> _similar(v, typ; unwrap = unwrap), Z) : similar(v, Z)

function StructArray{T}(::Base.UndefInitializer, sz::Dims; unwrap = t -> false) where {T}
    buildfromschema(typ -> _undef_array(typ, sz; unwrap = unwrap), T)
end
StructArray{T}(u::Base.UndefInitializer, d::Integer...; unwrap = t -> false) where {T} = StructArray{T}(u, convert(Dims, d); unwrap = unwrap)

function similar_structarray(v::AbstractArray, ::Type{Z}; unwrap = t -> false) where {Z}
    buildfromschema(typ -> _similar(v, typ; unwrap = unwrap), Z)
end

function StructArray(v::AbstractArray{T}; unwrap = t -> false) where {T}
    s = similar_structarray(v, T; unwrap = unwrap)
    for i in eachindex(v)
        @inbounds s[i] = v[i]
    end
    s
end
StructArray(s::StructArray) = copy(s)

Base.convert(::Type{StructArray}, v::AbstractArray) = StructArray(v)

function Base.similar(::Type{StructArray{T, N, C}}, sz::Dims) where {T, N, C}
    cols = map_params(typ -> similar(typ, sz), C)
    StructArray{T}(cols)
end

Base.similar(s::S, sz::Tuple) where {S<:StructArray} = similar(S, Base.to_shape(sz))
Base.similar(s::S, sz::Base.DimOrInd...) where {S<:StructArray} = similar(S, Base.to_shape(sz))
Base.similar(s::S) where {S<:StructArray} = similar(S, Base.to_shape(axes(s)))

"""
`fieldarrays(s::StructArray)`

Return the field arrays corresponding to the various entry of the struct as a named tuple.
If the struct has no names (e.g. a tuple) automatic names are assigned (`:x1, :x2, ...`).

## Examples

```julia
julia> s = StructArray(rand(ComplexF64, 4));

julia> fieldarrays(s)
(re = [0.396526, 0.486036, 0.459595, 0.0323561], im = [0.147702, 0.81043, 0.00993469, 0.487091])
```
"""
fieldarrays(s::StructArray) = getfield(s, :fieldarrays)

Base.getproperty(s::StructArray, key::Symbol) = getfield(fieldarrays(s), key)
Base.getproperty(s::StructArray, key::Int) = getfield(fieldarrays(s), key)
Base.propertynames(s::StructArray) = fieldnames(typeof(fieldarrays(s)))

Base.size(s::StructArray) = size(fieldarrays(s)[1])
Base.axes(s::StructArray) = axes(fieldarrays(s)[1])

@generated function Base.getindex(x::StructArray{T, N, NamedTuple{names, types}}, I::Int...) where {T, N, names, types}
    args = [:(getfield(cols, $i)[I...]) for i in 1:length(names)]
    return quote
        cols = fieldarrays(x)
        @boundscheck checkbounds(x, I...)
        @inbounds $(Expr(:call, :createinstance, :T, args...))
    end
end

function Base.getindex(s::StructArray{T, N, C}, I::Union{Int, AbstractArray, Colon}...) where {T, N, C}
    StructArray{T}(map(v -> getindex(v, I...), fieldarrays(s)))
end

function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}
    StructArray{T}(map(v -> view(v, I...), fieldarrays(s)))
end

function Base.setindex!(s::StructArray, vals, I::Int...)
    @boundscheck checkbounds(s, I...)
    @inbounds foreachcolumn((col, val) -> (col[I...] = val), s, vals)
    s
end

@inline getfieldindex(v::Tuple, field::Symbol, index::Integer) = getfield(v, index)
@inline getfieldindex(v, field::Symbol, index::Integer) = getproperty(v, field)

function Base.push!(s::StructArray, vals)
    foreachcolumn(push!, s, vals)
end

function Base.append!(s::StructArray, vals)
    foreachcolumn(append!, s, vals)
end

function Base.cat(args::StructArray...; dims)
    f = key -> cat((getproperty(t, key) for t in args)...; dims=dims)
    T = mapreduce(eltype, promote_type, args)
    StructArray{T}(map(f, fields(eltype(args[1]))))
end

function Base.resize!(s::StructArray, i::Integer)
    for a in fieldarrays(s)
        resize!(a, i)
    end
    return s
end

function Base.empty!(s::StructArray)
    foreachcolumn(empty!, s)
end

for op in [:hcat, :vcat]
    @eval begin
        function Base.$op(args::StructArray...)
            f = key -> $op((getproperty(t, key) for t in args)...)
            T = mapreduce(eltype, promote_type, args)
            StructArray{T}(map(f, fields(eltype(args[1]))))
        end
    end
end

Base.copy(s::StructArray{T,N,C}) where {T,N,C} = StructArray{T,N,C}(C(copy(x) for x in fieldarrays(s)))

function Base.reshape(s::StructArray{T}, d::Dims) where {T}
    StructArray{T}(map(x -> reshape(x, d), fieldarrays(s)))
end


text:
using Test
using MPI

MPI.Init()

usize = MPI.universe_size()
@test usize === nothing || usize >= 1

@test MPI.tag_ub() >= 32767

MPI.Finalize()


text:
using Documenter
using AtomicLevels

makedocs(
    modules = [AtomicLevels],
    sitename = "AtomicLevels",
    pages = [
        "Home" => "index.md",
        "Orbitals" => "orbitals.md",
        "Configurations" => "configurations.md",
        "Term symbols" => "terms.md",
        "CSFs" => "csfs.md",
        "Other utilities" => "utilities.md",
        "Internals" => "internals.md",
    ],
    format = Documenter.HTML(
        mathengine = MathJax(Dict(
            :TeX => Dict(
                :equationNumbers => Dict(:autoNumber => "AMS"),
                :Macros => Dict(
                    :defd => "≝",
                    :ket => ["|#1\\rangle", 1],
                    :bra => ["\\langle#1|", 1],
                    :braket => ["\\langle#1|#2\\rangle", 2],
                    :ketbra => ["|#1\\rangle\\!\\langle#2|", 2],
                    :matrixel => ["\\langle#1|#2|#3\\rangle", 3],
                    :vec => ["\\mathbf{#1}", 1],
                    :mat => ["\\mathsf{#1}", 1],
                    :conj => ["#1^*", 1],
                    :im => "\\mathrm{i}",
                    :operator => ["\\mathfrak{#1}", 1],
                    :Hamiltonian => "\\operator{H}",
                    :hamiltonian => "\\operator{h}",
                    :Lagrangian => "\\operator{L}",
                    :fock => "\\operator{f}",
                    :lagrange => ["\\epsilon_{#1}", 1],
                    :vary => ["\\delta_{#1}", 1],
                    :onebody => ["(#1|#2)", 2],
                    :twobody => ["[#1|#2]", 2],
                    :twobodydx => ["[#1||#2]", 2],
                    :direct => ["{\\operator{J}_{#1}}", 1],
                    :exchange => ["{\\operator{K}_{#1}}", 1],
                    :diff => ["\\mathrm{d}#1\\,", 1]
                )
            )
        ))
    ),
    strict = true,
    doctest = false
)

deploydocs(repo = "github.com/JuliaAtoms/AtomicLevels.jl.git")


text:
using Zygote
using Tullio

export Tikhonov, GR, TV
export generate_spatial_grad_square, generate_GR, generate_TV

 # General hint
 # for the creation of the regularizers we are using meta programming
 # because the fastest way for automatic differentiation and Zygote
 # is Tullio.jl at the moment.
 # Our metaprogramming code was initially based on
 # https://github.com/mcabbott/Tullio.jl/issues/11
 #

"""
    generate_indices(num_dims, d, ind1, ind2)
Generates a list of symbols which can be used to generate Tullio expressions
via metaprogramming.
`num_dims` is the total number of dimensions.
`d` is the dimension where there is a offset in the index.
`ind1` and `ind2` are the offsets each.
 # Examples
 ```julia-repl
julia> a, b = generate_indices(5, 2, 1, 1)
(Any[:i1, :(i2 + 1), :i3, :i4, :i5], Any[:i1, :(i2 + 1), :i3, :i4, :i5])
```
"""
function generate_indices(num_dims, d, ind1, ind2)
    # create initial symbol
    ind = :i
    # create the array of symbols for each dimension
    inds1 = map(1:num_dims) do di
        # map over numbers and append the number of the position to symbol i
        i = Symbol(ind, di)
        # at the dimension where we want to do the step, add $ind1
        di == d ? :($i + $ind1) : i
    end
    inds2 = map(1:num_dims) do di
        i = Symbol(ind, di)
        # here we do the step but in the other dimension. ind2 should be
        # (-1) * ind1 or negative
        di == d ? :($i + $ind2) : i
    end
    return inds1, inds2
end

"""
    generate_laplace(num_dims, sum_dims_arr, weights)
Generate the Tullio statement for computing the abs2 of Laplacian.
`num_dims` is the dimension of the array.
`sum_dims_arr` is a array indicating over which dimensions we must sum over.
`weights` is a array of a weight for the different dimension.
"""
function generate_laplace(num_dims, sum_dims_arr, weights; debug=false)
    # create out list for the final expression
    # add accumulates the different add expressions
    out, add = [], []
    # loop over all dimensions which we want to sum. for each dimension must
    # be a weight provided
    for (d, w) in zip(sum_dims_arr, weights)
        # get the two lists of indices
        inds1, inds2 = generate_indices(num_dims, d, 1, -1)
        # critical part where we actually add the two expressions for
        # the steps in the dimension to the add array
        push!(add, :($w * arr[$(inds1...)] + $w * arr[$(inds2...)]))
    end
    # for laplace we need one final negative term at the position itself
    inds = map(1:num_dims) do di
        i = Symbol(:i, di)
    end
    # subtract this final term
    pre_factor = 2 ^ num_dims * sum(weights)
    push!(add, :(-$:($pre_factor * arr[$(inds...)])))
    # create final expressions by adding all elements of the add list
    if debug
        push!(out, :(res = abs2(+$(add...))))
    else
        push!(out, :(@tullio threads=false res = abs2(+$(add...))))
    end
    return out
end

"""
    create_Ndim_regularizer(expr, num_dims, sum_dims_arr, weights, ind1, ind2)
    A helper function to create a N-dimensional regularizer. In principle
    the same as `generate_laplace` but more general
    `expr` needs to be a function which takes `inds1`, `inds2` and a weight `w`-
    `num_dims` is the total amount of dimensions
    `sum_dims_arr` is a array indicating over which dimensions we must sum over.
    `weights` is a array of a weight for the different dimension.
    ``
"""
function create_Ndim_regularizer(expr, num_dims, sum_dims_arr, weights,
                                 ind1, ind2)
    out, add = [], []
    for (d, w) in zip(sum_dims_arr, weights)
        inds1, inds2 = generate_indices(num_dims, d, ind1, ind2)
        push!(add, expr(inds1, inds2, w))
    end
    push!(out, :(@tullio threads=false res = +($(add...))))
    return out
end

"""
    generate_spatial_grad_square(num_dims, sum_dims_arr, weights)
Generate the Tullio statement for calculating the squared spatial gradient
over n dimensions.
`num_dims` is the dimension of the array.
`sum_dims_arr` is a array indicating over which dimensions we must sum over.
`weights` is a array of a weight for the different dimension.
`ind1` and `ind2` are the offsets for the difference.
"""
function generate_spatial_grad_square(num_dims, sum_dims_arr, weights)
    expr(inds1, inds2, w) = :($w * abs2(arr[$(inds1...)] - arr[$(inds2...)]))
    @eval x = arr -> ($(create_Ndim_regularizer(expr, num_dims, sum_dims_arr,
                        weights, 1, -1)...))
    return x
end

"""
    Tikhonov(; <keyword arguments>)
This function returns a function to calculate the Tikhonov regularizer
of a n-dimensional array.
# Arguments
- `num_dims=2`:
- `sum_dims=[1, 2]`: A array containing the dimensions we want to sum over
- `weights=nothing`: A array containing weights to weight the contribution of
    different dimensions. If `weights=nothing` all dimensions are weighted equally.
- `step=1`: A integer indicating the step width for the array indexing
- `mode="laplace"`: Either `"laplace"`, `"spatial_grad_square"`, `"identity"` accounting for different
    modes of the Tikhonov regularizer. Default is `"laplace"`.
# Examples
To create a regularizer for a 3D dataset where the third dimension
has different contribution.
```julia-repl
julia> reg = Tikhonov(num_dims=2, sum_dims=[1, 2], weights=[1, 1], mode="identity");

julia> reg([1 2 3; 4 5 6; 7 8 9])
285
```
"""
function Tikhonov(;num_dims=2, sum_dims=[1, 2], weights=[1, 1], step=1, mode="laplace")
    if weights == nothing
        weights = ones(num_dims)
    end
    if mode == "laplace"
        Γ = @eval arr -> ($(generate_laplace(num_dims, sum_dims, weights)...))
    elseif mode == "spatial_grad_square"
        expr(inds1, inds2, w) = :($w * abs2(arr[$(inds1...)] - arr[$(inds2...)]))
        Γ = @eval arr -> ($(create_Ndim_regularizer(expr, num_dims, sum_dims,
                            weights, step, (-1) * step)...))
    elseif mode == "identity"
        Γ = arr -> sum(abs2.(arr))
    else
        throw(ArgumentError("The provided mode is not valid."))
    end

    return Γ
end

"""
    generate_GR(num_dims, sum_dims_arr, weights, ind1, ind2, ϵ)
Generate the Tullio statement for computing the Good's roughness.
`num_dims` is the dimension of the array. `sum_dims_arr` is a array
indicating over which dimensions we must sum over.
`weights` is a array of a weight for the different dimension.
`ind1` and `ind2` are the offsets for the difference.
"""
function generate_GR(num_dims, sum_dims_arr, weights, ind1, ind2; debug=false)
    out, add = [], []
    inds = map(1:num_dims) do di
        i = Symbol(:i, di)
    end

    for (d, w) in zip(sum_dims_arr, weights)
        inds1, inds2 = generate_indices(num_dims, d, ind1, ind2)
        push!(add, :($w * (arr[$(inds1...)] + arr[$(inds2...)])))
    end
    prefactor = - 4 / (abs(ind1) + abs(ind2))
    diff_factor = -sum(weights) * 2
    push!(add, :($diff_factor *arr[$(inds...)]))
    if debug
        push!(out, :(res = $prefactor * arr[$(inds...)] * +($(add...))))
    else
        push!(out, :(@tullio threads=false res = $prefactor * arr[$(inds...)] * +($(add...))))
    end
    return out
end

"""
    GR(; <keyword arguments>)
This function returns a function to calculate the Good's roughness regularizer
of a n-dimensional array.
# Arguments
- `num_dims=2`: Dimension of the array that should be regularized
- `sum_dims=[1, 2]`: A array containing the dimensions we want to sum over
- `weights=nothing`: A array containing weights to weight the contribution of
    different dimensions. If `weights=nothing` all dimensions are weighted equally.
- `step=1`: A integer indicating the step width for the array indexing
- `mode="forward"`: Either `"central"` or `"forward"` accounting for different
    modes of the spatial gradient. Default is "central".
# Examples
To create a regularizer for a 3D dataset where the third dimension
has different contribution. For the derivative we use forward mode.
```julia-repl
julia> reg = GR(num_dims=2, sum_dims=[1, 2], weights=[1, 1], mode="forward");

julia> reg([1 2 3; 4 5 6; 7 8 9])
-26.36561871738898
```
"""
function GR(; num_dims=2, sum_dims=[1, 2], weights=[1, 1], step=1,
              mode="central", ϵ=1f-8)
    if weights == nothing
        weights = ones(num_dims)
    end
    if mode == "central"
        GRf = @eval arr -> ($(generate_GR(num_dims, sum_dims, weights,
                                        step, (-1) * step)...))
    elseif mode == "forward"
        GRf = @eval arr -> ($(generate_GR(num_dims, sum_dims, weights,
                                        step, 0)...))
    else
        throw(ArgumentError("The provided mode is not valid."))
    end

    # we need to add a ϵ to prevent NaN in the derivative of it
    return arr -> GRf(sqrt.(arr .+ ϵ))
end

"""
    generate_TV(num_dims, sum_dims_arr, weights, ind1, ind2, ϵ)
Generate the Tullio statement for computing the Good's roughness.
`num_dims` is the dimension of the array.
`sum_dims_arr` is a array
indicating over which dimensions we must sum over.
`weights` is a array of a weight for the different dimension.
`ind1` and `ind2` are the offsets for the difference.
`ϵ` is a numerical constant to prevent division by zero.
    this is important for the gradient
"""
function generate_TV(num_dims, sum_dims_arr, weights, ind1, ind2, ϵ=1f-8; debug=false)
    out, add = [], []
    for (d, w) in zip(sum_dims_arr, weights)
        inds1, inds2 = generate_indices(num_dims, d, ind1, ind2)
        push!(add, :($w * abs2(arr[$(inds1...)] - arr[$(inds2...)])))
    end
    push!(add, ϵ)
    if debug
        push!(out, :(res = sqrt(+($(add...)))))
    else
        push!(out, :(@tullio threads=false res = sqrt(+($(add...)))))
    end
    return out
end

"""
    TV(; <keyword arguments>)
This function returns a function to calculate the Total Variation regularizer
of a n-dimensional array.
# Arguments
- `num_dims=2`:
- `sum_dims=[1, 2]`: A array containing the dimensions we want to sum over
- `weights=nothing`: A array containing weights to weight the contribution of
    different dimensions. If `weights=nothing` all dimensions are weighted equally.
- `step=1`: A integer indicating the step width for the array indexing
- `mode="central"`: Either `"central"` or `"forward"` accounting for different
    modes of the spatial gradient. Default is "central".
# Examples
To create a regularizer for a 3D dataset where the third dimension
has different contribution. For the derivative we use forward mode.
```julia-repl
julia> reg = TV(num_dims=2, sum_dims=[1, 2], weights=[1, 1], mode="forward");

julia> reg([1 2 3; 4 5 6; 7 8 9])
12.649111f0
```
"""
function TV(; num_dims=2, sum_dims=[1, 2], weights=nothing, step=1, mode="central")

    if weights == nothing
        weights = ones(num_dims)
    end

    if mode == "central"
        total_var = @eval arr -> ($(generate_TV(num_dims, sum_dims, weights,
                                        step, (-1) * step)...))
    elseif mode == "forward"
        total_var = @eval arr -> ($(generate_TV(num_dims, sum_dims, weights,
                                        step, 0)...))
    else
        throw(ArgumentError("The provided mode is not valid."))
    end
    return total_var
end


text:
using Pkg
#include("/media/mat/HDD/EnKF/src/EnKF.jl")
Pkg.activate("/media/mat/HDD/EnKF/")
using Documenter, EnKF, ViscousFlow
makedocs(
    sitename = "EnKF.jl",
    pages = [
        "Home" => "index.md",
        "Manual" => [#"manual/inflation.md",
                     "manual/state.md",
                     #"manual/system.md",
                     #"manual/timemarching.md"
                     ],
#        "Internals" => [ "internals/properties.md"]
    ],
   format = Documenter.HTML(assets = ["assets/custom.css"]),
#   strict = true
)

deploydocs(
    repo = "github.com/mleprovost/EnKF.jl.git",
)

max_stars_repo_name: mleprovost/EnKF

max_issues_repo_name: mleprovost/EnKF

max_forks_repo_name: mleprovost/EnKF


text:
module CUFFT

using CUDAapi

using ..CuArrays
import ..CuArrays: unsafe_free!

using CUDAdrv
using CUDAdrv: CUstream

using CEnum

const libcufft = Ref("libcufft")

# core library
include("libcufft_common.jl")
include("error.jl")
include("libcufft.jl")

# low-level wrappers
include("util.jl")
include("wrappers.jl")

# high-level integrations
include("fft.jl")

end


text:
push!(LOAD_PATH, "/Users/megorov/Desktop/projects/stanford/sniper/egorov/sniperRepo/new-version")

using UrbanMap_
using SniperMDP_
using ValueIterations
using Simulator
using SimServers

# create a map and the MDP
map = UrbanMap(1, [0.5], [0.5], [0.42], [0.22], 10, 10)
mdp = SniperMDP(map, nSnipers = 1, nMonitors = 2, targets=[(2.0,8.0), (2.0,1.0)])

# load the policy
policy = load_policy("policies/mmdp-10-iter-1.policy")

# set up the server
socket = 4444
server = SniperServer(socket)

# create the simulation type
init = 470872
sim  = MDPSimulation(mdp, init, 100)

# begin simulation and send results to client
#simAndSend(server, sim, mdp, policy)

# begin batch sim
n = 500
rewards = zeros(n)
shot    = zeros(Int64, n)
obs     = zeros(Int64, n)
for i = 1:n
    res = simulate(policy, mdp, 100, initState=init)
    rewards[i] = sum(res.rewards)
    shot[i]    = res.nKills
    obs[i]     = res.nObs
end