text:
# Necessary to have initialize set u_modified to false if all don't do anything
# otherwise unnecessary save
INITIALIZE_DEFAULT(cb,u,t,integrator) = u_modified!(integrator, false)
FINALIZE_DEFAULT(cb,u,t,integrator) = nothing

@enum RootfindOpt::Int8 begin
  NoRootFind    = 0
  LeftRootFind  = 1
  RightRootFind = 2
end

function Base.convert(::Type{RootfindOpt}, b::Bool)
  return b ? LeftRootFind : NoRootFind
end

"""
```julia
ContinuousCallback(condition,affect!,affect_neg!;
                   initialize = INITIALIZE_DEFAULT,
                   finalize = FINALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=LeftRootFind,
                   save_positions=(true,true),
                   interp_points=10,
                   abstol=10eps(),reltol=0,repeat_nudge=1//100)
```

```julia
function ContinuousCallback(condition,affect!;
                   initialize = INITIALIZE_DEFAULT,
                   finalize = FINALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=LeftRootFind,
                   save_positions=(true,true),
                   affect_neg! = affect!,
                   interp_points=10,
                   abstol=10eps(),reltol=0,repeat_nudge=1//100)
```

Contains a single callback whose `condition` is a continuous function. The callback is triggered when this function evaluates to 0.

# Arguments
- `condition`: This is a function `condition(u,t,integrator)` for declaring when
  the callback should be used. A callback is initiated if the condition hits
  `0` within the time interval. See the [Integrator Interface](@ref integrator) documentation for information about `integrator`.
- `affect!`: This is the function `affect!(integrator)` where one is allowed to
  modify the current state of the integrator. If you do not pass an `affect_neg!`
  function, it is called when `condition` is found to be `0` (at a root) and
  the cross is either an upcrossing (from negative to positive) or a downcrossing
  (from positive to negative). You need to explicitly pass `nothing` as the
  `affect_neg!` argument if it should only be called at upcrossings, e.g.
  `ContinuousCallback(condition, affect!, nothing)`. For more information on what can
  be done, see the [Integrator Interface](@ref integrator) manual page. Modifications to
  `u` are safe in this function.
- `affect_neg!=affect!`: This is the function `affect_neg!(integrator)` where one is allowed to
  modify the current state of the integrator. This is called when `condition` is
  found to be `0` (at a root) and the cross is an downcrossing (from positive to
  negative). For more information on what can
  be done, see the [Integrator Interface](@ref integrator) manual page. Modifications to
  `u` are safe in this function.
- `rootfind=LeftRootFind`: This is a flag to specify the type of rootfinding to do for finding
  event location. If this is set to `LeftRootfind`, the solution will be backtracked to the point where
  `condition==0` and if the solution isn't exact, the left limit of root is used. If set to
  `RightRootFind`, the solution would be set to the right limit of the root. Otherwise the systems and
  the `affect!` will occur at `t+dt`. Note that these enums are not exported and thus one needs to
  reference them as `SciMLBase.LeftRootFind`, `SciMLBase.RightRootFind`, or `SciMLBase.NoRootFind`.
- `interp_points=10`: The number of interpolated points to check the condition. The
  condition is found by checking whether any interpolation point / endpoint has
  a different sign. If `interp_points=0`, then conditions will only be noticed if
  the sign of `condition` is different at `t` than at `t+dt`. This behavior is not
  robust when the solution is oscillatory, and thus it's recommended that one use
  some interpolation points (they're cheap to compute!).
  `0` within the time interval.
- `save_positions=(true,true)`: Boolean tuple for whether to save before and after the `affect!`.
  This saving will occur just before and after the event, only at event times, and
  does not depend on options like `saveat`, `save_everystep`, etc. (i.e. if
  `saveat=[1.0,2.0,3.0]`, this can still add a save point at `2.1` if true).
  For discontinuous changes like a modification to `u` to be
  handled correctly (without error), one should set `save_positions=(true,true)`.
- `idxs=nothing`: The components which will be interpolated into the condition. Defaults
  to `nothing` which means `u` will be all components.
- `initialize`: This is a function `(c,u,t,integrator)` which can be used to initialize
  the state of the callback `c`. It should modify the argument `c` and the return is
  ignored.
- `finalize`: This is a function `(c,u,t,integrator)` which can be used to finalize
  the state of the callback `c`. It can modify the argument `c` and the return is ignored.
- `abstol=1e-14` & `reltol=0`: These are used to specify a tolerance from zero for the rootfinder:
  if the starting condition is less than the tolerance from zero, then no root will be detected.
  This is to stop repeat events happening just after a previously rootfound event.
- `repeat_nudge = 1//100`: This is used to set the next testing point after a
  previously found zero. Defaults to 1//100, which means after a callback the next
  sign check will take place at t + dt*1//100 instead of at t to avoid repeats.
"""
struct ContinuousCallback{F1,F2,F3,F4,F5,T,T2,T3,I,R} <: AbstractContinuousCallback
  condition::F1
  affect!::F2
  affect_neg!::F3
  initialize::F4
  finalize::F5
  idxs::I
  rootfind::RootfindOpt
  interp_points::Int
  save_positions::BitArray{1}
  dtrelax::R
  abstol::T
  reltol::T2
  repeat_nudge::T3
  ContinuousCallback(condition::F1,affect!::F2,affect_neg!::F3,
                     initialize::F4,finalize::F5,idxs::I,rootfind,
                     interp_points,save_positions,dtrelax::R,abstol::T,reltol::T2,
                     repeat_nudge::T3) where {F1,F2,F3,F4,F5,T,T2,T3,I,R} =
                       new{F1,F2,F3,F4,F5,T,T2,T3,I,R}(condition,
                                               affect!,affect_neg!,
                                               initialize,finalize,idxs,rootfind,interp_points,
                                               BitArray(collect(save_positions)),
                                               dtrelax,abstol,reltol,repeat_nudge)
end

ContinuousCallback(condition,affect!,affect_neg!;
                   initialize = INITIALIZE_DEFAULT,
                   finalize = FINALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=LeftRootFind,
                   save_positions=(true,true),
                   interp_points=10,
                   dtrelax=1,
                   abstol=10eps(),reltol=0,
                   repeat_nudge = 1//100) = ContinuousCallback(
                              condition,affect!,affect_neg!,initialize,finalize,
                              idxs,
                              rootfind,interp_points,
                              save_positions,
                              dtrelax,abstol,reltol,repeat_nudge)

