text:
export AbstractModel, SampledModel, IntegratedModel, DiffusiveSampledModel, DiffusiveIntegratedModel, HyperbolicModel, ParabolicModel
export num_species, external_velocity, interaction, mobility, diffusion, eachspecies, species
export SM, IM, DSM, DIM, HM, PM

"""
$(TYPEDEF)

Abstract type representing a particle model.
"""
abstract type AbstractModel{N} end

"""
$(TYPEDSIGNATURES)

Returns the number of species of the model.
"""
num_species(mod::AbstractModel) = error("unimplemented")
num_species(mod::Type{<:AbstractModel}) = error("unimplemented")

"""
$(TYPEDSIGNATURES)

Returns the external velocity field associated to the species `i`.
"""
external_velocity(mod::AbstractModel, i::Integer) = error("unimplemented")

"""
$(TYPEDSIGNATURES)

Returns the interaction exerted on the species `i` by the species `j`.
"""
interaction(mod::AbstractModel, i::Integer, j::Integer) = error("unimplemented")

"""
$(TYPEDSIGNATURES)

Returns the mobility associated to the species `i`.
"""
mobility(mod::AbstractModel, i::Integer) = error("unimplemented")

"""
$(TYPEDSIGNATURES)

Returns the diffusion associated to the species `i`.
"""
diffusion(mod::AbstractModel, i::Integer) = error("unimplemented")

"""
$(TYPEDSIGNATURES)

Returns an iterator over the indices of the species of the model.
"""
eachspecies(mod::AbstractModel) = 1:num_species(mod)
eachspecies(mod::Type{<:AbstractModel}) = 1:num_species(mod)

"""
$(TYPEDSIGNATURES)

Returns the particles associated to the species `i`.
"""
species(state, i::Integer) = error("unimplemented")
species(a::ArrayPartition, i::Integer) = a.x[i]

"""
    SampledModel((V₁, ...), ((W′₁₁, ...), ...), (mob₁, ...))

Represents a particles system with:

- external velocities `Vᵢ`,
- sampled interactions `W′ᵢⱼ` (this is the effect of the species `j` on the species `i`),
- mobilities `mobᵢ`.

See also [`IntegratedModel`](@ref).

# Examples

```jldoctest
julia> using ConservationLawsParticles.Examples, RecursiveArrayTools

julia> model = SampledModel(
           (V, V),
           ((Wprime_attr, Wprime_rep),
            (Wprime_rep, Wprime_attr)),
           (mobρ, mobσ));

julia> x = ArrayPartition(gaussian_particles(2, 4), collect(range(-3, 4, length=5)))
([-2.000000000000003, -0.30305473018369145, 0.30305473018369145, 2.000000000000003], [-3.0, -1.25, 0.5, 2.25, 4.0])

julia> velocities(x, model, 0.)
([6.291136247066298, 0.2466663161150116, -0.7218917091339228, -7.22630670503873], [22.405129478914613, 1.1249366684885518, 1.5188519354999799, -7.87111869358889, -54.536397957423915])
```
"""
struct SampledModel{
    N,
    TVs         <: Tuple{Vararg{Any,N}},
    TWprimes    <: Tuple{Vararg{Tuple{Vararg{Any,N}},N}},
    Tmobilities <: Tuple{Vararg{Any,N}},
} <: AbstractModel{N}
    Vs::TVs
    Wprimes::TWprimes
    mobilities::Tmobilities
end

SampledModel(V, Wprime, mob) = SampledModel((V,), ((Wprime,),), (mob,))

"""
    IntegratedModel((V₁, ...), ((W₁₁, ...), ...), (mob₁, ...))

Represents a particles system with:

- external velocities `Vᵢ`,
- integrated interactions `Wᵢⱼ` (this is the effect of the species `j` on the species `i`),
- mobilities `mobᵢ`.

See also [`SampledModel`](@ref).

# Examples

```jldoctest
julia> using ConservationLawsParticles.Examples, RecursiveArrayTools

julia> model = IntegratedModel(
           (V, V),
           ((W_attr, W_rep),
            (W_rep, W_attr)),
           (mobρ, mobσ));

julia> x = ArrayPartition(gaussian_particles(2, 4), collect(range(-3, 4, length=5)))
([-2.000000000000003, -0.30305473018369145, 0.30305473018369145, 2.000000000000003], [-3.0, -1.25, 0.5, 2.25, 4.0])

julia> velocities(x, model, 0.)
([6.621647425385332, 0.30545649966452776, -0.42671180044513507, -6.914302510772401], [23.261098611816987, 0.9023583690557855, 0.7055593913708742, -8.774526834146023, -55.23308442021724])
```
"""
struct IntegratedModel{
    N,
    TVs         <: Tuple{Vararg{Any,N}},
    TWs         <: Tuple{Vararg{Tuple{Vararg{Any,N}},N}},
    Tmobilities <: Tuple{Vararg{Any,N}},
} <: AbstractModel{N}
    Vs::TVs
    Ws::TWs
    mobilities::Tmobilities
end

IntegratedModel(V, W, mob) = IntegratedModel((V,), ((W,),), (mob,))

num_species(mod::SampledModel{N}) where N = N
num_species(mod::IntegratedModel{N}) where N = N

num_species(mod::Type{<:SampledModel{N}}) where N = N
num_species(mod::Type{<:IntegratedModel{N}}) where N = N

external_velocity(mod::SampledModel, i::Integer) = mod.Vs[i]
external_velocity(mod::IntegratedModel, i::Integer) = mod.Vs[i]

interaction(mod::SampledModel, i::Integer, j::Integer) = mod.Wprimes[i][j]
interaction(mod::IntegratedModel, i::Integer, j::Integer) = mod.Ws[i][j]

mobility(mod::SampledModel, i::Integer) = mod.mobilities[i]
mobility(mod::IntegratedModel, i::Integer) = mod.mobilities[i]

diffusion(mod::SampledModel, i::Integer) = nothing
diffusion(mod::IntegratedModel, i::Integer) = nothing

"""
    DiffusiveSampledModel((V₁, ...), ((W′₁₁, ...), ...), (mob₁, ...), (D₁, ...))

Represents a particles system with:

- external velocities `Vᵢ`,
- sampled interactions `W′ᵢⱼ` (this is the effect of the species `j` on the species `i`),
- mobilities `mobᵢ`,
- diffusions `Dᵢ`.

See also [`DiffusiveIntegratedModel`](@ref).

# Examples

```jldoctest
julia> using ConservationLawsParticles.Examples, RecursiveArrayTools

julia> model = DiffusiveSampledModel(
           (V, V),
           ((Wprime_attr, Wprime_rep),
            (Wprime_rep, Wprime_attr)),
           (mobρ, mobσ),
           (Diffusion(1), nothing));

julia> x = ArrayPartition(gaussian_particles(2, 4), collect(range(-3, 4, length=5)))
([-2.000000000000003, -0.30305473018369145, 0.30305473018369145, 2.000000000000003], [-3.0, -1.25, 0.5, 2.25, 4.0])

julia> velocities_diff(x, model, 0.)
([5.701842050142524, -1.3162440100303214, 0.8410186170114102, -6.637012508114956], [22.405129478914613, 1.1249366684885518, 1.5188519354999799, -7.87111869358889, -54.536397957423915])
```
"""
struct DiffusiveSampledModel{
    N,
    TVs         <: Tuple{Vararg{Any,N}},
    TWprimes    <: Tuple{Vararg{Tuple{Vararg{Any,N}},N}},
    Tmobilities <: Tuple{Vararg{Any,N}},
    Tdiffusions <: Tuple{Vararg{Any,N}}
} <: AbstractModel{N}
    Vs::TVs
    Wprimes::TWprimes
    mobilities::Tmobilities
    diffusions::Tdiffusions
end

"""
    DiffusiveIntegratedModel((V₁, ...), ((W₁₁, ...), ...), (mob₁, ...), (D₁, ...))

Represents a particles system with:

- external velocities `Vᵢ`,
- integrated interactions `Wᵢⱼ` (this is the effect of the species `j` on the species `i`),
- mobilities `mobᵢ`,
- diffusions `Dᵢ`.

See also [`DiffusiveSampledModel`](@ref).

# Examples

```jldoctest
julia> using ConservationLawsParticles.Examples, RecursiveArrayTools

julia> model = DiffusiveIntegratedModel(
           (V, V),
           ((W_attr, W_rep),
            (W_rep, W_attr)),
           (mobρ, mobσ),
           (Diffusion(1), nothing));

julia> x = ArrayPartition(gaussian_particles(2, 4), collect(range(-3, 4, length=5)))
([-2.000000000000003, -0.30305473018369145, 0.30305473018369145, 2.000000000000003], [-3.0, -1.25, 0.5, 2.25, 4.0])

julia> velocities_diff(x, model, 0.)
([6.032353228461558, -1.2574538264808053, 1.1361985257001979, -6.3250083138486275], [23.261098611816987, 0.9023583690557855, 0.7055593913708742, -8.774526834146023, -55.23308442021724])
```
"""
struct DiffusiveIntegratedModel{
    N,
    TVs         <: Tuple{Vararg{Any,N}},
    TWs         <: Tuple{Vararg{Tuple{Vararg{Any,N}},N}},
    Tmobilities <: Tuple{Vararg{Any,N}},
    Tdiffusions <: Tuple{Vararg{Any,N}}
} <: AbstractModel{N}
    Vs::TVs
    Ws::TWs
    mobilities::Tmobilities
    diffusions::Tdiffusions
end

num_species(mod::DiffusiveSampledModel{N}) where N = N
num_species(mod::DiffusiveIntegratedModel{N}) where N = N

num_species(mod::Type{<:DiffusiveSampledModel{N}}) where N = N
num_species(mod::Type{<:DiffusiveIntegratedModel{N}}) where N = N

external_velocity(mod::DiffusiveSampledModel, i::Integer) = mod.Vs[i]
external_velocity(mod::DiffusiveIntegratedModel, i::Integer) = mod.Vs[i]

interaction(mod::DiffusiveSampledModel, i::Integer, j::Integer) = mod.Wprimes[i][j]
interaction(mod::DiffusiveIntegratedModel, i::Integer, j::Integer) = mod.Ws[i][j]

mobility(mod::DiffusiveSampledModel, i::Integer) = mod.mobilities[i]
mobility(mod::DiffusiveIntegratedModel, i::Integer) = mod.mobilities[i]

diffusion(mod::DiffusiveSampledModel, i::Integer) = mod.diffusions[i]
diffusion(mod::DiffusiveIntegratedModel, i::Integer) = mod.diffusions[i]

"""
    HyperbolicModel((V₁, ...), ((I₁₁, ...), ...), (mob₁, ...))

Represents a particles system with:

- external velocities `Vᵢ`,
- integrated interactions `Iᵢⱼ` (this is the effect of the species `j` on the species `i`),
- mobilities `mobᵢ`.

See also [`ParabolicModel`](@ref).

# Examples

```jldoctest
julia> using ConservationLawsParticles.Examples, RecursiveArrayTools

julia> model = HyperbolicModel(
           (V, V),
           ((SampledInteraction(Wprime_attr), IntegratedInteraction(W_rep)),
            (IntegratedInteraction(W_rep), SampledInteraction(Wprime_attr))),
           (mobρ, mobσ));

julia> x = ArrayPartition(gaussian_particles(2, 4), collect(range(-3, 4, length=5)))
([-2.000000000000003, -0.30305473018369145, 0.30305473018369145, 2.000000000000003], [-3.0, -1.25, 0.5, 2.25, 4.0])

julia> abstract_velocities(x, model, 0.)
([6.267901451064502, 0.2629261335925302, -0.3841814343731375, -6.560556536451573], [22.921033907881895, 0.8199108735979148, 0.7055593913708742, -8.681409489341364, -54.89301971628215])
```
"""
struct HyperbolicModel{
    N,
    TVs <: NTuple{N,Any},
    TIs <: NTuple{N,NTuple{N,Any}},
    TMs <: NTuple{N,Any}
} <: AbstractModel{N}
    vels::TVs
    ints::TIs
    mobs::TMs
