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


function solve_modelST_TR(H, g, nlp_stop, PData::PDataST, δ::T; cgtol::T = 0.1) where {T}
    # cas particulier Steihaug-Toint
    # ϵ = sqrt(eps(T)) # * 100.0 # old
    # cgtol = max(ϵ, min(cgtol, 9 * cgtol / 10, 0.01 * norm(g)^(1.0 + PData.ζ))) # old

    ζ, ξ, maxtol, mintol = PData.ζ, PData.ξ, PData.maxtol, PData.mintol
    n = length(g)
    gNorm2 = norm(g)
    # precision = max(1e-12, min(0.5, (gNorm2^ζ)))
    # Tolerance used in Assumption 2.6b in the paper ( ξ > 0, 0 < ζ ≤ 1 )
    cgatol = min(maxtol, ξ * gNorm2^(1 + ζ))
    cgatol = max(mintol, cgatol) # add some feasible limit
    cgrtol = min(maxtol, ξ * gNorm2^ζ)
    cgrtol = max(mintol, cgrtol) # add some feasible limit

    solver = PData.solver
    cg!(
        solver,
        H,
        -g,
        atol = cgatol,
        rtol = cgrtol,
        radius = δ,
        itmax = max(2 * n, 50),
        verbose = 0,
    )

    PData.d .= solver.x
    PData.OK = solver.stats.solved

    return PData.d, PData.λ
end


"""
timedev(data, rate; [frequency=false], [overlapping=true], [taus=Octave])
Calculates the time deviation

#parameters:
* `<data>`:			The data array to calculate the deviation from either as as phases or frequencies.
* `<rate>`:			The rate of the data given.
* `[frequency]`:		True if `data` contains frequency data otherwise (default) phase data is assumed.
* `[overlapping]`:	True (default) to calculate overlapping deviation, false otherwise.
* `[taus]`:			Taus to calculate the deviation at. This can either be an AllanTauDescriptor type `AllTaus, Decadade, HalfDecade, Octave, HalfOctave, QuarterOctave`, an array of taus to calculate at, a float number to build a custom log-scale on or an integer to build a specific number of log spaced points.

#returns: named tupple (tau, deviation, error, count)
* `tau`:		Taus which where used.
* `deviation`:	Deviations calculated.
* `error`:		Respective errors.
* `count`:		Number of contributing terms for each deviation.
"""
function timedev(
		data::Array{T, 1},
		rate::AbstractFloat;
		frequency::Bool = false,
		overlapping::Bool = true,
		taus::Union{Type{U}, Integer, AbstractFloat, Array{Float64}} = 192) where {T, U <: AllanTauDescriptor}

	n = length(data)
	if n < 4
		error("Length for `data` in timedev must be at least 4 or greater")
		#we check this here, so that we can output the right function name in case of the error
	end

	(mdtaus, mddeviation, mderror, mdcount) = mallandev(data, rate, frequency = frequency, overlapping = overlapping, taus = taus)
	mdm = mdtaus ./ sqrt(3)

	(tau = mdtaus, deviation = mdm .* mddeviation, error = mdm .* mderror, count = mdcount)
end


using .Hidato

hidat = """
__ 33 35 __ __  .  .  .
__ __ 24 22 __  .  .  .
__ __ __ 21 __ __  .  .
__ 26 __ 13 40 11  .  .
27 __ __ __  9 __  1  .
 .  . __ __ 18 __ __  .
 .  .  .  . __  7 __ __
 .  .  .  .  .  .  5 __"""

const kingmoves = [[-1, -1], [-1, 0], [-1, 1], [0, -1], [0, 1], [1, -1], [1, 0], [1, 1]]

board, maxmoves, fixed, starts = hidatoconfigure(hidat)
printboard(board)
hidatosolve(board, maxmoves, kingmoves, fixed, starts[1][1], starts[1][2], 1)
printboard(board)


""" The Cameras module is responsible for ray generation. It defines
types for each kind of camera (in the base assignment, this only
includes CanonicalCamera) and implements a pixel_to_ray function
that takes a camera and a set of pixel indices and returns a ray
along which to search for objects that appear at that pixel.
"""
module Cameras

import LinearAlgebra.normalize
import LinearAlgebra.norm
import LinearAlgebra.cross

#push!(LOAD_PATH, pwd())
using ..GfxBase

export CanonicalCamera
export pixel_to_ray

""" The most basic perspective camera. The eye is at (0, 0, 0),
the up vector is (0, 1, 0), and the view direction is (0, 0, -1).
The viewport has width 1 and whatever height makes pixels square.
The focal length (distance from eye to viewport) is 1. """
mutable struct CanonicalCamera
    canv_height::Int
    canv_width::Int
end

""" Given a camera and pixel coordinates 1-indexed, i=row,j=column],
return the ray along which to search for objects that project to
that pixel. """
function pixel_to_ray(camera::CanonicalCamera, i, j)
    # viewport height = 1
    vp_height = camera.canv_height / camera.canv_width

    # convert from i,j 2D array indices to (u, v) coordinates
    # with (0,0) in the center of the image
    # half pixel shift, scale down to size 1, shift origin to center
    u =  ((j - 0.5)    / camera.canv_width       - 0.5)

    # same as u, except i increases downward so flip with a negative sign
    # and multiply by vp_height to account for aspect ratio
    v = -((i - 0.5)    / camera.canv_height      - 0.5) * vp_height

    # focal length is 1, pointing down the -z axis
    w = -1.0

    Ray(Vec3(0, 0, 0), normalize(Vec3(u, v, w)))
end

""" A general perspective camera."""
mutable struct PerspectiveCamera

    # orthonormal basis specifies position and orientation:
    eye::Vec3  # position of the eye
    u_axis::Vec3 # points right in image space
    v_axis::Vec3 # points up in image space
    w_axis::Vec3 # opposite the view direction

    # the viewport is width 1, parallel to the uv plane and
    # centered at (0, 0, -focal):
    focal::Real # focal length (aka d), the distance to the image plane

    # image dimensions:
    canv_height::Int # image height in pixels
    canv_width::Int # image width in pixels
end

function PerspectiveCamera(eye::Vec3, view::Vec3, up::Vec3, focal::Real, canv_height::Int, canv_width::Int)
    w = normalize(-view)
    u = normalize(cross(up, w))
    v = normalize(cross(w, u))
    PerspectiveCamera(eye, u, v, w, focal, canv_height, canv_width)
end

function pixel_to_ray(camera::PerspectiveCamera, i, j)
    vp_height = camera.canv_height / camera.canv_width

    u =  ((j - 0.5) / camera.canv_width  - 0.5)
    v = -((i - 0.5) / camera.canv_height - 0.5) * vp_height
    w = -camera.focal

    Ray(camera.eye, normalize(u * camera.u_axis + v * camera.v_axis + w * camera.w_axis))
end

end # module Cameras


export  computeSP,
        computeSP2

function computeSP(a_::Vector{Int64},
                   α::Vector{Vector{Float64}},β::Vector{Vector{Float64}},
                   nodes::Vector{Vector{Float64}},weights::Vector{Vector{Float64}},
                   ind::Matrix{Int64};
                   issymmetric::BitArray=falses(length(α)),
                   zerotol::Float64=1e-10)
                  # display("war hier")
    @assert minimum(a_)>=0 "no negative degrees allowed"
    l, p = size(ind); # p-variate basis
    p == 1 && computeSP(a_,α[1],β[1],nodes[1],weights[1];issymmetric=issymmetric[1])
    l -= 1
    # m = length(a_) # m-dimensional scalar product
    @assert maximum(a_)<=l "not enough elements in multi-index (requested: $(maximum(a)), max: $l)"
    @assert length(α)==length(β)==length(nodes)==length(weights)==length(issymmetric)==p "inconsistent number of rec. coefficients and/or nodes/weights"

    a = Vector{Int64}()
    for aa in a_
        !iszero(aa) && push!(a,aa)
    end

    if  iszero(length(a))
        return prod(β[i][1] for i=1:p)
    elseif isone(length(a))
        return 0.
    else
        inds_uni = multi2uni(a,ind)
        val = 1.
        for i=1:p
            # compute univariate scalar product
            @inbounds v = computeSP(inds_uni[i,:],α[i],β[i],nodes[i],weights[i];issymmetric=issymmetric[i])
            isapprox(v,0,atol=zerotol) ? (return 0.) : (val*=v)
        end
    end
    return val