function ContinuousCallback(condition,affect!;
                   initialize = INITIALIZE_DEFAULT,
                   finalize = FINALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=LeftRootFind,
                   save_positions=(true,true),
                   affect_neg! = affect!,
                   interp_points=10,
                   dtrelax=1,
                   abstol=10eps(),reltol=0,repeat_nudge=1//100)

 ContinuousCallback(
            condition,affect!,affect_neg!,initialize,finalize,idxs,
            rootfind,interp_points,
            collect(save_positions),
            dtrelax,abstol,reltol,repeat_nudge)

end

"""
```julia
VectorContinuousCallback(condition,affect!,affect_neg!,len;
                         initialize = INITIALIZE_DEFAULT,
                         finalize = FINALIZE_DEFAULT,
                         idxs = nothing,
                         rootfind=LeftRootFind,
                         save_positions=(true,true),
                         interp_points=10,
                         abstol=10eps(),reltol=0,repeat_nudge = 1//100)
```

```julia
VectorContinuousCallback(condition,affect!,len;
                   initialize = INITIALIZE_DEFAULT,
                   finalize = FINALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=LeftRootFind,
                   save_positions=(true,true),
                   affect_neg! = affect!,
                   interp_points=10,
                   abstol=10eps(),reltol=0,repeat_nudge=1//100)
```

This is also a subtype of `AbstractContinuousCallback`. `CallbackSet` is not feasible when you have a large number of callbacks,
as it doesn't scale well. For this reason, we have `VectorContinuousCallback` - it allows you to have a single callback for
multiple events.

# Arguments

- `condition`: This is a function `condition(out, u, t, integrator)` which should save the condition value in the array `out`
   at the right index. Maximum index of `out` should be specified in the `len` property of callback. So this way you can have
   a chain of `len` events, which would cause the `i`th event to trigger when `out[i] = 0`.
- `affect!`: This is a function `affect!(integrator, event_index)` which lets you modify `integrator` and it tells you about
   which event occured using `event_idx` i.e. gives you index `i` for which `out[i]` came out to be zero.
- `len`: Number of callbacks chained. This is compulsory to be specified.

Rest of the arguments have the same meaning as in [`ContinuousCallback`](@ref).
"""
struct VectorContinuousCallback{F1,F2,F3,F4,F5,T,T2,T3,I,R} <: AbstractContinuousCallback
  condition::F1
  affect!::F2
  affect_neg!::F3
  len::Int
  initialize::F4
  finalize::F5
  idxs::I
  rootfind::RootfindOpt
  interp_points::Int
  save_positions::BitArray{1}
  dtrelax::R
  abstol::T
  reltol::T2
  repeat_nudge::T3
  VectorContinuousCallback(condition::F1,affect!::F2,affect_neg!::F3,len::Int,
                           initialize::F4,finalize::F5,idxs::I,rootfind,
                           interp_points,save_positions,dtrelax::R,
                           abstol::T,reltol::T2,repeat_nudge::T3) where {F1,F2,F3,F4,F5,T,T2,T3,I,R} =
                       new{F1,F2,F3,F4,F5,T,T2,T3,I,R}(condition,
                                               affect!,affect_neg!,len,
                                               initialize,finalize,idxs,rootfind,interp_points,
                                               BitArray(collect(save_positions)),
                                               dtrelax,abstol,reltol,repeat_nudge)
end

VectorContinuousCallback(condition,affect!,affect_neg!,len;
                         initialize = INITIALIZE_DEFAULT,
                         finalize = FINALIZE_DEFAULT,
                         idxs = nothing,
                         rootfind=LeftRootFind,
                         save_positions=(true,true),
                         interp_points=10,
                         dtrelax=1,
                         abstol=10eps(),reltol=0,repeat_nudge=1//100) = VectorContinuousCallback(
                              condition,affect!,affect_neg!,len,
                              initialize,finalize,
                              idxs,
                              rootfind,interp_points,
                              save_positions,dtrelax,
                              abstol,reltol,repeat_nudge)

function VectorContinuousCallback(condition,affect!,len;
                   initialize = INITIALIZE_DEFAULT,
                   finalize = FINALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=LeftRootFind,
                   save_positions=(true,true),
                   affect_neg! = affect!,
                   interp_points=10,
                   dtrelax=1,
                   abstol=10eps(),reltol=0,repeat_nudge=1//100)

 VectorContinuousCallback(
            condition,affect!,affect_neg!,len,initialize,finalize,idxs,
            rootfind,interp_points,
            collect(save_positions),
            dtrelax,abstol,reltol,repeat_nudge)

end

"""
```julia
DiscreteCallback(condition,affect!;
                 initialize = INITIALIZE_DEFAULT,
                 finalize = FINALIZE_DEFAULT,
                 save_positions=(true,true))
```

# Arguments

- `condition`: This is a function `condition(u,t,integrator)` for declaring when
  the callback should be used. A callback is initiated if the condition evaluates
  to `true`. See the [Integrator Interface](@ref integrator) documentation for information about `integrator`.
    - `affect!`: This is the function `affect!(integrator)` where one is allowed to
  modify the current state of the integrator. For more information on what can
  be done, see the [Integrator Interface](@ref integrator) manual page.
- `save_positions`: Boolean tuple for whether to save before and after the `affect!`.
  This saving will occur just before and after the event, only at event times, and
  does not depend on options like `saveat`, `save_everystep`, etc. (i.e. if
  `saveat=[1.0,2.0,3.0]`, this can still add a save point at `2.1` if true).
  For discontinuous changes like a modification to `u` to be
  handled correctly (without error), one should set `save_positions=(true,true)`.
- `initialize`: This is a function `(c,u,t,integrator)` which can be used to initialize
  the state of the callback `c`. It should modify the argument `c` and the return is
  ignored.
- `finalize`: This is a function `(c,u,t,integrator)` which can be used to finalize
  the state of the callback `c`. It should can the argument `c` and the return is
  ignored.
"""
struct DiscreteCallback{F1,F2,F3,F4} <: AbstractDiscreteCallback
  condition::F1
  affect!::F2
  initialize::F3
  finalize::F4
  save_positions::BitArray{1}
  DiscreteCallback(condition::F1,affect!::F2,
                   initialize::F3,finalize::F4,save_positions) where {F1,F2,F3,F4} = new{F1,F2,F3,F4}(condition,
                                                                                affect!,initialize,finalize,
                                                                                BitArray(collect(save_positions)))
end
DiscreteCallback(condition,affect!;
                 initialize = INITIALIZE_DEFAULT, finalize = FINALIZE_DEFAULT,
                 save_positions=(true,true)) = DiscreteCallback(condition,affect!,initialize,finalize,save_positions)

"""
$(TYPEDEF)

Multiple callbacks can be chained together to form a `CallbackSet`. A `CallbackSet`
is constructed by passing the constructor `ContinuousCallback`, `DiscreteCallback`,
`VectorContinuousCallback` or other `CallbackSet` instances:

    CallbackSet(cb1,cb2,cb3)

You can pass as many callbacks as you like. When the solvers encounter multiple
callbacks, the following rules apply:

* `ContinuousCallback`s and `VectorContinuousCallback`s are applied before `DiscreteCallback`s. (This is because
  they often implement event-finding that will backtrack the timestep to smaller
  than `dt`).
* For `ContinuousCallback`s and `VectorContinuousCallback`s, the event times are found by rootfinding and only
  the first `ContinuousCallback` or `VectorContinuousCallback` affect is applied.
* The `DiscreteCallback`s are then applied in order. Note that the ordering only
  matters for the conditions: if a previous callback modifies `u` in such a way
  that the next callback no longer evaluates condition to `true`, its `affect`
  will not be applied.

"""
struct CallbackSet{T1<:Tuple,T2<:Tuple} <: DECallback
  continuous_callbacks::T1
  discrete_callbacks::T2
end

CallbackSet(callback::AbstractDiscreteCallback) = CallbackSet((),(callback,))
CallbackSet(callback::AbstractContinuousCallback) = CallbackSet((callback,),())
CallbackSet() = CallbackSet((),())
CallbackSet(cb::Nothing) = CallbackSet()

# For Varargs, use recursion to make it type-stable
CallbackSet(callbacks::Union{DECallback,Nothing}...) = CallbackSet(split_callbacks((), (), callbacks...)...)

"""
    split_callbacks(cs, ds, args...)

Split comma seperated callbacks into sets of continous and discrete callbacks.
"""
@inline split_callbacks(cs, ds) = cs, ds
@inline split_callbacks(cs, ds, c::Nothing, args...) = split_callbacks(cs, ds, args...)
@inline split_callbacks(cs, ds, c::AbstractContinuousCallback, args...) = split_callbacks((cs..., c), ds, args...)
@inline split_callbacks(cs, ds, d::AbstractDiscreteCallback, args...) = split_callbacks(cs, (ds..., d), args...)
@inline function split_callbacks(cs, ds, d::CallbackSet, args...)
  split_callbacks((cs...,d.continuous_callbacks...), (ds..., d.discrete_callbacks...), args...)
end


text:
include(joinpath("gammaReg", "research_inverse.jl"))
include(joinpath("gammaReg", "research_log.jl"))
include(joinpath("gammaReg", "chosenVariables_log.jl"))
include(joinpath("gammaReg", "chosenVariables_log.jl"))

ext: jl

max_stars_repo_name: AlexLsn/CSO

max_issues_repo_name: AlexLsn/CSO

max_forks_repo_name: AlexLsn/CSO


text:
"""
The Candidate type represents a 5-character string as a tuple of 5 Char's.
"""
struct Candidate
    chars::NTuple{5, Char}
end

"""
Turn a String object into a Candidate object. The input must be a lowercase
string containing only ASCII characters.
"""
function Candidate(s::String)
    if length(s) != 5
        throw(ArgumentError("Candidate must contain exactly 5 characters"))
    end
    if !all(islowercase, collect(s))
        throw(ArgumentError("Candidate must contain only lowercase characters"))
    end
    if !isascii(s)
        throw(ArgumentError("Candidate must contain only ASCII characters"))
    end
    Candidate(ntuple(i -> s[i], 5))
end

"""
Turn a Candidate object into a String object.
"""
String(c::Candidate) = join(c.chars)

"""
Turn a Candidate object into a LetterCount object.
"""
function LetterCounts(candidate::Candidate)
    letter_counts = LetterCounts()
    for char in candidate.chars
        letter_counts[char] += 1
    end
    letter_counts
end

"""
Turn a Candidate object into a LetterFlags object.
"""
function LetterFlags(candidate::Candidate)
    letter_flags = LetterFlags()
    for char in candidate.chars
        letter_flags[char] = true
    end
    letter_flags
end


text:
"""
    pupil_θ(sz, pp::PSFParams, sampling)

returns the θ angle (to the optical axis) in the sample space as a pupil array.
"""
function pupil_θ(sz, pp::PSFParams, sampling)
    asin.(k_r(sz, pp, sampling)./k_0(pp))
end

"""
    pupil_ϕ(sz, pp::PSFParams, sampling)

returns the azimuthal angle ϕ  in the sample space as a pupil array.
"""
function pupil_ϕ(sz, pp::PSFParams, sampling)
    ϕ_tuple.(k_xy(sz, pp, sampling))
end

"""
    aplanatic_factor(sz, pp::PSFParams, sampling)

returns the aplanatic factor as specified in `pp.aplanatic` as a pupil array.
"""
function aplanatic_factor(sz, pp::PSFParams, sampling; is_proj=false)
    if is_proj
        pp.aplanatic.(pupil_θ(sz, pp, sampling)) ./ cos.(pupil_θ(sz, pp, sampling))
    else
        pp.aplanatic.(pupil_θ(sz, pp, sampling))
    end
end

"""
    field_pupil(sz, pp, sampling)

returns the pupil polarization as a 4D array with the XY polarization components stacked along the 4th dimension.
"""
function field_pupil(sz, pp, sampling)
    idx_to_dim(pp.polarization.(pp.dtype, k_xy_rel_pupil(sz,pp,sampling)), 4)
end

"""
    field_xy_to_xyz(field,pp,sampling)

converts a 2D field at the pupil to a 2D field behind the lens, containing E_x, E_Y and E_z.
"""
function field_xy_to_xyz(field,pp,sampling)
    if size(field)[4] == 1
        return field; # just stay scalar and apply nothing.
    end
    sz = size(field)
    ϕ = pupil_ϕ(sz, pp, sampling) # establish angles
    θ = pupil_θ(sz, pp, sampling)
    cosϕ = cos.(ϕ)
    sinϕ = sin.(ϕ)
    cosθ = cos.(θ)
    sinθ = sin.(θ)
    field_rad = field[:,:,:,1] .* cosϕ + field[:,:,:,2] .* sinϕ # radial field component in the pupil
    field_azi = -field[:,:,:,1] .* sinϕ + field[:,:,:,2] .* cosϕ # azimuthal field component in the pupil
    field_rad_cosθ = cosθ .* field_rad # what remains of the radial field component behind lens (with energy conservation)
    cat(field_rad_cosθ .* cosϕ .- field_azi .* sinϕ, field_rad_cosθ .* sinϕ .+ field_azi .* cosϕ, sinθ .* field_rad, dims=4)
end

"""
    field_xyz(sz, pp, sampling)

creates a 2D pupil field behind the lens, containing E_x, E_Y and E_z.
"""
function field_xyz(sz, pp, sampling)
    field_xy_to_xyz(field_pupil(sz, pp, sampling), pp, sampling)
end

"""
    pupil_xyz(sz, pp, sampling=nothing)

creates a pupil with electric field distributions in XYZ. Returns a 4D dataset with the electric field components along the 4th dimension.
"""
function pupil_xyz(sz, pp, sampling=nothing)
    if isnothing(sampling)
        sampling = get_Ewald_sampling(sz, pp)
    end
    if isnothing(pp.aberrations) || isempty(pp.aberrations.indices)
        field_xyz(sz, pp, sampling) .* aplanatic_factor(sz,pp,sampling, is_proj=true) .* ft(jinc_r_2d(sz[1:2], pp, sampling=sampling) .* my_disc(sz[1:2],pp))
    else
        field_xyz(sz, pp, sampling) .* aplanatic_factor(sz,pp,sampling, is_proj=true) .* get_zernike_pupil(sz, pp, sampling) .* ft(jinc_r_2d(sz[1:2], pp, sampling=sampling) .* my_disc(sz[1:2],pp))
    end
end

"""
    get_propagator(sz,pp,sampling)

retrieves the propagator phase, prpagating a single Z-slice.
"""
function get_propagator(sz,pp,sampling)
    if isnothing(sampling)
        sampling = get_Ewald_sampling(sz, pp)
    end
    k_max_rel = sampling[1:2] ./ (pp.λ / pp.n)
    # prop1 = propagator(pp.dtype, sz[1:2], Δz=1, k_max = k_max_rel)
    # prop = cispi.(2 .* phase_kz(pp.dtype, sz[1:2], scale = 1 ./ (k_max_rel .* sz[1:2])))
    scalar = (2π*sampling[3] / (pp.λ / pp.n))
    xy_scale = 1 ./ (k_max_rel .* sz[1:2])
    return scalar .* phase_kz(pp.dtype, sz[1:2], scale = xy_scale), scalar, xy_scale
end

"""
    get_propagator_gradient(prop_phase, scalar, xy_scale)

calculates the gradient of the propagator along the X and Y directions of the pupil.
"""
function get_propagator_gradient(prop_phase, scalar, xy_scale)
    #dr_prop_phase = ifelse.(prop_phase .== zero(pp.dtype), zero(pp.dtype), scalar.^2 ./ prop_term .* rr(pp.dtype, sz[1:2], scale = xy_scale))
    dtype = eltype(prop_phase)
    sz = size(prop_phase)[1:2]
    dx_prop_phase = ifelse.(prop_phase .== zero(dtype), zero(dtype), scalar.^2 ./ prop_phase .* xx(dtype, sz, scale = xy_scale.^2))
    dy_prop_phase = ifelse.(prop_phase .== zero(dtype), zero(dtype), scalar.^2 ./ prop_phase .* yy(dtype, sz, scale = xy_scale.^2))
    return dx_prop_phase, dy_prop_phase
end

"""
    get_zernike_pupil_phase(sz, pp, sampling)

calculates the phases in the pupil for a given set of aberrations as defined by `J` and `coefficients`.
By default this follows the OSA nomenclature. See the help file of `ZernikePolynomials.jl` for more information.
The pupil phase (up to the pupil border as defined by the `NA` in `pp`) is returned.

Arguments:
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple

Example:
```jdoctest
julia> using PSFs, FFTW

julia> aberr = PSFs.Aberrations([Zernike_Spherical, Zernike_ObliqueAstigmatism],[0.1, 0.2])
Aberrations([12, 3], [0.1, 0.2], :OSA)

julia> pp = PSFParams(580.0, 1.4, 1.518; aberrations=aberr)
PSFParams(580.0, 1.4, 1.518, Float32, ModeWidefield, PSFs.pol_scalar, PSFs.var"#42#43"(), PSFs.MethodPropagateIterative, nothing, Aberrations([12, 3], [0.1, 0.2], :OSA), nothing)

julia> sz = (10,10,64)
(10, 10, 64)

julia> sampling=(190,190,100)
(190, 190, 100)

julia> PSFs.get_zernike_pupil_phase(sz,pp,sampling)
10×10 Matrix{Float64}:
 0.0   0.0         0.0         0.0         0.0         0.0         0.0         0.0         0.0         0.0
 0.0   0.0         0.0         0.533599    0.202003   -0.0206203  -0.170662   -0.21173     0.0         0.0
 0.0   0.0         0.477274    0.186398    0.0287554  -0.104828   -0.250743   -0.3726     -0.361221    0.0
 0.0   0.533599    0.186398    0.0937363   0.0736565   0.016983   -0.112676   -0.278928   -0.3726     -0.21173
 0.0   0.202003    0.0287554   0.0736565   0.154747    0.162853    0.0615812  -0.112676   -0.250743   -0.170662
 0.0  -0.0206203  -0.104828    0.016983    0.162853    0.223607    0.162853    0.016983   -0.104828   -0.0206203
 0.0  -0.170662   -0.250743   -0.112676    0.0615812   0.162853    0.154747    0.0736565   0.0287554   0.202003
 0.0  -0.21173    -0.3726     -0.278928   -0.112676    0.016983    0.0736565   0.0937363   0.186398    0.533599
 0.0   0.0        -0.361221   -0.3726     -0.250743   -0.104828    0.0287554   0.186398    0.477274    0.0
 0.0   0.0         0.0        -0.21173    -0.170662   -0.0206203   0.202003    0.533599    0.0         0.0

```
"""
function get_zernike_pupil_phase(sz, pp, sampling)
    J = pp.aberrations.indices
    coefficients = pp.aberrations.coefficients
    index = pp.aberrations.index_style
    border = k_pupil_pos(sz[1:2],pp,sampling[1:2])
    X = ramp(1,sz[1],scale = 1/border[1])
    Y = ramp(1,sz[2],scale = 1/border[2])
    D = [[Zernike(j,coord=:cartesian,index=index)(x,y) for x in X, y in Y] for j in J ]
    return reduce(+,map(*,D,coefficients))
end

"""
    get_zernike_pupil(sz, pp, sampling)

calculates the phases in the pupil for a given set of aberrations as defined by `J` and `coefficients`.
By default this follows the OSA nomenclature. See the help file of `ZernikePolynomials.jl` for more information.
The complex-valued pupil (up to the pupil border as defined by the `NA` in `pp`) is returned.

Arguments:
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple

Example:
```jdoctest
julia> using PSFs, FFTW

julia> aberr = PSFs.Aberrations([Zernike_Spherical, Zernike_ObliqueAstigmatism],[0.1, 0.2])
Aberrations([12, 3], [0.1, 0.2], :OSA)

julia> pp = PSFParams(580.0, 1.4, 1.518; aberrations=aberr)
PSFParams(580.0, 1.4, 1.518, Float32, ModeWidefield, PSFs.pol_scalar, PSFs.var"#42#43"(), PSFs.MethodPropagateIterative, nothing, Aberrations([12, 3], [0.1, 0.2], :OSA), nothing)

julia> sz = (10,10,64)
(10, 10, 64)

julia> sampling=(190,190,100)
(190, 190, 100)

julia> PSFs.get_zernike_pupil(sz,pp,sampling)
10×10 Matrix{ComplexF64}:
 1.0+0.0im        1.0+0.0im               1.0+0.0im             1.0+0.0im               1.0+0.0im       …          1.0+0.0im             1.0+0.0im               1.0+0.0im             1.0+0.0im
 1.0+0.0im        1.0+0.0im               1.0+0.0im       -0.977799-0.209546im     0.297025+0.95487im         0.478105-0.878303im   0.238146-0.971229im          1.0+0.0im             1.0+0.0im
 1.0+0.0im        1.0+0.0im         -0.989823+0.142305im   0.389073+0.921207im     0.983723+0.179694im     -0.00466964-0.999989im  -0.696362-0.717691im    -0.643318-0.765599im        1.0+0.0im
 1.0+0.0im  -0.977799-0.209546im     0.389073+0.921207im   0.831518+0.555499im     0.894807+0.446453im        0.759688-0.650288im  -0.180764-0.983527im    -0.696362-0.717691im   0.238146-0.971229im
 1.0+0.0im   0.297025+0.95487im      0.983723+0.179694im   0.894807+0.446453im     0.563395+0.826187im        0.926073+0.377344im   0.759688-0.650288im  -0.00466964-0.999989im   0.478105-0.878303im
 1.0+0.0im   0.991619-0.129199im     0.790818-0.612051im   0.994312+0.106505im     0.520607+0.853797im  …     0.520607+0.853797im   0.994312+0.106505im     0.790818-0.612051im   0.991619-0.129199im
 1.0+0.0im   0.478105-0.878303im  -0.00466964-0.999989im   0.759688-0.650288im     0.926073+0.377344im        0.563395+0.826187im   0.894807+0.446453im     0.983723+0.179694im   0.297025+0.95487im
 1.0+0.0im   0.238146-0.971229im    -0.696362-0.717691im  -0.180764-0.983527im     0.759688-0.650288im        0.894807+0.446453im   0.831518+0.555499im     0.389073+0.921207im  -0.977799-0.209546im
 1.0+0.0im        1.0+0.0im         -0.643318-0.765599im  -0.696362-0.717691im  -0.00466964-0.999989im        0.983723+0.179694im   0.389073+0.921207im    -0.989823+0.142305im        1.0+0.0im
 1.0+0.0im        1.0+0.0im               1.0+0.0im        0.238146-0.971229im     0.478105-0.878303im        0.297025+0.95487im   -0.977799-0.209546im          1.0+0.0im             1.0+0.0im

```
"""
function get_zernike_pupil(sz, pp, sampling)
    cispi.(2 .*get_zernike_pupil_phase(sz, pp, sampling))
end


text:
module TestTypeDefs
    using Base.Test
    using NullableArrays

    x = NullableArray(
        [1, 2, 3],
        [false, false, true]
    )

    y = NullableArray(
        [
            1 2;
            3 4;
        ],
        [
            false false;
            true false;
        ],
    )

    @test isa(x, NullableArray{Int, 1})
    @test isa(x, NullableVector{Int})

    @test isa(y, NullableArray{Int, 2})
    @test isa(y, NullableMatrix{Int})
end


text:
#####
##### Convinience functions
#####

"""
    total_extent(topology, H, Δ, L)

Returns the total extent, including halo regions, of constant-spaced
`Periodic` and `Flat` dimensions with number of halo points `H`,
constant grid spacing `Δ`, and interior extent `L`.
"""
total_extent(topology, H, Δ, L) = L + (2H - 1) * Δ

"""
    total_extent(::Type{Bounded}, H, Δ, L)

Returns the total extent of, including halo regions, of constant-spaced
`Bounded` and `Flat` dimensions with number of halo points `H`,
constant grid spacing `Δ`, and interior extent `L`.
"""
total_extent(::Type{Bounded}, H, Δ, L) = L + 2H * Δ

"""
    total_length(loc, topo, N, H=0)

Returns the total length (number of nodes), including halo points, of a field
located at `Cell` centers along a grid dimension of length `N` and with halo points `H`.
"""
total_length(loc, topo, N, H=0) = N + 2H

"""
    total_length(::Type{Face}, ::Type{Bounded}, N, H=0)

Returns the total length, including halo points, of a field located at
cell `Face`s along a grid dimension of length `N` and with halo points `H`.
"""
total_length(::Type{Face}, ::Type{Bounded}, N, H=0) = N + 1 + 2H

#####
##### << Nodes >>
#####

@inline interior_indices(loc, topo, N) = 1:N
@inline interior_indices(::Type{Face}, ::Type{Bounded}, N) = 1:N+1

@inline interior_x_indices(loc, grid) = interior_indices(loc, topology(grid, 1), grid.Nx)
@inline interior_y_indices(loc, grid) = interior_indices(loc, topology(grid, 2), grid.Ny)
@inline interior_z_indices(loc, grid) = interior_indices(loc, topology(grid, 3), grid.Nz)

@inline interior_parent_indices(loc, topo, N, H) = 1+H:N+H
@inline interior_parent_indices(::Type{Face}, ::Type{Bounded}, N, H) = 1+H:N+1+H

# Node by node
@inline xnode(::Type{Cell}, i, grid) = @inbounds grid.xC[i, 1, 1]
@inline xnode(::Type{Face}, i, grid) = @inbounds grid.xF[i, 1, 1]

@inline ynode(::Type{Cell}, j, grid) = @inbounds grid.yC[1, j, 1]
@inline ynode(::Type{Face}, j, grid) = @inbounds grid.yF[1, j, 1]

@inline znode(::Type{Cell}, k, grid) = @inbounds grid.zC[1, 1, k]
@inline znode(::Type{Face}, k, grid) = @inbounds grid.zF[1, 1, k]

# Convenience is king
@inline xC(i, grid) = xnode(Cell, i, grid)
@inline xF(i, grid) = xnode(Face, i, grid)

@inline yC(j, grid) = ynode(Cell, j, grid)
@inline yF(j, grid) = ynode(Face, j, grid)

@inline zC(k, grid) = znode(Cell, k, grid)
@inline zF(k, grid) = znode(Face, k, grid)

"""
    xnodes(loc, grid)

Returns a view over the interior `loc=Cell` or `loc=Face` nodes
on `grid` in the x-direction. For `Bounded` directions,
`Face` nodes include the boundary points.

See `znodes` for examples.
"""
xnodes(::Type{Cell}, grid) = view(grid.xC, 1:grid.Nx, :, :)

"""
    ynodes(loc, grid)

Returns a view over the interior `loc=Cell` or `loc=Face` nodes
on `grid` in the y-direction. For `Bounded` directions,
`Face` nodes include the boundary points.

See `znodes` for examples.
"""
ynodes(::Type{Cell}, grid) = view(grid.yC, :, 1:grid.Ny, :)

"""
    znodes(loc, grid)

Returns a view over the interior `loc=Cell` or `loc=Face` nodes
on `grid` in the z-direction. For `Bounded` directions,
`Face` nodes include the boundary points.

Examples
========

```jldoctest
julia> using Oceananigans, Oceananigans.Grids

julia> horz_periodic_grid = RegularCartesianGrid(size=(3, 3, 3), extent=(2π, 2π, 1),
                                                 topology=(Periodic, Periodic, Bounded));

julia> zC = znodes(Cell, horz_periodic_grid)
1×1×3 view(OffsetArray(reshape(::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, 1, 1, 5), 1:1, 1:1, 0:4), :, :, 1:3) with eltype Float64 with indices 1:1×1:1×Base.OneTo(3):
[:, :, 1] =
 -0.8333333333333331

[:, :, 2] =
 -0.4999999999999999

[:, :, 3] =
 -0.16666666666666652

julia> zF = znodes(Face, horz_periodic_grid)
1×1×4 view(OffsetArray(reshape(::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, 1, 1, 6), 1:1, 1:1, 0:5), :, :, 1:4) with eltype Float64 with indices 1:1×1:1×Base.OneTo(4):
[:, :, 1] =
 -1.0

[:, :, 2] =
 -0.6666666666666666

[:, :, 3] =
 -0.33333333333333337

[:, :, 4] =
 -4.44089209850063e-17
```
"""
znodes(::Type{Cell}, grid) = view(grid.zC, :, :, 1:grid.Nz)

xnodes(::Type{Face}, grid) = view(grid.xF, interior_indices(Face, topology(grid, 1), grid.Nx), :, :)
ynodes(::Type{Face}, grid) = view(grid.yF, :, interior_indices(Face, topology(grid, 2), grid.Ny), :)
znodes(::Type{Face}, grid) = view(grid.zF, :, :, interior_indices(Face, topology(grid, 3), grid.Nz))

"""
    nodes(loc, grid)

Returns a 3-tuple of views over the interior nodes
at the locations in `loc` in `x, y, z`.

See `xnodes`, `ynodes`, and `znodes`.
"""
nodes(loc, grid::AbstractGrid) = (xnodes(loc[1], grid),
                                  ynodes(loc[2], grid),
                                  znodes(loc[3], grid))

#####
##### Convinience functions
#####

unpack_grid(grid) = grid.Nx, grid.Ny, grid.Nz, grid.Lx, grid.Ly, grid.Lz

#####
##### Input validation
#####

function validate_topology(topology)
    for T in topology
        if !isa(T(), AbstractTopology)
            e = "$T is not a valid topology! " *
                "Valid topologies are: Periodic, Bounded, Flat."
            throw(ArgumentError(e))
        end
    end

    return topology
end

"""Validate that an argument tuple is the right length and has elements of type `argtype`."""
function validate_tupled_argument(arg, argtype, argname)
    length(arg) == 3        || throw(ArgumentError("length($argname) must be 3."))
    all(isa.(arg, argtype)) || throw(ArgumentError("$argname=$arg must contain $argtype s."))
    all(arg .> 0)           || throw(ArgumentError("Elements of $argname=$arg must be > 0!"))
    return nothing
end

coordinate_name(i) = i == 1 ? "x" : i == 2 ? "y" : "z"

function validate_dimension_specification(i, c)
    name = coordinate_name(i)
    length(c) == 2       || throw(ArgumentError("$name length($c) must be 2."))
    all(isa.(c, Number)) || throw(ArgumentError("$name=$c should contain numbers."))
    c[2] >= c[1]         || throw(ArgumentError("$name=$c should be an increasing interval."))
    return nothing
end

function validate_regular_grid_size_and_extent(FT, size, extent, halo, x, y, z)
    validate_tupled_argument(size, Integer, "size")
    validate_tupled_argument(halo, Integer, "halo")

    # Find domain endpoints or domain extent, depending on user input:
    if !isnothing(extent) # the user has specified an extent!

        (!isnothing(x) || !isnothing(y) || !isnothing(z)) &&
            throw(ArgumentError("Cannot specify both length and x, y, z keyword arguments."))

        validate_tupled_argument(extent, Number, "extent")

        Lx, Ly, Lz = extent

        # An "oceanic" default domain:
        x = (  0, Lx)
        y = (  0, Ly)
        z = (-Lz,  0)

    else # isnothing(extent) === true implies that user has not specified a length

        (isnothing(x) || isnothing(y) || isnothing(z)) &&
            throw(ArgumentError("Must supply length or x, y, z keyword arguments."))

        for (i, c) in enumerate((x, y, z))
            validate_dimension_specification(i, c)
        end

        Lx = x[2] - x[1]
        Ly = y[2] - y[1]
        Lz = z[2] - z[1]
    end

    return FT(Lx), FT(Ly), FT(Lz), FT.(x), FT.(y), FT.(z)
end

function validate_vertically_stretched_grid_size_and_xy(FT, size, halo, x, y)
    validate_tupled_argument(size, Integer, "size")
    validate_tupled_argument(halo, Integer, "halo")

    for (i, c) in enumerate((x, y))
            validate_dimension_specification(i, c)
        end

        Lx = x[2] - x[1]
        Ly = y[2] - y[1]

    return FT(Lx), FT(Ly), FT.(x), FT.(y)
end


text:
# issue 1018
using Compose, DataFrames, Gadfly
Gadfly.push_theme(Theme(default_color=colorant"black", background_color="white"))

xscale = Scale.x_continuous(format=:plain)
gp = Geom.polygon(preserve_order=true)
D = DataFrame(x=[0, 0, 10^6, 10^6], y=[0, 10, 10, 0])

# Padding
# In absolute units:
pa = plot(D, x=:x, y=:y, gp, xscale, style(plot_padding=[10mm,10mm,5mm,5mm]))
# In relative units (relative to width & height of plot)
pb = plot(D, x=:x, y=:y, gp, xscale, style(plot_padding=[0.05w,0.05w,0.2h,0.2h]))
# 1mm padding on all sides, note x-axis right label gets cut deliberately
pc = plot(D, x=:x, y=:y, gp, xscale, style(plot_padding=[1mm]))
hstack(pa,pb,pc)


text:
"""
Hypergraph creation and manipulation.
"""

using MatrixNetworks
using StatsBase
include("Utilities.jl")

"""
`Hypergraph{T}`
=====================

Represents a vertex-labeled hypergraph as a list of edges.

Fields
------
   - `vals::Vector{T}`: The values corresponding to each node (or hyperedge, depending on use case)
   - `n::Int64`: The number of nodes
   - `m::Int64`: The number of hyperedges
   - `D::Vector{Int64}`: The degree sequence, where `D[i] = degree of node i`
   - `K::Vector{Int64}`: The edge dimension sequence, where `K[j] = size of edge j`
   - `edges::Vector{Vector{Int64}}`: The hyperedges and their members

A hypergraph is a generalization of a graph in which edges can contain any number
of nodes. So, for example, if we wanted to denote a three-way relationship between
nodes 1, 2, and 3, the hypergraph would contain the edge `{1, 2, 3}`.

Hypergraphs are useful because they allow us to more accurately indicate relationships
among groups of things. Consider a coauthorship network. In a dyadic graph,
we might represent a five-way collaboration as a 5-clique. While this does capture
the fact that all pairs of authors have now appeared in a paper together, it doesn't seem
right to categorize the event as 15 pairwise interactions. With a hypergraph, we
can more succinctly and intuitively represent the event as a single hyperedge of five nodes.
"""
mutable struct Hypergraph{T}
   vals::Vector{T}
   n::Int64
   m::Int64
   D::Vector{Int64}
   K::Vector{Int64}
   edges::Vector{Vector{Int64}}
end

"""
`Hypergraph_kernel`
===================

Verifies that:
   - The number of edges == `m`
   - The number of vals == `n`
   - All nodes are between 0 and `n`

If all conditions are met, returns degree and edge dimension sequences
"""
function Hypergraph_kernel(edges::Vector{Vector{Int64}}, vals::Vector{T},
                           n::Int64, m::Int64) where T
   @assert length(edges) == m # m is the number of edges
   @assert length(vals) == n # Each node has a val associated with it
   @assert all([0 < edges[i][j] < n + 1 for i = 1:m for j = 1:length(edges[i])]) # No node exceeds n

   D = zeros(Int64, n)
   K = zeros(Int64, m)
   for e = 1:m
      K[e] = size(edges[e],1)
      for v in edges[e]
         D[v] += 1
      end
   end
   edges = sort!.(edges, by=v -> D[v], rev=true) # Sort edges by descending node degree

   return D, K
end

"""
`Hypergraph` constructors
=============================

Functions
---------
   - `Hypergraph(edges, vals, n, m)`: Produces a hypergraph with the given edges, values, and size
   - `Hypergraph(edges, n, m)`: Produces a hypergraph with the given edge set and size, with all values as 1.0
   - `Hypergraph(edges)`: Produces a hypergraph with the given set of edges

Examples
--------
~~~~
Hypergraph([[1,2], [3,4], [1,4]], ["One", "Two", "Three", "Four"], 4, 3)
Hypergraph([[i,i+2] for i=1:3], 4, 3)
Hypergraph([[1,2,3], [2,4,6], [1,5]])
~~~~
"""
function Hypergraph(edges::Vector{Vector{Int64}}, vals::Vector{T},
                          n::Int64, m::Int64) where T
   D, K= Hypergraph_kernel(edges, vals, n, m)
   return Hypergraph(vals, n, m, D, K, edges)
end

function Hypergraph(edges::Vector{Vector{Int64}}, n::Int64, m::Int64)
   return Hypergraph(edges, ones(n), n, m)
end

function Hypergraph(edges::Vector{Vector{Int64}})
   return Hypergraph(edges, maximum([e[i] for e in edges for i = 1:length(e)]), length(edges))
end

Base.copy(H::Hypergraph) = Hypergraph(deepcopy(h.edges), deepcopy(h.vals), h.n, h.m)

"""
`add_node!`
===========

Add a node with a specific value to the hypergraph.

Arguments
---------
   - `H::Hypergraph`: The hypergraph to add to
   - `val::T`: The value to associate with the new node
"""
function add_node!(H::Hypergraph, val::T) where T
   H.n += 1
   push!(H.D, 0)
   push!(H.vals, val)
end

"""
`add_node!`
===========

Add a node to the hypergraph.
"""
function add_node!(H::Hypergraph)
   add_node!(H, nothing)
end

"""
`add_edge!`
===========

Add a hyperedge to a hypergraph.

Arguments
---------
   - `V::Hypergraph`: The hypergraph to add an edge to
   - `e::Vector{Int64}`: The edge to add
"""
function add_edge!(V::Hypergraph, e::Vector{Int64})
   push!(V.edges, e)
   push!(V.K, length(e))

   for v in e
      if v > V.n
         add_node!(V)
      end
      V.D[v] += 1
   end
end

"""
`add_value_edge!`
=================

Add a hyperedge to a hypergraph by specifying the corresponding values to be linked.

Arguments
---------
   - `V::Hypergraph`: The hypergraph to add an edge to
   - `e::Vector{T}`: The values to connect

Example
-------
~~~~
add_value_edge!(V, ["Kim", "Jim", "Tim"])
~~~~
"""
function add_edge_by_value!(V::Hypergraph, e::Vector{T}) where T
   node_edge = []

   for v in e
      i = findfirst(x -> x==v, V.vals)
      if i == -1
         add_node!(V, v)
         i = V.n + 1
      end
      V.D[i] += 1
      push!(node_edge, i)
   end

   push!(V.edges, node_edge)
   push!(V.K, length(node_edge))
end

"""
`remove_node!`
==============

Delete the node at index `n` from a hypergraph.
"""
function remove_node!(H::Hypergraph, n::Int64)
   for i = 1:H.m
      filter!(x -> x == n, H.edges[i]) # Remove from edge
      map!(x -> x > n ? x - 1 : x, H.edges[i]) # Adjust node identities
      if H.edges[i] == [] # Clear empty edges
         deleteat!(H.edges, i)
         deleteat!(H.K, i)
      end
      H.K[i] -= count(x -> x==n, H.edges[i]) # Edit edge dimension sequence
   end

   deleteat!(H.D, n) # Remove from degree sequence

   H.n -= 1
   H.m = length(H.edges)
end

"""
`remove_edge!`
==============

Delete the edge at index `e` from a hypergraph.
"""
function remove_edge!(H::Hypergraph, e::Int64)
   for n in H.edges[e]
      H.D[n] -= 1 # Edit degree sequence
   end
   deleteat!(H.edges, e) # Remove edge
   deleteat!(H.K, e)
   H.m -= 1
end

"""
`edge_intersect`
================

Get the set of intersection between two hyperdges. Returns a set of integers.

Arguments
---------
   - `H::Hypergraph`: The hypergraph where the edges live
   - `e1::Int64`: The index of the first edge
   - `e2::Int64`: The index of the second edge

Examples
--------
~~~~
edge_intersect(Hypergraph([[1,3,5], [2,3,5], [3,5]], 1, 2)) -> {3, 5}
~~~~
"""
function edge_intersect(H::Hypergraph, e1::Int64, e2::Int64)
   return intersect(H.edges[e1], H.edges[e2])
end

"""
`num_parallel`
==============

Get the number of hyperedges parallel (multiset equal) to a given edge.

Arguments
---------
   - `H::Hypergraph`: The hypergraph to be analyzed
   - `e::Int64`: The index of the edge
"""
function num_parallel(H::Hypergraph, e::Int64)
   return count(x -> x == H.edges[e], H.edges)
end

"""
`random_edges`
==============

Select `n` distinct random edges from a hypergraph.

Example
========
~~~~
random_edges(G, 3) -> [[e1], [e2], [e3]]
~~~~
"""
function random_edges(H::Hypergraph, n::Int64)
   return H.edges[random_edge_indices(H, n)]
end

"""
`random_edge_indices`
==============

Select `n` distinct random edge indices from a hypergraph.

Example
========
~~~~
random_edge_indices(G, 3) -> [e1, e2, e3]
~~~~
"""
function random_edge_indices(H::Hypergraph, n::Int64)
   return sample(1:H.m, n, replace=false)
end


text:
#=

  From srl-pp-tutorial-wasp-stockholm.pdf
  "Statistical Relational learning and Probabilistic Programming"
  by Luc De Raedt, Anton Dries, and Angelika Kimmig
  https://dtai.cs.kuleuven.be/problog/wasp17-tutorial.html

  slide 394f:
  """
  - Flipping a coin with unknown weight
  - Prior: uniform distribution on [0,1]
  - Observation: 5x heads in a row
  """

  The ProbLog model return the following which corresponds to the
  density of the probabilist
  """
   weight(c1,0.1): 3.3994697e-13
   weight(c1,0.3): 2.1679411e-06
   weight(c1,0.5): 0.0041497433
   weight(c1,0.7): 0.1317485
   weight(c1,0.9): 0.86409959         <----
   weight(c2,0.1): 3.2276726e-06
   weight(c2,0.3): 0.024109997
   weight(c2,0.5): 0.66724754         <----
   weight(c2,0.7): 0.30628626
   weight(c2,0.9): 0.0023529733
  """

  Here we observe 13 tosses and detecting the probability of throwing a head.

  Cf ~/problog/coins_learning.pl
     ~/blog/coins_learning.blog
     ~/webppl/coins_learning.wppl

  Note that this model is much simpler than the ProbLog model since ProbLog
  don't support continous distributions.
  Which makes it clear that I might have missed the point of the ProbLog model. :-)

  For a model which is much more close to the ProbLog model, see coins_learning2.jl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function coins_learning(obs)
    n = length(obs)
    p ~ Uniform(0,1)
    obs ~ filldist(flip(p),n)
    # post ~ filldist(flip(p),n)

end

# p: 0.9344
# obs = [true for _ in 1:13]
# p: 0.5984
obs = [true,false,true,true,true,true,true,false,false,true,false,false,true]
model = coins_learning(obs)
num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(5), 10_000)
chns = sample(model, SMC(1000), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
# -----------------------------------------------------------------------
# Data structure
# -----------------------------------------------------------------------
# -----------------------------------------------------------------------
# Cylinder position data structure
# -----------------------------------------------------------------------

# The two first values refer to a global coordinate system

# All the other values are local with the origin on the upwind cylinder center

# Global origin at the step edge

# These were made mutable in order to be changed without creating
# new variables, specially to get the current position.

# -----------------------------------------------------------------------
# Required packages
using AutoHashEquals, DataFrames

# 1 cylinder position
@auto_hash_equals mutable struct OneCylType
    x::Float64 # cylinder center - x coordinate (global)
    y::Float64 # cylinder center - y coordinate (global)
end

# 2 cylinder position
@auto_hash_equals mutable struct TwoCylType
    x::Float64 # first cylinder center - x coordinate (global)
    y::Float64 # first cylinder center - y coordinate (global)
    dx::Float64 # delta x to the second cylinder (related to the first)
    dy::Float64 # delta y to the second cylinder (related to the first)
end

# 3 cylinder position
@auto_hash_equals mutable struct ThreeCylType
    x::Float64 # first cylinder center - x coordinate (global)
    y::Float64 # first cylinder center - y coordinate (global)
    dx1::Float64 # delta x to the second cylinder (related to the first)
    dy1::Float64 # delta y to the second cylinder (related to the first)
    dx2::Float64 # delta x to the third cylinder (related to the second)
    dy2::Float64 # delta y to the third cylinder (related to the second)
end

# Union
CylType = Union{OneCylType,TwoCylType,ThreeCylType}
ArrayCylType = Union{Array{OneCylType,1},Array{TwoCylType,1},Array{ThreeCylType,1}}

# -----------------------------------------------------------------------
# Data structure for the stepper motor
# -----------------------------------------------------------------------
# No need for mutable data here.
# -----------------------------------------------------------------------

# Step data and strategy of mov. to pass to the stepper motor
@auto_hash_equals struct StepperMotionType
    step::Array{Int64,1}
    strategy::Array{Int64,1}
end

# -----------------------------------------------------------------------
# New method for convert
# -----------------------------------------------------------------------
# Make convert work for all the structures created
# -----------------------------------------------------------------------

# For one cylinder
function Base.convert(::Type{Array}, x::OneCylType)
    return [x.x, x.y]
end

function Base.convert(::Type{DataFrame}, x::OneCylType)
    return DataFrame(x=x.x, y=x.y)
end

function Base.convert(::Type{DataFrame}, x::Array{OneCylType,1})
    v = DataFrame(x=Float64[],y=Float64[])
    for i in x
        append!(v, convert(DataFrame, i))
    end
    return v
end

function Base.convert(::Type{OneCylType}, x::AbstractVector)
    if length(x) != 2
        throw(DimensionMismatch("vector length should be equal 2"))
    end
    return OneCylType(x[1],x[2])
end

function Base.convert(::Type{OneCylType}, x::DataFrame)
    if nrow(x) == 1
        return OneCylType(x.x, x.y)
    else
        e = eachrow(x)
        return [OneCylType(i.x, i.y) for i in e]
    end
end

# For two cylinder
function Base.convert(::Type{Array}, x::TwoCylType)
    return [x.x, x.y, x.dx, x.dy]
end

function Base.convert(::Type{DataFrame}, x::TwoCylType)
    return DataFrame(x=x.x, y=x.y, dx=x.dx, dy=x.dy)
end

function Base.convert(::Type{DataFrame}, x::Array{TwoCylType,1})
    v = DataFrame(x=Float64[],y=Float64[],dx=Float64[],dy=Float64[])
    for i in x
        append!(v, convert(DataFrame, i))
    end
    return v
end

function Base.convert(::Type{TwoCylType}, x::AbstractVector)
    if length(x) != 4
        throw(DimensionMismatch("vector length should be equal 4"))
    end
    return TwoCylType(x[1],x[2],x[3],x[4])
end

function Base.convert(::Type{TwoCylType}, x::DataFrame)
    if nrow(x) == 1
        return TwoCylType(x.x, x.y, x.dx, x.dy)
    else
        e = eachrow(x)
        return [TwoCylType(i.x, i.y, i.dx, i.dy) for i in e]
    end
end

# For three cylinder
function Base.convert(::Type{Array}, x::ThreeCylType)
    return [x.x, x.y, x.dx1, x.dy1, x.dx2, x.dy2]
end

function Base.convert(::Type{DataFrame}, x::ThreeCylType)
    return DataFrame(x=x.x, y=x.y, dx1=x.dx1, dy1=x.dy1, dx2=x.dx2, dy2=x.dy2)
end

function Base.convert(::Type{DataFrame}, x::Array{ThreeCylType,1})
    v = DataFrame(x=Float64[],y=Float64[],dx1=Float64[],dy1=Float64[],dx2=Float64[],dy2=Float64[])
    for i in x
        append!(v, convert(DataFrame, i))
    end
    return v
end

function Base.convert(::Type{ThreeCylType}, x::AbstractVector)
    if length(x) != 6
        throw(DimensionMismatch("vector length should be equal 6"))
    end
    return ThreeCylType(x[1],x[2],x[3],x[4],x[5],x[6])
end

function Base.convert(::Type{ThreeCylType}, x::DataFrame)
    if nrow(x) == 1
        return ThreeCylType(x.x, x.y, x.dx1, x.dy1, x.dx2, x.dy2)
    else
        e = eachrow(x)
        return [ThreeCylType(i.x, i.y, i.dx1, i.dy1, i.dx2, i.dy2) for i in e]
    end
end

# -----------------------------------------------------------------------
# New method for getindex
# -----------------------------------------------------------------------
# Make getindex work for all the structures created
# -----------------------------------------------------------------------
function Base.getindex(cyl::CylType, i::Integer)
    a = convert(Array, cyl)
    return a[i]
end

function Base.getindex(cyl::CylType, r::AbstractRange)
    a = convert(Array, cyl)
    return a[r]
end

# For StepperMotionType
function Base.getindex(sm::StepperMotionType, i::Integer)
    if i==1
        sm.step
    elseif i==2
        sm.strategy
    else
        throw(DimensionMismatch("Dimensions not matching"))
    end
end

# -----------------------------------------------------------------------
# New method for setindex!
# -----------------------------------------------------------------------
# Make getindex work for all the structures created
# -----------------------------------------------------------------------

# For OneCylType
function Base.setindex!(cyl::OneCylType, v::AbstractFloat, i::Integer)
    if i == 1
        cyl.x = v
    elseif i == 2
        cyl.y = v
    else
        throw(BoundsError)
    end
end

# For TwoCylType
function Base.setindex!(cyl::TwoCylType, v::AbstractFloat, i::Integer)
    if i==1
        cyl.x = v
    elseif i==2
        cyl.y = v
    elseif i==3
        cyl.dx = v
    elseif i==4
        cyl.dy = v
    else
        throw(BoundsError)
    end
end

# For ThreeCylType
function Base.setindex!(cyl::ThreeCylType, v::AbstractFloat, i::Integer)
    if i==1
        cyl.x = v
    elseif i==2
        cyl.y = v
    elseif i==3
        cyl.dx1 = v
    elseif i==4
        cyl.dy1 = v
    elseif i == 5
        cyl.dx2 = v
    elseif i == 6
        cyl.dy2 = v
    else
        throw(BoundsError)
    end
end

# -----------------------------------------------------------------------
# New method for length
# -----------------------------------------------------------------------
# Make length work for all the structures created
# -----------------------------------------------------------------------
Base.length(cyl::CylType) = convert(Array, cyl)

# -----------------------------------------------------------------------
# New method for length
# -----------------------------------------------------------------------
# Make length work for all the structures created
# -----------------------------------------------------------------------
Base.lastindex(cyl::CylType) = length(cyl)

Base.length(sm::StepperMotionType) = 2


text:
using Oceananigans.Solvers
using Oceananigans.Operators
using Oceananigans.Grids: with_halo
using Oceananigans.Architectures
using Oceananigans.Grids: AbstractGrid
using Oceananigans.Fields: ReducedField
using Oceananigans.Solvers: HeptadiagonalIterativeSolver
import Oceananigans.Solvers: solve!

"""
    struct MatrixImplicitFreeSurfaceSolver{V, S, R}

The hepta-diagonal matrix iterative implicit free-surface solver.

$(TYPEDFIELDS)
"""
struct MatrixImplicitFreeSurfaceSolver{V, S, R}
    "The vertically-integrated lateral areas"
    vertically_integrated_lateral_areas :: V
    "The matrix iterative solver"
    matrix_iterative_solver :: S
    "The right hand side of the free surface evolution equation"
    right_hand_side :: R
end

"""
    MatrixImplicitFreeSurfaceSolver(grid::AbstractGrid, gravitational_acceleration, settings)

Return a the framework for solving the elliptic equation with one of the iterative solvers of IterativeSolvers.jl
with a sparse matrix formulation.

```math
[∇ ⋅ H ∇ - Az / (g Δt²)] ηⁿ⁺¹ = (∇ʰ ⋅ Q★ - Az ηⁿ / Δt) / (g Δt)
```

representing an implicit time discretization of the linear free surface evolution equation
for a fluid with variable depth `H`, horizontal areas `Az`, barotropic volume flux `Q★`, time
step `Δt`, gravitational acceleration `g`, and free surface at time-step `n` `ηⁿ`.
"""
function MatrixImplicitFreeSurfaceSolver(grid::AbstractGrid, gravitational_acceleration, settings)

    # Initialize vertically integrated lateral face areas
    ∫ᶻ_Axᶠᶜᶜ = Field{Face, Center, Nothing}(grid)
    ∫ᶻ_Ayᶜᶠᶜ = Field{Center, Face, Nothing}(grid)

    vertically_integrated_lateral_areas = (xᶠᶜᶜ = ∫ᶻ_Axᶠᶜᶜ, yᶜᶠᶜ = ∫ᶻ_Ayᶜᶠᶜ)

    compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas)

    arch = architecture(grid)
    right_hand_side = zeros(grid, grid.Nx * grid.Ny) # linearized RHS for matrix operations

    # Set maximum iterations to Nx * Ny if not set
    settings = Dict{Symbol, Any}(settings)
    maximum_iterations = get(settings, :maximum_iterations, grid.Nx * grid.Ny)
    settings[:maximum_iterations] = maximum_iterations

    coeffs = compute_matrix_coefficients(vertically_integrated_lateral_areas, grid, gravitational_acceleration)

    solver = HeptadiagonalIterativeSolver(coeffs; reduced_dim = (false, false, true),
                                          grid = grid, settings...)

    return MatrixImplicitFreeSurfaceSolver(vertically_integrated_lateral_areas, solver, right_hand_side)