end

"""
    ParabolicModel((V₁, ...), ((I₁₁, ...), ...), (mob₁, ...), (D₁, ...))

Represents a particles system with:

- external velocities `Vᵢ`,
- interactions `Iᵢⱼ` (this is the effect of the species `j` on the species `i`),
- mobilities `mobᵢ`,
- diffusions `Dᵢ`.

See also [`HyperbolicModel`](@ref).

# Examples

```jldoctest
julia> using ConservationLawsParticles.Examples, RecursiveArrayTools

julia> model = ParabolicModel(
           (V, V),
           ((SampledInteraction(Wprime_attr), IntegratedInteraction(W_rep)),
            (IntegratedInteraction(W_rep), SampledInteraction(Wprime_attr))),
           (mobρ, mobσ),
           (Diffusion(1), nothing));

julia> x = ArrayPartition(gaussian_particles(2, 4), collect(range(-3, 4, length=5)))
([-2.000000000000003, -0.30305473018369145, 0.30305473018369145, 2.000000000000003], [-3.0, -1.25, 0.5, 2.25, 4.0])

julia> abstract_velocities(x, model, 0.)
([5.678607254140728, -1.2999841925528028, 1.1787288917721954, -5.971262339527799], [22.921033907881895, 0.8199108735979148, 0.7055593913708742, -8.681409489341364, -54.89301971628215])
```
"""
struct ParabolicModel{
    N,
    TVs <: NTuple{N,Any},
    TIs <: NTuple{N,NTuple{N,Any}},
    TMs <: NTuple{N,Any},
    TDs <: NTuple{N,Any}
} <: AbstractModel{N}
    vels::TVs
    ints::TIs
    mobs::TMs
    difs::TDs
end

num_species(mod::HyperbolicModel{N}) where N = N
num_species(mod::ParabolicModel{N}) where N = N

num_species(mod::Type{<:HyperbolicModel{N}}) where N = N
num_species(mod::Type{<:ParabolicModel{N}}) where N = N

external_velocity(mod::HyperbolicModel, i::Integer) = mod.vels[i]
external_velocity(mod::ParabolicModel, i::Integer) = mod.vels[i]

interaction(mod::HyperbolicModel, i::Integer, j::Integer) = mod.ints[i][j]
interaction(mod::ParabolicModel, i::Integer, j::Integer) = mod.ints[i][j]

mobility(mod::HyperbolicModel, i::Integer) = mod.mobs[i]
mobility(mod::ParabolicModel, i::Integer) = mod.mobs[i]

diffusion(mod::HyperbolicModel, i::Integer) = nothing
diffusion(mod::ParabolicModel, i::Integer) = mod.difs[i]

const SM = SampledModel
const IM = IntegratedModel
const DSM = DiffusiveSampledModel
const DIM = DiffusiveIntegratedModel
const HM = HyperbolicModel
const PM = ParabolicModel


text:
export Tableau, Victory

"""
A `Tableau` is a full position in FreeCell with four free cells,
a foundation, and a cascade.
+ `Tableau()` creates an empty Tableau
+ `Tableau(deck())` is the way to start with a shuffled deck
"""
struct Tableau
    Free::free_cells
    Found::foundation
    Casc::cascade
    function Tableau()
        FC = free_cells()
        FD = foundation()
        CC = cascade()
        new(FC, FD, CC)
    end
end

"""
`Victory()` returns the tableau in which all cards have been
moved into the foundation.
"""
function Victory()::Tableau
    T = Tableau()
    for k = 1:52
        add_card(T.Found, Card(k))
    end
    return T
end

function Tableau(d::Vector{Card})
    C = cascade(d)
    T = Tableau()
    for k = 1:8
        T.Casc.piles[k] = C.piles[k]
    end
    return T
end

function show(io::IO, T::Tableau)
    wide = 50
    println(io, "-"^wide)
    println(io, "Free Cells:\n\t$(free_string(T.Free))")
    println(io, "Foundation:")
    println(io, T.Found)
    println(io, "Cascade:")
    println(io, T.Casc)
    print(io, "="^wide)
end

"""
`score(T::Tableau)` is the number of cards moved into the foundation.
"""
function score(T::Tableau)::Number
    return -(10 * score(T.Free) + 100 * score(T.Found) + score(T.Casc))
end

function hash(T::Tableau, h::UInt)
    a = hash(T.Free, h)
    b = hash(T.Found.piles, h)
    c = hash(T.Casc.piles, h)

    hash([a, b, c])
end

function (==)(T::Tableau, TT::Tableau)
    T.Free == TT.Free && T.Found.piles == TT.Found.piles && T.Casc.piles == TT.Casc.piles
end

ext: jl


text:
"""
    module TimeDomain
"""
module TimeDomain

files = (
         "1D/TimeDomain1D.jl",
         "2D/TimeDomain2D.jl",
         # "3D/TimeDomain3D.jl",
        )

export Constant
export HelmholtzFDTD
export HelmholtzBlochFDTD
export HelmholtzPointSource
export propagate!
export initialize!

using ..Common
using Formatting
using Plots
using RecipesBase

import ..Common.DEFAULT_CFL_NUMBER
import LinearAlgebra: mul!

abstract type AbstractFields{N,M} end
abstract type AbstractFDTD{N,M,P} end

include("DivGradCurls.jl")

propagate!() = nothing
_αβ() = nothing

################################################################################
# Define Field Structures

struct HelmholtzWaveFields{N,M} <: AbstractFields{N,M} # N is dimension, M is memory depth
    φ::NTuple{M,VectorField{N,1,Float64}}
    ∇Φ::NTuple{M,NTuple{N,VectorField{N,1,Float64}}}
end

function Base.getproperty(hwf::HelmholtzWaveFields{N,2}, sym::Symbol) where N
    if Base.sym_in(sym,(:φ,:ϕ,:phi,:psi,:ψ))
        return getfield(hwf,sym)[1]
    else
        return getfield(hwf,sym)
    end
end

# struct HelmholtzMatterFields{N} <: AbstractFields{N,2}
#     ρ::VectorField{N,1,Float64}
#     ι::VectorField{N,1,Float64}
#     D::VectorField{N,1,Float64}
# end
#
# function Base.getproperty(hmf::HelmholtzMatterFields, sym::Symbol)
#     if sym==:P
#         return complex.(getfield(hmf,:ρ),getfield(hmf,:ι))
#     else
#         return getfield(hmf,sym)
#     end
# end
#
# function Base.propertynames(::HelmholtzMatterFields, private=false)
#     if private
#         return (:P, fieldnames(HelmholtzMatterFields)...)
#     else
#         return (:P,:D)
#     end
# end

# struct HelmholtzBlochFields{N} <: AbstractFields{N}
#     wave::HelmholtzWaveFields{N}
#     matter::HelmholtzMatterFields{N}
# end
#
# function Base.getproperty(hbf::HelmholtzBlochFields, sym::Symbol)
#     if Base.sym_in(sym, propertynames(getfield(hbf,:matter)), true)
#         return getproperty(getfield(hbf,:matter),sym)
#     elseif Base.sym_in(sym, propertynames(getfield(hbf,:wave)))
#         return getproperty(getfield(hbf,:wave),sym)
#     else
#         return getfield(hbf,sym)
#     end
# end
#
# function Base.propertynames(hbf::HelmholtzBlochFields, private=false)
#     if private
#         return (fieldnames(HelmholtzBlochField)..., propertynames(hbf.wave)..., propertynames(hbf.matter))
#     else
#         return (propertynames(hbf.wave)..., propertynames(hbf.matter))
#     end
# end

################################################################################
# HelmholtzFDTD

struct AnimationOptions{TB,TO}
    animation::Animation
    interval::Int
    start::Int
    by::TB
    ylims::Vector{Float64}
    plotoptions::TO
end

function AnimationOptions(;
            interval::Integer=100,
            start::Integer=1,
            by::Function=real,
            ylims::Vector=Float64[0,0],
            plotoptions...)
    AnimationOptions(Animation(), interval, start, by, ylims, plotoptions)
end

struct HelmholtzFDTD{N,M,P,TS,TSIM} <: AbstractFDTD{N,M,P} # N is dimension, M is memory depth
    fields::HelmholtzWaveFields{N,M}
    grad::Gradient{N}
    div::Divergence{N}
    n::Vector{Int}
    t::Vector{Float64}
    dt::Float64
    source::TS
    α::Vector{Vector{Float64}}
    β::Vector{Vector{Float64}}
    pmldims::NTuple{P,Int}
    options::AnimationOptions
    simulation::TSIM
end

function HelmholtzFDTD(
            sim::Simulation{N,Common.Symmetric},
            sc::TS,
            dt::Real;
            kwargs...
            ) where {N,TS}

    fields = HelmholtzWaveFields(sim)
    grad = Gradient(sim)
    div = Divergence(sim)
    n = Int[0]
    t = Float64[dt/2]
    α, β = _αβ(sim, dt)
    pmldims = _whichdimspml(sim)
    return HelmholtzFDTD{N,_memorydepth(fields),_ndimpml(sim),TS,typeof(sim)}(fields, grad, div, n, t, dt, sc, α, β, pmldims, AnimationOptions(kwargs...), sim)
end

function Base.getproperty(fdtd::HelmholtzFDTD, sym::Symbol)
    if Base.sym_in(sym,(:φ,:ϕ,:phi,:ψ,:psi))
        return getproperty(getfield(fdtd,:fields),:φ)
    else
        return getfield(fdtd,sym)
    end
end

function Base.propertynames(::HelmholtzFDTD, private=false)
    if private
        return (:φ, fieldnames(HelmholtzFDTD)...)
    else
        return (:φ, :dt, :n, :options, :t)
    end
end

struct HelmholtzPointSource{FX,FW,FA,FP}
    xoft::FX    # position of time
    aoft::FA    # amplitude of time
    ωoft::FW    # frequency of time
    ϕoft::FP    # phase of time
    σ²::Float64
    N::Float64

    function HelmholtzPointSource(xoft, aoft, ωoft, ϕoft, σ²::Real, N::Real)
        typeof(xoft)<:Real ? xoft = Constant(xoft) : nothing
        typeof(aoft)<:Real ? aoft = Constant(aoft) : nothing
        typeof(ωoft)<:Real ? ωoft = Constant(ωoft) : nothing
        typeof(ϕoft)<:Real ? ϕoft = Constant(ϕoft) : nothing
        return new{typeof(xoft),typeof(aoft),typeof(ωoft),typeof(ϕoft)}(xoft,aoft,ωoft,ϕoft,σ²,N)
    end
end

