abstract type AbstractBlochMatrix{T<:Real} end

"""
    BlochMatrix(a11, a21, a31, a12, a22, a32, a13, a23, a33)
    BlochMatrix{T}()
    BlochMatrix()

Create a mutable `BlochMatrix` object representing a fixed-size 3×3 matrix.

# Properties
- `aij::Real`: Matrix entry (i,j) ∈ {1, 2, 3}²
"""
mutable struct BlochMatrix{T<:Real} <: AbstractBlochMatrix{T}
    a11::T
    a21::T
    a31::T
    a12::T
    a22::T
    a32::T
    a13::T
    a23::T
    a33::T
end

BlochMatrix(a...) = BlochMatrix(promote(a...)...)
BlochMatrix{T}() where {T} = BlochMatrix(zero(T), zero(T), zero(T), zero(T),
                                    zero(T), zero(T), zero(T), zero(T), zero(T))
BlochMatrix() = BlochMatrix{Float64}()

function Base.show(io::IO, ::MIME"text/plain", A::BlochMatrix{T}) where {T}

    print(io, "BlochMatrix{$T}:\n")
    haskey(io, :compact) || (io = IOContext(io, :compact => true))
    Base.print_array(io, Matrix(A))

end

function Base.fill!(A::BlochMatrix, v)

    A.a11 = v
    A.a21 = v
    A.a31 = v
    A.a12 = v
    A.a22 = v
    A.a32 = v
    A.a13 = v
    A.a23 = v
    A.a33 = v
    return nothing

end

function Base.copyto!(dst::BlochMatrix, src::BlochMatrix)

    dst.a11 = src.a11
    dst.a21 = src.a21
    dst.a31 = src.a31
    dst.a12 = src.a12
    dst.a22 = src.a22
    dst.a32 = src.a32
    dst.a13 = src.a13
    dst.a23 = src.a23
    dst.a33 = src.a33
    return nothing

end

function Base.Matrix(A::BlochMatrix{T}) where {T}

    mat = Matrix{T}(undef, 3, 3)
    mat[1,1] = A.a11
    mat[2,1] = A.a21
    mat[3,1] = A.a31
    mat[1,2] = A.a12
    mat[2,2] = A.a22
    mat[3,2] = A.a32
    mat[1,3] = A.a13
    mat[2,3] = A.a23
    mat[3,3] = A.a33

    return mat

end

"""
    BlochDynamicsMatrix(R1, R2, Δω)
    BlochDynamicsMatrix{T}()
    BlochDynamicsMatrix()

Create a mutable `BlochDynamicsMatrix` object.

# Properties
- `R1::Real`: Spin-lattice relaxation rate
- `R2::Real`: Spin-spin relaxation rate
- `Δω::Real`: Off-resonance frequency
"""
mutable struct BlochDynamicsMatrix{T<:Real} <: AbstractBlochMatrix{T}
    R1::T
    R2::T
    Δω::T
end

BlochDynamicsMatrix{T}() where {T} = BlochDynamicsMatrix(zero(T), zero(T), zero(T))
BlochDynamicsMatrix() = BlochDynamicsMatrix{Float64}()
BlochDynamicsMatrix(R1, R2, Δω) = BlochDynamicsMatrix(promote(R1, R2, Δω)...)

Base.convert(::Type{BlochDynamicsMatrix{T}}, A::BlochDynamicsMatrix) where {T} =
    BlochDynamicsMatrix(T(A.R1), T(A.R2), T(A.Δω))

function Base.Matrix(A::BlochDynamicsMatrix{T}) where {T}

    mat = Matrix{T}(undef, 3, 3)
    mat[1,1] = A.R2
    mat[2,1] = -A.Δω
    mat[3,1] = zero(T)
    mat[1,2] = A.Δω
    mat[2,2] = A.R2
    mat[3,2] = zero(T)
    mat[1,3] = zero(T)
    mat[2,3] = zero(T)
    mat[3,3] = A.R1

    return mat

end

"""
    FreePrecessionMatrix(E1, E2cosθ, E2sinθ)
    FreePrecessionMatrix{T}()
    FreePrecessionMatrix()

Create a mutable `FreePrecessionMatrix` object encoding the effects of
relaxation and off-resonance precession.

# Properties
- `E1::Real`: T1 relaxation
- `E2cosθ::Real`: T2 relaxation and off-resonance precession
- `E2sinθ::Real`: T2 relaxation and off-resonance precession

# Examples
```jldoctest
julia> T1 = 1000; T2 = 100; Δω = π/600; t = 100;

julia> F = FreePrecessionMatrix(exp(-t / T1), exp(-t / T2) * cos(Δω * t), exp(-t / T2) * sin(Δω * t))
FreePrecessionMatrix{Float64}:
 E1 = 0.9048374180359595
 E2cosθ = 0.31859294158449203
 E2sinθ = 0.18393972058572114

julia> Matrix(F)
3×3 Matrix{Float64}:
  0.318593  0.18394   0.0
 -0.18394   0.318593  0.0
  0.0       0.0       0.904837
```
"""
mutable struct FreePrecessionMatrix{T<:Real} <: AbstractBlochMatrix{T}
    E1::T
    E2cosθ::T
    E2sinθ::T
end

FreePrecessionMatrix{T}() where {T} = FreePrecessionMatrix(zero(T), zero(T), zero(T))
FreePrecessionMatrix() = FreePrecessionMatrix{Float64}()
FreePrecessionMatrix(E1, E2cosθ, E2sinθ) = FreePrecessionMatrix(promote(E1, E2cosθ, E2sinθ)...)

function Base.show(io::IO, ::MIME"text/plain", A::FreePrecessionMatrix{T}) where {T}

    print(io, "FreePrecessionMatrix{$T}:")
    print(io, "\n E1 = ", A.E1)
    print(io, "\n E2cosθ = ", A.E2cosθ)
    print(io, "\n E2sinθ = ", A.E2sinθ)

end

function Base.Matrix(A::FreePrecessionMatrix{T}) where {T}

    mat = Matrix{T}(undef, 3, 3)
    mat[1,1] = A.E2cosθ
    mat[2,1] = -A.E2sinθ
    mat[3,1] = zero(T)
    mat[1,2] = A.E2sinθ
    mat[2,2] = A.E2cosθ
    mat[3,2] = zero(T)
    mat[1,3] = zero(T)
    mat[2,3] = zero(T)
    mat[3,3] = A.E1

    return mat

end

function Base.copyto!(dst::BlochMatrix{T}, src::FreePrecessionMatrix) where {T}

    dst.a11 = src.E2cosθ
    dst.a21 = -src.E2sinθ
    dst.a31 = zero(T)
    dst.a12 = src.E2sinθ
    dst.a22 = src.E2cosθ
    dst.a32 = zero(T)
    dst.a13 = zero(T)
    dst.a23 = zero(T)
    dst.a33 = src.E1
    return nothing

end

"""
    ExchangeDynamicsMatrix(r)
    ExchangeDynamicsMatrix{T}()
    ExchangeDynamicsMatrix()

Create a mutable `ExchangeDynamicsMatrix` object.

# Properties
- `r::Real`: Exchange rate from one compartment to another
"""
mutable struct ExchangeDynamicsMatrix{T<:Real} <: AbstractBlochMatrix{T}
    r::T
end

ExchangeDynamicsMatrix{T}() where {T} = ExchangeDynamicsMatrix(zero(T))
ExchangeDynamicsMatrix() = ExchangeDynamicsMatrix{Float64}()

Base.convert(::Type{ExchangeDynamicsMatrix{T}}, A::ExchangeDynamicsMatrix) where {T} =
    ExchangeDynamicsMatrix(T(A.r))

function Base.Matrix(A::ExchangeDynamicsMatrix{T}) where {T}

    mat = Matrix{T}(undef, 3, 3)
    mat[1,1] = A.r
    mat[2,1] = zero(T)
    mat[3,1] = zero(T)
    mat[1,2] = zero(T)
    mat[2,2] = A.r
    mat[3,2] = zero(T)
    mat[1,3] = zero(T)
    mat[2,3] = zero(T)
    mat[3,3] = A.r

    return mat

end

abstract type AbstractBlochMcConnellMatrix{T<:Real,N} end

"""
    BlochMcConnellDynamicsMatrix(A, E)
    BlochMcConnellDynamicsMatrix{T}(N)
    BlochMcConnellDynamicsMatrix(N)

Create a `BlochMcConnellDynamicsMatrix` object with `N` compartments.

# Properties
- `A::NTuple{N,BlochDynamicsMatrix{<:Real}}`: List of `BlochDynamicsMatrix`es
  that make up the main block diagonal of the `BlochMcConnellDynamicsMatrix`
- `E::NTuple{M,ExchangeDynamicsMatrix{<:Real}}`: List of
  `ExchangeDynamicsMatrix`es that describe exchange between the different
  compartments; these matrices make up the remaining `M = N^2 - N` blocks of the
  `BlochMcConnellDynamicsMatrix`, sorted by column-major ordering
"""
struct BlochMcConnellDynamicsMatrix{T<:Real,N,M} <: AbstractBlochMcConnellMatrix{T,N}
    A::NTuple{N,BlochDynamicsMatrix{T}}
    E::NTuple{M,ExchangeDynamicsMatrix{T}}

    function BlochMcConnellDynamicsMatrix(
        A::NTuple{N,BlochDynamicsMatrix{T1}},
        E::NTuple{M,ExchangeDynamicsMatrix{T2}}
    ) where {M,N,T1,T2}

        M == N * (N - 1) || error("exchange rates must be defined for each pair of compartments")
        Tnew = promote_type(T1, T2)
        A = convert.(BlochDynamicsMatrix{Tnew}, A)
        E = convert.(ExchangeDynamicsMatrix{Tnew}, E)
        new{Tnew,N,M}(A, E)

    end
end

function BlochMcConnellDynamicsMatrix{T}(N) where {T}

    A = ntuple(i -> BlochDynamicsMatrix{T}(), N)
    E = ntuple(i -> ExchangeDynamicsMatrix{T}(), N * (N - 1))
    BlochMcConnellDynamicsMatrix(A, E)

end

BlochMcConnellDynamicsMatrix(N) = BlochMcConnellDynamicsMatrix{Float64}(N)

# N = 3
#   [2,1]: (1-1)*3-1+2+(2<1) = 0-1+2+0 = 1 ✓
#   [3,1]: (1-1)*3-1+3+(3<1) = 0-1+3+0 = 2 ✓
#   [1,2]: (2-1)*3-2+1+(1<2) = 3-2+1+1 = 3 ✓
#   [3,2]: (2-1)*3-2+3+(3<2) = 3-2+3+0 = 4 ✓
#   [1,3]: (3-1)*3-3+1+(1<3) = 6-3+1+1 = 5 ✓
#   [2,3]: (3-1)*3-3+2+(2<3) = 6-3+2+1 = 6 ✓
function getblock(A::BlochMcConnellDynamicsMatrix{T,N,M}, i, j) where {T,N,M}

    return i == j ? A.A[i] : A.E[(j-1)*N-j+i+(i<j)]

end

function Base.Matrix(A::BlochMcConnellDynamicsMatrix{T,N,M}) where {T,N,M}

    mat = Matrix{T}(undef, 3N, 3N)
    index = 0
    for j = 1:N, i = 1:N

        if i == j

            mat[3i-2,3j-2] = A.A[i].R2
            mat[3i-1,3j-2] = -A.A[i].Δω
            mat[3i  ,3j-2] = zero(T)
            mat[3i-2,3j-1] = A.A[i].Δω
            mat[3i-1,3j-1] = A.A[i].R2
            mat[3i  ,3j-1] = zero(T)
            mat[3i-2,3j]   = zero(T)
            mat[3i-1,3j]   = zero(T)
            mat[3i  ,3j]   = A.A[i].R1

        else

            index += 1
            mat[3i-2,3j-2] = A.E[index].r
            mat[3i-1,3j-2] = zero(T)
            mat[3i  ,3j-2] = zero(T)
            mat[3i-2,3j-1] = zero(T)
            mat[3i-1,3j-1] = A.E[index].r
            mat[3i  ,3j-1] = zero(T)
            mat[3i-2,3j]   = zero(T)
            mat[3i-1,3j]   = zero(T)
            mat[3i  ,3j]   = A.E[index].r

        end

    end

    return mat

end

"""
    BlochMcConnellMatrix(A)
    BlochMcConnellMatrix{T}(N)
    BlochMcConnellMatrix(N)

Create a `BlochMcConnellMatrix` object with `N` compartments and representing a
fixed-size 3N×3N matrix.

# Properties
- `A::NTuple{N,NTuple{N,BlochMatrix{<:Real}}}`: List of 3×3 matrices that
  comprise the blocks of the `BlochMcConnellMatrix`; `A[i][j]` is the (i,j)th
  block

# Examples
```jldoctest
julia> B = BlochMcConnellMatrix(3)
BlochMcConnellMatrix{Float64,3}:
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

julia> fill!(B, 0.123456789)

julia> B
BlochMcConnellMatrix{Float64,3}:
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
 0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457  0.123457
```
"""
struct BlochMcConnellMatrix{T<:Real,N} <: AbstractBlochMcConnellMatrix{T,N}
    A::NTuple{N,NTuple{N,BlochMatrix{T}}}
end

function BlochMcConnellMatrix{T}(N) where {T}

    A = ntuple(i -> ntuple(i -> BlochMatrix{T}(), N), N)
    BlochMcConnellMatrix(A)

end

BlochMcConnellMatrix(N) = BlochMcConnellMatrix{Float64}(N)

function Base.show(io::IO, ::MIME"text/plain", A::BlochMcConnellMatrix{T,N}) where {T,N}

    print(io, "BlochMcConnellMatrix{$T,$N}:\n")
    haskey(io, :compact) || (io = IOContext(io, :compact => true))
    Base.print_array(io, Matrix(A))

end

getblock(A::BlochMcConnellMatrix, i, j) = A.A[i][j]

function Base.fill!(A::BlochMcConnellMatrix{T,N}, v) where {T,N}

    for i = 1:N, j = 1:N
        fill!(getblock(A, i, j), v)
    end

end

function Base.copyto!(dst::BlochMcConnellMatrix{T,N}, src::BlochMcConnellMatrix{S,N}) where {S,T,N}

    for j = 1:N, i = 1:N
        db = getblock(dst, i, j)
        sb = getblock(src, i, j)
        db.a11 = sb.a11
        db.a21 = sb.a21
        db.a31 = sb.a31
        db.a12 = sb.a12
        db.a22 = sb.a22
        db.a32 = sb.a32
        db.a13 = sb.a13
        db.a23 = sb.a23
        db.a33 = sb.a33
    end

end