end

build_implicit_step_solver(::Val{:HeptadiagonalIterativeSolver}, grid, gravitational_acceleration, settings) =
    MatrixImplicitFreeSurfaceSolver(grid, gravitational_acceleration, settings)

#####
##### Solve...
#####

function solve!(η, implicit_free_surface_solver::MatrixImplicitFreeSurfaceSolver, rhs, g, Δt)

    solver = implicit_free_surface_solver.matrix_iterative_solver

    solve!(η, solver, rhs, Δt)

    return nothing
end

function compute_implicit_free_surface_right_hand_side!(rhs,
                                                        implicit_solver::MatrixImplicitFreeSurfaceSolver,
                                                        g, Δt, ∫ᶻQ, η)

    solver = implicit_solver.matrix_iterative_solver
    grid = solver.grid
    arch = architecture(grid)

    event = launch!(arch, grid, :xy,
                    implicit_linearized_free_surface_right_hand_side!,
                    rhs, grid, g, Δt, ∫ᶻQ, η,
		            dependencies = device_event(arch))

    return event
end

# linearized right hand side
@kernel function implicit_linearized_free_surface_right_hand_side!(rhs, grid, g, Δt, ∫ᶻQ, η)
    i, j = @index(Global, NTuple)
    Az   = Azᶜᶜᶜ(i, j, 1, grid)
    δ_Q  = flux_div_xyᶜᶜᶜ(i, j, 1, grid, ∫ᶻQ.u, ∫ᶻQ.v)
    t = i + grid.Nx * (j - 1)
    @inbounds rhs[t] = (δ_Q - Az * η[i, j, 1] / Δt) / (g * Δt)