function (ps::HelmholtzPointSource)(x::Point,t::Real)
    x0 = ps.xoft(t)
    typeof(x0)<:Point ? nothing : x0 = Point(x0...)
    ω = ps.ωoft(t)
    a = ps.aoft(t)
    ϕ = ps.ϕoft(t)
    r² = (x-x0).r^2
    return -a*sin(ϕ + ω*t) * exp(-r²/2ps.σ²)/ps.N/ω
end

function Base.show(io::IO, mime::MIME"text/plain", ps::HelmholtzPointSource)
    printstyled(io,"HelmholtzPointSource: "; color=PRINTED_COLOR_DARK)
    println(io)
    printstyled(io,"\tx(t)\t"; color=PRINTED_COLOR_VARIABLE)
    show(io,mime,ps.xoft)
    printstyled(io,"\n\tω(t)\t"; color=PRINTED_COLOR_VARIABLE)
    show(io,mime,ps.ωoft)
    printstyled(io,"\n\ta(t)\t"; color=PRINTED_COLOR_VARIABLE)
    show(io,mime,ps.aoft)
    printstyled(io,"\n\tϕ(t)\t"; color=PRINTED_COLOR_VARIABLE)
    show(io,mime,ps.ϕoft)
end

struct Constant<:Function c::Float64 end

(C::Constant)(t::Real) = C.c

Base.iszero(ps::HelmholtzPointSource) = iszero(ps.aoft)
Base.iszero(C::Constant) = iszero(C.c)

function Base.show(io::IO, ::MIME"text/plain", C::Constant)
    printstyled(io,C.c; color=PRINTED_COLOR_NUMBER)
    print(io," (constant)")
end

################################################################################
# HelmholtzBlochFDTD
#
# struct HelmholtzBlochFDTD{N,TS} <: AbstractFDTD{N}
#     fields::HelmholtzBlochField{N}
#     differentials::HelmholtzBlochField{N}
#     grad::Gradient{N}
#     div::Divergence{N}
#     n::Vector{Int}
#     t::Vector{Float64}
#     dt::Float64
#     source::TS
#     ρE::Vector{Float64}
#     σEp1⁻¹dt::Vector{Float64}
#     ρG::Vector{Float64}
#     σGp1⁻¹dt::Vector{Float64}
#     ε⁻¹σEp1⁻¹dt::Vector{Float64}
#     Ω::Matrix{Float64}
# end
#
# """
#     HelmholtzBlochFDTD(sim; dt=sim.dx*$DEFAULT_CFL_NUMBER)
# """
# function HelmholtzBlochFDTD(sim::Simulation{N,Common.Symmetric}; dt::Real=sim.dx*DEFAULT_CFL_NUMBER) where N
#     fields = HelmholtzField(sim)
#     differentials = HelmholtzField(sim)
#     grad = Gradient(sim)
#     div = Divergence(sim)
#     n = Int[0]
#     t = Float64[dt/2]
#     sc = 0
#
#     σ = real(sim.σ[1])
#     σEp1⁻¹dt = dt ./(1 .+ σ*dt/2)
#     ρE = (1 .- σ*dt/2).*σEp1⁻¹
#
#     σ = real(sim.σ_half[1])
#     σGp1⁻¹dt = dt ./(1 .+ σ*dt/2)
#     ρG = (1 .- σ*dt/2).*σGp1⁻¹
#
#     return HelmholtzBlochFDTD{Ns,typeof(sc)}(fields, differentials, grad, div, n, t, dt, sc, ρE,
#                 σEp1⁻¹dt, ρG, σGp1⁻¹dt, σEp1⁻¹dt./sim.ε)
# end
#
#
# # @inline
# function propagate!(fdtd::HelmholtzBlochFDTD, n::Integer=1)
#     dt = fdtd.dt
#
#     E = fdtd.fields.E
#     dE = fdtd.differentials.E
#     ρE = fdtd.ρE
#     grad = fdtd.grad.components
#
#     G = fdtd.fields.G
#     dG = fdtd.differentials.G
#     ρG = fdtd.ρG
#     div = fdtd.div
#
#     ε⁻¹σEp1⁻¹dt = fdtd.ε⁻¹σEp1⁻¹dt
#     σGp1⁻¹dt = fdtd.σGp1⁻¹dt
#
#     foreach(1:n) do i
#         # @fastmath
#         mul!(dE, div, G)
#         # @fastmath @inbounds @simd
#         for j ∈ eachindex(dE) E[j] = ρE[j]*E[j] + ε⁻¹σEp1⁻¹dt[j]*dE[j] end
#         # @fastmath @inbounds
#         fdtd.t[1] += dt/2
#
#         # @fastmath
#         mul!(dG, grad, E)
#         # @fastmath @inbounds @simd
#         for j ∈ eachindex(dG) G[j] = ρG[j]*G[j] + σGp1⁻¹dt[j]*dG[j] end
#         # @fastmath @inbounds
#         fdtd.t[1] += dt/2
#     end
#     fdtd.n[1] += n
#     return nothing
# end

################################################################################

Base.eltype(af::AbstractFields) = Float64
Base.length(af::AbstractFields) = length(af.φ)
Base.ndims(af::AbstractFields{N}) where N = N

Base.eltype(fdtd::AbstractFDTD) = Float64
Base.length(fdtd::AbstractFDTD) = length(fdtd.fields)
Base.ndims(fdtd::AbstractFDTD{N}) where N = N

_memorydepth(::AbstractFields{N,M}) where {N,M} = M
_memorydepth(::AbstractFDTD{N,M}) where {N,M} = M

_ndimpml(::AbstractFDTD{N,M,P}) where {N,M,P} = P
_whichdimspml(fdtd::AbstractFDTD) = _whichdimspml(fdtd.simulation)

################################################################################
# Pretty Printing

import ..Common.PRINTED_COLOR_DARK
import ..Common.PRINTED_COLOR_LIGHT
import ..Common.PRINTED_COLOR_NUMBER
import ..Common.PRINTED_COLOR_VARIABLE

function Base.show(io::IO,af::AbstractFields{N,M}) where {N,M}
    print(io,"$(N)D ")
    if typeof(af)<:HelmholtzWaveFields
        printstyled(io,"HelmholtzFields",color=PRINTED_COLOR_DARK)
    elseif typeof(af)<:HelmholtzBlochFields
        printstyled(io,"HelmholtzBlochFields",color=PRINTED_COLOR_DARK)
    end
    print(io," (")
    printstyled(io, length(af), color=PRINTED_COLOR_NUMBER)
    print(io," points, depth ")
    printstyled(io, M, color=PRINTED_COLOR_NUMBER)
    print(io,")")
end

function Base.show(io::IO,fdtd::AbstractFDTD{N,M}) where {N,M}
    print(io,"$(N)D ")
    if typeof(fdtd)<:HelmholtzFDTD
        printstyled(io,"HelmholtzFDTD",color=PRINTED_COLOR_LIGHT)
    end
    print(io," (")
    printstyled(io, length(fdtd), color=PRINTED_COLOR_NUMBER)
    print(io," points, depth ")
    printstyled(io, M, color=PRINTED_COLOR_NUMBER)
    printstyled(io,", step ")
    printstyled(io,"$(fdtd.n[1])", color=PRINTED_COLOR_NUMBER)
    print(io,", time ")
    printstyled(io,"t=",fmt("1.2e",fdtd.t[1]), color=PRINTED_COLOR_NUMBER)
    print(io,")")
end

function Base.show(io::IO, options::AnimationOptions)
    printstyled(io, "AnimationOptions: ", color=PRINTED_COLOR_DARK)
    print(io, "\n\tinterval\t")
    printstyled(io, options.interval; color=PRINTED_COLOR_NUMBER)
    print(io, "\n\tstart   \t")
    printstyled(io, options.start; color=PRINTED_COLOR_NUMBER)
    print(io, "\n\tplot by \t")
    printstyled(io, options.by; color=PRINTED_COLOR_NUMBER)
    print(io, "\n\tylims   \t")
    printstyled(IOContext(io,:type_info=>Vector), options.ylims; color=PRINTED_COLOR_NUMBER)
    print(io, "\n\toptions \t")
    printstyled(io, options.plotoptions.data, color=PRINTED_COLOR_NUMBER)
end

################################################################################
# Plotting

@recipe f(af::AbstractFields; by=abs2) = af, by
@recipe f(by::Function, af::AbstractFields) = af, by

@recipe f(sim::Simulation, af::AbstractFields; by=abs2) = sim, af, by
@recipe f(af::AbstractFields, sim::Simulation; by=abs2) = sim, af, by
@recipe f(sim::Simulation, by::Function, af::AbstractFields) = sim, af, by
@recipe f(by::Function, sim::Simulation, af::AbstractFields) = sim, af, by
@recipe f(by::Function, af::AbstractFields, sim::Simulation) = sim, af, by

@recipe f(fdtd::AbstractFDTD; by=abs2) = fdtd, by
@recipe f(by::Function, fdtd::AbstractFDTD) = fdtd, by

@recipe f(sim::Simulation, fdtd::AbstractFDTD; by=abs2) = simulation, fdtd, by
@recipe f(fdtd::AbstractFDTD, sim::Simulation; by=abs2) = simulation, fdtd, by
@recipe f(sim::Simulation, by::Function, fdtd::AbstractFDTD) = simulation, fdtd, by
@recipe f(by::Function, sim::Simulation, fdtd::AbstractFDTD) = simulation, fdtd, by
@recipe f(by::Function, fdtd::AbstractFDTD, sim::Simulation) = simulation, fdtd, by

"""
    gif(fdtd, [filename; fps])
"""
gif(fdtd::AbstractFDTD; kwargs...) = gif(fdtd.options.animation; kwargs...)

################################################################################

foreach(include,files)

end # module


text: # import Atom


text:
module TestStructuredArrays

using Test, ArrayTools, StructuredArrays

@testset "Utilities        " begin
    dims = (Int8(2), Int16(3), Int32(4), Int64(5), 6)
    @test to_size(dims) === map(Int, dims)
    @test to_size(dims) === map(to_int, dims)
    @test to_size(dims[2]) === (to_int(dims[2]),)
end

