export AffineSubspace,
    AffineExtrinsic,
    AffineIntrinsic,
    Coordinates,
    Intrinsic,
    Extrinsic,
    intrinsic,
    extrinsic,
    dim,
    codim,
    ambient_dim,
    rand_affine_subspace,
    coord_change,
    geodesic_distance,
    geodesic

# References
const _LKK19 = """Lim, Lek-Heng, Ken Sze-Wai Wong, and Ke Ye. "Numerical algorithms on the affine Grassmannian." SIAM Journal on Matrix Analysis and Applications 40.2 (2019): 371-393"""

"""
    Coordinates

A type used for encoding the used coordinates and for performing coordinate changes.

Currently supported coordinates are:
- [`Intrinsic`](@ref)
- [`Extrinsic`](@ref)
"""
struct Coordinates{T} end
Base.broadcastable(C::Coordinates) = Ref(C)

"""
    Intrinsic <: Coordinates

Indicates the use of the intrinsic description of an affine subspace. See also [`AffineIntrinsic`](@ref).
"""
const Intrinsic = Coordinates{:Intrinsic}()

"""
    Extrinsic <: Coordinates

Indicates the use of the extrinsic description of an affine subspace. See also [`AffineExtrinsic`](@ref).
"""
const Extrinsic = Coordinates{:Extrinsic}()
# These are not part of the public API
const IntrinsicStiefel = Coordinates{:IntrinsicStiefel}()
const ExtrinsicStiefel = Coordinates{:ExtrinsicStiefel}()

"""
    AffineExtrinsic

Extrinsic description of an ``m``-dimensional affine subspace ``L`` in ``n``-dimensional space.
That is ``L = \\{ x | A x = b \\}``.
"""
struct AffineExtrinsic{T}
    A::Matrix{T}
    b::Vector{T}
end
(A::AffineExtrinsic)(x::AbstractVector) = A.A * x - A.b

"""
    dim(A::AffineExtrinsic)

Dimension of the affine subspace `A`.
"""
dim(A::AffineExtrinsic) = size(A.A, 2) - size(A.A, 1)

"""
    codim(A::AffineExtrinsic)

Codimension of the affine subspace `A`.
"""
codim(A::AffineExtrinsic) = size(A.A, 1)

Base.:(==)(A::AffineExtrinsic, B::AffineExtrinsic) = A.A == B.A && A.b == B.b
function Base.copy!(A::AffineExtrinsic, B::AffineExtrinsic)
    copy!(A.A, B.A)
    copy!(A.b, B.b)
    A
end
Base.copy(A::AffineExtrinsic) = AffineExtrinsic(copy(A.A), copy(A.b))

function Base.show(io::IO, A::AffineExtrinsic{T}) where {T}
    println(io, "AffineExtrinsic{$T}:")
    println(io, "A:")
    show(io, A.A)
    println(io, "\nb:")
    show(io, A.b)
end

Base.broadcastable(A::AffineExtrinsic) = Ref(A)

"""
    AffineIntrinsic

Intrinsic description of an ``m``-dimensional affine subspace ``L`` in ``n``-dimensional space.
That is ``L = \\{ u | A u + b₀ \\}``. Here, ``A`` and ``b₀`` are in orthogonal coordinates.
That is, the columns of ``A`` are orthonormal and ``A' b₀ = 0``.
"""
struct AffineIntrinsic{T}
    # orthogonal coordinates
    A::Matrix{T}
    b₀::Vector{T}
    # stiefel coordinates
    Y::Matrix{T}
end
AffineIntrinsic(A::Matrix, b₀::Vector) = AffineIntrinsic(A, b₀, stiefel_coordinates(A, b₀))

(A::AffineIntrinsic)(u::AbstractVector, ::Coordinates{:Intrinsic}) = A.A * u + A.b₀
(A::AffineIntrinsic)(u::AbstractVector, ::Coordinates{:IntrinsicStiefel}) = A.Y * u

function stiefel_coordinates(A::AbstractMatrix, b::AbstractVector)
    γ = sqrt(1 + sum(abs2, b))
    n, k = size(A)
    Y = zeros(eltype(A), n + 1, k + 1)
    Y[1:n, 1:k] .= A
    Y[1:n, k+1] .= b ./ γ
    Y[n+1, k+1] = 1 / γ
    Y
end

"""
    dim(A::AffineIntrinsic)

Dimension of the affine subspace `A`.
"""
dim(I::AffineIntrinsic) = size(I.A, 2)

"""
    codim(A::AffineIntrinsic)

Codimension of the affine subspace `A`.
"""
codim(I::AffineIntrinsic) = size(I.A, 1) - size(I.A, 2)

function Base.:(==)(A::AffineIntrinsic, B::AffineIntrinsic)
    A.A == B.A && A.b₀ == B.b₀ && A.Y == B.Y
end

function Base.show(io::IO, A::AffineIntrinsic{T}) where {T}
    println(io, "AffineIntrinsic{$T}:")
    println(io, "A:")
    show(io, A.A)
    println(io, "\nb₀:")
    show(io, A.b₀)
end

function Base.copy!(A::AffineIntrinsic, B::AffineIntrinsic)
    copy!(A.A, B.A)
    copy!(A.b₀, B.b₀)
    copy!(A.Y, B.Y)
    A
end
Base.copy(A::AffineIntrinsic) = AffineIntrinsic(copy(A.A), copy(A.b₀), copy(A.Y))

Base.broadcastable(A::AffineIntrinsic) = Ref(A)

function AffineIntrinsic(E::AffineExtrinsic)
    svd = LA.svd(E.A; full = true)
    b₀ = svd \ E.b
    m, n = size(E.A)
    A = Matrix((@view svd.Vt[m+1:end, :])')
    AffineIntrinsic(A, b₀)
end

function AffineExtrinsic(I::AffineIntrinsic)
    svd = LA.svd(I.A; full = true)
    m, n = size(I.A)
    A = Matrix((@view svd.U[:, n+1:end])')
    b = A * I.b₀
    AffineExtrinsic(A, b)
end

"""
    AffineSubspace(A, b)

An ``m``-dimensional affine subspace ``L`` in ``n``-dimensional space given by
the extrinsic description ``L = \\{ x | A x = b \\}``.

```julia-repl
julia> A = AffineSubspace([1 0 3; 2 1 3], [5, -2])
1-dim. affine subspace {x|Ax=b} with eltype Float64:
A:
2×3 Array{Float64,2}:
 1.0  0.0  3.0
 2.0  1.0  3.0
b:
2-element Array{Float64,1}:
  5.0
 -2.0

julia> dim(A)
1

julia> codim(A)
2

julia> ambient_dim(A)
3
```

An `AffineSubspace` holds always its [`extrinsic`](@ref) description, see also [`AffineIntrinsic`](@ref),
as well as its [`intrinsic`](@ref) description, see [`AffineExtrinsic`](@ref).

```julia-repl
julia> intrinsic(A)
AffineIntrinsic{Float64}:
A:
3×1 Array{Float64,2}:
 -0.6882472016116853
  0.6882472016116853
  0.22941573387056186
b₀:
3-element Array{Float64,1}:
 -3.0526315789473677
 -3.947368421052632
  2.684210526315789
```

An `AffineSubspace` can be evaluated with either using [`Intrinsic`](@ref)
or [`Extrinsic`](@ref) coordinates.

```julia-repl
julia> u = [0.5]
1-element Array{Float64,1}:
 0.5

julia> x = A(u, Intrinsic)
3-element Array{Float64,1}:
 -3.3967551797532103
 -3.6032448202467893
  2.79891839325107

julia> A(x, Extrinsic)
  2-element Array{Float64,1}:
   0.0
   0.0
```

To change the used coordinates you can use [`coord_change`](@ref).
```julia-repl
julia> coord_change(A, Extrinsic, Intrinsic, x)
1-element Array{Float64,1}:
 0.49999999999999994
```
"""
struct AffineSubspace{T}
    extrinsic::AffineExtrinsic{T}
    intrinsic::AffineIntrinsic{T}
end

AffineSubspace(I::AffineIntrinsic) = AffineSubspace(AffineExtrinsic(I), I)
AffineSubspace(E::AffineExtrinsic) = AffineSubspace(E, AffineIntrinsic(E))

function AffineSubspace(A::AbstractMatrix{T}, b::AbstractVector{T}) where {T}
    size(A, 1) == length(b) || throw(ArgumentError("Size of A and b not compatible."))
    0 < size(A, 1) < size(A, 2) ||
        throw(ArgumentError("Affine subspace has to be given in extrinsic coordinates, i.e., by A x = b."))

    AffineSubspace(AffineExtrinsic(Matrix(float.(A)), Vector(float.(b))))
end

Base.broadcastable(A::AffineSubspace) = Ref(A)

"""
    dim(A::AffineSubspace)

Dimension of the affine subspace `A`.
"""
dim(A::AffineSubspace) = dim(A.intrinsic)

"""
    codim(A::AffineSubspace)

Codimension of the affine subspace `A`.
"""
codim(A::AffineSubspace) = codim(A.intrinsic)

"""
    ambient_dim(A::AffineSubspace)

Dimension of ambient space of the affine subspace `A`.
"""
ambient_dim(A::AffineSubspace) = dim(A) + codim(A)

function Base.show(io::IO, A::AffineSubspace{T}) where {T}
    println(io, "$(dim(A))-dim. affine subspace {x|Ax=b} with eltype $T:")
    println(io, "A:")
    show(io, A.extrinsic.A)
    println(io, "\nb:")
    show(io, A.extrinsic.b)
end

"""
    intrinsic(A::AffineSubspace)

Obtain the intrinsic description of `A`, see also [`AffineIntrinsic`](@ref).
"""
intrinsic(A::AffineSubspace) = A.intrinsic

"""
    extrinsic(A::AffineSubspace)

Obtain the extrinsic description of `A`, see also [`AffineExtrinsic`](@ref).
"""
extrinsic(A::AffineSubspace) = A.extrinsic

function Base.copy!(A::AffineSubspace, B::AffineSubspace)
    copy!(A.intrinsic, B.intrinsic)
    copy!(A.extrinsic, B.extrinsic)
    A
end
Base.copy(A::AffineSubspace) = AffineSubspace(copy(A.extrinsic), copy(A.intrinsic))

function Base.:(==)(A::AffineSubspace, B::AffineSubspace)
    intrinsic(A) == intrinsic(B) && extrinsic(A) == extrinsic(B)
end

function (A::AffineSubspace)(x::AbstractVector, ::Coordinates{:Intrinsic})
    intrinsic(A)(x, Intrinsic)
end
function (A::AffineSubspace)(x::AbstractVector, ::Coordinates{:IntrinsicStiefel})
    intrinsic(A)(x, IntrinsicStiefel)
end
function (A::AffineSubspace)(x::AbstractVector, ::Coordinates{:Extrinsic})
    extrinsic(A)(x)
end

"""
    rand_affine_subspace(n::Integer; dim | codim, real = false)

Generate a random [`AffineSubspace`](@ref) with given dimension `dim` or codimension `codim`
(one of them has to be provided) in ambient space of dimension `n`.
If `real` is `true`, then the extrinsic description is real.
The subspace is generated by drawing each entry of the extrinsic description indepdently
from a normal distribuation using `randn`.

## Example
```julia-repl
julia> rand_affine_subspace(3; dim = 1)
1-dim. affine subspace {x|Ax=b} with eltype Complex{Float64}:
A:
2×3 Array{Complex{Float64},2}:
  -1.73825+1.27987im   -0.0871343+0.840408im  -0.551957+0.106397im
 -0.597132-0.343965im   -0.122543-0.172715im   -1.04949+0.370917im
b:
2-element Array{Complex{Float64},1}:
  0.47083334430689394 + 0.8099804422599071im
 -0.12018696822943896 + 0.11723026326952792im

julia> rand_affine_subspace(4; codim = 1)
3-dim. affine subspace {x|Ax=b} with eltype Complex{Float64}:
A:
1×4 Array{Complex{Float64},2}:
 0.345705+0.0893881im  -0.430867-0.663249im  0.979969-0.569378im  -0.29722-0.192493im
b:
1-element Array{Complex{Float64},1}:
 0.7749708228192062 + 0.9762873764567546im
```
"""
function rand_affine_subspace(
    n::Integer;
    dim::Union{Nothing,Integer} = nothing,
    codim::Union{Nothing,Integer} = nothing,
    real::Bool = false,
)
    !isnothing(dim) ||
        !isnothing(codim) ||
        throw(ArgumentError("Neither `dim` or `codim` specified."))

    if !isnothing(dim)
        0 < dim < n || throw(ArgumentError("`dim` has to be between 0 and `n`."))
        k = dim
    else
        0 < codim < n || throw(ArgumentError("`codim` has to be between 0 and `n`."))
        k = n - codim
    end
    T = real ? Float64 : ComplexF64
    A = randn(T, n - k, n)
    b = randn(T, n - k)

    AffineSubspace(A, b)
end

# Coordinate changes
"""
    coord_change(A::AffineSubspace, C₁::Coordinates, C₂::Coordinates, p)

Given an affine subspace `A` and a point `p` in coordinates `C₁` compute the point
`x` describing p in coordinates `C₂`.

## Example
```julia-repl
julia> A = AffineSubspace([1 0 3; 2 1 3], [5, -2]);

julia> u = [1.25];

julia> x = coord_change(A, Intrinsic, Extrinsic, u)
3-element Array{Float64,1}:
 -3.9129405809619744
 -3.087059419038025
  2.9709801936539915

julia> A(x, Extrinsic)
2-element Array{Float64,1}:
 0.0
 0.0

julia> x - A(u, Intrinsic)
3-element Array{Float64,1}:
 0.0
 0.0
 0.0
```
"""
coord_change(A::AffineSubspace, ::C, ::C, x) where {C<:Coordinates} = x
coord_change(A::AffineSubspace, ::Coordinates{:Intrinsic}, ::Coordinates{:Extrinsic}, u) =
    A(u, Intrinsic)
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:Intrinsic},
    ::Coordinates{:IntrinsicStiefel},
    u,
)
    [
        u
        1 / A.intrinsic.Y[end, end]
    ]
end
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:Intrinsic},
    ::Coordinates{:ExtrinsicStiefel},
    u,
)
    v = coord_change(A, Intrinsic(), IntrinsicStiefel(), u)
    coord_change(A, IntrinsicStiefel(), ExtrinsicStiefel(), v)
end

coord_change(A::AffineSubspace, ::Coordinates{:Extrinsic}, ::Coordinates{:Intrinsic}, x) =
    A.intrinsic.A' * (x - A.intrinsic.b₀)
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:Extrinsic},
    ::Coordinates{:IntrinsicStiefel},
    x,
)
    u = coord_change(A, Extrinsic(), Intrinsic(), x)
    coord_change(A, Intrinsic(), IntrinsicStiefel(), u)
end
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:Extrinsic},
    ::Coordinates{:ExtrinsicStiefel},
    x,
)
    [x; 1]
end
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:IntrinsicStiefel},
    ::Coordinates{:Intrinsic},
    u,
)
    γ = A.intrinsic.Y[end, end]
    [u[i] / (γ * u[end]) for i = 1:(length(u)-1)]