end

function compute_matrix_coefficients(vertically_integrated_areas, grid, gravitational_acceleration)

    arch = grid.architecture

    Nx, Ny = (grid.Nx, grid.Ny)

    C     = zeros(Nx, Ny, 1)
    diag  = arch_array(arch, zeros(eltype(grid), Nx, Ny, 1))
    Ax    = arch_array(arch, zeros(eltype(grid), Nx, Ny, 1))
    Ay    = arch_array(arch, zeros(eltype(grid), Nx, Ny, 1))
    Az    = zeros(Nx, Ny, 1)

    ∫Ax = vertically_integrated_areas.xᶠᶜᶜ
    ∫Ay = vertically_integrated_areas.yᶜᶠᶜ

    event_c = launch!(arch, grid, :xy, _compute_coefficients!,
                      diag, Ax, Ay, ∫Ax, ∫Ay, grid, gravitational_acceleration,
                      dependencies = device_event(arch))

    wait(event_c)

    return (Ax, Ay, Az, C, diag)
end

@kernel function _compute_coefficients!(diag, Ax, Ay, ∫Ax, ∫Ay, grid, g)
    i, j = @index(Global, NTuple)
    @inbounds begin
        Ay[i, j, 1]    = ∫Ay[i, j, 1] / Δyᶜᶠᶜ(i, j, 1, grid)
        Ax[i, j, 1]    = ∫Ax[i, j, 1] / Δxᶠᶜᶜ(i, j, 1, grid)
        diag[i, j, 1]  = - Azᶜᶜᶜ(i, j, 1, grid) / g
    end