# Used when computing matrix exponential
function Base.copyto!(dst::BlochMcConnellMatrix{T,N}, src::AbstractMatrix) where {T,N}

    for j = 1:N, i = 1:N
        b = getblock(dst, i, j)
        b.a11 = src[3i-2,3j-2]
        b.a21 = src[3i-1,3j-2]
        b.a31 = src[3i  ,3j-2]
        b.a12 = src[3i-2,3j-1]
        b.a22 = src[3i-1,3j-1]
        b.a32 = src[3i  ,3j-1]
        b.a13 = src[3i-2,3j]
        b.a23 = src[3i-1,3j]
        b.a33 = src[3i  ,3j]
    end

end

function Base.Matrix(A::BlochMcConnellMatrix{T,N}) where {T,N}

    mat = Matrix{T}(undef, 3N, 3N)
    for j = 1:N, i = 1:N

        b = getblock(A, i, j)
        mat[3i-2,3j-2] = b.a11
        mat[3i-1,3j-2] = b.a21
        mat[3i  ,3j-2] = b.a31
        mat[3i-2,3j-1] = b.a12
        mat[3i-1,3j-1] = b.a22
        mat[3i  ,3j-1] = b.a32
        mat[3i-2,3j]   = b.a13
        mat[3i-1,3j]   = b.a23
        mat[3i  ,3j]   = b.a33

    end

    return mat

end

function Base.eltype(
    ::Union{<:AbstractBlochMatrix{T},<:AbstractBlochMcConnellMatrix{T,N}}
) where {T,N}

    return T

end

function Base.:(==)(
    A::Union{<:AbstractBlochMatrix,<:AbstractBlochMcConnellMatrix},
    B::Union{<:AbstractBlochMatrix,<:AbstractBlochMcConnellMatrix}
)

    return Matrix(A) == Matrix(B)

end

function Base.isapprox(
    A::Union{<:AbstractBlochMatrix,<:AbstractBlochMcConnellMatrix},
    B::Union{<:AbstractBlochMatrix,<:AbstractBlochMcConnellMatrix};
    kwargs...
)

    return isapprox(Matrix(A), Matrix(B); kwargs...)

end

"""
    ExcitationMatrix(A)
    ExcitationMatrix{T}()
    ExcitationMatrix()

Create an `ExcitationMatrix` object. Multiplying by a `MagnetizationMC` object
has the effect of multiplying each component of the multi-compartment
magnetization by the `ExcitationMatrix`.

# Properties
- `A::BlochMatrix{<:Real}`: Matrix used to describe the excitation

# Examples
```jldoctest
julia> E = ExcitationMatrix(BlochMatrix(0, 1, 0, 1, 0, 0, 0, 0, 1))
ExcitationMatrix{Int64}:
 0  1  0
 1  0  0
 0  0  1

julia> M = MagnetizationMC((1, 2, 3), (4, 5, 6))
2-compartment Magnetization vector with eltype Int64:
 Compartment 1:
  Mx = 1
  My = 2
  Mz = 3
 Compartment 2:
  Mx = 4
  My = 5
  Mz = 6

julia> E * M
2-compartment Magnetization vector with eltype Int64:
 Compartment 1:
  Mx = 2
  My = 1
  Mz = 3
 Compartment 2:
  Mx = 5
  My = 4
  Mz = 6
```
"""
struct ExcitationMatrix{T<:Real}
    A::BlochMatrix{T}
end

ExcitationMatrix{T}() where {T} = ExcitationMatrix(BlochMatrix{T}())
ExcitationMatrix() = ExcitationMatrix(BlochMatrix())

function Base.show(io::IO, ::MIME"text/plain", A::ExcitationMatrix{T}) where {T}

    print(io, "ExcitationMatrix{$T}:\n")
    haskey(io, :compact) || (io = IOContext(io, :compact => true))
    Base.print_array(io, Matrix(A.A))

end

"""
    idealspoiling = IdealSpoilingMatrix()

Matrix representing ideal spoiling. Multiplying by a `Magnetization` or
`MagnetizationMC` has the effect of setting the x and y components to 0.

# Examples
```jldoctest
julia> idealspoiling * MagnetizationMC((1, 1, 1), (2, 2, 2))
2-compartment Magnetization vector with eltype Int64:
 Compartment 1:
  Mx = 0
  My = 0
  Mz = 1
 Compartment 2:
  Mx = 0
  My = 0
  Mz = 2
```
"""
struct IdealSpoilingMatrix end
const idealspoiling = IdealSpoilingMatrix()

# ==============================================================================
# Math operations
# ==============================================================================

# +
Base.:+(M1::Magnetization, M2::Magnetization) = Magnetization(M1.x + M2.x, M1.y + M2.y, M1.z + M2.z)
Base.:+(M1::MagnetizationMC{T1,N}, M2::MagnetizationMC{T2,N}) where {T1,T2,N} = MagnetizationMC(ntuple(i -> M1[i] + M2[i], N)...)

# -
Base.:-(M1::Magnetization, M2::Magnetization) = Magnetization(M1.x - M2.x, M1.y - M2.y, M1.z - M2.z)
Base.:-(M1::MagnetizationMC{T1,N}, M2::MagnetizationMC{T2,N}) where {T1,T2,N} = MagnetizationMC(ntuple(i -> M1[i] - M2[i], N)...)

function Base.:-(::UniformScaling, B::BlochMatrix{T}) where {T}

    return BlochMatrix(
        one(T) - B.a11,
        -B.a21,
        -B.a31,
        -B.a12,
        one(T) - B.a22,
        -B.a32,
        -B.a13,
        -B.a23,
        one(T) - B.a33
    )

end

# *
Base.:*(M::Magnetization, a) = Magnetization(M.x * a, M.y * a, M.z * a)
Base.:*(M::MagnetizationMC{T,N}, a) where {T,N} = MagnetizationMC(ntuple(i -> M[i] * a, N)...)

function Base.:*(A::BlochMatrix, M::Magnetization)

    return Magnetization(
        A.a11 * M.x + A.a12 * M.y + A.a13 * M.z,
        A.a21 * M.x + A.a22 * M.y + A.a23 * M.z,
        A.a31 * M.x + A.a32 * M.y + A.a33 * M.z
    )

end

function Base.:*(A::FreePrecessionMatrix, M::Magnetization)

    return Magnetization(
        A.E2cosθ * M.x + A.E2sinθ * M.y,
        A.E2cosθ * M.y - A.E2sinθ * M.x,
        A.E1 * M.z
    )

end

Base.:*(A::ExcitationMatrix, M::Magnetization) = A.A * M
Base.:*(::IdealSpoilingMatrix, M::Magnetization{T}) where {T} = Magnetization(zero(T), zero(T), M.z)

function Base.:*(A::BlochMcConnellMatrix{T1,N}, M::MagnetizationMC{T2,N}) where {T1,T2,N}

    return MagnetizationMC(ntuple(N) do i
        Mc = getblock(A, i, 1) * M[1]
        for j = 2:N
            muladd!(Mc, getblock(A, i, j), M[j])
        end
        Mc
    end...)

end

Base.:*(A::ExcitationMatrix, M::MagnetizationMC) = MagnetizationMC((A * Mc for Mc in M)...)
Base.:*(A::IdealSpoilingMatrix, M::MagnetizationMC) = MagnetizationMC((A * Mc for Mc in M)...)

function Base.:*(A::BlochMatrix, B::BlochMatrix)

    return BlochMatrix(
        A.a11 * B.a11 + A.a12 * B.a21 + A.a13 * B.a31,
        A.a21 * B.a11 + A.a22 * B.a21 + A.a23 * B.a31,
        A.a31 * B.a11 + A.a32 * B.a21 + A.a33 * B.a31,
        A.a11 * B.a12 + A.a12 * B.a22 + A.a13 * B.a32,
        A.a21 * B.a12 + A.a22 * B.a22 + A.a23 * B.a32,
        A.a31 * B.a12 + A.a32 * B.a22 + A.a33 * B.a32,
        A.a11 * B.a13 + A.a12 * B.a23 + A.a13 * B.a33,
        A.a21 * B.a13 + A.a22 * B.a23 + A.a23 * B.a33,
        A.a31 * B.a13 + A.a32 * B.a23 + A.a33 * B.a33
    )

end

function Base.:*(A::BlochMatrix, B::FreePrecessionMatrix)

    return BlochMatrix(
        A.a11 * B.E2cosθ - A.a12 * B.E2sinθ,
        A.a21 * B.E2cosθ - A.a22 * B.E2sinθ,
        A.a31 * B.E2cosθ - A.a32 * B.E2sinθ,
        A.a11 * B.E2sinθ + A.a12 * B.E2cosθ,
        A.a21 * B.E2sinθ + A.a22 * B.E2cosθ,
        A.a31 * B.E2sinθ + A.a32 * B.E2cosθ,
        A.a13 * B.E1,
        A.a23 * B.E1,
        A.a33 * B.E1
    )

end

function Base.:*(A::BlochMatrix, ::IdealSpoilingMatrix)

    T = eltype(A)
    return BlochMatrix(zero(T), zero(T), zero(T), zero(T), zero(T), zero(T), A.a13, A.a23, A.a33)

end

Base.:*(A::ExcitationMatrix, ::IdealSpoilingMatrix) = A.A * idealspoiling

function Base.:*(A::BlochMcConnellMatrix{T1,N}, B::BlochMcConnellMatrix{T2,N}) where {T1,T2,N}

    C = BlochMcConnellMatrix{promote_type(T1, T2)}(N)
    mul!(C, A, B)
    return C

end

function Base.:*(A::ExcitationMatrix{T1}, B::BlochMcConnellMatrix{T2,N}) where {T1,T2,N}

    C = BlochMcConnellMatrix{promote_type(T1, T2)}(N)
    mul!(C, A, B)
    return C

end

function Base.:*(A::IdealSpoilingMatrix, B::BlochMcConnellMatrix{T,N}) where {T,N}

    C = BlochMcConnellMatrix{T}(N)
    mul!(C, A, B)
    return C

end

# /
Base.:/(M::Magnetization, a) = Magnetization(M.x / a, M.y / a, M.z / a)
Base.:/(M::MagnetizationMC{T,N}, a) where {T,N} = MagnetizationMC(ntuple(i -> M[i] / a, N)...)

# \
Base.:\(A::BlochMatrix, M::Magnetization) = Magnetization((Matrix(A) \ Vector(M))...)

function Base.:\(A::BlochMcConnellMatrix{T1,N}, M::MagnetizationMC{T2,N}) where {T1,T2,N}

    result = Matrix(A) \ Vector(M)
    return MagnetizationMC(ntuple(N) do i
        @inbounds r1 = result[3i-2]
        @inbounds r2 = result[3i-1]
        @inbounds r3 = result[3i]
        Magnetization(r1, r2, r3)
    end...)

end

"""
    add!(A, B)

Compute `A + B`, storing the result in `A`.
"""
function add!(M1::Magnetization, M2::Magnetization)

    M1.x += M2.x
    M1.y += M2.y
    M1.z += M2.z
    return nothing

end

function add!(M1::MagnetizationMC{T1,N}, M2::MagnetizationMC{T2,N}) where {T1,T2,N}

    for i = 1:N
        add!(M1[i], M2[i])
    end

end

"""
    add!(C, A, B)

Compute `A + B`, storing the result in `C`.
"""
function add!(C::AbstractVector, M1::Magnetization, M2::Magnetization)

    C[1] = M1.x + M2.x
    C[2] = M1.y + M2.y
    C[3] = M1.z + M2.z
    return nothing

end

function add!(C::AbstractVector, M1::MagnetizationMC{T1,N}, M2::MagnetizationMC{T2,N}) where {T1,T2,N}

    for i = 1:N
        add!(view(C, 3i-2:3i), M1[i], M2[i])
    end

end

function add!(C::AbstractMatrix, A::BlochMcConnellMatrix{T1,N}, B::BlochMcConnellMatrix{T2,N}) where {T1,T2,N}

    for j = 1:N, i = 1:N
        Ab = getblock(A, i, j)
        Bb = getblock(B, i, j)
        C[3i-2,3j-2] = Ab.a11 + Bb.a11
        C[3i-1,3j-2] = Ab.a21 + Bb.a21
        C[3i  ,3j-2] = Ab.a31 + Bb.a31
        C[3i-2,3j-1] = Ab.a12 + Bb.a12
        C[3i-1,3j-1] = Ab.a22 + Bb.a22
        C[3i  ,3j-1] = Ab.a32 + Bb.a32
        C[3i-2,3j]   = Ab.a13 + Bb.a13
        C[3i-1,3j]   = Ab.a23 + Bb.a23
        C[3i  ,3j]   = Ab.a33 + Bb.a33
    end

end

"""
    subtract!(A, B)

Compute `A - B`, storing the result in `A`.
"""
function subtract!(M1::Magnetization, M2::Magnetization)

    M1.x -= M2.x
    M1.y -= M2.y
    M1.z -= M2.z
    return nothing

end

function subtract!(M1::MagnetizationMC{T1,N}, M2::MagnetizationMC{T2,N}) where {T1,T2,N}

    for i = 1:N
        subtract!(M1[i], M2[i])
    end

end

"""
    subtract!(C, A, B)

Compute `A - B`, storing the result in `C`.
"""
function subtract!(C::AbstractMatrix{T}, ::UniformScaling, B::BlochMatrix) where {T}

    C[1,1] = one(T) - B.a11
    C[2,1] = -B.a21
    C[3,1] = -B.a31
    C[1,2] = -B.a12
    C[2,2] = one(T) - B.a22
    C[3,2] = -B.a32
    C[1,3] = -B.a13
    C[2,3] = -B.a23
    C[3,3] = one(T) - B.a33
    return nothing

end

function subtract!(C::AbstractMatrix, A::BlochMcConnellMatrix{T1,N}, B::BlochMcConnellMatrix{T2,N}) where {T1,T2,N}

    for j = 1:N, i = 1:N
        Ab = getblock(A, i, j)
        Bb = getblock(B, i, j)
        C[3i-2,3j-2] = Ab.a11 - Bb.a11
        C[3i-1,3j-2] = Ab.a21 - Bb.a21
        C[3i  ,3j-2] = Ab.a31 - Bb.a31
        C[3i-2,3j-1] = Ab.a12 - Bb.a12
        C[3i-1,3j-1] = Ab.a22 - Bb.a22
        C[3i  ,3j-1] = Ab.a32 - Bb.a32
        C[3i-2,3j]   = Ab.a13 - Bb.a13
        C[3i-1,3j]   = Ab.a23 - Bb.a23
        C[3i  ,3j]   = Ab.a33 - Bb.a33
    end

end

