# simd version is defined in simd.jl
@generated function _map(f, xs::Vararg{AbstractTensor, N}) where {N}
    S = promote_space(map(Space, xs)...)
    exps = map(indices(S)) do i
        vals = [:(xs[$j][$i]) for j in 1:N]
        :(f($(vals...)))
    end
    TT = tensortype(S)
    return quote
        @_inline_meta
        @inbounds $TT($(exps...))
    end
end

@inline Base.:+(x::AbstractTensor, y::AbstractTensor) = _map(+, x, y)
@inline Base.:-(x::AbstractTensor, y::AbstractTensor) = _map(-, x, y)
@inline Base.:*(y::Real, x::AbstractTensor) = _map(x -> x*y, x)
@inline Base.:*(x::AbstractTensor, y::Real) = _map(x -> x*y, x)
@inline Base.:/(x::AbstractTensor, y::Real) = _map(x -> x/y, x)
@inline Base.:+(x::AbstractTensor) = x
@inline Base.:-(x::AbstractTensor) = _map(-, x)

# with AbstractArray
@generated Base.:+(x::AbstractTensor, y::AbstractArray) = :(@_inline_meta; x + convert($(tensortype(Space(size(x)))), y))
@generated Base.:+(x::AbstractArray, y::AbstractTensor) = :(@_inline_meta; convert($(tensortype(Space(size(y)))), x) + y)
@generated Base.:-(x::AbstractTensor, y::AbstractArray) = :(@_inline_meta; x - convert($(tensortype(Space(size(x)))), y))
@generated Base.:-(x::AbstractArray, y::AbstractTensor) = :(@_inline_meta; convert($(tensortype(Space(size(y)))), x) - y)

@inline _add_uniform(x::AbstractSquareTensor, λ::Real) = x + λ * one(x)
@inline Base.:+(x::AbstractTensor, y::UniformScaling) = _add_uniform( x,  y.λ)
@inline Base.:-(x::AbstractTensor, y::UniformScaling) = _add_uniform( x, -y.λ)
@inline Base.:+(x::UniformScaling, y::AbstractTensor) = _add_uniform( y,  x.λ)
@inline Base.:-(x::UniformScaling, y::AbstractTensor) = _add_uniform(-y,  x.λ)
@inline dot(x::Union{AbstractVec, AbstractMat, AbstractSquareTensor}, y::UniformScaling) = x * y.λ
@inline dot(x::UniformScaling, y::Union{AbstractVec, AbstractMat, AbstractSquareTensor}) = x.λ * y
@inline double_contraction(x::AbstractSquareTensor, y::UniformScaling) = tr(x) * y.λ
@inline double_contraction(x::UniformScaling, y::AbstractSquareTensor) = x.λ * tr(y)

# error for standard multiplications
error_multiply() = error("use `⋅` (`\\cdot`) for single contraction and `⊡` (`\\boxdot`) for double contraction instead of `*`")
Base.:*(::AbstractTensor, ::AbstractTensor) = error_multiply()
Base.:*(::AbstractTensor, ::UniformScaling) = error_multiply()
Base.:*(::UniformScaling, ::AbstractTensor) = error_multiply()

function contraction_exprs(x::Type{<: AbstractTensor}, y::Type{<: AbstractTensor}, N::Int)
    nx = ndims(x)
    ny = ndims(y)
    ij_x = UnitRange(1, nx)
    ij_y = UnitRange(nx + 1 - N, nx + ny - N)
    freeinds = find_freeindices([ij_x; ij_y])
    einsum_contraction_expr(Tuple(freeinds), (x, y), (ij_x, ij_y))
end

"""
    contraction(::AbstractTensor, ::AbstractTensor, ::Val{N})

Conduct contraction of `N` inner indices.
For example, `N=2` contraction for third-order tensors ``A_{ij} = B_{ikl} C_{klj}``
can be computed in Tensorial.jl as

```jldoctest
julia> B = rand(Tensor{Tuple{3,3,3}});

julia> C = rand(Tensor{Tuple{3,3,3}});

julia> A = contraction(B, C, Val(2))
3×3 Tensor{Tuple{3, 3}, Float64, 2, 9}:
 1.36912   1.86751  1.32531
 1.61744   2.34426  1.94101
 0.929252  1.89656  1.79015
```

Following symbols are also available for specific contractions:

- `x ⊗ y` (where `⊗` can be typed by `\\otimes<tab>`): `contraction(x, y, Val(0))`
- `x ⋅ y` (where `⋅` can be typed by `\\cdot<tab>`): `contraction(x, y, Val(1))`
- `x ⊡ y` (where `⊡` can be typed by `\\boxdot<tab>`): `contraction(x, y, Val(2))`
"""
@generated function contraction(x::AbstractTensor, y::AbstractTensor, ::Val{N}) where {N}
    TT, exps = contraction_exprs(x, y, N)
    quote
        @_inline_meta
        tensors = (x, y)
        @inbounds $TT($(exps...))
    end
end

"""
    otimes(x::AbstractTensor, y::AbstractTensor)
    x ⊗ y

Compute tensor product such as ``A_{ij} = x_i y_j``.
`x ⊗ y` (where `⊗` can be typed by `\\otimes<tab>`) is a synonym for `otimes(x, y)`.

# Examples
```jldoctest
julia> x = rand(Vec{3})
3-element Vec{3, Float64}:
 0.5908446386657102
 0.7667970365022592
 0.5662374165061859

julia> y = rand(Vec{3})
3-element Vec{3, Float64}:
 0.4600853424625171
 0.7940257103317943
 0.8541465903790502

julia> A = x ⊗ y
3×3 Tensor{Tuple{3, 3}, Float64, 2, 9}:
 0.271839  0.469146  0.504668
 0.352792  0.608857  0.654957
 0.260518  0.449607  0.48365
```
"""
@inline otimes(x1::AbstractTensor, x2::AbstractTensor) = contraction(x1, x2, Val(0))

"""
    dot(x::AbstractTensor, y::AbstractTensor)
    x ⋅ y

Compute dot product such as ``a = x_i y_i``.
This is equivalent to [`contraction(::AbstractTensor, ::AbstractTensor, Val(1))`](@ref).
`x ⋅ y` (where `⋅` can be typed by `\\cdot<tab>`) is a synonym for `dot(x, y)`.

# Examples
```jldoctest
julia> x = rand(Vec{3})
3-element Vec{3, Float64}:
 0.5908446386657102
 0.7667970365022592
 0.5662374165061859

julia> y = rand(Vec{3})
3-element Vec{3, Float64}:
 0.4600853424625171
 0.7940257103317943
 0.8541465903790502

julia> a = x ⋅ y
1.3643452781654772
```
"""
@inline dot(x1::AbstractTensor, x2::AbstractTensor) = contraction(x1, x2, Val(1))
@inline double_contraction(x1::AbstractTensor, x2::AbstractTensor) = contraction(x1, x2, Val(2))

"""
    norm(::AbstractTensor)

Compute norm of a tensor.

# Examples
```jldoctest
julia> x = rand(Mat{3, 3})
3×3 Tensor{Tuple{3, 3}, Float64, 2, 9}:
 0.590845  0.460085  0.200586
 0.766797  0.794026  0.298614
 0.566237  0.854147  0.246837

julia> norm(x)
1.7377443667834922
```
"""
@inline norm(x::AbstractTensor) = sqrt(contraction(x, x, Val(ndims(x))))

@inline normalize(x::AbstractTensor) = x / norm(x)

# v_k * S_ikjl * u_l
@inline function dotdot(v1::Vec{dim}, S::SymmetricFourthOrderTensor{dim}, v2::Vec{dim}) where {dim}
    S′ = SymmetricFourthOrderTensor{dim}((i,j,k,l) -> @inbounds S[j,i,l,k])
    v1 ⋅ S′ ⋅ v2
end

"""
    tr(::AbstractSecondOrderTensor)
    tr(::AbstractSymmetricSecondOrderTensor)

Compute the trace of a square tensor.

# Examples
```jldoctest
julia> x = rand(Mat{3,3})
3×3 Tensor{Tuple{3, 3}, Float64, 2, 9}:
 0.590845  0.460085  0.200586
 0.766797  0.794026  0.298614
 0.566237  0.854147  0.246837

julia> tr(x)
1.6317075356075135
```
"""
@inline tr(x::AbstractSquareTensor) = @einsum x[i,i]

# symmetric
@inline symmetric(x::AbstractSymmetricSecondOrderTensor) = x
@inline function symmetric(x::AbstractSymmetricSecondOrderTensor, uplo::Symbol)
    if uplo == :U || uplo == :L
        return x
    end
    throw(ArgumentError("uplo argument must be either :U (upper) or :L (lower)"))
end
## :U or :L
for (uplo, filter) in ((:U, (inds,i,j) -> i ≤ j ? inds[i,j] : inds[j,i]),
                       (:L, (inds,i,j) -> i ≥ j ? inds[i,j] : inds[j,i]))
    @eval @generated function $(Symbol(:symmetric, uplo))(x::AbstractSecondOrderTensor{dim}) where {dim}
        exps = [:(Tuple(x)[$($filter(independent_indices(x), Tuple(I)...))]) for I in CartesianIndices(x)]
        TT = SymmetricSecondOrderTensor{dim}
        quote
            @_inline_meta
            @inbounds $TT($(exps[indices(TT)]...))
        end
    end
end
@inline function symmetric(x::AbstractSecondOrderTensor{dim}, uplo::Symbol) where {dim}
    uplo == :U && return symmetricU(x)
    uplo == :L && return symmetricL(x)
    throw(ArgumentError("uplo argument must be either :U (upper) or :L (lower)"))
end
## no uplo
@inline symmetric(x::AbstractSecondOrderTensor{dim}) where {dim} =
    SymmetricSecondOrderTensor{dim}((i,j) -> @inbounds i == j ? x[i,j] : (x[i,j] + x[j,i]) / 2)
@inline symmetric(x::AbstractSecondOrderTensor{1}) =
    @inbounds SymmetricSecondOrderTensor{1}(x[1,1])
@inline symmetric(x::AbstractSecondOrderTensor{2}) =
    @inbounds SymmetricSecondOrderTensor{2}(x[1,1], (x[2,1]+x[1,2])/2, x[2,2])
@inline symmetric(x::AbstractSecondOrderTensor{3}) =
    @inbounds SymmetricSecondOrderTensor{3}(x[1,1], (x[2,1]+x[1,2])/2, (x[3,1]+x[1,3])/2, x[2,2], (x[3,2]+x[2,3])/2, x[3,3])