end

function computeSP(a::Vector{Int64},op::Vector{OrthoPoly},q::Vector{Quad},ind::Matrix{Int64})
    @assert length(op)==length(q) "Inconsistent lengh of polynomials and quadrature rules"
    @assert issymmetric.(op)==issymmetric.(q)
    α,β = coeffs(op)
    n,w = nw(q)
    computeSP(a,α,β,n,w,ind;issymmetric=issymmetric.(op))
end

function computeSP(a::Vector{Int64},mOP::MultiOrthoPoly)
    @assert typeof(mOP.uni)==Vector{OrthoPolyQ} "no quadrature rules provided"
    ops = [ mOP.uni[i].op for i=1:length(mOP.uni) ]
    quads = [ mOP.uni[i].quad for i=1:length(mOP.uni) ]
    computeSP(a,ops,quads,mOP.ind)
end

"""
__Univariate__
```
computeSP(a::Vector{Int64},α::Vector{Float64},β::Vector{Float64},nodes::Vector{Float64},weights::Vector{Float64};issymmetric::Bool=false)
computeSP(a::Vector{Int64},op::OrthoPoly,q::Quad;issymmetric=issymmetric(op))
computeSP(a::Vector{Int64},opq::OrthoPolyQ)
```
__Multivariate__
```
computeSP( a::Vector{Int64},
           α::Vector{Vector{Float64}},β::Vector{Vector{Float64}},
           nodes::Vector{Vector{Float64}},weights::Vector{Vector{Float64}},
           ind::Matrix{Int64};
           issymmetric::BitArray=falses(length(α)))
computeSP(a::Vector{Int64},op::Vector{OrthoPoly},q::Vector{Quad},ind::Matrix{Int64})
computeSP(a::Vector{Int64},mOP::MultiOrthoPoly)
```

Computes the scalar product
```math
\\langle \\phi_{a_1},\\phi_{a_2},\\cdots,\\phi_{a_n} \\rangle,
```
where `n = length(a)`.
This requires to provide the recurrence coefficients `(α,β)` and the quadrature rule
`(nodes,weights)`, as well as the multi-index `ind`.
If provided via the keyword `issymmetric`, symmetry of the weight function is exploited.
All computations of the multivariate scalar products resort back to efficient computations
of the univariate scalar products.
Mathematically, this follows from Fubini's theorem.

The function is multiply dispatched to facilitate its use with `OrthoPolyQ` or a
suitable combination of `OrthoPoly` and its quadrature rule `Quad`.

!!! note
    - Zero entries of ``a`` are removed automatically to simplify computations, which follows from
    ```math
    \\langle \\phi_i, \\phi_j, \\phi_0,\\cdots,\\phi_0 \\rangle = \\langle \\phi_i, \\phi_j \\rangle,
    ```
    because `\\phi_0 = 1`.
    - It is checked whether enough quadrature points are supplied to solve the integral exactly.
"""
function computeSP(a_::Vector{Int64},α::Vector{Float64},β::Vector{Float64},nodes::Vector{Float64},weights::Vector{Float64};issymmetric::Bool=false)
    @assert minimum(a_) >= 0 "no negative degrees allowed"
    # remove zero entries
    # this works because ``<Φ_i,Φ_j,Φ_0,...,Φ_0> = <Φ_i,Φ_j>``
    # hence, avoiding unnecessary operations that introduced numerical gibberish
    a = Vector{Int64}()
    for aa in a_
        !iszero(aa) && push!(a,aa)
    end
    # a = filter(!iszero,a_)
    # simplification in case the measure is symmetric w.r.t t=0
    # exploit symmetry of the density, see Theorem 1.17 in W. Gautschi's book
    # and exploit the fact that
    (issymmetric && isodd(sum(a))) && return 0.
    # Gauss quadrature rules have exactness 2N-1
    @assert Int(ceil(0.5*(sum(a)+1)))<=length(nodes) "not enough nodes to integrate exactly ($(length(nodes)) provided, where $(Int(ceil(0.5*(sum(a)+1)))) are needed)."
    if iszero(length(a))
        return first(β)
    elseif isone(length(a))
        return 0.
    elseif length(a) == 2
        @inbounds a[1] == a[2] ? (return computeSP2(first(a),β)[end]) : (return 0.)
    else
        # display("war hier")
        f = ones(Float64,length(nodes))
        @simd for i = 1:length(a)
            @inbounds f = f.*evaluate(a[i],nodes,α,β)
        end
        # f = reduce(*,hcat(map(i->evaluate(a[i],nodes,α,β),1:length(a))...),dims=2)
        # display(f-f_)
        return dot(weights,f)
    end
end
function computeSP(a::Vector{Int64},op::OrthoPoly,q::Quad)
    @assert issymmetric(op)==issymmetric(q) "inconsistent symmetries"
    computeSP(a,op.α,op.β,q.nodes,q.weights;issymmetric=issymmetric(op))
end
computeSP(a::Vector{Int64},opq::OrthoPolyQ) = computeSP(a,opq.op,opq.quad)

"""
    computeSP2(n::Int64,β::Vector{Float64})
    computeSP2(n::Int64,op::OrthoPoly) = computeSP2(n,op.β)
    computeSP2(op::OrthoPoly) = computeSP2(op.deg,op.β)
    computeSP2(opq::OrthoPolyQ) = computeSP2(opq.op)

Computes the `n` *regular* scalar products aka 2-norms of the orthogonal polynomials, namely
```math
\\|ϕ_i\\|^2 = \\langle \\phi_i,\\phi_i\\rangle \\quad \\forall i \\in \\{ 0,\\dots,n \\}.
```
Notice that only the values of `β` of the recurrence coefficients `(α,β)` are required.
The computation is based on equation (1.3.7) from Gautschi, W. "Orthogonal Polynomials: Computation and Approximation".
Whenever there exists an analytic expressions for `β`, this function should be used.

The function is multiply dispatched to facilitate its use with the composite types
`OrthoPoly` and `OrthoPolyQ`.
"""
function computeSP2(n::Int64,β::Vector{Float64})
    @assert n>=0 "can only compute scalar products for non-negative degrees"
    @assert length(β) >= n + 1
    n == 0 && return β[1]
    s = ones(Float64,n+1)
    @inbounds s[1] = β[1]  # order 0computeSP2
    for i=2:n+1
        @inbounds s[i] = s[i-1]*β[i]
    end
    return s
end
computeSP2(n::Int64,op::OrthoPoly) = computeSP2(n,op.β)
computeSP2(op::OrthoPoly) = computeSP2(op.deg,op.β)
computeSP2(opq::OrthoPolyQ) = computeSP2(opq.op)

# function computeSP2(n::Int64,α::Vector{Float64},β::Vector{Float64},nodes::Vector{Float64},weights::Vector{Float64})
#     @assert n>=0 "degree has to be non-negative"
#     @assert n<=length(α) "not enough recurrence coefficients"
#     Nquad, Nrec = length(nodes), length(α)
#     n==0 ? (return β[1]) : (dot(weights,evaluate(n,nodes,α,β).^2))
# end
# computeSP2(n::Int64,opq::OrthoPolyQ) = computeSP2(n,opq.op.α,opq.op.β,opq.quad.nodes,opq.quad.weights)