function subtract!(C::AbstractMatrix{T1}, A::UniformScaling, B::BlochMcConnellMatrix{T2,N}) where {T1,T2,N}

    for j = 1:N, i = 1:N
        b = getblock(B, i, j)
        if i == j
            C[3i-2,3j-2] = one(T1) - b.a11
            C[3i-1,3j-2] = -b.a21
            C[3i  ,3j-2] = -b.a31
            C[3i-2,3j-1] = -b.a12
            C[3i-1,3j-1] = one(T1) - b.a22
            C[3i  ,3j-1] = -b.a32
            C[3i-2,3j]   = -b.a13
            C[3i-1,3j]   = -b.a23
            C[3i  ,3j]   = one(T1) - b.a33
        else
            C[3i-2,3j-2] = -b.a11
            C[3i-1,3j-2] = -b.a21
            C[3i  ,3j-2] = -b.a31
            C[3i-2,3j-1] = -b.a12
            C[3i-1,3j-1] = -b.a22
            C[3i  ,3j-1] = -b.a32
            C[3i-2,3j]   = -b.a13
            C[3i-1,3j]   = -b.a23
            C[3i  ,3j]   = -b.a33
        end
    end

end

"""
    mul!(A, a)

Compute `A * a`, storing the result in `A`.
"""
LinearAlgebra.mul!(M::Magnetization, a) = (M.x *= a; M.y *= a; M.z *= a; nothing)
LinearAlgebra.mul!(M::MagnetizationMC{T,N}, a) where {T,N} = foreach(M -> mul!(M, a), M)

function LinearAlgebra.mul!(A::BlochDynamicsMatrix, t::Real)

    A.R1 *= t
    A.R2 *= t
    A.Δω *= t
    return nothing

end

function LinearAlgebra.mul!(E::ExchangeDynamicsMatrix, t::Real)

    E.r *= t
    return nothing

end

function LinearAlgebra.mul!(A::BlochMcConnellDynamicsMatrix, t::Real)

    for A in A.A
        mul!(A, t)
    end
    for E in A.E
        mul!(E, t)
    end

end

"""
    mul!(C, A, B)

Compute `A * B`, storing the result in `C`.
"""
function LinearAlgebra.mul!(M2::Magnetization, A::BlochMatrix, M1::Magnetization)

    M2.x = A.a11 * M1.x + A.a12 * M1.y + A.a13 * M1.z
    M2.y = A.a21 * M1.x + A.a22 * M1.y + A.a23 * M1.z
    M2.z = A.a31 * M1.x + A.a32 * M1.y + A.a33 * M1.z
    return nothing

end

function LinearAlgebra.mul!(M2::Magnetization, A::FreePrecessionMatrix, M1::Magnetization)

    M2.x = A.E2cosθ * M1.x + A.E2sinθ * M1.y
    M2.y = A.E2cosθ * M1.y - A.E2sinθ * M1.x
    M2.z = A.E1 * M1.z
    return nothing

end

function LinearAlgebra.mul!(M2::Magnetization, A::ExcitationMatrix, M1::Magnetization)

    mul!(M2, A.A, M1)

end

function LinearAlgebra.mul!(M2::Magnetization{T}, ::IdealSpoilingMatrix, M1::Magnetization) where {T}

    M2.x = zero(T)
    M2.y = zero(T)
    M2.z = M1.z
    return nothing

end

function LinearAlgebra.mul!(M2::Vector, A::BlochMatrix, M1::Magnetization)

    M2[1] = A.a11 * M1.x + A.a12 * M1.y + A.a13 * M1.z
    M2[2] = A.a21 * M1.x + A.a22 * M1.y + A.a23 * M1.z
    M2[3] = A.a31 * M1.x + A.a32 * M1.y + A.a33 * M1.z
    return nothing

end

function LinearAlgebra.mul!(M2::Vector, A::ExcitationMatrix, M1::Magnetization)

    mul!(M2, A.A, M1)

end

function LinearAlgebra.mul!(M2::MagnetizationMC{T1,N}, A::BlochMcConnellMatrix{T2,N}, M1::MagnetizationMC{T3,N}) where {T1,T2,T3,N}

    for i = 1:N
        M = M2[i]
        mul!(M, getblock(A, i, 1), M1[1])
        for j = 2:N
            muladd!(M, getblock(A, i, j), M1[j])
        end
    end

end

function LinearAlgebra.mul!(M2::MagnetizationMC{T1,N}, A::ExcitationMatrix, M1::MagnetizationMC{T2,N}) where {T1,T2,N}

    for i = 1:N
        mul!(M2[i], A, M1[i])
    end

end

function LinearAlgebra.mul!(M2::MagnetizationMC{T1,N}, ::IdealSpoilingMatrix, M1::MagnetizationMC{T2,N}) where {T1,T2,N}

    for i = 1:N
        mul!(M2[i], idealspoiling, M1[i])
    end

end

function LinearAlgebra.mul!(C::BlochMatrix, A::BlochMatrix, t::Real)

    C.a11 = A.a11 * t
    C.a21 = A.a21 * t
    C.a31 = A.a31 * t
    C.a12 = A.a12 * t
    C.a22 = A.a22 * t
    C.a32 = A.a32 * t
    C.a13 = A.a13 * t
    C.a23 = A.a23 * t
    C.a33 = A.a33 * t
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix, A::BlochMatrix, B::BlochMatrix)

    C.a11 = A.a11 * B.a11 + A.a12 * B.a21 + A.a13 * B.a31
    C.a21 = A.a21 * B.a11 + A.a22 * B.a21 + A.a23 * B.a31
    C.a31 = A.a31 * B.a11 + A.a32 * B.a21 + A.a33 * B.a31
    C.a12 = A.a11 * B.a12 + A.a12 * B.a22 + A.a13 * B.a32
    C.a22 = A.a21 * B.a12 + A.a22 * B.a22 + A.a23 * B.a32
    C.a32 = A.a31 * B.a12 + A.a32 * B.a22 + A.a33 * B.a32
    C.a13 = A.a11 * B.a13 + A.a12 * B.a23 + A.a13 * B.a33
    C.a23 = A.a21 * B.a13 + A.a22 * B.a23 + A.a23 * B.a33
    C.a33 = A.a31 * B.a13 + A.a32 * B.a23 + A.a33 * B.a33
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix, A::BlochMatrix, B::FreePrecessionMatrix)

    C.a11 = A.a11 * B.E2cosθ - A.a12 * B.E2sinθ
    C.a21 = A.a21 * B.E2cosθ - A.a22 * B.E2sinθ
    C.a31 = A.a31 * B.E2cosθ - A.a32 * B.E2sinθ
    C.a12 = A.a11 * B.E2sinθ + A.a12 * B.E2cosθ
    C.a22 = A.a21 * B.E2sinθ + A.a22 * B.E2cosθ
    C.a32 = A.a31 * B.E2sinθ + A.a32 * B.E2cosθ
    C.a13 = A.a13 * B.E1
    C.a23 = A.a23 * B.E1
    C.a33 = A.a33 * B.E1
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix, A::FreePrecessionMatrix, B::BlochMatrix)

    C.a11 = A.E2cosθ * B.a11 + A.E2sinθ * B.a21
    C.a21 = A.E2cosθ * B.a21 - A.E2sinθ * B.a11
    C.a31 = A.E1 * B.a31
    C.a12 = A.E2cosθ * B.a12 + A.E2sinθ * B.a22
    C.a22 = A.E2cosθ * B.a22 - A.E2sinθ * B.a12
    C.a32 = A.E1 * B.a32
    C.a13 = A.E2cosθ * B.a13 + A.E2sinθ * B.a23
    C.a23 = A.E2cosθ * B.a23 - A.E2sinθ * B.a13
    C.a33 = A.E1 * B.a33
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix, A::BlochDynamicsMatrix, B::BlochMatrix)

    C.a11 = A.R2 * B.a11 + A.Δω * B.a21
    C.a21 = A.R2 * B.a21 - A.Δω * B.a11
    C.a31 = A.R1 * B.a31
    C.a12 = A.R2 * B.a12 + A.Δω * B.a22
    C.a22 = A.R2 * B.a22 - A.Δω * B.a12
    C.a32 = A.R1 * B.a32
    C.a13 = A.R2 * B.a13 + A.Δω * B.a23
    C.a23 = A.R2 * B.a23 - A.Δω * B.a13
    C.a33 = A.R1 * B.a33
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix, A::ExchangeDynamicsMatrix, B::BlochMatrix)

    C.a11 = A.r * B.a11
    C.a21 = A.r * B.a21
    C.a31 = A.r * B.a31
    C.a12 = A.r * B.a12
    C.a22 = A.r * B.a22
    C.a32 = A.r * B.a32
    C.a13 = A.r * B.a13
    C.a23 = A.r * B.a23
    C.a33 = A.r * B.a33
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix, A::ExcitationMatrix, B::BlochMatrix)

    mul!(C, A.A, B)

end

function LinearAlgebra.mul!(C::BlochMatrix{T}, ::IdealSpoilingMatrix, B::BlochMatrix) where {T}

    C.a11 = zero(T)
    C.a21 = zero(T)
    C.a31 = B.a31
    C.a12 = zero(T)
    C.a22 = zero(T)
    C.a32 = B.a32
    C.a13 = zero(T)
    C.a23 = zero(T)
    C.a33 = B.a33
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix{T}, A::FreePrecessionMatrix, B::FreePrecessionMatrix) where {T}

    C.a11 = A.E2cosθ * B.E2cosθ - A.E2sinθ * B.E2sinθ
    C.a12 = A.E2cosθ * B.E2sinθ + A.E2sinθ * B.E2cosθ
    C.a13 = zero(T)
    C.a21 = -C.a12
    C.a22 = C.a11
    C.a23 = zero(T)
    C.a31 = zero(T)
    C.a32 = zero(T)
    C.a33 = A.E1 * B.E1
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix, A::FreePrecessionMatrix, B::ExcitationMatrix)

    mul!(C, A, B.A)

end

function LinearAlgebra.mul!(C::BlochMatrix, A::ExcitationMatrix, B::FreePrecessionMatrix)

    mul!(C, A.A, B)

end

function LinearAlgebra.mul!(C::BlochMatrix{T}, ::IdealSpoilingMatrix, B::FreePrecessionMatrix) where {T}

    C.a11 = zero(T)
    C.a21 = zero(T)
    C.a31 = zero(T)
    C.a12 = zero(T)
    C.a22 = zero(T)
    C.a32 = zero(T)
    C.a13 = zero(T)
    C.a23 = zero(T)
    C.a33 = B.E1
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix{T}, A::BlochDynamicsMatrix, B::BlochDynamicsMatrix) where {T}

    C.a11 = A.R2 * B.R2 - A.Δω * B.Δω
    C.a21 = -A.Δω * B.R2 - A.R2 * B.Δω
    C.a31 = zero(T)
    C.a12 = -C.a21
    C.a22 = C.a11
    C.a32 = zero(T)
    C.a13 = zero(T)
    C.a23 = zero(T)
    C.a33 = A.R1 * B.R1
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix{T}, A::BlochDynamicsMatrix, B::ExchangeDynamicsMatrix) where {T}

    C.a11 = A.R2 * B.r
    C.a21 = -A.Δω * B.r
    C.a31 = zero(T)
    C.a12 = A.Δω * B.r
    C.a22 = A.R2 * B.r
    C.a32 = zero(T)
    C.a13 = zero(T)
    C.a23 = zero(T)
    C.a33 = A.R1 * B.r
    return nothing

end

function LinearAlgebra.mul!(C::BlochMatrix, A::ExchangeDynamicsMatrix, B::BlochDynamicsMatrix)

    mul!(C, B, A)

end

function LinearAlgebra.mul!(C::BlochMatrix{T}, A::ExchangeDynamicsMatrix, B::ExchangeDynamicsMatrix) where {T}

    C.a11 = A.r * B.r
    C.a21 = zero(T)
    C.a31 = zero(T)
    C.a12 = zero(T)
    C.a22 = A.r * B.r
    C.a32 = zero(T)
    C.a13 = zero(T)
    C.a23 = zero(T)
    C.a33 = A.r * B.r
    return nothing

end

function LinearAlgebra.mul!(C::BlochMcConnellMatrix{T1,N}, A::BlochMcConnellMatrix{T2,N}, t::Real) where {T1,T2,N}

    for i = 1:N, j = 1:N
        mul!(getblock(C, i, j), getblock(A, i, j), t)
    end

end

function LinearAlgebra.mul!(
    C::BlochMcConnellMatrix{T1,N},
    A::ExcitationMatrix,
    B::BlochMcConnellMatrix{T2,N}
) where {T1,T2,N}

    for j = 1:N, i = 1:N
        mul!(getblock(C, i, j), A.A, getblock(B, i, j))
    end

end

function LinearAlgebra.mul!(
    C::BlochMcConnellMatrix{T1,N},
    ::IdealSpoilingMatrix,
    B::BlochMcConnellMatrix{T2,N}
) where {T1,T2,N}

    for j = 1:N, i = 1:N
        mul!(getblock(C, i, j), idealspoiling, getblock(B, i, j))
    end

end

function LinearAlgebra.mul!(
    C::AbstractBlochMcConnellMatrix{T1,N},
    A::AbstractBlochMcConnellMatrix{T2,N},
    B::AbstractBlochMcConnellMatrix{T3,N}
) where {T1,T2,T3,N}

    for j = 1:N, i = 1:N
        mul!(getblock(C, i, j), getblock(A, i, 1), getblock(B, 1, j))
        for k = 2:N
            muladd!(getblock(C, i, j), getblock(A, i, k), getblock(B, k, j))
        end
    end

end

"""
    div!(A, a)

Compute `A / a`, storing the result in `A`.
"""
div!(M::Magnetization, a) = (M.x /= a; M.y /= a; M.z /= a; nothing)
div!(M::MagnetizationMC{T,N}, a) where {T,N} = foreach(M -> div!(M, a), M)

"""
    muladd!(C, A, B)

Compute `A * B + C`, storing the result in `C`.
"""
function muladd!(M2::Magnetization, A::BlochMatrix, M1::Magnetization)

    M2.x += A.a11 * M1.x + A.a12 * M1.y + A.a13 * M1.z
    M2.y += A.a21 * M1.x + A.a22 * M1.y + A.a23 * M1.z
    M2.z += A.a31 * M1.x + A.a32 * M1.y + A.a33 * M1.z
    return nothing

end

function muladd!(M2::Magnetization, A::FreePrecessionMatrix, M1::Magnetization)

    M2.x += A.E2cosθ * M1.x + A.E2sinθ * M1.y
    M2.y += A.E2cosθ * M1.y - A.E2sinθ * M1.x
    M2.z += A.E1 * M1.z
    return nothing

end

function muladd!(M2::MagnetizationMC{T1,N}, A::BlochMcConnellMatrix{T2,N}, M1::MagnetizationMC{T3,N}) where {T1,T2,T3,N}

    for i = 1:N
        M = M2[i]
        for j = 1:N
            muladd!(M, getblock(A, i, j), M1[j])
        end
    end

end

function muladd!(C::BlochMatrix, A::BlochMatrix, t::Real)

    C.a11 += A.a11 * t
    C.a21 += A.a21 * t
    C.a31 += A.a31 * t
    C.a12 += A.a12 * t
    C.a22 += A.a22 * t
    C.a32 += A.a32 * t
    C.a13 += A.a13 * t
    C.a23 += A.a23 * t
    C.a33 += A.a33 * t
    return nothing