"""
    skew(::AbstractSecondOrderTensor)
    skew(::AbstractSymmetricSecondOrderTensor)

Compute skew-symmetric (anti-symmetric) part of a second order tensor.
"""
@inline skew(x::AbstractSecondOrderTensor) = (x - x') / 2
@inline skew(x::AbstractSymmetricSecondOrderTensor{dim, T}) where {dim, T} = zero(SecondOrderTensor{dim, T})

"""
    skew(ω::Vec{3})

Construct a skew-symmetric (anti-symmetric) tensor `W` from a vector `ω` as

```math
\\bm{\\omega} = \\begin{Bmatrix}
    \\omega_1 \\\\
    \\omega_2 \\\\
    \\omega_3
\\end{Bmatrix}, \\quad
\\bm{W} = \\begin{bmatrix}
     0         & -\\omega_3 &  \\omega_2 \\\\
     \\omega_3 & 0          & -\\omega_1 \\\\
    -\\omega_2 &  \\omega_1 &  0
\\end{bmatrix}
```

# Examples
```jldoctest
julia> skew(Vec(1,2,3))
3×3 Tensor{Tuple{3, 3}, Int64, 2, 9}:
  0  -3   2
  3   0  -1
 -2   1   0
```
"""
@inline function skew(ω::Vec{3})
    z = zero(eltype(ω))
    @inbounds @Mat [ z    -ω[3]  ω[2]
                     ω[3]     z -ω[1]
                    -ω[2]  ω[1]     z]
end

# transpose/adjoint
@inline transpose(x::AbstractTensor{Tuple{@Symmetry{dim, dim}}}) where {dim} = x
@inline transpose(x::AbstractTensor{Tuple{m, n}}) where {m, n} = Tensor{Tuple{n, m}}((i,j) -> @inbounds x[j,i])
@inline adjoint(x::AbstractTensor) = transpose(x)

# det
@generated function extract_vecs(x::AbstractSquareTensor{dim}) where {dim}
    exps = map(1:dim) do j
        :(Vec($([getindex_expr(x, :x, i, j) for i in 1:dim]...)))
    end
    quote
        @_inline_meta
        @inbounds tuple($(exps...))
    end
end
@inline function det(x::AbstractSquareTensor{1})
    @inbounds x[1,1]
end
@inline function det(x::AbstractSquareTensor{2})
    @inbounds x[1,1] * x[2,2] - x[1,2] * x[2,1]
end
@inline function det(A::AbstractSquareTensor{3})
    a1, a2, a3 = extract_vecs(A)
    (a1 × a2) ⋅ a3
end
@inline function det(x::AbstractSquareTensor{dim}) where {dim}
    det(SMatrix{dim, dim}(x))
end

"""
    cross(x::Vec{3}, y::Vec{3}) -> Vec{3}
    cross(x::Vec{2}, y::Vec{2}) -> Vec{3}
    cross(x::Vec{1}, y::Vec{1}) -> Vec{3}
    x × y

Compute the cross product between two vectors.
The vectors are expanded to 3D frist for dimensions 1 and 2.
The infix operator `×` (written `\\times`) can also be used.
`x × y` (where `×` can be typed by `\\times<tab>`) is a synonym for `cross(x, y)`.

# Examples
```jldoctest
julia> x = rand(Vec{3})
3-element Vec{3, Float64}:
 0.5908446386657102
 0.7667970365022592
 0.5662374165061859

julia> y = rand(Vec{3})
3-element Vec{3, Float64}:
 0.4600853424625171
 0.7940257103317943
 0.8541465903790502

julia> x × y
3-element Vec{3, Float64}:
  0.20535000738340053
 -0.24415039787171888
  0.11635375677388776
```
"""
@inline cross(x::Vec{1, T1}, y::Vec{1, T2}) where {T1, T2} = zero(Vec{3, promote_type(T1, T2)})
@inline function cross(x::Vec{2, T1}, y::Vec{2, T2}) where {T1, T2}
    z = zero(promote_type(T1, T2))
    @inbounds Vec(z, z, x[1]*y[2] - x[2]*y[1])
end
@inline function cross(x::Vec{3}, y::Vec{3})
    @inbounds Vec(x[2]*y[3] - x[3]*y[2],
                  x[3]*y[1] - x[1]*y[3],
                  x[1]*y[2] - x[2]*y[1])
end

# power
@inline Base.literal_pow(::typeof(^), x::AbstractSquareTensor, ::Val{-1}) = inv(x)
@inline Base.literal_pow(::typeof(^), x::AbstractSquareTensor, ::Val{0})  = one(x)
@inline Base.literal_pow(::typeof(^), x::AbstractSquareTensor, ::Val{1})  = x
@inline function Base.literal_pow(::typeof(^), x::AbstractSquareTensor, ::Val{p}) where {p}
    p > 0 ? (y = x; q = p) : (y = inv(x); q = -p)
    z = y
    for i in 2:q
        y = _powdot(y, z)
    end
    y
end
## helper functions
@inline _powdot(x::AbstractSecondOrderTensor, y::AbstractSecondOrderTensor) = dot(x, y)
@generated function _powdot(x::AbstractSymmetricSecondOrderTensor{dim}, y::AbstractSymmetricSecondOrderTensor{dim}) where {dim}
    _, exps = contraction_exprs(x, y, 1)
    quote
        @_inline_meta
        tensors = (x, y)
        @inbounds SymmetricSecondOrderTensor{dim}($(exps[indices(x)]...))
    end
end

# rotate
"""
    rotmat(θ::Real; degree::Bool = false)

Construct 2D rotation matrix.

# Examples
```jldoctest
julia> rotmat(30, degree = true)
2×2 Tensor{Tuple{2, 2}, Float64, 2, 4}:
 0.866025  -0.5
 0.5        0.866025
```
"""
@inline function rotmat(θ::Real; degree::Bool = false)
    if degree
        θ = deg2rad(θ)
    end
    sinθ = sin(θ)
    cosθ = cos(θ)
    @Mat [cosθ -sinθ
          sinθ  cosθ]
end

"""
    rotmat(θ::Vec{3}; sequence::Symbol, degree::Bool = false)
    rotmatx(θ::Real)
    rotmaty(θ::Real)
    rotmatz(θ::Real)

Convert Euler angles to rotation matrix.
Use 3 characters belonging to the set (X, Y, Z) for intrinsic rotations,
or (x, y, z) for extrinsic rotations.

# Examples
```jldoctest
julia> α, β, γ = rand(Vec{3});

julia> rotmat(Vec(α,β,γ), sequence = :XYZ) ≈ rotmatx(α) ⋅ rotmaty(β) ⋅ rotmatz(γ)
true

julia> rotmat(Vec(α,β,γ), sequence = :xyz) ≈ rotmatz(γ) ⋅ rotmaty(β) ⋅ rotmatx(α)
true

julia> rotmat(Vec(α,β,γ), sequence = :XYZ) ≈ rotmat(Vec(γ,β,α), sequence = :zyx)
true
```
"""
function rotmat(θ::Vec{3}; sequence::Symbol, degree::Bool = false)
    @inbounds α, β, γ = θ[1], θ[2], θ[3]
    if degree
        α, β, γ = map(deg2rad, (α, β, γ))
    end
    # intrinsic
    sequence == :XZX && return rotmatx(α) ⋅ rotmaty(β) ⋅ rotmatz(γ)
    sequence == :XYX && return rotmatx(α) ⋅ rotmaty(β) ⋅ rotmatx(γ)
    sequence == :YXY && return rotmaty(α) ⋅ rotmatx(β) ⋅ rotmaty(γ)
    sequence == :YZY && return rotmaty(α) ⋅ rotmatz(β) ⋅ rotmaty(γ)
    sequence == :ZYZ && return rotmatz(α) ⋅ rotmaty(β) ⋅ rotmatz(γ)
    sequence == :ZXZ && return rotmatz(α) ⋅ rotmatx(β) ⋅ rotmatz(γ)
    sequence == :XZY && return rotmatx(α) ⋅ rotmatz(β) ⋅ rotmaty(γ)
    sequence == :XYZ && return rotmatx(α) ⋅ rotmaty(β) ⋅ rotmatz(γ)
    sequence == :YXZ && return rotmaty(α) ⋅ rotmatx(β) ⋅ rotmatz(γ)
    sequence == :YZX && return rotmaty(α) ⋅ rotmatz(β) ⋅ rotmatx(γ)
    sequence == :ZYX && return rotmatz(α) ⋅ rotmaty(β) ⋅ rotmatx(γ)
    sequence == :ZXY && return rotmatz(α) ⋅ rotmatx(β) ⋅ rotmaty(γ)
    # extrinsic
    sequence == :xzx && return rotmatx(γ) ⋅ rotmaty(β) ⋅ rotmatz(α)
    sequence == :xyx && return rotmatx(γ) ⋅ rotmaty(β) ⋅ rotmatx(α)
    sequence == :yxy && return rotmaty(γ) ⋅ rotmatx(β) ⋅ rotmaty(α)
    sequence == :yzy && return rotmaty(γ) ⋅ rotmatz(β) ⋅ rotmaty(α)
    sequence == :zyz && return rotmatz(γ) ⋅ rotmaty(β) ⋅ rotmatz(α)
    sequence == :zxz && return rotmatz(γ) ⋅ rotmatx(β) ⋅ rotmatz(α)
    sequence == :xzy && return rotmaty(γ) ⋅ rotmatz(β) ⋅ rotmatx(α)
    sequence == :xyz && return rotmatz(γ) ⋅ rotmaty(β) ⋅ rotmatx(α)
    sequence == :yxz && return rotmatz(γ) ⋅ rotmatx(β) ⋅ rotmaty(α)
    sequence == :yzx && return rotmatx(γ) ⋅ rotmatz(β) ⋅ rotmaty(α)
    sequence == :zyx && return rotmatx(γ) ⋅ rotmaty(β) ⋅ rotmatz(α)
    sequence == :zxy && return rotmaty(γ) ⋅ rotmatx(β) ⋅ rotmatz(α)
    throw(ArgumentError("sequence $sequence is not supported"))
end
@inline function rotmatx(θ::Real; degree::Bool = false)
    if degree
        θ = deg2rad(θ)
    end
    o = one(θ)
    z = zero(θ)
    sinθ = sin(θ)
    cosθ = cos(θ)
    @Mat [o z     z
          z cosθ -sinθ
          z sinθ  cosθ]
end
@inline function rotmaty(θ::Real; degree::Bool = false)
    if degree
        θ = deg2rad(θ)
    end
    o = one(θ)
    z = zero(θ)
    sinθ = sin(θ)
    cosθ = cos(θ)
    @Mat [ cosθ z sinθ
           z    o z
          -sinθ z cosθ]
end
@inline function rotmatz(θ::Real; degree::Bool = false)
    if degree
        θ = deg2rad(θ)
    end
    o = one(θ)
    z = zero(θ)
    sinθ = sin(θ)
    cosθ = cos(θ)
    @Mat [cosθ -sinθ z
          sinθ  cosθ z
          z     z    o]
end

"""
    rotmat(a => b)

Construct rotation matrix rotating vector `a` to `b`.
The norms of two vectors must be the same.

# Examples
```jldoctest
julia> a = normalize(rand(Vec{3}))
3-element Vec{3, Float64}:
 0.526847334217759
 0.683741457787621
 0.5049054419691867

julia> b = normalize(rand(Vec{3}))
3-element Vec{3, Float64}:
 0.36698690362212083
 0.6333543148133657
 0.6813097125956302

julia> R = rotmat(a => b)
3×3 Tensor{Tuple{3, 3}, Float64, 2, 9}:
 -0.594528   0.597477   0.538106
  0.597477  -0.119597   0.792917
  0.538106   0.792917  -0.285875

julia> R ⋅ a ≈ b
true
```
"""
function rotmat(pair::Pair{Vec{dim, T}, Vec{dim, T}})::Mat{dim, dim, T} where {dim, T}
    # https://math.stackexchange.com/questions/180418/calculate-rotation-matrix-to-align-vector-a-to-vector-b-in-3d/2672702#2672702
    a = pair.first
    b = pair.second
    dot(a, a) ≈ dot(b, b) || throw(ArgumentError("the norms of two vectors must be the same"))
    a ==  b && return  one(Mat{dim, dim, T})
    a == -b && return -one(Mat{dim, dim, T})
    c = a + b
    2 * (c ⊗ c) / (c ⋅ c) - one(Mat{dim, dim, T})
end

"""
    rotmat(θ, n::Vec; degree::Bool = false)

Construct rotation matrix from angle `θ` and direction `n`.

# Examples
```jldoctest
julia> x = Vec(1.0, 0.0, 0.0)
3-element Vec{3, Float64}:
 1.0
 0.0
 0.0

julia> n = Vec(0.0, 0.0, 1.0)
3-element Vec{3, Float64}:
 0.0
 0.0
 1.0

julia> rotmat(π/2, n) ⋅ x
3-element Vec{3, Float64}:
 1.1102230246251565e-16
 1.0
 0.0
```
"""
function rotmat(θ::Real, x::Vec{3}; degree::Bool = false)
    if degree
        θ = deg2rad(θ)
    end
    n = normalize(x)
    W = skew(n)
    I + W*sin(θ) + W^2*(1-cos(θ))
end

"""
    rotate(x::Vec, R::SecondOrderTensor)
    rotate(x::SecondOrderTensor, R::SecondOrderTensor)
    rotate(x::SymmetricSecondOrderTensor, R::SecondOrderTensor)

Rotate `x` by rotation matrix `R`.
This function can hold the symmetry of `SymmetricSecondOrderTensor`.

# Examples
```jldoctest
julia> A = rand(SymmetricSecondOrderTensor{3})
3×3 SymmetricSecondOrderTensor{3, Float64, 6}:
 0.590845  0.766797  0.566237
 0.766797  0.460085  0.794026
 0.566237  0.794026  0.854147

julia> R = rotmatz(π/4)
3×3 Tensor{Tuple{3, 3}, Float64, 2, 9}:
 0.707107  -0.707107  0.0
 0.707107   0.707107  0.0
 0.0        0.0       1.0

julia> rotate(A, R)
3×3 SymmetricSecondOrderTensor{3, Float64, 6}:
 -0.241332   0.0653796  -0.161071
  0.0653796  1.29226     0.961851
 -0.161071   0.961851    0.854147

julia> R ⋅ A ⋅ R'
3×3 Tensor{Tuple{3, 3}, Float64, 2, 9}:
 -0.241332   0.0653796  -0.161071
  0.0653796  1.29226     0.961851
 -0.161071   0.961851    0.854147
```
"""
@inline rotate(v::Vec, R::SecondOrderTensor) = R ⋅ v
@inline rotate(v::Vec{2}, R::SecondOrderTensor{3}) = (R2x2 = Mat{2, 2}((i,j) -> @inbounds R[i,j]); R2x2 ⋅ v)
@inline rotate(A::SecondOrderTensor, R::SecondOrderTensor) = @einsum R[i,j] * A[j,k] * R[l,k]
@generated function rotate(A::SymmetricSecondOrderTensor{dim}, R::SecondOrderTensor{dim}) where {dim}
    _, exps = contraction_exprs(SecondOrderTensor{dim}, SecondOrderTensor{dim}, 1)
    TT = SymmetricSecondOrderTensor{dim}
    quote
        @_inline_meta
        ARᵀ = @einsum A[i,j] * R[k,j]
        tensors = (R, ARᵀ)
        @inbounds $TT($(exps[indices(TT)]...))
    end
end

function angleaxis(R::SecondOrderTensor{3})
    # https://math.stackexchange.com/questions/893984/conversion-of-rotation-matrix-to-quaternion
    θ = acos((tr(R)-1) / 2)
    n = Tensor(Real.(eigvecs(SArray(R))[:,3]))
    θ, n
end

# ----------------------------------------------#
# operations calling methods in StaticArrays.jl #
# ----------------------------------------------#

# eigvals/eigen
@inline function eigvals(x::AbstractSymmetricSecondOrderTensor; permute::Bool = true, scale::Bool = true)
    Tensor(eigvals(Symmetric(SArray(x)); permute, scale))
end
@inline function eigen(x::AbstractSymmetricSecondOrderTensor)
    eig = eigen(Symmetric(SArray(x)))
    Eigen(Tensor(eig.values), Tensor(eig.vectors))
end

# exp
@inline Base.exp(x::AbstractSecondOrderTensor) = typeof(x)(exp(SArray(x)))
@inline Base.exp(x::AbstractSymmetricSecondOrderTensor) = typeof(x)(exp(Symmetric(SArray(x))))

# diag/diagm
@inline diag(x::Union{AbstractMat, AbstractSymmetricSecondOrderTensor}, ::Val{k} = Val(0)) where {k} = Tensor(diag(SArray(x), Val{k}))
@inline diagm(kvs::Pair{<: Val, <: AbstractVec}...) = Tensor(diagm(map(kv -> kv.first => SArray(kv.second), kvs)...))
@inline diagm(x::Vec) = diagm(Val(0) => x)

# qr
struct QR{Q, R, P}
    Q::Q
    R::R
    p::P
end

# iteration for destructuring into components
Base.iterate(S::QR) = (S.Q, Val(:R))
Base.iterate(S::QR, ::Val{:R}) = (S.R, Val(:p))
Base.iterate(S::QR, ::Val{:p}) = (S.p, Val(:done))
Base.iterate(S::QR, ::Val{:done}) = nothing

for pv in (:true, :false)
    @eval function qr(A::AbstractMat, pivot::Val{$pv})
        F = qr(SArray(A), pivot)
        QR(Tensor(F.Q), Tensor(F.R), Tensor(F.p))
    end
end
qr(A::AbstractMat) = qr(A, Val(false))

# lu
struct LU{L, U, p}
    L::L
    U::U
    p::p
end

# iteration for destructuring into components
Base.iterate(S::LU) = (S.L, Val(:U))
Base.iterate(S::LU, ::Val{:U}) = (S.U, Val(:p))
Base.iterate(S::LU, ::Val{:p}) = (S.p, Val(:done))
Base.iterate(S::LU, ::Val{:done}) = nothing

for pv in (:true, :false)
    @eval function lu(A::AbstractMat, pivot::Val{$pv}; check = true)
        F = lu(SArray(A), pivot; check)
        LU(LowerTriangular(Tensor(parent(F.L))), UpperTriangular(Tensor(parent(F.U))), Tensor(F.p))
    end
end
lu(A::AbstractMat; check = true) = lu(A, Val(true); check)

# svd
struct SVD{T, TU, TS, TVt} <: Factorization{T}
    U::TU
    S::TS
    Vt::TVt
end
SVD(U::AbstractArray{T}, S::AbstractVector, Vt::AbstractArray{T}) where {T} = SVD{T, typeof(U), typeof(S), typeof(Vt)}(U, S, Vt)

@inline function Base.getproperty(F::SVD, s::Symbol)
    if s === :V
        return getfield(F, :Vt)'
    else
        return getfield(F, s)
    end
end
Base.propertynames(::SVD) = (:U, :S, :V, :Vt)

# iteration for destructuring into components
Base.iterate(S::SVD) = (S.U, Val(:S))
Base.iterate(S::SVD, ::Val{:S}) = (S.S, Val(:V))
Base.iterate(S::SVD, ::Val{:V}) = (S.V, Val(:done))
Base.iterate(S::SVD, ::Val{:done}) = nothing

function svd(A::AbstractMat; full=Val(false))
    F = svd(SArray(A); full)
    SVD(Tensor(F.U), Tensor(F.S), Tensor(F.Vt))
end


using CompScienceMeshes, BEAST
o, x, y, z = euclidianbasis(3)

Γ = readmesh(joinpath(dirname(@__FILE__),"sphere2.in"))
X, Y = raviartthomas(Γ), buffachristiansen(Γ)

ω = κ = 1.0
μ = ϵ = 1.0

K, N = MWDoubleLayer3D(im*κ), NCross()

E = PlaneWaveMW(z, x, κ, complex(1.0))
H = -1/(im*μ*ω)*curl(E)
h = (n × H) × n

@hilbertspace j; @hilbertspace m;
mfie = @discretise (K+0.5N)[m,j] == h[m]  j∈X m∈Y
u = solve(mfie)

using PlotlyJS
include(Pkg.dir("CompScienceMeshes","examples","plotlyjs_patches.jl"))

fcr, geo = facecurrents(u, X)
p = patch(geo, real.(norm.(fcr)))


module Jeeves
    using LinearAlgebra, Combinatorics
    using Optim, Distributions
    using DataFrames: DataFrame
    using TexTables, OrderedCollections
    using RDatasets

    export fit
    export fit!
    export coef
    export summary
    export tidy

    export vcovIID, vcovCluster, vcovHetero, vcovBoot

    ## Types ##
    abstract type Model end
    abstract type LinearModel <: Model  end
    abstract type GeneralisedLinearModel <: Model end

    abstract type Fit <: Model end
    abstract type LinearModelFit <: Fit end

    abstract type vcov end

    ## Modules ##
    include("model.jl")
    include("ols.jl")
    include("tsls.jl")
    include("probit.jl")
    include("vcov.jl")
    include("vcovCluster.jl")
    include("vcovHetero.jl")
    include("vcovBoot.jl")
    include("utils.jl")
end # module


using DifferentialEquations, DataFrames
using Plots
using CSV
include("odeFileExample.jl")
include("variableNames.jl")
include("scanIncludes.jl")

f=DDEFunction(ddeFile,syms=Symbol.(syms))
#f=ODEFunction(odeFile,syms=Symbol.(syms))
y0=ones(length(syms))
p=1
h(p,t)=ones(length(syms))
params=1
prob=DDEProblem(f,ones(length(syms)),h,(0.0,10.0),syms=Symbol.(syms))
#prob=ODEProblem(f,ones(length(syms)),(0.0,10.0))
print("Solving Equations \n")
@timev sol=solve(prob)
plot(sol)

modify["k_ABBinding"]=1.5

prob=DDEProblem(f,ones(length(syms)),h,(0.0,10.0),syms=Symbol.(syms))
print("Solving Equations \n")
@timev sol=solve(prob)
plot!(sol)
png("testOutScan.png")


function mdl_line_ode!(
    device_states,
    output_ode,
    V_r_from,
    V_i_from,
    V_r_to,
    V_i_to,
    current_r_from,
    current_i_from,
    current_r_to,
    current_i_to,
    sys_f,
    branch::PSY.DynamicBranch,
)
    L = PSY.get_x(branch)
    R = PSY.get_r(branch)
    ω_b = sys_f * 2 * π

    Il_r = device_states[1]
    Il_i = device_states[2]
    output_ode[1] = (ω_b / L) * ((V_r_from[1] - V_r_to[1]) - (R * Il_r - L * Il_i))
    output_ode[2] = (ω_b / L) * ((V_i_from[1] - V_i_to[1]) - (R * Il_i + L * Il_r))

    current_r_from[1] -= Il_r
    current_i_from[1] -= Il_i
    current_r_to[1] += Il_r
    current_i_to[1] += Il_i
end


# poweriter.jl

using MIRT: poweriter

using Random: seed!
using LinearAlgebra: opnorm, norm
using Test: @test, @inferred

seed!(0)
M,N = 30,20

A = randn(M,N) # real
s0 = opnorm(A)
chat = false
_,s1 = @inferred poweriter(A ; tol=1e-9, chat=chat)
@test s0 ≈ s1

A = randn(ComplexF32, M, N) # complex
s0 = opnorm(A)
_,s1 = @inferred poweriter(A ; tol=1e-9, chat=chat)
@test s0 ≈ s1


"""
    tf!(dtm::AbstractMatrix{Real}, tf::AbstractMatrix{AbstractFloat})

Overwrite `tf` with the term frequency of the `dtm`.

Works correctly if `dtm` and `tf` are same matrix.

See also: [`tf`](@ref), [`tf_idf`](@ref), [`tf_idf!`](@ref)
"""
function tf!(dtm::AbstractMatrix{T1}, tf::AbstractMatrix{T2}) where {T1 <: Real, T2 <: AbstractFloat}
    n, p = size(dtm)

    # TF tells us what proportion of a document is defined by a term
    for i in 1:n
        words_in_document = 0
        for j in 1:p
            words_in_document += dtm[i, j]
        end
        tf[i, :] = dtm[i, :] ./ max(words_in_document, one(T1))
    end

    return tf
end

"""
    tf!(dtm::SparseMatrixCSC{Real}, tf::SparseMatrixCSC{AbstractFloat})

Overwrite `tf` with the term frequency of the `dtm`.

`tf` should have the has same nonzeros as `dtm`.

See also: [`tf`](@ref), [`tf_idf`](@ref), [`tf_idf!`](@ref)
"""
function tf!(dtm::SparseMatrixCSC{T}, tf::SparseMatrixCSC{F}) where {T <: Real, F <: AbstractFloat}
    rows = rowvals(dtm)
    dtmvals = nonzeros(dtm)
    tfvals = nonzeros(tf)
    @assert size(dtmvals) == size(tfvals)

    # TF tells us what proportion of a document is defined by a term
    words_in_documents = sum(dtm,dims=2)

    n, p = size(dtm)
    for i = 1:p
       for j in nzrange(dtm, i)
          row = rows[j]
          tfvals[j] = dtmvals[j] / max(words_in_documents[row], one(T))
       end
    end
    return tf
end

tf!(dtm::AbstractMatrix{T}) where {T <: Real} = tf!(dtm, dtm)

tf!(dtm::SparseMatrixCSC{T}) where {T <: Real} = tf!(dtm, dtm)

"""
    tf(dtm::DocumentTermMatrix)
    tf(dtm::SparseMatrixCSC{Real})
    tf(dtm::Matrix{Real})

Compute the `term-frequency` of the input.

# Example

```julia-repl
julia> crps = Corpus([StringDocument("To be or not to be"),
              StringDocument("To become or not to become")])

julia> update_lexicon!(crps)

julia> m = DocumentTermMatrix(crps)

julia> tf(m)
2×6 SparseArrays.SparseMatrixCSC{Float64,Int64} with 10 stored entries:
  [1, 1]  =  0.166667
  [2, 1]  =  0.166667
  [1, 2]  =  0.333333
  [2, 3]  =  0.333333
  [1, 4]  =  0.166667
  [2, 4]  =  0.166667
  [1, 5]  =  0.166667
  [2, 5]  =  0.166667
  [1, 6]  =  0.166667
  [2, 6]  =  0.166667
```

See also: [`tf!`](@ref), [`tf_idf`](@ref), [`tf_idf!`](@ref)
"""
tf(dtm::DocumentTermMatrix) = tf(dtm.dtm)

tf(dtm::Matrix{T}) where {T <: Real} = tf!(dtm, Array{Float64}(undef, size(dtm)...))

tf(dtm::SparseMatrixCSC{T}) where {T <: Real} =  tf!(dtm, similar(dtm, Float64))

"""
    tf_idf!(dtm::AbstractMatrix{Real}, tf_idf::AbstractMatrix{AbstractFloat})

Overwrite `tf_idf` with the tf-idf (Term Frequency - Inverse Doc Frequency) of the `dtm`.

`dtm` and `tf-idf` must be matrices of same dimensions.

See also: [`tf`](@ref), [`tf!`](@ref) , [`tf_idf`](@ref)
"""
function tf_idf!(dtm::AbstractMatrix{T1}, tfidf::AbstractMatrix{T2}) where {T1 <: Real, T2 <: AbstractFloat}
    n, p = size(dtm)

    # TF tells us what proportion of a document is defined by a term
    tf!(dtm, tfidf)

    # IDF tells us how rare a term is in the corpus
    documents_containing_term = vec(sum(dtm .> 0, dims=1))
    idf = log.(n ./ documents_containing_term)

    # TF-IDF is the product of TF and IDF
    for i in 1:n
        for j in 1:p
           tfidf[i, j] = tfidf[i, j] * idf[j]
        end
    end

    return tfidf
end

"""
    tf_idf!(dtm::SparseMatrixCSC{Real}, tfidf::SparseMatrixCSC{AbstractFloat})

Overwrite `tfidf` with the tf-idf (Term Frequency - Inverse Doc Frequency) of the `dtm`.

The arguments must have same number of nonzeros.

See also: [`tf`](@ref), [`tf_idf`](@ref), [`tf_idf!`](@ref)
"""
function tf_idf!(dtm::SparseMatrixCSC{T}, tfidf::SparseMatrixCSC{F}) where {T <: Real, F <: AbstractFloat}
    rows = rowvals(dtm)
    dtmvals = nonzeros(dtm)
    tfidfvals = nonzeros(tfidf)
    @assert size(dtmvals) == size(tfidfvals)

    n, p = size(dtm)

    # TF tells us what proportion of a document is defined by a term
    words_in_documents = F.(sum(dtm, dims=2))
    oneval = one(F)

    # IDF tells us how rare a term is in the corpus
    documents_containing_term = vec(sum(dtm .> 0, dims=1))
    idf = log.(n ./ documents_containing_term)

    for i = 1:p
       for j in nzrange(dtm, i)
          row = rows[j]
          tfidfvals[j] = dtmvals[j] / max(words_in_documents[row], oneval) * idf[i]
       end
    end

    return tfidf
end

"""
    tf_idf!(dtm)

Compute tf-idf for `dtm`
"""
tf_idf!(dtm::AbstractMatrix{T}) where {T <: Real} = tf_idf!(dtm, dtm)

tf_idf!(dtm::SparseMatrixCSC{T}) where {T <: Real} = tf_idf!(dtm, dtm)

# This does not make sense, since DocumentTermMatrix is based on an array of integers
#tf_idf!(dtm::DocumentTermMatrix) = tf_idf!(dtm.dtm)

"""
    tf(dtm::DocumentTermMatrix)
    tf(dtm::SparseMatrixCSC{Real})
    tf(dtm::Matrix{Real})

Compute `tf-idf` value (Term Frequency - Inverse Document Frequency) for the input.

In many cases, raw word counts are not appropriate for use because:

- Some documents are longer than other documents
- Some words are more frequent than other words

A simple workaround this can be done by performing `TF-IDF` on a `DocumentTermMatrix`

# Example

```julia-repl
julia> crps = Corpus([StringDocument("To be or not to be"),
              StringDocument("To become or not to become")])

julia> update_lexicon!(crps)

julia> m = DocumentTermMatrix(crps)

julia> tf_idf(m)
2×6 SparseArrays.SparseMatrixCSC{Float64,Int64} with 10 stored entries:
  [1, 1]  =  0.0
  [2, 1]  =  0.0
  [1, 2]  =  0.231049
  [2, 3]  =  0.231049
  [1, 4]  =  0.0
  [2, 4]  =  0.0
  [1, 5]  =  0.0
  [2, 5]  =  0.0
  [1, 6]  =  0.0
  [2, 6]  =  0.0
```

See also: [`tf!`](@ref), [`tf_idf`](@ref), [`tf_idf!`](@ref)
"""
tf_idf(dtm::DocumentTermMatrix) = tf_idf(dtm.dtm)

tf_idf(dtm::SparseMatrixCSC{T}) where {T <: Real} =  tf_idf!(dtm, similar(dtm, Float64))

tf_idf(dtm::Matrix{T}) where {T <: Real} = tf_idf!(dtm, Array{Float64}(undef, size(dtm)...))


squared(x::Real) = x*x
squared(x::Vector) = x.*x
maxzero(x::Real) = max(x,zero(x))


@testset "Codegen" begin
    @testset "monomial_product" begin
        @test SP.monomial_product(Float64, [2, 3, 1], :(c[5]))[1] == :(c[5] * (x[1] * x[1]) * (x[2] * x[2] * x[2]) * x[3])
        #
        #
        @test SP.monomial_product(Float64, [2, 3, 1], :(c[5]), 2)[1] == :((3 * c[5] * (x[1] * x[1]) * (x[2] * x[2])) * x[3])

        # optimizations
        # no ^1
        @test SP.monomial_product(Float64, [1, 3, 1], :(c[5]))[1] == :(c[5] * x[1] * (x[2] * x[2] * x[2]) * x[3])
        # omit ^0
        @test SP.monomial_product(Float64, [1, 0, 1], :(c[5]))[1] == :(c[5] * x[1] * x[3])
        @test SP.monomial_product(Float64, [0, 0, 0], :(c[5]))[1] == :(c[5])
        # omit 1 * and ^0
        @test SP.monomial_product(Float64, [1, 3, 2], :(c[5]), 1)[1] == :(c[5] * (x[2] * x[2] * x[2]) * (x[3] * x[3]))
        # omit ^1 in derivative
        @test SP.monomial_product(Float64, [2, 2, 2], :(c[5]), 2)[1] == :((2 * c[5] * (x[1] * x[1]) * x[2]) * (x[3] * x[3]))
    end

    @testset "evaluate_codegen low_level" begin
        E = [ 4  4  1  3  5
              2  4  2  2  5
              0  1  2  2  2 ]
        degrees, subs = SP.degrees_submatrices(E)
        @test length(degrees) == 3
        @test degrees == [0, 1, 2]

        degrees2, subs2 = SP.degrees_submatrices(subs[3])
        @test subs2 == [[1 3], reshape([5], 1, 1)]
        @test degrees2 == [2, 5]

        @test SP.pow(:x, 0) == :(one(x))
    end

    @testset "evalpoly_deriv" begin
        z = rand(Complex{Float64})
        tol = 1e-12
        for T in [Float64, Complex{Float64}]
            c0, c1, c2, c3, c4, c5 = rand(T, 6)

            true_val = c0 + c1 * z + c2 * z^2 + c3 * z^3 + c4 * z^4 + c5 * z^5
            true_dval = c1 + 2c2 * z + 3c3 * z^2 + 4c4 * z^3 + 5c5 * z^4
            val, dval = SP.@goertzel_deriv z c0 c1 c2 c3 c4 c5
            @test abs(val - true_val) < tol
            @test abs(dval - true_dval) < tol
            val, dval = SP.@horner_deriv z c0 c1 c2 c3 c4 c5
            @test abs(val - true_val) < tol
            @test abs(dval - true_dval) < tol
            val, dval = SP.@evalpoly_derivative z c0 c1 c2 c3 c4 c5
            @test abs(val - true_val) < tol
            @test abs(dval - true_dval) < tol

            true_val = c0 + c1 * z + c2 * z^2
            true_dval = c1 + 2c2 * z
            val, dval = SP.@goertzel_deriv z c0 c1 c2
            @test abs(val - true_val) < tol
            @test abs(dval - true_dval) < tol
            val, dval = SP.@horner_deriv z c0 c1 c2
            @test abs(val - true_val) < tol
            @test abs(dval - true_dval) < tol

            true_val = c0 + c1 * z
            true_dval = c1
            val, dval = SP.@goertzel_deriv z c0 c1
            @test abs(val - true_val) < tol
            @test abs(dval - true_dval) < tol
            val, dval = SP.@horner_deriv z c0 c1
            @test abs(val - true_val) < tol
            @test abs(dval - true_dval) < tol
        end
    end

    @testset "gradient_codegen" begin
          # bugfix
          @polyvar a b t z
          f = -0.25a^5*t+0.25b^5*t+0.5z^6
          g = Polynomial(f)
          x = rand(4)
          vars = [a, b, t, z]
          @test map(v -> (MP.differentiate(f, v))(vars=>x), vars) ≈ gradient(g, x)
    end

end


using ASCIIPlots, Test

x = randn(20)
y = 21.0 * x + randn(20)
lineplot(x, y)


### A Pluto.jl notebook ###
# v0.17.2

# using Markdown
# using InteractiveUtils

# ╔═╡ 74975885-9a4e-4857-8135-9e4f69061caf
begin
	using DocStringExtensions
	using StaticArrays
	using LinearAlgebra
end

# ╔═╡ c0a30957-4c7b-4d7b-bfa9-c2fb691a077b
#=╠═╡ notebook_exclusive
begin
	using PlutoUtils
	using PlotlyBase
	using BenchmarkTools
end
  ╠═╡ notebook_exclusive =#

# ╔═╡ 379613ec-0973-4000-ae8c-d7c33ddca18e
#=╠═╡ notebook_exclusive
md"""
# Packages
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 736b0cf6-bec2-4226-8ef4-70f6a865d34a
#=╠═╡ notebook_exclusive
ToC()
  ╠═╡ notebook_exclusive =#

# ╔═╡ f8243a65-9f5e-464e-bb06-0bb4f5131b8b
#=╠═╡ notebook_exclusive
md"""
# Exports
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 8660a7c4-eb78-4e7c-966b-d759df7f3dfa
#=╠═╡ notebook_exclusive
md"""
# Lattice Functions
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 71163795-9695-4f11-acc2-6e3838c8a158
#=╠═╡ notebook_exclusive
md"""
## generate\_regular_lattice
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ da97848f-a7ff-4f2d-b98d-e8bf1ccc3038
function lattice_generator(dx::T, dy::T, ds::T;x0::T = T(0), y0::T = T(0), M::Int = 70,N::Int = M) where T<:Real
	# Function to generate x position as function of row,column number m,n
	x(m, n) = m * dx + n * ds + x0
	# Function to generate y position as function of row,column number m,n
	y(n) = n * dy + y0
	# Generate the elements. For each row, shift the columns to always have the search domain around x=0
	gen = ((x(m - round(Int,n * ds / dx), n), y(n)) for n in -N:N,m in -M:M)
	return gen
end

# ╔═╡ f8a53711-e07f-4b6b-84ea-803679496571
"""
Generate a regular lattice of points
$(TYPEDSIGNATURES)

# Arguments
- `dx` → element spacing on the x axis
- `dy` → element spacing on the y axis
- `ds` → displacement along x between rows of elements
- `f_cond::Function` → function of two arguments (`f(x,y) = ...`) returning `true` if element at position `x,y` must be kept and `false` otherwise

# Keyord Arguments
- `x0 = 0` → x coordinate of the origin of the lattice
- `y0 = 0` → y coordinate of the origin of the lattice
- `M::Int = 70` → Number of elements to generate per row of points before appliying the filtering function `f_cond`
- `N::Int = M` → Number of rows of points to generate before appliying the filtering function `f_cond`
"""
function generate_regular_lattice(dx::T, dy::T, ds::T, f_cond::Function = (x, y) -> true;x0::T = T(0), y0::T = T(0), M::Int = 70,N::Int = M) where T<:Real
	gen = lattice_generator(dx, dy, ds; x0, y0, M, N)
	return [x for x ∈ gen if f_cond(x...)]
end

# ╔═╡ ce22b91e-6bba-4312-a89b-1a78f84034d3
function regular_lattice_nelements(dx::T, dy::T, ds::T, f_cond::Function = (x, y) -> true;x0::T = T(0), y0::T = T(0), M::Int = 70,N::Int = M) where T<:Real
	gen = lattice_generator(dx, dy, ds; x0, y0, M, N)
	sum(x -> f_cond(x...), gen)
end

# ╔═╡ 0eb5d19d-b535-4cda-a89b-26ba295e2711
generate_regular_lattice(dx::Real,dy::Real,ds::Real,args...;kwargs...) = generate_regular_lattice(promote(dx,dy,ds)...,args...;kwargs...)

# ╔═╡ 2ba73b51-ecb9-4632-9f39-bdaeb8c5bd34
#=╠═╡ notebook_exclusive
@benchmark generate_square_lattice(1, (x,y) -> x^2 + y^2 < 100)
  ╠═╡ notebook_exclusive =#

# ╔═╡ f0834b38-8efe-4e77-b0f9-47e5b7595191
#=╠═╡ notebook_exclusive
md"""
## generate\_rect_lattice
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 9e5db472-f96a-4acb-96ae-024b5c73a93d
"""
    generate_rect_lattice(spacing_x::Real,spacing_y::Real[,f_cond];kwargs...)
# Summary
Generate a rectangular lattice of points (with different spacing among x and y directions)
# Arguments
- `spacing_x` → spacing between points on the x axis
- `spacing_y` → spacing between points on the y axis

See [`generate_regular_lattice`](@ref) for a description of `f_cond` and of  the keyword arguments
"""
generate_rect_lattice(spacing_x::Real,spacing_y::Real,args...;kwargs...) = generate_regular_lattice(spacing_x,spacing_y,0,args...;kwargs...)

# ╔═╡ c41f9f41-d8bd-4001-98cb-2ab788404b1b
#=╠═╡ notebook_exclusive
md"""
## generate\_square\_lattice
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 8787134f-9d14-4329-8dda-72557e3175b8
"""
`generate_square_lattice(spacing::Real[,f_cond];kwargs...)`
# Summary
Generate a square lattice of points (with equal spacing among x and y directions)
# Arguments
- `spacing` → spacing between points on both x and y axis

See [`generate_regular_lattice`](@ref) for a description of `f_cond` and of  the keyword arguments
"""
generate_square_lattice(spacing::Real,args...;kwargs...) = generate_regular_lattice(spacing,spacing,0,args...;kwargs...)

# ╔═╡ f589306c-919d-468a-a0fd-9367acc36a7b
#=╠═╡ notebook_exclusive
md"""
## generate\_hex\_lattice
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 6ca05079-4c0d-4c45-8486-a4291310189d
"""
`generate_hex_lattice(spacing::Real[,f_cond];kwargs...)`
# Summary
Generate a hexagonal lattice of points (with equal distance between neighboring points).
The hexagonal lattice generated by this function has distance between points on the same
column √3 times greater than the distance between points on the same row.
# Arguments
- `spacing` → spacing between points

See [`generate_regular_lattice`](@ref) for a description of `f_cond` and of  the keyword arguments
"""
generate_hex_lattice(spacing::Real,args...;kwargs...) = generate_regular_lattice(spacing .* (1,√3/2,.5)...,args...;kwargs...)

# ╔═╡ 243621c4-245a-4267-9bb4-568e673450fa
#=╠═╡ notebook_exclusive
generate_hex_lattice(1;M=20, x0 = .5) |> x -> scatter(x; mode="markers") |> Plot
  ╠═╡ notebook_exclusive =#

# ╔═╡ 059045e0-9acc-438d-b3f5-602f8d5892f7
#=╠═╡ notebook_exclusive
md"""
# Misc Functions
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 2f1f02c5-3ea5-40c1-8fae-704d150036e6
# Get the conversion from linear to db and viceversa
"""
$(TYPEDSIGNATURES)
Convert a number from linear to dB
"""
lin2db(x::Real) = 10log10(x)

# ╔═╡ 5be397fe-a531-423c-8be0-5d31df79dd2f
"""
$(TYPEDSIGNATURES)
Convert a number from dB to linear
"""
db2lin(x::Real) = 10^(x/10)

# ╔═╡ b2e80c33-bbfe-43ca-8795-c9d8d6fa52a9
# Convert between frequency and wavelength
"""
$(TYPEDSIGNATURES)
Get the wavelength (in m) starting from the frequency (in Hz)
"""
f2λ(f::Real) = c₀/f

# ╔═╡ 9165c4d4-69b5-456c-813c-4725feeb5b52
"""
$(TYPEDSIGNATURES)
Get the frequency (in Hz) starting from the wavelength (in m)
"""
λ2f(λ::Real) = c₀/λ

# ╔═╡ a5ca5a8a-8497-41e2-9af0-92db5db9ce73
#=╠═╡ notebook_exclusive
md"""
# Generate Colors
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 3f2a31d4-0fa8-40fa-9dc4-bd6a26d2ddc9
# Initialize the vector that contains the matrix to compute the beam coloring. We limit ourselves at 500 colors to start
const F_reuse_matrix = (square = SMatrix{2,2,Float64,4}[], triangular = SMatrix{2,2,Float64,4}[])

# ╔═╡ 6b9beb62-dc7e-4b8b-9b7c-8fee5b1da98f
function _coloring_inner_bruteforce!(T_mat, rot_mat, grid_max)
    max_colours = length(T_mat)
    check_vec = fill((typemax(Int),typemax(Int)),max_colours)
    @inline norm2(x) = sum(abs2.(x))
    @inbounds for x1 = 0:grid_max, y1 = 0:grid_max, x2 = -grid_max:grid_max, y2 = 0:grid_max
        mat = @SMatrix [x1 y1;x2 y2]
        # Compute the determinant
        t_det = Int(det(mat))
        # Skip points which are not likely to give useful results
        if (t_det < 1)  || (t_det > max_colours) || (t_det > grid_max^2) || (maximum(abs.(mat)) > ceil(sqrt(t_det) + 3))
            continue
        end
        # Compute the angle between the basis vectors identified by [x1,y1] and [x2,y2]
        angle = abs(atan(y1,x1) - atan(y2,x2))
        # Skip cases where the angle between vectors is either too acute or too obtuse
        if abs(π/2-angle) > π/4
            continue
        end
        # Create temp variables for computation of the minimum distance
        dmat = mat*rot_mat
        # Compute frobenius norm and minimum squared distance for the candidate lattice generating matrix
        frobe = round(Int,norm2(dmat))
        # display(frobe)
        # Minimum squared distance is either the modulo of one of the vectors or the modulo of sum or difference of them
        dmin = round(Int,minimum((norm2(dmat[1,:]), norm2(dmat[2,:]), norm2(sum(dmat;dims=1)), norm2(diff(dmat;dims=1)))))
        # Check if the current realization is better than the saved one
        if isless((-dmin,frobe),check_vec[t_det])
            # Update the check_vec
            check_vec[t_det] = (-dmin,frobe)
            # Update the vector containing the generating matrices
            T_mat[t_det] = mat'
        end
    end
end

# ╔═╡ 14cb2a0b-2ea8-471b-987f-1647f1516992
## Here we have the functions for the coloring computation
function compute_F_cell(max_colours::Int;grid_max::Int=25)
    #=
    This function is used to compute all the possible 2x2 lattice generating matrices for possible coloring schemes up to 'max_colours' colors
    Computation is done with a brute-force approach, generating all possible 2x1 vectors with maximum elements up to grid_max
    =#

    # Find the current length of the pre-computed vector
    current_length = length(F_reuse_matrix.square)
    if current_length >= max_colours
        # We already computed the function for the required number of colors
        return
    end
    n_missing = max_colours - current_length
    append!(F_reuse_matrix.square,Vector{SMatrix{2,2,Float64,4}}(undef,n_missing))
    append!(F_reuse_matrix.triangular,Vector{SMatrix{2,2,Float64,4}}(undef,n_missing))
    # Compute the matrix for the square lattice
    _coloring_inner_bruteforce!(F_reuse_matrix.square,I,grid_max)
    # Compute the matrix for the triangular lattice
    _coloring_inner_bruteforce!(F_reuse_matrix.triangular,@SMatrix([1 0;cosd(60) sind(60)]),grid_max)
end

# ╔═╡ 1d44cf1c-11a5-4366-94f3-85b695c6ca12
function generate_F_reuse_matrix(lattice_type::Symbol=:triangular,N_Colours::Int=4;max_colours::Int=max(10,N_Colours))
    compute_F_cell(max_colours)
    return getproperty(F_reuse_matrix,lattice_type)[N_Colours]
end

# ╔═╡ 7e68054e-4268-424e-b413-ef18baf832ac
"""
    generate_colors(BeamCenters::AbstractVector,N_Colours::Int=4;lattice_type::Symbol=:triangular)

Provide a the coloring breakdown for a given set of lattice points.
# Arguments
- `BeamCenters` → Vector of Tuple or StaticVectors expressing the U-V coordinates of each point in the lattice for which the coloring is being computed
- `N_Colours` → Number of colors to divide the lattice in. Defaults to `4`

# keyword Arguments
- `lattice_type` → Symbol that can either be `:triangular` or `:square`, idenifying the type of lattice. Defaults to `:triangular`
"""
function generate_colors(BeamCenters::AbstractVector,N_Colours::Int=4;first_color_coord=nothing,first_color_idx=nothing,precision_digits::Int=7,lattice_type::Symbol=:triangular)
    #=
    **************************************************************************
       Generate frequency colouring file
    **************************************************************************

     References:
     [1]   "On the frequency allocation for mobile radio telephone systems", C. de
           Almeida; R. Palazzo, Proceedings of 6th International Symposium on
           Personal, Indoor and Mobile Radio Communications, Year: 1995, Volume: 1, Pages: 96 - 99 vol.1
     [2]   P. Angeletti, "Simple implementation of vectorial modulo operation based
           on fundamental parallelepiped," in Electronics Letters, vol. 48, no. 3, pp. 159-160,
           February 2 2012. doi: 10.1049/el.2011.3667

     Authors: Alberto Mengali, 2020, European Space Agency

    Input Arguments:
    BeamCenters:         A vector containing the U,V coordinates of the beam centers as tuples or staticvectors
    N_Colours:           An integer number specifying the number of colors to generate in the association
    first_color_coord:   A tuple or static vector containing the U,V coordinates of the beam that will have the first color
    first_color_idx:     The beam index of the beam containing the first color, either this variable of first_order_coord can be specified, not together
    precision_digits:    The number of digits to be used in the flooring operation
    =#

    # Check if either the first color coordinates or first color idx are given
    if first_color_coord === nothing
        if first_color_idx === nothing
            # Initialize to coorinates of the first beam
            first_color_idx = 1;
        end
        first_color_coord = BeamCenters[first_color_idx]
    else
        if first_color_idx !== nothing
            @warn "Both first_color_idx and first_color_coord were given, disregarding the idx variable"
        end
    end
    if lattice_type ∉ (:triangular, :square)
        @error "The specified lattice type ($lattice_type) is not recognized, it should be either :triangular or :square"
    end

    n_beams = length(BeamCenters)

    # If only 1 color is requested, simply return the trivial result
    if N_Colours == 1
        Colours = ones(n_beams)
        nbeams_per_color = n_beams
        idxs = fill(fill(true,nbeams))
    end

    # Find the minimum distance in U and V
    minU_dist = minimum(diff(sort(first.(BeamCenters)) |> x -> unique!(y -> round(y;digits=precision_digits),x)))
    minV_dist = minimum(diff(sort(last.(BeamCenters)) |> x -> unique!(y -> round(y;digits=precision_digits),x)))
    if lattice_type === :triangular
        beamU_spacing = 2minU_dist
        beamV_spacing = 2minV_dist
        # Matrix to normalize the grid points in u-v into a integer grid with y(v) axis not pointing north but north-east with 60° inclination
        D = @SMatrix [1 -1/2;0 1]
    elseif lattice_type === :square
        beamU_spacing = minU_dist
        beamV_spacing = minV_dist
        D = @SMatrix [1 0;0 1]
    end
    # Get the coloring generating matrix
    F_reuse_matrix = generate_F_reuse_matrix(lattice_type,N_Colours)
    # Create the set that will contain the unique results
    unique_colors_vecs = SVector{2,Int}[]
    # Initialize the colors vector
    Colors = similar(BeamCenters, Int)
    @inbounds for (n,p₀) ∈ enumerate(BeamCenters)
        # Compute the integer beam centers
        p₀_normalized = p₀ ./ SA_F64[beamU_spacing, minV_dist]
        # Transform the beam centers from u-v coordinates in radians into integer indexes
        beam_index_vector = round.(Int,D*p₀_normalized);
        # Find the values of the beam indexes modulo F_reuse_matrix
        unique_beam_index = round.(Int,beam_index_vector .- (F_reuse_matrix*floor.(round.(inv(F_reuse_matrix)*beam_index_vector,digits=precision_digits))))
        # Check if this color has already been assigned/registered
        idx = findfirst(x -> x == unique_beam_index,unique_colors_vecs)
        if idx isa Nothing
            push!(unique_colors_vecs, unique_beam_index)
            cidx = length(unique_colors_vecs)
        else
            cidx = idx
        end
        Colors[n] = cidx
    end

    return Colors
end

# ╔═╡ d8584d03-8eb4-4864-b646-a6a0656a2e12
begin
	export generate_regular_lattice, generate_square_lattice, generate_hex_lattice, generate_rect_lattice
	export generate_colors
	export f2λ, λ2f, db2lin, lin2db
end

# ╔═╡ 8160086a-6349-447c-87ae-880b02fa97f5
#=╠═╡ notebook_exclusive
md"""
# Tests
"""
  ╠═╡ notebook_exclusive =#

# ╔═╡ 3ea0415c-af14-430c-bf7c-2c7d71b7a333
#=╠═╡ notebook_exclusive
let
	N_colors = 6
	lat = generate_hex_lattice(1; M = 10)
	colors = generate_colors(lat,N_colors)
	data = scatter(lat;mode="markers", marker_color = colors)
	Plot(data)
end
  ╠═╡ notebook_exclusive =#

# ╔═╡ 65b74f4a-e2af-4caf-8719-5a59c6349bb9
#=╠═╡ notebook_exclusive
let
	N_colors = 4
	lat = generate_square_lattice(1; M = 10)
	colors = generate_colors(lat,N_colors; lattice_type = :square)
	data = scatter(lat;mode="markers", marker_color = colors)
	Plot(data)
end
  ╠═╡ notebook_exclusive =#

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
PlotlyBase = "a03496cd-edff-5a9b-9e67-9cda94a718b5"
PlutoUtils = "ed5d0301-4775-4676-b788-cf71e66ff8ed"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

[compat]
BenchmarkTools = "~1.2.0"
DocStringExtensions = "~0.8.6"
PlotlyBase = "~0.8.18"
PlutoUtils = "~0.4.13"
StaticArrays = "~1.2.13"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.0-rc2"
manifest_format = "2.0"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "0bc60e3006ad95b4bb7497698dd7c6d649b9bc06"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.1"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.BenchmarkTools]]
deps = ["JSON", "Logging", "Printf", "Profile", "Statistics", "UUIDs"]
git-tree-sha1 = "61adeb0823084487000600ef8b1c00cc2474cd47"
uuid = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
version = "1.2.0"

[[deps.Chain]]
git-tree-sha1 = "cac464e71767e8a04ceee82a889ca56502795705"
uuid = "8be319e6-bccf-4806-a6f7-6fae938471bc"
version = "0.4.8"

[[deps.ColorSchemes]]
deps = ["ColorTypes", "Colors", "FixedPointNumbers", "Random"]
git-tree-sha1 = "a851fec56cb73cfdf43762999ec72eff5b86882a"
uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
version = "3.15.0"

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "024fe24d83e4a5bf5fc80501a314ce0d1aa35597"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.0"

[[deps.Colors]]
deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
git-tree-sha1 = "417b0ed7b8b838aa6ca0a87aadf1bb9eb111ce40"
uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
version = "0.12.8"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[deps.Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[deps.DataAPI]]
git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.9.0"

[[deps.DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[deps.DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "b19534d1895d702889b219c382a6e18010797f0b"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.6"

[[deps.Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.FileWatching]]
uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

[[deps.Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[deps.Glob]]
git-tree-sha1 = "4df9f7e06108728ebf00a0a11edee4b29a482bb2"
uuid = "c27321d9-0574-5035-807b-f59d2c89b15c"
version = "1.3.0"

[[deps.Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[deps.HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

[[deps.IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[deps.IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[deps.LaTeXStrings]]
git-tree-sha1 = "f2355693d6778a178ade15952b7ac47a4ff97996"
uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
version = "1.3.0"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[deps.MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "3d3e902b31198a27340d0bf00d6ac452866021cf"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.9"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[deps.Parameters]]
deps = ["OrderedCollections", "UnPack"]
git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe"
uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a"
version = "0.12.3"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "ae4bbcadb2906ccc085cf52ac286dc1377dceccc"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.1.2"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[deps.PlotlyBase]]
deps = ["ColorSchemes", "Dates", "DelimitedFiles", "DocStringExtensions", "JSON", "LaTeXStrings", "Logging", "Parameters", "Pkg", "REPL", "Requires", "Statistics", "UUIDs"]
git-tree-sha1 = "180d744848ba316a3d0fdf4dbd34b77c7242963a"
uuid = "a03496cd-edff-5a9b-9e67-9cda94a718b5"
version = "0.8.18"

[[deps.PlutoDevMacros]]
deps = ["MacroTools", "PlutoHooks"]
git-tree-sha1 = "7392720177703062cb2e2a0115efb77dc5dc818c"
uuid = "a0499f29-c39b-4c5c-807c-88074221b949"
version = "0.3.7"

[[deps.PlutoHooks]]
deps = ["FileWatching", "InteractiveUtils", "Markdown", "UUIDs"]
git-tree-sha1 = "f297787f7d7507dada25f6769fe3f08f6b9b8b12"
uuid = "0ff47ea0-7a50-410d-8455-4348d5de0774"
version = "0.0.3"

[[deps.PlutoTest]]
deps = ["HypertextLiteral", "InteractiveUtils", "Markdown", "Test"]
git-tree-sha1 = "92b8ae1eee37c1b8f70d3a8fb6c3f2d81809a1c5"
uuid = "cb4044da-4d16-4ffa-a6a3-8cad7f73ebdc"
version = "0.2.0"

[[deps.PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "e071adf21e165ea0d904b595544a8e514c8bb42c"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.19"

[[deps.PlutoUtils]]
deps = ["Chain", "Glob", "HypertextLiteral", "InteractiveUtils", "Markdown", "PlutoDevMacros", "PlutoHooks", "PlutoTest", "PlutoUI", "PrettyTables", "Reexport", "Requires", "UUIDs"]
git-tree-sha1 = "3d3856ecfea340b4ee0c77e5c3228dd1b4478ae1"
uuid = "ed5d0301-4775-4676-b788-cf71e66ff8ed"
version = "0.4.13"

[[deps.PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "d940010be611ee9d67064fe559edbb305f8cc0eb"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.2.3"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.Profile]]
deps = ["Printf"]
uuid = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[deps.Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "4036a3bd08ac7e968e27c203d45f5fff15020621"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.1.3"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.StaticArrays]]
deps = ["LinearAlgebra", "Random", "Statistics"]
git-tree-sha1 = "3c76dde64d03699e074ac02eb2e8ba8254d428da"
uuid = "90137ffa-7385-5640-81b9-e52037218182"
version = "1.2.13"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[deps.TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[deps.Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "TableTraits", "Test"]
git-tree-sha1 = "fed34d0e71b91734bf0a7e10eb1bb05296ddbcd0"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.6.0"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[deps.Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[deps.UnPack]]
git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b"
uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
version = "1.0.2"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─379613ec-0973-4000-ae8c-d7c33ddca18e
# ╠═74975885-9a4e-4857-8135-9e4f69061caf
# ╠═c0a30957-4c7b-4d7b-bfa9-c2fb691a077b
# ╠═736b0cf6-bec2-4226-8ef4-70f6a865d34a
# ╟─f8243a65-9f5e-464e-bb06-0bb4f5131b8b
# ╠═d8584d03-8eb4-4864-b646-a6a0656a2e12
# ╟─8660a7c4-eb78-4e7c-966b-d759df7f3dfa
# ╟─71163795-9695-4f11-acc2-6e3838c8a158
# ╠═f8a53711-e07f-4b6b-84ea-803679496571
# ╠═ce22b91e-6bba-4312-a89b-1a78f84034d3
# ╠═da97848f-a7ff-4f2d-b98d-e8bf1ccc3038
# ╠═0eb5d19d-b535-4cda-a89b-26ba295e2711
# ╠═2ba73b51-ecb9-4632-9f39-bdaeb8c5bd34
# ╟─f0834b38-8efe-4e77-b0f9-47e5b7595191
# ╠═9e5db472-f96a-4acb-96ae-024b5c73a93d
# ╟─c41f9f41-d8bd-4001-98cb-2ab788404b1b
# ╠═8787134f-9d14-4329-8dda-72557e3175b8
# ╟─f589306c-919d-468a-a0fd-9367acc36a7b
# ╠═6ca05079-4c0d-4c45-8486-a4291310189d
# ╠═243621c4-245a-4267-9bb4-568e673450fa
# ╟─059045e0-9acc-438d-b3f5-602f8d5892f7
# ╠═2f1f02c5-3ea5-40c1-8fae-704d150036e6
# ╠═5be397fe-a531-423c-8be0-5d31df79dd2f
# ╠═b2e80c33-bbfe-43ca-8795-c9d8d6fa52a9
# ╠═9165c4d4-69b5-456c-813c-4725feeb5b52
# ╟─a5ca5a8a-8497-41e2-9af0-92db5db9ce73
# ╠═3f2a31d4-0fa8-40fa-9dc4-bd6a26d2ddc9
# ╠═14cb2a0b-2ea8-471b-987f-1647f1516992
# ╠═6b9beb62-dc7e-4b8b-9b7c-8fee5b1da98f
# ╠═1d44cf1c-11a5-4366-94f3-85b695c6ca12
# ╠═7e68054e-4268-424e-b413-ef18baf832ac
# ╟─8160086a-6349-447c-87ae-880b02fa97f5
# ╠═3ea0415c-af14-430c-bf7c-2c7d71b7a333
# ╠═65b74f4a-e2af-4caf-8719-5a59c6349bb9
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


# We illustrate here the use of `ListofStates` in dealing with a warm start procedure.

# `ListofStates` is designed to store the of the iteration process.
# In this tutorial, we compare the resolution of a convex unconstrained problem with 3 variants:
#  - a steepest descent method
#  - an inverse-BFGS method
#  - a mix of 5 steps of steepest descent and then switching to a BFGS initialized with the 5 previous steps.

using Stopping, NLPModels, LinearAlgebra, Test, Printf

# First, we introduce our two implementations that both uses an backtracking Armijo linesearch.

import Stopping.armijo
function armijo(xk, dk, fk, slope, f)
  t = 1.0
  fk_new = f(xk + dk)
  while f(xk + t * dk) > fk + 1.0e-4 * t * slope
    t /= 1.5
    fk_new = f(xk + t * dk)
  end
  return t, fk_new
end

function steepest_descent(stp :: NLPStopping)

  xk = stp.current_state.x
  fk, gk = objgrad(stp.pb, xk)

  OK = update_and_start!(stp, fx = fk, gx = gk)

  @printf "%2s %9s %7s %7s %7s\n" "k" "fk" "||∇f(x)||" "t" "λ"
  @printf "%2d %7.1e %7.1e\n" stp.meta.nb_of_stop fk norm(stp.current_state.current_score)
  while !OK
    dk = - gk
    slope = dot(dk, gk)
    t, fk = armijo(xk, dk, fk, slope, x->obj(stp.pb, x))
    xk += t * dk
    fk, gk = objgrad(stp.pb, xk)

    OK = update_and_stop!(stp, x = xk, fx = fk, gx = gk)

    @printf "%2d %9.2e %7.1e %7.1e %7.1e\n" stp.meta.nb_of_stop fk norm(stp.current_state.current_score) t slope
  end
  return stp
end

function bfgs_quasi_newton_armijo(stp :: NLPStopping; Hk = nothing)

  xk = stp.current_state.x
  fk, gk = objgrad(stp.pb, xk)
  gm = gk

  dk, t = similar(gk), 1.
  if isnothing(Hk)
    Hk = I #start from identity matrix
  end

  OK = update_and_start!(stp, fx = fk, gx = gk)

  @printf "%2s %7s %7s %7s %7s\n" "k" "fk" "||∇f(x)||" "t" "cos"
  @printf "%2d %7.1e %7.1e\n" stp.meta.nb_of_stop fk norm(stp.current_state.current_score)

  while !OK
    if stp.meta.nb_of_stop != 0
      sk = t * dk
      yk = gk - gm
      ρk = 1/dot(yk, sk)
      #we need yk'*sk > 0 for instance yk'*sk ≥ 1.0e-2 * sk' * Hk * sk
      Hk = ρk ≤ 0.0 ? Hk : (I - ρk * sk * yk') * Hk * (I - ρk * yk * sk') + ρk * sk * sk'
      if norm(sk) ≤ 1e-14
        break
      end
      #H2 = Hk + sk * sk' * (dot(sk,yk) + yk'*Hk*yk )*ρk^2 - ρk*(Hk * yk * sk' + sk * yk'*Hk)
    end
    dk = - Hk * gk
    slope = dot(dk, gk) # ≤ -1.0e-4 * norm(dk) * gnorm
    t, fk = armijo(xk, dk, fk, slope, x->obj(stp.pb, x))

    xk = xk + t * dk
    gm = copy(gk)
    gk = grad(stp.pb, xk)

    OK = update_and_stop!(stp, x = xk, fx = fk, gx = gk)
    @printf "%2d %7.1e %7.1e %7.1e %7.1e\n" stp.meta.nb_of_stop fk norm(stp.current_state.current_score) t slope
  end
  stp.stopping_user_struct = Dict( :Hk => Hk)
  return stp
end

# We consider the following convex unconstrained problem model using `ADNLPModels.jl` and defines a related `NLPStopping`.

fH(x) = (x[2]+x[1].^2-11).^2+(x[1]+x[2].^2-7).^2
nlp = ADNLPModel(fH, [10., 20.])
stp = NLPStopping(nlp, optimality_check = unconstrained_check,
                 atol = 1e-6, rtol = 0.0, max_iter = 100)

reinit!(stp, rstate = true, x = nlp.meta.x0)
steepest_descent(stp)

@test status(stp) == :Optimal
@test stp.listofstates == VoidListofStates()

elapsed_time(stp)
nlp.counters

reinit!(stp, rstate = true, x = nlp.meta.x0, rcounters = true)
bfgs_quasi_newton_armijo(stp)

@test status(stp) == :Optimal
@test stp.listofstates == VoidListofStates()

elapsed_time(stp)
nlp.counters

NLPModels.reset!(nlp)
stp_warm = NLPStopping(nlp, optimality_check = unconstrained_check,
                      atol = 1e-6, rtol = 0.0, max_iter = 5,
                      n_listofstates = 5) #shortcut for list = ListofStates(5, Val{NLPAtX{Float64,Array{Float64,1},Array{Float64,2}}}()))
steepest_descent(stp_warm)
@test status(stp_warm) == :IterationLimit
@test length(stp_warm.listofstates) == 5

Hwarm = I
for i=2:5
  sk = stp_warm.listofstates.list[i][1].x - stp_warm.listofstates.list[i-1][1].x
  yk = stp_warm.listofstates.list[i][1].gx - stp_warm.listofstates.list[i-1][1].gx
  ρk = 1/dot(yk, sk)
  if ρk > 0.0
    global Hwarm = (I - ρk * sk * yk') * Hwarm * (I - ρk * yk * sk') + ρk * sk * sk'
  end
end

reinit!(stp_warm)
stp_warm.meta.max_iter = 100
bfgs_quasi_newton_armijo(stp_warm, Hk = Hwarm)
status(stp_warm)

elapsed_time(stp_warm)
nlp.counters


using ArgParse, LinearMaps, HDF5, JLD

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
      "--hfield"
        help = "Magnetic field parameter."
        arg_type = Real
        default = 2.0
      "--Jcoupling", "-J"
        help = "Coupling constant."
        arg_type = Real
        default = 1.0
      "--Nx", "-x"
        help = "Number of columns in the lattice."
        arg_type = Int
        default = 10
      "--Ny", "-y"
        help = "Number of rows in the lattice."
        arg_type = Int
        default = 1
      "--outputpath", "-o"
        help = "Where to save output. Default is ./ising. The full output
                string is formatted to also include the arguments; e.g.
                ./ising_h1_0_J1_0Nx3Ny3.out"
        arg_type = AbstractString
        default =  "./ising"
    end

    return parse_args(s)
end

"""
This makes the nearest-neighbour Hamiltonian, which is a 4x4 matrix.
"""
function buildisingham(J, h)
    X = [0 1; 1 0]
    Z = [1 0; 0 -1]
    I2 = eye(2)
    XX = kron(X, X)
    ZI = kron(Z, I2)
    IZ = kron(I2, Z)
    H = -(J*XX + h/2*(ZI+IZ))
    return H
end

"""
This is the function that actually multiplies the input by the hamiltonian.
It does this once for each cycle in the permutation; the length of the cycle
needs to be supplied..
"""
function applycycleperm(v, out, ham, cycleperm, cyclelength)
    D = length(v)
    N = convert(Int64, log2(D))
    siteshp = tuple(collect(Iterators.repeated(2,N))...)#siteshp=(2,2,2..) i.e. each dim is a site
                                              #  with one entry per spin state
    for n=1:cyclelength
        #reshape so we can apply ham
        v, out = reshape(v, (4, 2^(N-2))), reshape(out, (4,2^(N-2)))
        out += ham*v
        #reshape so we can permute sites (since each gets its own dim)
        v, out = reshape(v, siteshp), reshape(out, siteshp)
        v, out = permutedims(v, cycleperm), permutedims(out, cycleperm)
    end
    out = reshape(out, D)
    return v, out
end

"""
Constructs a cyclic permutation tuple of length N, which is a 1:N range
whose entries at the indices permuteidx are replaced by the values in
permutedval. For example, N=5, permuteidx=(1,3,5), permutedval=(5,1,3)
would return cycleperm=(5,2,1,4,3).
"""
function makecycleperm(permuteidx, permutedval, N)
    L = length(permuteidx)
    @assert L==length(permutedval) "$L, $length(permutedval)"
    cyclepermarr = collect(1:N)
    for p =  1:L
      idx = permuteidx[p]
      cyclepermarr[idx] = permutedval[p]
    end
    cycleperm = tuple(cyclepermarr...)
    return cycleperm
end

function addperm!(first::Int, stride::Int, last::Int, permuteidxs, permutedvals)
      theseidxs = collect(first:stride:last) #e.g. [4,5,6] if 2nd col of 3x3
      thesevals = circshift(theseidxs, -1) #e.g. [5,6,4]

      push!(permuteidxs, theseidxs)
      push!(permutedvals, thesevals)
      # #need to also push a cyclic permutation to cover the PBCs
      # push!(permuteidxs, circshift(theseidxs, 1))
      # push!(permutedvals, circshift(thesevals, 1))
end

"""
This function applies the full Hamiltonian to the input state v, by
repeatedly applying the nearest-neighbour Hamiltonian to appropriate
permutations of v such that each nearest-neighbour pair is coupled.
"""
function applyspinham(v::AbstractVector, ham::Matrix, Nx::Int, Ny::Int)
    D = length(v)
    N = Nx*Ny
    @assert D==2^N "D ($D) didn't equal 2^N (2^$N)"
    siteshp = tuple(collect(Iterators.repeated(2,N))...)#siteshp=(2,2,2..) i.e. each dim is a site
                                              #  with one entry per spin state
    out = zeros(siteshp)

    #can probably pre-allocate these, but they are gonna be short enough that
    #little would be gained
    permuteidxs = []
    permutedvals = []

    #It's probably possible to subsume these two for loops into a single
    #more general function, but meh.
    #Make the column permutations
    if Ny>1
      for col = 1:Nx
        first = (col-1)*Ny+1
        last = (col-1)*Ny+Ny
        addperm!(first, 1, last, permuteidxs, permutedvals)
      end
    end

    #Make the row permutations
    if Nx>1
      for row = 1:Ny
        first = row
        last = (Nx-1)*Ny+row
        addperm!(first, Ny, last, permuteidxs, permutedvals)
      end
    end
    # println("permuteidxs:")
    # println(permuteidxs)
    # println("permutedvals:")
    # println(permutedvals)

    #Apply the Hamiltonian once for each permutation
    for p in zip(permuteidxs,  permutedvals)
      #println(p[1], p[2])
      cycleperm = makecycleperm(p[1], p[2], N)
      v, out = applycycleperm(v, out, ham, cycleperm, length(p[1]))
    end
    out = reshape(out, D)
    return out
end

"""
This wraps all the various calls needed to find the Ising ground state because
I don't like long main functions. The Julia function eigs() implements
the Lancoz eigensolver, which conveniently requires the target operator to
be defined only through its action on vectors (i.e. Julia's LinearMap type).
"""
function findgroundstate_ising(parsed_args::Dict)
  Nx = parsed_args["Nx"]
  Ny = parsed_args["Ny"]
  N = Nx*Ny
  # rng = MersenneTwister()
  # v = randn(rng, 2^N) + randn(rng, 2^N)*im
  isingham = buildisingham(parsed_args["Jcoupling"],  parsed_args["hfield"])
  applyisingham(v) = applyspinham(v, isingham, Nx, Ny)
  H = LinearMap(applyisingham, 2^N)
  u, v = eigs(H; nev=1, which=:SR)
  v/=sqrt(sum(abs2(v))) #normalize
  energy = real(u[1])#/N
  edensity = energy/N
  #infN_e_density = -4/pi
  println("Energy:", energy)
  println("Energy density:", edensity)
  #info((edensity - infN_e_density)/infN_e_density)
  return energy, v
end

"""
Little function to save the output to disk.
"""
function saveoutput(parsed_args::Dict, energy::Real,
                    wavefunction::AbstractArray)
    outbuffer = IOBuffer()
    print(outbuffer, parsed_args["outputpath"])
    print(outbuffer, string("_h",
                            replace(string(parsed_args["hfield"]), ".", "_")))
    print(outbuffer, string("J",
                            replace(string(parsed_args["Jcoupling"]), ".",  "_")))
    print(outbuffer, string("Nx", parsed_args["Nx"]))
    print(outbuffer, string("Ny", parsed_args["Ny"]))
    print(outbuffer, ".jld")
    outstring = takebuf_string(outbuffer)
    println("Saving to ", outstring)
    save(outstring, "energy", energy)
    save(outstring, "wavefunction", wavefunction)
end

function main()
  parsed_args = parse_commandline()
  println("Hi!! Finding the Ising model ground state with the following arguments:")
  for (arg, val) in parsed_args
    println("   $arg    =>    $val")
  end
  energy, wavefunction = findgroundstate_ising(parsed_args)
  saveoutput(parsed_args, energy, wavefunction)
end
main()


# ---
# title: 1018. Binary Prefix Divisible By 5
# id: problem1018
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/binary-prefix-divisible-by-5/description/>
# hidden: true
# ---
#
# Given an array `A` of `0`s and `1`s, consider `N_i`: the i-th subarray from
# `A[0]` to `A[i]` interpreted as a binary number (from most-significant-bit to
# least-significant-bit.)
#
# Return a list of booleans `answer`, where `answer[i]` is `true` if and only if
# `N_i` is divisible by 5.
#
# **Example 1:**
#
#
#
#     Input: [0,1,1]
#     Output: [true,false,false]
#     Explanation:
#     The input numbers in binary are 0, 01, 011; which are 0, 1, and 3 in base-10.  Only the first number is divisible by 5, so answer[0] is true.
#
#
# **Example 2:**
#
#
#
#     Input: [1,1,1]
#     Output: [false,false,false]
#
#
# **Example 3:**
#
#
#
#     Input: [0,1,1,1,1,1]
#     Output: [true,false,false,false,true,false]
#
#
# **Example 4:**
#
#
#
#     Input: [1,1,1,0,1]
#     Output: [false,false,false,false,false]
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 30000`
#   2. `A[i]` is `0` or `1`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#
# Create functions if a string is passed as summary_statistic
#
function build_summary_statistic(summary_statistic_string::String)

    if summary_statistic_string == "mean"
        return mean_summary_statistic
    elseif summary_statistic_string == "variance"
        return var_summary_statistic
    elseif summary_statistic_string == "max"
        return max_summary_statistic
    elseif summary_statistic_string == "min"
        return min_summary_statistic
    elseif summary_statistic_string == "range"
        return range_summary_statistic
    elseif summary_statistic_string == "median"
        return median_summary_statistic
    elseif summary_statistic_string == "q1"
        return q1_summary_statistic
    elseif summary_statistic_string == "q3"
        return q3_summary_statistic
    elseif summary_statistic_string == "iqr"
        return iqr_summary_statistic
    elseif summary_statistic_string == "keep_all"
        return keep_all_summary_statistic
    else
        error("$(summary_statistic_string) is not a valid summary statistic")
    end

end

function build_summary_statistic(summary_statistic_strings::Vector{String})
    function summary_statistic(data)
        sum_stat_funcs = [build_summary_statistic(sum_stat) for sum_stat in summary_statistic_strings]
        out = [func(data) for func in sum_stat_funcs]
        return vcat(out...)
    end
    return summary_statistic
end

#
# Check a user-defined summary statistic function
#
function build_summary_statistic(summary_statistic_func::Function)
    # Check output shape is right
    test_output = summary_statistic_func(rand(3,100))

    if ndims(test_output) != 1
        error("The summary statistic function provided does not return a 1D array")
    end

    return summary_statistic_func
end

#
# Built-in summary statistics. Data should be in the format returned by
# DifferentialEquations.solve i.e. an array with size (n_trajectories, n_time_points)
#

# Mean of each trajectory
function mean_summary_statistic(data::AbstractArray{Float64,2})
    return mean(data, dims=2)[:]
end

# Variance of each trajectory
function var_summary_statistic(data::AbstractArray{Float64,2})
    return var(data, dims=2)[:]
end

# Max of each trajectory
function max_summary_statistic(data::AbstractArray{Float64,2})
    return maximum(data, dims=2)[:]
end

# Min of each trajectory
function min_summary_statistic(data::AbstractArray{Float64,2})
    return minimum(data, dims=2)[:]
end

# Range of each trajectory
function range_summary_statistic(data::AbstractArray{Float64,2})
    return max_summary_statistic(data) - min_summary_statistic(data)
end

# Median of each trajectory
function median_summary_statistic(data::AbstractArray{Float64,2})
    return median(data, dims=2)[:]
end

# First quartile of each trajectory
function q1_summary_statistic(data::AbstractArray{Float64,2})
    return [quantile(data[i,:], 0.25) for i = 1:size(data,1)]
end

# Third quartile of each trajectory
function q3_summary_statistic(data::AbstractArray{Float64,2})
    return [quantile(data[i,:], 0.75) for i = 1:size(data,1)]
end

# Inter-quartile range of each trajectory
function iqr_summary_statistic(data::AbstractArray{Float64,2})
    return q3_summary_statistic(data) - q1_summary_statistic(data)
end

# Keep all the data from all the trajectoriess (in a single 1d array)
function keep_all_summary_statistic(data::AbstractArray{Float64,2})
    return reshape(data, prod(size(data)))
end


if basename(pwd()) == "aoc"
    cd("2017/2")
end

function part1(filename::AbstractString)
    sum(-(reverse(extrema(parse.(Int, split(line))))...) for line in eachline(filename))
end

@assert part1("example.txt") == 18
part1("input.txt")

findround(str::AbstractString) = findround(parse.(Int, split(str)))

function findround(xs::Vector{Int})
    for i in firstindex(xs):lastindex(xs)-1
        for j in i+1:lastindex(xs)
            lo, hi = minmax(xs[i], xs[j])
            if isinteger(hi / lo)
                return hi ÷ lo
            end
        end
    end
end

part2(filename::AbstractString) = sum(findround.(eachline(filename)))

@assert part2("example2.txt") == 9
part2("input.txt")


# ---
# title: 903. Valid Permutations for DI Sequence
# id: problem903
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Divide and Conquer, Dynamic Programming
# link: <https://leetcode.com/problems/valid-permutations-for-di-sequence/description/>
# hidden: true
# ---
#
# We are given `S`, a length `n` string of characters from the set `{'D', 'I'}`.
# (These letters stand for "decreasing" and "increasing".)
#
# A  _valid permutation_  is a permutation `P[0], P[1], ..., P[n]` of integers
# `{0, 1, ..., n}`, such that for all `i`:
#
#   * If `S[i] == 'D'`, then `P[i] > P[i+1]`, and;
#   * If `S[i] == 'I'`, then `P[i] < P[i+1]`.
#
# How many valid permutations are there?  Since the answer may be large,
# **return your answer modulo`10^9 + 7`**.
#
#
#
# **Example 1:**
#
#
#
#     Input: "DID"
#     Output: 5
#     Explanation:
#     The 5 valid permutations of (0, 1, 2, 3) are:
#     (1, 0, 3, 2)
#     (2, 0, 3, 1)
#     (2, 1, 3, 0)
#     (3, 0, 2, 1)
#     (3, 1, 2, 0)
#
#
#
#
# **Note:**
#
#   1. `1 <= S.length <= 200`
#   2. `S` consists only of characters from the set `{'D', 'I'}`.
#
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


function test_aux()
  # test Givens reflector corner cases
  (c, s, ρ) = Krylov.sym_givens(0.0, 0.0)
  @test (c == 1.0) && (s == 0.0) && (ρ == 0.0)

  a = 3.14
  (c, s, ρ) = Krylov.sym_givens(a, 0.0)
  @test (c == 1.0) && (s == 0.0) && (ρ == a)
  (c, s, ρ) = Krylov.sym_givens(-a, 0.0)
  @test (c == -1.0) && (s == 0.0) && (ρ == a)

  b = 3.14
  (c, s, ρ) = Krylov.sym_givens(0.0, b)
  @test (c == 0.0) && (s == 1.0) && (ρ == b)
  (c, s, ρ) = Krylov.sym_givens(0.0, -b)
  @test (c == 0.0) && (s == -1.0) && (ρ == b)

  # test roots of a quadratic
  roots = Krylov.roots_quadratic(0.0, 0.0, 0.0)
  @test length(roots) == 1
  @test roots[1] == 0.0

  roots = Krylov.roots_quadratic(0.0, 0.0, 1.0)
  @test length(roots) == 0

  roots = Krylov.roots_quadratic(0.0, 3.14, -1.0)
  @test length(roots) == 1
  @test roots[1] == 1.0 / 3.14

  roots = Krylov.roots_quadratic(1.0, 0.0, 1.0)
  @test length(roots) == 0

  roots = Krylov.roots_quadratic(1.0, 0.0, 0.0)
  @test length(roots) == 2
  @test roots[1] == 0.0
  @test roots[2] == 0.0

  roots = Krylov.roots_quadratic(1.0, 3.0, 2.0)
  @test length(roots) == 2
  @test roots[1] ≈ -2.0
  @test roots[2] ≈ -1.0

  roots = Krylov.roots_quadratic(1.0e+8, 1.0, 1.0)
  @test length(roots) == 0

  # ill-conditioned quadratic
  roots = Krylov.roots_quadratic(-1.0e-8, 1.0e+5, 1.0, nitref=0)
  @test length(roots) == 2
  @test roots[1] == 1.0e+13
  @test roots[2] == 0.0

  # iterative refinement is crucial!
  roots = Krylov.roots_quadratic(-1.0e-8, 1.0e+5, 1.0, nitref=1)
  @test length(roots) == 2
  @test roots[1] == 1.0e+13
  @test roots[2] == -1.0e-05

  # not ill-conditioned quadratic
  roots = Krylov.roots_quadratic(-1.0e-7, 1.0, 1.0, nitref=0)
  @test length(roots) == 2
  @test isapprox(roots[1],  1.0e+7, rtol=1.0e-6)
  @test isapprox(roots[2], -1.0, rtol=1.0e-6)

  roots = Krylov.roots_quadratic(-1.0e-7, 1.0, 1.0, nitref=1)
  @test length(roots) == 2
  @test isapprox(roots[1], 1.0e+7, rtol=1.0e-6)
  @test isapprox(roots[2], -1.0, rtol=1.0e-6)

  # test trust-region boundary
  x = ones(5)
  d = ones(5); d[1:2:5] .= -1
  @test_throws ErrorException Krylov.to_boundary(x, d, -1.0)
  @test_throws ErrorException Krylov.to_boundary(x, d, 0.5)
  @test_throws ErrorException Krylov.to_boundary(x, zeros(5), 1.0)
  @test maximum(Krylov.to_boundary(x, d, 5.0)) ≈ 2.209975124224178
  @test minimum(Krylov.to_boundary(x, d, 5.0)) ≈ -1.8099751242241782
  @test maximum(Krylov.to_boundary(x, d, 5.0, flip=true)) ≈ 1.8099751242241782
  @test minimum(Krylov.to_boundary(x, d, 5.0, flip=true)) ≈ -2.209975124224178
end

test_aux()


function compare_curves(args...; tmin = 1e-8, tmax = 1, npoints = 50)
    ts = 10 .^ range(log(10, tmin), stop = log(10, tmax), length = npoints)

    curves = Dict(Symbol(r) => [] for r in args if r isa Function)
    curves[:slope2] = []
    curves[:slope3] = []

    for (i, t) in enumerate(ts)
        for curve in args
            if curve isa Function
                val = curve(t)
                @assert val >= 0 "curve should be positive, value is $(curve(t))."
                push!(curves[Symbol(curve)], (t, val))
            end
        end
        push!(curves[:slope2], (t, t^2))
        push!(curves[:slope3], (t, t^3))
    end
    return curves
end

function display_curvescomparison(curves)
    ps = []

    for (k, v) in curves
        legendentry = string(k)
        opts = occursin("slope", legendentry) ? Dict("no marks" => nothing) : Dict()
        push!(ps, PlotInc(PGFPlotsX.Options(opts...), Coordinates(v)))
        push!(ps, LegendEntry(replace(string(k), "_" => " ")))
    end

    ## Plot errors:
    return fig_time_subopt = @pgf Axis(
        {
            xmode = "log",
            ymode = "log",
            height = "12cm",
            width = "12cm",
            xlabel = "t",
            # ylabel = L"f(R_x(t\eta))-f(x)-t\langle gradf(x),\eta\rangle",
            legend_pos = "outer north east",
        },
        ps...,
    )
end

function check_regularizer_gradient_hessian(M, reg; x=nothing, ξ=nothing, η=nothing)

    isnothing(x) && (x = randomMPoint(M))
    isnothing(ξ) && (ξ = randomTVector(M, x))
    isnothing(η) && (η = randomTVector(M, x))

    f(x) = g(reg, x)
    gradf_x = ∇M_g(reg, M, x)
    hessf_x_ξ = ∇²M_g_ξ(reg, M, x, ξ)
    hessf_x_η = ∇²M_g_ξ(reg, M, x, η)

    # 1. Gradient is tangent.
    println("- gradf(x) ∈ T_x M:\t\t\t\t", is_tangent_vector(M, x, gradf_x; atol = 1e-10))

    # 2. Hessian is symetric
    println(
        "- Hess f(x)[η] ∈ T_x M:\t\t\t\t",
        is_tangent_vector(M, x, hessf_x_η, atol = 1e-10),
    )
    println(
        "- Hess f(x)[ξ] ∈ T_x M:\t\t\t\t",
        is_tangent_vector(M, x, hessf_x_ξ, atol = 1e-10),
    )
    println(
        "- ⟨Hess f(x)[ξ], η⟩ - ⟨Hess f(x)[η], ξ⟩:\t",
        inner(M, x, hessf_x_η, ξ) - inner(M, x, hessf_x_ξ, η),
    )

    f_x = f(embed(M, x))
    ηgradfx = inner(M, x, η, gradf_x)
    ηHessf_xη = inner(M, x, η, hessf_x_η)

    @show ηgradfx, ηHessf_xη

    frgrad(t) = abs(f(embed(M, retract(M, x, t * η))) - (f(embed(M, x)) + t * ηgradfx))
    function frhess(t)
        return abs(
            f(embed(M, retract(M, x, t * η))) - (f_x + t * ηgradfx + 0.5 * t^2 * ηHessf_xη),
        )
    end

    return compare_curves(frgrad, frhess)
end

function check_e2r_gradient_hessian(M)
    x = randomMPoint(M)
    p = embed(M, retract(M, x, randomTVector(M, x)))

    ## function, rgrad, rhess
    f(x) = exp(sum(x .* (x .- p) .^ 2))

    function get_rgrad(M::Manifold, x)
        egrad = ForwardDiff.gradient(f, embed(M, x))
        return egrad_to_rgrad(M, x, egrad)
    end

    function get_rhess(M::Manifold, x, η)
        x_emb = embed(M, x)
        η_emb = embed(M, x, η)
        ∇f_x = ForwardDiff.gradient(f, x_emb)
        ∇²f_x_η =
            reshape(ForwardDiff.hessian(f, x_emb) * vec(η_emb), representation_size(M))
        return ehess_to_rhess(M, x, ∇f_x, ∇²f_x_η, η)
    end

    ## Gradient check
    η = randomTVector(M, x)

    # f1(t) =
    # f_x = f(x)
    gradf_x = get_rgrad(M, x)

    # 1. gradf ∈ T_x M
    println("- gradf(x) ∈ T_x M:\t\t\t\t", is_tangent_vector(M, x, gradf_x; atol = 1e-10))
    println("- η ∈ T_x M:\t\t\t\t\t", is_tangent_vector(M, x, η; atol = 1e-10))

    # 2. Hessian is symetric
    η = randomTVector(M, x)
    ξ = randomTVector(M, x)

    hessf_x_η = get_rhess(M, x, η)
    hessf_x_ξ = get_rhess(M, x, ξ)

    println(
        "- Hess f(x)[η] ∈ T_x M:\t\t\t\t",
        is_tangent_vector(M, x, hessf_x_η, atol = 1e-10),
    )
    println(
        "- Hess f(x)[ξ] ∈ T_x M:\t\t\t\t",
        is_tangent_vector(M, x, hessf_x_ξ, atol = 1e-10),
    )
    println(
        "- ⟨Hess f(x)[ξ], η⟩ - ⟨Hess f(x)[η], ξ⟩:\t",
        inner(M, x, hessf_x_η, ξ) - inner(M, x, hessf_x_ξ, η),
    )

    f_x = f(embed(M, x))
    ηgradfx = inner(M, x, η, gradf_x)
    ηHessf_xη = inner(M, x, get_rhess(M, x, η), η)

    @show ηgradfx, ηHessf_xη

    frgrad(t) = abs(f(embed(M, retract(M, x, t * η))) - (f(embed(M, x)) + t * ηgradfx))
    function frhess(t)
        return abs(
            f(embed(M, retract(M, x, t * η))) - (f_x + t * ηgradfx + 0.5 * t^2 * ηHessf_xη),
        )
    end

    return compare_curves(frgrad, frhess)
end

function check_retraction(M)
    Random.seed!(1234)
    x = randomMPoint(M)
    Random.seed!(4312)
    η = randomTVector(M, x)

    x_emb = embed(M, x)
    η_emb = embed(M, x, η)

    function retractiontη(t)
        retracted = retract(M, x, t * η)
        return embed(M, retracted)
    end
    xtη(t) = x_emb + t * η_emb

    function retractionerror(t; x = x)
        return norm(project(M, x, retractiontη(t) - xtη(t)))
    end

    return compare_curves(retractionerror)
end

function remove_small_functionvals(values; threshold = 1e-12)
    return filter(x -> x[2] >= threshold, values)
end

function get_sloperesidual(data; threshold = 1e-12)
    n = length(data)
    x = [data[i][1] for i in 1:n]
    logy = [log(data[i][2]) for i in 1:n]

    ## Fit linear regressor
    # model : y = βx + c
    # explain Y = [log.(y) 1] by X = [log.(x)]
    # Y = X β
    # β ∈ ℝ 2x1 - slope and absciss at origin
    # X ∈ ℝ nx2 - absciss and intercept
    # Y ∈ ℝ nx1 - ordonate
    X = ones(n, 2)
    X[:, 1] = log.(x)
    Y = logy

    F = factorize(X' * X)
    β = F \ X' * Y
    predictedslope = β[1]
    residual = 0.5 * norm((Y - X * β)[:, 1])^2
    return predictedslope, residual
end

export get_sloperesidual, remove_small_functionvals


"""
A uniform mesh in one space dimension of `N` cells
"""
struct Uniform1DFVMesh{N,T}
  Δx :: T
  cell_centers :: Vector{T}
  cell_faces :: Vector{T}
end

getncells(mesh::Uniform1DFVMesh{N}) where {N} = N
getnnodes(mesh::Uniform1DFVMesh{N}) where {N} = N+1
getnodecoords(mesh::Uniform1DFVMesh, node_idx::Int) = mesh.cell_faces[node_idx]
getdimension(mesh::Uniform1DFVMesh) = 1

Base.summary(::Uniform1DFVMesh{N}) where {N} = string("FV 1D Mesh of ",N," cells")

function Base.show(io::IO, A::Uniform1DFVMesh)
  println(io,summary(A))
  print(io,"cells: ")
  show(io,getncells(A))
  println(io)
  print(io,"Δx: ")
  show(io, A.Δx)
  println(io)
  println(io,"Interval: [",A.cell_faces[1],",",A.cell_faces[end],"]")
end

TreeViews.hastreeview(x::Uniform1DFVMesh) = true
function TreeViews.treelabel(io::IO,x::Uniform1DFVMesh,
                             mime::MIME"text/plain" = MIME"text/plain"())
  show(io,mime,Text(Base.summary(x)))
end

function Uniform1DFVMesh(N::Int,bounds::Vector{T}) where {T}
    xend = bounds[end]
    xinit = bounds[1]
    L = xend - xinit
    dx = L/N
    xx = [i*dx+dx/2+xinit for i in 0:(N-1)]
    faces = [xinit + dx*i for i in 0:N]
    Uniform1DFVMesh{N,T}(dx,xx,faces)
end

@inline cell_volume(mesh::Uniform1DFVMesh) = mesh.Δx
@inline cell_volume(mesh::Uniform1DFVMesh,cell::Int) = mesh.Δx
cell_volumes(mesh::Uniform1DFVMesh) = mesh.Δx * ones(mesh.N)
get_nodes_matrix(mesh::Uniform1DFVMesh) = mesh.cell_faces
cell_centers(mesh::Uniform1DFVMesh) = mesh.cell_centers
cell_facets(mesh::Uniform1DFVMesh) = mesh.cell_faces


"""
    map_column_pairs(f, X::Matrix{Float64}[, Y::Matrix{Float64} = X])

Create a matrix by applying function `f` to each pair of columns of input matrices
`X` and `Y`.
"""
function map_column_pairs(f, X::MatF64, Y::MatF64)
    nobs1 = size(X,2)
    nobs2 = size(Y,2)
    D = Array{Float64}(undef, nobs1, nobs2)
    map_column_pairs!(D, f, X, Y)
    return D
end

function map_column_pairs(f, X::MatF64)
    dim, nobsv = size(X)
    D = Array{Float64}(undef, nobsv, nobsv)
    for i in 1:nobsv
        for j in 1:i
            @inbounds D[i,j] = f(X[:,i], X[:,j])
            if i != j; @inbounds D[j,i] = D[i,j]; end;
        end
    end
    return D
end

"""
    map_column_pairs!(D::Matrix{Float64}, f, X::Matrix{Float64}[, Y::Matrix{Float64} = X])

Like [`map_column_pairs`](@ref), but stores the result in `D` rather than a new matrix.
"""
function map_column_pairs!(D::MatF64, f, X::MatF64, Y::MatF64)
    dim, nobs1 = size(X)
    nobs2 = size(Y,2)
    dim == size(Y,1) || throw(ArgumentError("Input observation matrices must have consistent dimensions"))
    size(D,1) == nobs1 || throw(ArgumentError(@sprintf("D has %d rows, while X has %d columns (should be same)",
                                                       size(D,1), nobs1)))
    size(D,2) == nobs2 || throw(ArgumentError(@sprintf("D has %d columns, while Y has %d columns (should be same)",
                                                       size(D,2), nobs2)))
    for i in 1:nobs1, j in 1:nobs2
        @inbounds D[i,j] = f(view(X,:,i), view(Y,:,j))
    end
    return D
end

function map_column_pairs!(D::MatF64, f, X::MatF64)
    dim, nobsv = size(X)
    size(D,1) == nobsv || throw(ArgumentError(@sprintf("D has %d rows, while X has %d columns (should be same)",
                                                       size(D,1), nobsv)))
    size(D,2) == nobsv || throw(ArgumentError(@sprintf("D has %d columns, while X has %d columns (should be same)",
                                                       size(D,2), nobsv)))
    @inbounds for i in 1:nobsv
        for j in 1:i
            D[i,j] = f(view(X,:,i), view(X,:,j))
            if i != j; D[j,i] = D[i,j]; end;
        end
    end
    return D
end


# ---
# title: "`tophat`"
# id: op_tophat
# cover: assets/tophat.png
# ---

# The (white) tophat operator is defined as `img - opening(img)`. Intuitively, this filter
# can be used to extract small white elements and details from an image.

using ImageMorphology
using TestImages
using ImageBase
using ImageShow

#-

img = restrict(testimage_dip3e("fig1005")) # wirebond mask

out1 = tophat(img) # default: all spatial dimensions, r=1, a box-shape SE
out2 = tophat(img; dims=(2,)) # only apply to the second dimension
out3 = tophat(img; r=5) # half-size r=5
out4 = tophat(img, strel_diamond((5, 5))) # generic SE -- this version doesn't accept `r` and `dims`

mosaic(img, out1, out2, out3, out4; nrow=1)

# ## See also

# To extract black small details, use [`bothat`](@ref op_bothat).

bothat(img) == tophat(complement.(img))

# For a comprehensive and more accurate documentation, please check the [`tophat`](@ref)
# reference page.

# save cover image #src
using FileIO #src
mkpath("assets") #src
img = Gray.(restrict(testimage("blobs"))) #src
save("assets/tophat.png", tophat(img; r=2)) #src


using Gadfly, BBOBFunctions

##

function plot_fun(f)
    r = 5

    c, n, nline = Geom.contour(levels=exp10.(range(-6, stop=10, length=40))), 800, 1000
    if any(f .== [BBOBFunctions.F16,BBOBFunctions.F17,BBOBFunctions.F18,BBOBFunctions.F19])
        c, n, nline = Geom.contour(levels=exp10.(range(-6, stop=8, length=20))), 500, 1000
    end

    p1 = plot(
        layer(x=[f.x_opt[1]],y=[f.x_opt[2]],Geom.point,Theme(default_color=colorant"red")),
        layer(
            z=(x,y)->f([x,y])-f.f_opt, x=range(-r, stop=r, length=n), y=range(-r, stop=r, length=n),
            c,
        ),
        Coord.cartesian(xmin=-r,xmax=r,ymin=-r,ymax=r),
        Guide.title(string(f))
    )

    # line plot
    x = collect(range(-r, stop=r, length=nline))

    p2 = plot(
        layer(x=x,y=[ f([xi,f.x_opt[2],f.x_opt[3]]) for xi in x]-f.f_opt, Geom.line),
        layer(x=x,y=[ f([f.x_opt[1],xi,f.x_opt[3]]) for xi in x]-f.f_opt, Geom.line,Theme(default_color=colorant"orange")),
        layer(x=x,y=[ f([f.x_opt[1],f.x_opt[2],xi]) for xi in x]-f.f_opt, Geom.line,Theme(default_color=colorant"red")),
#        layer(x=[f.x_opt[1]],y=[f([f.x_opt[1],0,0])]-f.f_opt,Geom.point),
#        layer(x=[f.x_opt[2]],y=[f([0,f.x_opt[2],0])]-f.f_opt,Geom.point,Theme(default_color=colorant"orange")),
#        layer(x=x,y=[f([xi,xi,xi]) for xi in x]-f.f_opt,Geom.line,Theme(default_color=colorant"red")),
        Coord.cartesian(xmin=-r,xmax=r),
        Guide.title(string(f)),
        Scale.y_log10
    )
    p1,p2
end

outdir = joinpath(Pkg.dir(),"BBOBFunctions","data","plots","functions")

for i in 20:length(enumerate(BBOBFunctions.BBOBFunction))

    f = getfield(BBOBFunctions,Symbol("F$i"))
    info(f)
    p1,p2 = plot_fun(f)

    draw(PDF(joinpath(outdir,"F$i.pdf"),16cm,16cm),p1)
    draw(SVG(joinpath(outdir,"F$i.svg"),16cm,16cm),p1)
    draw(PDF(joinpath(outdir,"1D","F$i.pdf"),16cm,12cm),p2)
    draw(SVG(joinpath(outdir,"1D","F$i.svg"),16cm,12cm),p2)
end

##


cfs = [:s2, :l2, :c2, :surfsurf, :surfvoid]
dirs = Dict(
    1 => [:x],
    2 => [:x, :y, :xy, :yx],
    3 => [:x, :y, :z]
)

function test_cf_on_img(img, cf_map, cf_dir)
    for periodic in [true, false]
        cmap = cf_map(img; periodic)
        cdir = cf_dir(img, true; len=size(img, 1), periodic, directions=dirs[ndims(img)])

        for (direction, data) in cdir
            @test Map.dir_from_map(cmap, direction) ≈ data
        end
    end
end

function test_all()
    dims = [1, 2, 3]
    n = 10

    for dim in dims
        ns = Tuple(ones(Int, dim) * n)
        img = rand(Bool, ns)

        for cf in cfs
            @testset "random $(dim)D image, $(cf)" begin
                test_cf_on_img(img, Map.eval(cf), Directional.eval(cf))
            end
        end
    end
end

test_all()


using PolarimetryCalibration, Plots, LaTeXStrings

PArange = -60:1.0:60
model = Float64[(müller_rec(ϵ = -0.002,
                            ϕ = 0.003,
                            ΔG = -0.002,
                            ψ = 0.11,
#                    α = fit_param[7])* # Circular feed
                            α = deg2rad(45.0))'* # Circular feed
                 müller_sky(deg2rad(PAaz))'*
                 lin_pol_source(flux = 10.11,
                        p = 0.1,
                        pa = deg2rad(-64.0)))[i]
         for PAaz in PArange, i in 1:4]

# plot(plot(PArange, model[:,1], label = "I"), scatter!(pa, I, markersize = 3),
#      plot(PArange, model[:,2],label = "V"), scatter!(pa, V, markersize = 3),
#      plot(PArange, model[:,3],label = "U"), scatter!(pa, U, markersize = 3),
#      plot(PArange, -model[:,4],label = "Q"), scatter!(pa, Q, markersize = 3))

p1 = plot(PArange, model[:,1],title = "I")
p2 = plot(PArange, model[:,2],title = "V")
p3 = plot(PArange, model[:,3],title = "U")
p4 = plot(PArange, -model[:,4],title = "Q")
plot(p1, p2, p3, p4, legend = :none)


@defcomp climaten2ocycle begin
    # Global N2O emissions in Mt of N
    globn2o = Parameter(index=[time])

    # Atmospheric N2O concentration
    acn2o = Variable(index=[time])

    # N2O decay
    n2odecay = Variable()

    #
    lifen2o = Parameter(default = 114.0)

    # N2o pre industrial
    n2opre = Parameter(default = 285.0)

    function run_timestep(p, v, d, t)

        if is_first(t)
            v.n2odecay = 1.0 / p.lifen2o

            v.acn2o[t] = 296
        else
            # Calculate N2O concentrations
            v.acn2o[t] = v.acn2o[t - 1] + 0.2079 * p.globn2o[t] - v.n2odecay * (v.acn2o[t - 1] - p.n2opre)

            if v.acn2o[t] < 0
                error("n2o atmospheric concentration out of range")
            end
        end
    end
end


struct FDJONSWAP <: FDSpectralModel
    θ::Vector{Float64}
end

parameter(J::FDJONSWAP) = J.θ

function sdf(J::Union{FDJONSWAP,FDSwellJS}, ω)
    α = J.θ[1]
    ωₚ = J.θ[2]
    γ = J.θ[3]
    r = J.θ[4]
    ω = abs(ω)
    s = 4.0
    σ1² = 0.0049 + 0.0032 * (ω > ωₚ)
    δ = exp(-1 / (2 * σ1²) * (ω / ωₚ - 1)^2)
    sdf = α * ω^(-r) * exp(-(r / s) * (ω / ωₚ)^(-s)) * γ^δ / 2
    sdf *= (!isnan(sdf))
    return sdf
end

function spreading(J::FDJONSWAP, ω, ϕ)
    ω = abs(ω)
    halfPeakSep = J.θ[6] / exp(J.θ[7] * ((ω > J.θ[2])*((J.θ[2] / ω) - 1.0) + 1.0)) / 2
    φₘ₁ = J.θ[5] + halfPeakSep
    φₘ₂ = J.θ[5] - halfPeakSep
    σ = J.θ[8] - (J.θ[9]/3)*(4*(J.θ[2]/ω)^2 - (J.θ[2]/ω)^8)
    return 0.5 * (WrappedGaussian(ϕ, φₘ₁, σ) + WrappedGaussian(ϕ, φₘ₂, σ))
end