"""
function beam_test()
Test the time-dependent and steady state beam problem.

"""
function beam_test()
    #
    dt = 0.02
    n = 20
    stepnum = 5
    (t, se, xe, fhist, fhistt) = ivpBeam(n, dt, stepnum)
    beamtdout = (length(fhist) == 6) && (norm(fhistt, Inf) < 5.e-5)
    beamtdout || println("error in beam_test beamtdout")
    (pout, nout) = ptcBeam(10, 100)
    pout2 = ptcBeam(10,100; jknowsdt=true)
    kdtdiff=norm(pout.solution-pout2.solution,Inf) +
     norm(pout2.history-pout.history,Inf)
    kdtok=(kdtdiff < 1.e-14)
    kdtok || println("error is knowsdt test")
    nsolp = norm(pout.solution)
    nsoln = norm(nout.solution)
    itp = length(pout.history)
    pnormok = (nsolp > 5.0) && (nsoln < 1.e-15)
    pnormok || println("error in beam_test pnromok")
    presok = (itp < 100) && (pout.history[itp] < 1.e-10)
    presok || println("error in beam_test presok")
    return beamtdout && pnormok && presok && kdtok
end


#=

  Marathon puzzle in Julia ConstraintSolver.jl

  From Xpress example
  http://www.dashoptimization.com/home/cgi-bin/example.pl?id=mosel_puzzle_5_3
  """
  Dominique, Ignace, Naren, Olivier, Philippe, and Pascal
  have arrived as the first six at the Paris marathon.
  Reconstruct their arrival order from the following
  information:
  a) Olivier has not arrived last
  b) Dominique, Pascal and Ignace have arrived before Naren
     and Olivier
  c) Dominique who was third last year has improved this year.
  d) Philippe is among the first four.
  e) Ignace has arrived neither in second nor third position.
  f) Pascal has beaten Naren by three positions.
  g) Neither Ignace nor Dominique are on the fourth position.

     (c) 2002 Dash Associates
    author: S. Heipcke, Mar. 2002
  """

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function marathon2(print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n = 6
    names = ["Dominique", "Ignace", "Naren", "Olivier", "Philippe", "Pascal"]

    @variable(model, 1 <= runners[1:n] <= n, Int)
    dominique, ignace, naren, olivier, philippe, pascal = runners

    @constraint(model, runners in CS.AllDifferentSet())

    # a: olivier not last
    @constraint(model, olivier != n)

    # b: dominique, pascal and ignace before naren and olivier
    @constraint(model, dominique  <= naren)
    @constraint(model, dominique  <= olivier)
    @constraint(model, pascal     <= naren)
    @constraint(model, pascal     <= olivier)
    @constraint(model, ignace     <= naren)
    @constraint(model, ignace     <= olivier)

    # c: dominique better than third
    @constraint(model, dominique  <= 3)

    # d: philippe is among the first four
    @constraint(model, philippe   <= 4)

    # e: ignace neither second nor third
    @constraint(model, ignace     != 2)
    @constraint(model, ignace     != 3)

    # f: pascal three places earlier than naren
    @constraint(model, pascal + 3 == naren)

    # g: neither ignace nor dominique on fourth position
    @constraint(model, ignace     != 4)
    @constraint(model, dominique  != 4)

    # for the presentation (inverse)
    @variable(model,1 <= runners_inv[1:n] <= n, Int)
    assignment_ctr(model,runners, runners_inv)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                runners_val = convert.(Integer,JuMP.value.(runners; result=sol))
                runners_inv_val = convert.(Integer,JuMP.value.(runners_inv; result=sol))
                println("runners:$runners_val")
                println("runners_inv:$runners_inv_val")
                for i in 1:n
                    println("Place $i: $(names[runners_inv_val[i]])")
                end
            end
        end
    else
        println("status:$status")
    end

    return status
end

@time marathon2(true,true)


@doc raw"""
    JONSWAP{K}(α,ωₚ,γ,r)
    JONSWAP{K}(x)

4 parameter JONSWAP model for vertical displacement.

# Type parameter
The type parameter `K` is a non-negative integer representing the ammount of aliasing to be done.
This is useful as some devices have high-pass filters which limit the ammount of aliasing that will be present in a recorded series.
In particular, a `JONSWAP{K}` model for a series recorded every `Δ` seconds would be appropriate for a series sampled every `Δ` seconds and filtered with a high-pass filter at `(2K-1)π/Δ`.

# Arguments
- `α`: The scale parameter.
- `ωₚ`: The peak angular frequency.
- `γ`: The peak enhancement factor.
- `r`: The tail decay index.

# Vector constructor arguments
- `x`: Vector of parameters `[α,ωₚ,γ,r]`.

# Background
The JONSWAP spectral density function is
```math
f(ω) = αω^{-r}\exp\left \{-\frac{r}{4}\left(\frac{|ω|}{ωₚ}\right)^{-4}\right \}γ^{δ(|ω|)}
```
where
```math
δ(ω) = \exp\left\{-\tfrac{1}{2 (0.07+0.02\cdot\mathbb{1}_{ωₚ>|ω|})^2}\left (\tfrac{|ω|}{ωₚ}-1\right )^2\right\}.
```

"""
struct JONSWAP{K} <: WhittleLikelihoodInference.UnknownAcvTimeSeriesModel{1,Float64}
    α::Float64
    ωₚ::Float64
    γ::Float64
    r::Float64

    r_over4::Float64
    ωₚ²::Float64
    ωₚ³::Float64
    ωₚ⁴::Float64
    ωₚ⁶::Float64
    logγ::Float64
    function JONSWAP{K}(α,ωₚ,γ,r) where {K}
        α > 0 || throw(ArgumentError("JONSWAP requires α > 0"))
        ωₚ > 0 || throw(ArgumentError("JONSWAP requires ωₚ > 0"))
        γ >= 1 || throw(ArgumentError("JONSWAP requires γ > 1"))
        r > 1 || throw(ArgumentError("JONSWAP requires r > 1"))
        new{K}(α,ωₚ,γ,r,r/4,ωₚ^2,ωₚ^3,ωₚ^4,ωₚ^6,log(γ))
    end
    function JONSWAP{K}(x::AbstractVector{Float64}) where {K}
        @boundscheck checkparameterlength(x,JONSWAP{K})
        @inbounds JONSWAP{K}(x[1], x[2], x[3], x[4])
    end
end

# functions to throw informative error if type parameter not provided
JONSWAP(x::AbstractVector{Float64}) = JONSWAP(1,1,1,1)
JONSWAP(α,ωₚ,γ,r) = error("JONSWAP process requires the ammount of aliasing specified as a type parameter. Use JONSWAP{K}() where K ∈ N.")

WhittleLikelihoodInference.npars(::Type{JONSWAP{K}}) where {K} = 4
WhittleLikelihoodInference.nalias(::JONSWAP{K}) where {K} = K

WhittleLikelihoodInference.lowerbounds(::Type{JONSWAP{K}}) where {K} = [0,0,1,1]
WhittleLikelihoodInference.upperbounds(::Type{JONSWAP{K}}) where {K} = [Inf,Inf,Inf,Inf]

@inline @fastmath function WhittleLikelihoodInference.sdf(model::JONSWAP{K}, ω::Real) where {K}
    α,ωₚ,γ,r = model.α,model.ωₚ,model.γ,model.r
    ω = abs(ω)
    if ω > 1e-10
        σ1² = 0.0049 + 0.0032 * (ω > ωₚ)
        δ = exp(-1 / (2σ1²) * (ω / ωₚ - 1)^2)
        ωₚ⁴_over_ω⁴ = model.ωₚ⁴ / (ω^4)
        return (α * ω^(-r) * exp(-(model.r_over4) * ωₚ⁴_over_ω⁴) * γ^δ) / 2
    else
        return 0.0
    end
end

@propagate_inbounds @fastmath function WhittleLikelihoodInference.grad_add_sdf!(out, model::JONSWAP{K}, ω::Real) where {K}
    @boundscheck checkbounds(out,1:4)
    @inbounds begin
        ω = abs(ω)
        if ω > 1e-10
            α,ωₚ,γ,r = model.α,model.ωₚ,model.γ,model.r
            σ1² = 0.0049 + 0.0032 * (ω > ωₚ)
            δ = exp(-1 / (2σ1²) * (ω / ωₚ - 1)^2)
            ω⁻⁴ = ω^(-4)
            ωₚ⁴_over_ω⁴ = model.ωₚ⁴ * ω⁻⁴
            sdf = (α * ω^(-r) * exp(-(model.r_over4) * ωₚ⁴_over_ω⁴) * γ^δ) / 2

            ∂S∂α = sdf / α
            ∂S∂ωₚ = sdf * (δ*model.logγ * ω / σ1² *(ω-ωₚ) / model.ωₚ³ - r*model.ωₚ³*ω⁻⁴)
            ∂S∂γ = sdf * δ / γ
            ∂S∂r = sdf * (-log(ω)-ωₚ⁴_over_ω⁴/4)

            out[1] += ∂S∂α
            out[2] += ∂S∂ωₚ
            out[3] += ∂S∂γ
            out[4] += ∂S∂r
        end # 0 otherwise
    end
    return nothing
end

@propagate_inbounds @fastmath function WhittleLikelihoodInference.hess_add_sdf!(out, model::JONSWAP{K}, ω::Real) where {K}
    @boundscheck checkbounds(out,1:10)
    @inbounds begin
        ω = abs(ω)
        if ω > 1e-10
            α,ωₚ,γ,r = model.α,model.ωₚ,model.γ,model.r
            σ1² = 0.0049 + 0.0032 * (ω > ωₚ)
            δ = exp(-1 / (2σ1²) * (ω / ωₚ - 1)^2)
            ω⁻⁴ = ω^(-4)
            ωₚ⁴_over_ω⁴ = model.ωₚ⁴ * ω⁻⁴
            ωₚ³_over_ω⁴ = model.ωₚ³ * ω⁻⁴
            sdf = (α * ω^(-r) * exp(-(model.r_over4) * ωₚ⁴_over_ω⁴) * γ^δ) / 2

            δωlogγ_over_σ1² = δ*model.logγ * ω / σ1²
            ∂S∂ωₚUsepart = (δ*model.logγ * ω / σ1² *(ω-ωₚ) / model.ωₚ³ - r*ωₚ³_over_ω⁴)
            ∂S∂ωₚ = sdf * ∂S∂ωₚUsepart
            δ_over_γ = δ / γ
            ∂S∂γ = sdf * δ_over_γ
            ∂r_part = (-log(ω)-ωₚ⁴_over_ω⁴/4)
            ∂S∂r = sdf * ∂r_part

            ∂S∂αωₚ = ∂S∂ωₚ / α
            ∂S∂αγ = ∂S∂γ / α
            ∂S∂αr = ∂S∂r / α

            ∂S∂ωₚ2 = ∂S∂ωₚ * ∂S∂ωₚUsepart + sdf * (δωlogγ_over_σ1² *((-3ω + 2ωₚ)/model.ωₚ⁴ + ω * (ω-ωₚ)^2/model.ωₚ⁶/σ1²) - 3r*model.ωₚ²*ω⁻⁴)
            ∂S∂ωₚγ = ∂S∂γ * ∂S∂ωₚUsepart + sdf * δ_over_γ * ω / σ1²*(ω-ωₚ)/model.ωₚ³
            ∂S∂ωₚr = ∂S∂r * ∂S∂ωₚUsepart - sdf * model.ωₚ³*ω⁻⁴
            ∂S∂γ2 = δ_over_γ * (∂S∂γ - sdf/γ)
            ∂S∂γr = ∂S∂r * δ_over_γ
            ∂S∂r2 = ∂S∂r * ∂r_part

            # out[1] += 0
            out[2] += ∂S∂αωₚ
            out[3] += ∂S∂αγ
            out[4] += ∂S∂αr
            out[5] += ∂S∂ωₚ2
            out[6] += ∂S∂ωₚγ
            out[7] += ∂S∂ωₚr
            out[8] += ∂S∂γ2
            out[9] += ∂S∂γr
            out[10] += ∂S∂r2
        end # 0 otherwise
    end

    return nothing
end


# spherical quadrature
KitBase.legendre_quadrature(6)
KitBase.octa_quadrature(8)

KitBase.VSpace1D() |> show
KitBase.VSpace2D() |> show
KitBase.VSpace3D() |> show
KitBase.MVSpace1D() |> show
KitBase.MVSpace2D() |> show
KitBase.MVSpace3D() |> show

KitBase.VSpace1D(-5, 5, 16, "newton")
KitBase.VSpace1D(-5, 5, 16, "algebra")
KitBase.VSpace2D(-5, 5, 16, -5, 5, 16, "newton")
KitBase.VSpace2D(-5, 5, 16, -5, 5, 16, "algebra")
KitBase.VSpace2D(-5, 5, 16, -5, 5, 16, "maxwell")
KitBase.VSpace3D(-5, 5, 16, -5, 5, 16, -5, 5, 16, "newton")
KitBase.VSpace3D(-5, 5, 16, -5, 5, 16, -5, 5, 16, "algebra")
KitBase.MVSpace1D(-5, 5, -5, 5, 16, "newton")
KitBase.MVSpace2D(-5, 5, -5, 5, 16, -5, 5, -5, 5, 16, "newton")
KitBase.MVSpace3D(-5, 5, -5, 5, 8, -5, 5, -5, 5, 8, -5, 5, -5, 5, 8, "newton")
KitBase.UnstructVSpace(-1, 1, 16, rand(16), ones(16))


using ProceduralNoise
using Plots

scale = (2, 2) .^ 2
resolution = (2, 2) .^ 9
A = perlin(resolution, scale)
heatmap(A)


import Distributions
using Random
using SpecialFunctions: besseli

struct CPNormal{T,S} <: Distributions.ContinuousUnivariateDistribution
    μ::T
    σ::S
end

Base.minimum(::CPNormal) = -Inf
Base.maximum(::CPNormal) = Inf

const log2π = log(2*π)
function Distributions.logpdf(dist::CPNormal, x::Real)
    μ,σ = dist.μ, dist.σ
    s,c = sincos(x-μ)
    dθ = atan(s, c)
    return  -(abs2(dθ/σ) + log2π)/2 - log(σ)
end

function Base.rand(rng::AbstractRNG, d::CPNormal)
    return d.μ + d.σ*randn(rng)
end

Base.rand(rng::AbstractRNG, ::Type{Float64}, d::CPNormal) = rand(rng, d)


import Pkg;
Pkg.activate(".")
# %%
path = pwd() * "/functions"
push!(LOAD_PATH,path)
# %%
using PyCall
sym = pyimport("sympy")
# %%
p, ξ, δ, ω, μᵖ, μᶠ, c₀ᶜ , c₀⁰ , n₀ᶠ, Nᶜ , N⁰ , Nᶠ, α₁, α₂, nᶜ, ρ₀ᶜ, ρ₁ᶜ, ρ₂ᶜ, n⁰, ρ₀⁰, ρ₁⁰, ρ₂⁰, nᶠ, ρ₁ᶠ, ρ₂ᶠ, E₀ᶜ, Eᶜ, E₀⁰, E⁰, α₀, β₀, Wᶜ, W⁰, n₀ᶜ, n₀⁰, ν₀ᶜ, ν₁ᶜ, ν₂ᶜ, ν₀⁰, ν₁⁰, ν₂⁰, W₀ᶜ, W₀⁰, W∞ᶜ, W∞⁰ = sym.symbols("p, ξ, δ, ω , μᵖ , μᶠ, c₀ᶜ , c₀⁰ , n₀ᶠ, Nᶜ , N⁰ , Nᶠ, α₁, α₂, nᶜ, ρ₀ᶜ, ρ₁ᶜ, ρ₂ᶜ, n⁰, ρ₀⁰, ρ₁⁰, ρ₂⁰, nᶠ, ρ₁ᶠ, ρ₂ᶠ, E₀ᶜ, Eᶜ, E₀⁰, E⁰, α₀, β₀, Wᶜ, W⁰, n₀ᶜ, n₀⁰, ν₀ᶜ, ν₁ᶜ, ν₂ᶜ, ν₀⁰, ν₁⁰, ν₂⁰, W₀ᶜ, W₀⁰, W∞ᶜ, W∞⁰")
# %%
# ====================================================== #
# ================= Dual to dual ======================= #
# ====================================================== #
system = [(p+ξ)*Nᶜ + (ω+ξ)*N⁰ - c₀ᶜ, μᵖ*Nᶜ + (p + μᵖ - ω)*N⁰ + μᵖ*Nᶠ - ( μᵖ/p + c₀⁰ ), μᶠ*Nᶜ + μᶠ*N⁰ + (p+δ+μᶠ)*Nᶠ - (μᶠ/p + n₀ᶠ) ]
sol = sym.linsolve(system, [Nᶜ, N⁰, Nᶠ])
# %%
sol = sym.factor(sol, p)
# %%
Nᶜ = sol.args[1][1]
N⁰ = sol.args[1][2]
Nᶠ = sol.args[1][3]
# %%
top_Nᶜ, bot_Nᶜ = sym.fraction(Nᶜ)
top_Nᶜ =  sym.Poly(top_Nᶜ,p)
for c in top_Nᶜ.coeffs()
    println(sym.factor(c, [c₀ᶜ, c₀⁰, n₀ᶠ]))
end
# %%
top_N⁰, bot_N⁰ = sym.fraction(N⁰)
top_N⁰ =  sym.Poly(top_N⁰,p)
for c in top_N⁰.coeffs()
    println(sym.factor(c, [c₀ᶜ, c₀⁰, n₀ᶠ]))
end
# %%
top_Nᶠ, bot_Nᶠ = sym.fraction(Nᶠ)
top_Nᶠ =  sym.Poly(top_Nᶠ,p)
for c in top_Nᶠ.coeffs()
    println(sym.factor(c, [c₀ᶜ, c₀⁰, n₀ᶠ]))
end
# Laplace transforms seem correct !
# %%
# ======================= Nᶜ ========================= #
bot_Nᶜ = sym.Poly(bot_Nᶜ, p)
roots = sym.roots(bot_Nᶜ)
# %%
expr = sym.cancel(top_Nᶜ*p/(p*(p-ω)*(p-α₁)*(p-α₂)))
expr = expr.subs(p , 0)
# nᶜ : Checks out!
# %%
expr = sym.factor(sym.cancel(top_Nᶜ/(p*(p-ω)*(p-α₁)*(p-α₂)) - nᶜ/p),p)
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(0)
expr = sym.cancel(top*(p-ω)/bot)
expr = sym.factor(expr.subs(p , ω),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ])
# ρ₀ᶜ : Checks out!
# %%
expr = sym.cancel(top/bot - ρ₀ᶜ/(p-ω))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(ω)
expr = sym.cancel(top*(p-α₁)/bot)
expr = sym.factor(expr.subs(p , α₁),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ, ρ₀ᶜ])
# ρ₁ᶜ : Checks out !
# %%
expr = sym.cancel(top/bot - ρ₁ᶜ/(p-α₁))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(α₁)
expr = sym.cancel(top*(p-α₂)/bot)
expr = sym.factor(expr.subs(p , α₂),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ, ρ₀ᶜ, ρ₁ᶜ])
# %%
# ======================= N⁰ ========================= #
bot_N⁰ = sym.Poly(bot_N⁰, p)
roots = sym.roots(bot_N⁰)
# %%
expr = sym.cancel(top_N⁰*p/(p*(p-ω)*(p-α₁)*(p-α₂)))
expr = expr.subs(p , 0)
# n⁰ : Checks out!
# %%
expr = sym.factor(sym.cancel(top_N⁰/(p*(p-ω)*(p-α₁)*(p-α₂)) - n⁰/p),p)
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(0)
expr = sym.cancel(top*(p-ω)/bot)
expr = sym.factor(expr.subs(p , ω),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ])
# ρ₀⁰ : Checks out!
# %%
expr = sym.cancel(top/bot - ρ₀⁰/(p-ω))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(ω)
expr = sym.cancel(top*(p-α₁)/bot)
expr = sym.factor(expr.subs(p , α₁),[c₀ᶜ, c₀⁰, n₀ᶠ, n⁰, ρ₀⁰])
# ρ₁⁰ : Checks out !
# %%
expr = sym.cancel(top/bot - ρ₁⁰/(p-α₁))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(α₁)
expr = sym.cancel(top*(p-α₂)/bot)
expr = sym.factor(expr.subs(p , α₂),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ, ρ₀ᶜ, ρ₁ᶜ])
# ρ₂⁰ : Checks out !
# %%
# ======================= Nᶠ ========================= #
bot_Nᶠ = sym.Poly(bot_Nᶠ, p)
roots = sym.roots(bot_Nᶠ)
# %%
expr = sym.cancel(top_Nᶠ*p/(p*(p-α₁)*(p-α₂)))
expr = expr.subs(p , 0)
# nᶠ : Checks out!
# %%
expr = sym.factor(sym.cancel(top_Nᶠ/(p*(p-α₁)*(p-α₂)) - nᶠ/p),p)
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(0)
expr = sym.cancel(top*(p-α₁)/bot)
expr = sym.factor(expr.subs(p, α₁),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ])
# ρ₁ᶠ : Checks out!
# %%
expr = sym.cancel(top/bot - ρ₁ᶠ/(p-α₁))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(α₁)
expr = sym.cancel(top*(p-α₂)/bot)
expr = sym.factor(expr.subs(p , α₂),[c₀ᶜ, c₀⁰, n₀ᶠ, n⁰, ρ₀⁰])
# ρ₂ᶠ : Checks out!
# %%
# ====================================================== #
# ================= Dual to classic ==================== #
# ====================================================== #
system = [(p+ξ)*Nᶜ + (ω+ξ)*N⁰ - c₀ᶜ, μᵖ*Nᶜ + (p + μᵖ - ω)*N⁰ + μᵖ*Nᶠ - ( μᵖ/p + c₀⁰ ), (p+δ)*Nᶠ - n₀ᶠ]
sol = sym.linsolve(system, [Nᶜ, N⁰, Nᶠ])
# %%
sol = sym.factor(sol, p)
# %%
Nᶜ = sol.args[1][1]
N⁰ = sol.args[1][2]
Nᶠ = sol.args[1][3]
# %%
top_Nᶜ, bot_Nᶜ = sym.fraction(Nᶜ)
top_Nᶜ =  sym.Poly(top_Nᶜ,p)
for c in top_Nᶜ.coeffs()
    println(sym.factor(c, [c₀ᶜ, c₀⁰, n₀ᶠ]))