end

function muladd!(C::BlochMatrix, ::UniformScaling, t::Real)

    C.a11 += t
    C.a22 += t
    C.a33 += t
    return nothing

end

function muladd!(C::BlochMatrix, A::BlochMatrix, B::BlochMatrix)

    C.a11 += A.a11 * B.a11 + A.a12 * B.a21 + A.a13 * B.a31
    C.a21 += A.a21 * B.a11 + A.a22 * B.a21 + A.a23 * B.a31
    C.a31 += A.a31 * B.a11 + A.a32 * B.a21 + A.a33 * B.a31
    C.a12 += A.a11 * B.a12 + A.a12 * B.a22 + A.a13 * B.a32
    C.a22 += A.a21 * B.a12 + A.a22 * B.a22 + A.a23 * B.a32
    C.a32 += A.a31 * B.a12 + A.a32 * B.a22 + A.a33 * B.a32
    C.a13 += A.a11 * B.a13 + A.a12 * B.a23 + A.a13 * B.a33
    C.a23 += A.a21 * B.a13 + A.a22 * B.a23 + A.a23 * B.a33
    C.a33 += A.a31 * B.a13 + A.a32 * B.a23 + A.a33 * B.a33
    return nothing

end

function muladd!(C::BlochMatrix, A::BlochDynamicsMatrix, B::BlochMatrix)

    C.a11 += A.R2 * B.a11 + A.Δω * B.a21
    C.a21 += A.R2 * B.a21 - A.Δω * B.a11
    C.a31 += A.R1 * B.a31
    C.a12 += A.R2 * B.a12 + A.Δω * B.a22
    C.a22 += A.R2 * B.a22 - A.Δω * B.a12
    C.a32 += A.R1 * B.a32
    C.a13 += A.R2 * B.a13 + A.Δω * B.a23
    C.a23 += A.R2 * B.a23 - A.Δω * B.a13
    C.a33 += A.R1 * B.a33
    return nothing

end

function muladd!(C::BlochMatrix, A::ExchangeDynamicsMatrix, B::BlochMatrix)

    C.a11 += A.r * B.a11
    C.a21 += A.r * B.a21
    C.a31 += A.r * B.a31
    C.a12 += A.r * B.a12
    C.a22 += A.r * B.a22
    C.a32 += A.r * B.a32
    C.a13 += A.r * B.a13
    C.a23 += A.r * B.a23
    C.a33 += A.r * B.a33
    return nothing

end

function muladd!(C::BlochMatrix, A::BlochDynamicsMatrix, B::BlochDynamicsMatrix)

    C.a11 += A.R2 * B.R2 - A.Δω * B.Δω
    C.a21 -= A.Δω * B.R2 + A.R2 * B.Δω
    C.a12 += A.R2 * B.Δω + A.Δω * B.R2
    C.a22 += A.R2 * B.R2 - A.Δω * B.Δω
    C.a33 += A.R1 * B.R1
    return nothing

end

function muladd!(C::BlochMatrix, A::BlochDynamicsMatrix, B::ExchangeDynamicsMatrix)

    C.a11 += A.R2 * B.r
    C.a21 -= A.Δω * B.r
    C.a12 += A.Δω * B.r
    C.a22 += A.R2 * B.r
    C.a33 += A.R1 * B.r
    return nothing

end

function muladd!(C::BlochMatrix, A::ExchangeDynamicsMatrix, B::BlochDynamicsMatrix)

    muladd!(C, B, A)

end

function muladd!(C::BlochMatrix, A::ExchangeDynamicsMatrix, B::ExchangeDynamicsMatrix)

    C.a11 += A.r * B.r
    C.a22 += A.r * B.r
    C.a33 += A.r * B.r
    return nothing

end

function muladd!(C::BlochMcConnellMatrix{T1,N}, A::BlochMcConnellMatrix{T2,N}, t::Real) where {T1,T2,N}

    for i = 1:N, j = 1:N
        muladd!(getblock(C, i, j), getblock(A, i, j), t)
    end

end

function muladd!(C::BlochMcConnellMatrix{T,N}, I::UniformScaling, t::Real) where {T,N}

    for i = 1:N
        muladd!(getblock(C, i, i), I, t)
    end

end

"""
    subtractmul!(C, X, A, B)

Compute `(X - A) * B`, storing the result in `C`.
"""
function subtractmul!(
    M2::Magnetization,
    ::Nothing,
    A::BlochMatrix,
    M1::Magnetization
)

    M2.x = -A.a11 * M1.x - A.a12 * M1.y - A.a13 * M1.z
    M2.y = -A.a21 * M1.x - A.a22 * M1.y - A.a23 * M1.z
    M2.z = -A.a31 * M1.x - A.a32 * M1.y - A.a33 * M1.z
    return nothing

end

function subtractmul!(
    M2::Magnetization,
    ::UniformScaling,
    A::BlochMatrix,
    M1::Magnetization
)

    M2.x =  (1 - A.a11) * M1.x - A.a12 * M1.y - A.a13 * M1.z
    M2.y = -A.a21 * M1.x + (1 - A.a22) * M1.y - A.a23 * M1.z
    M2.z = -A.a31 * M1.x - A.a32 * M1.y + (1 - A.a33) * M1.z
    return nothing

end

function subtractmul!(
    M2::MagnetizationMC{T1,N},
    I::UniformScaling,
    A::BlochMcConnellMatrix{T2,N},
    M1::MagnetizationMC{T3,N}
) where {T1,T2,T3,N}

    for i = 1:N
        M = M2[i]
        if i == 1
            subtractmul!(M, I, getblock(A, i, 1), M1[1])
        else
            subtractmul!(M, nothing, getblock(A, i, 1), M1[1])
        end
        for j = 2:N
            if j == i
                subtractmuladd!(M, I, getblock(A, i, j), M1[j])
            else
                subtractmuladd!(M, nothing, getblock(A, i, j), M1[j])
            end
        end
    end

end

"""
    subtractmuladd!(C, X, A, B)

Compute `(X - A) * B + C`, storing the result in `C`.
"""
function subtractmuladd!(
    M2::Magnetization,
    ::Nothing,
    A::BlochMatrix,
    M1::Magnetization
)

    M2.x -= A.a11 * M1.x + A.a12 * M1.y + A.a13 * M1.z
    M2.y -= A.a21 * M1.x + A.a22 * M1.y + A.a23 * M1.z
    M2.z -= A.a31 * M1.x + A.a32 * M1.y + A.a33 * M1.z
    return nothing

end

function subtractmuladd!(
    M2::Magnetization,
    ::UniformScaling,
    A::BlochMatrix,
    M1::Magnetization
)

    M2.x +=  (1 - A.a11) * M1.x - A.a12 * M1.y - A.a13 * M1.z
    M2.y += -A.a21 * M1.x + (1 - A.a22) * M1.y - A.a23 * M1.z
    M2.z += -A.a31 * M1.x - A.a32 * M1.y + (1 - A.a33) * M1.z
    return nothing

end

"""
    absolutesum(A)

Compute the sum of the absolute values of the elements of `A`.
"""
function absolutesum(A::BlochDynamicsMatrix)

    return 2 * abs(A.R2) + 2 * abs(A.Δω) + abs(A.R1)

end

function absolutesum(E::ExchangeDynamicsMatrix)

    return 3 * abs(E.r)

end

function absolutesum(A::BlochMcConnellDynamicsMatrix{T,N,M}) where {T,N,M}

    result = absolutesum(A.A[1])
    for i = 2:N
        result += absolutesum(A.A[i])
    end
    for E in A.E
        result += absolutesum(E)
    end
    return result

end


import Entropies: SymbolicPermutation, CountOccurrences, symbolize, symbolize!
export SymbolicPermutation

function transferentropy(s, t, est::SymbolicPermutation;
        base = 2, q = 1, τT = -1, τS = -1, η𝒯 = 1, dT = 1, dS = 1, d𝒯 = 1)

    # Generalized embeddings of each marginal separately, using the same dimension for each
    # marginal. Then symbolize the embedded time series, which gives integer symbol sequences.
    symb_s = symbolize(s, est)
    symb_t = symbolize(t, est)

    # After symbolization, use the general interface with occurrence frequency entropy
    # estimator (probabilities are obtained from histograms of the symbols).
    transferentropy(symb_s, symb_t, CountOccurrences(),
        τT = τT, τS = τS, η𝒯 = η𝒯, dT = dT, dS = dS, d𝒯 = d𝒯)
end

function transferentropy!(symb_s, symb_t, s, t, est::SymbolicPermutation; base = 2, q = 1,
        τT = -1, τS = -1, η𝒯 = 1, dT = 1, dS = 1, d𝒯 = 1)

    symbolize!(symb_s, s, est)
    symbolize!(symb_t, t, est)

    transferentropy(symb_s, symb_t, CountOccurrences(),
        τT = τT, τS = τS, η𝒯 = η𝒯, dT = dT, dS = dS, d𝒯 = d𝒯)
end

function transferentropy(s, t, c, est::SymbolicPermutation; base = 2, q = 1,
        τT = -1, τS = -1, τC = -1, η𝒯 = 1, dT = 1, dS = 1, d𝒯 = 1, dC = 1)

    symb_s = symbolize(s, est)
    symb_t = symbolize(t, est)
    symb_c = symbolize(c, est)

    transferentropy(symb_s, symb_t, symb_c, CountOccurrences(),
        τT = τT, τS = τS, τC = τC, η𝒯 = η𝒯, dT = dT, dS = dS, d𝒯 = d𝒯, dC = dC)
end

function transferentropy!(symb_s, symb_t, symb_c, s, t, c, est::SymbolicPermutation;
        base = 2, q = 1,
        τT = -1, τS = -1, τC = -1, η𝒯 = 1, dT = 1, dS = 1, d𝒯 = 1, dC = 1)

    symbolize!(symb_s, s, est)
    symbolize!(symb_t, t, est)
    symbolize!(symb_c, c, est)

    transferentropy(symb_s, symb_t, symb_c, CountOccurrences(),
        τT = τT, τS = τS, τC = τC, η𝒯 = η𝒯, dT = dT, dS = dS, d𝒯 = d𝒯, dC = dC)
end


function isvalidcircle(i,j,r,m,n,options)
	return (options.bounds_row[1] < i/m < options.bounds_row[2]) &&
		(options.bounds_column[1] < j/n < options.bounds_column[2]) &&
		(options.bounds_radius[1] < r/m < options.bounds_radius[2])
end

"""
	initvals(img,thresh=0.5)
Try to find good initial guesses for the cornea position in image `img`, which should be full color. The `thresh` parameter is the fraction of the max in the green channel intensity that a pixel should have to be considered part of the sclera. Returns a vector of zero to two tuples of (center_i,center_j,radius)
"""
function initvals(X::AbstractMatrix{T} where T <: AbstractRGB,purk=[];options=get_defaults())
	m,n = size(X)
	if isempty(purk)
		purk = findpurkinje(X;options)
	end

	# vertically, use the purkinje if possible
	i_c = length(purk) > 50 ? median(i[1] for i in purk) : m/2

	# horizontal guess 1: use the sclera's brightness
	G = green.(X)
	G[purk] .= 0  # ignore the purkinje

	idx = findall(G.>0.5*maximum(G))
	## get distribution peaks
	j_c = r_c = jleft = jright = wleft = wright = 0
	try
		(jleft,wleft),(jright,wright) = findpeaks(idx,(m,n),2)
		@debug "jleft,jright = $jleft,$jright"
		j_c = (jleft+jright)/2
		r_c = (jright-jleft)/2
	catch
	end

	## reject if the radius is unrealistic
	@debug "init1 = ($i_c,$j_c,$r_c)"
	guess = isvalidcircle(i_c,j_c,r_c,m,n,options) ? [(i_c,j_c,r_c)] : []

	# horizontal guess 2: use the purkinje location heuristically
	if length(purk) > 50
		j_p = median(i[2] for i in purk)
		γ = options.purkinje_location
		if wleft > wright
			r_c = (jright - j_p)/(1-γ)
			j_c = jright - r_c
		else
			r_c = (j_p-jleft)/(1+γ)
			j_c = jleft + r_c
		end
		@debug "init2 = ($i_c,$j_c,$r_c)"
		if isvalidcircle(i_c,j_c,r_c,m,n,options)
			push!(guess,(i_c,j_c,r_c))
		end
	end

	return guess
end

"""
	totalgrad(icen,jcen,r,Z[,trange])
Computes the sum of the radial derivative in an image along the specified arc of a circle. `Z` must be a callable function that returns numerical values for real arguments (i.e., allow interpolation). If given, `trange` is a 2-vector defining the range of angles to be used, measured ccw from "straight down" in the usual image visualization (i.e., vertically flipped).
"""
totalgrad(ic,jc,r,Z,θ) = sum(circlegrad(ic,jc,r,Z,θ))
function circlegrad(ic,jc,r,Z,θ)
	g = zeros(size(θ))
	p2c = (r,t) -> (ic+r*cos(t),jc+r*sin(t))
    for (k,t) in enumerate(θ)
		Zplus = Z(p2c(r+2,t)...)
		Zminus = Z(p2c(r-2,t)...)
		g[k] = Zplus - Zminus
	end
    return g
end

"""
	fitcircle(imgfun,i0,j0,r0[,t_range][;options=missing])
Find the circle that optimizes the total r-difference criterion, given initial values for the center and radius. Optionally specify the range of angles to be used (measured from straight down in the image) and options for the `optimize` solver.
"""
function fitcircle(imgfun,m,n,i0,j0,r0,method,θ=[-π,π];options=get_defaults())
	# don't let the center get too close to the edges, or the radius get too small or too close to the nearest edge
	function penalty(u)
		f = (x,lo,hi) -> log( max(1e-12,4(x-lo)*(hi-x)/(hi-lo)^2) )^4
		i,j,r = u
		b1 = options.bounds_row
		b2 = options.bounds_column
		b3 = options.bounds_radius
		return f(i/m,b1...) + f(j/n,b2...) + f(r/m,b3...)
	end

	if length(θ)==2  # given a range
		# space angles at 1 degree, regardless of the size of the range
		θ = filter( t->θ[1] ≤ t ≤ θ[2], 2π*(-180:180)/360)
	end

	objective = u -> -totalgrad(u...,imgfun,θ) + 0.1penalty(u)
	res = optimize(objective,[i0,j0,r0],method,Optim.Options(x_tol = 1e-2,f_tol = 1e-8))
	if !Optim.converged(res)
		@warn "Optimization did not converge"
		#show(res)
	end
	return res.minimizer,res.minimum
end


using Statistics
include("asvar.jl")