end
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:IntrinsicStiefel},
    ::Coordinates{:Extrinsic},
    u,
)
    x = coord_change(A, IntrinsicStiefel(), ExtrinsicStiefel(), u)
    coord_change(A, ExtrinsicStiefel(), Extrinsic(), x)
end
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:IntrinsicStiefel},
    ::Coordinates{:ExtrinsicStiefel},
    u,
)
    A(u, IntrinsicStiefel)
end
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:ExtrinsicStiefel},
    ::Coordinates{:Intrinsic},
    x,
)
    u = coord_change(A, ExtrinsicStiefel(), IntrinsicStiefel(), x)
    coord_change(A, IntrinsicStiefel(), Intrinsic(), u)
end
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:ExtrinsicStiefel},
    ::Coordinates{:Extrinsic},
    x,
)
    x[1:end-1] ./ x[end]
end
function coord_change(
    A::AffineSubspace,
    ::Coordinates{:ExtrinsicStiefel},
    ::Coordinates{:IntrinsicStiefel},
    x,
)
    A.intrinsic.Y' * x
end

"""
    geodesic_distance(A::AffineSubspace, B::AffineSubspace)

Compute the geodesic distance between `A` and `B` in the affine Grassmanian `Graff(k, n)`
where `k = dim(A)` and `n` is the amebient dimension.
This follows the derivation in [^LKK19].

[^LKK19]: $_LKK19.
"""
geodesic_distance(A::AffineSubspace, B::AffineSubspace) =
    geodesic_distance(A.intrinsic, B.intrinsic)
function geodesic_distance(A::AffineIntrinsic, B::AffineIntrinsic)
    sqrt(sum(σᵢ -> acos(min(σᵢ, 1.0))^2, LA.svdvals(A.Y' * B.Y)))
end

# geodesic
"""
    geodesic_svd(A::AffineSubspace, B::AffineSubspace)

Computes the factors ``Q``, ``Θ`` and ``U`` from Corollary 4.3 in [^LKK19].
These values are necessary to construct the geodesic between `A` and `B`.

[^LKK19]: $_LKK19
"""
geodesic_svd(A::AffineSubspace, B::AffineSubspace) = geodesic_svd(A.intrinsic, B.intrinsic)
function geodesic_svd(A::AffineIntrinsic, B::AffineIntrinsic)
    U, Σ, V = LA.svd(A.Y' * B.Y)
    # Θ = acos.(min.(Σ, 1.0))
    # We have acos(1-eps()) = 2.1073424255447017e-8
    # So this is numerically super unstable if the singular value is off by only eps()
    # We try to correct this fact by treating σ >= 1 - 2eps() as 1.0
    Θ = map(σ -> σ + 2 * eps() > 1.0 ? 0.0 : acos(σ), Σ)

    n, k = size(A.Y)
    # M = (LA.I - A.Y * A.Y') * B.Y * inv(A.Y' * B.Y)
    # Have to compute an SVD of M s.t. M = Q tanΘ U'
    # Equivalently M * U = Q tan(Θ)
    # We can achieve this by using a *pivoted* QR
    # since then R will be a diagonal matrix s.t. the absolute value of R_ii is θ_{k-i}

    # inv(A.Y' * B.Y) * U = V * LA.diagm(inv.(Σ))
    MU = (LA.I - A.Y * A.Y') * B.Y * V * LA.diagm(inv.(Σ))
    Q, R = LA.qr(MU, Val(true))
    # correct signs and ordering of Q
    Q′ = Q[:, k:-1:1]
    for j = 1:k
        # Look if we need to flip signs
        real(R[k-j+1, k-j+1]) < 0 || continue
        for i = 1:n
            Q′[i, j] = -Q′[i, j]
        end
    end

    Q′, Θ, U
end

"""
    geodesic(A::AffineSubspace, B::AffineSubspace)

Returns the geodesic ``γ(t)`` connecting `A` and `B` in the Grassmanian ``Gr(k+1,n+1)``
where ``k`` is the dimension of ``A`` and ``n`` is the ambient dimension.
See also Corollary 4.3 in [^LKK19].

[^LKK19]: $_LKK19
"""
geodesic(A::AffineSubspace, B::AffineSubspace) = geodesic(A.intrinsic, B.intrinsic)
function geodesic(A::AffineIntrinsic, B::AffineIntrinsic)
    Q, Θ, U = geodesic_svd(A, B)
    t -> A.Y * U * LA.diagm(cos.(t .* Θ,)) * U' + Q * LA.diagm(sin.(t .* Θ,)) * U'
end
#


include(joinpath(pwd(), "examples/implicit_dynamics/utils.jl"))

"""
    Hopper2D
    	model inspired by "Dynamically Stable Legged Locomotion"
		s = (x, z, t, r)
			x - lateral position
			z - vertical position
			t - body orientation
			r - leg length
"""
mutable struct Hopper2D{T}
    dim::Dimensions

    mb::T # mass of body
    ml::T # mass of leg
    Jb::T # inertia of body
    Jl::T # inertia of leg

    μ_world::T  # coefficient of friction
    μ_joint::T  # gravity
	g::T

	joint_friction::SVector
end

lagrangian(model::Hopper2D, q, q̇) = 0.0

# Kinematics
function kinematics(::Hopper2D, q)
	[q[1] + q[4] * sin(q[3]),
	 q[2] - q[4] * cos(q[3])]
end

# Methods
function M_func(model::Hopper2D, q)
	Diagonal([model.mb + model.ml,
    					   model.mb + model.ml,
    					   model.Jb + model.Jl,
    					   model.ml])
 end

function C_func(model::Hopper2D, q, q̇)
	[0.0,
	  (model.mb + model.ml) * model.g,
	  0.0,
	  0.0]
end

function ϕ_func(model::Hopper2D, q)
    [q[2] - q[4] * cos(q[3])]
end

function J_func(::Hopper2D, q)
    [1.0 0.0 (q[4] * cos(q[3])) sin(q[3]);
		          0.0 1.0 (q[4] * sin(q[3])) (-1.0 * cos(q[3]))]
end

function B_func(::Hopper2D, q)
	[0.0 0.0 1.0 0.0;
                -sin(q[3]) cos(q[3]) 0.0 1.0]
end

# Working Parameters
gravity = 9.81 # gravity
μ_world = 0.8 # coefficient of friction
μ_joint = 0.0

# TODO: change to Raibert parameters
mb = 3.0 # body mass
ml = 0.3  # leg mass
Jb = 0.75 # body inertia
Jl = 0.075 # leg inertia

# Dimensions
nq = 4
nu = 2
nw = 2
nc = 1
nb = 2

model = Hopper2D(Dimensions(nq, nu, nw, nc),
			   mb, ml, Jb, Jl,
			   μ_world, μ_joint, gravity,
			   SVector{4}(zeros(4)))

function lagrangian_derivatives(model::Hopper2D, q, v)
	D1L = -1.0 * C_func(model, q, v)
    D2L = M_func(model, q) * v
	return D1L, D2L
end

lagrangian_derivatives(model, rand(nq), rand(nq))

function dynamics(model::Hopper2D, h, q0, q1, u1, λ1, q2)
    # evalutate at midpoint
    qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

    D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
    D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

    return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
    	+ transpose(B_func(model, qm2)) * u1
        + transpose(J_func(model, q2)) * λ1)
        # -h[1] * model.joint_friction .* vm2)
end

function residual(model, z, θ, κ)
    nq = model.dim.q
    nu = model.dim.u
    nc = model.dim.c
    nb = nc * 2

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    μ = model.μ_world

    q2 = z[1:nq]
    γ1 = z[nq .+ (1:nc)]
    b1 = z[nq + nc .+ (1:nb)]
    ψ1 = z[nq + nc + nb .+ (1:nc)]
    η1 = z[nq + nc + nb + nc .+ (1:nb)]
    s1 = z[nq + nc + nb + nc + nb .+ (1:nc)]
    s2 = z[nq + nc + nb + nc + nb + nc .+ (1:nc)]

	ϕ = ϕ_func(model, q2)

	λ1 = [b1[1] - b1[2]; γ1]
    vT = (J_func(model, q2) * (q2 - q1) / h[1])[1]
	vT_stack = [vT; -vT]
	ψ_stack = ψ1 .* ones(nb)

	[
     dynamics(model, h, q0, q1, u1, λ1, q2);
	 s1 - ϕ;
	 vT_stack + ψ_stack - η1;
	 s2 .- (μ[1] * γ1 .- sum(b1));
	 γ1 .* s1 .- κ;
	 b1 .* η1 .- κ;
	 ψ1 .* s2 .- κ
    ]
end

nz = nq + nc + nb + nc + nb + nc + nc
nθ = nq + nq + nu + 1

idx_ineq = collect(nq .+ (1:(nc + nb + nc + nb + nc + nc)))
z_subset_init = 0.1 * ones(nc + nb + nc + nb + nc + nc)

# Declare variables
@variables z[1:nz]
@variables θ[1:nθ]
@variables κ[1:1]

# Residual
r = residual(model, z, θ, κ)
r = Symbolics.simplify.(r)
rz = Symbolics.jacobian(r, z, simplify = true)
rθ = Symbolics.jacobian(r, θ, simplify = true) # TODO: sparse version

# Build function
r_func = eval(build_function(r, z, θ, κ)[2])
rz_func = eval(build_function(rz, z, θ)[2])
rθ_func = eval(build_function(rθ, z, θ)[2])

rz_array = similar(rz, Float64)
rθ_array = similar(rθ, Float64)

@save joinpath(@__DIR__, "dynamics/residual.jl") r_func rz_func rθ_func rz_array rθ_array
@load joinpath(@__DIR__, "dynamics/residual.jl") r_func rz_func rθ_func rz_array rθ_array


Random.seed!(2019)

# Dynamics
dt=0.1
function dynamics(x,u)
    v,ϕ = u
    xp = x
    xp += dt*([v*cos(x[3]);v*sin(x[3]); ϕ])
    xp[3] = xp[3]
    return xp
end

W = dt*0.05*Matrix{Float64}(I,3,3)

dmodel = NonlinearDynamicsModel(dynamics,W)

# Measurement
observe(x,u) = [norm(x[1:2])]

V = 0.1*Matrix{Float64}(I,1,1)

omodel = NonlinearObservationModel(observe,V)

# Build UKF
ukf = UnscentedKalmanFilter(dmodel,omodel);

# Simulation
times = 0:dt:5
action_sequence = [[1.0, 2*sin(pi/2*t)] for t in times]

b0 = GaussianBelief([0.0,0.0,0.0], 0.01*Matrix{Float64}(I,3,3))

sim_states, sim_measurements = simulation(ukf,b0,action_sequence);

# Filter
filtered_beliefs = run_filter(ukf, b0, action_sequence, sim_measurements)
μ, Σ = unpack(filtered_beliefs)

# Tests
@test ukf.d isa NonlinearDynamicsModel
@test ukf.o isa NonlinearObservationModel
@test length(sim_measurements) == length(action_sequence)

points, w_μ, w_Σ = @inferred unscented_transform(b0, ukf.λ, ukf.α, ukf.β)
bp = @inferred unscented_transform_inverse(points, w_μ, w_Σ)

first_update = update(ukf, b0, action_sequence[1], sim_measurements[1])
first_predict = predict(ukf, b0, action_sequence[1])
first_predict_measure = measure(ukf, first_predict, sim_measurements[1],
                        u = action_sequence[1])
array_isapprox(first_update.μ, first_predict_measure.μ)
array_isapprox(first_update.Σ, first_predict_measure.Σ)

@test length(filtered_beliefs[end].μ) == length(b0.μ)
@test size(filtered_beliefs[end].Σ) == size(b0.Σ)
@test size(μ) == (length(filtered_beliefs), length(filtered_beliefs[1].μ))
@test size(Σ) == (length(filtered_beliefs), size(filtered_beliefs[1].Σ)...)

# change with algorithm / rng changes

array_isapprox(sim_states[end], [0.16459, 3.34923, 2.65721];atol=0.001)
array_isapprox(sim_measurements[end], [3.17086];atol=0.001)
array_isapprox(filtered_beliefs[end].μ, [1.83418, 2.78963, 1.37579];
                atol=0.001)
array_isapprox(filtered_beliefs[end].Σ, [0.739291 -0.472135 -0.282821;
                                        -0.472135 0.342115 0.197056;
                                        -0.282821 0.197056 0.185713]; atol=0.001)


using Alistair

# PERFORMANCE COMPARISON OPTIONS
const TEST_SIZE_1 = 1000

Profile.clear()

#=
println("---- LINEAR REGRESSION ----")

beta = [0.7; 0.45]
X = hcat(ones(TEST_SIZE_1,1), collect(1.0:TEST_SIZE_1))
Y = X * beta + 0.1*randn(TEST_SIZE_1, 1)

olsreg = OLS()

olsreg(X, Y)
@time res = olsreg(X, Y)

linregress(X, Y, olsreg)
@time res = linregress(X, Y, olsreg)

linregress(X, Y, OLS(intercept=true, robust=HCEVariance()))
@time res = linregress(X, Y, OLS(intercept=true, robust=HCEVariance()))

linregress(X, Y, FGLS(intercept=true, robust=HCEVariance()))
@time res = linregress(X, Y, FGLS(intercept=true, robust=HCEVariance()))

linregress(X, Y, IteratedFGLS(10, intercept=true, robust=HCEVariance()))
@time res = linregress(X, Y, IteratedFGLS(10, intercept=true, robust=HCEVariance()))

println("---- NONLINEAR REGRESSION ----")

f(X_, beta_) = exp.(beta_[1]*X_[:,1] + beta_[2]*log.(X_[:,2]))

Ynlin = f(X, beta) + 0.1*randn(TEST_SIZE_1, 1)

OptimizeRegression = Optimize(f, [1.0, 1.0])

res = nlinfit(X, Ynlin, OptimizeRegression)
@time res = nlinfit(X, Ynlin, OptimizeRegression)
=#

println("---- GENERAL LINEAR MODEL REGRESSION ----")

beta_cat = [0.5; 2.0]
X_cat = hcat(ones(TEST_SIZE_1,1), collect(linspace(-2.2,2.0,TEST_SIZE_1)))
Y_cat = round.(Alistair.logistic(X_cat * beta_cat))

res_logit = glmregress(X_cat, Y_cat, Logit())
@time res_logit = glmregress(X_cat, Y_cat, Logit())
println(res_logit)


# ---
# title: 719. Find K-th Smallest Pair Distance
# id: problem719
# author: Qling
# date: 2021-04-07
# difficulty: Hard
# categories: Array, Binary Search, Heap
# link: <https://leetcode.com/problems/find-k-th-smallest-pair-distance/description/>
# hidden: true
# ---
#
# Given an integer array, return the k-th smallest **distance** among all the
# pairs. The distance of a pair (A, B) is defined as the absolute difference
# between A and B.
#
# **Example 1:**
#
#
#
#     Input:
#     nums = [1,3,1]
#     k = 1
#     **Output: 0**
#     Explanation:
#     Here are all the pairs:
#     (1,3) -> 2
#     (1,1) -> 0
#     (3,1) -> 2
#     Then the 1st smallest distance pair is (1,1), and its distance is 0.
#
#
# **Note:**
#
#   1. `2 <= len(nums) <= 10000`.
#   2. `0 <= nums[i] < 1000000`.
#   3. `1 <= k <= len(nums) * (len(nums) - 1) / 2`.
#
#
## @lc code=start
using LeetCode

function smallest_distance_pair(nums::Vector{Int}, k::Int)::Int
    sort!(nums)
    n = length(nums)
    function count_not_greater(diff)
        i, ans = 1, 0
        for j in 2: n
            while nums[j] - nums[i] > diff
                i += 1
            ans += j - i
            end
        end

        return ans
    end

    left, right = 0, nums[end] - nums[1]
    while left <= right
        mid = (left + right) >> 1
        if count_not_greater(mid) >= k
            right = mid - 1
        else
            left = mid + 1
        end
    end

    return left
end
## @lc code=end


using Optimize
using NLPModels

using Plots
# pyplot()
plotly()

n_min = 2
n_max = 200

TestCUTEst = false

# Math Prog Base collection
include("MPBProblems.jl")
test_probs = mpb_probs

# Select solvers
using LineSearch
using Stopping
using LSDescentMethods

#stop_norm = ARCTR.stop_norm
atol = 1.0e-5
rtol = 1.0e-10

solvers = [CG_HZ, CG_HZ, CG_HZ, CG_HZ, CG_HZ]
labels = ["CG_HZ Armijo backtracking";"CG_HZ HZ"; "CG_HZ MT"; "CG_HZ TR_SecA"; "CG_HZ TR_Nwt"]

stp = TStopping(atol = atol, rtol = rtol, max_iter = 1000000, max_eval = 1000000, max_time = 100.0)

options = [Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => Newarmijo_wolfe)
           Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => _hagerzhang2!, :τ₀ => 0.01, :mayterminate => true, :τ₁ => 0.1)
           Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => _morethuente2!, :τ₀ => 0.01, :τ₁ => 0.1)
           Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => TR_SecA_ls, :τ₀ => 0.000000010000000, :τ₁ => 0.092999999998146)
           Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => TR_Nwt_ls, :τ₀ => 0.0001, :τ₁ => 0.001)
           ]