end
# %%
top_N⁰, bot_N⁰ = sym.fraction(N⁰)
top_N⁰ =  sym.Poly(top_N⁰,p)
for c in top_N⁰.coeffs()
    println(sym.factor(c, [c₀ᶜ, c₀⁰, n₀ᶠ]))
end
# %%
top_Nᶠ, bot_Nᶠ = sym.fraction(Nᶠ)
top_Nᶠ =  sym.Poly(top_Nᶠ,p)
for c in top_Nᶠ.coeffs()
    println(sym.factor(c, [c₀ᶜ, c₀⁰, n₀ᶠ]))
end
# Laplace transforms seem correct !
# %%
# ======================= Nᶜ ========================= #
bot_Nᶜ = sym.Poly(bot_Nᶜ, p)
roots = sym.roots(bot_Nᶜ)
# %%
expr = sym.cancel(top_Nᶜ*p/(p*(p-ω)*(p + μᵖ + ξ)*(p + δ)))
expr = expr.subs(p , 0)
# nᶜ : Checks out!
# %%
expr = sym.factor(sym.cancel(top_Nᶜ/(p*(p-ω)*(p + μᵖ + ξ)*(p + δ)) - nᶜ/p),p)
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(0)
expr = sym.cancel(top*(p-ω)/bot)
expr = sym.factor(expr.subs(p , ω),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ])
# ρ₀ᶜ : Checks out!
# %%
expr = sym.cancel(top/bot - ρ₀ᶜ/(p-ω))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(ω)
expr = sym.cancel(top*(p + μᵖ + ξ)/bot)
expr = sym.factor(expr.subs(p , -(μᵖ + ξ)),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ, ρ₀ᶜ])
# ρ₁ᶜ : Checks out !
# %%
expr = sym.cancel(top/bot - ρ₁ᶜ/(p + μᵖ + ξ))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(-(μᵖ + ξ))
expr = sym.cancel(top*(p + δ)/bot)
expr = sym.factor(expr.subs(p , -δ),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ, ρ₀ᶜ, ρ₁ᶜ])
# %%
# ======================= N⁰ ========================= #
bot_N⁰ = sym.Poly(bot_N⁰, p)
roots = sym.roots(bot_N⁰)
# %%
expr = sym.cancel(top_N⁰*p/(p*(p-ω)*(p + μᵖ + ξ)*(p + δ)))
expr = expr.subs(p , 0)
# n⁰ : Checks out!
# %%
expr = sym.factor(sym.cancel(top_N⁰/(p*(p-ω)*(p + μᵖ + ξ)*(p + δ)) - n⁰/p),p)
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(0)
expr = sym.cancel(top*(p-ω)/bot)
expr = sym.factor(expr.subs(p , ω),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ])
# ρ₀⁰ : Checks out!
# %%
expr = sym.cancel(top/bot - ρ₀⁰/(p-ω))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(ω)
expr = sym.cancel(top*(p + μᵖ + ξ)/bot)
expr = sym.factor(expr.subs(p , -(μᵖ + ξ)),[c₀ᶜ, c₀⁰, n₀ᶠ, n⁰, ρ₀⁰])
# ρ₁⁰ : Checks out !
# %%
expr = sym.cancel(top/bot - ρ₁⁰/(p + μᵖ + ξ))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(-(μᵖ + ξ))
expr = sym.cancel(top*(p+δ)/bot)
expr = sym.factor(expr.subs(p , -δ),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ, ρ₀ᶜ, ρ₁ᶜ])
# ρ₂⁰ : Checks out !
# %%
# ====================================================== #
# ================= Classic to classic ================= #
# ====================================================== #
system = [(p+ξ)*Nᶜ + (ω+ξ)*N⁰ - c₀ᶜ, μᵖ*Nᶜ + (p + μᵖ - ω)*N⁰ - ( μᵖ/p + c₀⁰ )]
sol = sym.linsolve(system, [Nᶜ, N⁰])
# %%
sol = sym.factor(sol, p)
# %%
Nᶜ = sol.args[1][1]
N⁰ = sol.args[1][2]
# %%
top_Nᶜ, bot_Nᶜ = sym.fraction(Nᶜ)
top_Nᶜ =  sym.Poly(top_Nᶜ,p)
for c in top_Nᶜ.coeffs()
    println(sym.factor(c, [c₀ᶜ, c₀⁰, n₀ᶠ]))