function switchingtime(a::Float64,b::Float64,u::Float64=rand())
# generate switching time for rate of the form max(0, a + b s) + c
# under the assumptions that b > 0, c > 0
  if (b > 0)
    if (a < 0)
      return -a/b + switchingtime(0.0, b, u);
    else # a >= 0
      return -a/b + sqrt(a^2/b^2 - 2 * log(u)/b);
    end
  elseif (b == 0) # degenerate case
    if (a < 0)
      return Inf;
    else # a >= 0
      return -log(u)/a;
    end
  else # b <= 0
    if (a <= 0)
      return Inf;
    else # a > 0
      y = -log(u); t1=-a/b;
      if (y >= a * t1 + b *t1^2/2)
        return Inf;
      else
        return -a/b - sqrt(a^2/b^2 + 2 * y /b);
      end
    end
  end
end

function switchingtime_quadratic(a::Float64, b::Float64, c::Float64,u::Float64=rand())
  dummy::Float64 = (-b^3 + 6*a*b*c + 2*(-6*c^2*log(u) + c * sqrt(-3*a^2*b^2 + 16*a^3*c + 6*b^3*log(u) - 36*a*b*c*log(u) + 36*c^2*(log(u))^2)))^(1/3)
  return 1/(2c)*(-b + (b^2-4*a*c)/dummy + dummy)
end

function EstimateMomentAlongTrajectory(t_skeleton::Vector{Float64}, x_skeleton::Vector{Vector{Float64}}, v_skeleton::Vector{Vector{Float64}}; p = 1, ellipticQ::Bool = false, x_ref::Vector{Float64} = Vector{Float64}(undef,0))
  # assume every component follows dynamics x_t = x_ref + (x_0-x_ref)*cos(t) + v_0*sin(t)

  dim = length(x_skeleton[1])

  if ellipticQ && length(x_ref) == 0
    println("WARNING : EstimateMomentAlongTrajectory : No valid reference point provided.")
    x_ref = zeros(dim)
  end
  T = (t_skeleton[end]-t_skeleton[1])
  estimate = 0.0
  for i=2:length(t_skeleton)
    v = v_skeleton[i-1]
    x = x_skeleton[i-1]
    t0 = t_skeleton[i-1]
    t1 = t_skeleton[i]
    Δt = t1 - t0
    if ellipticQ
      if p==1
        estimate = estimate .+ 1/T * (x_ref * Δt + (x - x_ref) * sin(Δt) + v * (1-cos(Δt)))
      elseif p==2
        y = x - x_ref
        estimate = estimate .+ 1/(2*T) * (v .*(4 * x_ref + y) + Δt * (v.^2 + 2 * x_ref.^2 + y.^2) - (4 * x_ref + y * cos(Δt) + v * sin(Δt)).*(v*cos(Δt) - y *sin(Δt)))
      else
        error("Higher order moments (p >2) along Elliptic trajectories currently not implemented.")
      end
    else # assume linear trajectories
      estimate =  estimate .+ 1 ./((p+1)*v_skeleton[i-1] * T) .* ( x_skeleton[i].^(p+1) - x_skeleton[i-1].^(p+1))
    end
  end
  return estimate

end

function EllipticDynamics(t::Real, y0::Vector{Float64}, w0::Vector{Float64})
    # simulate dy/dt = w, dw/dt = -y

    y_new = y0 * cos(t) + w0 * sin(t);
    w_new = -y0 * sin(t) + w0 * cos(t);

    return (y_new, w_new);

end

function SkeletonIntoBatches(t_skeleton::Vector{Float64}, x_skeleton::Vector{Vector{Float64}}, v_skeleton::Vector{Vector{Float64}}, n_batches::Int, ellipticQ::Bool = false, x_ref::Vector{Float64} = Vector{Float64}(undef,0))

  t_skeletons::Vector{Vector{Float64}} = Vector{Vector{Float64}}(undef,0)
  x_skeletons::Vector{Vector{Vector{Float64}}} = Vector{Vector{Vector{Float64}}}(undef,0)
  v_skeletons = similar(x_skeletons)

  if n_batches < 1
    error("n_batches < 1")
  end

  if (ellipticQ && length(x_ref) == 0)
    x_ref = zeros(length(x_skeleton[1]))
  end
  T_max = t_skeleton[end]
  j_max = length(t_skeleton)
  T_batch = T_max / n_batches
  t_current = t_skeleton[1]
  x_current = x_skeleton[1]
  v_current = v_skeleton[1]
  j = 2
  for i = 1:n_batches
    t_skeleton_batch = Vector{Float64}(undef,0)
    x_skeleton_batch = Vector{Vector{Float64}}(undef,0)
    v_skeleton_batch = similar(x_skeleton_batch)
    # add initial position to skeleton
    push!(t_skeleton_batch, t_current)
    push!(x_skeleton_batch, x_current)
    push!(v_skeleton_batch, v_current)

    while (t_skeleton[j] < i*T_batch && j < j_max)
      t_current = t_skeleton[j]
      x_current = x_skeleton[j]
      v_current = v_skeleton[j]
      push!(t_skeleton_batch, t_current)
      push!(x_skeleton_batch, x_current)
      push!(v_skeleton_batch, v_current)
      j = j + 1
    end
    # j = min (j : t_skeleton[j] > i * T_batch)

    # add final position of this batch
    push!(t_skeleton_batch, i*T_batch)
    Δt = (i*T_batch - t_current)
    if  (ellipticQ)
      (y,v_current) = EllipticDynamics(Δt,x_current-x_ref, v_current)
      x_current = x_ref + y
    else
      x_current = x_current + Δt * v_current
    end
    t_current = i*T_batch
    push!(x_skeleton_batch, x_current)
    push!(v_skeleton_batch, v_current)

    # add skeleton to list of skeletons
    push!(t_skeletons, t_skeleton_batch)
    push!(x_skeletons, x_skeleton_batch)
    push!(v_skeletons, v_skeleton_batch)
  end
  return (t_skeletons, x_skeletons, v_skeletons)

end

function BatchMeansPDMP(t_skeleton::Vector{Float64}, x_skeleton::Vector{Vector{Float64}}, v_skeleton::Vector{Vector{Float64}}; n_batches::Int = 100, ellipticQ::Bool = false, x_ref::Vector{Float64} = Vector{Float64}(undef,0))

  dimension = length(x_skeleton[1])
  if (n_batches <= 0)
    error("A valid number of batches is required.")
  end
  T_batch = (t_skeleton[end]-t_skeleton[1])/n_batches
  (t_skeletons, x_skeletons, v_skeletons) = SkeletonIntoBatches(t_skeleton, x_skeleton, v_skeleton, n_batches, ellipticQ, x_ref)
  Z = Array{Float64,2}(undef,dimension,n_batches) # batch means
  for i=1:n_batches
    Z[:,i] = EstimateMomentAlongTrajectory(t_skeletons[i],x_skeletons[i],v_skeletons[i],ellipticQ=ellipticQ,x_ref=x_ref,p=1)
  end
  vec(var(Z, dims=2)*T_batch)
end

function ess_pdmp_components(t_skeleton::Vector{Float64}, x_skeleton::Vector{Vector{Float64}}, v_skeleton::Vector{Vector{Float64}}; n_batches::Int = 50, ellipticQ::Bool = false, x_ref::Vector{Float64} = Vector{Float64}(undef,0))

  asymptotic_variances = BatchMeansPDMP(t_skeleton, x_skeleton, v_skeleton, n_batches = n_batches, ellipticQ = ellipticQ, x_ref = x_ref)
  if any(asymptotic_variances .< 0)
    println("WARNING : ess_pdmp : at least one component of the asymptotic variance is negative.")
    println("asymptotic_variances = ", asymptotic_variances)
  end
  variances = EstimateMomentAlongTrajectory(t_skeleton,x_skeleton,v_skeleton, ellipticQ = ellipticQ, x_ref = x_ref, p=2) .- (EstimateMomentAlongTrajectory(t_skeleton,x_skeleton,v_skeleton, ellipticQ = ellipticQ, x_ref = x_ref).^2)
  if any(variances .< 0)
    println("WARNING : ess_pdmp : at least one component of the variance is negative.")
    println("variances = ", variances)
  end
  ess = variances./asymptotic_variances*(t_skeleton[end]-t_skeleton[1])
end

function ess_pdmp_squared_radius(t_skeleton::Vector{Float64}, x_skeleton::Vector{Vector{Float64}}, v_skeleton::Vector{Vector{Float64}}; n_batches::Int = 50, ellipticQ::Bool = false, x_ref::Vector{Float64} = Vector{Float64}(undef,0), true_mean::Float64 = NaN, true_variance::Float64 = NaN)

  n_samples = n_batches * 1000
  samples = ExtractSamples(t_skeleton, x_skeleton, v_skeleton, n_samples, ellipticQ=ellipticQ, x_ref = x_ref);
  X = sum(samples.^2, dims = 1)
  if isnan(true_mean)
    mean = Vector{Float64}(undef,0)
  else
    mean = [true_mean]
  end
  if isnan(true_variance)
    variance = Vector{Float64}(undef,0)
  else
    variance = [true_variance]
  end
  ESS(X,n_batches = n_batches, true_mean = mean, true_variance = variance)[1]
end

function plot_pdmp(t_skeleton, x_skeleton, v_skeleton; ellipticQ::Bool = false, x_ref::Vector{Float64} = Vector{Float64}(undef,0), name::String = "", n_samples::Int = 0)

  if (n_samples == 0)
    n_samples = Int(round(t_skeleton[end] - t_skeleton[1]))
  end
  samples = ExtractSamples(t_skeleton, x_skeleton, v_skeleton, n_samples, ellipticQ=ellipticQ, x_ref = x_ref);
  plot_ref = plot(samples[1,:], samples[2,:],legend = false, ratio=:equal);
  savefig(plot_ref, string(name,"-trace.pdf"))
  # plot_ref = bar(autocor(samples'))
  # if (length(name) > 0)
  #     savefig(plot_ref, string(name,"-acf.pdf"))
  # end
end

function ExtractSamples(t_skeleton::Vector{Float64}, x_skeleton::Vector{Vector{Float64}}, v_skeleton::Vector{Vector{Float64}}, n_samples::Int; x_ref::Vector{Float64} = Vector{Float64}(undef,0), ellipticQ::Bool=false)

    T = t_skeleton[end]
    Δt = T/(n_samples-1)
    dim = size(x_skeleton[1])[1]
    samples = Matrix{Float64}(undef, dim, n_samples)
    t = 0.0
    if (ellipticQ && length(x_ref) == 0)
        println("WARNING : ExtractSamples : No valid reference point provided.")
        x_ref = zeros(dim)
    end
    x = x_skeleton[1]
    v = v_skeleton[1]
    samples[:,1] = x
    skeleton_index = 2
    counter = 1
    for i in range(2, stop=n_samples)
        t_max = (i-1) * Δt
        while (counter + 1 < length(t_skeleton) && t_skeleton[counter + 1] < t_max)
            counter = counter + 1
            x = x_skeleton[counter]
            v = v_skeleton[counter]
            t = t_skeleton[counter]
        end
        if (ellipticQ)
            (y,v) = EllipticDynamics(t_max - t, x-x_ref, v)
            x = y + x_ref
        else
            x = x + v * (t_max - t)
        end
        t = t_max
        samples[:,i] = x
    end

    return samples
end


using MathGL
mgl = MathGL

gr = mgl.Graph(800, 500)
mgl.xlabel(gr, "x")
mgl.ylabel(gr, "y")
mgl.box(gr)
mgl.axis(gr)
mgl.grid(gr)
mgl.plot(gr, 0.8*sin(linspace(-4pi, 4pi, 200)))
mgl.write(gr, "../graphs/sin_julia.png")


export ShapeProperties, material, transform
export Shape, normal, getcolor
export Sphere

mutable struct ShapeProperties
    material::Material
    transf::TransfMatrix

    ShapeProperties(mat) = new(mat, TransfMatrix())
    ShapeProperties(mat, transf) = new(mat, transf)
end

abstract type Shape end
material(s::Shape) = s.prop.material
transform(s::Shape) = s.prop.transf

mutable struct Sphere <: Shape
    center::Vec3d
    radius::Float64
    prop::ShapeProperties

    Sphere(c, r, mat) = new(c, r, ShapeProperties(mat))
    Sphere(c, r, mat, transf) = new(c, r, ShapeProperties(mat, transf))
end

"""
    normal(pt::Vec3d, sphere::Sphere)

Given a point `pt` on the surface of the sphere, return a `Vec3d`
object containing the normalize outward vector outwards the surface at
point `pt`.
"""
function normal(pt::Vec3d, sphere::Sphere)
    normalize(pt - sphere.center)
end

function getcolor(sphere::Sphere, point::Vec3d)
    normpoint = normalize(point - sphere.center)
    theta = acos(normpoint[3])
    phi = atan(normpoint[2], normpoint[1])
    getcolor(material(sphere), mod2pi(phi) / 2π, theta / π)
end


export SpelledInterval, spelled, spelledp
export SpelledIC, sic, spc
export octaves, internalocts, fifths
export degree, generic, diasteps, alteration, letter
export parsespelled, parsespelledpitch, @i_str, @p_str

import Base: +, -, *, ==

# helpers

const perfectints = Set{Int}([0,3,4])
accstr(n, u, d) =
    if n > 0; repeat(u,n) elseif n < 0; repeat(d,-n) else "" end
qualpf(n, a, p, d) =
    if n > 0; repeat(a,n) elseif n < 0; repeat(d,-n) else p end
qualimpf(n, a, mj, mn, d) =
    if n > 0; repeat(a,n) elseif n < -1; repeat(d,-n-1) elseif n == -1; mn else mj end

# common functions (special interface)

"""
    degree(i)
    degree(p)

Return the "relative scale degree" (0-6) to which the interval points
(unison=`0`, 2nd=`1`, octave=`0`, 2nd down=`6`, etc.).
For pitches, return the integer that corresponds to the letter (C=`0`, D=`1`, ...).

See also: [`generic`](@ref), [`diasteps`](@ref), [`letter`](@ref)
"""
function degree end
degree(fifths::Int) = mod(fifths*4,7)
degree(p::Pitch) = degree(p.pitch)

"""
    generic(i)

Return the generic interval, i.e. the number of diatonic steps modulo octave.
Unlike [`degree`](@ref), the result respects the sign of the interval
(unison=`0`, 2nd up=`1`, 2nd down=`-1`).
For pitches, use [`degree`](@ref).

See also: [`degree`](@ref), [`diasteps`](@ref)
"""
function generic end

"""
    letter(p)

Returns the letter of a pitch as a character.

See also: [`degree`](@ref).
"""
letter(p::Pitch) = 'A' + mod(degree(p)+2, 7)

"""
    diasteps(i)

Return the diatonic steps of the interval (unison=`0`, 2nd=`1`, ..., octave=`7`).
Respects both direction and octaves.

See also [`degree`](@ref) and [`generic`](@ref).
"""
function diasteps end