end


text:
# Product task maps share domain but not codomain
domain_manifold(task_map::PTMX) = domain_manifold(base_task_map(task_map).map1)

function codomain_manifold(task_map::PTMX)
    base_map = base_task_map(task_map)
    PM{codomain_manifold(base_map.map1), codomain_manifold(base_map.map2)}
end
function codomain_manifolds(task_map::PTMX)
    base_map = base_task_map(task_map)
    codomain_manifold(base_map.map1), codomain_manifold(base_map.map2)
end

choose_chart_emb(xme, task_map::PTM, CM, CN) =
    choose_chart_emb(task_map_emb(xme, task_map, CM, CN), CN)
choose_chart_chart(xm, task_map::PTM, CM, CN) =
    choose_chart_chart(task_map_chart(xm, task_map, CM, CN), CN)

base_task_map(task_map::PTM) = task_map

function task_map_emb(xme, task_map::PTM, CM, CN)
    CN1, CN2 = product_chart_split(CN)
    xn1e = task_map_emb(xme, task_map.map1, CM, CN1)
    xn2e = task_map_emb(xme, task_map.map2, CM, CN2)
    [xn1e; xn2e]
end

function task_map_emb_chart(xme, task_map::PTM, CM, CN)
    CN1, CN2 = product_chart_split(CN)
    xn1 = task_map_emb_chart(xme, task_map.map1, CM, CN1)
    xn2 = task_map_emb_chart(xme, task_map.map2, CM, CN2)
    [xn1; xn2]