end
# %%
top_N⁰, bot_N⁰ = sym.fraction(N⁰)
top_N⁰ =  sym.Poly(top_N⁰,p)
for c in top_N⁰.coeffs()
    println(sym.factor(c, [c₀ᶜ, c₀⁰, n₀ᶠ]))
end
# Laplace transforms seem correct !
# %%
# ======================= Nᶜ ========================= #
bot_Nᶜ = sym.Poly(bot_Nᶜ, p)
roots = sym.roots(bot_Nᶜ)
# %%
expr = sym.cancel(top_Nᶜ*p/(p*(p-ω)*(p + μᵖ + ξ)))
expr = expr.subs(p , 0)
# nᶜ : Checks out!
# %%
expr = sym.factor(sym.cancel(top_Nᶜ/(p*(p-ω)*(p + μᵖ + ξ)) - nᶜ/p),p)
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(0)
expr = sym.cancel(top*(p-ω)/bot)
expr = sym.factor(expr.subs(p , ω),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ])
# ρ₀ᶜ : Checks out!
# %%
expr = sym.cancel(top/bot - ρ₀ᶜ/(p-ω))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(ω)
expr = sym.cancel(top*(p + μᵖ + ξ)/bot)
expr = sym.factor(expr.subs(p , -(μᵖ + ξ)),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ, ρ₀ᶜ])
# ρ₁ᶜ : Checks out !
# %%
# ======================= N⁰ ========================= #
bot_N⁰ = sym.Poly(bot_N⁰, p)
roots = sym.roots(bot_N⁰)
# %%
expr = sym.cancel(top_N⁰*p/(p*(p-ω)*(p + μᵖ + ξ)))
expr = expr.subs(p , 0)
# n⁰ : Checks out!
# %%
expr = sym.factor(sym.cancel(top_N⁰/(p*(p-ω)*(p + μᵖ + ξ)) - n⁰/p),p)
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(0)
expr = sym.cancel(top*(p-ω)/bot)
expr = sym.factor(expr.subs(p , ω),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ])
# ρ₀⁰ : Checks out!
# %%
expr = sym.cancel(top/bot - ρ₀⁰/(p-ω))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(ω)
expr = sym.cancel(top*(p + μᵖ + ξ)/bot)
expr = sym.factor(expr.subs(p , -(μᵖ + ξ)),[c₀ᶜ, c₀⁰, n₀ᶠ, n⁰, ρ₀⁰])
# ρ₁⁰ : Checks out !
# %%
expr = sym.cancel(top/bot - ρ₁⁰/(p + μᵖ + ξ))
top,bot = sym.fraction(expr)
top = sym.Poly(top, p)
top = top - top(-(μᵖ + ξ))
expr = sym.cancel(top*(p+δ)/bot)
expr = sym.factor(expr.subs(p , -δ),[c₀ᶜ, c₀⁰, n₀ᶠ, nᶜ, ρ₀ᶜ, ρ₁ᶜ])
# ρ₂⁰ : Checks out !
# %%
# ====================================================== #
# ======================= Wages ======================== #
# ====================================================== #
system = [(p-ω)*Wᶜ - W₀ᶜ + (ω+ξ)*Eᶜ*(nᶜ/p + ρ₀ᶜ/(p-ω) + ρ₁ᶜ/(p-α₁) + ρ₂ᶜ/(p-α₂) + n⁰/p + ρ₀⁰/(p-ω) + ρ₁⁰/(p-α₁) + ρ₂⁰/(p-α₂)) + α₀*E₀ᶜ*n₀ᶜ, (p-ω)*W⁰ - W₀⁰ - μᵖ*E⁰*(1/p - (nᶜ/p + ρ₀ᶜ/(p-ω) + ρ₁ᶜ/(p-α₁) + ρ₂ᶜ/(p-α₂)) - (n⁰/p + ρ₀⁰/(p-ω) + ρ₁⁰/(p-α₁) + ρ₂⁰/(p-α₂)) - (nᶠ/p + ρ₁ᶠ/(p-α₁) + ρ₂ᶠ/(p-α₂))) + β₀*E₀⁰*n₀⁰]
sol = sym.linsolve(system, [Wᶜ, W⁰])
# %%
Wᶜ = sol.args[1][1]
W⁰ = sol.args[1][2]
# %%
top_Wᶜ, bot_Wᶜ = sym.fraction(Wᶜ)
top_Wᶜ =  sym.Poly(top_Wᶜ,p)
# %%
top_W⁰, bot_W⁰ = sym.fraction(W⁰)
top_W⁰ =  sym.Poly(top_W⁰,p)
# Laplace transforms seem correct !
# %%
# ======================= Wᶜ ========================= #
bot_Wᶜ = sym.Poly(bot_Wᶜ, p)
roots = sym.roots(bot_Wᶜ)
# %%
expr = sym.cancel(top_Wᶜ/((p-ω)^2*(p-α₁)*(p-α₂)))
W∞ᶜ = sym.cancel(expr.subs(p , 0))
# %%
expr = sym.cancel(top_Wᶜ/(p*(p-α₁)*(p-α₂)))
ν₀ᶜ = sym.cancel(expr.subs(p , ω))
# %%
expr = sym.cancel(top_Wᶜ/(p*(p-ω)*(p-α₁)*(p-α₂)) - ν₀ᶜ/(p-ω))
expr = expr.subs(p, ω)
sym.apart(expr, ω)
# %%
expr = sym.cancel(top_Wᶜ/(p*(p-ω)^2*(p-α₂)))
ν₁ᶜ = sym.cancel(expr.subs(p , α₁))
# ν₁ᶜ : Checks out !
# %%
expr = sym.cancel(top_Wᶜ/(p*(p-ω)^2*(p-α₁)))
ν₂ᶜ = sym.cancel(expr.subs(p , α₂))
# ν₂ᶜ : Checks out !
# %%
# ======================= W⁰ ========================= #
bot_W⁰ = sym.Poly(bot_W⁰, p)
roots = sym.roots(bot_W⁰)
# %%
expr = sym.cancel(top_W⁰/((p-ω)^2*(p-α₁)*(p-α₂)))
sym.cancel(expr.subs(p , 0))
# %%
expr = sym.cancel(top_W⁰/(p*(p-α₁)*(p-α₂)))
ν₀⁰ = sym.cancel(expr.subs(p , ω))
# %%
expr = sym.cancel(top_W⁰/(p*(p-ω)*(p-α₁)*(p-α₂)) - ν₀⁰/(p-ω))
expr = expr.subs(p, ω)
sym.apart(expr, ω)
# %%
expr = sym.cancel(top_W⁰/(p*(p-ω)^2*(p-α₂)))
ν₁⁰ = sym.cancel(expr.subs(p , α₁))
# ν₁ᶜ : Checks out !
# %%
expr = sym.cancel(top_W⁰/(p*(p-ω)^2*(p-α₁)))
ν₂⁰ = sym.cancel(expr.subs(p , α₂))
#