"""
    alteration(i)
    alteration(p)

Return the number of semitones by which the interval is altered from its the perfect or major variant.
Positive alteration always indicates augmentation,
negative alteration indicates diminution (minor or smaller) of the interval.
For pitches, return the accidentals (positive=sharps, negative=flats, `0`=natural).
"""
function alteration end
alteration(p::Pitch) = alteration(ic(p.pitch))

"""
    octaves(i)
    octaves(p)

Return the number of octaves the interval spans.
Positive intervals start at 0 octaves, increasing.
Negative intervals start at -1 octaves, decreasing.
(You might want to use `octaves(abs(i))` instead).

For a pitch, return its octave.
"""
function octaves end
octaves(p::Pitch) = octaves(p.pitch)

"""
    internalocts(i)

Return the internal octaves (i.e. dependent on the 5ths dimension) of an interval.
"""
function internalocts end

"""
    fifths(i)
    fifths(p)

Return the octave-invariant part of the interval in fifths
(unison=`0`, 5th up=`1`, 4th up/5th down=`-1`).
For a pitch, return the pitch class on the line of fifths.
"""
function fifths end
fifths(p::Pitch) = fifths(p.pitch)

# spelled interval
# ----------------

"""
    SpelledInterval <: Interval

Spelled intervals represented as pairs of fifths and octaves.
E.g., `SpelledInterval(-3, 2)` represents a minor 3rd upwards
(3 fifths down, 2 octaves up).
"""
struct SpelledInterval <: Interval
    fifths :: Int
    octaves :: Int
end

"""
    spelled(fifths, octaves)

Creates a spelled interval from `fifths` and `octaves`.
"""
spelled(f, o) = SpelledInterval(f, o)

"""
    spelledp(fifths, octaves)

Creates a spelled pitch from `fifths` and `octaves`.
"""
spelledp(f, o) = Pitch(spelled(f, o))

# accessors

degree(i::SpelledInterval) = degree(i.fifths)
generic(i::SpelledInterval) = if sign(i) < 0; -degree(-i.fifths) else degree(i.fifths) end
diasteps(i::SpelledInterval) = i.fifths*4 + i.octaves*7
alteration(i::SpelledInterval) = let absi = abs(i); fld(absi.fifths+1, 7) end
octaves(i::SpelledInterval) = i.octaves + fld(i.fifths*4, 7)
internalocts(i::SpelledInterval) = i.octaves
fifths(i::SpelledInterval) = i.fifths

# general interface

function Base.show(io::IO, i::SpelledInterval)
   # negative? print as -abs(i)
    if sign(i) == -1
        print(io, "-")
        print(io, abs(i))
        return
    end

    dia = degree(i)
    diff = alteration(i) # interval is always positive, so direction is correct
    qual = if dia ∈ perfectints
        qualpf(diff, 'a', 'P', 'd')
    else
        qualimpf(diff, 'a', 'M', 'm', 'd')
    end

    octs = octaves(i)

    print(io, qual * string(dia+1) * ':' * string(octs))
end

function Base.show(io::IO, p::Pitch{SpelledInterval})
    i = p.pitch
    dia = degree(i)
    alter = alteration(i)
    print(io, string('A' + mod(dia + 2, 7), accstr(alter, '♯', '♭'), string(octaves(i))))
end

Base.isless(i1::SpelledInterval, i2::SpelledInterval) =
    (diasteps(i1), alteration(i1)) <
    (diasteps(i2), alteration(i2))

Base.isequal(i1::SpelledInterval, i2::SpelledInterval) =
    i1.octaves == i2.octaves && i1.fifths == i2.fifths

Base.hash(i::SpelledInterval, x::UInt) = hash(i.octaves, hash(i.fifths, x))

+(i1::SpelledInterval, i2::SpelledInterval) =
    spelled(i1.fifths + i2.fifths, i1.octaves + i2.octaves)
-(i1::SpelledInterval, i2::SpelledInterval) =
    spelled(i1.fifths - i2.fifths, i1.octaves - i2.octaves)
-(i::SpelledInterval) = spelled(-i.fifths, -i.octaves)
Base.zero(::Type{SpelledInterval}) = spelled(0,0)
Base.zero(::SpelledInterval) = spelled(0,0)
*(i::SpelledInterval, n::Integer) = spelled(i.fifths*n, i.octaves*n)
*(n::Integer,i::SpelledInterval) = spelled(i.fifths*n, i.octaves*n)

octave(::Type{SpelledInterval}) = spelled(0,1)

Base.sign(i::SpelledInterval) = sign(diasteps(i))
Base.abs(i::SpelledInterval) = if sign(i) < 0; -i else i end

ic(i::SpelledInterval) = sic(i.fifths)
embed(i::SpelledInterval) = i
intervaltype(::Type{SpelledInterval}) = SpelledInterval
intervalclasstype(::Type{SpelledInterval}) = SpelledIC

isstep(i::SpelledInterval) = abs(diasteps(i)) <= 1
chromsemi(::Type{SpelledInterval}) = spelled(7,-4)

# conversion

tomidi(i::SpelledInterval) = midi(7*i.fifths + 12*i.octaves)
tomidi(p::Pitch{SpelledInterval}) = Pitch(tomidi(p.pitch) + midi(12))  # C4 = 48 semitones above C0 = midi(60)

# spelled interval class
# ----------------------

"""
    SpelledIC <: IntervalClass

Spelled interval class represented on the line of 5ths with `0 = C`.
E.g., `SpelledIC(3)` represents a major 6th upwards or minor 3rd downwards
(i.e., three 5ths up modulo octave).
"""
struct SpelledIC <: IntervalClass
    fifths :: Int
end

"""
    sic(fifths)

Creates a spelled interval class going `fifths` 5ths upwards.
"""
sic(fs) = SpelledIC(fs)

"""
    spc(fifths)

Creates a spelled pitch class.
In analogy to `sic`, this function takes a number of 5ths.
"""
spc(fs) = Pitch(sic(fs))

# accessors

degree(i::SpelledIC) = degree(i.fifths)
generic(i::SpelledIC) = degree(i.fifths)
diasteps(i::SpelledIC) = degree(i.fifths)
alteration(i::SpelledIC) = fld(i.fifths + 1, 7)
octaves(i::SpelledIC) = 0
internalocts(i::SpelledIC) = 0
fifths(i::SpelledIC) = i.fifths

# interface functions

function Base.show(io::IO, ic::SpelledIC)
    dia = degree(ic)
    diff = alteration(ic)
    qual = if dia ∈ perfectints
        qualpf(diff, 'a', 'P', 'd')
    else
        qualimpf(diff, 'a', 'M', 'm', 'd')
    end

    print(io, qual * string(dia+1))
end

function Base.show(io::IO, p::Pitch{SpelledIC})
    i = p.pitch
    dia = degree(i)
    alter = alteration(i)
    print(io, ('A' + mod(dia+2, 7)) * accstr(alter, '♯', '♭'))
end

Base.isless(i1::SpelledIC, i2::SpelledIC) = isless(i1.fifths,i2.fifths)
Base.isequal(i1::SpelledIC, i2::SpelledIC) = isequal(i1.fifths,i2.fifths)
Base.hash(i::SpelledIC, x::UInt) = hash(i.fifths, x)

+(i1::SpelledIC, i2::SpelledIC) = sic(i1.fifths + i2.fifths)
-(i1::SpelledIC, i2::SpelledIC) = sic(i1.fifths - i2.fifths)
-(i::SpelledIC) = sic(-i.fifths)
Base.zero(::Type{SpelledIC}) = sic(0)
Base.zero(::SpelledIC) = sic(0)
*(i::SpelledIC, n::Integer) = sic(i.fifths * n)
*(n::Integer,i::SpelledIC) = sic(i.fifths * n)

octave(::Type{SpelledIC}) = sic(0)
function Base.sign(i::SpelledIC)
    dia = degree(i)
    if dia == 0; 0 elseif dia > 3; -1 else 1 end
end
Base.abs(i::SpelledIC) = if sign(i) < 0; -i else i end

ic(i::SpelledIC) = i
embed(i::SpelledIC) = spelled(i.fifths, -fld(i.fifths*4, 7))
intervaltype(::Type{SpelledIC}) = SpelledInterval
intervalclasstype(::Type{SpelledIC}) = SpelledIC

isstep(i::SpelledIC) = degree(i) ∈ [0,1,6]
chromsemi(::Type{SpelledIC}) = sic(7)

# conversion

tomidi(i::SpelledIC) = midic(i.fifths * 7)
tomidi(p::Pitch{SpelledIC}) = midipc(p.pitch.fifths * 7)

# parsing
# -------

const rgsic = r"^(-?)(a+|d+|[MPm])([1-7])$"
const rgspelled = r"^(-?)(a+|d+|[MPm])([1-7]):(-?)(\d+)$"

function matchinterval(modifier, num)
    dia = parse(Int, num) - 1
    perfect = dia ∈ perfectints

    alt = if modifier == "M" && !perfect
        0
    elseif modifier == "m" && !perfect
        -1
    elseif lowercase(modifier) == "p" && perfect
        0
    elseif occursin(r"^a+$", modifier)
        length(modifier)
    elseif occursin(r"^d+$", modifier)
        -length(modifier) - (perfect ? 0 : 1)
    else
        error("cannot parse interval \"$modifier$num\"")
    end

    mod(dia*2+1, 7) - 1 + 7*alt
end

"""
    parsespelled(str)

Parse a spelled interval or interval class string.
The type is determined from the string,
so `i"M3:0"` returns an interval while `i"M3"` returns an interval class.

See also: [`@i_str`](@ref), [`parsespelledpitch`](@ref).
"""
function parsespelled(str)
    m = match(rgsic, str)
    if m != nothing
        int = sic(matchinterval(m[2], m[3]))
    else
        m = match(rgspelled, str)
        if m != nothing
            fifths = matchinterval(m[2], m[3])
            octs = parse(Int, m[4]*m[5])
            int = spelled(fifths, octs - fld(fifths*4, 7))
        else
            error("cannot parse interval \"$str\"")
        end
    end

    # invert if necessary
    if m[1] == "-"
        -int
    else
        int
    end
end

"""
    i"str"

Parse a spelled interval or interval class string.
The type is determined from the string,
so `i"M3:0"` returns an interval while `i"M3"` returns an interval class.

See also: [`@p_str`](@ref), [`parsespelled`](@ref), [`parsespelledpitch`](@ref).
"""
macro i_str(str)
    parsespelled(str)
end

const rgspelledpc = r"^([a-g])(♭+|♯+|b+|#+)?$"i
const rgspelledp = r"^([a-g])(♭+|♯+|b+|#+)?(-?\d+)$"i

function matchpitch(letter, accs)
    letter = uppercase(letter)[1]
    if letter >= 'A' && letter <= 'G'
        dia = mod(letter - 'A' - 2, 7)
    else
        error("cannot parse pitch letter \"$letter\"")
    end

    alt = if accs == nothing || accs == ""
        0
    elseif occursin(r"^♭+|b+$"i, accs)
        -length(accs)
    elseif occursin(r"^♯+|#+$"i, accs)
        length(accs)
    else
        error("cannot parse accidentals \"$accs\"")
    end

    mod(dia*2 + 1, 7) - 1 + 7*alt
end

"""
    parsespelledpitch(str)

Parse a spelled pitch or pitch class string.
The type is determined from the string,
so `p"G4"` returns a pitch while `p"G"` returns a pitch class.

See also: [`@p_str`](@ref), [`parsespelled`](@ref).
"""
function parsespelledpitch(str)
    m = match(rgspelledpc, str)
    if m != nothing
        spc(matchpitch(m[1], m[2]))
    else
        m = match(rgspelledp, str)
        if m != nothing
            octs = parse(Int, m[3])
            fifths = matchpitch(m[1], m[2])
            spelledp(fifths, octs - fld(fifths*4, 7))
        else
            error("cannot parse pitch \"$str\"")
        end
    end
end

"""
    p"str"

Parse a spelled pitch or pitch class string.
The type is determined from the string,
so `p"G4"` returns a pitch while `p"G"` returns a pitch class.

See also: [`@i_str`](@ref), [`parsespelledpitch`](@ref), [`parsespelled`](@ref).
"""
macro p_str(str)
    parsespelledpitch(str)
end


using GPMLj, Test, Random

include("kernels.jl")
include("likelihoods.jl")
include("mean_functions.jl")
include("models.jl")
include("parameter_priors.jl")
include("train.jl")


module DeepShiba
    include("core/core_structures.jl")
    include("core/core_functions.jl")
    include("utils/output_txt.jl")
    include("utils/plot.jl")
    include("functions/math_functions.jl")
    include("functions/trans_functions.jl")
    export Variable, Func, variable, @variable, backward!, numerical_diff, cleargrad!, show, +, -, *, /, reshape, sin, cos, tanh, plot
end


# core plotting functions

plotbeta(α, β)
plotbeta() = plotBeta(1, 1)

plotunif(a, b)
plotunif() = plotUnif(0.0, 1.0)

plotbinom(n, p)
plotbinom() = plotBinom(1, 0.5)

plotbernoulli(p)
plotbernoulli() = plotBernolli(0.5)

plottest(t)


using UnitfulColors
using Plots, Unitful, UnitfulRecipes, Colors, Latexify, UnitfulLatexify

function showUnitColor(title, x, varargs...)
    return heatmap(
        x,
        [-1, 1],
        repeat(x, 2, 1);
        background=:black,
        color=colormatch.(x, varargs...),
        legend=false,
        colorbar=false,
        xlabel="\\textrm{$title}",
        yticks=[],
        yshowaxis=false,
        fontfamily="Computer Modern",
        unitformat=latexify,
    )
end

length = 1000;
pl = plot(
    showUnitColor("Wavelength", range(350, 750; length) .* u"nm"),
    showUnitColor("Frequency", range(400, 800; length) .* u"THz"),
    showUnitColor("Energy", range(1.5, 3.5; length) .* u"eV"),
    showUnitColor("Temperature (peak)", range(4000, 8000; length) .* u"K"),
    showUnitColor("Temperature (spectrum)", range(1000, 27000; length) .* u"K", :spectrum);
    layout=(5, 1),
);

savefig(pl, "$(pkgdir(UnitfulColors))/doc/examples.png")


include("../src/LarSurf.jl")
using Plasm
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
# using LarSurf

block_size = [5, 5, 5]
margin_size = 0

## Artifical data
data3d = ones(Int16, 19, 15, 31)
data3d[2:14,8:13,9:11] .= 10
# data3d

voxelsize_mm = [0.5, 0.9, 0.8]
threshold=0

blocks_number, blocks_number_axis = LarSurf.number_of_blocks_per_axis(
    size(data3d), block_size)

block1 = LarSurf.get_block(
    1, data3d, block_size, margin_size, blocks_number_axis, false
)

verts, trifaces = LarSurf.import_data3d(block1, voxelsize_mm)

