using LinearAlgebra
using KrylovKit
using OMEinsum
using Random
using Zygote

"""
tensor order graph: from left to right, top to bottom.
```
a ────┬──── c    a──────┬──────c     a─────b
│     b     │    │      │      │     │     │
├─ d ─┼─ e ─┤    │      b      │     ├──c──┤
│     g     │    │      │      │     │     │
f ────┴──── h    d──────┴──────e     d─────e
```
"""

"""
    i, j = ktoij(k,Ni,Nj)
    LinearIndices -> CartesianIndices
"""
function ktoij(k,Ni,Nj)
    Cart = CartesianIndices((1:Ni,1:Nj))
    Index = Cart[k]
    Index[1],Index[2]
end

safesign(x::Number) = iszero(x) ? one(x) : sign(x)

"""
    qrpos(A)

Returns a QR decomposition, i.e. an isometric `Q` and upper triangular `R` matrix, where `R`
is guaranteed to have positive diagonal elements.
"""
qrpos(A) = qrpos!(copy(A))
function qrpos!(A)
    mattype = _mattype(A)
    F = qr!(mattype(A))
    Q = mattype(F.Q)
    R = F.R
    phases = safesign.(diag(R))
    rmul!(Q, Diagonal(phases))
    lmul!(Diagonal(conj!(phases)), R)
    return Q, R
end