end

function task_map_chart_emb(xm, task_map::PTM, CM, CN)
    CN1, CN2 = product_chart_split(CN)
    xn1e = task_map_chart_emb(xm, task_map.map1, CM, CN1)
    xn2e = task_map_chart_emb(xm, task_map.map2, CM, CN2)
    [xn1e; xn2e]
end

function task_map_chart(xm, task_map::PTM, CM, CN)
    CN1, CN2 = product_chart_split(CN)
    xn1 = task_map_chart(xm, task_map.map1, CM, CN1)
    xn2 = task_map_chart(xm, task_map.map2, CM, CN2)
    [xn1; xn2]
end


text:
module test_millboard_tuples

using Millboard
using Test

Millboard.set_table_mode(:markdown)

board = ()
@test sprint(show, table(board)) == """
"""

board = ([],)
@test sprint(show, table(board)) == """
|   | 1 |
|---|---|
| 1 |   |"""

board = ([1,2,3],)
@test sprint(show, table(board)) == """
|   | 1 |
|---|---|
|   | 1 |
|   | 2 |
| 1 | 3 |"""

board = [()]
@test sprint(show, table(board)) == """
|   | 1  |
|---|----|
| 1 | () |"""

board = ((1,2,3),)
@test sprint(show, table(board)) == """
|   | 1         |
|---|-----------|
| 1 | (1, 2, 3) |"""

board = [(1,2,3)]
@test sprint(show, table(board)) == """
|   | 1         |
|---|-----------|
| 1 | (1, 2, 3) |"""

board = [(1,2,3) (4,5,6)]
@test sprint(show, table(board)) == """
|   | 1         | 2         |
|---|-----------|-----------|
| 1 | (1, 2, 3) | (4, 5, 6) |"""

end # module test_tuples


text:
include("makeStretchData.jl")
include("../../../src/utils.jl")

using PlotlyJS, SeisIO, SeisNoise, JLD2, FFTW, ORCA

foname="/Users/jared/SCECintern2019/data/validation/verificationData.jld2"

types = ["dampedSinusoid", "rickerConv", "realData", "spectSynth"]

finame_xcorr="/Users/jared/SCECintern2019/data/reference/BPnetwork_2003/reference_xcorr/BPnetwork_2003_ref.jld2"
dataset_xcorr="BP.CCRB..BP1.BP.VARB..BP1"
xcf=jldopen(finame_xcorr)
real_xcorr=xcf[dataset_xcorr]
save=true
plot=false

if size(real_xcorr.corr)[2]>1 stack!(real_xcorr, allstack=true) end

dvVlist = collect(-0.03:0.005:0.03)
noiselist = collect(0:0.01:0.1)

if save
    f=jldopen(foname, "a+")
    f["info/dvVlist"] = dvVlist
    f["info/noiselist"] = noiselist
end

dampedSinParams = Dict( "A"    => 1.0,
                        "ω"    => 0.75,
                        "ϕ"    => 0.0,
                        "λ"    => 0.025,
                        "dt"   => 0.05,
                        "η"    => 0.0,
                        "t0"   => 0.0,
                        "npts" => 4001 )

sincParams = Dict( "A"    => 1.0,
                   "ω"    => 0.1,
                   "ϕ"    => 0.0,
                   "dt"   => 0.05,
                   "t0"   => 0.0,
                   "npts" => 4001 )

rickerParams = Dict( "f"      => 0.25,
                     "dt"     => 0.05,
                     "npr"    => 4001,
                     "npts"   => 4001,
                     "m"      => 1,
                     "sparse" => 0,
                     "stretchSource" => 0.0 )

chirpParams = Dict( "c"     => collect(range(0.15, stop=15.0, length=100)),
                    "tp"    => collect(range(0.04, stop=4.0, length=100)),
                    "mintp" => 0.04,
                    "maxtp" => 4.0,
                    "dist"  => 1000.0,
                    "n"     => 500,
                    "dt"    => 0.002,
                    "t0"    => 0.0 )

spectParams = Dict( "A"    => 1.0,
                    "dt"   => 0.05,
                    "npts" => 4001,
                    "t0"   => 0.0 )

lags = -real_xcorr.maxlag:1/real_xcorr.fs:real_xcorr.maxlag

for dvV in dvVlist
    for noiselvl in noiselist
        if "dampedSinusoid" in types
            # Example of damped sinusoid generation, stretching, and noise addition
            signal1_ds, t_ds = generateSignal("dampedSinusoid", dampedSinParams)
            normalize!(signal1_ds)
            signal2_ds, st_ds = stretchData(signal1_ds, dampedSinParams["dt"], dvV, n=noiselvl)

            addNoise!(signal1_ds, noiselvl)
            addNoise!(signal2_ds, noiselvl, seed=664739)
        end

        if "rickerConv" in types
            # Example of ricker wavelet generation and convolution with random reflectivity
            # series, stretching, and noise addition
            signal1_rc, t_rc = generateSignal("ricker", rickerParams, sparse=rickerParams["sparse"], stretchSource=0.0)
            normalize!(signal1_rc)
            if rickerParams["stretchSource"] != 0.0
                inSignal, t_rc = generateSignal("ricker", rickerParams, sparse=rickerParams["sparse"], stretchSource=rickerParams["stretchSource"])
                normalize!(inSignal)
            else
                inSignal = signal1_rc
            end
            signal2_rc, st_rc = stretchData(inSignal, rickerParams["dt"], dvV, n=noiselvl)

            addNoise!(signal1_rc, noiselvl)
            addNoise!(signal2_rc, noiselvl, seed=664739)
        end

        if "spectSynth" in types
            # Example of synthetic generation in frequency domain and irfft to get into time domain
            signal1_ss, t_ss = generateSignal("spectSynth", spectParams)
            normalize!(signal1_ss)
            signal2_ss, st_ss = stretchData(signal1_ss, spectParams["dt"], dvV, n=noiselvl)

            addNoise!(signal1_ss, noiselvl)
            addNoise!(signal2_ss, noiselvl, seed=664739)
        end

        if "realData" in types
            # Example of stretching real cross-correlations
            xcorr = real_xcorr.corr[:,1]
            normalize!(xcorr)
            stretch_xcorr, st_xcorr = stretchData(xcorr, 1/real_xcorr.fs, dvV, starttime=-real_xcorr.maxlag, stloc=0.0, n=noiselvl)

            addNoise!(xcorr, noiselvl)
            addNoise!(stretch_xcorr, noiselvl, seed=664739)
        end

        if save
            f["dampedSinusoid/$dvV.$noiselvl/ref"] = signal1_ds
            f["dampedSinusoid/$dvV.$noiselvl/cur"] = signal2_ds
            f["rickerConv/$dvV.$noiselvl/ref"] = signal1_rc
            f["rickerConv/$dvV.$noiselvl/cur"] = signal2_rc
            f["realData/$dvV.$noiselvl/ref"] = xcorr
            f["realData/$dvV.$noiselvl/cur"] = stretch_xcorr
            f["spectSynth/$dvV.$noiselvl/ref"] = signal1_ss
            f["spectSynth/$dvV.$noiselvl/cur"] = signal2_ss
        end

        if plot
            plots = Array{PlotlyJS.Plot, 1}()
            if "dampedSinusoid" in types
                p1 = PlotlyJS.Plot([PlotlyJS.scatter(;x=t_ds, y=signal1_ds, name="Unstretched", line_color="black"),
                                    PlotlyJS.scatter(;x=t_ds, y=signal2_ds, name="Stretched $(dvV*(-100))%", line_color="red")])
            else
                p1=PlotlyJS.Plot([PlotlyJS.scatter(;y=ones(10))])
            end
            if "rickerConv" in types
                p2 = PlotlyJS.Plot([PlotlyJS.scatter(;x=t_rc, y=signal1_rc, name="Unstretched", line_color="black"),
                                    PlotlyJS.scatter(;x=t_rc, y=signal2_rc, name="Stretched $(dvV*(-100))%", line_color="red")])
            else
                p2=PlotlyJS.Plot([PlotlyJS.scatter(;y=ones(10))])
            end
            if "realData" in types
                p3 = PlotlyJS.Plot([PlotlyJS.scatter(;x=lags, y=xcorr, name="Unstretched", line_color="black"),
                                    PlotlyJS.scatter(;x=lags, y=stretch_xcorr, name="Stretched $(dvV*(-100))%", line_color="red")])
            else
                p3=PlotlyJS.Plot([PlotlyJS.scatter(;y=ones(10))])
            end
            # temporarily, you cannot plot only certain pairs of data types. You must plot all.
            # these conditionals will be more meaningful later
            plots=[p3]

            p=PlotlyJS.plot(plots)
            display(p)
            readline()
        end
    end