println(verts)
println(trifaces)

println("======= To LAR ========")

V, EV, FE = LarSurf.to_lar(verts, trifaces)

println(V)
println(EV)
println(FE)

println("V: ", typeof(V), size(V))
println("EV: ", typeof(EV))
print("FV: ", typeof(FE))

@show EV;


using LowRankModels
Random.seed!(1);

test_losses = Loss[
QuadLoss(),
L1Loss(),
HuberLoss(),
PeriodicLoss(1),
OrdinalHingeLoss(1,10),
LogisticLoss(),
WeightedHingeLoss()
]

for test_iteration = 1:500
	# Create the configuration for the model (random losses)
	config = round.(Int, abs(round(4*rand(length(test_losses)))));
	#config = [0 0 1 0 10 0 100]
	losses, doms = Array(Loss,1), Array(Domain,1);
	for (n,l) in zip(config, test_losses)
		for i=1:n
			push!(losses, l);
			push!(doms, l.domain);
		end
	end
	losses, doms = losses[2:end], doms[2:end]; # this is because the initialization leaves us with an #undef
	my_error_metric(glrm::GLRM, X::Array{Float64,2}, Y::Array{Float64,2}) = error_metric(glrm, X, Y, doms, standardize=true) # embed the domains into the error function.

	# Make a low rank matrix as our data precursor
	m, n, true_k = 100, length(doms), round.(Int, round(length(losses)/2))+1;

	X_real, Y_real = randn(true_k,m), randn(true_k,n);
	A_real = X_real'*Y_real;

	# Impute over the low rank-precursor to make our heterogenous dataset
	A = impute(doms, losses, A_real);				# our data with noise

	# Create a glrm using these losses and data
	p = Params(1e-2, max_iter=1000, abs_tol=0.00000001, min_stepsize=1e-15)
	rx, ry = ZeroReg(), ZeroReg();

	k_range = [int(round(true_k/2)), true_k]
	train_err_at_k = Dict()
	println("\n########################################################")
	println("Model:\nlosses = $(losses)\nrx,ry = $(rx), $(ry)")
	for k in k_range
		model = GLRM(A, losses, rx,ry,k, scale=false, offset=false);
		println("\nk = $(k)")
		# Test that our imputation is consistent
		if my_error_metric(model, X_real, Y_real) != 0
			error("Imputation failed.")
		end

		real_obj = objective(model, X_real, Y_real, include_regularization=false);
		X_fit,Y_fit,ch = fit!(model, params=p, verbose=false);
		println("Starting objective: $(ch.objective[1])\t Ending objective: $(ch.objective[end])")

		train_err, test_err, trainers, testers = cross_validate(model, nfolds=3, error_fn=my_error_metric, verbose=false);
		train_err_at_k[k] = my_error_metric(model, model.X, model.Y)
		println("train err: $train_err")
		println("Test err: $test_err")
	end
	train_err_at_k = [train_err_at_k[k] for k in k_range]
	if !all([de<0 for de in diff(train_err_at_k)]) #check that errors monotonically decrease as k increases
		@warn("==================================================================================")
		@warn("ERRORS WENT UP FOR THIS CONFIGURATION")
		@warn("Model:\nlosses = $(losses)\nrx,ry = $(rx), $(ry)")
		@warn("Ranks: $(string(k_range))")
		@warn("Training errors: $(string(train_err_at_k))")
		@warn("==================================================================================")
	end
end


#=

  Balanced brackets in ConstraintSolver.jl

  This model all balanced brackets of size m*2.

  See https://en.wikipedia.org/wiki/Catalan_number

  The number of generated solutions for m:
   m        #
   ----------
    1       1
    2       2
    3       5
    4      14
    5      42
    6     132
    7     429
    8    1430
    9    4862
   10   16796
   11   58786
   12  208012
   13  742900

  Which - of course - is the Catalan numbers.

  http://oeis.org/search?q=1%2C2%2C5%2C14%2C42%2C132%2C429%2C1430%2C4862%2C16796%2C58786%2C208012&language=english&go=Search
  http://oeis.org/A000108

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

=#

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

function balanced_brackets(m=4,print_solutions=true,all_solutions=true,timeout=6)

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

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

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

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

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

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

                                                            "time_limit"=>timeout,

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

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            "lp_optimizer" => ipopt_optimizer,
                                        ))
    n = m*2
    t = [1,-1] # +1 for "[", -1 for "]"
    @variable(model, 1 <= x[1:n] <= 2, Int) # 1="[",  2="]"

    # Counter (cumulative sum)
    @variable(model, 0 <= c[1:n] <= n, Int)

    # Leading "["
    @constraint(model, x[1] == 1)
    @constraint(model, c[1] == 1)

    for i in 2:n
        tt = @variable(model, [1:1], CS.Integers([-1,1]))
        my_element(model,x[i], t, tt[1])
        @constraint(model, c[i] == c[i-1] + tt[1])
    end

    # Concluding "]"
    @constraint(model, x[n] == 2)
    @constraint(model, c[n] == 0)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                if print_solutions
                    x_val = convert.(Integer,JuMP.value.(x; result=sol))
                    c_val = convert.(Integer,JuMP.value.(c; result=sol))
                    println("x:$x_val")
                    println("c:$c_val")
                    println()
                end

            end
        end
    else
        println("status:$status")
    end

    return status, num_sols
end

begin
    @time local status, num_sols = balanced_brackets(4,true,true)
    println("num_sols: $num_sols")

    for m in 1:20
        println("m:$m")
        @time local status, num_sols = balanced_brackets(m,false,true)
        if status == MOI.TIME_LIMIT
            break
        end
    end
end


type OrnsteinUhlenbeckProcess <: StochasticProcess1D
  speed::Float64
  vol::Float64
  x0::Float64
  level::Float64

  OrnsteinUhlenbeckProcess(speed::Float64, vol::Float64, x0::Float64 = 0.0, level::Float64 = 0.0) = new(speed, vol, x0, level)
end

expectation(process::OrnsteinUhlenbeckProcess, ::Float64, x0::Float64, dt::Float64) = process.level + (x0 - process.level) * exp(-process.speed * dt)
get_x0(process::OrnsteinUhlenbeckProcess) = process.x0

function variance(process::OrnsteinUhlenbeckProcess, ::Float64, ::Float64, dt::Float64)
  v = process.vol
  if process.speed < sqrt(eps())
    # algebraic limits for small speed
    return v * v * dt
  else
    return 0.5 * v * v / process.speed * (-expm1(-2.0 * process.speed * dt))
  end
end

std_deviation(process::OrnsteinUhlenbeckProcess, t::Float64, x0::Float64, dt::Float64) = sqrt(variance(process, t, x0, dt))

type GsrProcess <: StochasticProcess1D
  times::Vector{Float64}
  vols::Vector{Float64}
  reversions::Vector{Float64}
  T::Float64
  revZero::Vector{Bool}
  cache1::Dict{Pair{Float64, Float64}, Float64}
  cache2::Dict{Pair{Float64, Float64}, Float64}
  cache3::Dict{Pair{Float64, Float64}, Float64}
  cache4::Dict{Float64, Float64}
  cache5::Dict{Pair{Float64, Float64}, Float64}
end

function GsrProcess(times::Vector{Float64}, vols::Vector{Float64}, reversions::Vector{Float64}, T::Float64 = 60.0)
  revZero = falses(length(reversions))
  cache1 = Dict{Pair{Float64, Float64}, Float64}()
  cache2 = Dict{Pair{Float64, Float64}, Float64}()
  cache3 = Dict{Pair{Float64, Float64}, Float64}()
  cache4 = Dict{Float64, Float64}()
  cache5 = Dict{Pair{Float64, Float64}, Float64}()

  return GsrProcess(times, vols, reversions, T, revZero, cache1, cache2, cache3, cache4, cache5)
end

get_x0(::GsrProcess) = 0.0

function flush_cache!(gsrP::GsrProcess)
  for i in eachindex(gsrP.reversions)
    if abs(gsrP.reversions[i]) < 1e-4
      gsrP.revZero[i] = true
    else
      gsrP.revZero[i] = false
    end
  end
  empty!(gsrP.cache1)
  empty!(gsrP.cache2)
  empty!(gsrP.cache3)
  empty!(gsrP.cache4)
  empty!(gsrP.cache4)

  return gsrP
end

lower_index(process::GsrProcess, x::Float64) = upper_bound(process.times, x)
upper_index(process::GsrProcess, x::Float64) = upper_bound(process.times, x - 1e-15) + 1

function time2(process::GsrProcess, idx::Int)
  if idx == 1
    return 0.0
  end

  if idx > length(process.times) + 1
    return process.T
  end

  return process.times[idx - 1]
end

vol(process::GsrProcess, x::Int) = x > length(process.vols) ? process.vols[end] : process.vols[x]
rev_zero(process::GsrProcess, x::Int) = x > length(process.revZero) ? process.revZero[end] : process.revZero[x]
rev(process::GsrProcess, x::Int) = x > length(process.reversions) ? process.reversions[end] : process.reversions[x]

floored_time(process::GsrProcess, x::Int, flr::Float64) = flr != -1.0 ? max(flr, time2(process, x)) : time2(process, x)
capped_time(process::GsrProcess, x::Int, cap::Float64) = cap != -1.0 ? min(cap, time2(process, x)) : time2(process, x)

function variance(process::GsrProcess, w::Float64, ::Float64, dt::Float64)
  t = w + dt
  key = Pair(w, t)

  p = get(process.cache3, key, typemax(Float64))
  if p != typemax(Float64)
    return p
  end

  res = 0.0
  for k = lower_index(process, w):upper_index(process, t) - 1
    res2 = vol(process, k) * vol(process, k)
    res2 *= rev_zero(process, k) ? -(floored_time(process, k, w) - capped_time(process, k + 1, t)) :
                                  (-expm1(2.0 * rev(process, k) * (floored_time(process, k, w) - capped_time(process, k + 1, t)))) /
                                  (2.0 * rev(process, k))

    for i = k + 1:upper_index(process, t) - 1
      res2 *= exp(-2.0 * rev(process, i) * (capped_time(process, i + 1, t) - time2(process, i)))
    end

    res += res2
  end

  process.cache3[key] = res

  return res
end

function expectationp1(process::GsrProcess, w::Float64, xw::Float64, dt::Float64)
  t = w + dt
  key = Pair(w, t)
  k = get(process.cache1, key, typemax(Float64))
  if k != typemax(Float64)
    return xw * k
  end

  res2 = 1.0
  for i = lower_index(process, w):upper_index(process, t) - 1
    res2 *= exp(-rev(process, i) * (capped_time(process, i + 1, t) - floored_time(process, i, w)))
  end

  process.cache1[key] = res2
  return res2 * xw
end

function expectationp2(process::GsrProcess, w::Float64, dt::Float64)
  t = w + dt
  key = Pair(w, t)

  p = get(process.cache2, key, typemax(Float64))
  if p != typemax(Float64)
    return p
  end

  T = process.T

  res = 0.0
  # int A(s, t)y(s)
  for k = lower_index(process, w):upper_index(process, t) - 1
    # l < k
    for l = 1:k - 1
      res2 = 1.0
      # alpha 1
      if rev_zero(process, l)
        res2 *= vol(process, l) * vol(process, l) * (time2(process, l + 1) - time2(process, l))
      else
        res2 *= vol(process, l) * vol(process, l) / (2.0 * rev(process, l)) * (1.0 - exp(-2.0 * rev(process, l) *
                (time2(process, l + 1) - time2(process, l))))
      end
      # zeta_i (i > k)
      for i = k+1:upper_index(process, t) - 1
        res2 *= exp(-rev(process, i) * (capped_time(process, i + 1, t) - time2(process, i)))
      end
      # beta_i (j<k)
      for j = l+1:k - 1
        res2 *= exp(-2.0 * rev(process, j) * (time2(process, j + 1) - time2(process, j)))
      end

      # zeta_k beta_k
      if rev_zero(process, k)
        res2 *= 2.0 * time2(process, k) - floored_time(process, k, w) - capped_time(process, k + 1, t) -
                2.0 * (time2(process, k) - capped_time(process, k + 1, t))
      else
        res2 *= (exp(rev(process, k) * (2.0 * time2(process, k) - floored_time(process, k, w) - capped_time(process, k + 1, t))) -
                exp(2.0 * rev(process, k) * (time2(process, k) - capped_time(process, k + 1, t)))) / rev(process, k)
      end
      res += res2
    end
    # l = k
    res2 = 1.0
    # alpha_k zeta_k
    if rev_zero(process, k)
      res2 *= vol(process, k) * vol(process, k) / 4.0 * (4.0 * ^(capped_time(process, k + 1, t) - time2(process, k), 2.0) -
              (^(floored_time(process, k, w) - 2.0 * time2(process, k) + capped_time(process, k + 1, t), 2.0) +
              ^(capped_time(process, k + 1, t) - floored_time(process, k, w), 2.0)))
    else
      res2 *= vol(process, k) * vol(process, k) / (2.0 * rev(process, k) * rev(process, k)) * (exp(-2.0 * rev(process, k) *
              (capped_time(process, k + 1, t) - time2(process, k))) + 1.0 - (exp(-rev(process, k) * (floored_time(process, k, w) -
              2.0 * time2(process, k) + capped_time(process, k + 1, t))) + exp(-rev(process, k) *
              (capped_time(process, k + 1, t) - floored_time(process, k, w)))))
    end

    # zeta_i (i>k)
    for i = k+1:upper_index(process, t) - 1
      res2 *= exp(-rev(process, i) * (capped_time(process, i + 1, t) - time2(process, i)))
    end
    # no beta_j in this case
    res += res2
  end

  # int -A(s, t) sigma^2 G(s, T)
  for k = lower_index(process, w):upper_index(process, t) - 1
    res2 = 0.0
    # l > k
    for l = k+1:upper_index(process, T) - 1
      res3 = 1.0
      # eta_l
      res3 *= rev_zero(process, l) ? capped_time(process, l + 1, T) - time2(process, l) :
                                      (1.0 - exp(-rev(process, l) * (capped_time(process, l + 1, T) - time2(process, l)))) /
                                      rev(process, l)

      # zeta_i (i > k)
      for i = k+1:upper_index(process, t) - 1
        res3 *= exp(-rev(process, i) * (capped_time(process, i + 1, t) - time2(process, i)))
      end

      # gamma_j (j > k)
      for j = k+1:l-1
        res3 *= exp(-rev(process, j) * (time2(process, j + 1) - time2(process, j)))
      end

      # zeta_k gamma_k
      if rev_zero(process, k)
        res3 *= (capped_time(process, k + 1, t) - time2(process, k + 1) - (2.0 * floored_time(process, k, w) - capped_time(process, k + 1, t) -
                time2(process, k + 1))) / 2.0
      else
        res3 *= (exp(rev(process, k) * (capped_time(process, k + 1, t) - time2(process, k + 1))) - exp(rev(process, k) *
                (2.0 * floored_time(process, k, w) - capped_time(process, k + 1, t) - time2(process, k + 1)))) / (2.0 * rev(process, k))
      end

      res2 += res3
    end
    # l = k
    res3 = 1.0
    # eta_k zeta_k
    if rev_zero(process, k)
      res3 *= (-^(capped_time(process, k + 1, t) - capped_time(process, k + 1, T), 2.0) - 2.0 * ^(capped_time(process, k + 1, t) - floored_time(process, k, w), 2.0) +
              ^(2.0 * floored_time(process, k, w) - capped_time(process, k + 1, T) - capped_time(process, k + 1, t), 2.0)) / 4.0
    else
      res3 *= (2.0 - exp(rev(process, k) * (capped_time(process, k + 1, t) - capped_time(process, k + 1, T))) -
              (2.0 * exp(-rev(process, k) * (capped_time(process, k + 1, t) - floored_time(process, k, w))) -
              exp(rev(process, k) * (2.0 * floored_time(process, k, w) - capped_time(process, k + 1, T) -
              capped_time(process, k + 1, t))))) / (2.0 * rev(process, k) * rev(process, k))
    end

    # zeta_i (i > k)
    for i = k+1:upper_index(process, t) - 1
      res3 *= exp(-rev(process, i) * (capped_time(process, i + 1, t) - time2(process, i)))
    end

    # no gamma_j in this case
    res2 += res3

    # add to main accumulator
    res += -vol(process, k) * vol(process, k) * res2
  end
  process.cache2[key] = res

  return res
