using Cliffords
import LinearAlgebra

# helper to remove the first "+/-" char from Pauli string labels
snip(s::String) = s[nextind(s,1):end];

###############################################################################
## Useful tools for manipulation of matrix representations in a tomography ####
###############################################################################

"""
    unitary2choi(U::AbstractArray)

Returns the Choi super operator representation of a unitary.

# Arguments
- `U`: the unitary matrix representation of the process

# Returns
- `choi`: a choi matrix representation

# Examples
julia> choi = unitary2choi(str2unitary("1QX90p"))
4×4 Array{Complex{Float64},2}:
 0.25+0.0im    0.0+0.25im   0.0+0.25im  0.25+0.0im
  0.0-0.25im  0.25+0.0im   0.25+0.0im    0.0-0.25im
  0.0-0.25im  0.25+0.0im   0.25+0.0im    0.0-0.25im
 0.25+0.0im    0.0+0.25im   0.0+0.25im  0.25+0.0im
"""
function unitary2choi(U::AbstractArray)

    d=size(U,1);
    choi=zeros(ComplexF64,d^2,d^2);
    for ii in 1:d
        for jj in 1:d
            proj = zeros(ComplexF64,d,d);
            proj[ii,jj] = 1;
            choi = choi + kron(proj,U*proj*U')/d;
        end
    end
    return choi
end

"""
    choi2pauliMap(choi::AbstractArray)

Converts a Choi representation to a Pauli Map representation.

# Arguments
- `choi`: the matrix representation of the process

# Returns
- `pauli_map`: a Pauli representation of the choi matrix

# Examples
julia> choi = unitary2choi(str2unitary("1QX90p"))
julia> choi2pauliMap(choi)
4×4 Array{Float64,2}:
  1.0          0.0   0.0          0.0
  0.0          1.0   0.0          0.0
  0.0          0.0   3.33067e-16  1.0
 -5.55112e-17  0.0  -1.0          3.33067e-16
"""
function choi2pauliMap(choi::AbstractArray)

    #Dimension of superoperator
    d2 = size(choi,1);

    #Create the Pauli opearators for n qubits
    num_qubits = Int(log2(sqrt(d2)))
    pauliOps = map(x -> complex(x), vec(permutedims(allpaulis(num_qubits),
                                    reverse(collect(1:num_qubits)))));

    pauliMap = zeros(d2,d2);
    for ct1 in 1:d2
        p1 = pauliOps[ct1]
        for ct2 in 1:d2
            p2 = pauliOps[ct2]
            pauliMap[ct2,ct1] = real(LinearAlgebra.tr(choi*kron(transpose(p2), p1)));
        end
    end
    return pauliMap
end

"""
    unitary2pauli(U:AbstractArray)

Compute the Pauli representation of an arbitrary unitary

# Arguments
- `U`: the matrix representation of the unitary.

# Returns
- pauli_map: Matrix{ComplexF64}

# Examples
julia> unitary2pauli(str2unitary("1QY90p"))
4×4 Array{Float64,2}:
  1.0          0.0          0.0   0.0
  0.0          3.33067e-16  0.0  -1.0
  0.0          0.0          1.0   0.0
 -5.55112e-17  1.0          0.0   3.33067e-16
"""
function unitary2pauli(U::AbstractArray)
    return choi2pauliMap(unitary2choi(U));
end

"""
    str2unitary(strIn::String)

Converts a string description of a gate into a unitary matrix.  Supported
unitaries are: X_4I, IX_4, IX_2, InvCNOT12, CNOT12, CNOT21, 1QXp, 1QX22p,
IY_2, IY, ZXm90_2, 1QHad, XX, 1QId, IX_8, 1QT, YI, X90ZX90, CNOT12_y, X90X90,
1QX45p, ZXm90, 1QX90p, ZX90, 1QY90p, Y90Y90, 1QYp, X_8I, 1QY90m, X_2I, XZm90,
XI, XZ90, 1QX90m, 1QZ90, II, IX, Y_2I

# Arguments
- `strIn`: string id of the unitary

# Returns
- `U`: Matrix{ComplexF64} representing the unitary transformation

# Examples
julia> str2unitary("1QX90p")
2×2 Array{Complex{Float64},2}:
 0.707107+0.0im            0.0-0.707107im
      0.0-0.707107im  0.707107+0.0im
"""
function str2unitary(strIn::String)

    #Basic Pauli operators
    X = complex(Pauli(1));
    Y = complex(Pauli(3));
    Z = complex(Pauli(2));
    I = complex(Pauli(4));

    Uout = Dict("CNOT12" => Matrix{ComplexF64}([[1,0,0,0] [0,1,0,0] [0,0,0,1] [0,0,1,0]]),
                "InvCNOT12" => Matrix{ComplexF64}([[0,1,0,0] [1,0,0,0] [0,0,1,0] [0,0,0,1]]),
                "CNOT21" => Matrix{ComplexF64}([[1,0,0,0] [0,0,0,1] [0,0,1,0] [0,1,0,0]]),
                "1QId" => I,
                "1QX90p" => exp(-1im*pi*X/4),
                "1QX90m" => exp(1im*pi*X/4),
                "1QY90p" => exp(-1im*pi*Y/4),
                "1QY90m" => exp(1im*pi*Y/4),
                "1QXp" => exp(-1im*pi*X/2),
                "1QYp" => exp(-1im*pi*Y/2),
                "1QX45p" => exp(-1im*pi*X/8),
                "1QX22p" => exp(-1im*pi*X/16),
                "1QHad" => exp(-1im*(pi/2)*(1/sqrt(2))*(X+Z)),
                "1QZ90" => exp(-1im*(pi/4)*Z),
                "1QT" => exp(-1im*(pi/8)*Z),
                "II" => kron(I,I),
                "XI" => exp(-1im*kron(X,I)*pi/2),
                "IX" => exp(-1im*kron(I,X)*pi/2),
                "YI" => exp(-1im*kron(Y,I)*pi/2),
                "IY" => exp(-1im*kron(I,Y)*pi/2),
                "X_2I" => exp(-1im*kron(X,I)*pi/4),
                "IX_2" => exp(-1im*kron(I,X)*pi/4),
                "Y_2I" => exp(-1im*kron(Y,I)*pi/4),
                "IY_2" => exp(-1im*kron(I,Y)*pi/4),
                "XZ90" => exp(-1im*kron(X,Z)*pi/4),
                "ZX90" => exp(-1im*kron(Z,X)*pi/4),
                "XZm90" => exp(1im*kron(X,Z)*pi/4),
                "ZXm90" => exp(1im*kron(Z,X)*pi/4),
                "X_8I" => exp(-1im*kron(X,I)*pi/16),
                "X_4I" => exp(-1im*kron(X,I)*pi/8),
                "IX_8" => exp(-1im*kron(I,X)*pi/16),
                "IX_4" => exp(-1im*kron(I,X)*pi/8),
                "CNOT12_y" => Matrix{ComplexF64}([[1,0,0,0] [0,1,0,0] [0,0,-1im,0] [0,0,0,1im]]),
                "X90ZX90" => Matrix{ComplexF64}([[1,0,0,0] [0,0,0,-1im] [0,0,1,0] [0,-1im,0,0]]),
                "XX" => exp(-1im*kron(X,I)*pi/2)*exp(-1im*kron(I,X)*pi/2),
                "X90X90" => exp(-1im*kron(X,I)*pi/4)*exp(-1im*kron(I,X)*pi/4),
                "Y90Y90" => exp(-1im*kron(Y,I)*pi/4)*exp(-1im*kron(I,Y)*pi/4)
            );
    return Uout[strIn]
end

"""
    choi2chi(choi::AbstractArray)

Converts a Choi superoperator to a chi represtation in Pauli basis.

# Arguments
- `choi`: the matrix representation of the process

# Returns
- `chi`: chi matrix representation of the process

# Examples
julia> using RandomQuantum
julia> U = rand(RandomClosedEvolution(2,2))
julia> choi = unitary2choi(U)
julia> choi2chi(choi)
"""
function choi2chi(choi::AbstractArray)

    # Some dimensions
    d2 = size(choi,1);
    d = Int(sqrt(d2));

    # Get the Kraus operators from the eigen decomposition of the Choi
    vals, vecs = LinearAlgebra.eigen(choi);

    # note the ordering diff between Matlab
    vals = LinearAlgebra.diagm(reverse(vals));
    vecs = reverse(vecs, dims=2);

    chi = zeros(ComplexF64,d2,d2);

    #Create the Pauli opearators for n qubits
    num_qubits = Int(log2(sqrt(d2)));
    pauliOps = map(x -> complex(x), vec(permutedims(allpaulis(num_qubits),
                                    reverse(collect(1:num_qubits)))));

    # Transform from the Krauss basis to the Pauli basis
    for kraussct in 1:size(vals,1)
        tmpKrauss = reshape(vecs[:,kraussct],d,d)*sqrt(d); # Krauss operator should have norm of d
        for paulict1 in 1:d2
            pauliLeft = LinearAlgebra.tr(pauliOps[paulict1]*tmpKrauss)/d;
            for paulict2 in 1:d2
                pauliRight = LinearAlgebra.tr(pauliOps[paulict2]*transpose(tmpKrauss))/d;
                chi[paulict1, paulict2] = chi[paulict1, paulict2] + vals[kraussct]*pauliLeft*pauliRight;
            end
        end
    end
    return chi
end

###############################################################################
## Fidelity and plotting tools for tomography #################################
###############################################################################
"""
    getProcessFidelity(choi::AbstractArray, idealProcess::String)

Compute the gate and process fidelity for a give Choi matrix and a given ideal
process matrix.

# Arguments
- `choi`: the matrix representation of the process
- `idealProcess`: string representing the ideal process in str2unitary()

# Returns
- `processFidelity`: -> real(tr(choi*choiIdeal))
- `gateFidelity`: -> (d*processFidelity +1)/(d + 1)

# Examples
julia> getProcessFidelity(choi, "XI")
"""
function getProcessFidelity(choi::AbstractArray, idealProcess::String)
    nbrQubits = log2(sqrt(size(choi, 1)))
    # Calculate the overlaps with the ideal gate
    if typeof(idealProcess) == String
        unitaryIdeal = str2unitary(idealProcess);
    else
        # assume we are passed a matrix otherwise
        unitaryIdeal = idealProcess;
    end
    choiIdeal = unitary2choi(unitaryIdeal);

    processFidelity = real(LinearAlgebra.tr(choi*choiIdeal));
    gateFidelity = (2^nbrQubits*processFidelity+1)/(2^nbrQubits+1);

    return processFidelity, gateFidelity
end

"""
    getProcessFidelity(choi::AbstractArray, idealProcess::AbstractMatrix)

Compute the gate and process fidelity for a give Choi matrix and a given ideal
unitary.

# Arguments
- `choi`: the matrix representation of the process
- `idealProcess`: matrix representing the ideal process unitary

# Returns
- `processFidelity`: -> real(tr(choi*choiIdeal))
- `gateFidelity`: -> (d*processFidelity +1)/(d + 1)

# Examples
julia> getProcessFidelity(choi, U)
"""
function getProcessFidelity(choi::AbstractArray, idealProcess::AbstractMatrix)
    nbrQubits = log2(sqrt(size(choi, 1)))
    # Calculate the overlaps with the ideal gate
    choiIdeal = unitary2choi(idealProcess);

    processFidelity = real(LinearAlgebra.tr(choi*choiIdeal));
    gateFidelity = (2^nbrQubits*processFidelity+1)/(2^nbrQubits+1);

    return processFidelity, gateFidelity
end

"""
    PlotProcess(U::AbstractMatrix; p_map::String="pauli", save_fig::Bool=false)

Plot the Pauli transfer matrix of an operator U.  Only one and two qubit
operators are supported.

# Arguments
- `U`: the matrix representation of the unitary you want to plot.
- `p_map`: which process map to use {"pauli", "chi"}.  "pauli" will plot the
           Pauli transfer matrix and "chi" will plot the chi matrix.
- `save_fig`: set to true if you want a .png of the figure saved in the current
              directory.

# Returns
- void: The function outputs a plot

# Examples
julia> PauliPlot(str2unitary("1QY90p"), save_fig=true)
"""
function PlotProcess(U::AbstractMatrix;
                     p_map::String="pauli",
                     save_fig::Bool=false)

    num_qubits = Int(log2(size(U,1)))
    op_num = 0:(4^num_qubits - 1)
    # This code snippet works for n-qubit unitaries, it's just not that
    # verbose...
    labels = map(p -> snip(string(p)),
                 vec(permutedims(allpaulis(num_qubits),
                 reverse(collect(1:num_qubits)))))
    # a more verbose and brittle option would be something like:
    # if num_qubits == 1
    #     labels = ["I", "X", "Y", "Z"]
    # elseif num_qubits == 2
    #     labels = ["II","IX","IY","IZ",
    #               "XI","XX","XY","XZ",
    #               "YI","YX","YY","YZ",
    #               "ZI","ZX","ZY","ZZ"]
    # else
    #     labels = []
    # end
    fig = figure()
    ax = fig.add_subplot(111)
    if p_map == "pauli"
        cax = ax.matshow(unitary2pauli(U), cmap="bwr")
        ax.set_title("Pauli transfer matrix")
    elseif p_map == "chi"
        cax = ax.matshow(real(choi2chi(unitary2choi(U))), cmap="bwr")
        ax.set_title(L"$\chi$ process matrix")
    end

    fig.colorbar(cax)

    ax.set_xticklabels(labels)
    ax.set_yticklabels(labels)
    ax.set_xticks(op_num)
    ax.set_yticks(op_num)
    ax.xaxis.set_ticks_position("bottom")
    ax.grid(color="grey", linestyle="--", alpha=0.1)
    ax.spines["right"].set_visible(false)
    ax.spines["top"].set_visible(false)
    ax.spines["left"].set_visible(false)
    ax.spines["bottom"].set_visible(false)
    ax.set_xlabel("Input Pauli Operator")
    ax.set_ylabel("Output Pauli Operator")

    if save_fig
        fig.savefig("ProcessPlot.png")
        println("Please rename the saved figure!")
    end
    fig.show()
end

"""
    PauliToCounts(Z::AbstractArray,nbrQubits)

Convert a set of 2^nbrQubits-1 of SigmaZ_j and correlator SigmaZ_iSigmaZ_j... expectation values into state populations (expectation value of projection operators)

# Arguments
- `Z`: array of SigmaZ expectation values. Z[k] is the expectation value of a product of I and Z operators corresponding to the binary value of k.
        So for example if k=0110 , Z[k]=<IZZI> . Z needs to have size of 2^nbrQubits-1
- `nbrQubits`: number of Qubits.

# Returns
- Array of projector expectation values

# Examples
julia> PauliToCounts([1,1,1],2)
"""
function PauliToCounts(Z::AbstractArray,nbrQubits)
  P = ones(1,2^nbrQubits)

  for k in 0:2^nbrQubits-1
       kbin = digits!(zeros(Int64,nbrQubits),k,base=2)
       kbin = -2*kbin .+ 1

       for h in 1:2^nbrQubits-1
           hbin = digits!(zeros(Int64,nbrQubits),h,base=2)
           P[k+1] += prod(sign.(kbin.*hbin .+ 0.5))*Z[h]
       end

   end
   return P/2^nbrQubits
end


# Reference journal article:
# Nikkhoo, M., Walter, T. R., Lundgren, P. R., Prats-Iraola, P. (2016):
# Compound dislocation models (CDMs) for volcano deformation analyses.
# Submitted to Geophysical Journal International
#
# Copyright (c) 2016 Mehdi Nikkhoo
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the
# following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# I appreciate any comments or bug reports.
#
# Mehdi Nikkhoo
# created: 2013.1.25
# Last modified: 2016.10.18
#
# Section 2.1, Physics of Earthquakes and Volcanoes
# Department 2, Geophysics
# Helmholtz Centre Potsdam
# German Research Centre for Geosciences (GFZ)
#
# email:
# mehdi.nikkhoo@gfz-potsdam.de
# mehdi.nikkhoo@gmail.com
#
# website:
# http://www.volcanodeformation.com
#
# Translated by Pengcheng Shi (shipengcheng1230@gmail.com) 12/2018
# The broadcast here causes tremendous amount of compiliation time, needs rewriting!

# export disp_rect_fs, strain_rect_fs, stress_rect_fs

# using LinearAlgebra

# function disp_rect_fs(x::V, y::V, z::V, x₀::T, y₀::T, depth::T, L::T, W::T, plunge::T, dip::T, strike::T, rake::T, slip::T, opening::T, nu::T, ref::S) where {V<:AbstractVector, T, S}
#     @assert size(x) == size(y) == size(z) "Unmatched length of observational coordinates."

#     bx = opening
#     by = slip * cosd(rake)
#     bz = slip * sind(rake)
#     Rz1 = [cosd(plunge) sind(plunge) zero(T); -sind(plunge) cosd(plunge) zero(T); zero(T) zero(T) one(T)]
#     Ry = [cosd(dip) zero(T) sind(dip); zero(T) one(T) zero(T); -sind(dip) zero(T) cosd(dip)]
#     Rz2 = [cosd(strike) sind(strike) zero(T); -sind(strike) cosd(strike) zero(T); zero(T) zero(T) one(T)]
#     Rt = Rz2 * Ry * Rz1

#     Pt1 = [-W/2, L/2, zero(T)]
#     Pt2 = [-W/2, -L/2, zero(T)]
#     Pt3 = [W/2, -L/2, zero(T)]
#     Pt4 = [W/2, L/2, zero(T)]

#     Ptr = _ref_pts(ref, L, W)
#     Pr = [x₀, y₀, -depth] - Rt * Ptr
#     P1 = Rt * Pt1 + Pr
#     P2 = Rt * Pt2 + Pr
#     P3 = Rt * Pt3 + Pr
#     P4 = Rt * Pt4 + Pr

#     ez = [zero(T), zero(T), one(T)]
#     Vnorm = Rt * ez
#     Vstrike = [sind(strike), cosd(strike), zero(T)]
#     Vdip = Vnorm × Vstrike

#     Pm = (P1 + P2 + P3 + P4) / 4
#     p1, p2, p3, p4 = [zeros(T, 3) for _ in 1: 4]
#     At = hcat(Vnorm, Vstrike, Vdip)'

#     x, y, z = coord_trans(x .- Pm[1], y .- Pm[2], z .- Pm[3], At)
#     p1[1], p1[2], p1[3] = coord_trans(P1[1] - Pm[1], P1[2] - Pm[2], P1[3] - Pm[3], At)
#     p2[1], p2[2], p2[3] = coord_trans(P2[1] - Pm[1], P2[2] - Pm[2], P2[3] - Pm[3], At)
#     p3[1], p3[2], p3[3] = coord_trans(P3[1] - Pm[1], P3[2] - Pm[2], P3[3] - Pm[3], At)
#     p4[1], p4[2], p4[3] = coord_trans(P4[1] - Pm[1], P4[2] - Pm[2], P4[3] - Pm[3], At)

#     e12 = normalize(p2 - p1)
#     e23 = normalize(p3 - p2)
#     e34 = normalize(p4 - p3)
#     e14 = normalize(p4 - p1)

#     Rectmode = @views rectmodefInder(y, z, x, p1[2:3], p2[2:3], p3[2:3], p4[2:3])

#     # use list comprehension causes being unable infer type (until Julia v1.2), see one example mentioned in the link below
#     # https://discourse.julialang.org/t/comprehension-type-inference/4324
#     # u, v, w = [zeros(T, length(x)) for _ ∈ 1: 3]
#     u = zeros(T, length(x)); v = zeros(T, length(x)); w = zeros(T, length(x))

#     casepLog = Rectmode .== 1
#     casenLog = Rectmode .== -1
#     casezLog = Rectmode .== 0

#     if any(casepLog)
#         xp, yp, zp = x[casepLog], y[casepLog], z[casepLog]
#         utmp, vtmp, wtmp = [Vector{T}(undef, length(xp)) for _ ∈ 1: 3]

#         utmp, vtmp, wtmp = RDSetupD(xp, yp, zp, bx, by, bz, nu, p1, e14)
#         u[casepLog] .+= utmp; v[casepLog] .+= vtmp; w[casepLog] .+= wtmp

#         utmp, vtmp, wtmp = RDSetupD(xp, yp, zp, bx, by, bz, nu, p2, -e12)
#         u[casepLog] .+= utmp; v[casepLog] .+= vtmp; w[casepLog] .+= wtmp

#         utmp, vtmp, wtmp = RDSetupD(xp, yp, zp, bx, by, bz, nu, p3, -e23)
#         u[casepLog] .+= utmp; v[casepLog] .+= vtmp; w[casepLog] .+= wtmp

#         utmp, vtmp, wtmp = RDSetupD(xp, yp, zp, bx, by, bz, nu, p4, -e34)
#         u[casepLog] .+= utmp; v[casepLog] .+= vtmp; w[casepLog] .+= wtmp
#     end
#     if any(casenLog)
#         xn, yn, zn = x[casenLog], y[casenLog], z[casenLog]
#         utmp, vtmp, wtmp = [Vector{T}(undef, length(xn)) for _ ∈ 1: 3]

#         utmp, vtmp, wtmp = RDSetupD(xn, yn, zn, bx, by, bz, nu, p1, -e14)
#         u[casenLog] .+= utmp; v[casenLog] .+= vtmp; w[casenLog] .+= wtmp

#         utmp, vtmp, wtmp = RDSetupD(xn, yn, zn, bx, by, bz, nu, p2, e12)
#         u[casenLog] .+= utmp; v[casenLog] .+= vtmp; w[casenLog] .+= wtmp

#         utmp, vtmp, wtmp = RDSetupD(xn, yn, zn, bx, by, bz, nu, p3, e23)
#         u[casenLog] .+= utmp; v[casenLog] .+= vtmp; w[casenLog] .+= wtmp

#         utmp, vtmp, wtmp = RDSetupD(xn, yn, zn, bx, by, bz, nu, p4, e34)
#         u[casenLog] .+= utmp; v[casenLog] .+= vtmp; w[casenLog] .+= wtmp
#     end
#     if any(casezLog)
#         u[casezLog] .= NaN; v[casezLog] .= NaN; w[casezLog] .= NaN
#     end

#     Fi = BurgersFuncRD(x, y, z, p1, p2, p3, p4)
#     @. u += Fi * bx
#     @. v += Fi * by
#     @. w += Fi * bz
#     u, v, w = coord_trans(u, v, w, At')
#     return u, v, w
# end

# function rd_disp_hs()
#     #TODO
# end

# function strain_rect_fs(x::V, y::V, z::V, x₀::T, y₀::T, depth::T, L::T, W::T, plunge::T, dip::T, strike::T, rake::T, slip::T, opening::T, λ::T, μ::T, ref::S) where {V<:AbstractVector, T, S}
#     nu = λ / (λ + μ) / 2

#     # lots of repeated codes here
#     bx = opening
#     by = slip * cosd(rake)
#     bz = slip * sind(rake)
#     Rz1 = [cosd(plunge) sind(plunge) zero(T); -sind(plunge) cosd(plunge) zero(T); zero(T) zero(T) one(T)]
#     Ry = [cosd(dip) zero(T) sind(dip); zero(T) one(T) zero(T); -sind(dip) zero(T) cosd(dip)]
#     Rz2 = [cosd(strike) sind(strike) zero(T); -sind(strike) cosd(strike) zero(T); zero(T) zero(T) one(T)]
#     Rt = Rz2 * Ry * Rz1

#     Pt1 = [-W/2, L/2, zero(T)]
#     Pt2 = [-W/2, -L/2, zero(T)]
#     Pt3 = [W/2, -L/2, zero(T)]
#     Pt4 = [W/2, L/2, zero(T)]

#     Ptr = _ref_pts(ref, L, W)
#     Pr = [x₀, y₀, -depth] - Rt * Ptr
#     P1 = Rt * Pt1 + Pr
#     P2 = Rt * Pt2 + Pr
#     P3 = Rt * Pt3 + Pr
#     P4 = Rt * Pt4 + Pr

#     ez = [zero(T), zero(T), one(T)]
#     Vnorm = Rt * ez
#     Vstrike = [sind(strike), cosd(strike), zero(T)]
#     Vdip = Vnorm × Vstrike

#     Pm = (P1 + P2 + P3 + P4) / 4
#     p1, p2, p3, p4 = [zeros(T, 3) for _ in 1: 4]
#     At = hcat(Vnorm, Vstrike, Vdip)'

#     x, y, z = coord_trans(x .- Pm[1], y .- Pm[2], z .- Pm[3], At)
#     p1[1], p1[2], p1[3] = coord_trans(P1[1] - Pm[1], P1[2] - Pm[2], P1[3] - Pm[3], At)
#     p2[1], p2[2], p2[3] = coord_trans(P2[1] - Pm[1], P2[2] - Pm[2], P2[3] - Pm[3], At)
#     p3[1], p3[2], p3[3] = coord_trans(P3[1] - Pm[1], P3[2] - Pm[2], P3[3] - Pm[3], At)
#     p4[1], p4[2], p4[3] = coord_trans(P4[1] - Pm[1], P4[2] - Pm[2], P4[3] - Pm[3], At)

#     e12 = normalize(p2 - p1)
#     e23 = normalize(p3 - p2)
#     e34 = normalize(p4 - p3)
#     e14 = normalize(p4 - p1)

#     Rectmode = @views rectmodefInder(y, z, x, p1[2:3], p2[2:3], p3[2:3], p4[2:3])

#     casepLog = Rectmode .== 1
#     casenLog = Rectmode .== -1
#     casezLog = Rectmode .== 0
#     lenx = length(x)

#     # use list comprehension causes being unable infer type (until Julia v1.2), see one example mentioned in the link below
#     # https://discourse.julialang.org/t/comprehension-type-inference/4324
#     # exx, eyy, ezz, exy, exz, eyz = [zeros(T, length(x)) for _ in 1: 6]
#     exx, eyy, ezz, exy, exz, eyz = zeros(T, lenx), zeros(T, lenx), zeros(T, lenx), zeros(T, lenx), zeros(T, lenx), zeros(T, lenx)

#     if any(casepLog)
#         xp, yp, zp = x[casepLog], y[casepLog], z[casepLog]
#         exxt, eyyt, ezzt, exyt, exzt, eyzt = [Vector{T}(undef, length(xp)) for _ ∈ 1: 6]

#         exxt, eyyt, ezzt, exyt, exzt, eyzt = RDSetupS(xp, yp, zp, bx, by, bz, nu, p1, e14)
#         exx[casepLog] .+= exxt; eyy[casepLog] .+= eyyt; ezz[casepLog] .+= ezzt; exy[casepLog] .+= exyt; exz[casepLog] .+= exzt; eyz[casepLog] .+= eyzt

#         exxt, eyyt, ezzt, exyt, exzt, eyzt = RDSetupS(xp, yp, zp, bx, by, bz, nu, p2, -e12)
#         exx[casepLog] .+= exxt; eyy[casepLog] .+= eyyt; ezz[casepLog] .+= ezzt; exy[casepLog] .+= exyt; exz[casepLog] .+= exzt; eyz[casepLog] .+= eyzt

#         exxt, eyyt, ezzt, exyt, exzt, eyzt = RDSetupS(xp, yp, zp, bx, by, bz, nu, p3, -e23)
#         exx[casepLog] .+= exxt; eyy[casepLog] .+= eyyt; ezz[casepLog] .+= ezzt; exy[casepLog] .+= exyt; exz[casepLog] .+= exzt; eyz[casepLog] .+= eyzt

#         exxt, eyyt, ezzt, exyt, exzt, eyzt = RDSetupS(xp, yp, zp, bx, by, bz, nu, p4, -e34)
#         exx[casepLog] .+= exxt; eyy[casepLog] .+= eyyt; ezz[casepLog] .+= ezzt; exy[casepLog] .+= exyt; exz[casepLog] .+= exzt; eyz[casepLog] .+= eyzt
#     end
#     if any(casenLog)
#         xn, yn, zn = x[casenLog], y[casenLog], z[casenLog]
#         exxt, eyyt, ezzt, exyt, exzt, eyzt = [Vector{T}(undef, length(xn)) for _ ∈ 1: 6]

#         exxt, eyyt, ezzt, exyt, exzt, eyzt = RDSetupS(xn, yn, zn, bx, by, bz, nu, p1, -e14)
#         exx[casenLog] .+= exxt; eyy[casenLog] .+= eyyt; ezz[casenLog] .+= ezzt; exy[casenLog] .+= exyt; exz[casenLog] .+= exzt; eyz[casenLog] .+= eyzt

#         exxt, eyyt, ezzt, exyt, exzt, eyzt = RDSetupS(xn, yn, zn, bx, by, bz, nu, p2, e12)
#         exx[casenLog] .+= exxt; eyy[casenLog] .+= eyyt; ezz[casenLog] .+= ezzt; exy[casenLog] .+= exyt; exz[casenLog] .+= exzt; eyz[casenLog] .+= eyzt

#         exxt, eyyt, ezzt, exyt, exzt, eyzt = RDSetupS(xn, yn, zn, bx, by, bz, nu, p3, e23)
#         exx[casenLog] .+= exxt; eyy[casenLog] .+= eyyt; ezz[casenLog] .+= ezzt; exy[casenLog] .+= exyt; exz[casenLog] .+= exzt; eyz[casenLog] .+= eyzt

#         exxt, eyyt, ezzt, exyt, exzt, eyzt = RDSetupS(xn, yn, zn, bx, by, bz, nu, p4, e34)
#         exx[casenLog] .+= exxt; eyy[casenLog] .+= eyyt; ezz[casenLog] .+= ezzt; exy[casenLog] .+= exyt; exz[casenLog] .+= exzt; eyz[casenLog] .+= eyzt
#     end
#     if any(casezLog)
#         exx[casezLog] .= NaN; eyy[casezLog] .= NaN; ezz[casezLog] .= NaN; exy[casezLog] .= NaN; exz[casezLog] .= NaN; eyz[casezLog] .= NaN
#     end
#     exx, eyy, ezz, exy, exz, eyz = TensTrans(exx, eyy, ezz, exy, exz, eyz, At')
# end

# function rd_strain_hs()
#     #TODO
# end

# function stress_rect_fs(x::V, y::V, z::V, x₀::T, y₀::T, depth::T, L::T, W::T, plunge::T, dip::T, strike::T, rake::T, slip::T, opening::T, λ::T, μ::T, ref::S) where {V, T, S}
#     exx, eyy, ezz, exy, exz, eyz = strain_rect_fs(x, y, z, x₀, y₀, depth, L, W, plunge, dip, strike, rake, slip, opening, λ, μ, ref)
#     ekk = exx + eyy + ezz
#     Sxx = @. 2μ * exx + λ * ekk
#     Syy = @. 2μ * eyy + λ * ekk
#     Szz = @. 2μ * ezz + λ * ekk
#     Sxy = @. 2μ * exy
#     Sxz = @. 2μ * exz
#     Syz = @. 2μ * eyz
#     return Sxx, Syy, Szz, Sxy, Sxz, Syz
# end

# function rd_stress_hs()
#     #TODO
# end

# @inline _ref_pts(::Val{:pc}, L::T, W::T) where T = zeros(T, 3)
# @inline _ref_pts(::Val{:p1}, L::T, W::T) where T = [-W/2, L/2, zero(T)]
# @inline _ref_pts(::Val{:p2}, L::T, W::T) where T = [-W/2, -L/2, zero(T)]
# @inline _ref_pts(::Val{:p3}, L::T, W::T) where T = [W/2, -L/2, zero(T)]
# @inline _ref_pts(::Val{:p4}, L::T, W::T) where T = [W/2, L/2, zero(T)]
# @inline _ref_pts(::Val{:mp12}, L::T, W::T) where T = [-W/2, zero(T), zero(T)]
# @inline _ref_pts(::Val{:mp23}, L::T, W::T) where T = [zero(T), -L/2, zero(T)]
# @inline _ref_pts(::Val{:mp34}, L::T, W::T) where T = [W/2, zero(T), zero(T)]
# @inline _ref_pts(::Val{:mp41}, L::T, W::T) where T = [zero(T), L/2, zero(T)]
# @inline _ref_pts(::Val{:mp21}, L::T, W::T) where T = _ref_pts(Val(:mp12), L, W)
# @inline _ref_pts(::Val{:mp32}, L::T, W::T) where T = _ref_pts(Val(:mp23), L, W)
# @inline _ref_pts(::Val{:mp43}, L::T, W::T) where T = _ref_pts(Val(:mp34), L, W)
# @inline _ref_pts(::Val{:mp14}, L::T, W::T) where T = _ref_pts(Val(:mp14), L, W)

# @inline function rectmodefInder(x::T, y::T, z::T, p1::U, p2::U, p3::U, p4::U) where {T, U}
#     pm = (p1 + p2 + p3 + p4) / 4
#     e21 = normalize(p1 - p2)
#     e41 = normalize(p1 - p4)
#     A = hcat(e21, e41)
#     rectmode = ones(Int, length(x))
#     r = A' * hcat(p1 - pm, p2 - pm, p3 - pm, p4 - pm)
#     P1, P2, P3, P4 = @views r[:,1], r[:,2], r[:,3], r[:,4]

#     @inbounds @simd for i ∈ eachindex(rectmode)
#         _x, _y = x[i] - pm[1], y[i] - pm[2]
#         r1 = A[1,1] * _x + A[2,1] * _y
#         r2 = A[1,2] * _x + A[2,2] * _y
#         _x, _y = r1, r2
#         if (_x ≥ 0 && _y ≥ 0 && ((_y - P1[2]) < (_x - P1[1]))) ||
#             (_x ≤ 0 && _y ≥ 0 && ((_y - P2[2]) > -(_x - P2[1]))) ||
#             (_x ≤ 0 && _y ≤ 0 && ((_y - P3[2]) > (_x - P3[1]))) ||
#             (_x ≥ 0 && _y ≤ 0 && ((_y - P4[2]) < -(_x - P4[1]))) ||
#             (_x < P1[1] && _x > P3[1] && _y < P1[2] && _y > P3[2])
#             rectmode[i] = -1
#         elseif ((_x == P1[1] || _x == P3[1]) && _y ≤ P1[2] && _y ≥ P3[2] && z[i] == 0) ||
#             ((_y == P1[2] || _y == P3[2]) && _x ≤ P1[1] && _x ≥ P3[1] && z[i] == 0)
#             rectmode[i] = 0
#         end
#     end
#     rectmode
# end

# @inline function RDSetupD(x::U, y::U, z::U, bx::T, by::T, bz::T, nu::T, RDVertex::V, SideVec::V) where {T, U, V}
#     y1 = @. SideVec[3] * (y - RDVertex[2]) - SideVec[2] * (z - RDVertex[3])
#     z1 = @. SideVec[2] * (y - RDVertex[2]) + SideVec[3] * (z - RDVertex[3])
#     by1 = @. SideVec[3] * by - SideVec[2] * bz
#     bz1 = @. SideVec[2] * by + SideVec[3] * bz

#     u, v0, w0 = AngDisDisp(x, y1, z1, -π/2, bx, by1, bz1, nu)
#     v = @. SideVec[3] * v0 + SideVec[2] * w0
#     w = @. -SideVec[2] * v0 + SideVec[3] * w0
#     return u, v, w
# end

# @inline function BurgersFuncRD(x::T, y::T, z::T, p1::U, p2::U, p3::U, p4::U) where {T, U}
#     Fi = zeros(size(x))
#     Ind = @. (abs(y) ≤ abs(p1[2] - p2[2])/2) & (abs(z) ≤ abs(p1[3] - p4[3])/2)

#     if any(Ind)
#         xI, yI, zI = x[Ind], y[Ind], z[Ind]

#         FiD1 = @. √(xI ^ 2 + (yI - p1[2]) ^ 2 + (zI - p1[3]) ^ 2) - (zI - p1[3]) - (yI - p1[2])
#         FiD2 = @. √(xI ^ 2 + (yI - p2[2]) ^ 2 + (zI - p2[3]) ^ 2) - (zI - p2[3]) - (yI - p2[2])
#         FiD3 = @. √(xI ^ 2 + (yI - p3[2]) ^ 2 + (zI - p3[3]) ^ 2) - (zI - p3[3]) - (yI - p3[2])
#         FiD4 = @. √(xI ^ 2 + (yI - p4[2]) ^ 2 + (zI - p4[3]) ^ 2) - (zI - p4[3]) - (yI - p4[2])

#         FiNt = @. xI * (FiD1 + FiD3) * (FiD2 * FiD4 - xI ^ 2) - xI * (FiD2 + FiD4) * (FiD1 * FiD3 - xI ^ 2)
#         FiDt = @. (FiD1 * FiD3 - xI ^ 2) * (FiD2 * FiD4 - xI ^ 2) + xI ^ 2 * (FiD1 + FiD3) * (FiD2 + FiD4)
#         Fi[Ind] = @. atan(FiNt, FiDt) / 2 / π
#     elseif !all(Ind)
#         bool = .!Ind
#         xO, yO, zO = x[bool], y[bool], z[bool]
#         a = cat(-xO, p1[2] .- yO, p1[3] .- zO; dims=2)
#         b = cat(-xO, p2[2] .- yO, p2[3] .- zO; dims=2)
#         c = cat(-xO, p3[2] .- yO, p3[3] .- zO; dims=2)
#         d = cat(-xO, p4[2] .- yO, p4[3] .- zO; dims=2)
#         na, nb, nc, nd = vec(.√(sum(abs2, a; dims=2))), vec(.√(sum(abs2, b; dims=2))), vec(.√(sum(abs2, c; dims=2))), vec(.√(sum(abs2, d; dims=2)))
#         @views begin
#             FiN1t = @. a[:,1] * (b[:,2] * c[:,3] - b[:,3] * c[:,2]) - a[:,2] * (b[:,1] * c[:,3]-b[:,3] * c[:,1]) + a[:,3] * (b[:,1] * c[:,2]-b[:,2] * c[:,1])
#             FiD1t = @. na * nb * nc + $vec($sum(a * b; dims=2)) * nc + $vec($sum(a * c; dims=2)) * nb + $vec($sum(b * c; dims=2)) * na
#             FiN2t = @. a[:,1] * (c[:,2] * d[:,3] - c[:,3] * d[:,2]) - a[:,2] * (c[:,1] * d[:,3] - c[:,3] * d[:,1]) + a[:,3] * (c[:,1] * d[:,2]-c[:,2] * d[:,1])
#             FiD2t = @. na * nc * nd + $vec($sum(a * c; dims=2)) * nd + $vec($sum(a * d; dims=2)) * nc + $vec($sum(c * d; dims=2)) * na
#         end
#         Fi[bool] = @. -2 * atan(FiN1t * FiD2t + FiN2t * FiD1t, FiD1t * FiD2t - FiN1t * FiN2t) / 4 / π
#     end
#     return Fi
# end

# @inline function RDSetupS(x::T, y::T, z::T, bx::U, by::U, bz::U, nu::U, RDVertex::V, SideVec::V) where {T, U, V}
#     y1 = @. SideVec[3] * (y - RDVertex[2]) - SideVec[2] * (z - RDVertex[3])
#     z1 = @. SideVec[2] * (y - RDVertex[2]) + SideVec[3] * (z - RDVertex[3])
#     by1 = @. SideVec[3] * by - SideVec[2] * bz
#     bz1 = @. SideVec[2] * by + SideVec[3] * bz

#     exx, eyy, ezz, exy, exz, eyz = AngDisStrain(x, y1, z1, -pi/2, bx, by1, bz1, nu)

#     B = [1 0 0; 0 SideVec[3] SideVec[2]; 0 -SideVec[2] SideVec[3]]
#     exx, eyy, ezz, exy, exz, eyz = TensTrans(exx, eyy, ezz, exy, exz, eyz, B)
#     return exx, eyy, ezz, exy, exz, eyz
# end

# function disp_rect_fs(x::V, y::V, z::V, x₀::T, y₀::T, depth::T, L::T, W::T, plunge::T, dip::T, strike::T, rake::T, slip::T, opening::T, nu::T, ref::S) where {V<:Number, T, S}
#     u, v, w = disp_rect_fs([x], [y], [z], x₀, y₀, depth, L, W, plunge, dip, strike, rake, slip, opening, nu, ref)
#     u[1], v[1], w[1]
# end

# function strain_rect_fs(x::V, y::V, z::V, x₀::T, y₀::T, depth::T, L::T, W::T, plunge::T, dip::T, strike::T, rake::T, slip::T, opening::T, λ::T, μ::T, ref::S) where {V<:Number, T, S}
#     out = strain_rect_fs([x], [y], [z], x₀, y₀, depth, L, W, plunge, dip, strike, rake, slip, opening, λ, μ, ref)
#     [q[1] for q in out]
# end


@doc raw"""
    HessianProblem <: Problem

specify a problem for hessian based algorithms.

# Fields
* `M`            : a manifold $\mathcal M$
* `cost` : a function $F\colon\mathcal M\to\mathbb R$ to minimize
* `gradient`     : the gradient $\nabla F\colon\mathcal M
  \to \mathcal T\mathcal M$ of the cost function $F$
* `hessian`      : the hessian $\operatorname{Hess}[F] (\cdot)_ {x} \colon \mathcal T_{x} \mathcal M
  \to \mathcal T_{x} \mathcal M$ of the cost function $F$
* `precon`       : the symmetric, positive deﬁnite
    preconditioner (approximation of the inverse of the Hessian of $F$)

# See also
[`truncated_conjugate_gradient_descent`](@ref), [`trust_regions`](@ref)
"""
struct HessianProblem{mT <: Manifold} <: Problem
    M::mT
    cost::Function
    gradient::Function
    hessian::Union{Function,Missing}
    precon::Function
    HessianProblem(M::mT,cost::Function,grad::Function,hess::Union{Function,Missing},pre::Function) where {mT <: Manifold} = new{mT}(M,cost,grad,hess,pre)
end

abstract type HessianOptions <: Options end
#
# Options
#
@doc raw"""
    TruncatedConjugateGradientOptions <: HessianOptions

describe the Steihaug-Toint truncated conjugate-gradient method, with

# Fields
a default value is given in brackets if a parameter can be left out in initialization.

* `x` : a point, where the trust-region subproblem needs
    to be solved
* `stop` : a function s,r = @(o,iter,ξ,x,xnew) returning a stop
    indicator and a reason based on an iteration number, the gradient and the
    last and current iterates
* `η` : a tangent vector (called update vector), which solves the
    trust-region subproblem after successful calculation by the algorithm
* `δ` : search direction
* `Δ` : the trust-region radius
* `residual` : the gradient
* `useRand` : indicates if the trust-region solve and so the algorithm is to be
        initiated with a random tangent vector. If set to true, no
        preconditioner will be used. This option is set to true in some
        scenarios to escape saddle points, but is otherwise seldom activated.

# Constructor

    TruncatedConjugateGradientOptions(x, stop, eta, delta, Delta, res, uR)

construct a truncated conjugate-gradient Option with the fields as above.

# See also
[`truncated_conjugate_gradient_descent`](@ref), [`trust_regions`](@ref)
"""
mutable struct TruncatedConjugateGradientOptions{P,T} <: HessianOptions
    x::P
    stop::StoppingCriterion
    η::T
    δ::T
    Δ
    residual::T
    useRand::Bool
    function TruncatedConjugateGradientOptions(
        x::P,
        stop::StoppingCriterion,
        η::T,
        δ::T,
        Δ,
        residual::T,
        uR::Bool
        ) where {P, T}
        return new{typeof(x),typeof(η)}(x,stop,η,δ,Δ,residual,uR)
    end
end

@doc raw"""
    TrustRegionsOptions <: HessianOptions

describe the trust-regions solver, with

# Fields
a default value is given in brackets if a parameter can be left out in initialization.

* `x` : a point as starting point
* `stop` : a function s,r = @(o,iter) returning a stop
    indicator and a reason based on an iteration number and the gradient
* `Δ` : the (initial) trust-region radius
* `Δ_bar` : the maximum trust-region radius
* `useRand` : indicates if the trust-region solve is to be initiated with a
        random tangent vector. If set to true, no preconditioner will be
        used. This option is set to true in some scenarios to escape saddle
        points, but is otherwise seldom activated.
* `ρ_prime` : a lower bound of the performance ratio for the iterate that
        decides if the iteration will be accepted or not. If not, the
        trust-region radius will have been decreased. To ensure this,
        ρ'>= 0 must be strictly smaller than 1/4. If ρ' is negative,
        the algorithm is not guaranteed to produce monotonically decreasing
        cost values. It is strongly recommended to set ρ' > 0, to aid
        convergence.
* `ρ_regularization` : Close to convergence, evaluating the performance ratio ρ
        is numerically challenging. Meanwhile, close to convergence, the
        quadratic model should be a good fit and the steps should be
        accepted. Regularization lets ρ go to 1 as the model decrease and
        the actual decrease go to zero. Set this option to zero to disable
        regularization (not recommended). When this is not zero, it may happen
        that the iterates produced are not monotonically improving the cost
        when very close to convergence. This is because the corrected cost
        improvement could change sign if it is negative but very small.

# Constructor

    TrustRegionsOptions(x, stop, delta, delta_bar, uR, rho_prime, rho_reg)

construct a trust-regions Option with the fields as above.

# See also
[`trust_regions`](@ref)
"""
mutable struct TrustRegionsOptions <: HessianOptions
    x
    stop::StoppingCriterion
    Δ
    Δ_bar
    retraction::Function
    useRand::Bool
    ρ_prime
    ρ_regularization
    function TrustRegionsOptions(
        x,
        stop::StoppingCriterion,
        Δ,
        Δ_bar,
        retr::Function,
        useRand::Bool,
        ρ_prime,
        ρ_regularization
        )
        return new(x, stop, Δ, Δ_bar, retr, useRand, ρ_prime, ρ_regularization)
    end
end

@doc raw"""
    getHessian(p,x,ξ)

evaluate the Hessian of a [`HessianProblem`](@ref) `p` at the point `x`
applied to a tangent vector `ξ`.
"""
function getHessian(p::HessianProblem, x, ξ; stepsize=2*10^(-14))
    if ismissing(p.hessian)
        return approxHessianFD(p.M, x -> get_gradient(p.M,x) ,x,ξ; stepsize=stepsize)
    else
        return p.hessian(p.M,x,ξ)
    end
end

@doc raw"""
    get_gradient(p,x)

evaluate the gradient of a [`HessianProblem`](@ref)`p` at the
point `x`.
"""
get_gradient(p::HessianProblem,x) = p.gradient(p.M,x)
@doc raw"""
    get_preconditioner(p,x,ξ)

evaluate the symmetric, positive deﬁnite preconditioner (approximation of the
inverse of the Hessian of the cost function `F`) of a
[`HessianProblem`](@ref) `p` at the point `x`applied to a
tangent vector `ξ`.
"""
get_preconditioner(p::Pr, x, ξ) where {Pr <: HessianProblem} = p.precon(p.M, x, ξ)

@doc raw"""
    approxHessianFD(p,x,ξ,[stepsize=2.0^(-14)])

return an approximated solution of the Hessian of the cost function applied to
a tangent vector `ξ` by using a generic finite difference approximation
based on computations of the gradient.

Input
* `p` – a Manopt problem structure (already containing the manifold and enough
        information to compute the cost gradient)
* `x` – a point where the Hessian is to be approximated
* `ξ` – a tangent vector on which the approximated Hessian is to be applied

# Optional
* `stepsize` – the length of the step with which the method should work

# Output
* a tangent vector generated by applying the approximated Hessian to the
    tangent vector ξ
"""
function approxHessianFD(M::MT, x, gradFct::Function, ξ; stepsize=2.0^(-14), transport=ParallelTransport()) where {MT <: Manifold}
    norm_xi = norm(M,x,ξ)
    if norm_xi < eps(Float64)
        return zero_tangent_vector(M, x)
    end
    c = stepsize / norm_xi
    grad = gradFct(x)
    x1 = exp(M, x, ξ, c)
    grad1 = gradFct(x1)
    grad1 = vector_transport_to(M, x1, grad1, x, transport)
    return (1/c)*(grad1 - grad)
end

@doc raw"""
    stopIfResidualIsReducedByFactor <: StoppingCriterion

A functor for testing if the norm of residual at the current iterate is reduced
by a factor compared to the norm of the initial residual, i.e.
$\Vert r_k \Vert_x \leqq \kappa \Vert r_0 \Vert_x$.
In this case the algorithm reached linear convergence.

# Fields
* `κ` – the reduction factor
* `initialResidualNorm` - stores the norm of the residual at the initial vector
    $\eta$ of the Steihaug-Toint tcg mehtod [`truncated_conjugate_gradient_descent`](@ref)
* `reason` – stores a reason of stopping if the stopping criterion has one be
  reached, see [`get_reason`](@ref).

# Constructor

    stopIfResidualIsReducedByFactor(iRN, κ)

initialize the stopIfResidualIsReducedByFactor functor to indicate to stop after
the norm of the current residual is lesser than the norm of the initial residual
iRN times κ.

# See also
[`truncated_conjugate_gradient_descent`](@ref), [`trust_regions`](@ref)
"""
mutable struct stopIfResidualIsReducedByFactor <: StoppingCriterion
    κ::Float64
    initialResidualNorm::Float64
    reason::String
    stopIfResidualIsReducedByFactor(iRN::Float64,κ::Float64) = new(κ,iRN,"")
end
function (c::stopIfResidualIsReducedByFactor)(p::P,o::O,i::Int) where {P <: HessianProblem, O <: TruncatedConjugateGradientOptions}
    if norm(p.M, o.x, o.residual) <= c.initialResidualNorm*c.κ && i > 0
        c.reason = "The algorithm reached linear convergence (residual at least reduced by κ=$(c.κ)).\n"
        return true
    end
    return false
end

@doc raw"""
    stopIfResidualIsReducedByPower <: StoppingCriterion

A functor for testing if the norm of residual at the current iterate is reduced
by a power of 1+θ compared to the norm of the initial residual, i.e.
$\Vert r_k \Vert_x \leqq  \Vert r_0 \Vert_{x}^{1+\theta}$. In this case the
algorithm reached superlinear convergence.

# Fields
* `θ` – part of the reduction power
* `initialResidualNorm` - stores the norm of the residual at the initial vector
    $\eta$ of the Steihaug-Toint tcg mehtod [`truncated_conjugate_gradient_descent`](@ref)
* `reason` – stores a reason of stopping if the stopping criterion has one be
    reached, see [`get_reason`](@ref).

# Constructor

    stopIfResidualIsReducedByPower(iRN, θ)

initialize the stopIfResidualIsReducedByFactor functor to indicate to stop after
the norm of the current residual is lesser than the norm of the initial residual
iRN to the power of 1+θ.

# See also
[`truncated_conjugate_gradient_descent`](@ref), [`trust_regions`](@ref)
"""
mutable struct stopIfResidualIsReducedByPower <: StoppingCriterion
    θ::Float64
    initialResidualNorm::Float64
    reason::String
    stopIfResidualIsReducedByPower(iRN::Float64,θ::Float64) = new(θ,iRN,"")
end
function (c::stopIfResidualIsReducedByPower)(p::P,o::O,i::Int) where {P <: HessianProblem, O <: TruncatedConjugateGradientOptions}
    if norm(p.M, o.x, o.residual) <= c.initialResidualNorm^(1+c.θ) && i > 0
        c.reason = "The algorithm reached superlinear convergence (residual at least reduced by power 1 + θ=$(1+(c.θ))).\n"
        return true
    end
    return false
end

@doc raw"""
    StopWhenTrustRegionIsExceeded <: StoppingCriterion

A functor for testing if the norm of the next iterate in the  Steihaug-Toint tcg
mehtod is larger than the trust-region radius, i.e. $\Vert η_{k}^{*} \Vert_x
≧ Δ$. terminate the algorithm when the trust region has been left.

# Fields
* `reason` – stores a reason of stopping if the stopping criterion has one be
    reached, see [`get_reason`](@ref).
* `storage` – stores the necessary parameters `η, δ, residual` to check the
    criterion.

# Constructor

    StopWhenTrustRegionIsExceeded([a])

initialize the StopWhenTrustRegionIsExceeded functor to indicate to stop after
the norm of the next iterate is greater than the trust-region radius using the
[`StoreOptionsAction`](@ref) `a`, which is initialized to store
`:η, :δ, :residual` by default.

# See also
[`truncated_conjugate_gradient_descent`](@ref), [`trust_regions`](@ref)
"""
mutable struct StopWhenTrustRegionIsExceeded <: StoppingCriterion
    reason::String
    storage::StoreOptionsAction
    StopWhenTrustRegionIsExceeded(a::StoreOptionsAction=StoreOptionsAction( (:η, :δ, :residual) )) = new("", a)
end
function (c::StopWhenTrustRegionIsExceeded)(p::P,o::O,i::Int) where {P <: HessianProblem, O <: TruncatedConjugateGradientOptions}
    if has_storage(c.storage,:δ) && has_storage(c.storage,:η) && has_storage(c.storage,:residual)
        η = get_storage(c.storage,:η)
        δ = get_storage(c.storage,:δ)
        residual = get_storage(c.storage,:residual)
        a1 = inner(p.M, o.x, o.useRand ? get_preconditioner(p, o.x, residual) : residual, residual)
        a2 = inner(p.M, o.x, δ, getHessian(p, o.x, δ))
        a3 = inner(p.M, o.x, η, get_preconditioner(p, o.x, δ))
        a4 = inner(p.M, o.x, δ, get_preconditioner(p, o.x, δ))
        norm = inner(p.M, o.x, η, η) - 2*( a1 / a2 ) * a3 + (a1 / a2)^2 * a4
        if norm >= o.Δ^2 && i >= 0
            c.reason = "Trust-region radius violation (‖η‖² = $norm >= $(o.Δ^2) = Δ²). \n"
            c.storage(p,o,i)
            return true
        end
    end
    c.storage(p,o,i)
    return false
end

@doc raw"""
    StopWhenCurvatureIsNegative <: StoppingCriterion

A functor for testing if the curvature of the model is negative, i.e.
$\langle \delta_k, \operatorname{Hess}[F](\delta_k)\rangle_x \leqq 0$.
In this case, the model is not strictly convex, and the stepsize as computed
does not give a reduction of the model.

# Fields
* `reason` – stores a reason of stopping if the stopping criterion has one be
    reached, see [`get_reason`](@ref).
* `storage` – stores the necessary parameter `δ` to check the
    criterion.

# Constructor

    StopWhenCurvatureIsNegative([a])

initialize the StopWhenCurvatureIsNegative functor to indicate to stop after
the inner product of the search direction and the hessian applied on the search
dircetion is less than zero using the [`StoreOptionsAction`](@ref) `a`, which
is initialized to just store `:δ` by default.

# See also
[`truncated_conjugate_gradient_descent`](@ref), [`trust_regions`](@ref)
"""
mutable struct StopWhenCurvatureIsNegative <: StoppingCriterion
    reason::String
    storage::StoreOptionsAction
    StopWhenCurvatureIsNegative(a::StoreOptionsAction=StoreOptionsAction( (:δ, ) )) = new("", a)
end
function (c::StopWhenCurvatureIsNegative)(p::P,o::O,i::Int) where {P <: HessianProblem, O <: TruncatedConjugateGradientOptions}
    if has_storage(c.storage,:δ)
        δ = get_storage(c.storage,:δ)
        if inner(p.M, o.x, δ, getHessian(p, o.x, δ)) <= 0 && i > 0
            c.reason = "Negative curvature. The model is not strictly convex (⟨δ,Hδ⟩_x = $(inner(p.M, o.x, δ, getHessian(p, o.x, δ))) <= 0).\n"
            c.storage(p,o,i)
            return true
        end
    end
    c.storage(p,o,i)
    return false
end


using Klara

covariates, = dataset("swiss", "measurements")
ndata, npars = size(covariates)

covariates = (covariates.-mean(covariates, 1))./repmat(std(covariates, 1), ndata, 1)

outcome, = dataset("swiss", "status")
outcome = vec(outcome)

function ploglikelihood(p::Vector{Float64}, v::Vector)
  Xp = v[2]*p
  dot(Xp, v[3])-sum(log(1+exp(Xp)))
end

plogprior(p::Vector{Float64}, v::Vector) = -0.5*(dot(p, p)/v[1]+length(p)*log(2*pi*v[1]))

p = BasicContMuvParameter(:p, loglikelihood=ploglikelihood, logprior=plogprior, nkeys=4)

model = likelihood_model([Hyperparameter(:λ), Data(:X), Data(:y), p], isindexed=false)

sampler = RAM(ones(npars))

mcrange = BasicMCRange(nsteps=10000, burnin=1000)

v0 = Dict(:λ=>100, :X=>covariates, :y=>outcome, :p=>[5.1, -0.9, 8.2, -4.5])

job = BasicMCJob(model, sampler, mcrange, v0)

run(job)

chain = output(job)

mean(chain)

acceptance(chain, diagnostics=false)


######
##FEM Heat Nonlinear Test
######
using FiniteElementDiffEq

#Define a parabolic problem
prob = prob_femheat_birthdeathinteractingsystem
#@code_warntype solve(prob,alg=:Euler)
sol = solve(prob,FEMDiffEqHeatEuler())

#plot(sol,plot_analytic=false,zlim=(0,2),cbar=false)

prob = prob_femheat_birthdeathsystem
sol = solve(prob,FEMDiffEqHeatImplicitEuler(),iterations=1000,autodiff=true)

#plot(sol,plot_analytic=false,zlim=(0,2),cbar=false)

true


# %% Import NaMaster for the power spectrum analysis on a partial sky
# Reference: Alonso et al.,  MNRAS, 484, 4127 (2019), https://github.com/LSSTDESC/NaMaster
nmt = pyimport("pymaster")
# Read in and apodize mask (C^2 apodization, 10 degrees)
m = readMapFromFITS(maskfile, 1, Float64)
mask = nmt.mask_apodization(m.pixels, 10.0, apotype = "C2")
# Create binning scheme. We will use Δℓ multipoles per bandpower.
b = nmt.NmtBin.from_nside_linear(nside, Δℓ, is_Dell = true)
# Array with effective multipole per bandpower
ell_eff = b.get_effective_ells()
nbands = b.get_n_bands()
# The function defined below will compute the power spectrum between two
# NmtFields f_a and f_b, using the coupling matrix stored in the
# NmtWorkspace wsp and subtracting the deprojection bias clb.
# Note that the most expensive operations in the MASTER algorithm are
# the computation of the coupling matrix and the deprojection bias. Since
# these two objects are precomputed, this function should be pretty fast!
function compute_master(f_a, f_b, wsp)
    # Compute the power spectrum (a la anafast) of the masked fields
    # Note that we only use n_iter=0 here to speed up the computation,
    # but the default value of 3 is recommended in general.
    cl_coupled = nmt.compute_coupled_cell(f_a, f_b)
    # Decouple power spectrum into bandpowers inverting the coupling matrix
    cl_decoupled = wsp.decouple_cell(cl_coupled)
    return cl_decoupled
end


function α_func(params, α, n)
    (; r, D, W, dim) = params

    m = length(D)
    bc1 = zeros(size(α))
    bc2 = zeros(size(α))

    for i = 1:length(α)
        bc = compute_bc(params, α[i], n)
        bc1[i] = bc[m, 1]
        bc2[i] = bc[m, 2]
    end

    rD = r[m] / √D[m]

    if dim == 2
        f1 = (
            √D[m] * α .* (besselj(n - 1, α * rD) - besselj(n + 1, α * rD)) / 2 +
            W[m] * besselj(n, α * rD)
        )
        f2 = (
            √D[m] * α .* (bessely(n - 1, α * rD) - bessely(n + 1, α * rD)) / 2 +
            W[m] * bessely(n, α * rD)
        )
    else
        f1 = (
            √D[m] * α .* (
                sphericalbesselj(n - 1, α * rD) - sphericalbesselj(n + 1, α * rD) -
                sphericalbesselj(n, α * rD) ./ (α * rD)
            ) / 2 + W[m] * sphericalbesselj(n, α * rD)
        )
        f2 = (
            √D[m] * α .* (
                sphericalbessely(n - 1, α * rD) - sphericalbessely(n + 1, α * rD) -
                sphericalbessely(n, α * rD) ./ (α * rD)
            ) / 2 + W[m] * sphericalbessely(n, α * rD)
        )
    end

    bc1 .* f1 + bc2 .* f2
end


using MDPTools, NestedFixedPoint
# Set RNG seed
srand(124)

# State space
n = 175
o_max = 450
Xval = linspace(0, o_max, n)
p = [0.0937; 0.4475; 0.4459; 0.0127; 0.0002];

F1 = zeros(n, n)
offset = -1
for i = 1:n, j = 1:length(p)
    i+offset+j > n && continue
    F1[i,i+offset+j] = i+offset+j == n ? sum(p[j:end]) : p[j]
end

# We can handle sparse matrices
sparse_F = [sparse(F1), sparse(F1)[ones(Int64, n), :]]
F = [F1, F1[ones(Int64, n), :]]

S = State(Xval, F)
sparse_S = State(Xval, sparse_F)
bdisc = 0.9999
Z1 = [zeros(n) -0.001*Xval]
Z2 = [-ones(n) zeros(n)]

U = LinearUtility([Z1, Z2], bdisc, [11.;2.5])
solve(U,S)
T = 120
N = 50

D = simulate(U, S, 1, T, N)

U.Θ*=0.
rust_nfxp = fit_nfxp(U, S, D)


# Unit to calculate solar geometry including solar elevation, declination,
#  azimuth etc using TSolar class. Data are hidden. 03/15/00 SK
# - 1st Revision 10/10/00: Changed to dealing only upto the top of the canopy. Radiation transter within the canopy is now a separate module.
# - added functins to calculate global radiation, atmospheric emissivity, etc as in Spitters et al. (1986), 3/18/01, SK
# 24Dec03, SK
# - added a function to calculate day length based on Campbell and Norman (1998) p 170,
# - added overloads to SetVal
# 2Aug04, SK
# - Translated to C++ from Delphi
# - revised some functions according to "An introduction to solar radiaiton" by Iqbal (1983)
# - (To Do) Add algorithms for instantaneous diffuse and direct radiation predictions from daily global solar radiation for given time
# - (To Do) This can be done by first applying sinusoidal model to the daily data to simulate hourly global solar radiation
# - (To Do) Then the division model of diffuse and direct radiations was applied
# - added direct and diffuse separation functions according to Weiss and Norman (1985), 3/16/05

@system Location begin
    lat: latitude => 36u"°" ~ preserve(u"°", parameter)
    long: longitude => 128u"°" ~ preserve(u"°", parameter)
    alt: altitude => 20u"m" ~ preserve(u"m", parameter)
end

import Dates

@system Sun begin
    #TODO override Location
    loc(context): location ~ ::Location #(override)
    calendar ~ ::Calendar(override)
    weather ~ ::Weather(override)

    # @derive time? -- takes account different Julian day conventions (03-01 vs. 01-01)
    t(calendar.time): datetime ~ track::datetime
    d(t): day => Dates.dayofyear(t) ~ track::int(u"d")
    h(t): hour => Dates.hour(t) ~ track::int(u"hr")

    ϕ(loc.lat): latitude ~ track(u"°") # DO NOT convert to radians for consistency
    λ(loc.long): longitude ~ track(u"°") # leave it as in degrees, used only once for solar noon calculation
    alt(loc.alt): altitude ~ track(u"m")

    solrad(weather.solrad): solar_radiation ~ track(u"W/m^2")
    τ: transmissivity => 0.5 ~ preserve(parameter) # atmospheric transmissivity, Goudriaan and van Laar (1994) p 30

    #####################
    # Solar Coordinates #
    #####################

    #HACK always use degrees for consistency and easy tracing
    #FIXME pascal version of LightEnv uses iqbal()
    δ(declination_angle_spencer): declination_angle ~ track(u"°")

    # Goudriaan 1977
    declination_angle_goudriaan(d) => begin
        g = 2pi * (d + 10u"d") / 365u"d"
        -23.45u"°" * cos(g)
    end ~ track(u"°")

    # Resenberg, blad, verma 1982
    declination_angle_resenberg(d) => begin
        g = 2pi * (d - 172u"d") / 365u"d"
        23.5u"°" * cos(g)
    end ~ track(u"°")

    # Iqbal (1983) Pg 10 Eqn 1.3.3, and sundesign.com
    declination_angle_iqbal(d) => begin
        g = 2pi * (d + 284u"d") / 365u"d"
        23.45u"°" * sin(g)
    end ~ track(u"°")

    # Campbell and Norman, p168
    declination_angle_campbell(d) => begin
        a = deg2rad(356.6 + 0.9856u"d^-1" * d)
        b = deg2rad(278.97 + 0.9856u"d^-1" * d + 1.9165sin(a))
        asind(0.39785sin(b))
    end ~ track(u"°")

    # Spencer equation, Iqbal (1983) Pg 7 Eqn 1.3.1. Most accurate among all
    declination_angle_spencer(d) => begin
        # gamma: day angle
        g = 2pi * (d - 1u"d") / 365u"d"
        0.006918 - 0.399912cos(g) + 0.070257sin(g) - 0.006758cos(2g) + 0.000907sin(2g) -0.002697cos(3g) + 0.00148sin(3g)
    end ~ track(u"rad")

    dph: degree_per_hour => 360u"°" / 24u"hr" ~ preserve(u"°/hr")

    # LC is longitude correction for Light noon, Wohlfart et al, 2000; Campbell & Norman 1998
    LC(λ, dph): longitude_correction => begin
        # standard meridian for pacific time zone is 120 W, Eastern Time zone : 75W
        # LC is positive if local meridian is east of standard meridian, i.e., 76E is east of 75E
        #standard_meridian = -120
        meridian = round(u"hr", λ / dph) * dph
        #FIXME use standard longitude sign convention
        #(long - meridian) / dph
        #HACK this assumes inverted longitude sign that MAIZSIM uses
        (meridian - λ) / dph
    end ~ track(u"hr")

    ET(d): equation_of_time => begin
        f = (279.575 + 0.9856u"d^-1" * d)*u"°"
        (-104.7sin(f) + 596.2sin(2f) + 4.3sin(3f) - 12.7sin(4f) -429.3cos(f) - 2.0cos(2f) + 19.3cos(3f)) / (60 * 60)
    end ~ track(u"hr")

    solar_noon(LC, ET) => 12u"hr" - LC - ET ~ track(u"hr")

    # θs: zenith angle
    hour_angle_at(ϕ, δ; θs(u"°")) => begin
        # this value should never become negative because -90 <= latitude <= 90 and -23.45 < decl < 23.45
        #HACK is this really needed for crop models?
        # preventing division by zero for N and S poles
        #denom = fmax(denom, 0.0001)
        # sunrise/sunset hour angle
        #TODO need to deal with lat_bound to prevent tan(90)?
        #lat_bound = radians(68)? radians(85)?
        # cos(h0) at cos(theta_s) = 0 (solar zenith angle = 90 deg == elevation angle = 0 deg)
        #-tan(ϕ) * tan(δ)
        c = (cos(θs) - sin(ϕ) * sin(δ)) / (cos(ϕ) * cos(δ))
        # c > 1: in the polar region during the winter, sun does not rise
        # c < -1: white nights during the summer in the polar region
        c = clamp(c, -1, 1)
        acosd(c)
    end ~ call(u"°")

    hour_angle_at_horizon(hour_angle_at) => hour_angle_at(90u"°") ~ track(u"°")

    # from Iqbal (1983) p 16
    half_day_length(hour_angle_at_horizon, dph) => (hour_angle_at_horizon / dph) ~ track(u"hr")
    day_length(half_day_length) => 2half_day_length ~ track(u"hr")

    sunrise(solar_noon, half_day_length) => (solar_noon - half_day_length) ~ track(u"hr")
    sunset(solar_noon, half_day_length) => (solar_noon + half_day_length) ~ track(u"hr")

    hour_angle(h, solar_noon, dph) => ((h - solar_noon) * dph) ~ track(u"°")

    αs(h=hour_angle, δ, ϕ): elevation_angle => begin
        #FIXME When time gets the same as solarnoon, this function fails. 3/11/01 ??
        asind(cos(h) * cos(δ) * cos(ϕ) + sin(δ) * sin(ϕ))
    end ~ track(u"°")

    ts(αs): positive_elevation_angle ~ track(u"°", min=0)

    θs(αs): zenith_angle => (90u"°" - αs) ~ track(u"°")

    # The solar azimuth angle is the angular distance between due South and the
    # projection of the line of sight to the sun on the ground.
    # View point from south, morning: +, afternoon: -
    # See An introduction to solar radiation by Iqbal (1983) p 15-16
    # Also see https://www.susdesign.com/sunangle/
    ϕs(αs, δ, ϕ): azimuth_angle => begin
        acosd((sin(δ) - sin(αs) * sin(ϕ)) / (cos(αs) * cos(ϕ)))
    end ~ track(u"°")

    ###################
    # Solar Radiation #
    ###################

    # atmospheric pressure in kPa (default = 100 kPa?)
    p(altitude): atmospheric_pressure => begin
        # campbell and Norman (1998), p 41
        101.3exp(-altitude / 8200u"m")
    end ~ track(u"kPa")

    m(p, ts): optical_air_mass_number => begin
        #FIXME check 101.3 is indeed in kPa
        #iszero(t_s) ? 0. : p / (101.3u"kPa" * sin(t_s))
        p / (101.3u"kPa" * sin(ts))
    end ~ track

    SC: solar_constant => 1370 ~ preserve(u"W/m^2", parameter)

    # Campbell and Norman's global solar radiation, this approach is used here
    insolation(ts, d, SC) => begin
        # solar constant, Iqbal (1983)
        #FIXME better to be 1361 or 1362 W/m-2?
        g = 2pi * (d - 10u"d") / 365u"d"
        SC * sin(ts) * (1 + 0.033cos(g))
    end ~ track(u"W/m^2")

    directional_solar_radiation(Fdir, solar_radiation) => begin
        Fdir * solar_radiation
    end ~ track(u"W/m^2")

    diffusive_solar_radiation(Fdif, solar_radiation) => begin
        Fdif * solar_radiation
    end ~ track(u"W/m^2")

    Fdir(τ, m): directional_coeff => begin
        # Goudriaan and van Laar's global solar radiation
        #FIXME should be goudriaan() version
        goudriaan(τ) = τ * (1 - diffusive_coeff)
        #FIXME: check if equation is same as campbell()
        # Takakura (1993), p 5.11
        takakura(τ, m) = τ^m
        # Campbell and Norman (1998), p 173
        campbell(τ, m) = τ^m
        campbell(τ, m)
    end ~ track

    # Fdif: Fraction of diffused light
    Fdif(τ, m): diffusive_coeff => begin
        # Goudriaan and van Laar's global solar radiation
        goudriaan(τ) = begin
            # clear sky : 20% diffuse
            if τ >= 0.7
                0.2
            # cloudy sky: 100% diffuse
            elseif τ <= 0.3
                1
            # inbetween
            else
                1.6 - 2τ
            end
        end
        # Takakura (1993), p 5.11
        takakura(τ, m) = (1 - τ^m) / (1 - 1.4log(τ)) / 2
        # Campbell and Norman (1998), p 173
        campbell(τ, m) = 0.3(1 - τ^m)
        campbell(τ, m)
    end ~ track

    directional_fraction(Fdif, Fdir) => (1 / (1 + Fdif/Fdir)) ~ track
    diffusive_fraction(Fdir, Fdif) => (1 / (1 + Fdir/Fdif)) ~ track

    # PARfr
    #TODO better naming: extinction? transmitted_fraction?
    PARfr(τ): photosynthetic_coeff => begin
        #if self.elevation_angle <= 0:
        #    0
        #TODO: implement Weiss and Norman (1985), 3/16/05
        weiss() = nothing
        # Goudriaan and van Laar (1994)
        goudriaan(τ) = begin
            # clear sky (τ >= 0.7): 45% is PAR
            if τ >= 0.7
                0.45
            # cloudy sky (τ <= 0.3): 55% is PAR
            elseif τ <= 0.3
                0.55
            else
                0.625 - 0.25τ
            end
        end
        goudriaan(τ)
    end ~ track

    Q: photosynthetic_active_radiation_conversion_factor => begin
        # 4.55 is a conversion factor from W to photons for solar radiation, Goudriaan and van Laar (1994)
        # some use 4.6 i.e., Amthor 1994, McCree 1981, Challa 1995.
        4.6
    end ~ preserve(u"μmol/J", parameter)

    # PARtot: total PAR (umol m-2 s-1) on horizontal surface (PFD)
    PARtot(solar_radiation, PARfr, Q): photosynthetic_active_radiation_total => begin
        # conversion factor from W/m2 to PFD (umol m-2 s-1) for PAR waveband (median 550 nm of 400-700 nm) of solar radiation,
        # see Campbell and Norman (1994) p 149
        solar_radiation * PARfr * Q
    end ~ track(u"μmol/m^2/s") # Quanta

    # PARdir
    PARdir(PARtot, directional_fraction): directional_photosynthetic_radiation => (directional_fraction * PARtot) ~ track(u"μmol/m^2/s") # Quanta

    # PARdif
    PARdif(PARtot, diffusive_fraction): diffusive_photosynthetic_radiation => (diffusive_fraction * PARtot) ~ track(u"μmol/m^2/s") # Quanta
end

@system SunController(Controller) begin
    calendar(context) ~ ::Calendar
    weather(context, calendar) ~ ::Weather
    s(context, calendar, weather): sun ~ ::Sun
end

plot_sun(v, d=3u"d"; kw...) = begin
    o = (
        :Calendar => (:init => ZonedDateTime(2007, 9, 1, tz"Asia/Seoul")),
        :Weather => (:store => loadwea(datapath("2007.wea"), tz"Asia/Seoul")),
    )
    r = simulate(SunController, stop=d, config=o, base=:sun)
    visualize(r, :time, v; kw...)
end

test_sun(d=3u"d") = foreach(v -> display(plot_sun(v, d)), [
    :δ, # declination_angle
    :αs, # elevation_angle
    :Fdir, # directional_coeff
    :Fdif, # diffusive_coeff
])


### construct a type of GSB mixture
type GSB
  mu::Vector{Float64}
  lambda::Vector{Float64}
  z::Float64
  w::Vector{Float64}
  W::Vector{Float64}
  Z::Vector{Float64}
end

### function for the generation of length(sample_sizes) data sets coming frmo mixture densities
### with means and stds given by means and sigmas. The mixture selection probabilities are specified in mixprobs
function toy_data(sample_sizes::Array{Int64, 1}, means::Array{Float64, 2}, sigmas::Array{Float64, 2}, mixprobs = Array{Float64, 2}, seed=1)
  srand(seed)
  m = length(sample_sizes)
  cumprobs = cumsum(mixprobs, dims=2)
  x = fill(Float64[], m)
  for j in 1:1:m
    x[j] = zeros(sample_sizes[j])
    for i in 1:1:sample_sizes[j]
      u = rand()
      for l in 1:1:m
        if u < cumprobs[j,l]
          x[j][i] = rand(Normal(means[j,l], sigmas[j,l]))
          break
        end
      end
    end
  end
  return x
end

### sample the clustering variables in density f_j
function sampleClustAllocs(R::Array{Any, 2}, y, pr::Array{Float64, 2}, S::Array{Int64, 1})
  n = length(y)
  m = length(pr)
  clusters = Array(Int64, n)
  allocations = zeros(Int64, n, m)

  for i in 1:1:n
    nc = 0.0
    for l in 1:1:m
      for j in 1:1:S[i]
        nc += pr[l] * R[l].lambda[j]^0.5 * exp(-0.5 * R[l].lambda[j] * (y[i] - R[l].mu[j])^2)
      end
    end
    prob = 0.0
    r = rand()
    flag = false
    for l in 1:1:m
      for j in 1:1:S[i]
        prob += pr[l] * R[l].lambda[j]^0.5 * exp(-0.5 * R[l].lambda[j] * (y[i] - R[l].mu[j])^2) / nc
        if r < prob
          clusters[i] = j
          allocations[i, l] = 1
          flag = true
          break
        end
      end
      if flag
        break
      end
    end
  end
  return clusters, allocations
end

### function for simulating truncated geometric random variables
function stgeornd(p::Float64, k::Int64)
  z::Int64
  z = floor(log(rand(Uniform(0.0,1.0))) / log(1 - p)) + k
  return z
end

### function for the sampling of geometric slice variables N
function sampleGeomSlice(R::Array{Any, 2}, clusters::Array{Int64, 1}, allocations::Array{Int64, 2})
  n = length(clusters)
  s = zeros(Int64, n)
  for i in 1:1:n
    for l in 1:1:size(allocations, 2)
      if allocations[i,l].==1
        s[i] = stgeornd(R[l].z, clusters[i])
        break
      end
    end
  end
  return s
end

### simulates from the Dirichlet posterior of the selection probabilities
function sampleSelectionProbabilities(a::Array{Float64, 2}, alloc::Array{Int64, 2})
  p = Array(Float64, size(alloc, 2))
  adir = zeros(1, size(alloc, 2))
  for i in 1:1:size(alloc, 1)
    for l in 1:1:size(alloc, 2)
      if alloc[i, l] == 1
        adir[l] += 1.0
      end
    end
  end
  p = rand(Dirichlet(vec(a + adir)))'
  return p
end

### function for the calculation of the number of unique clusters in group j
function uniqueClusters(clusters::Array{Int64, 1}, allocations::Array{Int64, 2})
  ucl = 0
  l = size(allocations, 2)
  for i in 1:1:l
    ucl += length(unique(clusters[allocations[:,i].==1]))
  end
  return ucl
end

### density estimation with sampling from the predictive
function samplePredictive(R::Array{Any, 2}, pr::Array{Float64, 2}, μ, σ, α, β)
  m = length(pr)
  Pr = cumsum(pr)
  xp = rand(Normal(rand(Normal(μ, σ)), sqrt(1 / rand(Gamma(α, 1 / β)))))
  for l in 1:1:m
    R[l].W = cumsum(R[l].w)
    rc = rand()
    rf = rand()
    if rc < Pr[l]
      for i in 1:1:length(R[l].W)
        if rf < R[l].W[i]
          xp = rand(Normal(R[l].mu[i], sqrt(1 / R[l].lambda[i])))
          break
        end
      end
      break
    end
  end
  return xp
end


module StandardizedMatrices
export StandardizedMatrix

#-----------------------------------------------------------------------------# types
typealias AVec{T} AbstractVector{T}
typealias AMat{T} AbstractMatrix{T}
typealias VecF Vector{Float64}

#----------------------------------------------------------------# StandardizedMatrix
"""
Treat a matrix as standardized, similar to `z = StatsBase.zscore(x, 1)`, without
altering the original data.

`z = StandardizedMatrix(x::AbstractMatrix)`
`z = StandardizedMatrix(x::AbstractMatrix, μ::AbstractVector, σ::AbstractVector)`
"""
immutable StandardizedMatrix{T <: AbstractFloat, S <: AMat} <: AMat{T}
	data::S
	μ::VecF			# column means
	σinv::VecF		# inverse of column stdevs
end
# This acts as inner constructor.  All constructors should call this.  See:
# http://docs.julialang.org/en/latest/manual/constructors/#parametric-constructors
function StandardizedMatrix(x::AMat, μ::AVec, σ::AVec)
	T = eltype((x[1] - μ[1]) / σ[1])
	n, p = size(x)
	μp = length(μ)
	σp = length(σ)
	@assert p == μp "size(x, 2) == $p doesn't match length(μ) == $μp"
	@assert p == σp "size(x, 2) == $p doesn't match length(σ) == $σp"
	StandardizedMatrix{T, typeof(x)}(x, μ, one(eltype(σ)) ./ σ)
end
StandardizedMatrix(x::AMat) = StandardizedMatrix(x, vec(mean(x, 1)), vec(std(x, 1)))

#----------------------------------------------------------------------# Base methods
# http://docs.julialang.org/en/release-0.4/manual/interfaces/#man-interfaces-abstractarray
Base.size(o::StandardizedMatrix) 				= size(o.data)
Base.linearindexing(o::StandardizedMatrix)		= Base.linearindexing(o.data)
function Base.getindex(o::StandardizedMatrix, i::Int, j::Int)
	v = getindex(o.data, i, j)
	return (v - o.μ[j]) * o.σinv[j]
end
function Base.getindex(o::StandardizedMatrix, i::Int)
	v = getindex(o.data, i)
	j = ceil(Int, i / size(o, 1))
	return (v - o.μ[j]) * o.σinv[j]
end
function Base.:*{T <: Real}(A::StandardizedMatrix, B::AVec{T})
	y = zeros(typeof(A[1] * B[1]), size(A, 1))
	A_mul_B!(y, A, B)
	y
end
function Base.:*{T <: Real}(A::StandardizedMatrix, B::AMat{T})
	y = zeros(typeof(A[1] * B[1]), size(A, 1), size(B, 2))
	A_mul_B!(y, A, B)
	y
end

#------------------------------------------------------# Matrix-Vector multiplication
function Base.A_mul_B!(y::AVec, A::StandardizedMatrix, b::AVec)
	A_mul_B!(y, A.data, Diagonal(A.σinv) * b)
	center!(y)
end
function Base.At_mul_B!(y::AVec, A::StandardizedMatrix, b::AVec)
	At_mul_B!(y, A.data, b - mean(b))
	_scale!(y, A.σinv)
end

#------------------------------------------------------# Matrix-Matrix multiplication
function Base.A_mul_B!(y::AMat, A::StandardizedMatrix, b::AMat)
	A_mul_B!(y, A.data, Diagonal(A.σinv) * b)
	center!(y)
end
function Base.At_mul_B!(y::AMat, A::StandardizedMatrix, b::AMat)
	At_mul_B!(y, A.data, b .- mean(b, 1))
	_scale!(y, A.σinv)
end

#---------------------------------------------------------------------------# helpers
function center!(v::Vector)
	μ = mean(v)
	for i in eachindex(v)
		@inbounds v[i] -= μ
	end
	v
end
function center!(m::Matrix)
	μ = mean(m, 1)
	@inbounds for j in 1:size(m, 2)
		μj = μ[j]
		for i in 1:size(m, 1)
			m[i, j] -= μj
		end
	end
	m
end
function _scale!(a::Vector, b::Vector)
	for i in eachindex(a)
		@inbounds a[i] *= b[i]
	end
	a
end
function _scale!(a::Matrix, b::Vector)
	for j in 1:size(a, 2)
		for i in 1:size(a, 1)
			a[i, j] *= b[i]
		end
	end
	a
end

end #module


module ConvAnalysisModule

using LinearAlgebra

using ..ObjectAgentModule

export ConvAnalysis_Data, errf, errf_sum, conv_measure

struct AgentData{T<:Real}
    cost::T;
    xk::Vector{T};
    zk::Vector{T};
    uk::Vector{T};
end

mutable struct ConvAnalysis_Data
    expData::Array{Array{AgentData,1}, 1}

    function ConvAnalysis_Data()
        new([]);
    end
end

function Base.getindex(A::ConvAnalysis_Data, elements...)
    return Base.getindex(A.expData, elements...);
end

function Base.lastindex(A::ConvAnalysis_Data)
    return Base.lastindex(A.expData)
end

function Base.push!(cad::ConvAnalysis_Data, agents::Array{OA, 1}) where {OA<:ObjectAgent}
    exp_iter = AgentData[];
    for agent in agents
        push!(exp_iter, agent);
    end
    push!(cad.expData, exp_iter);
end

function Base.push!(aga::Array{AgentData, 1}, agent::ObjectAgent)
    push!(aga, AgentData(costFun(agent, agent.zk), copy(agent.xk), copy(agent.zk), copy(agent.uk)));
end

function numOfAgents(cad::ConvAnalysis_Data)
    if !isempty(cad.expData)
        return size(cad.expData[1])[1]
    else
        return 0;
    end
end

function numOfIterations(cad::ConvAnalysis_Data)
    if !isempty(cad.expData)
        return size(cad.expData)[1]
    else
        return 0;
    end
end

function errf(cad::ConvAnalysis_Data)
    return [[cad[i][k].cost for i in 1:numOfIterations(cad)] for k in 1:numOfAgents(cad)]
end

function errf_sum(cad::ConvAnalysis_Data)
    """Returns average cost (F_des - f_model(x))ˆ2"""
    return [sum([cad[i][k].cost for k in 1:numOfAgents(cad)]) / numOfAgents(cad) for i in 1:numOfIterations(cad)]
end

function conv_measure(cad::ConvAnalysis_Data)
    return [ sqrt(sum( [ norm(cad[i][k].zk - cad[i-1][k].zk)^2 + norm(cad[i][k].uk - cad[i-1][k].uk)^2 for k in 1:numOfAgents(cad)] ) / numOfAgents(cad)) for i in 2:numOfIterations(cad)]
end

end


nz = 3 + 2 + 1 + 1 + 3 + 3 + 3 + 2
nθ = 3 + 1 + 1 + 1

idx = [3; 1; 2]
function residual(z, θ, κ)
    m = 3

    f = z[1:m]
    s = z[m .+ (1:2)]
    y1 = z[m + 2 .+ (1:1)]
    y2 = z[m + 2 + 1 .+ (1:1)]
    y3 = z[m + 2 + 1 + 1 .+ (1:3)]
    β = z[m + 2 + 1 + 1 + 3 .+ (1:3)]
    η = z[m + 2 + 1 + 1 + 3 + 3 .+ (1:3)]
    z = z[m + 2 + 1 + 1 + 3 + 3 + 3 .+ (1:2)]

    f̄ = θ[1:m]
    f_min = θ[m .+ (1:1)]
    f_max = θ[m + 1 .+ (1:1)]
    μ = θ[m + 1 + 1 .+ (1:1)]

    A = [0.0 0.0 μ; 1.0 0.0 0.0; 0.0 1.0 0.0]

    return [f - f̄ + [0.0; 0.0; -y1[1] + y2[1]] + transpose(A) * y3;
            [-y1[1]; -y2[1]] - z;
            f_max[1] - f[3] - s[1];
            f[3] - f_min[1] - s[2];
            A * f - β;
            -y3 - η;
            second_order_cone_product(β, η) - [κ; 0.0; 0.0];
            s .* z .- κ]
end

@variables z[1:nz], θ[1:nθ], κ[1:1]

r = residual(z, θ, κ)
r .= simplify(r)
r_func = eval(Symbolics.build_function(r, z, θ, κ)[2])

rz = Symbolics.jacobian(r, z)
rz = simplify.(rz)
rz_func = eval(Symbolics.build_function(rz, z, θ)[2])

rθ = Symbolics.jacobian(r, θ)
rθ = simplify.(rθ)
rθ_func = eval(Symbolics.build_function(rθ, z, θ)[2])

idx_ineq = collect([4, 5, 17, 18])
idx_soc = [collect([11, 12, 13]), collect([14, 15, 16])]

# ul = -1.0 * ones(m)
# uu = 1.0 * ones(m)
ū = [0.0; 0.0; 0.0]
u_min = 0.0
u_max = 1.0
scaling = 1.0
z0 = [ū; 1.0; 1.0; 0.0; 0.0; 0.0; 0.0; 0.0; 1.0; 0.1; 0.1; 1.0; 0.1; 0.1; 1.0; 1.0]

θ0 = [copy(ū); u_min; u_max; scaling]

# solver
opts_con = InteriorPointOptions(
    κ_init = 1.0,
    κ_tol = 1.0e-4,
    r_tol = 1.0e-8,
    diff_sol = false)

ip_con = interior_point(z0, θ0,
    r! = r_func, rz! = rz_func,
    rz = similar(rz, Float64),
    rθ! = rθ_func,
    rθ = similar(rθ, Float64),
    idx_ineq = idx_ineq,
    idx_soc = idx_soc,
    opts = opts_con)

opts_jac = InteriorPointOptions(
	κ_init = 1.0,
	κ_tol = 1.0e-3,
	r_tol = 1.0e-8,
	diff_sol = true)

ip_jac = interior_point(z0, θ0,
	r! = r_func, rz! = rz_func,
	rz = similar(rz, Float64),
	rθ! = rθ_func,
	rθ = similar(rθ, Float64),
	idx_ineq = idx_ineq,
  idx_soc = idx_soc,
	opts = opts_jac)

# interior_point_solve!(ip_con)
# interior_point_solve!(ip_jac)

function soc_projection(x, x_min, x_max, scaling)
    ip_con.z .= [x; 1.0; 1.0; 0.0; 0.0; 0.0; 0.0; 0.0; 1.0; 0.1; 0.1; 1.0; 0.1; 0.1; 1.0; 1.0]
    ip_con.θ .= [x; x_min; x_max; scaling]

    status = interior_point_solve!(ip_con)

    !status && (@warn "projection failure (res norm: $(norm(ip_con.r, Inf))) \n
    	               z = $(ip_con.z), \n
    				   θ = $(ip_con.θ)")

	return @view ip_con.z[1:3]
end

# soc_projection([100.0, 0.0, -1.0], 0.0, 1.0, 1.0)

function soc_projection_jacobian(x, x_min, x_max, scaling)
    ip_con.z .= [x; 1.0; 1.0; 0.0; 0.0; 0.0; 0.0; 0.0; 1.0; 0.1; 0.1; 1.0; 0.1; 0.1; 1.0; 1.0]
    ip_con.θ .= [x; x_min; x_max; scaling]

	interior_point_solve!(ip_jac)

	return view(ip_jac.δz, 1:3, 1:3)
end

# soc_projection_jacobian([0.0, 0.0, 1.0], 0.0, 10.0, 1.0)


struct Heisenberg{D} <: AbstractHeisenberg{D}
    size::NTuple{D, Int}
    periodic::Bool
    Heisenberg(size::Int...; periodic::Bool) = new{length(size)}(size, periodic)
end

Base.size(model::Heisenberg) = model.size

heisenberg_ij(nbit::Int, i::Int, j::Int=i+1) = put(nbit, i=>X)*put(nbit, j=>X) + put(nbit, i=>Y)*put(nbit, j=>Y) + put(nbit, i=>Z)*put(nbit, j=>Z)
const heisenberg_term = repeat(2, X, 1:2) + repeat(2, Y, 1:2) + repeat(2, Z, 1:2)

function hamiltonian(model::Heisenberg)
    nbit = nspin(model)
    sum(x->heisenberg_ij(nbit, x[1], x[2]), get_bonds(model))*0.25
end

function get_bonds(model::Heisenberg{2})
    m, n = model.size
    cis = LinearIndices(model.size)
    bonds = Tuple{Int, Int, Float64}[]
    for i=1:m, j=1:n
        (i!=m || model.periodic) && push!(bonds, (cis[i,j], cis[i%m+1,j], 1.0))
        (j!=n || model.periodic) && push!(bonds, (cis[i,j], cis[i,j%n+1], 1.0))
    end
    bonds
end

function get_bonds(model::Heisenberg{1})
    nbit, = model.size
    [(i, i%nbit+1, 1.0) for i in 1:(model.periodic ? nbit : nbit-1)]
end


using DifferentialEquations, Plots
import StatsBase, JSON, Random

"""
This file contains the definition of the Adaptive Filter model as well as
the Lazar model as an ODE to be solved by the DifferentialEquations library.
"""

#=== meta execution ========#
exec = false            # if false, no solving or plotting from this file
algo = :naf             # options: :af, :naf, :lazar
task = "relu"           # options: recovery:     "sin", "mix2", "mix3",
                        # linear computation:    "1of3", "shift", "noise"
                        # nonlinear computation: "abs", "relu", "sign"
mode = "train"          # options: "train", "test" (disables learning rule)
show = "mse_evol"       # options: "out", "y", "y*w", "zyw", "imp", "impi",
                        # "w_evol", "mse_evol" (measured during training)
mark = ""               # any mark to put in the figure filename

#=== input/output setup ====#
lp = 3                  # lowest = shortest period (fastest frequency)
mp = 4                  # middle period (for mix3 and 1of3 task)
hp = 5                  # highest = longest period (slowest frequency)
phase = lp/4            # phase for task="shift" (additional to latency)
offset = 0              # y-offset of signal (important for nonlinearity)
latency = -.60          # accounts for latency of information due to spikes
noise_std = .5          # std for gaussian white noise

#=== encoder parameters ====#
rate = 10               # the resting firing rate (spikes/sec)
b = 3.0                 # bias (> |minimal input|, s.t. x(t) + b > 0)
θ = 1.0                 # threshold (when reached, we spike)
κ = (b+offset)/rate/θ   # capacitance (integrator weight)
Δ = 0.0                 # TODO refractory?
l = 0.0                 # TODO leak?

#=== decoder parameters ====#
n = 10                  # n filters in total (equal for both models)
η = 1e-4                # learning rate      (equal for both models)
α = 2                   # fastest decay rate
λ = 2                   # fastest decay rate
λs = range(1, length=n, stop=0.01)

ramp(z) = log(1 + ℯ^z)  # non-linearity (ramp)
sigmoid(z) = ℯ^z / (1 + ℯ^z) # derivative of ramp
sig_der(z) = ℯ^(-z) / (1 + ℯ^(-z))^2 # derivative of sigmoid

#=== simulation ============#
training_cycles = 10000  # how many total periods to see in training
test_cycles = 3          # and how many in the test phase

function init_task(task)
    """
    Returns an io tuple (of input x and target z) and the target period.
    The target period is just the lowest period (lp) for a sine-based task
    and the least common multiple of the periods in a mix task.
    """
    Random.Random.seed!(1) # rng seed, for model comparability

    io = () # a tuple of functions x (input) and z (target)

    if task == "sin"       # reconstruct sine
        let input
            input(t) = sin(2π/lp*t) + offset
            io = (x=input, z=t->input(t+latency))
        end
        tp = lp
    elseif task == "mix2"  # reconstruct a mix of 2 sines
        let input
            input(t) = sin(2π/hp*t) + sin(2π/lp*t) + offset
            io = (x=input,  z=t->input(t+latency))
        end
        tp = lcm([lp, hp]) # total period i
    elseif task == "mix3"  # reconstruct a mix of 3 sines
        let input
            input(t) = sin(2π/hp*t) + sin(2π/mp*t) + sin(2π/lp*t) + offset
            io = (x=input,  z=t->input(t+latency))
        end
        tp = lcm([lp, mp, hp]) # total period is lcm of periods
    elseif task == "1of3"  # demixing: extract 1 component from a mix3
        let input, target
            input(t) = sin(2π/hp*t) + sin(2π/mp*t) + sin(2π/lp*t) + offset
            target(t) = sin(2π/mp*t) + offset
            io = (x=input, z=t->target(t+latency))
        end
        tp = mp
    elseif task == "shift" # shift sine by phase
        let input
            input(t) = sin(2π/lp*t) + offset
            io = (x=input, z=t->input(t+latency+phase))
        end
        tp = lp
    elseif task == "noise" # denoise: reconstruct sine with white noise
        let input
            input(t) = sin(2π/lp*t) + offset
            noise = randn(1000) .* noise_std # noise with µ=0
            io = (x=t->input(t) + noise[Int(floor(t*100))%1000+1],
                  z=t->input(t+latency))
        end
        tp = lp
    elseif task == "abs"   # take absolute value
        let input
            input(t) = sin(2π/lp*t) + offset
            io = (x=input, z=t->abs(input(t+latency)-offset)+offset)
        end
        tp = lp
    elseif task == "relu"  # take maximum of 0 and signal
        let input
            input(t) = sin(2π/lp*t) + offset
            io = (x=input, z=t->Base.max(0., input(t+latency)-offset)+offset)
        end
        tp = lp
    elseif task == "sqr"   # square of input
        let input
            input(t) = sin(2π/lp*t) + offset
            io = (x=input, z=t->(input(t+latency)-offset)^2+offset)
        end
        tp = lp
    elseif task == "sign"  # binary transformation: sign of input
        let input
            input(t) = sin(2π/lp*t) + offset
            io = (x=input, z=t->sign(input(t+latency)-offset)+offset)
        end
        tp = lp
    end

    global tp = tp
    return io, tp
end

function init_times(cycles, tp)
    """
    `Cycles` is how many total periods (`tp`) of the signal are presented.
    According to the `show` variable, the saving times (`tsave`) are set.
    Returns the list of save points `tsave`.
    """
    tmax = tp * cycles

    if show in ["w_evol", "mse_evol"]
        tsave = 0.:tp:tmax    # save regularly (once / period)
    else
        tsave = [0.,tmax]     # save only first and last state
    end

    return tsave
end

function init_state(; range=.05)
    """
    Initializes taps and weights randomly from a normal disrtibution
    centered around zero with standard deviation of `range`.
    Note that there are n+1 weights, the last one is the bias weight.
    Returns the initial state being in the form [ys; ws; 0].
    """
    Random.Random.seed!(1)         # rng seed, for comparability
    w0 = (.5 .- rand(n+1))*2*range # random initial weights
    y0 = (.5 .- rand(n))*2*range   # initial taps
    u0 = [y0; w0; 0.0]             # initial state: ys, ws, X
    return u0
end

function init_training(task, cycles)
    """
    Wrapper function around `init_task`, `init_times` and `init_state`.
    """
    # some global settings
    global mode = "train"
    global λs = range(λ, length=n, stop=0.01) # for the lazar decoder

    io, tp = init_task(task)
    u0 = init_state()
    tsave = init_times(cycles, tp)
    return io, tp, u0, tsave
end

#=== ODE DEFINITIONS =======#

function af(du, u, io, t)
    """
    System of ODE to model communication between two AF neurons,
    one reduced on the encoding, the other reduced on the decoding part.

    input signal → enc → channel (spikes) → dec → reconstructed signal

    The state has the following form: (y_1, …, y_n, w_1, …, w_n, w_b, X).
    Note that the spiking part (X) belongs to the presynaptic
    and the filter bank (y_1, …, w_n, w_b) to the postsynaptic neuron.
    """

    # decoder (postsynaptic)
    banklen = n ÷ 2                  # two filter banks of the same size

    # filter bank 1
    du[1] = -α * u[1]                # no input here (generated by events)
    du[2:banklen] .= -α * u[2:banklen] .+ α * u[1:banklen-1]

    # decoder output ("potential")
    z = u[1:n]' * u[n+1:2n] + u[2n+1]

    # filter bank 2
    du[banklen+1] = -α * u[banklen+1] + α * z # input from previous output z
    du[banklen+2:n] .= -α * u[banklen+2:n] .+ α * u[banklen+1:n-1]

    # weight updates
    if mode == "train"
        du[n+1:2n] = 2 * η * (io.z(t) - z) .* u[1:n]
        du[2n+1]   = 2 * η * (io.z(t) - z) # bias weight
    end

    # encoder (presynaptic)
    du[end] = (io.x(t) + b) / κ    # rectangular integration of signal
end

function naf(du, u, io, t)
    """
    Nonlinear AF, same as AF but with ramp function.

    System of ODE to model communication between two AF neurons,
    one reduced on the encoding, the other reduced on the decoding part.

    input signal → enc → channel (spikes) → dec → reconstructed signal

    The state has the following form: (y_1, …, y_n, w_1, …, w_n, w_b, X).
    Note that the spiking part (X) belongs to the presynaptic
    and the filter bank (y_1, …, w_n, w_b) to the postsynaptic neuron.
    """

    # decoder (postsynaptic)
    banklen = n ÷ 2                  # two filter banks of the same size

    # filter bank 1
    du[1] = -α * u[1]                # no input here (generated by events)
    du[2:banklen] .= -α * u[2:banklen] .+ α * u[1:banklen-1]

    # decoder output ("potential")
    z = u[1:n]' * u[n+1:2n] + u[2n+1]

    # filter bank 2
    du[banklen+1] = -α * u[banklen+1] + α*z # input from previous output z
    du[banklen+2:n] .= -α * u[banklen+2:n] .+ α * u[banklen+1:n-1]

    # weight updates
    if mode == "train"
        w_b = u[end-1] # bias weight for shift of ramp
        du[n+1:2n] = (2 * η * (io.z(t) - ramp(z+w_b)) .* u[1:n]
                      * sigmoid(z+w_b))
        # update for w_b (multiplied with n for stronger learning)
        du[end-1]  = 2 * n * η * (io.z(t) - ramp(z+w_b)) * sigmoid(z+w_b)
    end

    # encoder (presynaptic)
    du[end] = (io.x(t) + b) / κ    # rectangular integration of signal
end

function lazar(du, u, io, t)
    """
    Implementation of [Lazar 2006: A simple model for spike processing].
    A zeroth-order approximation of the stimulus is computed
    by spike weighting and filtering the signal.

    A decoding filter bank is constructed from n parallel filters
    each of which gets the same input from the presynaptic neuron.

    Filters have different time constants λ_i stored in λs.
    The output is the weighted sum of the filters, ∑_i w_i · y_i.

    The state has the following form: (y_1, …, y_n, w_1, …, w_n, w_b, X).
    Note that the spiking part (X) belongs to the presynaptic
    and the filter bank (y_1, …, w_n, w_b) to the postsynaptic neuron.
    """

    # decoder (postsynaptic)
    du[1:n] .= -λs .* u[1:n]

    # output
    z = u[1:n]' * u[n+1:2n] + u[2n+1]

    # weight updates
    if mode == "train"
        du[n+1:2n] .= 2 * η * (io.z(t) - z) .* u[1:n]
        du[2n+1]    = 2 * η * (io.z(t) - z) # bias weight
    end

    # encoder (presynaptic)
    du[end] = (io.x(t) + b) / κ # rectangular integration of signal
end

#=== ACCESSOR UTILS ==========#

function z(sol, t)
    u = sol(t)
    y = u[1:n]
    w = u[n+1:2n]
    w_b = u[2n+1] # bias weight
    w' * y + w_b
end

function r(sol, t)
    u = sol(t)
    ramp(z(sol, t))
end

function out(sol, t)
    if algo == :naf
        r(sol, t)
    else
        z(sol, t)
    end
end

function y_i(sol, i)
    function y(t)
        sol(t)[i]
    end
end

function w_i(sol, i)
    function w(t)
        sol(t)[n+i]
    end
end

function mse(sol, io, period)
    """
    Calculate mean squared error (normalized by signal amplitude)
    of the solution output w.r.t. the target function
    over the given period (array of time points)
    and normalize it by the std of the target amplitude.
    """
    mse = sum([(out(sol,t) - io.z(t))^2 for t in period]) /length(period)

    # compute std of target (on 1000 points over tp) to normalize mse
    var = StatsBase.var([io.z(t) for t in period])
    error = mse / var
    round(error, digits=5) # round on 5 decimal digits
end

#=== SOLVER FUNCTIONS ======#

function run(io, u0, tsave, save_spikes=false)
    """
    Runs the model `algo` in `mode` (train or test), i.e.
    it solves the specified ODE over the time specified in `tsave`.
    Returns the solution of the solver and if appropriate
    (`save_spikes == true`), also a list of all spike times.
    """
    tspan = (tsave[1], tsave[end])       # run over the whole save period
    tspike = undef                       # time of last spike, no initial
    spikes = []                          # list to save all spike times
    function threshold(u, t, integrator) # callback condition
        return u[end] - θ
    end
    function spike!(integrator)          # callback handler
        if algo == :lazar
            if tspike != undef           # ignore first spike (weight issue)
                isi = integrator.t - tspike  # inter-spike interval
                c_k = κ * θ - b * isi        # weighting for spikes
                integrator.u[1:n] .+= c_k    # add spike upon each tap
            end
            tspike = integrator.t        # save time of spike
        else                             # for AF/nAF, just do:
            integrator.u[1]  += 1.0      # input to postsynaptic neuron
        end
        integrator.u[end] = 0.0          # reset X integrator
        if save_spikes
            push!(spikes, integrator.t)  # save spike time
        end
    end
    fire = ContinuousCallback(threshold, spike!, # not to save u at spikes:
                              save_positions=(false, false))
    prob = ODEProblem(eval(algo), u0, tspan, io, callback=fire)
    sol = solve(prob, saveat=tsave, maxiters=Inf)
    if save_spikes
        return (sol, spikes) # solution and all spikes
    end
    return sol
end

function test_error(sol, io, tp)
    """
    Tests the last state of a solution `sol`,
    saves for `100/tp * test_cycles` times and
    returns the new solution object and the error.
    """
    global mode = "test"
    u_end, t_end = sol[end], sol.t[end]
    tsave2 = t_end:.01:t_end+test_cycles*tp
    sol2   = run(io, u_end, tsave2)
    return sol2, mse(sol2, io, tsave2)
end

function impulse_response(sol; dur=10)
    """
    Compute impulse response of a trained model (`sol`).
    The model is fully specified by its weight vector.
    Returns a solution for the i.r. over the duration `dur`.
    """
    w_end = sol[end][n+1:2n+1]
    u = zeros(2n+2)       # make new fresh state
    u[n+1:2n+1] = w_end   # take learned model
    if algo == :lazar
        u[1:n] .= 1       # put initial spike into all taps
    else
        u[1] = 1          # initial spike y_1 = 1
    end
    global mode = "test"  # do not adapt weights in rerun
    io = (x=t->0, )       # use no input and no z, because mode == "test"
    # Note that no encoder is necessary, therefore no callback.
    prob = ODEProblem(eval(algo), u, (0.,dur), io)
    solve(prob, saveat=0:.01:dur, maxiters=Inf)
end

function running_mse(sol, io, save_periods=[])
    """
    Takes a solution object and tests the model for each stored state.
    `save_periods` is a list of periods for which the solution is saved.
    Returns a list of MSEs and a list of the saved solutions, if any.
    """
    mses   = []
    sols   = []
    period = 0
    global mode = "test"
    for (u,t) in zip(sol.u, sol.t)
        tsave2 = t:tp/100:t+test_cycles*tp
        sol2 = run(io, u, tsave2)
        push!(mses, mse(sol2, io, tsave2))
        if period in save_periods
            push!(sols, sol2)
        end
        period += 1
    end
    return mses, sols
end

#=== BASIC PLOTS =============#

function palgo()
    """ Prettifies the algo name for plots. """
    if algo == :af
        "AF"
    elseif algo == :lazar
        "Lazar"
    elseif algo == :naf
        "nAF"
    end
end

function evol_plot(show, sol, io)
    """ Returns a plot. """
    if !(show in ["w_evol", "y_evol", "mse_evol"]) return end
    if show == "w_evol"
        plot([w_i(sol, i) for i in 1:n+1], sol.t,
             labels=["w_$i" for i in 1:n+1],
             xlabel="time",
             title="$(palgo()), weights", titleloc=:left)
    elseif show == "y_evol"
        plot([w_i(sol, i) for i in 1:n+1], sol.t,
             labels=["y_$i" for i in 1:n+1],
             xlabel="time",
             title="$(palgo()), taps")
    elseif show == "mse_evol"
        mses, _ = running_mse(sol, io)
        plot(mses,
             leg=false, title="$algo (final MSE: $(mses[end]))",
             ylabel="mse", xlabel="periods")
    end
end

function end_plot(show, sol, io, tp; return_mse=false)
    """ Returns a plot and, if `return_mse==true` also the error. """
    if !(show in ["out", "zyw", "y*w", "y"]) return end

    sol2, error = test_error(sol, io, tp)

    plt = nothing
    if show == "out"
        plt = plot([t->out(sol2, t), io.z], sol2.t,
                   labels=["approx", "target"],
                   title="$algo (MSE: $error)")
    elseif show == "y"
        plt = plot([y_i(sol2, i) for i in 1:n], sol2.t,
                   labels=["y_$i" for i in 1:n], leg=:bottomright,
                   xticks=sol2.t[1]:tp:sol2.t[end], xlabel="time",
                   title="$(palgo()), taps", titleloc=:left)
    elseif show == "y*w"
        plt = plot([t->(y_i(sol2, i)(t) * w_i(sol2, i)(t)) for i in 1:n],
                   sol2.t,
                   labels=["y_$i" for i in 1:n], leg=:bottomright,
                   xticks=sol2.t[1]:tp:sol2.t[end], xlabel="time",
                   title="$(palgo()), taps")
    elseif show == "zyw"
        plt = plot([t->z(sol2, t), y_i(sol2, 1), w_i(sol2, 1),
                    io.z,          y_i(sol2, n), w_i(sol2, n)],
                   sol2.t,
                   label=["approx" "y_1" "w_1" "target" "y_n" "w_n"],
                   title=["$algo (MSE: $error)" "taps" "weights"],
                   layout=(3,1))
    end

    if return_mse
        return plot, error
    end
    return plt
end

function imp_plot(show, sol)
    """ Returns a plot. """
    if !(show in ["imp", "impi"]) return end
    sol_imp = impulse_response(sol)

    if show == "imp"
        plot(t->z(sol_imp, t),
             sol_imp.t,
             label="total imp", xlabel="t",
             title="$algo, impulse response")
    elseif show == "impi"
        plot([y_i(sol_imp, i) for i in 1:n],
             sol_imp.t,
             labels=["y$i" for i in 1:n], xlabel="t",
             title="$algo, individual impulse response")
    end
end

#=== RUNNING + SAVING ======#

function serialize(u_end, filename, error)
    dict = Dict("algo"    => algo,
                "task"    => task,
                "n"       => n,
                "eta"     => η,
                "alpha"   => α,
                "latency" => latency,
                "mse"     => error,
                "u"       => u_end
                )
    string = JSON.json(dict)
    open(filename, "w") do f
        write(f, string)
    end
end

if exec
    # init
    io, tp, u0, tsave = init_training(task, training_cycles)

    # train
    sol = run(io, u0, tsave)

    # plot
    evol_plot(show, sol, io)
    _, error = end_plot(show, sol, io, tp, return_mse=true)
    imp_plot(show, sol)
    println("MSE: ", error)

    # save figure
    fn = "$algo-$task-$show-$mark-alpha$α-latency$latency-eta$η-offset$offset-t$(tsave[end])"
    dir = task
    savefig("../figures/$dir/$fn.png")

    # serialize the model
    #if mode == "train"
    #    serialize(sol1, "../models/$fn", error)
    #end
end


using PtFEM

ProjDir = dirname(@__FILE__)

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

data |> display
println()

res_df, g_coord, g_num, disp = p63_skyline(data)
println()

res_df


import LinearAlgebra.tr
export tr

function tr(e::AbstractExpr)
    return sum(diag(e))
end


sorted_unique_count(m) = isempty(m) ? 0 : 1+sum(m[2:end].≠m[1:end-1])

inverse_power_intensity_function(α::Real) = m -> 1/sorted_unique_count(m)^α
all_or_nothing_intensity_function(α::Real, β::Real) = m -> all(i -> i == first(m), m) ? α : β


module HDiscontinuousGalerkin

using Tensors
using FastGaussQuadrature
using SparseArrays
using LinearAlgebra

import Base:@propagate_inbounds

abstract type AbstractRefShape end
abstract type AbstractQuadratureRule end

abstract type Interpolation{dim,shape,order} end
abstract type DiscreteFunctionSpace{dim,T,FE} end

include("utils.jl")
include("shapes.jl")
include("mesh.jl")
include("generate_mesh.jl")
include("triangle_mesh.jl")
include("basis.jl")
include("FiniteElement.jl")
include("LagrangeFE.jl")
include("quadrature.jl")
include("GrundmannMoellerQuad.jl")
include("StrangQuad.jl")
include("ScalarFunctionSpaces.jl")
include("VectorFunctionSpaces.jl")
include("TraceFunctionSpaces.jl")
#include("ParametricFunctionSpaces.jl")
include("assembler.jl")
include("DiscreteFunctions.jl")
include("dofhandler.jl")
include("boundary.jl")
include("iterator.jl")

# Function exports
# mesh
export parse_mesh_triangle, cell_diameter, get_coordinates
export nodes, faces, cells, getcells, get_faces, getnodes, node
export Cell, Node, PolygonalMesh
export getncells, getnfaces, n_faces_per_cell, get_normal
export face_orientation, getfaceset
export getnodeset, getnodesets
export rectangle_mesh
export TriangleCell, RectangleCell
export reference_edge_nodes, getnnodes, get_cell_name
export get_vertices_matrix, getcells_matrix
export n_nodes_per_cell, get_cell_coordinates!
export cell_centroid, cell_volume

# Boundaries
export Dirichlet
export apply!

# Shapes
export get_nodal_points, volume
export get_num_faces
export get_num_vertices

# Quadratures
export QuadratureRule
export RefTetrahedron
export DefaultQuad, GrundmannMoeller, Strang, GaussLegendre
export getpoints, getweights, integrate

# basis
export Dubiner, Lagrange, jacobi, Legendre
export dubiner_basis, ∇dubiner_basis
export value, derivative, gradient_value
export getnbasefunctions
export get_default_geom_interpolator
export getorder, getlowerdiminterpol
export gettopology, get_interpolation

# finite Elements
export ContinuousLagrange, GenericFiniteElement

#utils
export get_affine_map

#Assembler
export start_assemble, assemble!, end_assemble

#FunctionSpaces
export VectorFunctionSpace, ScalarFunctionSpace, ScalarTraceFunctionSpace
export getnquadpoints, getdetJdV, shape_value
export shape_gradient, shape_divergence
export getnfacequadpoints, getfacedetJdS
export face_shape_value
export spatial_coordinate, reference_coordinate
export getnlocaldofs
export reinit!

#Discrete Functions
export function_value
export TrialFunction
export errornorm
export function_value, nodal_avg

#Handlers
export DofHandler, ndofs, ndofs_per_cell, celldofs!
export create_sparsity_pattern, reconstruct!
export dof_range

#ITerator
export CellIterator

end # module


# basic pin hole camera model

#TODO: create a package that defines a camera model, there are lots of different ones all over the place
"""
	$(TYPEDEF)
Data structure for a Camera model with parameters.
Use `CameraModel(width,height,fc,cc,skew,kc)` for easy construction.
"""
struct CameraModel
    width::Int		# image width
    height::Int		# image height
    fc::Vector{Float64}	# focal length in x and y
	f::Float64			# if one focal lenght is used
    cc::Vector{Float64}	# camera center
    skew::Float64	    # skew value
    kc::Vector{Float64} # distortion coefficients up to fifth order or NTuple{5,Float64}
    K::Matrix{Float64} # 3x3 camera calibration matrix (Camera intrinsics)
    Ki::Matrix{Float64} # inverse of a 3x3 camera calibratio matrix
end

"""
	$(SIGNATURES)
Constructor helper for creating a camera model.
"""
function CameraModel(width::Int,height::Int,fc::Vector{Float64},cc::Vector{Float64},skew::Float64,kc::Vector{Float64})::CameraModel
    KK = [fc[1]      skew  cc[1];
             0       fc[2] cc[2];
             0		    0     1]
    Ki = inv(KK)
    CameraModel(width,height,fc, mean(fc),cc,skew,kc,KK,Ki)
end

"""
	$(SIGNATURES)
Constructor helper for creating a camera model.
"""
function CameraModel(width::Int,height::Int,f::Float64,cc::Vector{Float64})::CameraModel
    KK = [f 0 cc[1];
          0 f cc[2];
          0	0     1.]
    Ki = inv(KK)
    return CameraModel(width,height,[f,f], f, cc, 0.0, [0.], KK,Ki)
end

mutable struct DepthCamera
  K::Array{Float64,2}
  shape::Tuple{Int, Int}
  skip::Int
  D::Vector{Float64}
  xs::Array{Float64}
  ys::Array{Float64}
  DepthCamera(K::Array{Float64,2};
      shape::Tuple{Int,Int}=(480,640),
      skip::Int=1,
      D::Vector{Float64}=zeros(5) ) = new( K, shape, skip, D, Array{Float64,2}(), Array{Float64,2}() )
end