end
if save close(f) end
close(xcf)

max_stars_repo_name: kura-okubo/SeisXcorrelation

max_issues_repo_name: kura-okubo/SeisXcorrelation

max_forks_repo_name: kura-okubo/SeisXcorrelation


text:
using BOMoD
using Random
#include("Combinatroics.jl")

# construct moduels

mt_Symbol = Mod(:sym4)
mt_a = Mod("a")
mt_b = Mod("b")
mt_c = Mod("c")
mt_d = Mod("d")

Group_Mod([Mod("a"),Mod("b"),Mod("c"),Mod("d")])

mt_Symbol_array = group_mod([:sym, :sym1, :sym3])
mt_String_array = group_mod(["a", "b", "c", "d","e"])

# set constrains
#single constrains
pos_con = Possition_Constrain(1,[Mod("a")])
Unordered_con_c_e = UnOrdered_Constrain([Mod("c") , Mod("e")])
Ordered_con_a_b = Ordered_Constrain([1 3],[Mod("a") Mod("b")])
Ordered_con_c_b = Ordered_Constrain([2 3],[Mod("c") Mod("b")])
#combine different constrains
Compose_Ordered =   Ordered_con_c_b +  Ordered_con_a_b
Compose_Constuct =   Ordered_con_c_b +  Ordered_con_a_b + Unordered_con_c_e
Compose_Construct_backup = Compose_Ordered + Unordered_con_c_e
Compose_multi = pos_con + Ordered_con_c_b + Ordered_con_a_b + Unordered_con_c_e
Compose_multi_backup = pos_con + Compose_Constuct
#easy combine
Compose_Constuct_easy = sum([Unordered_con_c_e, Ordered_con_a_b, Ordered_con_c_b ])
Compose_multi_easy = sum([Unordered_con_c_e, Ordered_con_a_b, Ordered_con_c_b, pos_con])

#test_design construction
#1 lenght no constrain
first_design =construct_ordered_design(mt_String_array,3)

#2 multiple length no constrain
construct_ordered_design(mt_String_array,2,4)
# test iteror over space
for i in first_design.space
    @show i
end
#take a constronct
aconstruct = first_design.space.space[2]
aconstruct2=Ordered_Construct([Mod("a") Mod("c") Mod("b")])
aconstrain = Ordered_Constrain([1 3],[Mod("a") Mod("b")])
#3 one lenght and construct contrain
construct_ordered_design(mt_String_array,3,Compose_Constuct)
construct_ordered_design(mt_String_array,2,3,Compose_Constuct)

rng = MersenneTwister()
sp = Random.Sampler(rng, 1:20)


text:
module CAS_module

push!(LOAD_PATH, "./")
push!(LOAD_PATH, "../")
push!(LOAD_PATH, "../ACASXU_Dependencies")
push!(LOAD_PATH, "../../VICAS/src")
push!(LOAD_PATH, "../../DeepCorrection")
push!(LOAD_PATH, "../../DeepCorrection/src")

using Simulation_Env_module
using JLD, CAS_QMDP, DeepQLearning

export CAS
export ACAS_Xu_Network, Simple_CAS, No_CAS, VICAS, Corrected_VICAS, VICASWeighted # NNCAS
export CAS_DICT
export acas_xu_network_data, evaluate_network
export alphas, get_qvals, norm_angle
export policy, STATE_DIM, AUG_STATE_DIM, TERM_VAR
export correction_version, VICAS_policyname, correction_type, correction_weight

abstract type CAS end

include("../ACASXU_Dependencies/load_network.jl")
const nnfile_name = "../ACASXU_Dependencies/ACASXU_TF12_run3_DS_Minus15_Online_08172017_200Epochs_largerRange.nnet"
const acas_xu_network_data = NNet(nnfile_name)

mutable struct ACAS_Xu_Network <: CAS
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function ACAS_Xu_Network()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :COC, 2 => :WL, 3 => :WR, 4 => :SL, 5 => :SR)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => 0., :WL => 5., :WR => -5., :SL => 10., :SR => -10.)
		return this
	end
end

# const VICAS_policyname = "test_mid3"
const VICAS_policyname = "mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0"
const alphas_filename =
	"../../VICAS/policies/" * VICAS_policyname * ".jld" # requires disc = mid3
const alphas = load_alphas(alphas_filename)

mutable struct VICAS <: CAS
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function VICAS()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :SR, 2 => :WR, 3 => :KEEP, 4 => :WL, 5 => :SL, 6 => :COC)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => -1.0, :WL => 5., :WR => -5., :SL => 10., :SR => -10., :KEEP => 0.0)
		return this
	end
end

mutable struct VICASWeighted <: CAS
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function VICASWeighted()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :SR, 2 => :WR, 3 => :KEEP, 4 => :WL, 5 => :SL, 6 => :COC)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => -1.0, :WL => 5., :WR => -5., :SL => 10., :SR => -10., :KEEP => 0.0)
		return this
	end
end

# const correction_type = :sector
# const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-13T18\:55\:43.491" # best sector
# const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-15T22\:38\:58.856"
# const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-28T10\:55\:10.697"
const correction_type = :closest
const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-28T10\:42\:31.084" # best closest
# const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-29T22\:18\:30.068"
# requires disc = mid3
const problem_file = "../../DeepCorrection/logs/" * correction_version * "/final_problem.jld"
const weights_file = "../../DeepCorrection/logs/" * correction_version * "/final_weights.jld"
policy = restore(problem_file=problem_file, weights_file=weights_file)
correction_weight = policy.env.correction_weight

mutable struct Corrected_VICAS <: CAS
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function Corrected_VICAS()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :SR, 2 => :WR, 3 => :KEEP, 4 => :WL, 5 => :SL, 6 => :COC)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => -1.0, :WL => 5., :WR => -5., :SL => 10., :SR => -10., :KEEP => 0.0)
		return this
	end
end

# const nn_version = ""
# const nn_problem_file = "../../DeepCorrection/logs/" * nn_version * "/final_problem.jld"
# const nn_weights_file = "../../DeepCorrection/logs/" * nn_version * "/final_weights.jld"
# nn_policy = restore(problem_file=nn_problem_file, weights_file=nn_weights_file)
# nn_policy.env.correction = true

# mutable struct NNCAS <: CAS
# 	Ind_to_Advisory_dict::Dict{Int64, Symbol}
# 	Advisory_to_Ind_dict::Dict{Symbol, Int64}
# 	Advisory_to_Action_dict::Dict{Symbol, Float64}

# 	function NNCAS()
# 		this = new()
# 		this.Ind_to_Advisory_dict =
# 				Dict(1 => :SR, 2 => :WR, 3 => :KEEP, 4 => :WL, 5 => :SL, 6 => :COC)
# 		this.Advisory_to_Ind_dict =
# 				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
# 		this.Advisory_to_Action_dict =
# 				Dict(:COC => -1.0, :WL => 5., :WR => -5., :SL => 10., :SR => -10., :KEEP => 0.0)
# 		return this
# 	end
# end

"""
	outdated
"""
mutable struct Simple_CAS <: CAS
	alert_threshold::Int64
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function Simple_CAS()
		this = new()
		set_alert_threshold = 1320. * 0.6 # [ft] ≈ 400 [m]
		this.alert_threshold = set_alert_threshold
		this.Ind_to_Advisory_dict =
				Dict(1 => :COC, 2 => :WL, 3 => :WR, 4 => :SL, 5 => :SR)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => 0., :WL => 5., :WR => -5., :SL => 10., :SR => -10.)
		return this
	end
end

mutable struct No_CAS <: CAS
	dt::Float64
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function No_CAS()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :COC, 2 => :WL, 3 => :WR, 4 => :SL, 5 => :SR)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => 0., :WL => 5., :WR => -5., :SL => 10., :SR => -10.)
		return this
	end
end

# construct all CAS' in advance to avoid repetitive contruction
CAS_DICT = Dict{DataType, CAS}()
for i = 1:length(subtypes(CAS))
	CAS_DICT[subtypes(CAS)[i]] = subtypes(CAS)[i]()
end

println("All CAS loaded!")

end # module

max_stars_repo_name: parachutel/DeepCorrectionCAS


text:
"""
    ActivationContribution{L,M} <: AbstractMutableComp
    ActivationContribution(l)
    ActivationContribution(l, method)

Calculate neuron utility based on activations and gradients using `method`.

Can be a performance bottleneck in cases with large activations. Use [`NeuronUtilityEvery`](@ref) to mitigate.

Default `method` is described in <https://arxiv.org/abs/1611.06440>.

Short summary is that the first order taylor approximation of the optimization problem: "which neurons shall I remove to minimize impact on the loss function?"
boils down to: "the ones which minimize `abs(gradient * activation)`" (assuming parameter independence).

"""
mutable struct ActivationContribution{L,M} <: AbstractMutableComp
    layer::L
    contribution # Type of activation not known yet :( Also leave some room for experimenting with things like storing the metric on the GPU
    method::M
end
ActivationContribution(l::AbstractMutableComp, method = Ewma(0.05f0)) = ActivationContribution(l, zeros(Float32, nout(l)), method)
ActivationContribution(l, method = Ewma(0.05f0)) = ActivationContribution(l, missing, method)

layer(m::ActivationContribution) = layer(m.layer)
layertype(m::ActivationContribution) = layertype(m.layer)
wrapped(m::ActivationContribution) = m.layer

Flux.trainable(m::ActivationContribution) = Flux.trainable(wrapped(m))

@functor ActivationContribution

function(m::ActivationContribution)(x...)
    act = wrapped(m)(x...)

    return Flux.Zygote.hook(act) do grad
        grad === nothing && return grad
        m.contribution = m.method(m.contribution, act, grad)
        return grad
    end
end

actdim(nd::Integer) = nd - 1

function NaiveNASlib.Δsize!(m::ActivationContribution, inputs::AbstractVector, outputs::AbstractVector; kwargs...)
    if m.contribution !== missing
        # This tends to happen when we are measuring contribution for a concatenation and we have added an extra input edge
        # TODO: Try to find another fix, perhaps we need to ensure that nout(v) if v wraps an ActivationContribution always return
        # the length of m.contribution
        outputs[outputs .> length(m.contribution)] .= -1
    end
    m.contribution = select(m.contribution, 1 => outputs; newfun = (args...) -> 0)
    NaiveNASlib.Δsize!(wrapped(m), inputs, outputs; kwargs...)