@testset "Uniform arrays   " begin
    dims = (Int8(2), Int16(3), Int32(4))
    @test to_size(dims) === map(Int, dims)
    N = length(dims)
    vals = (Float64(2.1), UInt32(7), UInt16(11),
            Float32(-6.2), Float64(pi), Int16(-4))
    for k in 1:6
        x = vals[k]
        T = typeof(x)

        A = (k == 1 ? UniformArray(x, dims) :
             k == 2 ? UniformArray(x, dims...) :
             k == 3 ? UniformArray{T}(x, dims) :
             k == 4 ? UniformArray{T}(x, dims...) :
             k == 5 ? UniformArray{T,N}(x, dims) :
             k == 6 ? UniformArray{T,N}(x, dims...) : break)
        @test eltype(A) === T
        @test ndims(A) == N
        @test size(A) == dims
        @test ntuple(i -> size(A,i), ndims(A)+1) == (size(A)..., 1)
        @test_throws ErrorException size(A,0)
        @test axes(A) === map(Base.OneTo, size(A))
        @test ntuple(i -> axes(A,i), ndims(A)+1) == (axes(A)..., Base.OneTo(1))
        @test_throws ErrorException axes(A,0)
        @test Base.has_offset_axes(A) == false
        @test IndexStyle(A) === IndexLinear()
        @test IndexStyle(typeof(A)) === IndexLinear()
        @test A == fill!(Array{T}(undef, size(A)), A[1])
        @test_throws ErrorException A[1] = zero(T)
        @test_throws BoundsError A[0]

        B = (k == 1 ? MutableUniformArray(x, dims) :
             k == 2 ? MutableUniformArray(x, dims...) :
             k == 3 ? MutableUniformArray{T}(x, dims) :
             k == 4 ? MutableUniformArray{T}(x, dims...) :
             k == 5 ? MutableUniformArray{T,N}(x, dims) :
             k == 6 ? MutableUniformArray{T,N}(x, dims...) : break)
        @test eltype(B) === T
        @test ndims(B) == N
        @test size(B) == dims
        @test ntuple(i -> size(B,i), ndims(B)+1) == (size(B)..., 1)
        @test_throws ErrorException size(B,0)
        @test axes(B) === map(Base.OneTo, size(B))
        @test ntuple(i -> axes(B,i), ndims(B)+1) == (axes(B)..., Base.OneTo(1))
        @test_throws ErrorException axes(B,0)
        @test Base.has_offset_axes(B) == false
        @test IndexStyle(B) === IndexLinear()
        @test IndexStyle(typeof(B)) === IndexLinear()
        @test B == fill!(Array{T}(undef, size(B)), B[1])
        @test_throws ErrorException B[1] = zero(T)
        @test_throws BoundsError B[0]
        x = x
        x -= one(x)
        B[1:end] = x
        @test B[end] == B[1] == x
        x -= one(x)
        B[:] = x
        @test B[end] == B[1] == x
        x -= one(x)
        B[1:length(B)] = x
        @test B[end] == B[1] == x
    end

    # Array with zero dimensions.
    C = UniformArray{Int,0}(17)
    @test C[1] == 17
    @test Base.axes1(C) == 1:1
    @test Base.has_offset_axes(C) == false

    # Using index 1 to set a single-element mutable uniform array is allowed.
    D = MutableUniformArray(11, 1)
    @test D[1] == 11
    D[1] = 6
    @test D[1] == 6

    # Call constructors with illegal dimension.
    @test_throws ErrorException UniformArray{Int}(17, 1, -1)
    @test_throws ErrorException MutableUniformArray{Int}(17, 1, -1)
end

@testset "Structured arrays" begin
    dims = (Int8(3), Int16(4))
    N = length(dims)
    f1(i) = -2i
    f2(i,j) = (-1 ≤ i - j ≤ 2)
    T1 = typeof(f1(1))
    T2 = typeof(f2(1,1))
    S1 = IndexLinear
    S2 = IndexCartesian
    for k in 1:16
        A = (k ==  1 ? StructuredArray(S2, f2, dims) :
             k ==  2 ? StructuredArray(S2(), f2, dims...) :
             k ==  3 ? StructuredArray{T2}(S2(), f2, dims) :
             k ==  4 ? StructuredArray{T2}(S2, f2, dims...) :
             k ==  5 ? StructuredArray{T2,N}(S2, f2, dims) :
             k ==  6 ? StructuredArray{T2,N}(S2(), f2, dims...) :
             k ==  7 ? StructuredArray{T2,N,S2}(f2, dims) :
             k ==  8 ? StructuredArray{T2,N,S2}(f2, dims...) :
             k ==  9 ? StructuredArray{T2,N,S2}(S2, f2, dims) :
             k == 10 ? StructuredArray{T2,N,S2}(S2(), f2, dims...) :
             k == 11 ? StructuredArray(f2, dims) :
             k == 12 ? StructuredArray(f2, dims...) :
             k == 13 ? StructuredArray{T2}(f2, dims) :
             k == 14 ? StructuredArray{T2}(f2, dims...) :
             k == 15 ? StructuredArray{T2,N}(f2, dims) :
             k == 16 ? StructuredArray{T2,N}(f2, dims...) : break)
        @test eltype(A) === T2
        @test ndims(A) == N
        @test size(A) == dims
        @test ntuple(i -> size(A,i), ndims(A)+1) == (size(A)..., 1)
        @test_throws ErrorException size(A,0)
        @test axes(A) === map(Base.OneTo, size(A))
        @test ntuple(i -> axes(A,i), ndims(A)+1) == (axes(A)..., Base.OneTo(1))
        @test_throws ErrorException axes(A,0)
        @test Base.has_offset_axes(A) == false
        @test IndexStyle(A) === IndexCartesian()
        @test IndexStyle(typeof(A)) === IndexCartesian()
        @test A == [f2(i,j) for i in 1:dims[1], j in 1:dims[2]]
        #@test_throws ErrorException A[1,1] = zero(T)
        @test_throws BoundsError A[0,1]
    end
    for k in 1:16
        A = (k ==  1 ? StructuredArray(S1, f1, dims) :
             k ==  2 ? StructuredArray(S1(), f1, dims...) :
             k ==  3 ? StructuredArray{T1}(S1(), f1, dims) :
             k ==  4 ? StructuredArray{T1}(S1, f1, dims...) :
             k ==  5 ? StructuredArray{T1,N}(S1, f1, dims) :
             k ==  6 ? StructuredArray{T1,N}(S1(), f1, dims...) :
             k ==  7 ? StructuredArray{T1,N,S1}(f1, dims) :
             k ==  8 ? StructuredArray{T1,N,S1}(f1, dims...) :
             k ==  9 ? StructuredArray{T1,N,S1}(S1, f1, dims) :
             k == 10 ? StructuredArray{T1,N,S1}(S1(), f1, dims...) : break)
        I = LinearIndices(A)
        @test eltype(A) === T1
        @test ndims(A) == N
        @test size(A) == dims
        @test ntuple(i -> size(A,i), ndims(A)+1) == (size(A)..., 1)
        @test_throws ErrorException size(A,0)
        @test axes(A) === map(Base.OneTo, size(A))
        @test ntuple(i -> axes(A,i), ndims(A)+1) == (axes(A)..., Base.OneTo(1))
        @test_throws ErrorException axes(A,0)
        @test Base.has_offset_axes(A) == false
        @test IndexStyle(A) === IndexLinear()
        @test IndexStyle(typeof(A)) === IndexLinear()
        @test A == [f1(I[CartesianIndex(i,j)])
                    for i in 1:dims[1], j in 1:dims[2]]
        #@test_throws ErrorException A[1,1] = zero(T)
        @test_throws BoundsError A[0,1]
    end

    # Call constructors with illegal dimension.
    @test_throws ErrorException StructuredArray{Bool}((i,j) -> i ≥ j, 1, -1)
end

end # module

nothing

max_stars_repo_name: emmt/UniformArrays

max_issues_repo_name: emmt/UniformArrays


text:
# we generate code in this module, so precompile where possible
__precompile__(true)

module Rotations

using LinearAlgebra
using StaticArrays
using Random

import Statistics

include("util.jl")
include("core_types.jl")
include("unitquaternion.jl")
include("mrps.jl")
include("euler_types.jl")
include("angleaxis_types.jl")
# TODO: Add method `mean_rotation` instead of `mean`
# include("mean.jl")
include("derivatives.jl")
include("principal_value.jl")

include("rodrigues_params.jl")
include("error_maps.jl")
include("rotation_error.jl")
include("log.jl")
include("eigen.jl")
include("deprecated.jl")

export
    Rotation, RotMatrix, RotMatrix2, RotMatrix3,
    Angle2d,
    Quat, UnitQuaternion,
    AngleAxis, RodriguesVec, RotationVec, RodriguesParam, MRP,
    RotX, RotY, RotZ,
    RotXY, RotYX, RotZX, RotXZ, RotYZ, RotZY,
    RotXYX, RotYXY, RotZXZ, RotXZX, RotYZY, RotZYZ,
    RotXYZ, RotYXZ, RotZXY, RotXZY, RotYZX, RotZYX,

    # Quaternion math ops
    logm, expm, ⊖, ⊕,

    # Quaternion maps
    ExponentialMap, QuatVecMap, CayleyMap, MRPMap, IdentityMap,

    # check validity of the rotation (is it close to unitary?)
    isrotation,

    # angle and axis introspection
    rotation_angle,
    rotation_axis,

    # quaternion from two vectors
    rotation_between,

    # principal value of a rotation
    principal_value

    # derivatives (names clash with ForwarDiff?)
    #jacobian, hessian

end # module


text:
# ========== Structs ==========
"""
    ClassData(type, s₁, s₂, s₃)

Based on the input `type`, generates 3 classes of signals with sample sizes
`s₁`, `s₂`, and `s₃` respectively. Accepted input types are:
- `:tri`: Triangular signals of length 32
- `:cbf`: Cylinder-Bell-Funnel signals of length 128

Based on N. Saito and R. Coifman in "Local Discriminant Basis and their Applications" in the
Journal of Mathematical Imaging and Vision, Vol. 5, 337-358 (1995).

**See also:** [`generateclassdata`](@ref)
"""
struct ClassData
    "Signal type, accepted inputs are `:tri` and `:cbf`"
    type::Symbol
    "Sample size for class 1"
    s₁::Int
    "Sample size for class 2"
    s₂::Int
    "Sample size for class 3"
    s₃::Int
    ClassData(type, s₁, s₂, s₃) = type ∈ [:tri, :cbf] ? new(type, s₁, s₂, s₃) :
        throw(ArgumentError("Invalid type. Accepted types are :tri and :cbf only."))
end

# ========== Functions ==========
# ----- Signal Generation -----
# Make a set of circularly shifted and noisy signals of original signal.
"""
    duplicatesignals(x, N, k[, noise=false, t=1])

Given a signal `x`, returns N shifted versions of the signal, each with shifts
of multiples of `k`.

Setting `noise = true` allows randomly generated Gaussian noises of μ = 0,
σ² = t to be added to the circularly shifted signals.

# Arguments
- `x::AbstractVector{T} where T<:Number`: 1D-signal to be duplicated.
- `N::Integer`:: Number of duplicated signals.
- `k::Integer`:: Circular shift size for each duplicated signal.
- `noise::Bool`: (Default: `false`) Whether or not to add Gaussian noise.
- `t::Real`: (Default: 1) Relative size of noise.

# Returns
`::Array{T}`: Duplicated signals.

# Examples
```@repl
using WaveletsExt

x = generatesignals(:blocks)
duplicatesignals(x, 5, 0)      # [x x x x x]
```

*See also:* [`generatesignals`](@ref)
"""
function duplicatesignals(x::AbstractVector{T},
                          N::Integer,
                          k::Integer,
                          noise::Bool = false,
                          t::Real = 1) where T<:Number

    n = length(x)
    X = Array{T, 2}(undef, (n, N))
    @inbounds begin
        for i in axes(X,2)
            X[:,i] = circshift(x, k*(i-1))
        end
    end
    X = noise ? X + t*randn(n,N) : X
    return X
end