using BasisFunctions, Test, DoubleFloats

for T in (Float64,LargeFloat)

    @testset "Discrete Gauss OPS Rule" begin
        for d in (Legendre{T}(10), Laguerre(10,rand(T)), Hermite{T}(10), ChebyshevT{T}(10),ChebyshevU{T}(10),Jacobi(10,rand(T),rand(T)),Fourier{T}(10))
            g = interpolation_grid(d)
            gr = gauss_rule(d)
            @test g ≈ points(gr)
            @test isorthogonal(d, gr)
            @test iscompatible(d, g)
            @test gram(d, gr) isa DiagonalOperator

            test_orthogonality_orthonormality(d, true, d isa Fourier, gr)

            @test BasisFunctions.quadweights(gauss_rule(d),measure(d)) ≈ ones(T, 10)
        end

        B = Fourier(10)
        d = B^2
        g = interpolation_grid(d)
        gr = gauss_rule(d)
        @test g ≈ points(gr)
        @test isorthogonal(d, gr)
        @test iscompatible(d, g)
        @test gram(d, gr) isa TensorProductOperator && all(map(x->isa(x, DiagonalOperator), components(gram(d,gr))))
        @test BasisFunctions.quadweights(gauss_rule(d),measure(d)) ≈ BasisFunctions.OuterProductArray(map(ones, axes(B^2))...)

    end