end

"""
    l2_squeeze(x, dimkeep)

Return l2 norm of `x` along all dimensions except `dimkeep` as a 1D array (singleton dimensions are removed).
"""
function l2_squeeze(x, dimskeep=1:ndims(x))
    dims = filter(i -> i ∉ dimskeep, 1:ndims(x))
    return sqrt.(dropdims(sum(x -> x^2, x, dims=dims), dims=Tuple(dims)))
end
l2_squeeze(z::Flux.Zeros, args...) = z

"""
    mean_squeeze(f, x, dimkeep)

Return mean value of `f.(x)` along all dimensions except `dimkeep` as a 1D array (singleton dimensions are removed).
"""
function mean_squeeze(f, x, dimskeep=1:ndims(x))
    dims = filter(i -> i ∉ dimskeep, 1:ndims(x))
    return dropdims(mean(f, x, dims=dims), dims=Tuple(dims))
end

# To peel the onion...
neuronutility(v::AbstractVertex) = neuronutility(base(v))
neuronutility(v::InputSizeVertex) = ones(nout(v))
neuronutility(v::CompVertex) = neuronutility(v.computation)
neuronutility(m::AbstractMutableComp) = neuronutility(wrapped(m))
function neuronutility(lm::LazyMutable)
   forcemutation(lm)
   neuronutility(wrapped(lm))
end
neuronutility(m::ActivationContribution) = m.contribution
neuronutility(l) = neuronutility(layertype(l), l)

# Default: mean of abs of weights + bias. Not a very good metric, but should be better than random
# Maybe do something about state in recurrent layers as well, but CBA to do it right now
neuronutility(::FluxParLayer, l) = l2_squeeze(weights(l), outdim(l)) .+ l2_squeeze(bias(l))
function neuronutility(::FluxDepthwiseConv, l)
    wm = l2_squeeze(weights(l), outdim(l))
    bm = l2_squeeze(bias(l))

    (length(wm) == 1 || length(wm) == length(bm)) && return wm .+ bm
    # use this to get insight on whether to repeat inner or outer:
    # cc = DepthwiseConv(reshape([1 1 1 1;2 2 2 2], 1, 1, 2, 4), [0,0,0,0,1,1,1,1])
    # cc(fill(10, (1,1,4,1)))
    return repeat(wm, length(bm) ÷ length(wm)) .+ bm
end

neuronutility(::FluxParNorm, l) = l.affine ? l2_squeeze(l.γ) .+ l2_squeeze(l.β) : missing

# Not possible to do anything since we don't know the size. Implementors can however use this to fallback to other ways if this is not an error
neuronutility(lt, l) = missing

neuronutility_safe(v) = neuronutility_safe(trait(v), v)
neuronutility_safe(t::DecoratingTrait, v) = neuronutility_safe(base(t), v)
neuronutility_safe(::Immutable, v) = 1
neuronutility_safe(::MutationSizeTrait, v) = clean_values(cpu(neuronutility(v)))
neuronutility_safe(m::AbstractMutableComp) = clean_values(cpu(neuronutility(m)))

clean_values(::Missing) = 1
clean_values(a::AbstractArray) = replace(a, NaN => -100, Inf => -100, -Inf => -100)

"""
    neuronutilitytaylor(currval, act, grad)

Calculate contribution of activations towards loss according to https://arxiv.org/abs/1611.06440.

Short summary is that the first order taylor approximation of the optimization problem: "which neurons shall I remove to minimize impact on the loss function?"
boils down to: "the ones which minimize abs(gradient * activation)" (assuming parameter independence).
"""
neuronutilitytaylor(currval, act, grad) = mean_squeeze(abs, (mean_squeeze(identity, act .* grad, (actdim(ndims(act)), ndims(act)))), 1)
# Kinda wished they had branded this better as 'taylor' can mean many things.

"""
    Ewma{R<:Real, M}
    Ewma(α::Real, method)

Exponential moving average of neuron utility calculated by `method`.

Parameter `α` acts as a forgetting factor, i.e larger values means faster convergence but more noisy estimate.
"""
struct Ewma{R<:Real, M}
    α::R
    method::M
    function Ewma(α::R, method::M) where {R,M}
        0 <= α <= 1 || error("α must be between 0 and 1, was $α")
        new{R,M}(α, method)
    end
end
Ewma(α) = Ewma(α, neuronutilitytaylor)

(m::Ewma)(currval, act, grad) = agg(m, currval, m.method(currval, act, grad))

function agg(m::Ewma, x, y)
    α = convert(float(eltype(x)), m.α)
    α .* x .+ (1 - α) .* y
end
agg(m::Ewma, ::Missing, y) = y

"""
    NeuronUtilityEvery{N,T}
    NeuronUtilityEvery(n::Int, method::T)

Calculate neuron utility using `method` every `n`:th call.

Useful to reduce runtime overhead.
"""
mutable struct NeuronUtilityEvery{N,T}
    cnt::Int
    method::T
    NeuronUtilityEvery(N::Int, method::T) where T = new{N, T}(0, method)
end
NeuronUtilityEvery(n::Int) = NeuronUtilityEvery(n, Ewma(0.05))

function (m::NeuronUtilityEvery{N})(currval, act, grad) where N
    ret = m.cnt % N == 0 ? m.method(currval, act, grad) : currval
    m.cnt += 1
    return ret
end


text:
# This file is a part of RadiationDetectorSignals.jl, licensed under the MIT License (MIT).

using RadiationDetectorSignals
using Test

using ArraysOfArrays, StaticArrays, TypedTables, Unitful

@testset "detector_hits" begin
    evtno_data = [[1, 1, 1], [2], [3, 4]]
    detno_data = [[2, 4, 4], [3], [2, 5]]
    thit_data = [[0.1, 0.2, 0.4], [0.3], [0.1, 0.7]]
    edep_data = [[45.2, 683.1, 137.4], [64.3], [84.2, 1043.0]]
    pos_data = [
        [SVector(0.4,-5.2,3.3), SVector(2.1,4.7,-5.5), SVector(-7.2,1.1,4.2)],
        [SVector(8.1,-0.9,-4.9)],
        [SVector(2.0,8.7,0.3), SVector(-2.3,9.8,7.4)]
    ]

    @test @inferred(DetectorHits((evtno = evtno_data[1], detno = detno_data[1], thit = thit_data[1], edep = edep_data[1], pos = pos_data[1]))) isa DetectorHits
    x = DetectorHits((evtno = evtno_data[1], detno = detno_data[1], thit = thit_data[1], edep = edep_data[1], pos = pos_data[1]))

    @test @inferred(DetectorHitEvents((evtno = evtno_data, detno = detno_data, thit = thit_data, edep = edep_data, pos = pos_data))) isa DetectorHitEvents
    X = DetectorHitEvents((evtno = evtno_data, detno = detno_data, thit = thit_data, edep = edep_data, pos = pos_data))

    @test @inferred(Table(X[1])) == x
end # testset

@testset "detector_hits with units" begin
    evtno_data = [[1, 1, 1], [2], [3, 4]]
    detno_data = [[2, 4, 4], [3], [2, 5]]
    thit_data = [[0.1, 0.2, 0.4]u"ns", [0.3]u"ns", [0.1, 0.7]u"ns"]
    edep_data = [[45.2, 683.1, 137.4]u"keV", [64.3]u"keV", [84.2, 1043.0]u"keV"]
    pos_data = [
        [SVector(0.4,-5.2,3.3)u"mm", SVector(2.1,4.7,-5.5)u"mm", SVector(-7.2,1.1,4.2)u"mm"],
        [SVector(8.1,-0.9,-4.9)u"mm"],
        [SVector(2.0,8.7,0.3)u"mm", SVector(-2.3,9.8,7.4)u"mm"]
    ]

    @test @inferred(DetectorHits((evtno = evtno_data[1], detno = detno_data[1], thit = thit_data[1], edep = edep_data[1], pos = pos_data[1]))) isa DetectorHits
    x = DetectorHits((evtno = evtno_data[1], detno = detno_data[1], thit = thit_data[1], edep = edep_data[1], pos = pos_data[1]))

    @test @inferred(DetectorHitEvents((evtno = evtno_data, detno = detno_data, thit = thit_data, edep = edep_data, pos = pos_data))) isa DetectorHitEvents
    X = DetectorHitEvents((evtno = evtno_data, detno = detno_data, thit = thit_data, edep = edep_data, pos = pos_data))

    @test @inferred(Table(X[1])) == x
end # testset


text:
@testset "Stress" begin
    gint = loadgraph(joinpath(testdir, "testdata", "graph-50-500.jgz"), "graph-50-500")

    c = vec(readdlm(joinpath(testdir, "testdata", "graph-50-500-sc.txt"), ','))
    for g in testdigraphs(gint)
        z  = @inferred(stress_centrality(g))
        @test z == c

        x  = @inferred(stress_centrality(g, 3))
        x2  = @inferred(stress_centrality(g, collect(1:20)))
        @test length(x) == 50
        @test length(x2) == 50
    end

    g1 = CycleGraph(4)
    add_vertex!(g1)
    add_edge!(g1, 4, 5)
    for g in testgraphs(g1)
        z  = @inferred(stress_centrality(g))
        @test z == [4, 2, 4, 10, 0]
    end
end


text:
import RIOPA
import Test: @testset, @test, @test_throws

@testset "ratios" begin
    @test parse(Rational{Int32}, "1/4") == 1 // 4
    @test parse(Rational{Int32}, "1//4") == 1 // 4
    @test parse(Rational, "1/4") == 1 // 4
    @test parse(Rational, "1//4") == 1 // 4

    @test RIOPA.get_ratio(0.2) == 0.2
    @test RIOPA.get_ratio("0.3") == 0.3
    @test RIOPA.get_ratio("1/4") == 0.25
    @test RIOPA.get_ratio("1//4") == 0.25
end


text:
using StaticArrays
using Rotations
using LinearAlgebra
using Test

import Rotations: RotationError, rotation_error, add_error, params

@testset "Rotation Error" begin
q1 = rand(UnitQuaternion)
dq = expm(0.1 * normalize(@SVector rand(3)))
q2 = q1*dq
@test q1\q2 ≈ dq
dg = RodriguesParam(dq)

emap = CayleyMap()
e1 = rotation_error(q2, q1, emap)
@test e1 ≈ q2 ⊖ q1
@test e1 ≈ params(dg)
@test add_error(q1, e1) ≈ q2
@test q1 ⊕ e1 isa UnitQuaternion
@test rotation_angle(UnitQuaternion(e1)) ≈ 0.1

aa1 = AngleAxis(q1)
@test add_error(aa1, e1) ≈ q2

e1 = rotation_error(q2, aa1, emap)
@test add_error(aa1, e1) ≈ AngleAxis(q2)
@test aa1 ⊕ e1 ≈ q2

# Test IdentityMap
g1 = rand(RodriguesParam)
g2 = RodriguesParam(g1*dg)
@test rotation_angle(g1\g2) ≈ 0.1

e1 = rotation_error(g2, g1, IdentityMap())
@test e1 ≈ g2 ⊖ g1
@test e1 ≈ params(dg)

p1 = rand(MRP)
dp = MRP(dg)
p2 = p1*dp
@test rotation_angle(p1\p2) ≈ 0.1

e1 = rotation_error(p2, p1, IdentityMap())
@test e1 ≈ p2 ⊖ p1
@test e1 ≈ params(dp)

@test_throws ArgumentError rotation_error(q2, q1, IdentityMap())
@test_throws ArgumentError q1 ⊕ e1

# Test inv
emap = CayleyMap()
imap = Rotations.InvCayleyMap()
g = inv(emap)(q1)
@test g ≈ imap(q1)
@test emap(g) ≈ q1
@test inv(imap)(g) ≈ q1
@test inv(inv(emap))(g) ≈ q1
@test inv(inv(imap))(q1) ≈ g

g = RodriguesParam(g)
@test UnitQuaternion(g) ≈ q1
@test imap(g) ≈ params(g)
end