# Generate 6 types of signals used in popular papers.
"""
    generatesignals(fn, L)

Generates a signal of length 2ᴸ given the function symbol `fn`. Current accepted inputs
below are based on D. Donoho and I. Johnstone in "Adapting to Unknown Smoothness via Wavelet
Shrinkage" Preprint Stanford, January 93, p 27-28.
- `:blocks`
- `:bumps`
- `:heavisine`
- `:doppler`
- `:quadchirp`
- `:mishmash`

The code for this function is adapted and translated based on MATLAB's Wavelet Toolbox's
`wnoise` function.

# Arguments
- `fn::Symbol`: Type of function/signal to generate.
- `L::Integer`: (Default = 7) Size of the signal to generate. Will return a signal of size
  2ᴸ.

# Returns
`::Vector{Float64}`: Signal of length 2ᴸ.

# Examples
```repl
using WaveletsExt

generatesignals(:bumps, 8)
```
"""
function generatesignals(fn::Symbol, L::Integer = 7)
    @assert L >= 1

    t = [0.1, 0.13, 0.15, 0.23, 0.25, 0.4, 0.44, 0.65, 0.76, 0.78, 0.81]
    h = [4, -5, 3, -4, 5, -4.2, 2.1, 4.3, -3.1, 5.1, -4.2]

    n = 1<<L
    if fn == :blocks
        tt = collect(range(0, 1, length=n))
        x = zeros(n)
        for j in eachindex(h)
            x += (h[j]*(1 .+ sign.(tt .- t[j]))/2)
        end
    elseif fn == :bumps
        h = abs.(h)
        w = 0.01*[0.5, 0.5, 0.6, 1, 1, 3, 1, 1, 0.5, 0.8, 0.5]
        tt = collect(range(0, 1, length=n))
        x = zeros(n)
        for j in eachindex(h)
            x += (h[j] ./ (1 .+ ((tt .- t[j]) / w[j]).^4))
        end
    elseif fn == :heavisine
        x = collect(range(0, 1, length=n))
        x = 4*sin.(4*pi*x) - sign.(x .- 0.3) - sign.(0.72 .- x)
    elseif fn == :doppler
        x = collect(range(0, 1, length=n))
        ϵ = 0.05
        x = sqrt.(x.*(1 .- x)) .* sin.(2*pi*(1+ϵ) ./ (x.+ϵ))
    elseif fn == :quadchirp
        tt = collect(range(0, 1, length=n))
        x = sin.((π/3) * tt .* (n * tt.^2))
    elseif fn == :mishmash
        tt = collect(range(0, 1, length=n))
        x = sin.((π/3) * tt .* (n * tt.^2))
        x = x + sin.(π * (n * 0.6902) * tt)
        x = x + sin.(π * tt .* (n * 0.125 * tt))
    else
        throw(ArgumentError("Unrecognised `fn`. Type `?generatefunction` to learn more."))
    end
    return x
end

# Generate 3 classes of signals used in Saito's LDB paper.
"""
    generateclassdata(c[, shuffle])

Generates 3 classes of data given a `ClassData` struct as an input. Returns a matrix
containing the 3 classes of signals and a vector containing their corresponding labels.

Based on N. Saito and R. Coifman in "Local Discriminant Basis and their Applications" in the
Journal of Mathematical Imaging and Vision, Vol. 5, 337-358 (1995).

# Arguments
- `c::ClassData`: Type of signal classes to generate.
- `shuffle::Bool`: (Default: `true`). Whether or not to shuffle the signals.

# Returns
- `::Matrix{Float64}`: Generated signals.
- `::Vector{Int64}`: Class corresponding to each column of generated signals.

# Examples
```@repl
using WaveletsExt

c = ClassData(:tri, 100, 100, 100)
generateclassdata(c)
```

**See also:** [`ClassData`](@ref)
"""
function generateclassdata(c::ClassData, shuffle::Bool = false)
    @assert c.s₁ >= 0
    @assert c.s₂ >= 0
    @assert c.s₃ >= 0

    if c.type == :tri
        n = 32
        i = collect(1:n)
        # Define classes
        y = Int64.(vcat(ones(Int, c.s₁), 2*ones(Int, c.s₂), 3*ones(Int, c.s₃)))

        # Random value generations
        u₁ = rand(Uniform(0,1), c.s₁)
        u₂ = rand(Uniform(0,1), c.s₂)
        u₃ = rand(Uniform(0,1), c.s₃)
        ϵ₁ = rand(Normal(0,1), (n, c.s₁))
        ϵ₂ = rand(Normal(0,1), (n, c.s₂))
        ϵ₃ = rand(Normal(0,1), (n, c.s₃))

        h₁ = max.(6 .- abs.(i.-7), 0)
        h₂ = max.(6 .- abs.(i.-15), 0)
        h₃ = max.(6 .- abs.(i.-11), 0)
        # Build signals
        H₁ = h₁ * u₁' + h₂ * (1 .- u₁)' + ϵ₁
        H₂ = h₁ * u₂' + h₃ * (1 .- u₂)' + ϵ₂
        H₃ = h₂ * u₃' + h₃ * (1 .- u₃)' + ϵ₃

        H = hcat(H₁, H₂, H₃)
    elseif c.type == :cbf
        n = 128
        ϵ = rand(Normal(0,1), (n, c.s₁+c.s₂+c.s₃))
        y = Int64.(vcat(ones(c.s₁), 2*ones(Int, c.s₂), 3*ones(Int, c.s₃)))

        d₁ = DiscreteUniform(16,32)
        d₂ = DiscreteUniform(32,96)

        # cylinder signals
        H₁ = zeros(n,c.s₁)
        a = rand(d₁,c.s₁)
        b = a+rand(d₁,c.s₁)
        η = randn(c.s₁)
        for k in 1:c.s₁
            H₁[a[k]:b[k],k]=(6+η[k])*ones(b[k]-a[k]+1)
        end

        # bell signals
        H₂ = zeros(n,c.s₂)
        a = rand(d₁,c.s₂)
        b = a+rand(d₂,c.s₂)
        η = randn(c.s₂)
        for k in 1:c.s₂
            H₂[a[k]:b[k],k]=(6+η[k])*collect(0:(b[k]-a[k]))/(b[k]-a[k])
        end

        # funnel signals
        H₃ = zeros(n,c.s₃)
        a = rand(d₁,c.s₃)
        b = a+rand(d₂,c.s₃)
        η = randn(c.s₃)
        for k in 1:c.s₃
            H₃[a[k]:b[k],k]=(6+η[k])*collect((b[k]-a[k]):-1:0)/(b[k]-a[k])
        end

        H = hcat(H₁, H₂, H₃) + ϵ
    else
        throw(ArgumentError("Invalid type. Accepted types are :tri and :cbf only."))
    end

    if shuffle
        idx = [1:(c.s₁+c.s₂+c.s₃)...]
        shuffle!(idx)
        H = H[:,idx]
        y = y[idx]
    end

    return H, y
end


text:
module TestTwoInertiasAndIdealGearWithMonteCarlo

using TinyModia
using DifferentialEquations
using ModiaPlot
using MonteCarloMeasurements
using Distributions
using Unitful

# The number of particles must be the same as for FloatType
const nparticles = 100
uniform(vmin,vmax) = StaticParticles(nparticles,Distributions.Uniform(vmin,vmax))

TwoInertiasAndIdealGearWithMonteCarlo = Model(
    J1 = 0.0025,
    J2 = uniform(50, 170),
    r  = 105,
    tau_max = 1,
    phi2 = Var(init = uniform(0.5,0.55)),
    w2 = Var(init = 0.0),
    equations = :[
        tau = if time < 1u"s"; tau_max elseif time < 2u"s"; 0 elseif time < 3u"s"; -tau_max else 0 end,

        # inertia1
        w1 = der(phi1),
        J1*der(w1) = tau - tau1,

        # gear
        phi1   = r*phi2,
        r*tau1 = tau2,

        # inertia2
        w2 = der(phi2),
        J2*der(w2) = tau2
    ]
)

twoInertiasAndIdealGearWithMonteCarlo = @instantiateModel(TwoInertiasAndIdealGearWithMonteCarlo,
                                                          FloatType = StaticParticles{Float64,nparticles})

simulate!(twoInertiasAndIdealGearWithMonteCarlo, Tsit5(), stopTime = 4.0,
          logParameters=true, logStates=true)

plot(twoInertiasAndIdealGearWithMonteCarlo, ["phi2", "w2"],
     MonteCarloAsArea=false)

end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct StorageAccountListKeysResult <: SwaggerModel
    keys::Any # spec type: Union{ Nothing, Vector{StorageAccountKey} } # spec name: keys

    function StorageAccountListKeysResult(;keys=nothing)
        o = new()
        validate_property(StorageAccountListKeysResult, Symbol("keys"), keys)
        setfield!(o, Symbol("keys"), keys)
        o
    end
end # type StorageAccountListKeysResult

const _property_map_StorageAccountListKeysResult = Dict{Symbol,Symbol}(Symbol("keys")=>Symbol("keys"))
const _property_types_StorageAccountListKeysResult = Dict{Symbol,String}(Symbol("keys")=>"Vector{StorageAccountKey}")
Base.propertynames(::Type{ StorageAccountListKeysResult }) = collect(keys(_property_map_StorageAccountListKeysResult))
Swagger.property_type(::Type{ StorageAccountListKeysResult }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_StorageAccountListKeysResult[name]))}
Swagger.field_name(::Type{ StorageAccountListKeysResult }, property_name::Symbol) =  _property_map_StorageAccountListKeysResult[property_name]

function check_required(o::StorageAccountListKeysResult)
    true
end

function validate_property(::Type{ StorageAccountListKeysResult }, name::Symbol, val)
end


text:
struct BBox{T}
    bmin::SVector{3,T}
    bmax::SVector{3,T}
end

struct Cluster{T}
    level::Int
    npar::Int
    bbox::BBox{T}
    pindex_start::Int
    xcoords::Union{Nothing,Vector{T}}
    ycoords::Union{Nothing,Vector{T}}
    zcoords::Union{Nothing,Vector{T}}
    gamma_hat::Union{Nothing,Array{T,3}}
    mom_hat::Union{Nothing,Array{SVector{3,T},3}}
    children::Union{Nothing,Tuple{Cluster{T},Cluster{T}}}
end

function make_cluster(particles::Particles{T}, parindices::Vector{Int}, lo, hi, level; n, threshold, stretch=SVector(1.0,1.0,1.0)) where {T}
    npar = (hi - lo + 1)
    pos = particles.positions
    bbox = find_bbox(pos, parindices, lo, hi)
    if npar <= threshold
        return Cluster(level, npar, bbox, lo, nothing, nothing, nothing, nothing, nothing, nothing)
    else
        split_dim = argmax( stretch .* (bbox.bmax - bbox.bmin) )
        k = split_median!(parindices, pos, split_dim, lo, hi)
        children = (make_cluster(particles, parindices, lo, k, level + 1; n=n, threshold=threshold, stretch=stretch), make_cluster(particles, parindices, k + 1, hi, level + 1; n=n, threshold=threshold, stretch=stretch))
        xcoords, ycoords, zcoords = cluster_coord(bbox; n=n)
        gamma_hat, mom_hat = cluster_weight(particles, parindices, lo, hi, bbox; n=n)
        return Cluster(level, npar, bbox, lo, xcoords, ycoords, zcoords, gamma_hat, mom_hat, children)
    end
end

function find_bbox(pos::AbstractVector{SVector{3,T}}, parindices::Vector{Int}, lo, hi) where {T}
    tmin = typemin(T)
    tmax = typemax(T)
    xmin, xmax = tmax, tmin
    ymin, ymax = tmax, tmin
    zmin, zmax = tmax, tmin
    for i in lo:hi
        (x, y, z) = pos[parindices[i]]
        xmin = x < xmin ? x : xmin
        xmax = x > xmax ? x : xmax
        ymin = y < ymin ? y : ymin
        ymax = y > ymax ? y : ymax
        zmin = z < zmin ? z : zmin
        zmax = z > zmax ? z : zmax
    end
    bmin = SVector(xmin, ymin, zmin)
    bmax = SVector(xmax, ymax, zmax)
    return BBox(bmin, bmax)