end


# Here are the main exported functions and types. Note there are other exported functions and types in files such as "effective_waves/effective_waves_export" and "average_waves/average_waves_export.jl."

export  EffectiveWave, AverageWave # the two main types
export  MatchWave # a combination of the two types above

# for MatchWave
export  match_error, x_mesh_match

# for discrete method
export  x_mesh

# for material and particle properties
export  Specie, Medium, volume_fraction, Zn, t_vectors, Nn, p_speed, maximum_hankel_order

# for effective waves
export  wavenumbers, wavenumber, effective_waves, transmission_angle,
        effective_wavevectors, scattering_amplitudes_average, scale_amplitudes_effective

export  reflection_coefficient, reflection_coefficients

export  effective_medium, reflection_coefficient_halfspace

# List of shorthand for some materials
export  Brick, IronArmco, LeadAnnealed, RubberGum, FusedSilica, GlassPyrex,
        ClayRock, WaterDistilled, Glycerol, Hexadecane, Acetone, Benzene,
        Nitrobenzene, OliveOil, SodiumNitrate, AirDry,
        LimeStone, Clay, Calcite, SilicaQuartz

import Base.isequal, Base.(==), Base.zero
import SpecialFunctions: besselj, hankelh1
import Statistics: mean, std


"""
    F = JopLog(spc)

where `F` is the log operator with domain and range given by `spc::JetSpace`.
"""
JopLog(spc::JetSpace) = JopNl(dom = spc, rng = spc, f! = JopLog_f!, df! = JopLog_df!, df′! = JopLog_df′!)
export JopLog

# TODO: assert denominator in linearization is > 0

JopLog_f!(d::AbstractArray, m::AbstractArray) = d .= log.(m)
JopLog_df!(δd::AbstractArray, δm::AbstractArray; mₒ) = δd .= δm ./ mₒ
JopLog_df′!(δm::AbstractArray, δd::AbstractArray; mₒ) = δm .= δd ./ conj(mₒ)


using LinearAlgebra
using QuantumGate

function single_qubit_gate(cite::Int, N::Int, Operator)
    A = σ⁰

    if cite == 1
        A = Operator
    end

    for i in 2:N
        if cite == i
            A = A ⊗ Operator
        else
            A = A ⊗ σ⁰
        end
    end

    return A
end

function Xgate(cite::Int, N::Int)
    single_qubit_gate(cite, N, σˣ)
end

function Ygate(cite::Int, N::Int)
    single_qubit_gate(cite, N, σʸ)
end

function Zgate(cite::Int, N::Int)
    single_qubit_gate(cite, N, σᶻ)
end

function Hadamardgate(cite::Int, N::Int)
    single_qubit_gate(cite, N, Had)
end

function Sgate(cite::Int, N::Int)
    single_qubit_gate(cite, N, Smat)
end

function Tgate(cite::Int, N::Int)
    single_qubit_gate(cite, N, Tmat)
end

function Rˣgate(cite::Int, N::Int, θ::Float64)
    single_qubit_gate(cite, N, Rˣmat(θ))
end

function Rʸgate(cite::Int, N::Int, θ::Float64)
    single_qubit_gate(cite, N, Rʸmat(θ))
end

function Rᶻgate(cite::Int, N::Int, θ::Float64)
    single_qubit_gate(cite, N, Rᶻmat(θ))
end

function controlled_Ugate(control::Int, target::Int, N::Int, U)
    A = single_qubit_gate(control, N, P⁰)

    if control == 1
        B = P¹
    elseif target == 1
        B = U
    else
        B = σ⁰
    end

    for i in 2:N
        if control == i
            B = B ⊗ P¹
        elseif target == i
            B = B ⊗ U
        else
            B = B ⊗ σ⁰
        end
    end
    return A + B
end

function CNOTgate(control::Int, target::Int, N::Int)
    controlled_Ugate(control, target, N, σˣ)
end

function SWAPgate(i, j, N)
    CNOTgate(i, j, N) * CNOTgate(j, i, N) * CNOTgate(i, j, N)
end


## Section definitions
#==========================================================================================#