"""
    lqpos(A)

Returns a LQ decomposition, i.e. a lower triangular `L` and isometric `Q` matrix, where `L`
is guaranteed to have positive diagonal elements.
"""
lqpos(A) = lqpos!(copy(A))
function lqpos!(A)
    mattype = _mattype(A)
    F = qr!(mattype(A'))
    Q = mattype(mattype(F.Q)')
    L = mattype(F.R')
    phases = safesign.(diag(L))
    lmul!(Diagonal(phases), Q)
    rmul!(L, Diagonal(conj!(phases)))
    return L, Q
end

function cellones(A)
    Ni, Nj = size(A)
    D = size(A[1,1],1)
    Cell = Array{_arraytype(A[1,1]){ComplexF64,2},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        Cell[i,j] = _mattype(A){ComplexF64}(I, D, D)
    end
    return Cell
end

function ρmap(ρ,Ai,J)
    Nj = size(Ai,1)
    for j = 1:Nj
        jr = J+j-1 - (J+j-1 > Nj)*Nj
        ρ = ein"(dc,csb),dsa -> ab"(ρ,Ai[jr],conj(Ai[jr]))
    end
    return ρ
end

function initialA(M, D)
    Ni, Nj = size(M)
    arraytype = _arraytype(M[1,1])
    A = Array{arraytype{ComplexF64,3},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        d = size(M[i,j], 4)
        A[i,j] = arraytype(rand(ComplexF64, D, d, D))
    end
    return A
end
"""
    getL!(A,L; kwargs...)

````
┌ A1─A2─    ┌      L ─
ρ │  │    = ρ   =  │
┕ A1─A2─    ┕      L'─
````

ρ=L'*L, return L, where `L`is guaranteed to have positive diagonal elements.
L = cholesky!(ρ).U
If ρ is not exactly positive definite, cholesky will fail
"""
function getL!(A,L; kwargs...)
    Ni,Nj = size(A)
    for j = 1:Nj, i = 1:Ni
        λ,ρs,info = eigsolve(ρ->ρmap(ρ,A[i,:],j), L[i,j]'*L[i,j], 1, :LM; ishermitian = false, maxiter = 1, kwargs...)
        @debug "getL eigsolv" λ info sort(abs.(λ))
        info.converged == 0 && @warn "getL not converged"
        ρ = ρs[1] + ρs[1]'
        ρ ./= tr(ρ)
        F = svd!(ρ)
        Lo = lmul!(Diagonal(sqrt.(F.S)), F.Vt)
        _, L[i,j] = qrpos!(Lo)
    end
    return L
end

"""
    getAL(A,L)

Given an MPS tensor `A` and `L` ，return a left-canonical MPS tensor `AL`, a gauge transform `R` and
a scalar factor `λ` such that ``λ AR R = L A``
"""
function getAL(A,L)
    Ni,Nj = size(A)
    arraytype = _arraytype(A[1,1])
    AL = Array{arraytype{ComplexF64,3},2}(undef, Ni, Nj)
    Le = Array{arraytype{ComplexF64,2},2}(undef, Ni, Nj)
    λ = zeros(Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        D, d, = size(A[i,j])
        Q, R = qrpos!(reshape(L[i,j]*reshape(A[i,j], D, d*D), D*d, D))
        AL[i,j] = reshape(Q, D, d, D)
        λ[i,j] = norm(R)
        Le[i,j] = rmul!(R, 1/λ[i,j])
    end
    return AL, Le, λ
end

function getLsped(Le, A, AL; kwargs...)
    Ni,Nj = size(A)
    L = Array{_arraytype(A[1,1]){ComplexF64,2},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        λ , Ls, info = eigsolve(X -> ein"(dc,csb),dsa -> ab"(X,A[i,j],conj(AL[i,j])), Le[i,j], 1, :LM; ishermitian = false, kwargs...)
        @debug "getLsped eigsolve" λ info sort(abs.(λ))
        info.converged == 0 && @warn "getLsped not converged"
        _, L[i,j] = qrpos!(Ls[1])
    end
    return L
end

"""
    leftorth(A,L=cellones(size(A,1),size(A,2),size(A[1,1],1)); tol = 1e-12, maxiter = 100, kwargs...)

Given an MPS tensor `A`, return a left-canonical MPS tensor `AL`, a gauge transform `L` and
a scalar factor `λ` such that ``λ AL L = L A``, where an initial guess for `L` can be
provided.
"""
function leftorth(A,L=cellones(A); tol = 1e-12, maxiter = 100, kwargs...)
    L = getL!(A,L; kwargs...)
    AL, Le, λ = getAL(A,L;kwargs...)
    numiter = 1
    while norm(L.-Le) > tol && numiter < maxiter
        L = getLsped(Le, A, AL; kwargs...)
        AL, Le, λ = getAL(A, L; kwargs...)
        numiter += 1
    end
    L = Le
    return AL, L, λ
end

"""
    rightorth(A,L=cellones(size(A,1),size(A,2),size(A[1,1],1)); tol = 1e-12, maxiter = 100, kwargs...)

Given an MPS tensor `A`, return a gauge transform R, a right-canonical MPS tensor `AR`, and
a scalar factor `λ` such that ``λ R AR^s = A^s R``, where an initial guess for `R` can be
provided.
"""
function rightorth(A,L=cellones(A); tol = 1e-12, maxiter = 100, kwargs...)
    Ni,Nj = size(A)
    arraytype = _arraytype(A[1,1])
    Ar = Array{arraytype{ComplexF64,3},2}(undef, Ni, Nj)
    Lr = Array{arraytype{ComplexF64,2},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        Ar[i,j] = permutedims(A[i,j],(3,2,1))
        Lr[i,j] = permutedims(L[i,j],(2,1))
    end
    AL, L, λ = leftorth(Ar,Lr; tol = tol, maxiter = maxiter, kwargs...)
    R = Array{arraytype{ComplexF64,2},2}(undef, Ni, Nj)
    AR = Array{arraytype{ComplexF64,3},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        R[i,j] = permutedims(L[i,j],(2,1))
        AR[i,j] = permutedims(AL[i,j],(3,2,1))
    end
    return R, AR, λ
end

"""
    LRtoC(L,R)

```
 ── Cᵢⱼ ──  =  ── Lᵢⱼ ── Rᵢⱼ₊₁ ──
```
"""
function LRtoC(L, R)
    Ni, Nj = size(L)
    arraytype = _arraytype(L[1,1])
    C = Array{arraytype{ComplexF64,2},2}(undef, Ni, Nj)
    for j in 1:Nj,i in 1:Ni
        jr = j + 1 - (j + 1 > Nj) * Nj
        C[i,j] = L[i,j] * R[i,jr]
    end
    return C
end

"""
    FLm = FLmap(ALi, ALip, Mi, FL, J)

ALip means ALᵢ₊₁
```
  ┌──        ┌──  ALᵢⱼ  ── ALᵢⱼ₊₁   ──   ...          a ────┬──── c
  │          │     │        │                         │     b     │
 FLm   =   FLᵢⱼ ─ Mᵢⱼ   ── Mᵢⱼ₊₁    ──   ...          ├─ d ─┼─ e ─┤
  │          │     │        │                         │     g     │
  ┕──        ┕──  ALᵢ₊₁ⱼ ─ ALᵢ₊₁ⱼ₊₁ ──   ...          f ────┴──── h
```
"""
function FLmap(ALi, ALip, Mi, FL, J)
    Nj = size(ALi,1)
    FLm = copy(FL)
    for j=1:Nj
        jr = J+j-1 - (J+j-1 > Nj)*Nj
        FLm = ein"((adf,abc),dgeb),fgh -> ceh"(FLm,ALi[jr],Mi[jr],ALip[jr])
    end
    return FLm
end

"""
    FRm = FRmap(ARi, ARip, Mi, FR, J)

ARip means ARᵢ₊₁
```
 ──┐       ... ─── ARᵢⱼ₋₁  ── ARᵢⱼ  ──┐            a ────┬──── c
   │                │          │      │            │     b     │
──FRm  =   ... ──── Mᵢⱼ₋₁  ── Mᵢⱼ  ──FRᵢⱼ          ├─ d ─┼─ e ─┤
   │                │          │      │            │     g     │
 ──┘       ... ─ ARᵢ₊₁ⱼ₋₁ ─ ARᵢ₊₁ⱼ  ──┘            f ────┴──── h
```
"""
function FRmap(ARi, ARip, Mi, FR, J)
    Nj = size(ARi,1)
    FRm = copy(FR)
    for j=1:Nj
        jr = J-(j-1) + (J-(j-1) < 1)*Nj
        FRm = ein"((ceh,abc),dgeb),fgh -> adf"(FRm,ARi[jr],Mi[jr],ARip[jr])
    end
    return FRm
end

function FLint(AL, M)
    Ni,Nj = size(AL)
    arraytype = _arraytype(AL[1,1])
    FL = Array{arraytype{ComplexF64,3},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        D = size(AL[i,j],1)
        dL = size(M[i,j],1)
        FL[i,j] = arraytype(rand(ComplexF64, D, dL, D))
    end
    return FL
end

function FRint(AR, M)
    Ni,Nj = size(AR)
    arraytype = _arraytype(AR[1,1])
    FR = Array{arraytype{ComplexF64,3},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        D = size(AR[i,j],1)
        dR = size(M[i,j],3)
        FR[i,j] = arraytype(rand(ComplexF64, D, dR, D))
    end
    return FR
end

"""
    λL, FL = leftenv(ALu, ALd, M, FL = FLint(ALu,M); kwargs...)

Compute the left environment tensor for MPS A and MPO M, by finding the left fixed point
of ALu - M - ALd contracted along the physical dimension.
```
 ┌──  ALuᵢⱼ  ── ALuᵢⱼ₊₁   ──   ...         ┌──
 │     │        │                          │
FLᵢⱼ ─ Mᵢⱼ   ── Mᵢⱼ₊₁     ──   ...  = λLᵢⱼ FLᵢⱼ
 │     │        │                          │
 ┕──  ALdᵢᵣⱼ  ─ ALdᵢᵣⱼ₊₁  ──   ...         ┕──
```
"""
leftenv(ALu, ALd, M, FL = FLint(ALu,M); kwargs...) = leftenv!(ALu, ALd, M, copy(FL); kwargs...)
function leftenv!(ALu, ALd, M, FL; kwargs...)
    Ni,Nj = size(ALu)
    λL = zeros(eltype(FL[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        ir = i + 1 - Ni * (i==Ni)
        λLs, FL1s, info= eigsolve(X->FLmap(ALu[i,:], conj(ALd[ir,:]), M[i,:], X, j), FL[i,j], 1, :LM; maxiter=1000 , ishermitian = false, kwargs...)
        @debug "leftenv! eigsolv" λLs info sort(abs.(λLs))
        info.converged == 0 && @warn "leftenv not converged"
        if length(λLs) > 1 && norm(abs(λLs[1]) - abs(λLs[2])) < 1e-12
            @show λLs
            if real(λLs[1]) > 0
                FL[i,j] = FL1s[1]
                λL[i,j] = λLs[1]
            else
                FL[i,j] = FL1s[2]
                λL[i,j] = λLs[2]
            end
        else
            FL[i,j] = FL1s[1]
            λL[i,j] = λLs[1]
        end
    end
    return λL, FL
end

"""
    λR, FR = rightenv(ARu, ARd, M, FR = FRint(ARu,M); kwargs...)

Compute the right environment tensor for MPS A and MPO M, by finding the left fixed point
of AR - M - conj(AR) contracted along the physical dimension.
```
   ... ─── ARuᵢⱼ₋₁ ── ARuᵢⱼ  ──┐          ──┐
            │          │       │            │
   ... ──── Mᵢⱼ₋₁  ── Mᵢⱼ   ──FRᵢⱼ  = λRᵢⱼ──FRᵢⱼ
            │          │       │            │
   ... ─   ARdᵢᵣⱼ₋₁ ─ ARdᵢᵣⱼ ──┘          ──┘
```
"""
rightenv(ARu, ARd, M, FR = FRint(ARu,M); kwargs...) = rightenv!(ARu, ARd, M, copy(FR); kwargs...)
function rightenv!(ARu, ARd, M, FR; kwargs...)
    Ni,Nj = size(ARu)
    λR = zeros(eltype(FR[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        ir = i + 1 - Ni * (i==Ni)
        λRs, FR1s, info= eigsolve(X->FRmap(ARu[i,:], conj(ARd[ir,:]), M[i,:], X, j), FR[i,j], 1, :LM;maxiter=1000 , ishermitian = false, kwargs...)
        @debug "rightenv! eigsolv" λRs info sort(abs.(λRs))
        info.converged == 0 && @warn "rightenv not converged"
        if length(λRs) > 1 && norm(abs(λRs[1]) - abs(λRs[2])) < 1e-12
            @show λRs
            if real(λRs[1]) > 0
                FR[i,j] = FR1s[1]
                λR[i,j] = λRs[1]
            else
                FR[i,j] = FR1s[2]
                λR[i,j] = λRs[2]
            end
        else
            FR[i,j] = FR1s[1]
            λR[i,j] = λRs[1]
        end
    end
    return λR, FR
end

"""
    ACm = ACmap(ACij, FLj, FRj, Mj, II)

```
                                ┌─────── ACᵢⱼ ─────┐
                                │        │         │             a ────┬──── c
┌─────── ACm  ─────┐      =     FLᵢⱼ ─── Mᵢⱼ ───── FRᵢⱼ          │     b     │
│        │         │            │        │         │             ├─ d ─┼─ e ─┤
                                FLᵢ₊₁ⱼ ─ Mᵢ₊₁ⱼ ──  FRᵢ₊₁ⱼ        │     g     │
                                │        │         │             f ────┴──── h
                                .        .         .
                                .        .         .
                                .        .         .
```
"""
function ACmap(ACij, FLj, FRj, Mj, II)
    Ni = size(FLj,1)
    ACm = copy(ACij)
    for i=1:Ni
        ir = II+i-1 - (II+i-1 > Ni)*Ni
        ACm = ein"((adf,abc),dgeb),ceh -> fgh"(FLj[ir],ACm,Mj[ir],FRj[ir])
    end
    return ACm
end

"""
    Cmap(Cij, FLjp, FRj, II)

```
                    ┌────Cᵢⱼ ───┐
                    │           │            a ─── b
┌──── Cm ───┐   =   FLᵢⱼ₊₁ ──── FRᵢⱼ         │     │
│           │       │           │            ├─ c ─┤
                    FLᵢ₊₁ⱼ₊₁ ── FRᵢ₊₁ⱼ       │     │
                    │           │            d ─── e
                    .           .
                    .           .
                    .           .
```
"""
function Cmap(Cij, FLjp, FRj, II)
    Ni = size(FLjp,1)
    Cm = copy(Cij)
    for i=1:Ni
        ir = II+i-1 - (II+i-1 > Ni)*Ni
        Cm = ein"(acd,ab),bce -> de"(FLjp[ir],Cm,FRj[ir])
    end
    return Cm
end

"""
    ACenv(AC, FL, M, FR;kwargs...)

Compute the up environment tensor for MPS `FL`,`FR` and MPO `M`, by finding the up fixed point
        of `FL - M - FR` contracted along the physical dimension.
```
┌─────── ACᵢⱼ ─────┐
│        │         │
FLᵢⱼ ─── Mᵢⱼ ───── FRᵢⱼ
│        │         │
FLᵢ₊₁ⱼ ─ Mᵢ₊₁ⱼ ──  FRᵢ₊₁ⱼ  =  λACᵢⱼ ┌──── ACᵢⱼ ───┐
│        │         │                │      │      │
.        .         .
.        .         .
.        .         .
```
"""
ACenv(AC, FL, M, FR; kwargs...) = ACenv!(copy(AC), FL, M, FR; kwargs...)
function ACenv!(AC, FL, M, FR; kwargs...)
    Ni,Nj = size(AC)
    λAC = zeros(eltype(AC[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        λACs, ACs, info = eigsolve(X->ACmap(X, FL[:,j], FR[:,j], M[:,j], i), AC[i,j], 1, :LM; maxiter=1000 ,ishermitian = false, kwargs...)
        @debug "ACenv! eigsolve" λACs info sort(abs.(λACs))
        info.converged == 0 && @warn "ACenv Not converged"
        if length(λACs) > 1 && norm(abs(λACs[1]) - abs(λACs[2])) < 1e-12
            @show λACs
            if real(λACs[1]) > 0
                AC[i,j] = ACs[1]
                λAC[i,j] = λACs[1]
            else
                AC[i,j] = ACs[2]
                λAC[i,j] = λACs[2]
            end
        else
            AC[i,j] = ACs[1]
            λAC[i,j] = λACs[1]
        end
    end
    return λAC, AC
end

"""
    Cenv(C, FL, FR;kwargs...)

Compute the up environment tensor for MPS `FL` and `FR`, by finding the up fixed point
    of `FL - FR` contracted along the physical dimension.
```
┌────Cᵢⱼ ───┐
│           │
FLᵢⱼ₊₁ ──── FRᵢⱼ
│           │
FLᵢ₊₁ⱼ₊₁ ── FRᵢ₊₁ⱼ   =  λCᵢⱼ ┌──Cᵢⱼ ─┐
│           │                │       │
.           .
.           .
.           .
```
"""
Cenv(C, FL, FR; kwargs...) = Cenv!(copy(C), FL, FR; kwargs...)
function Cenv!(C, FL, FR; kwargs...)
    Ni,Nj = size(C)
    λC = zeros(eltype(C[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        jr = j + 1 - (j==Nj) * Nj
        λCs, Cs, info = eigsolve(X->Cmap(X, FL[:,jr], FR[:,j], i), C[i,j], 1, :LM; maxiter=1000 ,ishermitian = false, kwargs...)
        @debug "Cenv! eigsolve" λCs info sort(abs.(λCs))
        info.converged == 0 && @warn "Cenv Not converged"
        if length(λCs) > 1 && norm(abs(λCs[1]) - abs(λCs[2])) < 1e-12
            @show λCs
            if real(λCs[1]) > 0
                C[i,j] = Cs[1]
                λC[i,j] = λCs[1]
            else
                C[i,j] = Cs[2]
                λC[i,j] = λCs[2]
            end
        else
            C[i,j] = Cs[1]
            λC[i,j] = λCs[1]
        end
    end
    return λC, C
end

function ACCtoAL(ACij,Cij)
    D,d, = size(ACij)
    QAC, RAC = qrpos(reshape(ACij,(D*d, D)))
    QC, RC = qrpos(Cij)
    errL = norm(RAC-RC)
    # @show errL
    reshape(QAC*QC', (D, d, D)), errL
end

function ACCtoAR(ACij,Cijr)
    D,d, = size(ACij)
    LAC, QAC = lqpos(reshape(ACij,(D, d*D)))
    LC, QC = lqpos(Cijr)
    errR = norm(LAC-LC)
    # @show errR
    reshape(QC'*QAC, (D, d, D)), errR
end

"""
    itoir(i,Ni,Nj)

````
i -> (i,j) -> (i,jr) -> ir
````
"""
function itoir(i,Ni,Nj)
    Liner = LinearIndices((1:Ni,1:Nj))
    Cart = CartesianIndices((1:Ni,1:Nj))
    Index = Cart[i]
    i,j = Index[1],Index[2]
    jr = j - 1 + (j==1)*Nj
    Liner[i,jr]
end

function ALCtoAC(AL,C)
    Ni,Nj = size(AL)
    ACij = [ein"asc,cb -> asb"(AL[i],C[i]) for i=1:Ni*Nj]
    reshape(ACij,Ni,Nj)
end

"""
    AL, AR = ACCtoALAR(AC, C)

QR factorization to get `AL` and `AR` from `AC` and `C`

````
──ALᵢⱼ──Cᵢⱼ──  =  ──ACᵢⱼ──  = ──Cᵢ₋₁ⱼ ──ARᵢⱼ──
  │                  │                  │
````
"""
function ACCtoALAR(AC, C)
    Ni,Nj = size(AC)
    ALijerrL = [ACCtoAL(AC[i],C[i]) for i=1:Ni*Nj]
    AL = reshape([ALijerrL[i][1] for i=1:Ni*Nj],Ni,Nj)
    errL = Zygote.@ignore sum([ALijerrL[i][2] for i=1:Ni*Nj])
    ARijerrR = [ACCtoAR(AC[i],C[itoir(i,Ni,Nj)]) for i=1:Ni*Nj]
    AR = reshape([ARijerrR[i][1] for i=1:Ni*Nj],Ni,Nj)
    errR = Zygote.@ignore sum([ARijerrR[i][2] for i=1:Ni*Nj])
    return AL, AR, errL, errR
end

"""
Compute the error through all environment `AL,C,AR,FL,M,FR`

````
        ┌── AC──┐
        │   │   │           ┌── AC──┐
MAC1 =  FL─ M ──FR  =  λAC  │   │   │
        │   │   │

        ┌── AC──┐
        │   │   │           ┌──C──┐
MAC2 =  FL─ M ──FR  =  λAC  │     │
        │   │   │
         ── AR──┘

── MAC1 ──    ≈    ── MAC2 ── AR ──
    │                         │
````
"""
function error(AL,C,AR,FL,M,FR)
    Ni,Nj = size(AL)
    AC = ALCtoAC(AL, C)
    err = 0
    for j = 1:Nj, i = 1:Ni
        MAC = ACmap(AC[i,j], FL[:,j], FR[:,j], M[:,j], i)
        MAC -= ein"(apc,dpc),dsb -> asb"(MAC,conj(AR[i,j]),AR[i,j])
        err += norm(MAC)
    end
    return err
end

"""
    λL, FL = obs_FL(ALu, ALd, M, FL = FLint(AL,M); kwargs...)

Compute the observable left environment tensor for MPS A and MPO M, by finding the left fixed point
of AL - M - conj(AL) contracted along the physical dimension.
```
 ┌──  ALuᵢⱼ  ── ALuᵢⱼ₊₁   ──   ...         ┌──
 │     │        │                          │
FLᵢⱼ ─ Mᵢⱼ   ── Mᵢⱼ₊₁     ──   ...  = λLᵢⱼ FLᵢⱼ
 │     │        │                          │
 ┕──  ALdᵢᵣⱼ  ─ ALdᵢᵣⱼ₊₁  ──   ...         ┕──
```
"""
obs_FL(ALu, ALd, M, FL = FLint(ALu,M); kwargs...) = obs_FL!(ALu, ALd, M, copy(FL); kwargs...)
function obs_FL!(ALu, ALd, M, FL; kwargs...)
    Ni,Nj = size(ALu)
    λL = zeros(eltype(FL[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        ir = Ni + 1 - i
        λLs, FL1s, info= eigsolve(X->FLmap(ALu[i,:], ALd[ir,:], M[i,:], X, j), FL[i,j], 1, :LM; ishermitian = false, kwargs...)
        @debug "obs_FL eigsolve" λLs info sort(abs.(λLs))
        info.converged == 0 && @warn "obs_FL Not converged"
        if length(λLs) > 1 && norm(abs(λLs[1]) - abs(λLs[2])) < 1e-12
            @show λLs
            if real(λLs[1]) > 0
                FL[i,j] = FL1s[1]
                λL[i,j] = λLs[1]
            else
                FL[i,j] = FL1s[2]
                λL[i,j] = λLs[2]
            end
        else
            FL[i,j] = FL1s[1]
            λL[i,j] = λLs[1]
        end
    end
    return λL, FL
end

"""
    λR, FR = obs_FR(ARu, ARd, M, FR = FRint(ARu,M); kwargs...)

Compute the observable right environment tensor for MPS A and MPO M, by finding the left fixed point
of AR - M - conj(AR) contracted along the physical dimension.
```
   ... ─── ARuᵢⱼ₋₁ ── ARuᵢⱼ  ──┐          ──┐
            │          │       │            │
   ... ──── Mᵢⱼ₋₁  ── Mᵢⱼ   ──FRᵢⱼ  = λRᵢⱼ──FRᵢⱼ
            │          │       │            │
   ... ─   ARdᵢᵣⱼ₋₁ ─ ARdᵢᵣⱼ ──┘          ──┘
```
"""
obs_FR(ARu, ARd, M, FR = FRint(ARu,M); kwargs...) = obs_FR!(ARu, ARd, M, copy(FR); kwargs...)
function obs_FR!(ARu, ARd, M, FR; kwargs...)
    Ni,Nj = size(ARu)
    λR = zeros(eltype(FR[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        ir = Ni + 1 - i
        λRs, FR1s, info= eigsolve(X->FRmap(ARu[i,:], ARd[ir,:], M[i,:], X, j), FR[i,j], 1, :LM; ishermitian = false, kwargs...)
        @debug "obs_FR! eigsolve" λRs info sort(abs.(λRs))
        info.converged == 0 && @warn "obs_FR! Not converged"
        if length(λRs) > 1 && norm(abs(λRs[1]) - abs(λRs[2])) < 1e-12
            @show λRs
            if real(λRs[1]) > 0
                FR[i,j] = FR1s[1]
                λR[i,j] = λRs[1]
            else
                FR[i,j] = FR1s[2]
                λR[i,j] = λRs[2]
            end
        else
            FR[i,j] = FR1s[1]
            λR[i,j] = λRs[1]
        end
    end
    return λR, FR
end

function BgFLint(AL, M)
    Ni,Nj = size(AL)
    arraytype = _arraytype(AL[1,1])
    BgFL = Array{arraytype{ComplexF64,4},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        ir = i + 1 - Ni * (i==Ni)
        irr = i + 2 - Ni * (i + 2 > Ni)
        D1 = size(AL[i,j],1)
        D2 = size(AL[irr,j],1)
        dL1 = size(M[i,j],1)
        dL2 = size(M[ir,j],1)
        BgFL[i,j] = arraytype(rand(ComplexF64, D1, dL1, dL2, D2))
    end
    return BgFL
end

"""
    BgFLm = BgFLmap(ALi, ALip, Mi, Mip, BgFLij, J)

ALip means ALᵢ₊₂
Mip means Mᵢ₊₁
```
  ┌──        ┌──  ALᵢⱼ  ── ALᵢⱼ₊₁    ──   ...       a ────┬──── c
  │          │     │        │                       │     b     │
  │          │  ─ Mᵢⱼ   ── Mᵢⱼ₊₁     ──   ...       ├─ d ─┼─ e ─┤
BgFLm   =  BgFLᵢⱼ  │        │                       │     f     │
  │          │  ─ Mᵢ₊₁ⱼ ── Mᵢ₊₁ⱼ₊₁   ──   ...       ├─ g ─┼─ h ─┤
  │          │     │        │                       │     j     │
  ┕──        ┕──  ALᵢ₊₂ⱼ ─ ALᵢ₊₂ⱼ₊₁  ──   ...       i ────┴──── k
```
"""
function BgFLmap(ALi, ALip, Mi, Mip, BgFLij, J)
    Nj = size(ALi,1)
    BgFLm = copy(BgFLij)
    for j=1:Nj
        jr = J+j-1 - (J+j-1 > Nj)*Nj
        BgFLm = ein"(((adgi,abc),dfeb),gjhf),ijk -> cehk"(BgFLm,ALi[jr],Mi[jr],Mip[jr],ALip[jr])
    end
    return BgFLm
end

"""
    λL, BgFL = bigleftenv(ALu, ALd, M, BgFL = BgFLint(AL,M); kwargs...)

Compute the up and down left environment tensor for MPS A and MPO M, by finding the left fixed point
of AL - M - M - conj(AL) contracted along the physical dimension.
```
   ┌──  ALuᵢⱼ ── ALuᵢⱼ₊₁   ──   ...           ┌──
   │     │        │                           │
   │  ─ Mᵢⱼ   ── Mᵢⱼ₊₁     ──   ...           │
 BgFLᵢⱼ  │        │                   = λLᵢⱼ BgFLᵢⱼ
   │  ─ Mᵢ₊₁ⱼ ── Mᵢ₊₁ⱼ₊₁   ──   ...           │
   │     │        │                           │
   ┕── ALdᵢ₊₂ⱼ ─ ALdᵢ₊₂ⱼ₊₁ ──   ...           ┕──
```
"""
bigleftenv(ALu, ALd, M, BgFL = BgFLint(ALu,M); kwargs...) = bigleftenv!(ALu, ALd, M, copy(BgFL); kwargs...)
function bigleftenv!(ALu, ALd, M, BgFL; kwargs...)
    Ni,Nj = size(ALu)
    λL = zeros(eltype(BgFL[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        ir = i + 1 - Ni * (i==Ni)
        irr = i + 2 - Ni * (i + 2 > Ni)
        λLs, BgFL1s, info= eigsolve(X->BgFLmap(ALu[i,:], ALd[irr,:], M[i,:], M[ir,:], X, j), BgFL[i,j], 1, :LM; ishermitian = false, kwargs...)
        @debug "bigleftenv! eigsolve" λLs info sort(abs.(λLs))
        if length(λLs) > 1 && norm(abs(λLs[1]) - abs(λLs[2])) < 1e-12
            @show λLs
            if real(λLs[1]) > 0
                BgFL[i,j] = BgFL1s[1]
                λL[i,j] = λLs[1]
            else
                BgFL[i,j] = BgFL1s[2]
                λL[i,j] = λLs[2]
            end
        else
            BgFL[i,j] = BgFL1s[1]
            λL[i,j] = λLs[1]
        end
    end
    return λL, BgFL
end

function BgFRint(AR, M)
    Ni,Nj = size(AR)
    arraytype = _arraytype(AR[1,1])
    BgFR = Array{arraytype{ComplexF64,4},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        ir = i + 1 - Ni * (i==Ni)
        irr = i + 2 - Ni * (i + 2 > Ni)
        D1 = size(AR[i,j],3)
        D2 = size(AR[irr,j],3)
        dR1 = size(M[i,j],3)
        dR2 = size(M[ir,j],3)
        BgFR[i,j] = arraytype(rand(ComplexF64, D1, dR1, dR2, D2))
    end
    return BgFR
end

"""
    FRm = FRmap(ARi, ARip, Mi, FR, J)

ARip means ARᵢ₊₁
```
 ──┐          ...  ─── ARᵢⱼ₋₁  ── ARᵢⱼ  ──┐        a ────┬──── c
   │                    │          │      │        │     b     │
 ──│          ... ──── Mᵢⱼ₋₁  ──  Mᵢⱼ   ──│        ├─ d ─┼─ e ─┤
  BgFRm   =             │          │     BgFRm     │     f     │
 ──│          ... ──── Mᵢ₊₁ⱼ₋₁ ── Mᵢ₊₁ⱼ ──│        ├─ g ─┼─ h ─┤
   │                    │          │      │        │     j     │
 ──┘          ...  ─ ARᵢ₊₂ⱼ₋₁ ─── ARᵢ₊₂ⱼ──┘        i ────┴──── k
```
"""
function BgFRmap(ARi, ARip, Mi, Mip, BgFR, J)
    Nj = size(ARi,1)
    BgFRm = copy(BgFR)
    for j=1:Nj
        jr = J-(j-1) + (J-(j-1) < 1)*Nj
        BgFRm = ein"(((cehk,abc),dfeb),gjhf),ijk -> adgi"(BgFRm,ARi[jr],Mi[jr],Mip[jr],ARip[jr])
    end
    return BgFRm
end

"""
    λR, BgFR = bigrightenv(ARu, ARd, M, BgFR = BgFRint(ARu,M); kwargs...)

Compute the up and down right environment tensor for MPS A and MPO M, by finding the left fixed point
of AR - M - M - conj(AR) contracted along the physical dimension.
```
     ──┐          ...  ─── ARᵢⱼ₋₁  ── ARᵢⱼ  ──┐
       │                    │          │      │
     ──│          ... ──── Mᵢⱼ₋₁  ──  Mᵢⱼ   ──│
λRᵢⱼ BgFRᵢⱼ   =             │          │     BgFRᵢⱼ
     ──│          ... ──── Mᵢ₊₁ⱼ₋₁ ── Mᵢ₊₁ⱼ ──│
       │                    │          │      │
     ──┘          ...  ─ ARᵢ₊₂ⱼ₋₁ ─── ARᵢ₊₂ⱼ──┘
```
"""
bigrightenv(ARu, ARd, M, BgFR = BgFRint(ARu,M); kwargs...) = bigrightenv!(ARu, ARd, M, copy(BgFR); kwargs...)
function bigrightenv!(ARu, ARd, M, BgFR; kwargs...)
    Ni,Nj = size(ARu)
    λR = zeros(eltype(BgFR[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        ir = i + 1 - Ni * (i==Ni)
        irr = i + 2 - Ni * (i + 2 > Ni)
        λRs, BgFR1s, info= eigsolve(X->BgFRmap(ARu[i,:], ARd[irr,:], M[i,:], M[ir,:], X, j), BgFR[i,j], 1, :LM; ishermitian = false, kwargs...)
        @debug "bigrightenv eigsolve" λRs info sort(abs.(λRs))
        if length(λRs) > 1 && norm(abs(λRs[1]) - abs(λRs[2])) < 1e-12
            @show λRs
            if real(λRs[1]) > 0
                BgFR[i,j] = BgFR1s[1]
                λR[i,j] = λRs[1]
            else
                BgFR[i,j] = BgFR1s[2]
                λR[i,j] = λRs[2]
            end
        else
            BgFR[i,j] = BgFR1s[1]
            λR[i,j] = λRs[1]
        end
    end
    return λR, BgFR
end

function norm_FLint(AL)
    Ni,Nj = size(AL)
    arraytype = _arraytype(AL[1,1])
    norm_FL = Array{arraytype{ComplexF64,2},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        D = size(AL[i,j],1)
        norm_FL[i,j] = arraytype(rand(ComplexF64, D, D))
    end
    return norm_FL
end

"""
    FL_normm = norm_FLmap(ALui, ALdi, FL_norm, J)

```
   ┌──        ┌──  ALuᵢⱼ ── ALuᵢⱼ₊₁ ──  ...     a───┬───c
  FLm   =    FLm     │        │                 │   b   │
   ┕──        ┕──  ALdᵢⱼ ── ALdᵢⱼ₊₁ ──  ...     d───┴───e
```
"""
function norm_FLmap(ALui, ALdi, FL_norm, J)
    Nj = size(ALui,1)
    FL_normm = copy(FL_norm)
    for j=1:Nj
        jr = J+j-1 - (J+j-1 > Nj)*Nj
        FL_normm = ein"(ad,abc),dbe -> ce"(FL_normm,ALui[jr],ALdi[jr])
    end
    return FL_normm
end

norm_FL(ALu, ALd, FL_norm = norm_FLint(ALu); kwargs...) = norm_FL!(ALu, ALd, FL_norm; kwargs...)
function norm_FL!(ALu, ALd, FL_norm; kwargs...)
    Ni,Nj = size(ALu)
    λL = zeros(eltype(FL_norm[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        λLs, FL_norms, info= eigsolve(X->norm_FLmap(ALu[i,:], ALd[i,:], X, j), FL_norm[i,j], 1, :LM; ishermitian = false, kwargs...)
        @debug "norm_FL eigsolve" λLs info sort(abs.(λLs))
        if length(λLs) > 1 && norm(abs(λLs[1]) - abs(λLs[2])) < 1e-12
            @show λLs
            if real(λLs[1]) > 0
                FL_norm[i,j] = FL_norms[1]
                λL[i,j] = λLs[1]
            else
                FL_norm[i,j] = FL_norms[2]
                λL[i,j] = λLs[2]
            end
        else
            FL_norm[i,j] = FL_norms[1]
            λL[i,j] = λLs[1]
        end
    end
    return λL, FL_norm
end

function norm_FRint(AR)
    Ni,Nj = size(AR)
    arraytype = _arraytype(AR[1,1])
    norm_FR = Array{arraytype{ComplexF64,2},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        D = size(AR[i,j],1)
        norm_FR[i,j] = arraytype(rand(ComplexF64, D, D))
    end
    return norm_FR
end

"""
    FR_normm = norm_FRmap(ARui, ARdi, FR_norm, J)

```
──┐       ... ─── ARᵢⱼ₋₁  ── ARᵢⱼ  ──┐      a───┬───c
 FRm  =            │          │     FRm     │   b   │
──┘       ... ─── ARᵢⱼ₋₁ ─── ARᵢⱼ  ──┘      d───┴───e
```
"""
function norm_FRmap(ARui, ARdi, FR_norm, J)
    Nj = size(ARui,1)
    FR_normm = copy(FR_norm)
    for j=1:Nj
        jr = J-(j-1) + (J-(j-1) < 1)*Nj
        FR_normm = ein"(ce,abc),dbe -> ad"(FR_normm,ARui[jr],ARdi[jr])
    end
    return FR_normm
end

norm_FR(ARu, ARd, FR_norm = norm_FRint(ARu); kwargs...) = norm_FR!(ARu, ARd, FR_norm; kwargs...)
function norm_FR!(ARu, ARd, FR_norm; kwargs...)
    Ni,Nj = size(ARu)
    λL = zeros(eltype(FR_norm[1,1]),Ni,Nj)
    for j = 1:Nj, i = 1:Ni
        λLs, FR_norms, info= eigsolve(X->norm_FRmap(ARu[i,:], ARd[i,:], X, j), FR_norm[i,j], 1, :LM; ishermitian = false, kwargs...)
        @debug "norm_FR! eigsolve" λLs info sort(abs.(λLs))
        if length(λLs) > 1 && norm(abs(λLs[1]) - abs(λLs[2])) < 1e-12
            @show λLs
            if real(λLs[1]) > 0
                FR_norm[i,j] = FR_norms[1]
                λL[i,j] = λLs[1]
            else
                FR_norm[i,j] = FR_norms[2]
                λL[i,j] = λLs[2]
            end
        else
            FR_norm[i,j] = FR_norms[1]
            λL[i,j] = λLs[1]
        end
    end
    return λL, FR_norm
end


"""
    RMSEA(sem_fit::SemFit)

Return the RMSEA.
"""
function RMSEA end

RMSEA(sem_fit::SemFit{Mi, So, St, Mo, O} where {Mi, So, St, Mo <: AbstractSemSingle, O}) =
    RMSEA(df(sem_fit), χ²(sem_fit), n_obs(sem_fit))

RMSEA(sem_fit::SemFit{Mi, So, St, Mo, O} where {Mi, So, St, Mo <: SemEnsemble, O}) =
    sqrt(length(sem_fit.model.sems))*RMSEA(df(sem_fit), χ²(sem_fit), n_obs(sem_fit))

function RMSEA(df, chi2, n_obs)
    rmsea = (chi2 - df) / (n_obs*df)
    rmsea > 0 ? nothing : rmsea = 0
    return sqrt(rmsea)
end


begin
    using Distributions
    using LinearAlgebra
    using Random
end;

Base.@kwdef struct Foo
    # Data
    y::Vector{Float64}
    X::Matrix{Float64}
    d::Vector{Int}
    # Transformed data
    N::Int = size(X, 1)
    K::Int = size(X, 2)
    J::Int = length(unique(d))
    Xvec::Vector{Vector{Float64}} = [X[i, :] for i in 1:N]
    # Hyperparameters
    B0::Symmetric{Float64, Matrix{Float64}} = Symmetric(9.0 * I(K))
    m0::Vector{Float64} = zeros(K)
    a0::Float64 = 0.01
    b0::Float64 = 0.01
    # Parameters
    β::Vector{Vector{Float64}} = [zeros(K) for _ in 1:J]
    v::Vector{Float64} = [1.0 for _ in 1:J]
    # Transformed parameters
    n::Vector{Int} = [sum(d .== j) for j in 1:J]
    XX::Vector{Matrix{Float64}} = [zeros(K, K) for _ in 1:J]
    Xy::Vector{Vector{Float64}} = [zeros(K) for _ in 1:J]
    yy::Vector{Float64} = [0.0 for _ in 1:J]
end;

function step!(model::Foo)
    (; N, J, d, y, Xvec, β, v, m0, B0, a0, b0, XX, Xy, yy, n) = model
    # Refresh XX and Xy
    for j in 1:J
        XX[j] .= 0.0
        Xy[j] .= 0.0
        yy[j] = 0.0
    end
    @inbounds for i in 1:N
        BLAS.syr!('L', 1.0, Xvec[i], XX[d[i]]) # XX[d[i]] += x[i] * x[i]'
        BLAS.axpy!(y[i], Xvec[i], Xy[d[i]])    # Xy[d[i]] += x[i] * y[i]
        yy[d[i]] += y[i]^2
    end

    # Compute the posterior hyperparameters
    for j in 1:J
        B1 = inv(Symmetric(XX[j], :L) + inv(B0))
        m1 = B1 * (Xy[j] + B0 \ m0)
        a1 = a0 + n[j] / 2
        b1 = b0 + (yy[j] + m0' * (B0 \ m0) - m1' * (B1 \ m1)) / 2
        v[j] = rand(InverseGamma(a1, b1))
        β[j] = rand(MvNormal(m1, v[j] * B1))
    end
    return nothing
end

function sample!(model::Foo; mcmcsize = 4000, burnin = mcmcsize ÷ 2, thin = 1)
    (; K, J, β, v) = model
    chainβ = [zeros(K, J) for _ in 1:(mcmcsize - burnin) ÷ thin]
    chainv = [zeros(J) for _ in 1:(mcmcsize - burnin) ÷ thin]
    for iter in 1:mcmcsize
        step!(model)
        if iter > burnin && iszero((iter - burnin) % thin)
            for j in 1:J
                chainβ[(iter - burnin) ÷ thin][:, j] .= β[j]
                chainv[(iter - burnin) ÷ thin][j] = v[j]
            end
        end
    end
    return chainβ, chainv
end

begin
    Random.seed!(1)
    N = 300
    K = 5
    J = 3
    ϵ = randn(N)
    X = randn(N, K)
    d = rand(1:J, N)
    β = ones(K)
    y = d .* (X * β + ϵ)
    model = Foo(; y, X, d)
    chainβ, chainv = sample!(model; mcmcsize = 10000)
end
mean(chainβ)
mean(chainv)


module STFT
"""
Compute the short-time Fourier transform of a signal
Parameters:
    -- data: The signal to find the STFT of
    -- windowFunc: An array sampled from the window function to be applied before taking the FFTs. The length of this function defines the window length
    -- windowStep: The step ( in samples) to move the window between FFTs
    -- low: The lowest harmonic to include in the output
    -- high: The highest harmonic to include in the output
    -- transformFunction: A function to apply to the windowed, FFT'd pieces of _data_
    -- processingFunction: A function to do the processing. This function is passed
       the window index and the windowed, FFT'd piece and must return a value of type
       outType. For instance, to compute the amplitude along some path, processingFunction should be (dat, i)->dat[path[i]]
"""
function stft_processed{dataType,windowType,outType}(data::AbstractArray{dataType, 1}, windowFunc::AbstractArray{windowType, 1}, windowStep::Integer, low::Integer, high::Integer, transformFunction::Function, processingFunction::Function, outT::Type{outType})
    dataLength = size(data, 1)
    windowLength = size(windowFunc)[1]
    nHarmonics = 1 + high - low
    nSnippets = Int(floor((dataLength-windowLength)/windowStep))
    endSample = nSnippets*windowStep
    if low > windowLength
	  throw(ArgumentError("The low harmonic (parameter 'low') must be less than the length of the window"))
    end
    if low < 1
	    throw(ArgumentError("The low harmonic (parameter 'low') must be at least 1"))
    end
    STFT_processed = Array{outType, 1}(nSnippets)
    temp = Array{Float64, 1}(windowLength)
    tempFFT = Array{Float64, 1}(windowLength)
    tempFFTProcessed = Array{Float64, 1}(nHarmonics)
    windowPos = 1
    harmonic = 1
    plan = plan_rfft(temp, flags = FFTW.MEASURE)

    for sample in 1:windowStep:endSample
        temp = windowFunc.*view(data, sample:(sample+windowLength-1))
        tempFFT = plan*temp

        harmonic = 1
        for i in low:high
            tempFFTProcessed[harmonic] = transformFunction(tempFFT[i].re, tempFFT[i].im)
	    STFT_processed[windowPos] = processingFunction(tempFFTProcessed, windowPos)
            harmonic += 1
        end

        windowPos += 1
    end

    return STFT_processed
end

"""
Compute the short-time Fourier transform of a signal
Parameters:
    -- data: The signal to find the STFT of
    -- windowFunc: An array sampled from the window function to be applied before taking the FFTs. The length of this function defines the window length
    -- windowStep: The step ( in samples) to move the window between FFTs
    -- low: The lowest harmonic to include in the output
    -- high: The highest harmonic to include in the output
    -- transformFunction: A function to apply to the windowed, FFT'd pieces of _data_
"""
function stft{dataType,windowType}(data::AbstractArray{dataType, 1}, windowFunc::AbstractArray{windowType, 1}, windowStep::Integer, low::Integer, high::Integer, transformFunction::Function)
    dataLength = size(data, 1)
    windowLength = size(windowFunc)[1]
    nHarmonics = 1 + high - low
    nSnippets = Int(floor((dataLength-windowLength)/windowStep))
    endSample = nSnippets*windowStep
    if low > windowLength
	  throw(ArgumentError("The low harmonic (parameter 'low') must be less than the length of the window"))
    end
    if low < 1
	    throw(ArgumentError("The low harmonic (parameter 'low') must be at least 1"))
    end
    STFT = Array{Float64, 2}(nHarmonics, nSnippets)
    temp = Array{Float64, 1}(windowLength)
    tempFFT = Array{Float64, 1}(windowLength)
    windowPos = 1
    harmonic = 1
    plan = plan_rfft(temp, flags = FFTW.MEASURE)

    for sample in 1:windowStep:endSample
        temp = windowFunc.*view(data, sample:(sample+windowLength-1))
        tempFFT = plan*temp

        harmonic = 1
        for i in low:high
            STFT[harmonic, windowPos] = transformFunction(tempFFT[i].re, tempFFT[i].im)
            harmonic += 1
        end

        windowPos += 1
    end

    return STFT
end

"""Compute the short-time Fourier transform, using a linear detrend on each windowed snippet"""
function detrended_stft{dataType<:Any,windowType<:Any}(data::AbstractArray{dataType, 1}, startSample::Integer, endSample::Integer, windowFunc::AbstractArray{windowType, 1}, windowStep::Integer, low::Integer, high::Integer, transformFunction::Function)
  if endSample == -1
    endSample = size(data, 1)
  end

    dataLength = size(data, 1)
    usedDataLength = endSample - startSample
    windowLength = size(windowFunc)[1]
    nHarmonics = 1 + high - low
    nSnippets = Int(floor((usedDataLength-windowLength)/windowStep))
    endSample = nSnippets*windowStep
    if low > windowLength
	  throw(ArgumentError("The low harmonic (parameter 'low') must be less than the length of the window"))
    end
    if low < 1
	    throw(ArgumentError("The low harmonic (parameter 'low') must be at least 1"))
    end
    STFT = Array{Float64, 2}(nHarmonics, nSnippets)
    temp = Array{Float64, 1}(windowLength)
    tempFFT = Array{Float64, 1}(windowLength)
    windowPos = 1
    harmonic = 1
    plan = plan_rfft(temp, flags = FFTW.MEASURE)

    for sample in startSample:windowStep:endSample
	temp = collect(view(data, sample:(sample+windowLength-1)))
	(a, b) = linreg(collect(1.0:length(temp)), temp)
	for i in eachindex(temp)
	    temp[i] -= (a+b*i) #FIXME: Do we really want to subtract out the intercept here?
	    temp[i] *= windowFunc[i]
	end
        tempFFT = plan*temp

        harmonic = 1
        for i in low:high
            STFT[harmonic, windowPos] = transformFunction(tempFFT[i].re, tempFFT[i].im)
            harmonic += 1
        end

        windowPos += 1
    end

    return STFT
end

include("centroid.jl")
include("envelope.jl")
include("windows.jl")
include("transformFunctions.jl")

export stft
export detrended_stft
end


"""
    drawdowns(pnl; compound)

Calculates the drawdown based on a returns time series.

# Arguments
- `returns`:    Vector of asset returns (usually log-returns).
- `geometric`:  Use geometric compounding (cumprod) if set to true, otherwise simple arithmetic sum (cumsum), e.g. for log-returns (default).
"""
function drawdowns(returns; geometric::Bool=false)
    if geometric
        returns = cumprod(1.0 .+ returns)
    else
        returns = 1.0 .+ cumsum(returns)
    end

    # cumulative returns needs to start at 1.0 for cumulative max to be correct
    returns_max = accumulate(max, [1.0; returns])[2:end]
    returns ./ returns_max .- 1.0
end

"""
    drawdowns_pnl(pnl)

Calculates the drawdown based on a Profit-and-Loss (PnL) time series, e.g. daily equity changes in USD.

# Arguments
- `pnl`:    Vector of Profit-and-Loss (PnL) values, e.g. daily equity changes in USD.
"""
function drawdowns_pnl(pnl)
    pnl = cumsum(pnl)

    # cumulative PnL needs to start at 0.0 for cumulative max to be correct
    pnl_max = accumulate(max, [0.0; pnl])[2:end]
    pnl .- pnl_max
end


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Tests related to simplified dipole model of the Earth geomagnetic field.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] http://helios.fmi.fi/~juusolal/geomagnetism/Lectures/Chapter3_dipole.pdf
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

# File: ./src/earth/geomagnetic_field_models/dipole
# =================================================

# Function: geomag_dipole
# -----------------------

@testset "Function geomag_dipole" begin

    d2r = π/180

    # Test 1
    # --------------------------------------------------------------------------
    #
    # Position aligned with the dipole moment vector.

    year     = SatelliteToolbox._dipole_pole[32,1]
    pole_lat = SatelliteToolbox._dipole_pole[32,2] * d2r
    pole_lon = SatelliteToolbox._dipole_pole[32,3] * d2r
    m        = SatelliteToolbox._dipole_mag[32,2]  * 1e22

    # Distance from the Earth center.
    r = R0 + 196e3

    # Compute using the simplified equations [1] [nT].
    k0     = 1e-7m
    B_norm = 2k0 / r^3 * 1e9

    # Rotate to ECEF.
    Deg = angle_to_dcm(-(π / 2 - pole_lat), -pole_lon, 0, :YZX)
    B_g = SVector{3}(0, 0, -B_norm)
    B_e_expected = Deg * B_g

    # Compute using the function.
    r_g = SVector{3}(0, 0, r)
    r_e = Deg * r_g

    B_e_result_f64 = geomag_dipole(r_e, year)
    B_e_result_f32 = geomag_dipole(Float32.(r_e), year)

    @test B_e_expected ≈ B_e_result_f64 atol=1e-9
    @test eltype(B_e_result_f64) === Float64

    @test B_e_expected ≈ B_e_result_f32 atol=1e-1
    @test eltype(B_e_result_f32) === Float32

    # Test 2
    # --------------------------------------------------------------------------
    #
    # Position at the magnetic Equator, which must have half the magnitude of
    # that of the test 1.

    B_g = SVector{3}(0, 0, B_norm / 2)
    B_e_expected = Deg * B_g

    r_g = SVector{3}(r, 0, 0)
    r_e = Deg * r_g

    B_e_result_f64 = geomag_dipole(r_e, year)
    B_e_result_f32 = geomag_dipole(Float32.(r_e), year)

    @test B_e_expected ≈ B_e_result_f64 atol=1e-9
    @test eltype(B_e_result_f64) === Float64

    @test B_e_expected ≈ B_e_result_f32 atol=1e-1
    @test eltype(B_e_result_f32) === Float32
end


#!/usr/bin/env julia
#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 1, Zadanie 4.
#

fl = Float64

# define the step
δ = fl(2^(-52))

x = fl(1 + δ)

# go through all the consecutive numbers
for i in range(1, length=2^52-1)
    if fl(x * fl(1/x)) != 1
        # if the condition is fulfilled print out the number in question and break the loop
        println(
            bitstring(x), "\n",
            bitstring(fl(1/x)), "\n",
            fl(x * fl(1/x)) != 1, "\n",
            x, "\n",
            fl(1/x), "\n",
            )
        break
    end
    # otherwise continue
    global x = fl(x + fl(δ * i))
end


### A Pluto.jl notebook ###
# v0.12.10

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 9bc349e8-efb0-11ea-1392-cd735828af33
# packages to load etc
begin
	using DelimitedFiles
	using NIfTI
	using Images
	using ImageView
	using LinearAlgebra
	using StatsBase
	using GLM
	using DataFrames
	using Plots
	using Distributions
	using StatsPlots
	using PlutoUI
	using Pipe
	using Printf
	md"""## loading packages...

	The notebook makes use of some existing functionality from other code which is stored in "packages" that you can install and use.

	"""
end

# ╔═╡ a04710a0-281b-11eb-3c11-5b53a085e59e
html"<button onclick=present()>Present</button>"

# ╔═╡ 8029070e-efb0-11ea-1d0a-7f402f2e9989
md"""# What are `spm`s

Statistical parametric maps... are maps / images that contain statistical values in the same x,y,z as the original image data...

They can be images of the β weights, or other statistics that are derived from them, such as

- the coefficient of determintation, $r^{2}$
- a t-statistic,  $t$, (like from a $t$-test you already know
- statistical significance, $p$-values, etc.

"""

# ╔═╡ 44032c62-2522-11eb-32d2-af427f6ce946
md"""## Doing linear regression

The linear regression analysis can be done very neatly using *linear algebra*. The details of this are beyond this 2nd year lab class, but just to give you an idea of what things look like:

$$\large \text{{\color{teal} \textbf{data}}} = \text{\bf \color{magenta} explanatory variables} \cdot \mathbf{w} $$

$$\large \mathbf{\color{teal} y} = \mathbf{\color{magenta} X}  \cdot w$$

In this equation $\mathbf{w}$, the weights, are what the want to determine. We have measured data -- and by designing the experiment in a particular way (what stimuli when, etc), we have determined the explanatory variables.

To find the weights we need to solve this equation for $\mathbf{w}$:

$$\begin{align}
y &= X \color{green}\mathbf{w}
\end{align}$$

"""

# ╔═╡ cf9fbcc4-2534-11eb-1a82-ad480057d2b0
d, h= readdlm("data/stimulus_timing.csv", ',', Float64, header = true );

# ╔═╡ 3df12514-2535-11eb-2cf3-fdfc727a993b
X = hcat(d[:,3:4], ones(size(d,1),1)) ;

# ╔═╡ b7f96162-25af-11eb-274a-7d368d09906b
md"""
## Using the `lm()` linear model function

*linear models* and regression are some of the most fundamental techniques used in stats, so they are implemented in many software tools, including `SPSS`, and packages for `R`, `Matlab`, `python`, `julia` and many others (as we have seen even in `Excel`)

Across all those tools, the models are often presented in terms of **formulas**, e.g.

  `y ~ face_response + object_response`

which stands for a model that tries to predict the data `y` as a linear combination of two explanatory variables, `face_response` and `object_response`

$$\begin{align}
\color{green}\mathbf{w} \color{black} &= X \backslash y
\end{align}$$

or in code
```julia
w = \(X,ȳ)
# or
w = pinv(X)*ȳ
```
"""

# ╔═╡ a66ed758-25bc-11eb-29ad-01f9ba1bd5be

# ╔═╡ dcb3ad64-281b-11eb-2a9c-5d9cfce9e2e2
md"""## Model:

The measured eesponse is combination of `F` and `O` response"""

# ╔═╡ ba55a268-281b-11eb-3bcd-5b948d03fdbf
md"""## Data, model, residuals..."""

# ╔═╡ dfb188e0-2836-11eb-0081-514c967cace5
md"""## How significant is the response

Turn the beta weight into a $t$ statistic and from there into a $p$ value

$$\begin{align}
t  &= \frac{\text{mean}}{\text{standard error}}
\end{align}$$

Then we need to figure out whether that  $t$-value we calculated is far enough away from 0 to be "significant".

"""

# ╔═╡ 1a0bfe62-2837-11eb-02e3-45eff0ed23f6
t_value = 5.29832/0.229458

# ╔═╡ 39ad1ca2-284a-11eb-3c64-e5e1317eb3d4
md"""$(@bind deg_freedom Slider(2:1:20; default=5, show_value=true)) """

# ╔═╡ e86ecc6e-2849-11eb-34da-0189163c776d
begin
	xx = -5.0:0.2:5.0
	plot(Distributions.TDist(deg_freedom), fill=(0, .5,:teal),
	size=(600,200), title="t Distribution (ν=$(deg_freedom))", label=false)
end

# ╔═╡ 148b4e22-25ba-11eb-28ac-4f90284e7e41
md""" ## Visualising many weights """

# ╔═╡ 6b22738e-281f-11eb-0be4-3195a4953529
md"""## And the `t` values

"""

# ╔═╡ c886d94e-25bc-11eb-0b39-8144659ec58f
md"""## Other versions of the model"""

# ╔═╡ f543c9da-281c-11eb-05bd-ad64b4005507
md"""## Only modelling the `F` response"""

# ╔═╡ 0a4f3af8-281d-11eb-3337-8fbb9998237c
md"""## Only modelling the `O` response"""

# ╔═╡ 41d2d3cc-281d-11eb-176a-c9230c75163c
md"""## Some timing / performance measurements"""

# ╔═╡ 2b2f6be4-f07f-11ea-1459-775749ca949f
md"""## Some helper functions """

# ╔═╡ 46fbc1de-13d1-11eb-2553-71dfd823a5a0
md"""...loading in the contents of a file into a **variable** called `data` works as follows:"""

# ╔═╡ cdf28f06-1449-11eb-34b1-65a1fb83f330
md"""
## Displaying images

To look at the images, we can stich them together into one big *mosaic* and turn them into *grayscale* image that the browswer can display. To do this, we need to find out what the largest (brightest pixel value) and smallest number in the data we want to display are. Because the `max` value might be an outlier, we'll use the 1st and 99th *percentile* instead -- this is often called a **robust range**.

The function `scaleFunc` then squeezes the numbers into a range $[0..1]$ which works well for displaying.
"""

# ╔═╡ 5aed662e-2525-11eb-12a3-c9c9118f6fec
demean = function(y)
	100 .* (y./mean(y) .- 1.0)
end

# ╔═╡ 88253376-1484-11eb-0967-27875c6da199
md"""because this is something we might want to do again and again, you can turn it into a **function**, which is something that takes an  *input*, changes it and spits it out again as an *output*
"""

# ╔═╡ 141dde62-efb1-11ea-34df-cd5e838d8129
begin
	# read data in from file
	data = NIfTI.niread("data/filtered_func_data.nii");
	# small technical point here: images are often stored in y,x,z
	# order (rather thatn x,y,z... so could permute data here (mutable!)

	md"""
	Succesfully loaded a dataset that has **$(ndims(data)) dimensions** of size $(size(data))

	If we look at the table of numbers for one slice (slice #12) for one timepoint (volume = 3 in the below), it looks a bit like the following:
	"""
end

# ╔═╡ f0ef494e-f80c-11ea-1163-b954c1739162
view(data, 1:10:64, 1:10:64, 12, 3);

# ╔═╡ 14f85bc0-2533-11eb-065b-93a0cccb036a
begin
	theSz = [ size(data)[1:3]... ; 3];
	β̂ = zeros(theSz...);
	tmap = zeros(theSz...);

	@time for ix = 1:theSz[1], iy = 1:theSz[2], iz = 1:theSz[3]
		# if mod(ii, 1000) == 1
		# 	println("tada")
		# end
		ols = lm(X, Float64.(demean(data[ix,iy,iz,:])))
		β̂[ix,iy,iz,:] = coef(ols)
		tmap[ix, iy, iz, :] = 	coeftable(ols).cols[3] # t values
	end

end

# ╔═╡ a8ab8ee6-f083-11ea-1316-7fbe4e125914
begin
md"""

### Calculate the beta **weights** at a voxel

We can pick the timeseries data at one particular location and see how long it takes to do that....

[plot residuals] $(@bind plotResiduals CheckBox()) – [plot model] $(@bind plotModel CheckBox())

x coord $(@bind xLoc Slider(1:size(data,1), default = 18, show_value=true))

y coord $(@bind yLoc Slider(1:size(data,2), default = 14, show_value=true))

z coord $(@bind zLoc Slider(1:size(data,3), default = 4, show_value=true))

"""
end

# ╔═╡ 04256b16-2525-11eb-30b3-eb750b5af1a2
y = data[xLoc,yLoc,zLoc,:];

# ╔═╡ 3c5138dc-25b9-11eb-058d-71a10f974857
ȳ = demean(y);

# ╔═╡ 42ace0ae-2538-11eb-3143-31d86905a747
md"""

slice: $(@bind whichSlice Slider( 1:size(data,3) ; default=1, show_value=true))

threshold: $(@bind whichThreshold Slider( 0:0.5:15 ; default=2.0, show_value=true))

"""

# ╔═╡ d0b91c60-2537-11eb-2ff2-eda55e302239
begin
	b1 = β̂[:,:,whichSlice,1]
	b2 = β̂[:,:,whichSlice,2]

	b1[abs.(b1) .<= whichThreshold] .= NaN
	b2[abs.(b2) .<= whichThreshold] .= NaN

	b3 = b1 .- b2

	heatmap(b1, color=:reds, clim = (0,5), layout = (1,2),
			title = "weights map 1",
			subplot=1,
			aspect_ratio=1, showaxis=false)
	heatmap!(b2, color=:blues, clim = (0,5), layout = (1,2),
			title = "weights map 2",
			subplot=2,
			aspect_ratio=1, showaxis=false)

		# heatmap!(b3, clim = (-15,15), layout = (1,3),
		# 	title = "w1 - w2 (f > o)",
		# 	subplot=3,
		# 	aspect_ratio=1, showaxis=false)
end

# ╔═╡ 86dbd432-281f-11eb-1aab-b53d4a226dd0
md"""

slice (t): $(@bind whichSliceT Slider( 1:size(data,3) ; default=1, show_value=true))

threshold (for t): $(@bind whichThresholdT Slider( 0:0.5:15 ; default=2.0, show_value=true))

"""

# ╔═╡ 757b4a0e-281f-11eb-32ad-cfff43e2ebe3
begin
	t1 = tmap[:,:,whichSliceT,1]
	t2 = tmap[:,:,whichSliceT,2]

	t1[abs.(t1) .<= whichThresholdT] .= NaN
	t2[abs.(t2) .<= whichThresholdT] .= NaN

	t3 = t1 .- t2
	t4 = t2 .- t1

	# heatmap(t1, color=:reds, clim = (0,5), layout = (1,1),
	# 		title = "t map 1",
	# 		subplot=1,
	# 		aspect_ratio=1, showaxis=false)
	# heatmap!(t2, color=:blues, clim = (0,5), layout = (1,1),
	# 		title = "t map 2",
	# 		subplot=2,
	# 		aspect_ratio=1, showaxis=false)
	heatmap(t3, color=cgrad([:blue, :white, :red]), clim = (-5,5), layout = (1,1),
			title = "t map (F - O)",
			subplot=1,
			aspect_ratio=1, showaxis=false)

end

# ╔═╡ 55f99666-25ba-11eb-1fd4-a15c6751527e
df = DataFrame(y = Float64.(demean(data[xLoc,yLoc, zLoc,:])),
				face_r = X[:,1], object_r = X[:,2]);

# ╔═╡ 3bf4e998-25a4-11eb-252b-6fdb6a29a92e
ols = lm(@formula(y ~ face_r + object_r), df)

# ╔═╡ 9af3adba-25ba-11eb-34a7-59079afd5231
md"""## How well does the model fit the data?

The coefficient of determination, $r^2$, tells us how much of the variance in the data is accounted for by our model. $r^2$ = $(round( r2(ols); digits=2)) means $( round( 100*r2(ols); digits=0)) % of the variance.

It can be calculated as

$$r^{2}= 1 - \frac{\text{var(residuals)}}{\text{var(data)}}$$

A value of $r^{2} = 1.0$  means that our model is perfect.

A value of $r^{2} \approx 0$ means that our model can't explain our data at all.

"""

# ╔═╡ ab4cc236-25b1-11eb-3f63-233af5cfea96
r² = r2(ols)

# ╔═╡ 10a0bd36-2837-11eb-2fe7-69a17334a3d9
ols

# ╔═╡ 821ffcee-2849-11eb-340d-2dc707fd0a8f
md""" ## But what about a $p$-value

The p-value corresponds to the area under the curve beyond the calculated t-value... i our example for `face_r` it was t = $(round(coeftable(ols).cols[3][2]; digits=2)) and a p-value of $(@sprintf "%.2g" coeftable(ols).cols[4][2])

"""

# ╔═╡ 765e7432-25a4-11eb-10b5-d3bdab410d87
olsFvO = lm(@formula(y ~ face_r - object_r), df)

# ╔═╡ 570626cc-25a5-11eb-123c-63c5b9aad411
olsF = lm(@formula(y ~ face_r), df)

# ╔═╡ 76cdcb54-25bb-11eb-0e29-9dfeae3f803d
r2(olsF)

# ╔═╡ 9fa22482-25a5-11eb-2b76-1949341e2630
olsO = lm(@formula(y ~ object_r), df)

# ╔═╡ 894cdc18-25bb-11eb-130d-93b864ddcc42
r2(olsO)

# ╔═╡ 1ef8659c-2539-11eb-184a-d95f169892a4
result = @timed @pipe data[xLoc,yLoc,zLoc,:] |>
	 					demean( _ ) |>
	 					vec( _ ) |>
	 					X \ _

# ╔═╡ 3483ee2c-2525-11eb-2967-d343eb3d16f0
begin

	p_ = plot(ȳ, lw=2, c=:black, label="data $([xLoc, yLoc, zLoc])", size=(600,200),
			xlabel="Time (volumes)", ylabel="fMRI response",
			title="r² = $(round(r²; digits=2))")

	modelFit = modelFit = X*result.value
	residuals = ȳ .- modelFit

	if plotModel
		plot!(modelFit, c=:red, lw=2.0, label="model fit")
	end

	if plotResiduals
		plot!(residuals, line=:stem, c=:blue, label="redisuals")
	end

	p_
end

# ╔═╡ 15ad8bf8-2835-11eb-1315-f5d8ea0d2bf9
md"""The *weights* for reconstructing are: **$(@sprintf "[%.2f. %.2f, %.2f]" result.value...)**"""

# ╔═╡ a48038d4-2539-11eb-3374-dbbeaed2a3a5
result.value

# ╔═╡ b3a74a5a-2539-11eb-017a-f9b3ffe38663
r_microseconds = result.time * 1000 * 1000; # in microseconds

# ╔═╡ 3e13431a-253a-11eb-265e-f11984b2b2fb
md"""The following expression feeds the data from [$(xLoc); $(yLoc); $(zLoc)] to a function called `demean()` and then to a function called `vec()` and finally does the maths using the `\` (backslash) operator:

Time taken: ≈ $(round(r_microseconds; digits=2)) microseconds

weights: $(@sprintf("[%.2f, %.2f, %.2f]", round.(result.value; digits=2)...))

"""

# ╔═╡ 13a8fe64-efb4-11ea-1ac6-dbd49975dddf
minMax = percentile(data[:], [1; 99])

# ╔═╡ 71606d26-efb4-11ea-12e2-75f85b9ea7d8
# helper function that knows how to scale images to a range that we can display
scaleFunc = scaleminmax(minMax[1], minMax[2])

# ╔═╡ c1d831e0-1448-11eb-0b18-a555cbcb18e1
@pipe data[:,:,:,3] |>
	mosaicview( _ , nrow=4; rowmajor=true) |>
	scaleFunc.( _ ) |>
	Gray.( _ )

# ╔═╡ b69a53ee-1484-11eb-0953-2b3ff81cf56c
function simpleMosaic(d, t=3, nr = 4)
	Gray.(scaleFunc.(mosaicview(d[:,:,:,t], nrow=nr; rowmajor=true)))
end

# ╔═╡ e2bea646-1484-11eb-3a26-ffc655b139c6
simpleMosaic(data) # now we can use the code this way

# ╔═╡ 072ce484-efb2-11ea-36ee-29a0147cfc04
"""
displaySlice(data, slicenum, orientation)

return and display a slice in a particular orientation. there are smarter
ways to implement the slicing through different orientations, but the code
using if/elseif is easy to folllow.

	displaySlice(data, 2, 1)

"""
displaySlice = function(e, s = 1, orientation = 3, scaling = scaleFunc)
	if  (orientation == 1)
		displayData = view(e, s, :,:,1)
	elseif (orientation == 2)
		displayData = view(e, :,s,:,1)
	elseif( orientation == 3)
		displayData = view(e, :, :, s, 1)
	else
		error("uhoh!")
	end
	rotl90(Gray.(scaling.(displayData)))
end

# ╔═╡ 8cbe6b98-144b-11eb-3d0a-35700773e22b
"""
showMosaic - arrange 4d data (at chosen timepoint) into a mosaic

    showMosaic(data,2)  # permute the dimensions (circshift twice)
						# and then arrange into roughly nxn mosaic

"""
function showMosaic(d, dim, t=1)
	@pipe d[:,:,:,t] |>
		permutedims(_, circshift([1,2,3],dim)) |>
		mosaicview(_,
			nrow=Int64(round(sqrt(size(_,3)...))),
			rowmajor=true,
			npad = 2,
			fillvalue = maximum(_)) |>
		scaleFunc.(_) |>
		Gray.(_)
end

# ╔═╡ Cell order:
# ╠═a04710a0-281b-11eb-3c11-5b53a085e59e
# ╟─8029070e-efb0-11ea-1d0a-7f402f2e9989
# ╟─f0ef494e-f80c-11ea-1163-b954c1739162
# ╟─44032c62-2522-11eb-32d2-af427f6ce946
# ╟─04256b16-2525-11eb-30b3-eb750b5af1a2
# ╟─3c5138dc-25b9-11eb-058d-71a10f974857
# ╟─cf9fbcc4-2534-11eb-1a82-ad480057d2b0
# ╟─3df12514-2535-11eb-2cf3-fdfc727a993b
# ╟─14f85bc0-2533-11eb-065b-93a0cccb036a
# ╟─b7f96162-25af-11eb-274a-7d368d09906b
# ╟─a66ed758-25bc-11eb-29ad-01f9ba1bd5be
# ╟─dcb3ad64-281b-11eb-2a9c-5d9cfce9e2e2
# ╠═3bf4e998-25a4-11eb-252b-6fdb6a29a92e
# ╟─ba55a268-281b-11eb-3bcd-5b948d03fdbf
# ╟─3483ee2c-2525-11eb-2967-d343eb3d16f0
# ╟─15ad8bf8-2835-11eb-1315-f5d8ea0d2bf9
# ╟─a8ab8ee6-f083-11ea-1316-7fbe4e125914
# ╟─9af3adba-25ba-11eb-34a7-59079afd5231
# ╠═ab4cc236-25b1-11eb-3f63-233af5cfea96
# ╟─dfb188e0-2836-11eb-0081-514c967cace5
# ╠═1a0bfe62-2837-11eb-02e3-45eff0ed23f6
# ╠═10a0bd36-2837-11eb-2fe7-69a17334a3d9
# ╟─821ffcee-2849-11eb-340d-2dc707fd0a8f
# ╠═e86ecc6e-2849-11eb-34da-0189163c776d
# ╟─39ad1ca2-284a-11eb-3c64-e5e1317eb3d4
# ╟─148b4e22-25ba-11eb-28ac-4f90284e7e41
# ╟─d0b91c60-2537-11eb-2ff2-eda55e302239
# ╟─42ace0ae-2538-11eb-3143-31d86905a747
# ╟─6b22738e-281f-11eb-0be4-3195a4953529
# ╟─757b4a0e-281f-11eb-32ad-cfff43e2ebe3
# ╟─86dbd432-281f-11eb-1aab-b53d4a226dd0
# ╟─c886d94e-25bc-11eb-0b39-8144659ec58f
# ╠═765e7432-25a4-11eb-10b5-d3bdab410d87
# ╟─f543c9da-281c-11eb-05bd-ad64b4005507
# ╠═570626cc-25a5-11eb-123c-63c5b9aad411
# ╠═76cdcb54-25bb-11eb-0e29-9dfeae3f803d
# ╟─0a4f3af8-281d-11eb-3337-8fbb9998237c
# ╠═9fa22482-25a5-11eb-2b76-1949341e2630
# ╠═894cdc18-25bb-11eb-130d-93b864ddcc42
# ╟─55f99666-25ba-11eb-1fd4-a15c6751527e
# ╟─41d2d3cc-281d-11eb-176a-c9230c75163c
# ╟─3e13431a-253a-11eb-265e-f11984b2b2fb
# ╠═1ef8659c-2539-11eb-184a-d95f169892a4
# ╠═a48038d4-2539-11eb-3374-dbbeaed2a3a5
# ╠═b3a74a5a-2539-11eb-017a-f9b3ffe38663
# ╟─2b2f6be4-f07f-11ea-1459-775749ca949f
# ╟─46fbc1de-13d1-11eb-2553-71dfd823a5a0
# ╟─cdf28f06-1449-11eb-34b1-65a1fb83f330
# ╟─c1d831e0-1448-11eb-0b18-a555cbcb18e1
# ╟─e2bea646-1484-11eb-3a26-ffc655b139c6
# ╠═9bc349e8-efb0-11ea-1392-cd735828af33
# ╠═13a8fe64-efb4-11ea-1ac6-dbd49975dddf
# ╟─71606d26-efb4-11ea-12e2-75f85b9ea7d8
# ╟─b69a53ee-1484-11eb-0953-2b3ff81cf56c
# ╟─5aed662e-2525-11eb-12a3-c9c9118f6fec
# ╟─88253376-1484-11eb-0967-27875c6da199
# ╟─141dde62-efb1-11ea-34df-cd5e838d8129
# ╟─072ce484-efb2-11ea-36ee-29a0147cfc04
# ╟─8cbe6b98-144b-11eb-3d0a-35700773e22b


export contains_expression

# SUBEXPRESSION SEARCH

# if haystack is a logical symbol, the kneedle must be equal to be a subexpression
"""
    contains_expression(haystack::AbstractExpression, kneedle::AbstractExpression)

Searches for a subexpression `kneedle` recursively within the `haystack` expression. Returns true if a match can be found
and false otherwise.
"""
contains_expression(haystack::LogicalSymbol, kneedle::AbstractExpression) = isequal(haystack, kneedle)

# TODO: come up with quicker subexpression algorithm for LogicalExpressions
# if kneedle is an expression, use a recursive equality check
function contains_expression(haystack::LogicalExpression, kneedle::LogicalExpression)
    if isequal(haystack, kneedle)
        return true
    end
    return any(contains_expression.(arguments(haystack), Iterators.repeat([kneedle], length(arguments(haystack)))))
end

# if kneedle is a symbol, use the variables set to quickly check if its a subexpression
contains_expression(haystack::LogicalExpression, kneedle::LogicalSymbol) = kneedle ∈ variables(haystack)

# dfs for expression search; returns path to first found expression, otherwise empty vector
locate_expression(haystack::LogicalSymbol, kneedle::AbstractExpression) = []
function locate_expression(haystack::LogicalExpression, kneedle::AbstractExpression)
    i = 1
    for arg ∈ arguments(haystack)
        if isequal(arg, kneedle)
            return [i]
        end
        located = locate_expression(arg, kneedle)
        if length(located) > 0
            return [i, located...]
        end
        i += 1
    end
    return []
end

function select_by_path(expr::AbstractExpression, path::Vector{Int})
    if length(path) == 0
        return expr
    end
    return select_by_path(arguments(expr)[path[1]], path[2:end])
end


# We want to explore fast ways to subset a large, square N*N matrix in julia dependent on
#   1. N = number of dimensions of large matrix / data set
#   2. S = subset size = floor(log(N)) - floor(10*log(N))
#   3. Method used to extract subset of dimensions/features

function extract_with_getindex(matrix, subset)
  matrix[subset,:][:,subset]
end

function extract_with_getindex2(matrix, subset)
  matrix[subset,subset]
end

function extract_in_loop(matrix, subset)
  m = length(subset)
  res = zeros(m, m)
  for i in 1:m
    si = subset[i]
    for j in 1:m
      res[i, j] = matrix[si, subset[j]]
    end
  end
  res
end

# Do some runs to ensure everything is compiled
m = randn(10,10)
ss = sort(shuffle(collect(1:10))[1:4])
r1 = extract_in_loop(m, ss)
r2 = extract_with_getindex(m, ss)
r3 = extract_with_getindex2(m, ss)

num_features(n) = rand(int(floor(log(n))):int(floor(10*log(n))))

NumRepeats = 10
for n in [100, 1000, 10000]
  times1 = zeros(NumRepeats)
  times2 = zeros(NumRepeats)
  times3 = zeros(NumRepeats)
  for i in 1:NumRepeats
    matrix = randn(n, n)
    ss = sort(shuffle(collect(1:n))[1:num_features(n)])
    tic()
      res1 = extract_in_loop(matrix, ss)
    times1[i] = toq()
    tic()
      res2 = extract_with_getindex(matrix, ss)
    times2[i] = toq()
    tic()
      res3 = extract_with_getindex2(matrix, ss)
    times3[i] = toq()
    if res1 != res2 || res1 != res3 || res2 != res3
      println("Results differ!")
    end
  end
  println("n = $(n), extract_in_loop = $(mean(times1)), extract_with_getindex = $(mean(times2)), extract_with_getindex2 = $(mean(times3)), getindex2 factor faster = $(mean(times1)/mean(times3))")
end


import Base.show
show(io::IO,eq::Equation)=print(io,ps(eq))
show(io::IO,ex::Factor)=print(io,ps(ex))

ps(eq::Equation)="$(ps(eq.lhs)) ≖ $(ps(eq.rhs))"
function ps(ex::Expression)
	s=""
	for t in ex
		s*=ps(t)*"+"
	end
	s[1:end-1]
end
function ps(t::Term)
	s=""
	for f in t
		s*=ps(f)*"*"
	end
	s[1:end-1]
end
function ps{T}(ar::Array{T,1})
	if isempty(ar)
		return "$T[]"
	end
	s="$T["
	for a in ar
		s*=ps(a)*","
	end
	s[1:end-1]*"]"
end
#function ps{T}(a::T)
#	"$T"
#end
function ps(a)
	io=IOBuffer()
	show(io,a)
	str=takebuf_string(io)
end
function ps(c::Component)
	ar=getargs(c)
	s="$(typeof(c))("
	for a in ar
		s*="$(ps(a)),"
	end
	s[1:end-1]*")"
end
ps(s::Symbol)=":$s"
ps(n::Number)="$n"

function print(io::IO,u::U)
	print(io,u.x)
	print(io," [")
	print(io,u.units)
	print(io,"]")
end
ps(a::Alt)="Alt($(a.indices))"
function show(io::IO,a::Alt)
	print(io,ps(a))
end
function print(io::IO,a::Alt)
	print(io,"ϵ$(a.indices)")
end


# How to print SimplePolynomials? We take advantage of the
# excellent printing available in the Polynomials module!

import Base: string, show
export integerize

# The all_ones_denom function checks if all the coefficients
# are (Gaussian) integers, but perhaps of Rational type.

function all_ones_denom(list::Vector{T})::Bool where {T<:IntegerX}
    return true
end
#
# function all_ones_denom(list::Vector{Complex{T}})::Bool where T <: Integer
#     return true
# end

function all_ones_denom(list::Vector{T})::Bool where {T<:Real}
    return all(denominator.(list) .== 1)
end

function all_ones_denom(list::Vector{T})::Bool where {T<:Complex}
    return all_ones_denom(real.(list)) && all_ones_denom(imag.(list))
end

function all_ones_denom(p::SimplePolynomial)::Bool
    return all_ones_denom(p.data)
end

function _top(x::T) where {T<:Union{Rational,Integer}}
    return numerator(x)
end

function _top(x::T) where {T<:Complex}
    return numerator(real(x)) + numerator(imag(x)) * im
end

_top(x::AbstractMod) = x

"""
`integerize(p::SimplePolynomial)` is useful when `p`
has `Rational` coefficients all of whose denominators are `1`.
In that case, `integerize(p)` returns a polynomial equal to the
original, but with coefficients that are integers (or Gaussian integers).
However, if any denominators are not `1`, the original polynomial
is returned.
"""
function integerize(p::SimplePolynomial)
    if all_ones_denom(p)
        return SimplePolynomial(_top.(p.data))
    end
    return p
end

"""
`string(p::SimplePolynomial)` returns a nice `String` representation
of `p`.
"""
function string(p::SimplePolynomial)
    return string(Polynomial(integerize(p)))
end

function show(io::IO, p::SimplePolynomial)
    print(io, string(p))
end

export string3

function pstring(p::SimplePolynomial)
    nterms = count((x) -> (x != 0), p.data)
    if nterms <= 1
        return string(p)
    end
    return "(" * string(p) * ")"
end

function string(f::SimpleRationalFunction)
    return pstring(f.num) * " / " * pstring(f.den)
end

function show(io::IO, f::SimpleRationalFunction)
    print(io, string(f))
end

function string3(f::SimpleRationalFunction)
    top = string(numerator(f))
    bot = string(denominator(f))

    n1 = length(top)
    n2 = length(bot)
    n = max(n1, n2)

    pad1 = 0
    pad2 = 0

    if n1 < n
        pad1 = div(n - n1, 2)
    else
        pad2 = div(n - n2, 2)
    end

    line1 = " "^pad1 * top
    line2 = "-"^n
    line3 = " "^pad2 * bot
    NL = "\n"

    return line1 * NL * line2 * NL * line3
end


using SwitchTimeOpt
using MathOptInterface
const MOI = MathOptInterface
using Ipopt

# Define Solver options
maxiter = 100
maxtime = 1000.0
verbose = 5
tolerance = 1e-06

#Define solver
solver = Ipopt.Optimizer()
MOI.set(solver, MOI.RawOptimizerAttribute("tol"), tolerance)
MOI.set(solver, MOI.RawOptimizerAttribute("print_level"), verbose)
MOI.set(solver, MOI.RawOptimizerAttribute("max_cpu_time"), maxtime)
MOI.set(solver, MOI.RawOptimizerAttribute("max_iter"), maxiter)

### Define system parameters
# Time vector
t0 = 0.0
tf = 12.0

# Integer input
uvec = [repeat([0.0; 1.0], 4, 1); 0.0]
uvec = copy(uvec')

# Number of switching times
N = size(uvec,2) - 1

# Cost funcction matrix
C = [1.0 0.0 -1.0;
     0.0 1.0 -1.0]
Q = C'*C

# Define initial state
x0 = [0.5; 0.7; 1]

### Define system dynamics
function nldyn(x::Array{Float64,1}, u::Array{Float64,1})
  n = length(x)
  f = zeros(n)
  f[1] = x[1] - x[1]*x[2] - 0.4*x[1]*u[1]
  f[2] = -x[2] + x[1]*x[2] - 0.2*x[2]*u[1]
  return f
end

function nldyn_deriv(x::Array{Float64,1}, u::Array{Float64,1})
  df = [1.0-x[2]-0.4*u[1]       -x[1]                   0;
        x[2]                     -1+x[1]-0.2*u[1]       0;
        0                        0                      0]
  return df
end

# Generate and solve problem
ngrid = 300

# Initialize the problem
m = stoproblem(
  x0,                 # Initial state
  nldyn,              # Nonlinear dynamics
  nldyn_deriv,        # Nonlinear dynamics derivative
  uvec,               # Vector of integer inputs
  ngrid = ngrid,      # Number of grid points
  t0 = t0,            # InitialtTime
  tf = tf,            # Final time
  Q = Q,              # Cost matrix
  solver = solver)

solve!(m)
objlin = getobjval(m)
tauopt = gettau(m)
status = getstat(m)
xsim, xpts, objode45, t = simulate(m, tauopt)

println("Status: ", status)
println("Objective value: ", objlin)

showPlots = true
savePlots = false
plotSolution(m, "Lotka-Volterra", "Lotka-Volterra", showPlots, savePlots)


using VertexEliminationOrder
using Test
using Graphs

ENV["JULIA_DEBUG"]=VertexEliminationOrder

G = SimpleGraph{Int64}(12)
# 1 to 5 is a clique
for i in 1:5
    for j in (i+1):5
        add_edge!(G, i, j)
    end
end
# 6 and 7 are midle nodes
add_edge!(G, 5, 6)
add_edge!(G, 6, 7)
add_edge!(G, 6, 8)
# 8 to 12 is a tree
add_edge!(G, 8, 9)
add_edge!(G, 8, 10)
add_edge!(G, 9, 11)
add_edge!(G, 9, 12)

@info "Testing iterative_dissection with custom graph"
res = iterative_dissection(G)
@info res
res_expected = treewidth_by_elimination!(G, res[1])
# 6 should be separator but t seems that it depends on the random s and
@info res_expected
@test abs(res[2] - res_expected) <= 1
@test res[2] >=  4# 4 is the tw but it returns an upper bound

for n in 2:20
    g = square_lattice_graph(n)
    @info "Testing iterative_dissection with square lattice graph n = $n"
    tmp = iterative_dissection(g)
    @info tmp
    tmp_expected = treewidth_by_elimination!(g, tmp[1])
    @info tmp_expected
    @test abs(tmp[2] - tmp_expected) <= 1
end


using Test
include("RandomHypergraphs.jl")
include("load_data.jl")

@testset "Stub Matching" begin
    # Degree and edge-dimension sequences are correct
    DK(H::Hypergraph) = H.D, H.K
    sequence_match(D::Vector{Int64}, K::Vector{Int64}) = return DK(stub_matching(D,K)) == (D, sort(K, rev=true))

    # Realizable hypergraphs
    @test sequence_match([1,1,1,1], [2,2])
    @test sequence_match([1,2,3,4,5,6],[1,2,3,4,5,6])
    @test sequence_match([6,3,3,3], [4,3,3,1,1,1,1,1])
    @test sequence_match([1,2,3,1,1,1],[3,2,3,1])

    # Unrealizable hypergraphs
    @test_throws ArgumentError sequence_match([0],[5])
    @test_throws ArgumentError sequence_match([5,5,5,5,5],[5])
    @test_throws ArgumentError sequence_match([1,1,1,1],[4,4,4])

end
@testset "MCMC" begin
    @testset "Vertex" begin
        @test_throws AssertionError MCMC(Hypergraph(Array{Array{Int64,1},1}(),4,0))
        @test_throws AssertionError MCMC(Hypergraph([[1,2]],2,1))
        @test typeof(MCMC(Hypergraph([[1,2],[3,4],[1,4]], 4, 3))) <: Hypergraph
        @test typeof(MCMC(stub_matching([1,2,3,1,1,1],[3,2,3,1]); samples=1000)) <: Hypergraph
        @test typeof(MCMC(load_hypergraph("Enron"); samples=10000)) <: Hypergraph
    end

    @testset "Stub" begin
        @test_throws AssertionError typeof(MCMC(Hypergraph(Array{Array{Int64,1},1}(),4,0); type="stub")) <: Hypergraph
        @test typeof(MCMC(Hypergraph([[1,2],[3,4],[1,4]], 4, 3); type="stub")) <: Hypergraph
        @test typeof(MCMC(stub_matching([1,2,3,1,1,1],[3,2,3,1]); type="stub")) <: Hypergraph
        @test typeof(MCMC(load_hypergraph("Enron"); samples=10000, type="stub")) <: Hypergraph
    end
end


function set_res(Delta_Qcon_hat, Delta_Qcon_hat_temp, cellxmax, cellymax, cellzmax)
    res = zeros(cellxmax, cellymax, cellzmax, 5)

    for i in 2:cellxmax-1
        for j in 2:cellymax-1
            for k in 2:cellzmax-1
                for l in 1:5
                    res[i,j,k,l] = abs(Delta_Qcon_hat[i,j,k,l]-Delta_Qcon_hat_temp[i,j,k,l])
                end
            end
        end
    end

    # println(res[3,3,:])
    # println(dt*RHS[3,3,:])
    return res
end

function check_converge(res, RHS, cellxmax, cellymax,cellzmax, init_small)
    norm2 = zeros(5)

    tempAxb = zeros(5)
    tempb = zeros(5)
    for i in 2:cellxmax-1
        for j in 2:cellymax-1
            for k in 2:cellzmax-1
                for l in 1:5
                    tempAxb[l] = tempAxb[l] + res[i,j,k,l]^2
                    tempb[l] = tempb[l] + RHS[i,j,k,l]^2
                end
            end
        end
    end

    #println(tempb)
    #println(tempAxb)
    #throw(UndefVarError(:x))

    for l in 1:5
        norm2[l] = (tempAxb[l]/(tempb[l]+init_small)) ^ 0.5
    end
    return norm2
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
using Tullio, LoopVectorization

I, J, K = 400, 200, 100
f1 = rand(I, 1, K)
f2 = rand(I, J, 1)
dev = rand(I, 1, 1)

function make_R(f1, f2, dev)
    I, J, K = axes(dev, 1), axes(f2, 2), axes(f1, 3)
    R = similar(dev, I, J, K)
    for i in I, j in J, k in K
        @inbounds R[i, j, k] = f1[i, 1, k] * f2[i, j, 1] / dev[i, 1, 1]
    end
    R
end

R1 = make_R(f1, f2, dev)
@tullio R2[i, j, k] := f1[i, 1, k] * f2[i, j, 1] / dev[i, 1, 1]
R3 = @. f1 * f2 / dev
R4 = @turbo @. f1 * f2 / dev
R5 = @tturbo @. f1 * f2 / dev

@show R1 ≈ R2 ≈ R3 ≈ R4 ≈ R5
@show typeof(R1) size(R1);

# %%
using BenchmarkTools
@btime R1 = make_R($f1, $f2, $dev)
@btime @tullio R2[i, j, k] := $f1[i, 1, k] * $f2[i, j, 1] / $dev[i, 1, 1]
@btime R3 = @. $f1 * $f2 / $dev
@btime R4 = @turbo @. $f1 * $f2 / $dev
@btime R5 = @tturbo @. $f1 * $f2 / $dev
;

# %%

0009/Tullio.jl and LoopVectorization.jl example.jl


# ==============================================================================
# Xavier Gandibleux - November 2021
#   Implemented in Julia 1.6

# ==============================================================================

println("""\nProject MOMH 2021" --------------------------------\n""")

const verbose = true
const graphic = true

include("parserMomkpZL.jl")
include("parserMomkpPG.jl")
include("vOptMomkp.jl")
include("displayGraphic.jl")

# ==============================================================================
# Datastructure of a generic bi-objective 0/1 IP where all coefficients are integer
#
#   Max  sum{j=1,...,n} C(k,j) x(j)                k=1,...,p
#   st   sum{j=1,...,n} A(i,j) x(j) <= b_{i}       i=1,...,m
#                       x(j) = 0 or 1

struct _bi01IP
    C  :: Matrix{Int} # objective functions, k=1..2, j=1..n
    A  :: Matrix{Int} # matrix of constraints, i=1..m, j=1..n
    b  :: Vector{Int} # right-hand side, i=1..m
end

# ==============================================================================

# Example ----------------------------------------------------------------------
fname = "instancesZL/knapsack.100.2"

# Read and load an instance of MO-01MKP from the collection of E. Zitzler / M. Laumanns
momkpZL = readInstanceMOMKPformatZL(verbose,fname)

# Reduce the MO-MKP instance to the two first objectives and store it as a generic Bi-01IP
dat1 = _bi01IP(momkpZL.P[1:2,:], momkpZL.W, momkpZL.ω)

# Example ----------------------------------------------------------------------
fname = "instancesPG/set1/ZL28.DAT"
#fname = "instancesPG/set2/kp28W-Perm.DAT"
#fname = "instancesPG/set3/W7BI-rnd1-1800.DAT"

# Read and load an instance of Bi-01BKP from the collection of O. Perederieieva / X. Gandibleux
momkpPG = readInstanceMOMKPformatPG(verbose,fname)

# Store it as a generic bi-01IP
dat2 = _bi01IP(momkpPG.P, momkpPG.W, momkpPG.ω)

# Example ----------------------------------------------------------------------
solverSelected = GLPK.Optimizer
#YN, XE = vSolveBi01IP(solverSelected, dat1.C, dat1.A, dat1.b)
YN, XE = vSolveBi01IP(solverSelected, dat2.C, dat2.A, dat2.b)

# ---- Displaying the results (XE and YN)
for n = 1:length(YN)
    X = value.(XE, n)
    print(findall(elt -> elt ≈ 1, X))
    println("| z = ",YN[n])
end

for n = 1:length(YN)
    println(YN[n])
end

# Example ----------------------------------------------------------------------
displayGraphics(fname,YN)


function binarysearch(a::Vector, item, isless=isless)
    leftindex = 1
    rightindex = length(a)
    while leftindex <= rightindex
        middleindex = (leftindex + rightindex) ÷ 2
        if isless(a[middleindex], item)
            leftindex = middleindex + 1
        elseif isless(item, a[middleindex])
            rightindex = middleindex - 1
        else
            return middleindex
        end
    end
    return false
end


# Getting length
print("Emter length: ")
length = readline()
length = parse(Float64, length) # Converting string to Float

# Getting width
print("Emter width: ")
width = readline()
width = parse(Float64, width) # Converting string to Float

# Calculating Area
area = length*width

println("Area of Rectangle is ", area)


get_sources(g::AbstractGraph{T}) where T <: Integer = vertices(g)[indegree(g) .== 0]
get_sinks(g::AbstractGraph{T})  where T <: Integer = vertices(g)[outdegree(g) .== 0]

function add_source_target!(g::AbstractGraph{T}) where T <: Integer
    sinks = get_sinks(g)
    sources = get_sources(g)

    add_vertex!(g)
    for source in sources
        add_edge!(g, vertices(g)[end], source)
    end

    add_vertex!(g)
    for sink in sinks
        add_edge!(g, sink, vertices(g)[end])
    end
end

function N⁻(g::Graphs.SimpleDiGraph{T}, vseqt::Vector{T}) where T <: Integer
    nV = length(vseqt)
    val = zeros(nV)
    val[vseqt[1]] = 1.0
    for v in vseqt
        nb = outneighbors(g, v)
        val[nb] .+= val[v]
    end
    return val
end

function N⁺(g::Graphs.SimpleDiGraph{T}, vseqt::Vector{T}) where T <: Integer
    nV = length(vseqt)
    val = zeros(nV)
    val[vseqt[nV]] = 1.0
    for v in reverse(vseqt)
        nb = inneighbors(g, v)
        val[nb] .+= val[v]
    end
    return val
end

function weights_matrix(g::AbstractGraph{T}, weights::AbstractVector{U}) where T <: Integer where U <: Real
    senders = [src(e) for e in edges(g)]
    receivers = [dst(e) for e in edges(g)]
    return sparse(senders, receivers, weights, nv(g), nv(g))
end

function _weights_spc_raw(g::AbstractGraph{T}) where T <: Integer
    g = Graphs.SimpleDiGraph(g)
    add_source_target!(g)
    vseqt = Graphs.topological_sort_by_dfs(g)
    st = vseqt[[1, end]]
    N_m = N⁻(g, vseqt)
    N_p = N⁺(g, vseqt)
    idx = [!(v in st) for v in vertices(g)]
    rem_vertices!(g, st)
    N_m[idx], N_p[idx], N_p[st[1]]
end

function _weights_spc_vertices(g, normalize)
    N_m, N_p, tf = _weights_spc_raw(g)
    vw = N_m .* N_p

    if normalize == :none
        return vw
    elseif normalize == :log
        return log.(vw)
    elseif normalize == :totalflow
        return vw /= tf
    else
        error("normalize needs to be :none, :log, or :totalflow.")
    end
end

function _weights_spc_edges(g, normalize)
    N_m, N_p, tf = _weights_spc_raw(g)

    ew = map(edges(g)) do e
        N_m[src(e)] * N_p[dst(e)]
    end

    if normalize == :none
        ew = ew
    elseif normalize == :log
        ew = log.(ew)
    elseif normalize == :totalflow
        ew = ew /= tf
    else
        error("normalize needs to be :none, :log, or :totalflow.")
    end

    weights_matrix(g, ew)
end

abstract type MainPathWeight end

abstract type MainPathEdgeWeight <: MainPathWeight end
abstract type MainPathVertexWeight <: MainPathWeight end

"""
    SPCEdge(normalize=:log)(g)

Struct representing Search Path Cout (SPC) edge weights, as defined in Batagelj (2003).
`normalize=:totalflow` indicates that weights should be normalized relative to the
total SPC flow, `normalize=:log` indicates that the logarithm of weights is returned.
Set `normalize` to :none to indicate that weights should not be transformed.

The struct can be called to compute weights directly or
can be passed to the `mainpath` function for dispatch.
"""
Base.@kwdef struct SPCEdge <: MainPathEdgeWeight
    normalize::Symbol = :log
end

(x::SPCEdge)(g) = _weights_spc_edges(g, x.normalize)

"""
    SPCVertex(normalize=:log)(g)

Struct representing Search Path Cout (SPC) vertex weights, as defined in Batagelj (2003).
`normalize=:totalflow` indicates that weights should be normalized relative to the
total SPC flow, `normalize=:log` indicates that the logarithm of weights is returned.
Set `normalize` to :none to indicate that weights should not be transformed.

The struct can be called to compute weights directly or
can be passed to the `mainpath` function for dispatch.
"""
Base.@kwdef struct SPCVertex <: MainPathVertexWeight
    normalize::Symbol = :log
end

(x::SPCVertex)(g) = _weights_spc_vertices(g, x.normalize)


#mutable struct DisjointSetNode
#    rank::Int
#    p::DisjointSetNode
#    DisjointSetNode() = (obj = new(0); obj.p = obj;)
#end

function findset(x::DisjointSetNode)
    if (x != x.p)
        x.p = findset(x.p)
    end
    return x.p
end

function my_union!(x::DisjointSetNode, y::DisjointSetNode)
    link(findset(x), findset(y))
end

function link(x::DisjointSetNode, y::DisjointSetNode)
    if (x.rank > y.rank)
        y.p = x
    else
        x.p = y
        if (y.rank == x.rank)
            y.rank = y.rank + 1
        end
    end
end

function hammingdistance(s1::String, s2::String)
    dist = 0 # accumulated "distance" between strings
    for i = 1:length(s1)
        if (Int(s1[i]) - Int(s2[i]) != 0) # if letters at element i dont match
            dist = dist + 1 # increase distance strings
        end
    end
    return dist
end

function findclusters(E::Vector{Tuple{Int, Int, Int}}, n::Int, k::Int)
    A = []
    sort!(E)
    nodes = []
    for i = 1:n
        push!(nodes, DisjointSetNode())
    end
    counter = 1
    while (counter <= n-k)
        edge = E[counter] # edge is tuple (w,u,v)
        if findset(nodes[edge[2]]) != findset(nodes[edge[3]]) # compare rootnode of node u and v
            push!(A,edge)
            my_union!(nodes[edge[2]], nodes[edge[3]])
        end
        counter = counter + 1
    end
    clusters = []
    c = 1 # cluster iterator
    push!(clusters, [])
    push!(clusters[1],[E[1][2]])
    clusterParent = findset(nodes[E[1][2]])
    println(nodes)
    for i = 2:n
        if findset(nodes[E[i][2]])==clusterParent
            push!(clusters[1][c], E[i][2])
        else
            c = c + 1
        end
    end
    if length(clusters[1]) < 2
        return clusters[1]
    else
    	return clusters
    end
end


using Qaintessent

"""
Identity map as sparse matrix with real-valued entries.
"""
sparse_identity(n) = sparse(1.0*I, n, n)

"""
Construct a parametrized quantum layer consisting of single qubit rotation gates.
"""
function single_qubit_layer(L::Integer)
    # local gates
    # [circuit_gate(i, RotationGate(0.05*π*randn(3))) for i in 1:L]
    cgc = CircuitGate[]
    append!(cgc, [circuit_gate(i, RxGate(0.05*π*randn())) for i in 1:L])
    append!(cgc, [circuit_gate(i, RyGate(0.05*π*randn())) for i in 1:L])
    append!(cgc, [circuit_gate(i, RzGate(0.05*π*randn())) for i in 1:L])
    cgc
end

"""
Construct a parametrized quantum layer consisting of two qubit entanglement gates for 1D chain.
"""
function trotterized_layer_1d(L::Integer, istart::Integer; boundary::Bool=false)
    # interaction gates
    vcat(
        [circuit_gate(i, i+1, EntanglementXXGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementXXGate(0.05π*randn()))], []),
        [circuit_gate(i, i+1, EntanglementYYGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementYYGate(0.05π*randn()))], []),
        [circuit_gate(i, i+1, EntanglementZZGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementZZGate(0.05π*randn()))], [])
    )
end

function create_cgc(L::Int; isboundary::Bool=false)
    cgc = CircuitGate[]
    layers = 3
    for _ in 1:layers
        append!(cgc, single_qubit_layer(L::Integer))
        append!(cgc, trotterized_layer_1d(L, 1; boundary=isboundary))
        append!(cgc, trotterized_layer_1d(L, 2; boundary=isboundary))

    end

    cgc
end


## Jacobi Polynomials
@registerN Jacobi AbstractCCOP2 α β
export Jacobi

"""
    Jacobi{α,  β, T}

Implements the [Jacobi](https://en.wikipedia.org/wiki/Jacobi_polynomials) polynomials. These have weight function `w(x) = (1-x)^α ⋅ (1+x)^β` over the domain `[-1,1]`. Many orthogonal polynomial types are special cases. The parameters are  specified to the constructors:

```jldoctest
julia> using Polynomials, SpecialPolynomials

julia> p = Jacobi{-1/2, -1/2}([0,0,1])
typename(Jacobi){-0.5,-0.5}(1⋅Jᵅᵝ₂(x))

julia> convert(Polynomial, p)
Polynomials.Polynomial(-0.375 + 0.75*x^2)

julia> monic(p) = (q=convert(Polynomial,p); q/q[end])
monic (generic function with 1 method)

julia> monic(p) ≈  monic(basis(Chebyshev, 2))
true

```

"""
Jacobi

basis_symbol(::Type{<:Jacobi{α,β}}) where {α,β} = "Jᵅᵝ"
Polynomials.domain(::Type{<:Jacobi{α,β}}) where {α,β} =
    Polynomials.Interval{β >= 0 ? Open : Closed,α >= 0 ? Open : Closed}(-1, 1)
abcde(::Type{<:Jacobi{α,β}}) where {α,β} =
    NamedTuple{(:a, :b, :c, :d, :e)}((-1, 0, 1, -(α + β + 2), β - α))

k0(P::Type{<:Jacobi}) = one(eltype(P))
function k1k0(P::Type{<:Jacobi{α,β}}, n::Int) where {α,β}
    n == -1 && return one(eltype(P)) / 1
    γ = 2n + α + β
    val = one(eltype(P))
    if n == 0
        val *= (α + β) / 2 + 1
    else
        num = (γ + 1) * (γ + 2)
        den = 2 * (n + 1) * (γ - n + 1)
        iszero(den) && return k1k0(P, Val(n))
        val *= num / den
    end
    val
end

#kn =  1/2^n ⋅ choose(2n+α+β, n) = (2n+α+β)_n / (2^b  n!)
function leading_term(P::Type{<:Jacobi{α,β}}, n::Int) where {α,β}
    a = α + β + n + 1
    nn = n
    tot = one(eltype(P)) / 1
    for i in 0:(n - 1)
        tot *= a / (2 * nn)
        a += 1
        nn -= 1
    end
    tot
end

weight_function(::Type{<:Jacobi{α,β}}) where {α,β} = x -> (1 - x)^α * (1 + x)^β
generating_function(::Type{<:Jacobi{α,β}}) where {α,β} =
    (t, x) -> begin
        R = sqrt(1 - 2x * t + t^2)
        2^(α + β) * 1 / R * (1 - t + R)^(-α) * (1 + t + R)^(-β)
    end

"""
    jacobi_eval(α, β, n, z)

Evaluate the nth basis element of `Pᵅᵝ` at `z` using

`Pₙᵅᵝ(z) = 2⁻ⁿ∑ (α+n, k) × (β+n, n-k) ⋅ (z+1)ᵏ ⋅ (z-1)ⁿ⁻ᵏ`

This alternate evaluation is  useful when `α` or `β ≤ -1` (and consequently the polynomials are not orthogonal)
"""
function jacobi_eval(α, β, n, z)
    T = eltype(z / 2)
    tot = zero(T)
    for k in 0:n
        tot +=
            generalized_binomial′(α + n, k) *
            generalized_binomial′(β + n, n - k) *
            ((z + 1) / 2)^k *
            ((z - 1) / 2)^(n - k)
    end
    tot
end

# gauss_nodes_weights(p::Type{P}, n) where {α, β, P <: Jacobi{α, β}} =
#     FastGaussQuadrature.gaussjacobi(n, α, β)

function classical_hypergeometric(::Type{<:Jacobi{α,β}}, n, x) where {α,β}
    (α ≤ -1 || β ≤ -1) && throw(ArgumentError("α and β must be > -1"))

    as = (-n, n + α + β + 1)
    bs = (α + 1,)

    Pochhammer_factorial(α + 1, n) * pFq(as, bs, (1 - x) / 2)
end

function norm2(::Type{<:Jacobi{α,β}}, n) where {α,β}
    α > -1 && β > -1 || throw(ArgumentError("α, β > -1 is necessary"))
    2^(α + β + 1) / (2n + α + β + 1) * (Γ(n + α + 1) * Γ(n + β + 1)) /
    (Γ(n + α + β + 1) * Γ(n + 1))
end
function ω₁₀(::Type{<:Jacobi{α,β}}, n) where {α,β}
    val = Γ(n + 1 + α + 1) / Γ(n + α + 1)
    val *= Γ(n + 1 + β + 1) / Γ(n + β + 1)
    val *= Γ(n + α + β + 1) / Γ(n + 1 + α + β + 1)
    val *= Γ(n + 1) / Γ(n + 1 + 1)
    val *= (2n + α + β + 1) / (2(n + 1) + α + β + 1)
    sqrt(val)
end

# overrides
B̃n(P::Type{<:Jacobi{α,β}}, ::Val{0}) where {α,β} =
    iszero(α + β) ? (α - β) * one(eltype(P)) / 2 : (α - β) * one(eltype(P)) / (2(α + β + 2))
C̃n(P::Type{<:Jacobi{α,β}}, ::Val{1}) where {α,β} =
    -one(eltype(P)) * ((α - β)^2 - (α + β + 2)^2) / ((α + β + 2)^2 * (α + β + 3))

b̂̃n(P::Type{<:Jacobi{α,β}}, ::Val{0}) where {α,β} = one(eltype(P)) * NaN
ĉ̃n(P::Type{<:Jacobi}, ::Val{0}) = zero(eltype(P))
ĉ̃n(P::Type{<:Jacobi{α,β}}, ::Val{1}) where {α,β} =
    one(eltype(P)) * ((α - β)^2 - (α + β + 2)^2) /
    ((α + β + 1) * (α + β + 2)^2 * (α + β + 3))

##
## --------------------------------------------------
##

@registerN MonicJacobi AbstractCCOP2 α β
export MonicJacobi
ϟ(::Type{<:MonicJacobi{α,β}}) where {α,β} = Jacobi{α,β}
ϟ(::Type{<:MonicJacobi{α,β,T}}) where {α,β,T} = Jacobi{α,β,T}

@register_monic(MonicJacobi)

@registerN OrthonormalJacobi AbstractCCOP2 α β
export OrthonormalJacobi
ϟ(::Type{<:OrthonormalJacobi{α,β}}) where {α,β} = Jacobi{α,β}
ϟ(::Type{<:OrthonormalJacobi{α,β,T}}) where {α,β,T} = Jacobi{α,β,T}

@register_orthonormal(OrthonormalJacobi)


include("graph.jl")
include("DFS.jl")
M = [0 1 0 1 0 0;
     0 0 0 0 1 0;
     0 0 0 0 1 1;
     0 1 0 0 0 0;
     0 0 0 1 0 0;
     0 0 0 0 0 1]
G = nonEmptyGraph(M)
DFS!(G)
printGraph(G)


using SurfaceGeometry

isdefined(:N) || (N=100)
isdefined(:PasiveStabilisation) || (PasiveStabilisation=true)
isdefined(:TopologyStabilisation) || (TopologyStabilisation=true)
isdefined (:NormalVelocities) || (NormalVelocities=true)
isdefined(:ViewerActive) || (ViewerActive=false)

using JLD
data = load("sphere.jld")
points = data["points"]
faces = data["faces"]

function interpolate(p1,v1,p2)
    v2 = zeros(p2)
    for xkey in 1:size(p2,2)
        d = Inf
        vi = 0
        for ykey in 1:size(p1,2)
            dist = norm(p1[:,ykey] - p2[:,xkey])
            if dist<d
                d = dist
                vi = ykey
            end
        end
        v2[:,xkey] = v1[:,vi]
    end
end

function velocity(t,points,faces)
    v = Array(Float64,size(points)...)
    if NormalVelocities==true
        normals = Array(Float64,size(points)...);
        NormalVectors!(normals,points,faces,i->FaceVRing(i,faces))

        for xkey in 1:size(points,2)
            v[:,xkey] = velocity(t,points[:,xkey])
            nx = normals[:,xkey]
            v[:,xkey] = nx*dot(nx,v[:,xkey])
        end
    else
        for xkey in 1:size(points,2)
            v[:,xkey] = velocity(t,points[:,xkey])
        end
    end
    return v
end

function velocity(t,pos)
    x,y,z = pos

    x = x*0.15 + 0.35
    y = y*0.15 + 0.35
    z = z*0.15 + 0.35

    u = 2*sin(pi*x)^2 * sin(2*pi*y) * sin(2*pi*z) * sin(2/3*pi*t)
    v = - sin(2*pi*x) * sin(pi*y)^2 * sin(2*pi*z) * sin(2/3*pi*t)
    w = - sin(2*pi*x) * sin(2*pi*y) * sin(pi*z)^2 * sin(2/3*pi*t)

    [u,v,w] /0.15
end

memory = []

h = 3/N

t = 0.001
p = points
f = faces
push!(memory,(t,p,f))

vp = zeros(p)
pp = zeros(p)

v1 = zeros(p)
v2 = zeros(p)

zc = SurfaceGeometry.Erdmanis2016(C=0.01,ftol=1e-3)

scale = 0.2
par = Elparameters(
                   m_use_fraction = false,
                   m_min_edge_length = 0.7*scale,
                   m_max_edge_length = 1.5*scale,
                   m_max_volume_change = 0.1*scale^3,
                   m_min_curvature_multiplier = 1.0,
                   m_max_curvature_multiplier = 1.0,
                   m_merge_proximity_epsilon = 0.5*scale,
                   m_proximity_epsilon = 0.00001,
m_perform_improvement = true,
m_collision_safety = false,
m_min_triangle_angle = 15,
m_max_triangle_angle = 120,
m_allow_vertex_movement = true, ###
m_use_curvature_when_collapsing = false,
m_use_curvature_when_splitting = false,
m_dt = h
)

for i in 1:35
    println("step $i just started")

    v1 = velocity(t,p,f)
    if PasiveStabilisation==true
        res = stabilise(p,f,v1;vp=interpolate(pp,vp,p),method=zc)
        println("Finit $(res.Finit) \t Fres $(res.Fres)")
        v1 = res.vres
    end

    v2 = velocity(t + h/2,p + h/2*v1,f)
    if PasiveStabilisation==true
        res = stabilise(p + h/2*v1,f,v2;vp=v1,method=zc)
        println("Finit $(res.Finit) \t Fres $(res.Fres)")
        v2 = res.vres
    end

    vp = copy(v2)
    pp = copy(p)

    ### Adams Bachforth oneliner
    #(t1,p1,v1), (t2,p2) = (t2,p2,v2), (t2+h,p2 + h/2/(t2-t1)*( (2*(t2-t1) + h)*v2 - h*v1))

    if TopologyStabilisation==true
        actualdt,p,f = improvemeshcol(p,f,p + h*v2,par)
        t = t + actualdt
    else
        p = p + h*v2
        t = t + h
    end

    push!(memory,(t,p,f))
end

using Escher
if !ViewerActive
    include(Pkg.dir("Escher", "src", "cli", "serve.jl"))
    @spawn escher_serve(5555,Pkg.dir("SurfaceGeometry","examples","viewers"))
    ViewerActive = true
end


using RecyclePermanent
using RNGPool
using JLD2
using BenchmarkTools

include("../models/gkModel.jl")
include("compare.jl")

@load "data/gkData.jld" ys θ
ys = ys[1:100]
model = makeGKModel(θ, ys)

N = 80*model.n

trials = 1000

setRNGs(12345)

n = 100
N = 50*n

rpio = RPIO{model.Tx}(N, n, 1)

simpleOutput = testEstimator(() -> simpleEstimate(model, rpio), trials)
biasedOutput = testEstimator(() -> biasedEstimate(model, rpio), trials)
rpOutput = testEstimator(() -> rpEstimate(model, rpio), trials)

kmio = KMIO{model.Tx}(N, n)
# ϵ = 1e-2
ϵ = 1e-5
kmEstimate(model, kmio, ϵ)
kmOutput = testEstimator(() -> kmEstimate(model, kmio, ϵ), trials)

@btime rpEstimate($model, $rpio)


include("gibbssampling.jl")
include("hmcsampling.jl")

is_stochastic(::SamplingInference) = false

function log_gp_prior(gp::SampledLatent, f::AbstractVector, X::AbstractVector)
    return logpdf(MvNormal(pr_mean(gp, X), pr_cov(gp)), f)
end

function logjoint(model::MCGP{T,L,<:SamplingInference}, x) where {T,L}
    fs = unsqueeze(model, x)
    return loglikelihood(model.likelihood, yview(model), fs) +
           sum(log_gp_prior.(model.f, fs))
end

function unsqueeze(model::MCGP, f)
    n = model.nFeatures
    return Tuple(f[((i - 1) * n + 1):(i * n)] for i in 1:(model.nLatent))
end

function ∇logjoint(model::MCGP{T,L,<:SamplingInference}, x) where {T,L}
    fs = unsqueeze(model, x)
    return vcat(∇loglikehood(likelihood(model), yview(model), fs)...) +
           vcat(∇logprior(model.f, fs)...)
end

function Distributions.loglikelihood(
    l::AbstractLikelihood, y::AbstractVector, f::Tuple{<:AbstractVector{T}}
) where {T<:Real}
    return sum(logpdf.(l, y, f...))
end

function ∇loglikehood(
    l::AbstractLikelihood, y::AbstractVector, f::Tuple{<:AbstractVector{T}}
) where {T}
    return ∇loglikehood.(l, y, f...)
end

function ∇logprior(gp::AbstractLatent, f)
    return -pr_cov(gp) / f # Remove μ₀ temp
end

function logprior(gp::AbstractLatent, f)
    return -logdet(pr_cov(gp)) / 2 - invquad(pr_cov(gp), f) / 2 # Remove μ₀ temp
end

function store_variables!(i::SamplingInference{T}, fs) where {T}
    return i.sample_store[(n_iter(i) - i.nBurnin) ÷ i.thinning, :, :] .= hcat(fs...)
end


type SDF_MYD_EP <: SDFDynamicSystem
    omg::Float64
    zta::Float64
    kpa::Float64
    phi::Float64
    alpha::Float64
    u0::Float64
    v0::Float64
    hv::AbstractArray{Float64,1}

    SDF_MYD_EP(omg::Float64, zta::Float64, kpa::Float64, phi::Float64, alpha::Float64) = new(omg, zta, kpa, phi, alpha, 0.0, 0.0, zeros(5))
    SDF_MYD_EP(omg::Float64, zta::Float64, kpa::Float64, phi::Float64) = new(omg, zta, kpa, phi, 0.0, 0.0, 0.0, zeros(5))
end

function fd(s::SDF_MYD_EP, u::Float64; update::Bool=false)
    up, dup, fp, kp, status = s.hv[1:5]
    k0 = s.kpa*s.omg*s.omg
    k1 = s.alpha*k0
    fy = s.phi*g
    uy = fy/k0
    du = u - up
    f_try = fp+du*kp
    bup = fy+k1*(u-uy)
    bdn = -fy+k1*(u+uy)

    if du == 0.0
        kc = kp
        f = fp
    elseif status == 0
        if du > 0.0
            if f_try > bup
                f = bup
                kc = k1
                status = 1
            else
                kc = k0
                f = f_try
            end
        else
            if f_try < bdn
                f = bdn
                kc = k1
                status = 1
            else
                kc = k0
                f = f_try
            end
        end
    elseif status == 1
        if dup*du > 0
            f = f_try
            kc = k1
        else
            kc = k0
            f = fp + kc*du
            status = 0
            if f > bup
                f = bup
                kc = k1
                status = 1
            elseif f < bdn
                f = bdn
                kc = k1
                status = 1
            end
        end
    end
    if update
        s.hv[1:5] = [u, du, f, kc, status]
    end
    f
end

function systemEquation(s::SDF_MYD_EP, p::Function)
    function dydt(y::AbstractArray{Float64,1}, t::Float64)
        u = y[1]
        v = y[2]
        omg2 = s.omg*s.omg
        dvdt = p(t)-2.0*s.zta*s.omg*v-fd(s,u)-omg2*u
        [v; dvdt]
    end
    dydt
end

function response(s::SDF_MYD_EP, p::Function, t::AbstractArray{Float64,1})
    y0 = [s.u0; s.v0]

    f = systemEquation(s, p)

    updateHistVar = (y) -> fd(s, y[1]; update=true)
    y = ode4hv(f, y0, t, updateHistVar)

    u = y[:,1]
    v = y[:,2]
    fdz = 2.0*s.zta*s.omg*v
    fdp = y[:,3]
    a = map(p,t)-fdz-fdp-s.omg*s.omg*u

    [u v a fdz fdp]
end