end

function cluster_coord(bbox::BBox{T}; n) where {T}
    # transform chebpts to [0,1]
    chebpts_shifted = [ (cos(pi * i / n) + 1.0) / 2.0 for i in 0:n]
    xmin, ymin, zmin = bbox.bmin
    xmax, ymax, zmax = bbox.bmax
    xcoords = chebpts_shifted .* (xmax - xmin) .+ xmin
    ycoords = chebpts_shifted .* (ymax - ymin) .+ ymin
    zcoords = chebpts_shifted .* (zmax - zmin) .+ zmin
    return xcoords, ycoords, zcoords
end

function cluster_weight(particles::Particles, parindices::Vector{Int}, lo, hi, bbox::BBox{T}; n) where {T}
    pos = particles.positions
    mom = particles.momentums
    bmin = bbox.bmin
    bmax = bbox.bmax
    chebpts = [cos(pi * i / n) for i in 0:n]
    w = [(-1.0)^(j) for j in 0:n]
    w[1] = 0.5
    w[n + 1] = 0.5 * (-1.0)^n
    diffx = zeros(n + 1)
    diffy = zeros(n + 1)
    diffz = zeros(n + 1)
    ax = zeros(n + 1)
    ay = zeros(n + 1)
    az = zeros(n + 1)
    sum_x = 0.0
    sum_y = 0.0
    sum_z = 0.0
    gamma_hat = zeros(n + 1, n + 1, n + 1)
    mom_hat  = fill(SVector(0.0, 0.0, 0.0), (n + 1, n + 1, n + 1))
    for i = lo:hi
        pindex = parindices[i]
        x, y, z = (pos[pindex] - bmin) ./ (bmax - bmin) * 2 .- 1.0
        diffx .= (x .- chebpts)
        diffy .= (y .- chebpts)
        diffz .= (z .- chebpts)
        indx = findfirst(x -> abs(x) < eps(0.0), diffx)
        indy = findfirst(x -> abs(x) < eps(0.0), diffy)
        indz = findfirst(x -> abs(x) < eps(0.0), diffz)
        if indx === nothing
            ax .= w ./ diffx
            sum_x = sum(ax)
        else
            ax .= 0.0
            ax[indx] = 1.0
            sum_x = 1.0
        end
        if indy === nothing
            ay .= w ./ diffy
            sum_y = sum(ay)
        else
            ay .= 0.0
            ay[indy] = 1.0
            sum_y = 1.0
        end
        if indz === nothing
            az .= w ./ diffz
            sum_z = sum(az)
        else
            az .= 0.0
            az[indz] = 1.0
            sum_z = 1.0
        end
        p = mom[pindex]
        gamma = sqrt(1.0 + dot(p, p))
        for k in 1:(n + 1)
            for j in 1:(n + 1)
                for i in 1:(n + 1)
                    weight = ax[i] * ay[j] * az[k] / (sum_x * sum_y * sum_z)
                    gamma_hat[i,j,k] += gamma * weight
                    mom_hat[i,j,k] += p * weight
                end
            end
        end
    end
    return gamma_hat, mom_hat
end

function print_bbox(cluster::Cluster{T}) where {T}
    bbox = cluter.bbox
    level = cluster.level
    xmin, ymin, zmin = bbox.bmin
    xmax, ymax, zmax = bbox.bmax
    npar = cluster.npar
    println("$(xmin) $(ymin) $(zmin) $(xmax) $(ymax) $(zmax) $(level) $(npar)")
    children = cluster.children
    if children !== nothing
        print_bbox(children[1])
        print_bbox(children[2])
    end
end

macro swap!(a, b)
    return quote
        $(esc(a)), $(esc(b)) = $(esc(b)), $(esc(a))
    end
end

function partition!(parindices::Vector{Int}, pos::AbstractVector{SVector{3,T}}, dim, lo, hi, pindex) where {T}
pval = pos[parindices[pindex]][dim]
    @swap!(parindices[pindex], parindices[hi])
    j = lo
    for i in lo:(hi - 1)
        if pos[parindices[i]][dim] < pval
            @swap!(parindices[j], parindices[i])
            j = j + 1
        end
    end
    @swap!(parindices[hi], parindices[j])
    return j
end

function split_median!(parindices::Vector{Int}, pos::AbstractVector{SVector{3,T}}, dim, lo, hi) where {T}
    k = floor(Int, (lo + hi) / 2.0)
    while true
        if lo == hi
            return lo
        end
        pindex = rand(lo:hi) # floor(Int, (lo + hi) / 2.0)
        pindex = partition!(parindices, pos, dim, lo, hi, pindex)
        if pindex == k
            return k
        elseif pindex > k
            hi = pindex - 1
        else
            lo = pindex + 1
        end
    end
end

function num_cluster(c::Cluster{T}) where {T}
    children = c.children
    if (children === nothing)
        return 1
    else
        return 1 + num_cluster(children[1]) + num_cluster(children[2])
    end
end


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

#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=64, pc=1, wd=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 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]
        I_λ = P*zarr  # complex valued ifft of zarr
        #im_rdc_0_1[:,:,f] = I_lambda
        zarr[f_i] .= 0   # reset zarr for next loop

        #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(ψ_λ)))
        dSdp[:,:,f] = 2 .* real.(conv(I_λ, ψ_λ))

    end
    return dSdp

end

function wst_S20_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
    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,ψ_λ))
            dS20dp[:,:,f1,f2] += I1dI2
            dS20dp[:,:,f2,f1] += I1dI2
        end
    end

    return dS20dp

end

function derivtest2(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

derivtest2(8)

Nx=128
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_S1_deriv(im, fhash)

# S1 deriv time, Jan 30
# 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)

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[:])

    # 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

# wst_S1_deriv agrees with brute force at 1e-10 level.
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)

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

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

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

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)

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)

#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:256,1:256]

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) #Using the local DHS here
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

#NEW
dust = Float64.(readdust())
dust = dust[1:16,1:16]
Nx    = 16
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) #Using the local DHS here
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)

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
module TestLinear
using Test
using LinearAlgebra
using GaussianDistributions
using GaussianDistributions: input_transformation, shift, scale, integral, gradient, value_gradient
using CovarianceFunctions
using ForwardDiff

function algebra_test(G, a, b)
    μ = mean(G)
    H = G + b
    @test mean(H) ≈ μ + b
    @test cov(H) ≈ cov(G)

    H = b + G
    @test mean(H) ≈ μ + b
    @test cov(H) ≈ cov(G)

    H = G - b
    @test mean(H) ≈ μ - b
    @test cov(H) ≈ cov(G)

    H = b - G
    @test mean(H) ≈ b - μ
    @test cov(H) ≈ cov(G)

    H = a*G
    @test mean(H) ≈ a*μ
    @test cov(H) ≈ a*cov(G)*a

    H = G*a
    @test mean(H) ≈ a*μ
    @test cov(H) ≈ a*cov(G)*a

    H = G / a
    @test mean(H) ≈ μ / a
    @test cov(H) ≈ cov(G)/a^2

    H = a \ G
    @test mean(H) ≈ μ / a
    @test cov(H) ≈ cov(G)/a^2
end