"""
HalfWingSection(; span, dihedral, sweep, taper, root_chord,
                  root_twist, tip_twist, root_foil, tip_foil,
                  position, angle, axis)

Define a `HalfWing` consisting of a single trapezoidal section.

# Arguments
- `span       :: Real         = 1.`: Span length
- `dihedral   :: Real         = 1.`: Dihedral angle (degrees)
- `sweep      :: Real         = 0.`: Leading-edge sweep angle (degrees)
- `taper      :: Real         = 1.`: Taper ratio of tip to root chord
- `root_chord :: Real         = 1.`: Root chord length
- `root_twist :: Real         = 0.`: Twist angle at root (degrees)
- `tip_twist  :: Real         = 0.`: Twist angle at tip (degrees)
- `root_foil  :: Foil         = naca4((0,0,1,2))`: Foil coordinates at root
- `tip_foil   :: Foil         = root_foil`: Foil coordinates at tip
- `position   :: Vector{Real} = zeros(3)`: Position
- `angle      :: Real         = 0.`: Angle of rotation (degrees)
- `axis       :: Vector{Real} = [0.,1.,0.]`: Axis of rotation
"""
function HalfWingSection(;
        span         = 1.,
        dihedral     = 0.,
        sweep        = 0.,
        w_sweep      = 0.,
        taper        = 1.,
        root_chord   = 1.,
        root_twist   = 0.,
        tip_twist    = 0.,
        root_control = (0., 0.75),
        tip_control  = root_control,
        root_foil    = naca4(0,0,1,2),
        tip_foil     = root_foil,
        position     = zeros(3),
        angle        = 0.,
        axis         = SVector(0., 1., 0.),
    )
    # Control surface deflections at root and tip
    root_foil = control_surface(root_foil;
                                angle = root_control[1],
                                hinge = root_control[2])

    tip_foil  = control_surface(tip_foil;
                                angle = tip_control[1],
                                hinge = tip_control[2])

    section = HalfWing(
        foils     = [root_foil, tip_foil],
        chords    = [root_chord, taper * root_chord],
        twists    = [root_twist, tip_twist],
        spans     = [span],
        dihedrals = [dihedral],
        sweeps    = [sweep],
        w_sweep   = w_sweep,
        position  = position,
        angle     = angle,
        axis      = axis
    )

    return section
end

function HalfWingSection(S, AR, λ, Λ, δ, τ_r, τ_t, w = 0.25)
    b    = S^2 / AR
    c_r  = 2 * S / (b * (1 + λ))

    HalfWingSection(span = AR, dihedral = δ, sweep = Λ, w_sweep = w, taper = λ, root_chord = c_r, root_twist = τ_r, tip_twist = τ_t)
end

"""
    WingSection(; span, dihedral, sweep, taper, root_chord,
                  root_twist, tip_twist, root_foil, tip_foil,
                  position, angle, axis)

Define a `Wing` consisting of two trapezoidal sections with reflection symmetry in the ``x``-``z`` plane.

# Arguments
- `span       :: Real         = 1.`: Span length
- `dihedral   :: Real         = 1.`: Dihedral angle (degrees)
- `LE_sweep   :: Real         = 0.`: Leading-edge sweep angle (degrees)
- `taper      :: Real         = 1.`: Taper ratio of tip to root chord
- `root_chord :: Real         = 1.`: Root chord length
- `root_twist :: Real         = 0.`: Twist angle at root (degrees)
- `tip_twist  :: Real         = 0.`: Twist angle at tip (degrees)
- `root_foil  :: Foil         = naca4((0,0,1,2))`: Foil coordinates at root
- `tip_foil   :: Foil         = root_foil`: Foil coordinates at tip
- `position   :: Vector{Real} = zeros(3)`: Position
- `angle      :: Real         = 0.`: Angle of rotation (degrees)
- `axis       :: Vector{Real} = [0.,1.,0.]`: Axis of rotation
"""
function WingSection(;
        span         = 1.,
        dihedral     = 0.,
        sweep        = 0.,
        w_sweep      = 0.,
        taper        = 1.,
        root_chord   = 1.,
        root_twist   = 0.,
        tip_twist    = 0.,
        root_control = (0., 0.75),
        tip_control  = root_control,
        root_foil    = naca4((0,0,1,2)),
        tip_foil     = root_foil,
        position     = zeros(3),
        angle        = 0.,
        axis         = SVector(1., 0., 0.),
    )

    right = HalfWingSection(
                span         = span / 2,
                dihedral     = dihedral,
                sweep        = sweep,
                w_sweep      = w_sweep,
                taper        = taper,
                root_chord   = root_chord,
                root_twist   = root_twist,
                tip_twist    = tip_twist,
                root_control = root_control,
                tip_control  = tip_control,
                root_foil    = root_foil,
                tip_foil     = tip_foil,
                position     = position,
                angle        = angle,
                axis         = axis
            )

    Wing(right)
end

## Standard stabilizers (non-exhaustive)
#==========================================================================================#

HorizontalTail(; root_chord = 1.0, taper = 1.0, span = 1.0, sweep = 0., w_sweep = 0., root_twist = 0., tip_twist = 0., root_foil = naca4(0,0,1,2), tip_foil = root_foil, angle = 0.) = WingSection(root_chord = root_chord, taper = taper, span = span, sweep = sweep, w_sweep = w_sweep, root_twist = root_twist, tip_twist = tip_twist, root_foil = naca4(0,0,1,2), tip_foil = tip_foil, angle = angle, axis = [0., 1., 0.])

VerticalTail(; root_chord = 1.0, taper = 1.0, span = 1.0, sweep = 0., w_sweep = 0., root_foil = naca4(0,0,1,2), tip_foil = root_foil) = HalfWingSection(root_chord = root_chord, taper = taper, span = span, sweep = sweep, w_sweep = w_sweep, tip_foil = tip_foil, angle = 90., axis = [1., 0., 0.])


using Test

using ScoreAD: score_AD, score_AD_log, reject_nonfinite

using Distributions: Normal, logpdf
using Statistics: mean
import ForwardDiff: Dual, value, derivative

"""
    BernoulliData(N, s)

`N` draws from a Bernoulli(θ) distribution, `s` of which are `1`, `N-s` are `0`.
"""
struct BernoulliData
    "Total number of observations."
    N::Int
    "Observations =1."
    s::Int
end

(data::BernoulliData)(θ) = data.s/data.N*score_AD(θ)

@testset "bernoulli" begin
    N = 105
    s = 61
    data = BernoulliData(N, s)
    θ = 0.5
    for θ in range(0.1; stop = 0.9, length = 100)
        @test data(θ) ≈ s/N
        @test derivative(data, θ) ≈ s/N * 1/θ
    end
end

"""
    NormalData(x)

A vector of draws from a Normal(0, σ) distribution.
"""
struct NormalData{T}
    x::Vector{T}
end

function (data::NormalData)(σ)
    d = Normal(0.0, σ)
    mean(x * score_AD_log(logpdf(d, x)) for x in data.x)
end

@testset "normal" begin
    σ = 1.0
    x = rand(Normal(0, σ), 100)
    data = NormalData(x)
    for σ in range(0.1; stop = 5, length = 100)
        @test data(σ) ≈ mean(x)
        @test derivative(data, σ) ≈ mean(x * (x^2/(σ^3) - 1/σ) for x in x)
    end
end

@testset "reject nonfinite" begin
    @test (@inferred reject_nonfinite(-1.0)) ≡ -1.0
    @test (@inferred reject_nonfinite(-Inf)) ≡ -Inf
    @test (@inferred reject_nonfinite(Inf)) ≡ -Inf
    @test (@inferred reject_nonfinite(NaN)) ≡ -Inf

    @test (@inferred reject_nonfinite(Dual(-1.0, (2.0, 3.0)))) ==
        Dual(-1.0, (2.0, 3.0))

    function test_rejected_Dual(ℓ)
        r = @inferred reject_nonfinite(ℓ)
        @test value(r) ≡ oftype(value(ℓ), -Inf)
        @test typeof(r) ≡ typeof(ℓ)
    end

    test_rejected_Dual(Dual(-Inf, (2.0, 3.0)))
    test_rejected_Dual(Dual(Inf, (2.0, 3.0)))
    test_rejected_Dual(Dual(NaN, (2.0, 3.0)))
    test_rejected_Dual(Dual(1.0, (NaN, 3.0)))
    test_rejected_Dual(Dual(1.0, (2.0, -Inf)))
end