end

function G!(process::GsrProcess, t::Float64, w::Float64, ::Float64)
  key = Pair(w, t)
  k = get(process.cache5, key, typemax(Float64))
  if k != typemax(Float64)
    return k
  end

  res = 0.0
  for i = lower_index(process, t):upper_index(process, w) - 1
    res2 = 1.0
    for j = lower_index(process, t):i - 1
      res2 *= exp(-rev(process, j) * (time2(process, j + 1) - floored_time(process, j, t)))
    end
    res2 *= rev_zero(process, i) ? capped_time(process, i + 1, w) - floored_time(process, i, t) :
                                  (1.0 - exp(-rev(process, i) * (capped_time(process, i + 1, w) - floored_time(process, i, t)))) / rev(process, i)

    res += res2
  end

  process.cache5[key] = res
  return res
end

function y!(process::GsrProcess, t::Float64)
  key = t
  k = get(process.cache4, key, typemax(Float64))
  if k != typemax(Float64)
    return k
  end

  res = 0.0
  for i = 1:upper_index(process, t) - 1
    res2 = 1.0
    for j = i + 1:upper_index(process, t) - 1
      res2 *= exp(-2.0 * rev(process, j) * (capped_time(process, j + 1, t) - time2(process, j)))
    end
    res2 *= rev_zero(process, i) ? vol(process, i) * vol(process, i) * (capped_time(process, i+1, t) - time2(process, i)) :
                                   (vol(process, i) * vol(process, i) / (2.0 * rev(process, i)) * (1.0 - exp(-2.0 * rev(process, i) * (capped_time(process, i+1, t) - time2(process, i)))))

    res += res2
  end

  process.cache4[key] = res
  return res
end

function expectation(process::GsrProcess, w::Float64, xw::Float64, dt::Float64)
  # t = w + dt

  # p1 = expectationp1(process, w, xw, dt)
  # p2 = expectationp2(process, w, dt)
  # println(p1)
  # println(p2)
  # error("DIE")

  return expectationp1(process, w, xw, dt) + expectationp2(process, w, dt)
end

std_deviation(process::GsrProcess, t0::Float64, x0::Float64, dt::Float64) = sqrt(variance(process, t0, x0, dt))

evolve(process::StochasticProcess1D, t0::Float64, x0::Float64, dt::Float64, dw::Float64) =
      apply(process, expectation(process, t0, x0, dt), std_deviation(process, t0, x0, dt) * dw)

apply(process::StochasticProcess1D, x0::Float64, dx::Float64) = x0 + dx

std_deviation(process::StochasticProcess1D, t0::Float64, x0::Float64, dt::Float64) = diffusion(process.disc, process, t0, x0, dt)

variance(process::StochasticProcess1D, t0::Float64, x0::Float64, dt::Float64) = variance(process.disc, process, t0, x0, dt)

get_size(process::StochasticProcess1D) = 1
get_factors(process::StochasticProcess) = get_size(process)


## Uniform

abstract type Uniform{T} <: Distribution{T} end

struct UniformType{T} <: Uniform{T} end

Uniform(::Type{T}) where {T} = UniformType{T}()

Base.getindex(::UniformType{T}) where {T} = T

struct UniformWrap{T,E} <: Uniform{E}
    val::T
end

Uniform(x::T) where {T} = Uniform(x, isuniform(T))

Uniform(x::T, ::Val{true}) where {T} = UniformWrap{T,gentype(T)}(x)

Base.getindex(x::UniformWrap) = x.val

const ImplicitUniform = Union{AbstractArray,AbstractDict,AbstractSet,
                              AbstractString,Tuple}

isuniform(::Type{T}) where {T} = Val(false)
isuniform(::Type{T}) where {T<:ImplicitUniform} = Val(true)

Base.show(io::IO, x::Uniform) = print(io, "Uniform($(x[]))")

## Wrap

struct Wrap{T,E} <: Distribution{E}
    val::T
end

Wrap(x::T) where {T} = Wrap{T,gentype(T)}(x)

Base.getindex(x::Wrap) = x.val

## wrap

wrap(::Type{T}) where {T} = Uniform(T)

wrap(x::T) where {T} = wrap(x, isuniform(T))

wrap(x::T, ::Val{true}) where {T} = Uniform(x)
wrap(x::T, ::Val{false}) where {T} = Wrap(x)

wrap(x::Distribution) = x
wrap(x::Sampler) = x

## sampling

Sampler(RNG::Type{<:AbstractRNG},
        d::Union{UniformWrap,UniformType,Wrap}, n::Repetition) =
            Sampler(RNG, d[], n)

## support / pmf

support(d::Union{Wrap,UniformWrap}) = support(d.val)
pmf(d::Union{Wrap,UniformWrap}, x) = pmf(d.val, x)


"""
    null1(N::BinaryNetwork)

Given a network `N`, `null1(N)` returns a network with the same dimensions, where
every interaction happens with a probability equal to the connectance of `N`.

Note that this does not guarantee that the network is not degenerate, so the
output of this analysis *should* be filtered using `is_degenerate`, or passed to
`simplify`. The output of this approach is *always* a probabilistic network of
the same partiteness as the original network.

#### References

Fortuna, M.A., Bascompte, J., 2006. Habitat loss and the structure of
plantanimal mutualistic networks. Ecol. Lett. 9, 281–286.
https://doi.org/10.1111/j.1461-0248.2005.00868.x
"""
function null1(N::BinaryNetwork)
    return linearfilter(N; α=[0.0, 0.0, 0.0, 1.0])
end

"""
    null3(N::BinaryNetwork; dims::Integer=1)

Given a network `N`, `null3(N)` returns a matrix with the same dimensions,
where every interaction happens with a probability equal to the out-degree
(`dims=1`) or to the in-degree (`dims=2`, number of predecessors) of each
species, divided by the total number of possible predecessors/successors.

Note that this does not guarantee that the network is not degenerate, so the
output of this analysis *should* be filtered using `is_degenerate`, or passed to
`simplify`. The output of this approach is *always* a probabilistic network of
the same partiteness as the original network.

#### References

Poisot, T., Stanko, M., Miklisová, D., Morand, S., 2013. Facultative and
obligate parasite communities exhibit different network properties. Parasitology
140, 1340–1345. https://doi.org/10.1017/S0031182013000851
"""
function null3(N::BinaryNetwork; dims::Integer=1)
  @assert dims ∈ [1,2]
  α = dims == 1 ? [0.0, 1.0, 0.0, 0.0] : [0.0, 0.0, 1.0, 0.0]
  return linearfilter(N; α=α)
end

"""
    null2(N::BinaryNetwork)

Given a network `N`, `null2(N)` returns a network with the same dimensions, where
every interaction happens with a probability equal to the degree of each
species.

Note that this does not guarantee that the network is not degenerate, so the
output of this analysis *should* be filtered using `is_degenerate`, or passed to
`simplify`. The output of this approach is *always* a probabilistic network of
the same partiteness as the original network.

#### References

Bascompte, J., Jordano, P., Melian, C.J., Olesen, J.M., 2003. The nested
assembly of plant-animal mutualistic networks. PNAS 100, 9383–9387.
https://doi.org/10.1073/pnas.1633576100
"""
function null2(N::BinaryNetwork)
  return linearfilter(N; α=[0.0, 0.5, 0.5, 0.0])
end

"""
    null4(N::BinaryNetwork)

Given a matrix `A`, `null4(A)` returns a matrix with the same dimensions, where
every interaction happens with a probability equal to the product of the degree
of each species.
"""
function null4(N::BinaryNetwork)
  Afiltered = sum(N, dims=1) .* sum(N, dims=2) ./ sum(N)^2
  ReturnType = typeof(N) <: AbstractBipartiteNetwork ? BipartiteProbabilisticNetwork : UnipartiteProbabilisticNetwork
  return ReturnType(Afiltered, _species_objects(N)...)
end


using PtFEM, Test, Compat

ProjDir = dirname(@__FILE__)

data = Dict(
  # Plane(ndim, nst, nxe, nye, nip, direction, finite_element(nod, nodof), axisymmetric)
  :struc_el => Plane(2, 4, 40, 20, 4, :y, Quadrilateral(8, 2), false),
  # ϕ (phi, friction angle), c (cohesion), ψ (psi, dilation angle),
  # γ (gamma, soil unit weight), E, ν (v, Poisson's ratio)
  # [ϕ, c, ψ, γ, E, ν]
  :properties => [20.0 10.0 20.0 16.0 1.0e5 0.3;],
  :x_coords => 0.0:0.25:10.0,
  :y_coords => 5.0:-0.25:0.0,
  :tol => 0.001,
  :limit => 200,
  :incs => 10,
  :nbo2 => 8,       # No of elements to be rigidly displaced in bearing capacity
  :qs => 20,        # Surface surcharge
  :presc => -0.001  # Magnitude of incremental vertical displacements rigid footing
)

@time res_df, g_coord, g_num, disp = p63(data)
println()
res_df

@test res_df[end,:load1] ≈ 274.3514626820244
@test res_df[end,:load2] ≈ 270.01838672066737
@test res_df[end,:iters] == 81


# %% imports
using Random
using LinearAlgebra
using Statistics
using Plots
using YAML
# %% classes

cfg = YAML.load(open("/Users/titou/Documents/PFE/RL-STDP/Julia/CartPole/CMAES/cfg_cmaes.yml"))

mutable struct CMAES
    N::Int
    μ::Int
    λ::Int
    τ::Float64
    τ_c::Float64
    τ_σ::Float64
    population::Array{Array{Float64}}
    offspring::Array{Array{Float64}}
    F_μ::Array{Float64}
    F_λ::Array{Float64}
    C::Array{Float64}
    s::Array{Float64}
    s_σ::Array{Float64}
    σ::Float64
    E::Array{Float64}
    W::Array{Float64}
    x::Array{Float64}
    cfg::Dict
end

function CMAES(init::Array{Float64} ;N=2, μ=1, λ=10, τ=sqrt(N), τ_c=N^2, τ_σ=sqrt(N), cfg::Dict = cfg)
    x = init
    x = clamp.(x, -1.0, 1.0)
    population = fill(x, µ)
    offspring = Array{Array{Float64}}(undef, λ)
    F_µ = Inf .* ones(µ)
    F_λ = Inf .* ones(λ)
    C = Array(Diagonal{Float64}(I, N))
    s = zeros(N)
    s_σ = zeros(N)
    σ = 1.0
    E = zeros(N, λ)
    W = zeros(N, λ);
    CMAES(N, μ, λ, τ, τ_c, τ_σ, population, offspring, F_µ, F_λ, C, s, s_σ, σ, E, W, x, cfg)
end

function step!(c::CMAES; obj=objective, visualize=false, anim=Nothing)
    # L1
    rng_offspring = Random.MersenneTwister(1234)
    sqrt_c = cholesky((c.C + c.C') / 2.0).U
    for i in 1:c.λ
        c.E[:,i] = randn(rng_offspring, c.N)
        c.W[:,i] = c.σ * (sqrt_c * c.E[:,i])
        c.offspring[i] = c.x + c.W[:,i]
        c.offspring[i] = clamp.(c.offspring[i], -1.0, 1.0)
        c.F_λ[i] = obj(c.offspring[i])
    end
    # Select new parent population
    idx = sortperm(c.F_λ)[1:c.μ]
    for i in 1:c.μ
        c.population[i] = c.offspring[idx[i]]
        c.F_μ[i] = c.F_λ[idx[i]]
    end
    # L2
    w = vec(mean(c.W[:,idx], dims=2))
    c.x += w
    # L3
    c.s = (1.0 - 1.0/c.τ)*c.s + (sqrt(c.μ/c.τ * (2.0 - 1.0/c.τ))/c.σ)*w
    # L4
    c.C = (1.0 - 1.0/c.τ_c).*c.C + (c.s./c.τ_c)*c.s'
    # L5
    ɛ = vec(mean(c.E[:,idx], dims=2))
    c.s_σ = (1.0 - 1.0/c.τ_σ)*c.s_σ + sqrt(c.μ/c.τ_σ*(2.0 - 1.0/c.τ_σ))*ɛ
    # L6
    c.σ = c.σ*exp(((c.s_σ'*c.s_σ)[1] - c.N)/(2*c.N*sqrt(c.N)))
    if visualize
        plot(xs, ys, fz, st=:contour)
        scatter!([c.offspring[i][1] for i in 1:c.λ], [c.offspring[i][2] for i in 1:c.λ],
            xlims=(-5, 5), ylims=(-5, 5), legend=:none)
        scatter!([c.x[1]], [c.x[2]], color=:black, marker=:rect,
            xlims=(-5, 5), ylims=(-5, 5), legend=:none)
        frame(anim)
    end
    c
end

function plot_obj()
    c = CMAES()
    println("x initial: ", c.x)
    anim = Animation()
    for i in 1:100
        v = mod(i, 1) == 0
        step!(c, visualize=v, anim=anim)
    end
    println("x final: ", c.x)
    gif(anim)
end

# %% test

# solution = [3.5, -0.2]
# sphere(x::Array{Float64}) = sum((x .- solution).^2)
#
# xs = -5.0:0.1:5.0
# ys = -5.0:0.1:5.0
#
# objective = sphere # sphere, himmelblau, styblinski_tang, rastrigin
# fz(x, y) = objective([x, y])
# println(solution) # optimal for sphere and rastrigin
# plot_obj()