# testing linear functionals of Gaussian distributions
@testset "linear transformations" begin
    @testset "univariate" begin
        μ, σ = randn(2)
        G = Gaussian(μ, σ^2)

        a, b = randn(2)
        algebra_test(G, a, b)
    end

    @testset "multivariate" begin
        n = 3
        μ = randn(n)
        Σ = randn(n, n)
        Σ = Σ'Σ
        G = Gaussian(μ, Σ)

        a = randn()
        b = randn(n)
        algebra_test(G, a, b)

        A = randn(2, n) # linear transformation
        H = A*G
        @test mean(H) ≈ A*μ
        @test cov(H) ≈ A*cov(G)*A'

        w = randn(n)
        @test dot(w, G) isa Univariate
        @test mean(dot(w, G)) ≈ dot(w, mean(G))
        @test cov(dot(w, G)) ≈ *(w', cov(G), w)

        @test sum(G) isa Univariate
        @test diff(G) isa Multivariate
        @test length(diff(G)) == length(G)-1
    end

    @testset "process" begin
        l = .5
        k = (x, y) -> exp(-sum(abs2, x-y)/(2l^2))
        G = Gaussian(sin, k)
        # testing scaling of GP
        a, b = randn(2)
        H = a*G
        @test H isa Gaussian
        x, y = randn(2)
        @test mean(H)(x) ≈ a * sin(x)
        @test cov(H)(x, y) ≈ a^2 * k(x, y)

        # testing division by scalar
        H = G / b
        @test H isa Gaussian
        @test mean(H)(x) ≈ sin(x) / b
        @test cov(H)(x, y) ≈ k(x, y) / b^2

        # testing addition to GP
        H = G + b
        @test H isa Gaussian
        @test mean(H)(x) ≈ sin(x) + b
        @test cov(H)(x, y) ≈ k(x, y)
    end

    @testset "input transformation" begin
        l = .5
        k = (x, y) -> exp(-sum(abs2, x-y)/(2l^2))
        G = Gaussian(sin, k)
        x, y = randn(2)

        # general input transformation
        f(x) = exp(-x^2)
        H = input_transformation(G, f)
        @test H(x) isa Gaussian
        @test mean(H(x)) ≈ mean(G(f(x)))
        @test cov(H)(x, y) ≈ cov(G)(f(x), f(y))

        # composition syntax
        H = G ∘ f
        @test H(x) isa Gaussian
        @test mean(H(x)) ≈ mean(G(f(x)))
        @test cov(H)(x, y) ≈ cov(G)(f(x), f(y))

        # shifting and scaling
        a, b = randn(2)
        H = shift(G, b)
        @test mean(H(x)) ≈ mean(G(x+b))
        @test cov(H)(x, y) ≈ cov(G)(x+b, y+b)

        H = scale(G, a)
        @test mean(H(x)) ≈ mean(G(a*x))
        @test cov(H)(x, y) ≈ cov(G)(a*x, a*y)
    end

    @testset "integral" begin
        l = .5
        k = (x, y) -> exp(-sum(abs2, x-y)/(2l^2))
        # testing 1d integral
        G = Gaussian(k)
        S = integral(G, 0, 2π)
        @test S isa Gaussian
        @test mean(S) == 0
        @test std(S) < 1

        # testing integral on conditional process
        n = 64 # enough points to determine process to high accuracy
        x = 2π * rand(n)
        σ = 1e-2
        y = @. sin(x) + σ*randn()
        C = G | (x, y, σ^2)
        S = integral(C, 0, 2π)
        @test isapprox(mean(S), 0, atol = σ)
        @test std(S) < 2σ
    end

    @testset "differential" begin
        k = CovarianceFunctions.RQ(1.)
        G = Gaussian(k)
        f = x->sum(sin, x)
        d, n = 3, 5
        X = randn(d, n)
        x = [c for c in eachcol(X)]
        y = f.(x)
        dy = (z->ForwardDiff.gradient(f, z)).(x)
        C = G | (x, gradient, dy, d)
        @test C isa Gaussian
        tol = 1e-4
        grad_cov = (x, y)->ForwardDiff.jacobian(z2->ForwardDiff.gradient(z1->C.Σ(z1, z2), x), y)
        for i in eachindex(x)
            @test isapprox(ForwardDiff.gradient(C.μ, x[i]), dy[i], atol = tol)
            @test C.Σ(x[i], x[i]) > tol # covariance in value does not collapse without value observations
            GC = grad_cov(x[i], x[i])
            @test norm(GC) < d^2*tol # gradient uncertainty collapses
        end

        ydy = [vcat(yi, dyi) for (yi, dyi) in zip(y, dy)]
        C = G | (x, value_gradient, ydy, d)
        @test C isa Gaussian
        tol = 1e-4
        grad_cov = (x, y)->ForwardDiff.jacobian(z2->ForwardDiff.gradient(z1->C.Σ(z1, z2), x), y)
        for i in eachindex(x) # these once failed, possibly due to ill conditioning
            @test isapprox(value_gradient(C.μ, x[i]), ydy[i], atol = tol)
            @test C.Σ(x[i], x[i]) < tol # covariance in value does collapse with value observations
            GC = grad_cov(x[i], x[i])
            @test norm(GC) < d^2*tol # gradient uncertainty collapses
        end

    end
end

end


text:
using Trilinos
using Test
using MPI
using SparseArrays

if Test.get_testset_depth() == 0
  MPI.Init()
end

comm = Teuchos.MpiComm(MPI.CComm(MPI.COMM_WORLD))
const n = 20
rowmap = Tpetra.Map(n, 0, comm)
A = Tpetra.CrsMatrix(spdiagm(0 => ones(n)),rowmap)

x = Tpetra.Vector(Tpetra.getDomainMap(A))
b = Tpetra.Vector(Tpetra.getRangeMap(A))
Tpetra.randomize(x)
Tpetra.randomize(b)

@show const OpT = supertype(supertype(supertype(typeof(A[]))))
@show const MVT = supertype(supertype(typeof(b[])))

solver_factory = Belos.SolverFactory{Float64, MVT, OpT}()

solparams = Teuchos.ParameterList()
solparams["Verbosity"] = Belos.StatusTestDetails + Belos.FinalSummary + Belos.TimingDetails
solver = Belos.create(solver_factory, "Block GMRES", solparams)
linprob = Belos.LinearProblem(A)

prec_factory = Ifpack2.Factory()
@show M = Ifpack2.create(prec_factory, "ILUT", A)
prec_params = Teuchos.ParameterList()
Ifpack2.setParameters(M, prec_params)
display(prec_params[])
Ifpack2.initialize(M)
Ifpack2.compute(M)

Belos.setRightPrec(linprob, M)
Belos.setProblem(linprob,x,b)
Belos.setProblem(solver, linprob)

@test Belos.solve(solver) == Belos.Converged
for (bi,xi) in zip(Tpetra.device_view(b),Tpetra.device_view(x))
  @test bi ≈ xi
end

println("Supported Belos solvers:")
foreach(println, Belos.supportedSolverNames(solver_factory))

if Test.get_testset_depth() == 0
  MPI.Finalize()
end


text:
module RedirectRequest

import ..Layer, ..request
using URIs
using ..Messages
using ..Pairs: setkv
import ..Header
import ..@debug, ..DEBUG_LEVEL

"""
    request(RedirectLayer, method, ::URI, headers, body) -> HTTP.Response

Redirects the request in the case of 3xx response status.
"""
abstract type RedirectLayer{Next <: Layer} <: Layer{Next} end
export RedirectLayer

function request(::Type{RedirectLayer{Next}},
                 method::String, url::URI, headers, body;
                 redirect_limit=3, forwardheaders=true, kw...) where Next
    count = 0
    while true

        # Verify the url before making the request. Verification is done in
        # the redirect loop to also catch bad redirect URLs.
        verify_url(url)

        res = request(Next, method, url, headers, body; reached_redirect_limit=(count == redirect_limit), kw...)

        if (count == redirect_limit
        ||  !isredirect(res)
        ||  (location = header(res, "Location")) == "")
            return res
        end

        kw = merge(merge(NamedTuple(), kw), (parent = res,))
        oldurl = url
        url = resolvereference(url, location)
        if forwardheaders
            headers = filter(headers) do h
                # false return values are filtered out
                header, value = h
                if header == "Host"
                    return false
                elseif (header in SENSITIVE_HEADERS
                    && !isdomainorsubdomain(url.host, oldurl.host))
                    return false
                else
                    return true
                end
            end
        else
            headers = Header[]
        end

        @debug 1 "➡  Redirect: $url"

        count += 1
    end

    @assert false "Unreachable!"
end

const SENSITIVE_HEADERS = Set([
    "Authorization",
    "Www-Authenticate",
    "Cookie",
    "Cookie2"
])

function isdomainorsubdomain(sub, parent)
    sub == parent && return true
    endswith(sub, parent) || return false
    return sub[length(sub)-length(parent)] == '.'
end

function verify_url(url::URI)
    if !(url.scheme in ("http", "https", "ws", "wss"))
        throw(ArgumentError("missing or unsupported scheme in URL (expected http(s) or ws(s)): $(url)"))
    end
    if isempty(url.host)
        throw(ArgumentError("missing host in URL: $(url)"))
    end
end

end # module RedirectRequest


text:
"""
    square-free decomposition of p
    Yun's algorithm

    based on Algorithm 6.2 in "Computer Algebra, Concepts and Techniques" by Edmund A. Lamangna
"""
function decompose(p::AbstractPolynomial)
    deg(p) < 1 && return FactoredPoly(false, p)
    p₀ = copy(p)
    p = rationalize(p)
    f = FactoredPoly()
    p′ = derivative(p)
    g = gcd(p, p′)
    r = p ÷ g
    i = 1
    x = var(p)

    while !isone(g) && deg(r, x) > 0
        s = gcd(g, r)
        p = r ÷ s
        if !isone(p)
            add_factor!(f, p / leading(p), i)
        end
        i += 1
        r = s
        g = g ÷ s
    end

    if !isone(r)
        add_factor!(f, r / leading(r), i)
    end

    c = p₀ ÷ poly(f)
    if !isone(c)
        add_factor!(f, c, 1)
    end

    return unrationalize(f)
end

decompose(eq) = wrap(decompose, eq)

function decompose2(p::AbstractPolynomial)
    deg(p) < 1 && return FactoredPoly(false, p)
    p₀ = copy(p)
    f = FactoredPoly()
    p′ = derivative(p)
    g = gcd(p, p′)
    r = integer_poly(p ÷ g)
    i = 1
    x = var(p)

    while !isone(g) && deg(r, x) > 0
        s = gcd(g, r)
        p = integer_poly(r ÷ s)
        if !isone(p)
            add_factor!(f, p ÷ leading(p), i)
        end
        i += 1
        r = s
        g = integer_poly(g ÷ s)
    end

    if !isone(r)
        add_factor!(f, r ÷ leading(r), i)
    end

    c = integer_poly(p₀ ÷ poly(f))
    if !isone(c)
        add_factor!(f, c, 1)
    end

    return unrationalize(f)
end

##################### Helper functions for Factorization #####################

# Lagrange interpolation algorithm for rational polynomials
function interp_poly(x, xs, ys)
    length(xs) != length(ys) && error("interp_poly needs the same number of x and y values")
    n = length(xs)
    Ps = 0
    for i = 1:n
        A = *([(x-xs[j]) for j=1:n if j!=i]...)
        B = *([(xs[i]-xs[j]) for j=1:n if j!=i]...)
        Ps += Rational{BigInt}(ys[i]) *  A / B
    end
    return Ps
end

function list_divisors(x)
    f = factor(x)
    l = [1]

    for (p, k) in f
        r = Int[]
        for i = 0:k
            append!(r, p^i .* l)
        end
        l = r
    end
    l
end

# note k is 1-based
function generate_cross_lists(D)
    n = *(length.(D)...)
    L = zeros(Int, (length(D), n))

    for k = 1:n
        j = k - 1
        for i = 1:length(D)
            n = length(D[i])
            L[i,k] = D[i][(j % n) + 1]
            j = j ÷ n
        end
    end
    return L
end

function assemble_factors(f, r, l)
    if !isone(r)
        add_factor!(f, r)
    end
    if !isone(l)
        add_factor!(f, l)
    end
    unrationalize(f)
end

########################## Main factorization algorithms #####################

"""
    Schubert-Kronecker Algorithm

    based on Algorithm 6.1 in "Computer Algebra, Concepts and Techniques" by Edmund A. Lamangna
"""
function factor_monic_sk(p::AbstractPolynomial)
    p = rationalize(p)
    r = copy(p)
    l = coef(p)
    p = integer_poly(p)
    x = var(p)
    f = FactoredPoly()
    A = Int[]
    D = Array{Int}[]

    i = 0
    while length(D) <= deg(r, x)÷2
        a = ((i+1)÷2) * (isodd(i) ? -1 : +1)
        u = convert(Int, p(a))

        if iszero(u)
            while iszero(r % (x-a))
                r = r ÷ (x-a)
                add_factor!(f, x-a)
            end
        else
            push!(A, a)
            divs = list_divisors(abs(u))
            push!(D, [divs; -divs])
        end
        i += 1
    end

    for d = 1:deg(r, x)÷2
        L = generate_cross_lists(D[1:d+1])
        xs = A[1:d+1]
        for i = 1:size(L,2)
            ys = L[:,i]
            q = interp_poly(x, xs, ys)

            while deg(q, x) > 0 && iszero(r % q)
                r = r ÷ q
                add_factor!(f, q)
                if deg(r) < 2d
                    return assemble_factors(f, r, l)
                end
            end
        end
        if deg(r, x) < 2*(d+1)
            return assemble_factors(f, r, l)
        end
    end
    assemble_factors(f, r, l)
end

function factor_schubert_kronecker(p::AbstractPolynomial)
    pᵢ = integer_poly(p)
    pₘ, c = to_monic(pᵢ)
    x = var(p)

    f = FactoredPoly()
    p₁ = one(p)

    for w in factors(decompose(pₘ))
        v, k = first(w), last(w)
        if deg(v, x) > 0
            for u in factors(factor_monic_sk(v))
                w₁ = prim(from_monic(first(u), c))
                add_factor!(f, w₁, k)
                p₁ = p₁ * w₁^k
            end
        end
    end

    q, r = divrem(p, p₁)

    !iszero(r) && @error "incorrect factorization"

    if !isone(q)
        add_factor!(f, q, 1)
    end

    return unrationalize(f)
end

factor_schubert_kronecker(eq) = wrap(factor, eq)

##############################################################################

"""
    Rational Fraction Decompositionusing the Hermite's method

    based on part of Algorithm 7.1 in "Computer Algebra, Concepts and Techniques" by Edmund A. Lamangna
"""
function factor_rational(n::AbstractPolynomialLike, d::AbstractPolynomialLike)
    !isequal(var(n), var(d)) && error("the numerator and denominator should have the same main variable")

    f = FactoredPoly(true)

    n = rationalize(n)
    d = rationalize(d)
    q, r = divrem(n, d)

    if !iszero(q)
        add_factor!(f, q)
    end

    for w in factors(decompose(d))
        v, k = rationalize(first(w)), last(w)
        d₁ = v ^ k
        d = d ÷ d₁
        g, s, t = gcdx(d₁, d)

        num = rationalize((t*r) % d₁)
        denom = v

        if !iszero(num)
            c = cont(num) // cont(denom)
            add_factor!(f, c * prim(num) / prim(denom), k)
        end
        r = s * r
    end

    f
end

factor_rational(r::RationalPoly) = factor_rational(numerator(r) / denominator(r))
factor_rational(p, q) = wrap(factor, p, q)


text:
# Run code for previous steps with plotting turned off.
make_plots_orig_4 = isdefined(Main,:make_plots) ? make_plots : true
 make_plots = false
 include("neid_solar_4_extract_chunks.jl")
 make_plots = make_plots_orig_4
 if make_plots
   using Plots
 end

# Set parameters for this analysis
oversample_fac_chunks = 2
 oversample_fac_orders = 1
 num_spectra_to_bin = 1   # 1 results in no binning in time
 idx_chunk_min = 20
 idx_chunk_max = 21
 plt_order = 42
 plt_order_pix = 3301:3800

chunk_grids = map(c->RvSpectML.make_grid_for_chunk(chunk_list_timeseries,c,oversample_factor=oversample_fac_chunks), 1:num_chunks(chunk_list_timeseries) )
 @time spectra_matrix = RvSpectML.pack_chunk_list_timeseries_to_matrix(chunk_list_timeseries,chunk_grids, alg=:Sinc, oversample_factor=oversample_fac_chunks)
 f_mean = calc_mean_spectrum(spectra_matrix.flux,spectra_matrix.var)
 deriv = calc_mean_dfluxdlnlambda(spectra_matrix.flux,spectra_matrix.var,spectra_matrix.λ,spectra_matrix.chunk_map)
 plt_times = (chunk_list_timeseries.times .-minimum(chunk_list_timeseries.times)).*24

if 2 <= num_spectra_to_bin <=20
  spectra_binned = RvSpectML.bin_spectra_consecutive(spectra_matrix,num_spectra_to_bin)
  f_mean = calc_mean_spectrum(spectra_binned.flux,spectra_binned.var)
  deriv = calc_mean_dfluxdlnlambda(spectra_binned.flux,spectra_binned.var,spectra_binned.λ,spectra_binned.chunk_map)
  times_binned = RvSpectML.bin_times_consecutive(chunk_list_timeseries.times,num_spectra_to_bin)
  plt_times = (times_binned.-minimum(chunk_list_timeseries.times)).*24
  spectra_matrix = spectra_binned
end
make_plots = true
if make_plots
   idx_chunk_min = 20
   idx_chunk_max = 20
   global idx_plt = (first(spectra_matrix.chunk_map[idx_chunk_min]):last(spectra_matrix.chunk_map[idx_chunk_max]))
   mean_in_plt = mean(f_mean[idx_plt])
   std_in_plt = stdm(f_mean[idx_plt],mean_in_plt)
   plt1 = plot()
   scatter!(plt1,spectra_matrix.λ[idx_plt],(spectra_matrix.flux[idx_plt,:].-mean_in_plt)./std_in_plt,linecolor=:black, markersize=0, color=:black,label=:none)
   plot!(plt1,spectra_matrix.λ[idx_plt],(f_mean[idx_plt].-mean_in_plt)./std_in_plt,linecolor=:red, label="Standardized mean spectrum")
   plot!(plt1,spectra_matrix.λ[idx_plt], deriv[idx_plt]./std(deriv[idx_plt]), label="Standardized deriv", linecolor=:green)
   display(plt1)
   #plot!(spectra_matrix.λ[idx_plt],./std_in_plt,linecolor=:red, label="Standardized mean spectrum")
end

println("Computing RVs using dflux/dlnlambda from chunks.")
 (rvs_1, σ_rvs_1) = RvSpectML.calc_rvs_from_taylor_expansion(spectra_matrix,mean=f_mean,deriv=deriv)
 rms_rvs_1 = std(rvs_1)
 mean_σ_rvs_1 = mean(σ_rvs_1)
 flush(stdout)
 println("# rms(RVs)  = ", rms_rvs_1, "  σ_RVs = ", mean_σ_rvs_1, "   N_obs = ", length(rvs_1) )
 flush(stdout)
 if make_plots
   plt2 = scatter(plt_times, rvs_1, yerr=σ_rvs_1, label="RV chunks", color=:blue)
 end

# Analyze spectra as a set of (large sections of) orders
order_grids = map(c->RvSpectML.make_grid_for_chunk(order_list_timeseries,c,oversample_factor=oversample_fac_orders), 1:num_chunks(order_list_timeseries) )

@time spectral_orders_matrix = RvSpectML.pack_chunk_list_timeseries_to_matrix(order_list_timeseries,order_grids)

if 2 <= num_spectra_to_bin <=20
    spectral_orders_matrix = RvSpectML.bin_spectra_consecutive(spectral_orders_matrix,num_spectra_to_bin)
  end
  f_mean_orders = calc_mean_spectrum(spectral_orders_matrix.flux,spectral_orders_matrix.var)
  deriv_orders = calc_mean_dfluxdlnlambda(spectral_orders_matrix.flux,spectral_orders_matrix.var,spectral_orders_matrix.λ,spectral_orders_matrix.chunk_map)

if make_plots
  plt_order = 40
  plt_order_pix = 3501:3800
  idx_plt = spectral_orders_matrix.chunk_map[plt_order][plt_order_pix]
  plt_λ = RvSpectML.get_λs(order_grids, idx_plt)
     mean_in_plt = mean(f_mean_orders[idx_plt])
     std_in_plt = stdm(f_mean_orders[idx_plt],mean_in_plt)
     local plt = plot()
     plot!(plt,plt_λ,(f_mean_orders[idx_plt].-mean_in_plt)./std_in_plt,label=:none,linecolor=:black)
     plot!(plt,plt_λ, deriv_orders[idx_plt]./std(deriv_orders[idx_plt]),label=:none, linecolor=:green)
     #scatter!(plt,plt_λ,(spectral_orders_matrix.flux[idx_plt,:].-mean_in_plt)./std_in_plt,markersize=1, label=:none)
     xlabel!(plt,"λ (Å)")
     ylabel!(plt,"Mean & Deriv")
     local plt3 = scatter(plt_λ,spectral_orders_matrix.flux[idx_plt,:].-f_mean_orders[idx_plt],markersize=1, label=:none)
     ylabel!(plt3,"Residuals")
     local pltall = plot(plt,plt3,layout=(2,1))
     display(pltall)
end

println("Computing RVs using dflux/dlnlambda from orders.")
  order_rvs = RvSpectML.calc_chunk_rvs_from_taylor_expansion(spectral_orders_matrix,mean=f_mean_orders,deriv=deriv_orders)
  ave_order_rvs = vec( sum(mapreduce(c ->order_rvs[c].rv./order_rvs[c].σ_rv.^2, hcat, 1:length(order_rvs)),dims=2) ./
                     sum(mapreduce(c ->(1.0./order_rvs[c].σ_rv.^2), hcat, 1:length(order_rvs)),dims=2) )
  rms_order_rvs = map(order->std(order.rv), order_rvs)
  mean_order_σrvs = map(order->mean(order.σ_rv), order_rvs)
  flush(stdout)
  println("# rms(ave_order_rvs)/√N = ", std(ave_order_rvs)/sqrt(length(ave_order_rvs)), "  <RMS RVs> = ", mean(rms_order_rvs)#=/sqrt(length(ave_order_rvs))=#, "  <σ_RVs> = ", mean(mean_order_σrvs) )
  if make_plots
    plt4 = histogram(abs.(ave_order_rvs),bins=40,label="|Ave Order RVs|", alpha=0.75)
    histogram!(plt4,rms_order_rvs,bins=40,label="RMS Order RVs", alpha=0.75)
    histogram!(plt4,mean_order_σrvs,bins=40,label="σ Order RVs", alpha=0.75)
    xlabel!("(m/s)")
    ylabel!("Counts")
  end

chunk_rms_cut_off = quantile(rms_order_rvs,0.9)
  idx_good_chunks = findall(x-> x<= chunk_rms_cut_off, rms_order_rvs)
  @assert length(idx_good_chunks)>=1
  ave_good_chunks_rvs = vec( sum(mapreduce(c ->order_rvs[c].rv./order_rvs[c].σ_rv.^2, hcat, idx_good_chunks),dims=2) ./
                      sum(mapreduce(c ->(1.0./order_rvs[c].σ_rv.^2), hcat, idx_good_chunks),dims=2) )
  sigma_good_chunks_rvs = sqrt.(vec( sum(mapreduce(c ->1.0./order_rvs[c].σ_rv.^2, hcat, idx_good_chunks),dims=2) ./
                      sum(mapreduce(c ->(1.0./order_rvs[c].σ_rv.^4), hcat, idx_good_chunks),dims=2) ))

 rms_rvs_ave_good_orders = std(ave_good_chunks_rvs)
 mean_sigma_good_chunks = mean(sigma_good_chunks_rvs)
 flush(stdout)
 println("# rms(RVs_good_orders) = ", rms_rvs_ave_good_orders,  "  <σ_RV good orders> = ", mean_sigma_good_chunks, "   N_obs = ", length(ave_good_chunks_rvs) )
 #map(o->plot!(plt,chunk_list_timeseries.times,order_rvs[o].rv,yerr=order_rvs[o].σ_rv,markersize=1,label=:none),idx_good_chunks )

if make_plots
   plt = plot()
   scatter!(plt,plt_times, rvs_1, yerr=σ_rvs_1, label="Equal weighted chunks", markersize=3, color=:blue, legend=:topleft)
   scatter!(plt,plt_times, ave_order_rvs,label="Ave Order RVs", markersize=3, color=:green)
   scatter!(plt,plt_times, ave_good_chunks_rvs, yerr=sigma_good_chunks_rvs, label="Ave good chunks", markersize=3, color=:red)
   xlabel!("Time (hours)")
   ylabel!("RV (m/s)")
   display(plt)
end


text:
module MultiRegion

export MultiRegionGrid, MultiRegionField
export XPartition, YPartition

using Oceananigans
using Oceananigans.Grids
using Oceananigans.Fields
using Oceananigans.Models
using Oceananigans.Architectures
using Oceananigans.BoundaryConditions
using Oceananigans.Utils
using CUDA
using Adapt
using OffsetArrays

using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
using Oceananigans.Utils: Reference, Iterate

using KernelAbstractions: Event, NoneEvent, @kernel, @index

import Base: show, length, size

import Oceananigans.Utils:
                getdevice,
                switch_device!,
                devices,
                isregional,
                getregion,
                _getregion,
                sync_all_devices!

abstract type AbstractMultiRegionGrid{FT, TX, TY, TZ, Arch} <: AbstractGrid{FT, TX, TY, TZ, Arch} end

abstract type AbstractPartition end

getname(type) = typeof(type).name.wrapper

include("multi_region_utils.jl")
include("x_partitions.jl")
include("y_partitions.jl")
include("multi_region_grid.jl")
include("multi_region_field.jl")
include("multi_region_abstract_operations.jl")
include("multi_region_boundary_conditions.jl")
include("multi_region_reductions.jl")
include("unified_heptadiagonal_iterative_solver.jl")
include("unified_implicit_free_surface_solver.jl")
include("multi_region_models.jl")
include("multi_region_output_writers.jl")

end #module