include("compare_solvers.jl")

P4 = compare_solvers_with_options2(solvers, options, labels, test_probs, n_min, n_max, printskip = false)[2]


using LinearAlgebra, Statistics, PosDefManifold, Test,
      Diagonalizations

# for more tests see the `Examples` folder

const err=1e-6

# compare two matrices
function compare(a::AbstractArray, b::AbstractArray)
     @test norm(a-b)/length(a) < err
end

# check that cov(F'XX'F)=D arranging terms for the two possible values of `dims`
function compare(d::Int, X::AbstractArray, F::AbstractArray, D::AbstractArray)
    d==1 ? Y=X*F : Y=F'*X
    d==1 ? C=((Y'*Y)/size(Y, 1))[1:size(D, 1), 1:size(D, 2)] :
           C=((Y*Y')/size(Y, 2))[1:size(D, 1), 1:size(D, 2)]
    compare(C, D)
end

# check that D is approximately diagonal
isApproxD(D)=norm(D-Diagonal(D))<err

####  Create data for testing the case k=1, m>1
# `t` is the number of samples, e.g.,
    # in Tangent space transfer learning is the number of trials,
# `m` is the number of datasets, e.g., subjects
# `n` is the number of variables, e.g.,
    # in Tangent space transfer learning is the number of elements of the tangent vectors
# `noise` must be smaller than 1.0. The smaller the noise, the more data are correlated
function getData(t, m, n, noise)
    # create m identical data matrices and rotate them by different
    # random orthogonal matrices V_1,...,V_m
    𝐕=[randU(n) for i=1:m] # random orthogonal matrices
    X=randn(n, t)  # data common to all subjects
    # each subject has this common part plus a random part
    𝐗=[𝐕[i]'*((1-noise)*X + noise*randn(n, t)) for i=1:m]
    return 𝐗, 𝐕
end

# As before but output k vectors of data as per the other method
# used for the case m>1, k>1
function getData(t, m, k, n, noise)
    # create m identical data matrices and rotate them by different
    # random orthogonal matrices V_1,...,V_m
    𝐕=[randU(n) for i=1:m] # random orthogonal matrices
    # variables common to all subjects with unique variance profile across k
    X=[(abs2.(randn(n))).*randn(n, t) for s=1:k]
    # each subject has this common part plus a random part
    𝐗=[[𝐕[i]*((1-noise)*X[s] + noise*randn(n, t)) for i=1:m] for s=1:k]
    return 𝐗, 𝐕
end

# Generate data for tests
n, t=3, 10
X=[ 0.0161263  -0.228679  -0.610958;
    0.884571    1.53687    0.569085;
   -1.26195    -1.35456   -0.303478;
   -0.275838    0.775055  -0.461387;
    1.35221    -0.890299  -0.224943;
   -0.204674    0.318228   0.266376;
   -0.203296   -1.83068   -0.869954;
   -0.726922   -0.529861  -0.609168;
   -0.0809636  -0.107088   0.0315276;
    0.177739    0.124636   0.495547]

Y=[ 2.16758   -2.29723   -1.92301;
    3.48593   -0.357984  -0.457085;
    3.02924   -1.5643    -1.38311;
    1.53344   -1.79264   -1.75126;
    0.990636  -0.618866  -0.643236;
    1.37825    0.704042   0.847845;
    0.324096   0.43795    0.142093;
   -3.68015    0.539658   0.436929;
    0.141023  -0.303311  -0.384726;
   -0.903474   1.20838    0.944415]

Cx=Symmetric((X'*X)/t)
Cy=Symmetric((Y'*Y)/t)
Cxy=(X'*Y)/t

𝐗 = [genDataMatrix(t, n) for i = 1:5]
𝐘 = [genDataMatrix(t, n) for i = 1:5]
𝐂x=Vector{Hermitian}([Hermitian((X'*X)/t) for X in 𝐗])
𝐂y=Vector{Hermitian}([Hermitian((Y'*Y)/t) for Y in 𝐘])
𝐂xy=Vector{Matrix}([(X'*Y)/size(X, 1) for (X, Y) in zip(𝐗, 𝐘)])

X_=[X, Matrix(X')]

############  Begin Test  ###################################

@testset "PCA" begin
    ## method (1)
    pC=pca(Cx; simple=true)
    @test Cx≈pC.F*pC.D*pC.iF

    ## method (2)
    for d=1:2
        p=pca(X_[d]; dims=d)
        compare(d, X_[d], p.F, p.D)
        p=pca(X_[d]; dims=d, eVar=3)
        compare(d, X_[d], p.F, p.D)
        p=pca(X_[d]; dims=d, eVar=0.5)
        compare(d, X_[d], p.F, p.D)
        p=pca(X_[d]; dims=d, covEst=LShrLW, eVar=0.5)
        # p=pca(X_[d]; covEst=NShrLW, eVar=0.5)
    end

    p=pca(X; simple=true)
    @test p==pC

    ## method (3)
    p=pca(𝐗; covEst=SCM, eVar=0.5, metric=logEuclidean)
    G=PosDefManifold.mean(logEuclidean, 𝐂x)
    D=p.F'*G*p.F
    @test isApproxD(D)
    @test Diagonal(D)≈Diagonal(reverse(eigvals(G))[1: size(D, 1)])
end

@testset "Whitening" begin
    ## method (1)
    wC=whitening(Cx; simple=true)
    @test wC.F'*Cx*wC.F≈I

    ## method (2)
    for d=1:2
        w=whitening(X_[d]; dims=d)
        compare(d, X_[d], w.F, Diagonal(ones(eltype(w.F), size(w.F, 2))))
        w=whitening(X_[d]; dims=d, eVar=3, eVarMeth=searchsortedlast)
        compare(d, X_[d], w.F, Diagonal(ones(eltype(w.F), size(w.F, 2))))
        w=whitening(X_[d]; dims=d, eVar=0.5, eVarMeth=searchsortedlast)
        compare(d, X_[d], w.F, Diagonal(ones(eltype(w.F), size(w.F, 2))))
        w=whitening(X_[d]; covEst=LShrLW, eVar=0.5)
        # w=whitening(X_[d]; covEst=NShrLW, eVar=0.5)
    end
    w=whitening(X; simple=true)
    @test w==wC

    ## method (3)
    w=whitening(𝐗; covEst=SCM, simple=true, metric=logEuclidean)
    G=PosDefManifold.mean(logEuclidean, 𝐂x)
    @test w.F'*G*w.F≈I
end

@testset "MCA" begin
    ## method (1)
    mC=mca(Cxy, simple=true)
    @test Cxy≈mC.F[1]*mC.D*mC.F[2]'

    ## method (2)
    m=mca(X, Y, simple=true)
    @test isApproxD(m.F[1]'*Cxy*m.F[2])
    @test m==mC

    ## method (3)
    m=mca(𝐗, 𝐘; eVar=0.99)
    m𝐂=mca(mean(𝐂xy); eVar=0.99)
    @test m==m𝐂
end

@testset "CCA" begin
    c=cca(X, Y)
    @test c.F[1]'*Cx*c.F[1]≈I
    @test c.F[2]'*Cy*c.F[2]≈I
    @test isApproxD(c.F[1]'*Cxy*c.F[2])

    # CCA of two vectors of data matrices
    c=cca(𝐗, 𝐘; metric=Euclidean, eVar=0.99)
    c𝐂=cca(Symmetric(mean(𝐂x)), Symmetric(mean(𝐂y)), mean(𝐂xy); eVar=0.99)
    @test c==c𝐂
end

@testset "CSP" begin
    t, n=20, 10
    X=genDataMatrix(t, n)
    Y=genDataMatrix(t, n)
    Cx=Symmetric((X'*X)/t)
    Cy=Symmetric((Y'*Y)/t)

    c=csp(X, Y)
    c=csp(Cx, Cy)
    @test isApproxD(c.F'*Cx*c.F)
    @test isApproxD(c.F'*Cy*c.F)
    @test c.F'*(Cx+Cy)*c.F≈I

    d=csp(Cx, Cy; eVarC=0.)
    @test isApproxD(d.F'*Cx*d.F)
    @test isApproxD(d.F'*Cy*d.F)
    @test d.F'*(Cx+Cy)*d.F≈I

    # the two solutions are identical out of a possible sign ambiguity on the evectors
    @test abs.(c.iF*d.F)≈I
end

@testset "CSTP" begin
    t, n=10, 20
    𝐗 = [genDataMatrix(t, n) for i = 1:10]
    Xfixed=randn(t, n)./1
    for i=1:length(𝐗) 𝐗[i]+=Xfixed end
    C1=Hermitian( mean((X'*X)/t for X∈𝐗) )
    C2=Hermitian( mean((X*X')/n for X∈𝐗) )
    c=cstp(Xfixed, C1, C2; simple=true)
    @test c.F[1]'*C2*c.F[1]≈I
    @test c.F[2]'*C1*c.F[2]≈I
    Z=c.F[1]'*Xfixed*c.F[2]
    n=minimum(size(Z))
    @test isApproxD(Z[1:n, 1:n])
end

@testset "AJD" begin
    t, n, k=50, 10, 50
    A=randn(n, n) # mixing matrix
    𝐗 = [genDataMatrix(t, n) for i = 1:k]
    Xfixed=randn(t, n)./1
    for i=1:length(𝐗) 𝐗[i]+=Xfixed end
    𝐂 = ℍVector([ℍ((𝐗[s]'*𝐗[s])/t) for s=1:k])
    aX=ajd(𝐗; simple=true)
    aC=ajd(𝐂; simple=true)
    @test aX≈aC

    # # # Test orthogonal AJD algorithms
    # create 20 random commuting matrices
    # they all have the same eigenvectors
    𝐂=randP(3, 20; eigvalsSNR=Inf, commuting=true)
    # estimate the approximate joint diagonalizer (ajd) using orthogonal solvers
    # the ajd must be equivalent to the eigenvector matrix of any of the matrices in 𝐂
    a=ajd(𝐂; algorithm=:OJoB)
    @test norm([spForm(a.F'*eigvecs(C)) for C ∈ 𝐂])/20<√err
    a=ajd(𝐂; algorithm=:JADE)
    @test norm([spForm(a.F'*eigvecs(C)) for C ∈ 𝐂])/20<√err
    a=ajd(𝐂; algorithm=:JADEmax)
    @test norm([spForm(a.F'*eigvecs(C)) for C ∈ 𝐂])/20<√err

    # # # Test non-orthogonal AJD algorithms
    # generate positive definite matrices with model A*D_κ*D, where A is an
    # invertible mixing matrix and D_κ, for all κ=1:k, are diagonal matrices.
    # The estimated AJD matrix must be the inverse of A
    # and all transformed matrices bust be diagonal
    n, k=3, 10
    Dest=PosDefManifold.randΛ(eigvalsSNR=Inf, n, k)
    # make the problem identifiable
    for i=1:k Dest[k][1, 1]*=i/(k/2) end
    for i=1:k Dest[k][3, 3]/=i/(k/2) end
    A=randn(n, n) # non-singular mixing matrix
    Cset3=Vector{Hermitian}([Hermitian(A*D*A') for D ∈ Dest])
    a=ajd(Cset3; algorithm=:NoJoB, eVarC=n)
    @test spForm(a.F'*A)<√err
    @test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
    a=ajd(Cset3; algorithm=:LogLike, eVarC=n)
    @test spForm(a.F'*A)<√err
    @test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
    a=ajd(Cset3; algorithm=:LogLikeR, eVarC=n)
    @test spForm(a.F'*A)<√err
    @test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
    a=ajd(Cset3; algorithm=:GAJD, eVarC=n)
    @test spForm(a.F'*A)<√err*10 # GAJD has problems sometimes
    @test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
    a=ajd(Cset3; algorithm=:QNLogLike, eVarC=n)
    @test spForm(a.F'*A)<√err
    @test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
end

@testset "gMCA" begin
    t, m, n, noise = 200, 2, 6, 0.1
    𝐗, 𝐕 = getData(t, m, n, noise)
    Cx=(𝐗[1]*𝐗[1]')/t
    Cy=(𝐗[2]*𝐗[2]')/t
    Cxy=(𝐗[1]*𝐗[2]')/t

    # check that for the case m=2 GMCA gives the same result as MCA
    gm=gmca(𝐗; simple=true)

    # do MCA
    m=mca(Cxy; simple=true)

    @test (m.F[1]'*Cxy*m.F[2]) ≈ (gm.F[1]'*Cxy*gm.F[2])
    # the following must be the identity matrix out of a possible sign ambiguity
    @test abs.(m.F[1]'*gm.F[1]) ≈ I
    @test abs.(m.F[2]'*gm.F[2]) ≈ I
end

@testset "gCCA" begin
    t, m, n, noise = 200, 2, 6, 0.1
    𝐗, 𝐕 = getData(t, m, n, noise)
    Cx=(𝐗[1]*𝐗[1]')/t
    Cy=(𝐗[2]*𝐗[2]')/t
    Cxy=(𝐗[1]*𝐗[2]')/t

    # check that for the case m=2 GCCA gives the same result as CCA
    gc=gcca(𝐗; simple=true)
    gm=gmca(𝐗; simple=true)
    # do CCA
    c=cca(𝐗[1], 𝐗[2]; simple=true)

    @test isApproxD(gc.F[1]'*Cxy*gc.F[2])

    @test (c.F[1]'*Cxy*c.F[2]) ≈ (gc.F[1]'*Cxy*gc.F[2])
    @test (gc.F[1]'*Cx*gc.F[1]) ≈ I
    @test (gc.F[2]'*Cy*gc.F[2]) ≈ I

    # the following must be the identity matrix out of a possible sign ambiguity
    @test abs.(c.iF[1]*gc.F[1]) ≈ I
    @test abs.(c.iF[2]*gc.F[2]) ≈ I

    # test that gCCA is equivalent to GMCA on pre-whitened data
    w1=whitening(𝐗[1])
    w2=whitening(𝐗[2])
    Y=[w1.F'*𝐗[1], w2.F'*𝐗[2]]
    C1=(Y[1]*Y[1]')/t
    C2=(Y[2]*Y[2]')/t
    C12=(Y[1]*Y[2]')/t
    gm2=gmca(Y; simple=true, algorithm=:OJoB)

    gm2=gcca(Y; simple=true, algorithm=:OJoB)
    @test gm2.F[1]'*C1*gm2.F[1] ≈ I
    @test gm2.F[2]'*C2*gm2.F[2] ≈ I
    @test isApproxD(gm2.F[1]'*C12*gm2.F[2])
    @test gm2.F[1]'*w1.F'*Cx*w1.F*gm2.F[1] ≈ I
    @test gm2.F[2]'*w2.F'Cy*w2.F*gm2.F[2] ≈ I
    @test isApproxD(gm2.F[1]'*w1.F'*Cxy*w2.F*gm2.F[2])
end

@testset "mAJD" begin
    # do joint blind source separation of non-stationary data
    t, m, n, k, noise = 200, 10, 3, 20, 0.01
    𝐗, 𝐕=getData(t, m, k, n, noise)
    aX=majd(𝐗; fullModel=true, algorithm=:OJoB)
    # the spForm index of the estimated demixing matrices times the true
    # mixing matrix must be low
    @test mean(spForm(aX.F[i]'*𝐕[i]) for i=1:m)<0.05
end


using OrdinaryDiffEq, Test, LinearAlgebra

T = 1000.0; Ttr = 0.0;
d0=1e-9; threshold=10^4*d0;
dt = 0.1; diff_eq_kwargs = Dict()

@inbounds function eom_lorenz!(du, u, p, t)
  σ = 10.0; ρ = 28.0; β = 8/3
  du[1] = σ*(u[2]-u[1])
  du[2] = u[1]*(ρ-u[3]) - u[2]
  du[3] = u[1]*u[2] - β*u[3]
end

prob = OrdinaryDiffEq.ODEProblem(eom_lorenz!, [0., 10.0, 0], (zero(T), T))
integ1 = init(prob, Tsit5(); diff_eq_kwargs...)
@. prob.u0 = prob.u0 + d0/sqrt(3)
integ2 = init(prob, Tsit5(); diff_eq_kwargs...)
integ1.opts.advance_to_tstop=true
integ2.opts.advance_to_tstop=true

dist = d0
λ = zero(eltype(integ1.u))
i = 0
tvector = dt:dt:T

for τ in tvector
  # evolve until rescaling:
  push!(integ1.opts.tstops, τ); step!(integ1)
  push!(integ2.opts.tstops, τ); step!(integ2)
  dist = norm(integ1.u .- integ2.u)
  # Rescale:
  if dist ≥ threshold
    # Rescale and reset everything:
    integ2.u .= integ1.u #.+ (integ2.u .- integ1.u)./a
    u_modified!(integ2, true)
    set_proposed_dt!(integ2,integ1)
    break
  end
end

τ = tvector[end]
push!(integ1.opts.tstops, τ); step!(integ1)
push!(integ2.opts.tstops, τ); step!(integ2)
@test dist = norm(integ1.u .- integ2.u) == 0


const beale = let res_init=zeros(3), jac_init=zeros(3,2), x_init=[0.1, 0.1]

    function res(x, r)
        r[1] = 1.5 - x[1] * (1.0 - x[2])
        r[2] = 2.25 - x[1] * (1.0 - x[2]^2)
        r[3] = 2.625 - x[1] * (1.0 -  x[2]^3)
        return r
    end

    function jac(x, jac)
        jac[1, 1] = x[2] - 1.0
        jac[1, 2] = x[1]
        jac[2, 1] = x[2] ^ 2 - 1.0
        jac[2, 2] = 2.0 * x[1] * x[2]
        jac[3, 1] = x[2] ^ 3 - 1.0
        jac[3, 2] = 3.0 * x[1] * (x[2] ^ 2)
        return jac
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("beale", res, jac,
                                                   res_init, jac_init, x_init;
                                                   scale=scale, verbose=verbose,
                                                   print_steps=print_steps)
end


export hanncoeffs, hannsum, hannsum!, hannfilter

"""
    hanncoeffs(tin,tout,L)
# Arguments
- `tin`: input dataset time (vector valued)
- `tout`: time where filtered output will be created
- `L`: filter length
# Output
- `h`: normalized coefficients of Hann(ing) filter
"""
function hanncoeffs(tin,tout,L)
    ctin = collect(tin)
    δ = abs.(ctin .- tout)./(2.0.*L)
    h = 0 .*ctin
    #ih = findall(δ.< 0.5)
    for hh = 1:length(h)
        if δ[hh] < 0.5
            h[hh] = (cos.(π.*δ[hh])).^2
        end
    end
    h ./= sum(h)  # normalize
    return h
end

"""
    hannsum(θin,tin,tout,L)
    use this when a timeseries is available
# Arguments
- `θ`: input value at all times t (e.g., θ(t))
- `tout`: time where filtered output will be created
- `tin`: input dataset time (vector valued)
- `L`: filter length
# Output
- `θout`: output value at time th (e.g., θh(th))
"""
function hannsum(θin,tin,tout,L)
    # function hannsum(θin,tin,tout,L)

    h = hanncoeffs(tin,tout,L)

    # iterate over all filter coeffs
    i = 1
    nh = length(tin)
    θout = 0.0
    for tt = 1:nh
        θout += h[tt]*θin[tt]
    end
    return θout
end

"""
    hannsum!(θout,θin,tin,tout,L)
    use this when a timeseries is available
# Arguments
- `θout`: output value at time th (e.g., θh(th))
- `θin`: input value at all times t (e.g., θ(t))
- `tin`: input dataset time (vector valued)
- `tout`: time where filtered output will be created
- `L`: filter length
"""
function hannsum!(θout,θin,tin,tout,L)
    # function hannsum(θin,tin,tout,L)
    # use this when a timeseries is available
    #
    # θ  = input value at all times t (e.g., θ(t))
    # th = time where filtered output will be created
    # t  = input dataset time (vector valued)
    # L = filter length
    #
    # θh = output value at time th (e.g., θh(th))

    h = hanncoeffs(tin,tout,L)

    # iterate over all filter coeffs
    i = 1
    nh = length(tin)
    #θout = 0.0
    for tt = 1:nh
        θout += h[tt]*θin[tt]
    end
end

"""
    function hannsum!(θout,θin,tin,tout,toutindex,L::Float64)

     Call this version for mesharrays
# Arguments
- `θout::MeshArrays.gcmarray{Float32,2,Array{Float32,2}}`: output value at time tout (e.g., θout(tout))
- `θin::MeshArrays.gcmarray{Float32,2,Array{Float32,2}}`: input value at all times t (e.g., θ(t))
- `tin::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}`: input dataset time (vector valued)
- `tout::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}`: time where filtered output will be created
- `toutindex::Integer`: output times?
- `L::Float64`: filter length
"""
function hannsum!(θout::MeshArrays.gcmarray{Float32,2,Array{Float32,2}},θin::MeshArrays.gcmarray{Float32,2,Array{Float32,2}},tin::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},tout::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},toutindex::Integer,L::Float64)

    h = hanncoeffs(tin,tout[toutindex],L)

    # iterate over all filter coeffs
    i = 1
    nin = length(tin)
    for tt = 1:nin
        if h[tt] > 0
            θout[:,toutindex] += h[tt]*θin[:,tt]
        end
    end
end

"""
    function hannfilter(θin,tin,tout,L)
# Arguments
- `θin::Array{Float32,1}`: input
- `tin`
- `tout`
- `L`
# Output
- `θout`: output
"""
function hannfilter(θin::Array{Float32,1},tin,tout,L)
    # assumes all values available in timeseries
    # initialize θh

    nt = length(tin)
    nh = length(tout)
    θout = zeros(Float32,nh)
    println("calling simple hanningfilter")
    for tt= 1:nh
        hannsum!(θout[tt],θin,tin,tout[tt],L)
    end
    return θout
end

"""
    function hannfilter(θin,tin,tout,L,γ)
# Arguments
- `θin::MeshArrays.gcmarray{Float32,2,Array{Float32,2}}`
- `tin::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}`
- `tout::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}`
- `L::Float64`
- `γ::gcmgrid`
# Output
- `θout`
"""
function hannfilter(θin::MeshArrays.gcmarray{Float32,2,Array{Float32,2}},tin::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},tout::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},L::Float64,γ::gcmgrid)
    # use this for gcmarray
    # initialize θout

    nin = length(tin)
    nout = length(tout)
    #println(nin,nout)
    #θout = zeros(Float32,nh)
    θout = MeshArray(γ,Float32,nout) # some nans here
    tmp1=zeros(Float32,Tuple(γ.ioSize))

    for tt= 1:nout
        println(tt)
        # initialize on the fly.
        θout[:,tt]=γ.read(tmp1,MeshArray(γ,Float32))
        #println("call hanningsum!")
        hannsum!(θout,θin,tin,tout,tt,L)
    end
    return θout
end


using Plots, LaTeXStrings
gr(colorbar=false,size=(600,150),legend=false,grid=false,xticks=false,yticks=false,axis=false)

using FourierGPE

## Initialize simulation
L = (60.0,)
N = (512,)
sim = Sim(L,N)
@unpack_Sim sim

## set simulation parameters
μ = 25.0
g = -0.01
γ = 0.0
Ns = 200
ξs = 2/abs(g)/Ns
us = 20
tf = π |> Float64
Nt = 150
ti = 0.0
t = LinRange(ti,tf,Nt)

## soliton wavefunction
x = X[1]
ψs(x) = sqrt(Ns/2ξs)*sech(x/ξs)*exp(im*us*x)
ψi = ψs.(x)
ϕi = kspace(ψi,sim);

## Set all fields
@pack_Sim! sim

## Evolve in k space
sol = runsim(sim)

## plot

y = abs2.(xspace(sol[Nt-76],sim))
plot(x,y,fill=(0, 0.2))

anim = @animate for i=1:Nt-8
    ψ = xspace(sol[i],sim)
    y = abs2.(ψ)
    plot(x,y,c=c3,fill=(0, 0.4,c3))
end

gif(anim, "./examples/brightsoliton.gif", fps = 25)


using DrWatson
@quickactivate "ExercisesRepo"
include(srcdir("colorscheme.jl"))

# %% The very first code snippet of the course
using DynamicalSystems # load the library

function lorenz_rule(u, p, t)
    σ, ρ, β = p
    x, y, z = u
    dx = σ*(y - x)
    dy = x*(ρ - z) - y
    dz = x*y - β*z
    return SVector(dx, dy, dz) # Static Vector
end

p  = [10.0, 28.0, 8/3] # parameters: σ, ρ, β
u₀ = [0, 10.0, 0]      # initial state
# create an instance of a `DynamicalSystem`
lorenz = ContinuousDynamicalSystem(lorenz_rule, u₀, p)

T  = 100.0 # total time
dt = 0.01  # sampling time
A  = trajectory(lorenz, T; dt)

# %% Animate evolution of trajectories in the Standard map
using InteractiveChaos
using DynamicalSystems
using Makie

# Standard map trajectories
ds = Systems.standardmap(; k = 1.0)
u0s = [[θ, p] for θ ∈ 0:2π for p ∈ 0:2π]
lims = ((0, 2π), (0, 2π))

scene, main, layout, obs = interactive_evolution(
    ds, u0s; tail = 1000, lims,
)
main.xlabel = "θ"
main.ylabel = "p"

# %% Lorenz system trajectories
using OrdinaryDiffEq: Tsit5, Vern9
ds = Systems.lorenz()

u0s =  [[10,20,40.0] .+ rand(3) for _ in 1:7]

diffeq = (alg = Tsit5(), dtmax = 0.01)

scene, main, layout, obs = interactive_evolution(
    ds, u0s; tail = 1000, diffeq, colors = COLORS,
)

# %% Add poincare plane
import AbstractPlotting
using AbstractPlotting.MakieLayout: LAxis
o = AbstractPlotting.Point3f0(-25, 0, 0)
w = AbstractPlotting.Point3f0(50, 0, 50)
p = AbstractPlotting.FRect3D(o, w)

# These lines are necessary for transparent planes
a = AbstractPlotting.RGBAf0(0,0,0,0)
c = AbstractPlotting.RGBAf0(0.2, 0.2, 0.8, 1.0)
img = AbstractPlotting.ImagePattern([c a; a c]);
AbstractPlotting.mesh!(main, p; color = img);

# %% Plot Poincare sos
psosplot = layout[:, 2] = LAxis(scene)
psos = poincaresos(ds, (2, 0.0), 2000.0)
AbstractPlotting.scatter!(psosplot, psos[:, 1], psos[:, 3])

display(scene)

# %% Henon heiles system trajectories
ds = Systems.henonheiles()

u0s = [[0.0, -0.25, 0.42081, 0.0],
[0.0, 0.1, 0.5, 0.0],
[0.0, -0.31596, 0.354461, 0.0591255]]

diffeq = (alg = Vern9(), dtmax = 0.01)
idxs = (1, 2)

scene, main, layout, obs = interactive_evolution(
    ds, u0s; idxs, tail = 2000, diffeq, colors = COLORS,
)
main.scene[AbstractPlotting.Axis][:names, :axisnames] = ("q₁", "q₂", "p₂")


@testset "Normal form f4 finite fields" begin

    R, (x, y, z) = PolynomialRing(GF(2^31 - 1), ["x", "y", "z"])

    # Regression test, Structural Identifiability
    @test Groebner.normalform([x], R(0)) == R(0)
    @test Groebner.normalform([x], R(1)) == R(1)

    Gs = [
        [x], [x, y], [x, y, z]
    ]

    @test Groebner.normalform(Gs[1], x) == 0
    @test Groebner.normalform(Gs[1], y) == y
    @test Groebner.normalform(Gs[3], x + y^2) == 0
    @test Groebner.normalform(Gs[2], 5x^2 - y^3 - 1) == -1
    @test Groebner.normalform(Gs[2], x + y + 3z) == 3z

    G = [
        x^2 + y*x + 1,
        y^2 - y
    ]

    @test Groebner.normalform(G, x + y + z) == x + y + z
    @test Groebner.normalform(G, x^2 + y^2) == y - y*x - 1
    @test Groebner.normalform(G, y^3) == Groebner.normalform(G, y^4) == y

    R, (x, y, z) = PolynomialRing(GF(2^31 - 1), ["x","y","z"], ordering=:degrevlex)
    G = [
        x^2 + y,
        y^2 + x
    ]
    @test Groebner.normalform(G, x^2 + y^2) == -x - y

    @test_nowarn Groebner.normalform([x, x*y], x)
end

@testset "Normal form -- a lot of variables" begin
    R, xs = PolynomialRing(QQ, ["x$i" for i in 1:63])

    GB = Groebner.groebner(xs)
    @test GB == reverse(xs)
    @test Groebner.normalform(GB, xs[1]) == R(0)

    R, xs = PolynomialRing(QQ, ["x$i" for i in 1:220])

    GB = Groebner.groebner(xs)
    @test GB == reverse(xs)
    @test Groebner.normalform(GB, xs[1]) == R(0)

    #=
    R, xs = PolynomialRing(QQ, ["x$i" for i in 1:1049])

    GB = Groebner.groebner(xs)
    @test GB == reverse(xs)
    @test Groebner.normalform(GB, xs[1]) == R(0)
    =#

end


#=
  Copyright (c) 2021
  https://github.com/Ismael-VC/Varvara.jl/blob/main/CONTRIBUTORS.md

  Permission to use, copy, modify, and distribute this software for any
  purpose with or without fee is hereby granted, provided that the above
  copyright notice and this permission notice appear in all copies.

  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
  WITH REGARD TO THIS SOFTWARE.
=#

module Tables

using Printf: @printf

uint8(x::Float64) = UInt8(trunc(x))
uint8(x::UInt8) = x

function main()::Int
  println("60 points on a circle128(bytex,bytey):\n")

  for i ∈ 0:59
    cx = cy = r = 128
    pos = (i - 15) % 60
    deg = (pos / 60.0) * 360.0
    rad = deg * (π / 180)
    x = cx + r * cos(rad)
    y = cy + r * sin(rad)

    if i > 0 && i % 8 == 0
      println()
    end

    @printf("%02x%02x ", uint8(clamp(x, 0x00, 0xff)), uint8(clamp(y, 0x00, 0xff)))
  end

  println("\n")
  return 0
end

end  # module

if abspath(PROGRAM_FILE) == @__FILE__
  using .Tables: main

  main()
end


using ApproxFun, MultivariateOrthogonalPolynomials, BlockArrays, SpecialFunctions, FillArrays, Plots
import ApproxFun: blockbandwidths, Vec, PiecewiseSegment
import MultivariateOrthogonalPolynomials: DirichletTriangle

#######
# Example 1 Poisson
#######

S = TriangleWeight(1,1,1,JacobiTriangle(1,1,1))
Δ = Laplacian(S)

N = 500
M = sparse(Δ[Block.(1:N), Block.(1:N)])
F = lu(M)
u_1 = F \ pad(coefficients(Fun(1, rangespace(Δ))), size(M,1))
u_2 = F \ pad(coefficients(Fun((x,y) -> x^2 + y^2, rangespace(Δ))), size(M,1))
u_3 = F \ pad(coefficients(Fun((x,y) -> x*y*(1-x-y), rangespace(Δ))), size(M,1))
u_4 = F \ pad(coefficients(Fun((x,y) -> x*y*(1-x-y) *cos(300x*y), rangespace(Δ))), size(M,1))
u_5 = F \ pad(coefficients(Fun((x,y) -> exp(-25((x-0.2)^2+(y-0.2)^2)), rangespace(Δ))), size(M,1))

gr()

plot(abs.(u_1)[1:10_000]; yscale=:log10, xscale=:log10, label="1")
plot!(abs.(u_2)[1:10_000]; yscale=:log10, xscale=:log10, label="x^2 + y^2")
plot!(abs.(u_3)[1:10_000]; yscale=:log10, xscale=:log10, label="xy(1-x-y)")
plot!(abs.(u_4)[1:10_000]; yscale=:log10, xscale=:log10, label="xy(1-x-y)cos(300xy)")

#######
# Example 2 Variable coefficient Helmholts
######

S = TriangleWeight(1,1,1,JacobiTriangle(1,1,1))
x,y = Fun(Triangle())

V = 1-(3*(x-1)^4 + 5*y^4)

f = Fun((x,y) -> x*y*exp(x) , JacobiTriangle(1.,1.,1.))

k = 100
L = Laplacian(S) + (k^2 * V)
N = ceil(Int,2k)
@time M̃ = L[Block.(1:N), Block.(1:N)]
@time M = sparse(M̃)
@time F = lu(M)
@time u_c = F \ pad(coefficients(f), size(M,1))
u = Fun(domainspace(L), u_c)

#######
# Example 3 Bi-harmonic
#######

S = TriangleWeight(2,2,2, JacobiTriangle(2,2,2))
Δ = Laplacian(S)
L = Δ^2

N = 1001
M = sparse(L[Block.(1:N), Block.(1:N)])
F = lu(M)

u_1 = F \ pad(coefficients(Fun(1, rangespace(L))), size(M,1))
u_2 = F \ pad(coefficients(Fun((x,y) -> x^2 + y^2, rangespace(L))), size(M,1))
u_3 = F \ pad(coefficients(Fun((x,y) -> x*y*(1-x-y), rangespace(L))), size(M,1))
# u_4 = F \ pad(coefficients(Fun((x,y) -> x*y*(1-x-y) *cos(300x*y), rangespace(L))), size(M,1))
u_4 = F \ pad(coefficients(Fun((x,y) -> exp(-25((x-0.2)^2+(y-0.2)^2)), rangespace(L))), size(M,1))

bnrms = u -> [norm(PseudoBlockArray(u, 1:N)[Block(K)],Inf) for K=1:N]

plot(bnrms(u_1); xscale=:log10, yscale=:log10, linewidth=3, label="1")
plot!(bnrms(u_2); xscale=:log10, yscale=:log10, linewidth=3, linestyle=:dot, label="x^2 + y^2")
plot!(bnrms(u_3); xscale=:log10, yscale=:log10, linewidth=3, linestyle=:dot, label="x^2 + y^2")
plot!(bnrms(u_4); xscale=:log10, yscale=:log10, linewidth=3, linestyle=:dot, label="x^2 + y^2")

plot(abs.(coefficients(Fun((x,y) -> exp(-25((x-0.2)^2+(y-0.2)^2)), rangespace(L)))); scale=:log10, yscale=:log10)

#######
# Example 4 Laplace's equation
#######

S = DirichletTriangle{1,1,1}()
B = Dirichlet() : S
Δ = Laplacian() : S
# L = [B; Δ]

fs = ((x,y) -> exp(x)*cos(y),
      (x,y) -> x^2,
      (x,y) ->  x^3*(1-x)^3*(1-y)^3,
      (x,y) -> x^3*(1-x)^3*(1-y)^3 * cos(100x*y))

N = 1000;
@time BM = sparse(B[Block.(1:N), Block.(1:N)]);
@time ΔM = sparse(Δ[Block.(1:N), Block.(1:N)]);
@time M = [BM; ΔM];
M = [sparse(I, size(M,1), 2) M] # add tau term
@time F = factorize(M);

# RHS
k = 1
∂u = Fun(fs[k], rangespace(B))
r = pad(coefficients(∂u, rangespace(B)), size(Fs[N],1))
@time u_cfs = F \ r

#######
# Example 5 Transport equation
#######

S = DirichletTriangle{0,1,0}()
∂S = Legendre(Segment(Vec(0.0,0), Vec(1.0,0)))
B = I : S → ∂S

Dx = Derivative([1,0]) : S
Dt = Derivative([0,1]) : S

x,_ = Fun(∂S)
u₀ = x*(1-x)*exp(x)

# u_t = u_x
L = [B; Dx - Dt]
N = 20; M = sparse(L[Block.(1:N), Block.(1:N)]);
u_c = M \ pad(coefficients([u₀; 0], rangespace(L)), size(M,1))
u = Fun(domainspace(L), u_c)

# 0.5u_t = u_x
S = DirichletTriangle{0,1,1}()
∂S = Legendre(Segment(Vec(0.0,0), Vec(1.0,0))) , Legendre(Segment(Vec(0.0,1.0),Vec(1.0,0)))
B_x = I : S → ∂S[1]
B_z = I : S → ∂S[2]
Dx = Derivative([1,0]) : S
Dt = Derivative([0,1]) : S

L = [B_x; B_z; Dx - 0.5Dt]
N = 100; M = sparse(L[Block.(1:N), Block.(1:N)])
u₀_x = Fun((x,_) -> x*exp(x-1), rangespace(B_x))
u₀_z = Fun((x,y) -> (1-y), rangespace(B_z))
u_c = M \ pad(coefficients([u₀_x; u₀_z; 0], rangespace(L)), size(M,1))
u = Fun(domainspace(L), u_c)

# u_t + u_x = 0
S = DirichletTriangle{1,1,0}()
∂S = Legendre(Segment(Vec(0.0,0), Vec(1.0,0))) , Legendre(Segment(Vec(0.0,0.0),Vec(0.0,1.0)))
B_x = I : S → ∂S[1]
B_y = I : S → ∂S[2]
Dx = Derivative([1,0]) : S
Dt = Derivative([0,1]) : S

L = [B_x; B_y; Dt + Dx]
N = 100; M = sparse(L[Block.(1:N), Block.(1:N)])
u₀_x = Fun((x,_) -> (1-x)*exp(x), rangespace(B_x))
u₀_y = Fun((x,y) -> (1-y), rangespace(B_y))

u_c = M \ pad(coefficients([u₀_x; u₀_y; 0], rangespace(L)), size(M,1))
u = Fun(domainspace(L), u_c)

#######
# Example 7 Helmholtz in a polygon
#######
d = Triangle() ∪ Triangle(Vec(1,1),Vec(1,0),Vec(0,1)) ∪
                 Triangle(Vec(1,1),Vec(0,1),Vec(0,2)) ∪
                 Triangle(Vec(0,0),Vec(0,1),Vec(-1,1.5))

f = Fun((x,y) -> cos(x*sin(y)), d)

∂d = PiecewiseSegment([Vec(0.,0), Vec(1.,0), Vec(1,1), Vec(0,2), Vec(0,1), Vec(-1,1.5), Vec(0,0)])

S = DirichletTriangle{1,1,1}.(components(d))
∂S = Legendre.(components(∂d))

∂S12 = Legendre(Segment(Vec(1.,0),Vec(0,1)))
∂S23 = Legendre(Segment(Vec(1.,1),Vec(0,1)))
∂S14 = Legendre(Segment(Vec(0.,0),Vec(0,1)))

Dx = Derivative([1,0])
Dy = Derivative([0,1])

k = 20

L = [(I : S[1] → ∂S[1]) 0                 0                     0                  0                0                       0                    0                  0                   0                    0                  0;
     0                  0                 0                    (I : S[2] → ∂S[2])  0                0                       0                    0                  0                    0                   0                  0;
     0                  0                 0                     0                  0                0                       (I : S[3] → ∂S[3])   0                  0                    0                   0                  0;
     0                  0                 0                     0                  0                0                       (I : S[3] → ∂S[4])   0                  0                    0                   0                  0;
     0                  0                 0                     0                  0                0                       0                    0                  0                   (I : S[4] → ∂S[5])   0                  0;
     0                  0                 0                     0                  0                0                       0                    0                  0                   (I : S[4] → ∂S[6])   0                  0;
     (I : S[1] → ∂S12)  0                 0                   -(I : S[2] → ∂S12)   0                0                       0                    0                  0                    0                   0                  0;
     0                  (I : S[1] → ∂S12) (I : S[1] → ∂S12)     0                -(I : S[2] → ∂S12) -(I : S[2] → ∂S12)      0                    0                  0                    0                   0                  0;
     0                  0                 0                    (I : S[2] → ∂S23)   0                0                      -(I : S[3] → ∂S23)    0                  0                    0                   0                  0;
     0                  0                 0                     0                  0                (I : S[2] → ∂S23)       0                    0                  -(I : S[3] → ∂S23)   0                   0                  0;
     (I : S[1] → ∂S14)  0                 0                     0                  0                0                       0                    0                  0                   -(I : S[4] → ∂S14)   0                  0;
     0                  (I : S[1] → ∂S14) 0                     0                  0                0                       0                    0                  0                    0                  -(I : S[4] → ∂S14)   0;
     Dx                -I                 0                     0                  0                0                       0                    0                  0                    0                   0                  0;
     Dy                 0                -I                     0                  0                0                       0                    0                  0                    0                   0                  0;
     0                  0                 0                     Dx                -I                0                       0                    0                  0                    0                   0                  0;
     0                  0                 0                     Dy                 0               -I                       0                    0                  0                    0                   0                  0;
     0                  0                 0                     0                  0                0                       (Dx : S[3])         -(I : S[3])         0                    0                   0                  0;
     0                  0                 0                     0                  0                0                       Dy                   0                 -I                    0                   0                  0;
     0                  0                 0                     0                  0                0                       0                    0                  0                    Dx                 -I                  0;
     0                  0                 0                     0                  0                0                       0                    0                  0                    (Dy : S[4])         0                 -(I : S[4])
    k^2*I              Dx                Dy                    0                  0                0                       0                    0                  0                    0                   0                  0;
     0                  0                 0                     k^2*I                  Dx               Dy                      0                    0                  0                    0                   0                  0;
     0                  0                 0                     0                  0                0                       k^2*I                    Dx                 Dy                   0                   0                  0;
     0                  0                 0                     0                  0                0                       0                    0                  0                    k^2*I                  Dx                 Dy
     ]

f = ones(∂d)

N = 101
@time M = sparse(L[Block.(1:N), Block.(1:N)]);
@time F = qr(M);
@time u_c = F \ pad(coefficients(vcat(components(f)..., Zeros(18)...), rangespace(L)), size(M,1));
u = Fun(domainspace(L), u_c)


""" Commutative diagrams in a category.

**NOTE**: This module is at an early stage of construction. At this time, only
commutative squares are implemented.
"""
module CommutativeDiagrams
export SquareDiagram, ob, hom, left, right, top, bottom

using AutoHashEquals

using ...GAT
using ...Theories: DoubleCategory
import ...Theories: ob, hom, dom, codom, compose, ⋅, ⋆, HomV, HomH,
  composeH, composeV, id, idH, idV, id2, id2H, id2V, left, right, top, bottom
using ..FinSets

# Commutative squares
#####################

"""    SquareDiagram(top, bottom, left, right)

Creates a square diagram in a category, which forms the 2-cells of the double
category Sq(C). The four 1-cells are given in top, bottom, left, right order, to
match the GAT of a double category.
"""
@auto_hash_equals struct SquareDiagram{Ob,Hom,Homs<:AbstractVector{Hom}}
  corners::Vector{Ob}
  sides::Homs
end

function SquareDiagram(homs::AbstractVector)
  length(homs) == 4 ||
    error("Square diagrams accept exactly 4 homs, in order top, bottom, left, right")
  obs = [dom(homs[1]), dom(homs[2]), dom(homs[4]), codom(homs[4])]

  # checking well-formedness
  # top-left share domains
  @assert dom(homs[1]) == dom(homs[3])
  # bottom-right share codomains
  @assert codom(homs[2]) == codom(homs[4])
  # left-bottom intersection
  @assert codom(homs[3]) == dom(homs[2])
  # top-right intersection
  @assert codom(homs[1]) == dom(homs[4])

  SquareDiagram(obs, homs)
end

SquareDiagram(top, bottom, left, right) =
  SquareDiagram([top, bottom, left, right])

ob(sq::SquareDiagram) = sq.corners
hom(sq::SquareDiagram) = sq.sides

top(sq::SquareDiagram) = sq.sides[1]
bottom(sq::SquareDiagram) = sq.sides[2]
left(sq::SquareDiagram) = sq.sides[3]
right(sq::SquareDiagram) = sq.sides[4]

""" Arrow category of FinSet as a double category of commutative squares.

For any category C, you can construct Sq(C) as a
[double category](https://ncatlab.org/nlab/show/double+category)
of commutative squares.

TODO: This construction has nothing specifically to do with finite sets and
functions and should be generalized to any category C.
"""
@instance DoubleCategory{FinSet, FinFunction, FinFunction, SquareDiagram} begin
  @import dom, codom, top, bottom, left, right, ⋅
  idH(A::FinSet) = FinFunction(identity, A, A)
  idV(A::FinSet) = FinFunction(identity, A, A)

  composeH(f::FinFunction, g::FinFunction) = compose(f,g)
  composeV(f::FinFunction, g::FinFunction) = compose(f,g)

  id2(A::FinSet) = SquareDiagram(idH(A), idH(A), idV(A), idV(A))
  # FIXME: how do you distinguish between vertical and horizontal if they are the same type?
  id2V(f::FinFunction) = SquareDiagram(f, f, idV(A), idV(A))
  id2H(f::FinFunction) = SquareDiagram(idH(A), idH(A), f, f)

  """    composeH(s₁, s₂)

  compose two squares horizontally as shown below:
      1   -f->   3  -g->   5
      |          |         |
      |          |         |
      v          v         v
      2  -f'->   4  -g'->  6
   """
  function composeH(s₁::SquareDiagram, s₂::SquareDiagram)
    @assert ob(s₁)[3] == ob(s₂)[1]
    @assert ob(s₁)[4] == ob(s₂)[2]
    @assert right(s₁) == left(s₂)

    f = top(s₁)
    f′= bottom(s₁)
    g = top(s₂)
    g′= bottom(s₂)
    return SquareDiagram(f⋅g, f′⋅g′, left(s₁), right(s₂))
  end

  """    composeV(s₁, s₂)

  compose two squares vertically as shown below:
      1   -->  3
      |        |
      |        |
      v        v
      2  -->   4
      |        |
      |        |
      v        v
      5  -->   6
  """
  function composeV(s₁::SquareDiagram, s₂::SquareDiagram)
    @assert ob(s₁)[2] == ob(s₂)[1]
    @assert ob(s₁)[4] == ob(s₂)[3]
    @assert bottom(s₁) == top(s₂)

    f = left(s₁)
    f′= right(s₁)
    g = left(s₂)
    g′= right(s₂)
    return SquareDiagram(top(s₁), bottom(s₂), f⋅g, f′⋅g′)
  end
end

end


import Base.@kwdef

abstract type AbstractSpline end

@kwdef struct Spline
    points::Vector{Point}
    duration::Real = 10000.
    sharpness::Real = 0.85
    centers = []
    controls = [[]]
    stepLength::Real = 60
    length::Real = length(points)
    delay::Real = 0
    steps = []

    function Spline(; points=points::Vector{Point}, duration=duration::Real=10000., sharpness=sharpness::Real=0.85,
        centers=centers=[],
        controls=controls=[[]],
        stepLength=stepLength::Real=60,
        length=length::Real = length(points),
        delay=delay::Real= 0,
        steps=steps= [])

        for point in points
            point.coordinates[3] = length(point) > 2 ? point.coordinates[3] : 0
        end

        for i in eachindex(points)
            p1 = points[i]
            p2 = points[i + 1]

            push!(centers,
                Point([(p1.coordinates[1] + p2.coordinates[1]) / 2,
                (p1.coordinates[2] + p2.coordinates[2]) / 2,
                (p1.coordinates[3] + p2.coordinates[3]) / 2]))
        end
        push!(controls, [points[1], points[1]])

        for j in eachindex(centers)
            c1 = centers[j]
            c2 = centers[j + 1]
            dx = points[j + 1].coordinates[1] - (c1.coordinates[1] + c2.coordinates[1]) / 2
            dy = points[j + 1].coordinates[2] - (c1.coordinates[2] + c2.coordinates[2]) / 2
            dz = points[j + 1].coordinates[3] - (c1.coordinates[3] + c2.coordinates[3]) / 2

            push!(controls, [
                Point([(1. - sharpness) * points[j + 1].coordinates[1] + sharpness * c1.coordinates[1] + dx,
                (1. - sharpness) * points[j + 1].coordinates[2] + sharpness * c1.coordinates[2] + dy,
                (1. - sharpness) * points[j + 1].coordinates[3] + sharpness * c1.coordinates[3] + dz]),
                Point([(1. - sharpness) * points[j + 1].coordinates[1] + sharpness * c2.coordinates[1] + dx,
                (1. - sharpness) * points[j + 1].coordinates[2] + sharpness * c2.coordinates[2] + dy,
                (1. - sharpness) * points[j + 1].coordinates[3] + sharpness * c2.coordinates[3] + dz])
                ])
        end

        push!(controls, [points[length - 1], points[length - 1]])
        steps = cacheSteps(stepLength)

        return new(points, duration, sharpness, centers, controls, stepLength, length, delay, steps)
    end

end # struct

function cacheSteps(minDist::Real)
    steps = Real[]
    lastStep = pos(0)
    push!(steps, 0)

    for i in 1:10:duration
        step = pos(i)
        dist = sqrt(
            (step.coordinates[1] - lastStep.coordinates[1])^2 +
            (step.coordinates[2] - lastStep.coordinates[2])^2 +
            (step.coordinates[3] - lastStep.coordinates[3])^2)

        if dist > minDist
            push!(steps, i)
            lastStep = step
        end
    end

    return steps

end

function pos(time::Real)
    t = time - delay
    t < 0 && (t = 0)
    t > duration && (t = duration - 1)

    t2 = t / duration
    t2 >= 1 && return points[length - 1]

    n = floor((length(points) - 1) * t2)
    t1 = (length - 1) * t2 - n

    return bezier(t1, points[n], controls[n][1], controls[n + 1][0], points[n + 1])
end

function bezier(; t::Real, p1::Point, c1::Point, c2::Point, p2::Point)
    b = [t^3, (3 * t^2 * (1-t)), (3 * t * (1 - t) * (1 - t)), ((1 - t) * (1 - t) * (1 - t))]

    return Point([
        p2.coordinates[1] * b[1] + c2.coordinates[1] * b[2] + c1.coordinates[1] * b[3] + p1.coordinates[1] * b[4],
        p2.coordinates[2] * b[1] + c2.coordinates[2] * b[2] + c1.coordinates[2] * b[3] + p1.coordinates[2] * b[4],
        p2.coordinates[3] * b[1] + c2.coordinates[3] * b[2] + c1.coordinates[3] * b[3] + p1.coordinates[3] * b[4]])
end


module G

using ValidatedNumerics;
include("SubDiff.jl")
include("sti.jl")

function equation(solver)
    sti.STI(solver.system.a*solver.previous.intervals) - solver.system.b.sti
end

function initialConditions(system)
    aMatrixMid = map(x -> mid(x), system.a)
    leftMatrix = sti.constituentMatrix(aMatrixMid)
    rightVector = system.b.sti
    initialInNumbers = leftMatrix \ rightVector
    initialInNumbers
end

function subDifferential(solver)
    identityMatrix = eye(solver.system.size)
    loSubDiff = map(
        index -> calulateLoSubdifferentialRow(slicedim(solver.system.a, 1, index), solver.previous.sti, slicedim(identityMatrix, 1, index)),
        range(1, solver.system.size)
    )
    hiSubDiff = map(
        index -> calulateHiSubdifferentialRow(slicedim(solver.system.a, 1, index), solver.previous.sti, slicedim(identityMatrix, 1, index)),
        range(1, solver.system.size)
    )
    vcat(loSubDiff..., hiSubDiff...)
end

function calulateLoSubdifferentialRow{T}(aVector :: Array{Interval{T}}, xVector :: Array{T}, _ :: Array{T})
    calulateSubdifferentialRowPart(SubDiff.productLoSubDifferential, -1, aVector, xVector, _)
end

function calulateHiSubdifferentialRow{T}(aVector :: Array{Interval{T}}, xVector :: Array{T}, _ :: Array{T})
    calulateSubdifferentialRowPart(SubDiff.productHiSubDifferential, 1, aVector, xVector, _)
end

function calulateSubdifferentialRowPart{T}(partFunction, sumFactor, aVector :: Array{Interval{T}}, xVector :: Array{T}, _ :: Array{T})
    systemSize = size(xVector, 1)
    systemSizeHalf = Int(systemSize/2)
    mapped =  map(
        index -> partFunction(aVector[index], xVector[index], xVector[index + systemSizeHalf]),
        range(1, systemSizeHalf)
    )
    calculatedSum = sumFactor * mapped
    stiedCalculatedSum = sti.STI(calculatedSum)
    transpose(stiedCalculatedSum)
end
end


struct RealCoding
    nbvar::Int
    lb::Vector{Float64}
    ub::Vector{Float64}
    nbbits::Vector{Int}
    nbbitstotal::Int
end
RealCoding(ϵ::Vector{Int}, lb, ub) = begin
    @assert length(lb)==length(ub)
    for i = 1:length(lb)
        @assert lb[i] < ub[i]
        @assert UInt128(10)^ϵ[i] <= UInt128(2)^127
    end
    nbvar = length(lb)
    nbbits = ones(Int, nbvar)
    for i = 1:nbvar
        while ((UInt128(10)^ϵ[i])*(ub[i]-lb[i]) >= UInt128(2)^nbbits[i])
            nbbits[i] += 1
        end
    end
    RealCoding(nbvar, lb, ub, nbbits, sum(nbbits))
end
RealCoding(ϵ::Int, lb, ub) = RealCoding([ϵ for _=1:length(lb)], lb, ub)

function decode(x, d::RealCoding)::Vector{Float64}

    res = zeros(d.nbvar)
    for i = 1:d.nbvar
        val = zero(UInt128)
        puis = one(UInt128)
        jstart = sum(d.nbbits[1:i])
        jend = jstart - d.nbbits[i] + 1
        for j = jstart:-1:jend
            x[j] && (val += puis)
            puis *= 2
        end
        res[i] = d.lb[i] + val * (d.ub[i] - d.lb[i]) / (UInt128(2)^d.nbbits[i] - 1)
    end
    res
end

function encode(x, d::RealCoding)::BitVector

    res = BitArray(0)
    for i = 1:d.nbvar
        tab = falses(d.nbbits[i])
        ind = 1
        puis = UInt128(2)^(d.nbbits[i] - 1)
        val = zero(UInt128)
        target = round(UInt128, (x[i] - d.lb[i]) / (d.ub[i] - d.lb[i]) * (UInt128(2)^d.nbbits[i] - 1))
        if target == UInt128(2)^d.nbbits[i] - 1
            target -= 1
        end
        while val < target
            if val + puis <= target
                val += puis
                tab[ind] = true
                val == target && break
            end
            puis ÷= 2
            ind += 1
        end
        append!(res, tab)
    end
    res
end


module TwoSpecies
using Distributions
using Parameters
using ...Info.MapInfo

"""
Parameters for the map.
"""
Parameters.@with_kw struct Params
    μXdist::Distributions.Distribution = Uniform(2, 4)
    μYdist::Distributions.Distribution = Uniform(3.7, 3.9)
    initdist::Distributions.Distribution = Uniform(0.1, 0.9)
    αdist::Distributions.Distribution = Uniform(0.02, 0.1)
    x₀::Float64 = rand(initdist)
    y₀::Float64 = rand(initdist)
    μX::Float64 = rand(μXdist)
    μY::Float64 = rand(μYdist)
    αXonY::Float64 = rand(αdist)
    αYonX::Float64 = rand(αdist)
end

"""
Logistic map object. Holds a realization of the twospecies logistic map, given by the
`itermap()` function. `X` and `Y` are time series of trajectories, `params` holds the
information used to generate the trajectiories, and `datestamp` is the time the realization
was generated.
"""
Parameters.@with_kw struct Map
    X::Vector{Float64} = Float64[]
    Y::Vector{Float64} = Float64[]
    info::MapInfo = MapInfo()
    params::Params = Params()
    datestamp::DateTime = now()
end

"""
    map(;params = LogisticMapParams)

Generate a trajectory of a bidirectionally coupled logistic map whose properties is
specified by a LogisticMapParams instance `p`.

The generated trajectory consists of two time series, X and Y, each `n_pts`
points long, starting from initial condition (x₀, y₀). If not specified otherwise, the
initial conditions are drawn according to `initdist`, which can be any
Distributions.Distribution object. The default is  `initdist= Uniform(0.01, 0.99)`.

The map is allowed to
evolve for `n_transient` time steps before we start to sample. After the
transient period is over, start sampling at every `stepsize`th iteration.

The coupling strength is given by `αXonY` and `αYonX`, which measures the
forcing strength of X on Y and the forcing strength of Y on X, in that order. If not
specified otherwise, `αXonY` and `αYonX` are drawn from `αdist`, which can be any
Distributions.Distribution object. The default is  `αdist= Uniform(0.05, 0.1)`.

The parameters `μX` and `μY` appear in the individual dynamics of X and Y, respectively.
If not specified otherwise, `μX` and `μY` are drawn from `μXdist` and `μYdist`, which both
can be any Distributions.Distribution object. The defaults are `μXdist = Uniform(2, 4)` and
`μYdist = Uniform(3.7, 3.9)`.

The `noise_frac` argument is the fraction of uniformly distributed noise relative to
1 standard deviation of the data.

If the model runs successfully given the chosen parameters, the function returns a
LogisticMap object. If the model collapses, a NaN is returned.
"""
function itermap(; p = Params(), info = MapInfo())

    X, Y = Vector{Float64}((info.n_pts*info.stepsize) + info.n_transient),
            Vector{Float64}((info.n_pts*info.stepsize) + info.n_transient)
    X[1] = p.x₀
    Y[1] = p.y₀

    # Run for `n_transient` steps, then start generating the map.
    for i = 2:info.n_transient+1
        X[i] = X[i-1] * (p.μX - p.μX*X[i-1] - p.αYonX*Y[i-1])
        Y[i] = Y[i-1] * (p.μY - p.μY*Y[i-1] - p.αXonY*X[i-1])
    end

    start_ind = info.n_transient + 1
    for i = start_ind:(info.n_pts*info.stepsize + (start_ind - 1))
        X[i] = X[i-1] * (p.μX - p.μX*X[i-1] - p.αYonX*Y[i-1])
        Y[i] = Y[i-1] * (p.μY - p.μY*Y[i-1] - p.αXonY*X[i-1])
    end

    if any(isnan, X) || any(isnan, Y)
        warn("Model collapsed. Retrying, keeping forcing strength and random seed, but randomising all other parameters.")
        return itermap(p = Params(αXonY = p.αXonY, αYonX = p.αYonX), info = info)
    end

    sample_inds = start_ind:info.stepsize:(info.n_pts*info.stepsize + (start_ind - 1))
    if info.noise_frac > 0
        X_final = X[sample_inds]
        Y_final = Y[sample_inds]
        noiseX = rand(Distributions.Uniform(-(std(X_final) * info.noise_frac), std(X_final) * info.noise_frac), info.n_pts)
        noiseY = rand(Distributions.Uniform(-(std(Y_final) * info.noise_frac), std(Y_final) * info.noise_frac), info.n_pts)
        return Map(
            X = X_final + noiseX,
            Y = Y_final + noiseY,
            params = p,
            info = info,
            datestamp = now())
    else
        return Map(X = X[sample_inds], Y = Y[sample_inds], params = p, info = info, datestamp = now())
    end
end

export Params, Map, itermap

end


#############################################################################
# operatornorm.jl
# Handles matrix operator norm (the maximum singular value of a matrix)
# and creates the alias sigmamax
# All expressions and atoms are subtypes of AbstractExpr.
# Please read expressions.jl first.
#############################################################################

### Operator norm

struct OperatorNormAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int,Int}

    function OperatorNormAtom(x::AbstractExpr)
        children = (x,)
        return new(:opnorm, hash(children), children, (1, 1))
    end
end

function sign(x::OperatorNormAtom)
    return Positive()
end

# The monotonicity
function monotonicity(x::OperatorNormAtom)
    return (NoMonotonicity(),)
end

function curvature(x::OperatorNormAtom)
    return ConvexVexity()
end

# in julia, `norm` on matrices is the operator norm
function evaluate(x::OperatorNormAtom)
    return opnorm(evaluate(x.children[1]), 2)
end

sigmamax(x::AbstractExpr) = OperatorNormAtom(x)

function LinearAlgebra.opnorm(x::AbstractExpr, p::Real = 2)
    if length(size(x)) <= 1 || minimum(size(x)) == 1
        throw(ArgumentError("argument to `opnorm` must be a matrix"))
    end
    if p == 1
        return maximum(sum(abs(x), dims = 1))
    elseif p == 2
        return OperatorNormAtom(x)
    elseif p == Inf
        return maximum(sum(abs(x), dims = 2))
    else
        throw(
            ArgumentError("matrix p-norms only defined for p = 1, 2, and Inf"),
        )
    end
end

# Create the equivalent conic problem:
#   minimize t
#   subject to
#            [tI_m A; A' tI_n] is positive semidefinite
# see eg Recht, Fazel, Parillo 2008 "Guaranteed Minimum-Rank Solutions of Linear Matrix Equations via Nuclear Norm Minimization"
# http://arxiv.org/pdf/0706.4138v1.pdf
function conic_form!(x::OperatorNormAtom, unique_conic_forms)
    if !has_conic_form(unique_conic_forms, x)
        A = x.children[1]
        m, n = size(A)
        t = Variable()
        p = minimize(t, [t*sparse(1.0I, m, m) A; A' t*sparse(1.0I, n, n)] ⪰ 0)
        cache_conic_form!(unique_conic_forms, x, p)
    end
    return get_conic_form(unique_conic_forms, x)
end


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# ╔═╡ bb6cae46-4108-11eb-2309-3d11f13a19fa
begin
	using Pkg
	Pkg.activate(mktempdir())
end

# ╔═╡ d8b334ae-4108-11eb-2546-99e75efc18d6
begin
	Pkg.add(["Plots", "PlutoUI","StatsBase", "DataStructures", "Unitful", "PhysicalConstants", "UnitfulEquivalences"])

	using Plots
	plotly()
	using PlutoUI
	using StatsBase
	using DataStructures
	using Unitful
	using PhysicalConstants.CODATA2018: c_0, h, ħ
	using UnitfulEquivalences
end

# ╔═╡ a78b9c3a-4109-11eb-26ca-1556cfeb2b8c
quantum_energy2 = uconvert(u"eV", 500u"THz", PhotonEnergy())

# ╔═╡ 5aa7ff98-410a-11eb-3848-895e37363b3c
period2 = 1/500u"THz"

# ╔═╡ 5990a14a-410b-11eb-086a-f129b7fdf82c
refracted_speed2 = uconvert(u"Mm/s", c_0 / 1.5)

# ╔═╡ 6209ea56-410c-11eb-0bb3-ff80947cdcb0
wavelength2 = uconvert(u"nm", refracted_speed2*period2)

# ╔═╡ 84c21a6c-4109-11eb-166e-151c460d28c8
md"""
## E2

Light, which possesses in vacuum the frequency of 500 THz and quantum energy of $quantum_energy2, passes through
sheet of glass with the refractive index of 1.5. The speed of such light in glass is $refracted_speed2, the frequency is
500 THz, the wavelength is $wavelength2 and quantum energy is 2.07 eV.
"""

# ╔═╡ fdae38fe-410c-11eb-18dd-9bae51fb4e00
freq3 = uconvert(u"THz", (1E13)u"Hz")

# ╔═╡ fdbe0894-410c-11eb-0ace-3d269b2942aa
n = 9

# ╔═╡ 220ec448-410d-11eb-17c3-d1682a1588c9
refractive3 = sqrt(n)

# ╔═╡ 74fa7748-410c-11eb-0688-a37608618c96
md"""
## E3.1

The electric field strength that oscillates with the frequency of **$freq3** is reduced **$n** times in some substance
with respect to the vacuum. The refractive index of this substance is **$refractive3**
"""

# ╔═╡ 31bce1b8-410d-11eb-008b-bf04b66bf555
slit_length4 = 0.03u"mm"

# ╔═╡ 2e0551ea-410d-11eb-176c-650b9c498d25
distance_from_slit = 2u"m"

# ╔═╡ 0aaf6748-410e-11eb-035a-8be0950091d6
minima_distance = 8u"cm"

# ╔═╡ 0aa8bb82-410e-11eb-163c-6de6eda9e616
minima_angle = uconvert(u"rad", atan(minima_distance/2, distance_from_slit))

# ╔═╡ 70a82a76-410e-11eb-2f0e-295b06d5d8fb
wavelength4 = uconvert(u"nm", sin(minima_angle) * slit_length4)

# ╔═╡ 3bb15c6c-410d-11eb-1b1b-4b13708ceb70
md"""
## E4.1

The monochromatic light is penetrating a single **$slit_length4** wide vertical slit. On the screen positioned perpendicular to the light propagating direction at the distance of **$distance_from_slit** behind the slit, a diffraction pattern appears in which the two first minima are located at the distance of **$minima_distance** from each other. The first minimum is located at the angle of **$minima_angle** with respect to the diffraction pattern centre. The wavelength of the light used was **$wavelength4**
"""

# ╔═╡ c08c3f10-4117-11eb-059e-1fa7c75e4240
absorption_coef_blue = 5u"cm^-1"

# ╔═╡ ffd6d630-4117-11eb-2200-a147d74af92f
absorbtion_coef_green = 3u"cm^-1"

# ╔═╡ 38b97e80-4118-11eb-34c5-65de18666aed
absorption_coef_red = 0.3u"cm^-1"

# ╔═╡ 45bb6c38-4118-11eb-0ba5-f55cde355ccd
blue_layer_thickness = uconvert(u"cm", 1/absorption_coef_blue)

# ╔═╡ 40d8a62c-4118-11eb-15c3-e9e0c5d87609
n_intensity_reduced = 3

# ╔═╡ 793b338a-4119-11eb-1916-6513f6eed696
ℯ

# ╔═╡ afcc4710-4119-11eb-28db-fdf5aa042d8e
substance_color1 = "red"

# ╔═╡ d0340d2e-4119-11eb-209d-7dcf7fa4c2e7
substance_color2 = "black"

# ╔═╡ c094cb3a-4117-11eb-168a-df762ad5739e
md"""
## E5
For some substance, the absorption coefficient for the blue light is **$absorption_coef_blue**. For the green light, it is **$absorbtion_coef_green** and for the red light **$absorption_coef_red**. This means that the intensity of the blue light will be reduced Euler number (e = 2.7183..) times by passing through the **$blue_layer_thickness** thick layer of the substance. The layer with the thickness of 1 cm reduces the intensity of green light e **$n_intensity_reduced** times. If we illuminate this substance with the white light, the colour of the substance will be **$substance_color1**. If we illuminate the substance with the blue light, the colour of the substance will be **$substance_color2**.

"""

# ╔═╡ 80572be0-411b-11eb-3f15-fbf4e6ed7cf3
electron_speed = 729u"km/s"

# ╔═╡ 5788010a-411e-11eb-2c69-bd73c5cd6d3f
ground_speed = 2188u"km/s"

# ╔═╡ 7fe77a98-411b-11eb-1835-75b0687dedfd
n6 = round(ground_speed / electron_speed)

# ╔═╡ 7fc27112-411b-11eb-3b4b-97599735cd31
angular_momentum = uconvert(u"kg*m^2*s^-1", 1.05E-34u"kg*m^2*s^-1" * n6)

# ╔═╡ 224bba9e-411f-11eb-2901-87f766e95a56
ground_frequency = 6580u"THz"

# ╔═╡ 3ed8da6c-411e-11eb-3a35-4164dfd2aa54
imagined_rotation_frequency = uconvert(u"THz", ground_frequency / n6 ^3)

# ╔═╡ 91b1a86c-411f-11eb-0700-a30281bd1f40
ground_energy = -13.60u"eV"

# ╔═╡ 8a2e2e76-411f-11eb-0da9-ed7c9382599f
net_energy = ground_energy / n6^2

# ╔═╡ 0396e5f2-4120-11eb-1502-f358ec34acd9
ground_radius = 0.529u"Å"

# ╔═╡ 03d6d638-4120-11eb-20d0-0fa5cffc6d3d
mean_distance = ground_radius * n6^2

# ╔═╡ 296b84ce-4121-11eb-2c07-1b1ff4cac46c
emitted_photon_wavelength = 656u"nm"

# ╔═╡ f74d6ff0-4120-11eb-2de1-9d7c40c9354c
quantum_energy6 = uconvert(u"eV", emitted_photon_wavelength, PhotonEnergy())

# ╔═╡ 66942310-4121-11eb-063c-37022902216f
energy_after_photon = net_energy - quantum_energy6

# ╔═╡ 81bc1422-4121-11eb-0cb5-215e6fb5c308
qn = round(sqrt(ground_energy/energy_after_photon))

# ╔═╡ 8238f47e-4121-11eb-06a7-a3f1bf8887cb
imag_rot_freq2 = uconvert(u"THz", ground_frequency / qn ^3)

# ╔═╡ 0e7c4e68-4122-11eb-3b1a-8f9808df6829
mean_distance2 = ground_radius * qn^2

# ╔═╡ 7fec90b4-411b-11eb-04db-a913940a9f22
md"""
## E6
In the atom of hydrogen, the electron moves according to the Bohr model. Its speed is **$electron_speed**. The angular momentum of this electron is **$angular_momentum**, the frequency of imagined rotation is **$imagined_rotation_frequency**, the net energy is **$net_energy** and the mean distance from the nucleus is **$mean_distance**. The electron comes closer to the nucleus and a photon with the wavelength of **$emitted_photon_wavelength** is emitted. The quantum energy of this radiation is **$quantum_energy6**. At the end of the transition, the frequency of the imagined rotation of the electron is **$imag_rot_freq2** and its mean distance from the nucleus is **$mean_distance2**.
"""

# ╔═╡ 46652b7e-4122-11eb-1576-29e34caf1b03
mass = 4E-25u"kg"

# ╔═╡ 78e0e354-4122-11eb-1a11-0b1d36a38b70
diam = 5E-10u"m"

# ╔═╡ 46798970-4122-11eb-31a8-ab967162b1ae
h_bar = 1E-34u"J*s"

# ╔═╡ 4691ff34-4122-11eb-1614-b5df5b328525
speed_increase = uconvert(u"m/s", h_bar / (mass * diam))

# ╔═╡ c9767010-4125-11eb-3529-43e03fc0ed7b
time_for_diam = diam / speed_increase

# ╔═╡ c9680bce-4125-11eb-1a30-93e346ab5ef8
time2 = uconvert(u"ns", time_for_diam)

# ╔═╡ 451edba2-4122-11eb-0ea9-b3e5a7021bd1
md"""
# E7

The uranium atom (a microscopic object) with the mass of **$mass** and diameter of **$diam** moves in a certain direction. The location of the atom is fixed by its passing through the slit possessing the width almost equal to the atom diameter. The screen containing this slit is positioned perpendicular to the velocity of the atom. As a result of penetrating this slit, the speed of the atom receives some unpredictable increase perpendicular to the direction of previous motion. Using the reduced Planck constant 10–34 J*s as a typical value of impact uncertainty, we can conclude that the increase of the speed is **$speed_increase**. After **$time_for_diam** seconds, the atom deviates from its classical straight trajectory by its own diameter. This result means that we can predict the motion of uranium atom in a classical way for maximally **$time2** nanoseconds.

"""

# ╔═╡ aa066c66-4126-11eb-1909-713aeda74edc
mass8 = 9E-31u"kg"

# ╔═╡ c4375a50-4126-11eb-2667-7b693682e9a8
well_width = 3u"Å"

# ╔═╡ aa323b5c-4126-11eb-18fa-a5aa100d0678
qn8 = 3

# ╔═╡ aa487444-4126-11eb-306c-1d98f391cfc6
planck8 = 6.6E-34u"J*s"

# ╔═╡ aa708326-4126-11eb-2129-df2f5666aa93
wavelength = 2*well_width / qn8

# ╔═╡ aa9ab420-4126-11eb-206b-2bd429f38b13
lin_momentum = uconvert(u"kg*m/s", planck8 / wavelength)

# ╔═╡ aa805792-4126-11eb-20eb-b127d6ae2898
energy8 = uconvert(u"eV", lin_momentum^2/(2*mass8))

# ╔═╡ c98691d4-4125-11eb-3e71-d1ffa4fc0f3a
md"""
## E8

An electron with the mass of **$mass8** is located in a single-dimensional potential well with the width of **$well_width** (approximate diameter of the atom). Let us describe this electron according to the laws of quantum mechanics as a standing wave in the string, whereas the length of the string is equal to the width of the potential well. In the state described by the quantum number **$qn8**, the wavelength of the electron is **$wavelength**, its
linear momentum is **$lin_momentum** and its energy is **$energy8**.

"""

# ╔═╡ aa5f27d4-4126-11eb-19db-9b02f89cf0f0

# ╔═╡ aa1d3ab8-4126-11eb-11d9-ab8a72c7e68f

# ╔═╡ Cell order:
# ╠═bb6cae46-4108-11eb-2309-3d11f13a19fa
# ╠═d8b334ae-4108-11eb-2546-99e75efc18d6
# ╟─84c21a6c-4109-11eb-166e-151c460d28c8
# ╠═a78b9c3a-4109-11eb-26ca-1556cfeb2b8c
# ╠═5aa7ff98-410a-11eb-3848-895e37363b3c
# ╠═5990a14a-410b-11eb-086a-f129b7fdf82c
# ╠═6209ea56-410c-11eb-0bb3-ff80947cdcb0
# ╟─74fa7748-410c-11eb-0688-a37608618c96
# ╠═fdae38fe-410c-11eb-18dd-9bae51fb4e00
# ╠═fdbe0894-410c-11eb-0ace-3d269b2942aa
# ╠═220ec448-410d-11eb-17c3-d1682a1588c9
# ╠═3bb15c6c-410d-11eb-1b1b-4b13708ceb70
# ╠═31bce1b8-410d-11eb-008b-bf04b66bf555
# ╠═2e0551ea-410d-11eb-176c-650b9c498d25
# ╠═0aaf6748-410e-11eb-035a-8be0950091d6
# ╠═0aa8bb82-410e-11eb-163c-6de6eda9e616
# ╠═70a82a76-410e-11eb-2f0e-295b06d5d8fb
# ╠═c094cb3a-4117-11eb-168a-df762ad5739e
# ╠═c08c3f10-4117-11eb-059e-1fa7c75e4240
# ╠═ffd6d630-4117-11eb-2200-a147d74af92f
# ╠═38b97e80-4118-11eb-34c5-65de18666aed
# ╠═45bb6c38-4118-11eb-0ba5-f55cde355ccd
# ╠═40d8a62c-4118-11eb-15c3-e9e0c5d87609
# ╠═793b338a-4119-11eb-1916-6513f6eed696
# ╠═afcc4710-4119-11eb-28db-fdf5aa042d8e
# ╠═d0340d2e-4119-11eb-209d-7dcf7fa4c2e7
# ╟─7fec90b4-411b-11eb-04db-a913940a9f22
# ╠═80572be0-411b-11eb-3f15-fbf4e6ed7cf3
# ╠═5788010a-411e-11eb-2c69-bd73c5cd6d3f
# ╠═7fe77a98-411b-11eb-1835-75b0687dedfd
# ╠═7fc27112-411b-11eb-3b4b-97599735cd31
# ╠═224bba9e-411f-11eb-2901-87f766e95a56
# ╠═3ed8da6c-411e-11eb-3a35-4164dfd2aa54
# ╠═91b1a86c-411f-11eb-0700-a30281bd1f40
# ╠═8a2e2e76-411f-11eb-0da9-ed7c9382599f
# ╠═0396e5f2-4120-11eb-1502-f358ec34acd9
# ╠═03d6d638-4120-11eb-20d0-0fa5cffc6d3d
# ╠═296b84ce-4121-11eb-2c07-1b1ff4cac46c
# ╠═f74d6ff0-4120-11eb-2de1-9d7c40c9354c
# ╠═66942310-4121-11eb-063c-37022902216f
# ╠═81bc1422-4121-11eb-0cb5-215e6fb5c308
# ╠═8238f47e-4121-11eb-06a7-a3f1bf8887cb
# ╠═0e7c4e68-4122-11eb-3b1a-8f9808df6829
# ╟─451edba2-4122-11eb-0ea9-b3e5a7021bd1
# ╠═46652b7e-4122-11eb-1576-29e34caf1b03
# ╠═78e0e354-4122-11eb-1a11-0b1d36a38b70
# ╠═46798970-4122-11eb-31a8-ab967162b1ae
# ╠═4691ff34-4122-11eb-1614-b5df5b328525
# ╠═c9767010-4125-11eb-3529-43e03fc0ed7b
# ╠═c9680bce-4125-11eb-1a30-93e346ab5ef8
# ╠═c98691d4-4125-11eb-3e71-d1ffa4fc0f3a
# ╠═aa066c66-4126-11eb-1909-713aeda74edc
# ╠═c4375a50-4126-11eb-2667-7b693682e9a8
# ╠═aa323b5c-4126-11eb-18fa-a5aa100d0678
# ╠═aa487444-4126-11eb-306c-1d98f391cfc6
# ╠═aa708326-4126-11eb-2129-df2f5666aa93
# ╠═aa9ab420-4126-11eb-206b-2bd429f38b13
# ╠═aa805792-4126-11eb-20eb-b127d6ae2898
# ╠═aa5f27d4-4126-11eb-19db-9b02f89cf0f0
# ╠═aa1d3ab8-4126-11eb-11d9-ab8a72c7e68f


using PyCall
using SymPy
using REPL

tosub = Dict(
    '1' => '₁',
    '2' => '₂',
    '3' => '₃',
    '4' => '₄',
    '5' => '₅',
    '6' => '₆',
    '7' => '₇',
    '8' => '₈',
    '9' => '₉',
    '0' => '₀',
)

"""
Generate Sym object its names are
x₁,x₂,...,xₙ
and return them as Sym[]
Example
julia> generate_symbol("x",3)
Sym[x₁, x₂, x₃]
"""
function generate_symbol(v::String, n::Int)
    data::Vector{Sym} = []
    for i = 1:n
        seqdigit = [tosub[c] for c in string(i)]
        v_s = "$v" * join(seqdigit)
        eval(Meta.parse("$(v_s)=Sym(\"$(v_s)\")"))
        push!(data, eval(Meta.parse("$(v_s)")))
    end
    return data
end

"""
This func cant use n more than 9
"""
function _generate_symbol(v, n)
    data = Sym[]
    for i = 1:n
        target = "$(v)\\_$i"
        _sub = REPL.REPLCompletions.completions(target, length(target))[1][1]
        sub = REPL.REPLCompletions.completion_text(_sub)
        v_s = v * sub
        eval(Meta.parse("$(v_s)=Sym(\"$(v_s)\")"))
        push!(data, eval(Meta.parse("$(v_s)")))
    end
    return data
end

n = 5
xs = generate_symbol("x", n)
ys = generate_symbol("y", n)

# Uncomment if you can switch randomized numerical data
#xs, ys=rand(n), rand(n)
#

@vars a b

model(x) = a * x + b
E = 1 / 2 * sum(@. (model(xs) - ys)^2)

function bruteforce()
    ∂E∂a, ∂E∂b = diff(E, a), diff(E, b)
    solution = solve([∂E∂a, ∂E∂b], (a, b))
    solution
end

function theoretical()
    denom = n * sum(xs .^ 2) - sum(xs)^2
    a_num = n * sum(xs .* ys) - sum(xs) * sum(ys)
    â = a_num / denom
    b_num = sum(xs .^ 2) * sum(ys) - sum(xs) * sum(xs .* ys)
    b̂ = b_num / denom
    return Dict(a => â, b => b̂)
end

function uselinsolve1()
    ∂E∂a, ∂E∂b = diff(E, a), diff(E, b)
    solution = linsolve((∂E∂a, ∂E∂b), (a, b))
    return solution
end

function uselinsolve2()
    A = [
        diff(E, a, 2) diff(diff(E, a), b)
        diff(diff(E, b), a) diff(E, b, 2)
    ] .|> simplify
    v = -[diff(E, a)(a => 0, b => 0), diff(E, b)(a => 0, b => 0)]
    solution = linsolve((A, v), a, b)
    return solution
end

py"""
def decompose(r):
    return [[eq for eq in eqs] for eqs in r.args]
"""

r0 = theoretical()
r1 = bruteforce()
r2 = py"decompose"(uselinsolve1())
r3 = py"decompose"(uselinsolve2())

@show r0[a]
@show r1[a]
@show r2[1] # a
@show r3[1] # a

@show r0[b]
@show r1[b]
@show r2[2] # b
@show r3[2] # b


#=
    `MVNSampler` is used to draw from multivariate normal distribution
=#

import Base: ==
import LinearAlgebra: BlasReal
import Random

struct MVNSampler{TM<:Real,TS<:Real,TQ<:BlasReal}
    mu::Vector{TM}
    Sigma::Matrix{TS}
    Q::Matrix{TQ}
end

function MVNSampler(mu::Vector{TM}, Sigma::Matrix{TS}) where {TM<:Real,TS<:Real}
    ATOL1, RTOL1 = 1e-8, 1e-8
    ATOL2, RTOL2 = 1e-8, 1e-14

    n = length(mu)

    if size(Sigma) != (n, n) # Check Sigma is n x n
        throw(ArgumentError(
            "Sigma must be 2 dimensional and square matrix of same length to mu"
        ))
    end

    issymmetric(Sigma) || throw(ArgumentError("Sigma must be symmetric"))

    C = cholesky(Symmetric(Sigma, :L), Val(true), check=false)
    A = C.factors
    r = C.rank
    p = invperm(C.piv)

    if r == n  # Positive definite
        Q = tril!(A)[p, p]
        return MVNSampler(mu, Sigma, Q)
    end

    non_PSD_msg = "Sigma must be positive semidefinite"

    for i in r+1:n
        A[i, i] >= -ATOL1 - RTOL1 * A[1, 1] ||
            throw(ArgumentError(non_PSD_msg))
    end

    tril!(view(A, :, 1:r))
    A[:, r+1:end] .= 0
    Q = A[p, p]
    isapprox(Q*Q', Sigma; rtol=RTOL2, atol=ATOL2) ||
        throw(ArgumentError(non_PSD_msg))

    return MVNSampler(mu, Sigma, Q)
end

# methods with the optional rng argument first
Random.rand(rng::AbstractRNG, d::MVNSampler) =
    d.mu + d.Q * randn(rng, length(d.mu))
Random.rand(rng::AbstractRNG, d::MVNSampler, n::Integer) =
    d.mu .+ d.Q * randn(rng, (length(d.mu), n))

# methods to draw from `MVNSampler`
Random.rand(d::MVNSampler) = rand(Random.GLOBAL_RNG, d)
Random.rand(d::MVNSampler, n::Integer) = rand(Random.GLOBAL_RNG, d, n)

==(f1::MVNSampler, f2::MVNSampler) =
    (f1.mu == f2.mu) && (f1.Sigma == f2.Sigma) && (f1.Q == f2.Q)


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ dcab9a18-5d58-4632-a325-d16af9c79344
using Pkg, DrWatson

# ╔═╡ c95f3f02-f179-4f81-9c53-6e6b87f25613
begin
	@quickactivate "NumericalMethodsforEngineers"
	using NumericalMethodsforEngineers, Symbolics, SymbolicUtils
	using StaticArrays
end

# ╔═╡ c9036632-adb9-11eb-33d1-c5193729e72a
function lagrangianpolynomial1(x::Vector{Float64},
		y::Vector{Float64}, xi::Vector{Float64})

	np = length(x)
 	yi = Vector{Float64}(undef, length(xi))
 	m = sort(hcat(x, y), dims=1)
	for l in 1:length(xi)
		for i in 1:np
			term = 1.0
			for j in 1:np
				j !== i && (term *= (xi[l] - m[j, 1])/(m[i, 1] - m[j, 1]))
			end
			yi[l] += term * m[i, 2]
		end
	end
	columns=[]
	push!(columns, m[:,1])
	push!(columns, m[:,2])
	df = DataFrame(columns, [:x, :y])

	columns=[]
	push!(columns, xi)
	push!(columns, yi)
	dfxi = DataFrame(columns, [:xi, :yi])

	(df, dfxi)
end

# ╔═╡ 298bbd55-d7b6-4229-8848-f0ea11f6448e

# ╔═╡ b8e7c2c7-5a72-4d62-9c2d-0b6cfa3e8fbb
@variables x y;

# ╔═╡ 057c029c-bd62-4ecb-9dc8-ed0f35684fd0
to_compute = [x^2 + y, y^2 + x];

# ╔═╡ d8e6cfe2-4cab-476d-8730-ec23470ac1b7
f_expr = build_function(to_compute, [x, y]);

# ╔═╡ dcb5841b-9ab1-4efe-928a-a69a02dffc95
begin
	myf = eval(f_expr[1])
	myf(SA[2.0, 3.0])
end

# ╔═╡ 2f360a8b-4ab6-4617-9481-1a10eab71579
Base.remove_linenums!(f_expr[2])

# ╔═╡ e5f5c803-d739-460f-bdf2-c9f5b55786a3
begin
	myf! = eval(f_expr[2])
	out = zeros(2)
	myf!(out, [2.0, 3.0])
	out
end

# ╔═╡ add02f10-37f9-49b1-80e4-205758be2605
f1 = build_function(to_compute, [x, y], expression=Val{false});

# ╔═╡ 2c817c8d-d108-432c-98b4-f849ec2c9c71
f1[1]([2.0, 3.0])

# ╔═╡ Cell order:
# ╠═dcab9a18-5d58-4632-a325-d16af9c79344
# ╠═c95f3f02-f179-4f81-9c53-6e6b87f25613
# ╠═c9036632-adb9-11eb-33d1-c5193729e72a
# ╠═298bbd55-d7b6-4229-8848-f0ea11f6448e
# ╠═b8e7c2c7-5a72-4d62-9c2d-0b6cfa3e8fbb
# ╠═057c029c-bd62-4ecb-9dc8-ed0f35684fd0
# ╠═d8e6cfe2-4cab-476d-8730-ec23470ac1b7
# ╠═dcb5841b-9ab1-4efe-928a-a69a02dffc95
# ╠═2f360a8b-4ab6-4617-9481-1a10eab71579
# ╠═e5f5c803-d739-460f-bdf2-c9f5b55786a3
# ╠═add02f10-37f9-49b1-80e4-205758be2605
# ╠═2c817c8d-d108-432c-98b4-f849ec2c9c71


# averaging functions, to operate on resistances
# cond_* functions convert to conductance, calulate mean, then convert to resistance
# res_* functions directly operate on the resistances
cond_cardinal_avg(x, y) = 1 / ((1/x + 1/y) / 2)
cond_diagonal_avg(x, y) = 1 / ((1/x + 1/y) / (2 * √2))
res_cardinal_avg(x, y) = (x + y) / 2
res_diagonal_avg(x, y) = ((x + y) * √2) / 2

function get_lat_lon_dims(A::GeoData.GeoArray)
    dim_names = String.(name(dims(A)))
    lat_first = findall(dim_names .== "Latitude")[1] < findall(dim_names .== "Longitude")[1]
    first_dim_type = lat_first ? Lat : Lon
    second_dim_type = lat_first ? Lon : Lat

    (dims(A, first_dim_type), dims(A, second_dim_type), Band(1:1, Categorical(order = Ordered())))
end


module Tutorial3
using FEniCS

T = 2.0            # final time
num_steps = 10     # number of time steps
dt = T / num_steps # time step size
alpha = 3          # parameter alpha
beta = 1.2
nx = ny = 8
mesh = UnitSquareMesh(nx, ny)
V = FunctionSpace(mesh, "P", 1)
u_D = Expression("1 + x[0]*x[0] + alpha*x[1]*x[1] + beta*t",
degree=2, alpha=alpha, beta=beta, t=0)

bc = DirichletBC(V, u_D, "on_boundary")
u_n = interpolate(u_D, V)
u = TrialFunction(V)
v = TestFunction(V)
f = Constant(beta - 2 - 2*alpha)
F = u*v*dx + dt*dot(grad(u), grad(v))*dx - (u_n + dt*f)*v*dx

a, L = lhs(F),rhs(F)

u=FeFunction(V)
global t = 0

for n = 0:(num_steps-1)
    global t=t+dt
    u_D.pyobject.t=t
    lvsolve(a,L,u,bc)
    u_e= interpolate(u_D,V)
    vv = get_array(u_e)
    ww = get_array(u)
    error = maximum(abs.(vv-ww))
    #@printf "t = %f: error = %10g" t error
    assign(u_n,u)
end
#colorbar()
end#module