function champs_4(t, fft_ubar, fft_vbar, fft_ug, fft_vg, A1, A2, matr,
                  conjmatr, sigma, llambda, Ktaubis, epsilon, Ntaumm)

    h1u, h1v = ftau(t,
		    fft_ubar, fft_vbar,
		    A1, A2,
		    matr, conjmatr,
		    sigma, llambda)

    h1u_fft = fft(h1u, dims=1)
    h1v_fft = fft(h1v, dims=1)

    dtu0u = h1u_fft[1, :] / Ntaumm
    dtu0v = h1v_fft[1, :] / Ntaumm

    h1u_fft[1, :] .= 0
    h1v_fft[1, :] .= 0

    h1u = epsilon * ifft(h1u_fft ./ (1im * Ktaubis), dims=1)
    h1v = epsilon * ifft(h1v_fft ./ (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t,
                              fft_ubar, fft_vbar,
			      dtu0u, dtu0v,
                              A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t,
			      fft_ubar, fft_vbar,
			      A1, A2,
                              matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu = champu1 .+ champu2
    champv = champv1 .+ champv2

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    dttu0u = champu_fft[1, :] / Ntaumm
    dttu0v = champv_fft[1, :] / Ntaumm

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth1u = epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dth1v = epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    champu, champv = ftau(t,
                          fft_ubar .+ h1u, fft_vbar .+ h1v,
			  A1, A2, matr,
                          conjmatr, sigma, llambda)
    h2u = champu .- dth1u
    h2v = champv .- dth1v

    h2u_fft = fft(h2u, dims=1)
    h2v_fft = fft(h2v, dims=1)

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    dtu1u = champu_fft[1, :] / Ntaumm
    dtu1v = champv_fft[1, :] / Ntaumm

    h2u_fft[1, :] .= 0.0
    h2v_fft[1, :] .= 0.0

    h2u = epsilon * ifft(h2u_fft ./ (1im * Ktaubis), dims=1)
    h2v = epsilon * ifft(h2v_fft ./ (1im * Ktaubis), dims=1)

    champu1, champv1 = d2uftau(t, fft_ubar, fft_vbar,
                               dtu0u, dtu0v, dtu0u, dtu0v,
                               A1, A2, matr, conjmatr, sigma,
                               llambda, Ntaumm)

    champu2, champv2 = dtduftau(t,
				fft_ubar, fft_vbar,
				dtu0u, dtu0v,
                                A1, A2,
				matr, conjmatr,
				sigma, llambda, Ntaumm)

    champu3, champv3 = duftau(t,
			      fft_ubar, fft_vbar,
			      dttu0u, dttu0v,
                              A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu4, champv4 = d2tftau(t,
			       fft_ubar, fft_vbar,
			       A1, A2,
                               matr, conjmatr,
			       sigma, llambda, Ntaumm)

    champu = champu1 .+ 2 * champu2 .+ champu3 .+ champu4
    champv = champv1 .+ 2 * champv2 .+ champv3 .+ champv4

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0.0
    champv_fft[1, :] .= 0.0

    dtth1u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dtth1v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t,
			      fft_ubar .+ h1u,
			      fft_vbar .+ h1v,
                              dtu1u .+ dth1u,
			      dtu1v .+ dth1v,
			      A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t,
			      fft_ubar .+ h1u,
			      fft_vbar .+ h1v,
			      A1, A2,
                              matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu = champu1 .+ champu2 .- dtth1u
    champv = champv1 .+ champv2 .- dtth1v

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth2u = epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dth2v = epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    h3u, h3v = ftau(t,
		    fft_ubar .+ h2u,
		    fft_vbar .+ h2v,
		    A1, A2,
		    matr, conjmatr,
		    sigma, llambda)

    h3u = h3u .- dth2u
    h3v = h3v .- dth2v

    h3u_fft = fft(h3u, dims=1)
    h3v_fft = fft(h3v, dims=1)

    h3u_fft[1, :] .= 0
    h3v_fft[1, :] .= 0

    dtauh3u = ifft(h3u_fft, dims=1)
    dtauh3v = ifft(h3v_fft, dims=1)

    h3u = epsilon * ifft(h3u_fft ./ (1im * Ktaubis), dims=1)
    h3v = epsilon * ifft(h3v_fft ./ (1im * Ktaubis), dims=1)

    C3u = fft_ubar .+ h3u
    C3v = fft_vbar .+ h3v

    ffu, ffv = ftau(t,
		    C3u .+ fft_ug,
		    C3v .+ fft_vg,
		    A1, A2,
		    matr, conjmatr,
                    sigma, llambda)

    champu, champv = ftau(t, C3u, C3v, A1, A2, matr, conjmatr,
			  sigma, llambda)

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champubaru = champu_fft[1, :] / Ntaumm
    champubarv = champv_fft[1, :] / Ntaumm

    champu1, champv1 = duftau(t, fft_ubar, fft_vbar,
                              champubaru, champubarv, A1, A2,
                              matr, conjmatr, sigma, llambda,
                              Ntaumm)

    champu2, champv2 = dtftau(t, fft_ubar, fft_vbar, A1, A2,
                              matr, conjmatr, sigma, llambda, Ntaumm)

    champu_fft = fft(champu1 .+ champu2, dims=1)
    champv_fft = fft(champv1 .+ champv2, dims=1)

    dttu0baru = champu_fft[1, :] / Ntaumm
    dttu0barv = champv_fft[1, :] / Ntaumm

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth1baru = epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dth1barv = epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t,
			      fft_ubar .+ h1u,
			      fft_vbar .+ h1v,
                              champubaru .+ dth1baru,
			      champubarv .+ dth1barv,
                              A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t,
			      fft_ubar .+ h1u,
			      fft_vbar .+ h1v,
			      A1, A2,
                              matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu_fft = fft(champu1 .+ champu2, dims=1)
    champv_fft = fft(champv1 .+ champv2, dims=1)

    dttu1baru = champu_fft[1, :] / Ntaumm
    dttu1barv = champv_fft[1, :] / Ntaumm

    champu1, champv1 = d2uftau(t,
			       fft_ubar, fft_vbar,
			       dtu0u, dtu0v,
                               champubaru, champubarv,
			       A1, A2, matr, conjmatr,
                               sigma, llambda, Ntaumm)

    champu2, champv2 = duftau(t,
			      fft_ubar, fft_vbar,
			      dttu0baru, dttu0barv,
                              A1, A2,
			      matr, conjmatr,
			      sigma, llambda,
                              Ntaumm)

    champu3, champv3 = dtduftau(t,
				fft_ubar, fft_vbar,
				dtu0u, dtu0v,
                                A1, A2,
				matr, conjmatr,
				sigma, llambda, Ntaumm)

    champu4, champv4 = dtduftau(t,
				fft_ubar, fft_vbar,
				champubaru, champubarv,
                                A1, A2,
				matr, conjmatr,
				sigma, llambda,
                                Ntaumm)

    champu5, champv5 = d2tftau(t,
			       fft_ubar, fft_vbar,
			       A1, A2,
			       matr, conjmatr,
                               sigma, llambda, Ntaumm)

    champu = champu1 .+ champu2 .+ champu3 .+ champu4 .+ champu5
    champv = champv1 .+ champv2 .+ champv3 .+ champv4 .+ champv5

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    dtttu0baru = champu_fft[1, :] / Ntaumm
    dtttu0barv = champv_fft[1, :] / Ntaumm

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dtth1baru = epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dtth1barv = epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    champu1, champv1 = d3uftau(t,
			       fft_ubar, fft_vbar,
			       dtu0u, dtu0v,
                               dtu0u, dtu0v,
			       champubaru, champubarv,
			       A1, A2,
			       matr, conjmatr,
			       sigma, llambda, Ntaumm)

    champu2, champv2 = dtd2uftau(t,
				 fft_ubar, fft_vbar,
				 dtu0u, dtu0v,
                                 dtu0u, dtu0v,
				 A1, A2,
				 matr, conjmatr,
				 sigma, llambda, Ntaumm)

    champu3, champv3 = d2uftau(t,
			       fft_ubar, fft_vbar,
                               dttu0baru, dttu0barv,
			       dtu0u, dtu0v,
			       A1, A2,
                               matr, conjmatr,
			       sigma, llambda, Ntaumm)

    champu4, champv4 = dtd2uftau(t,
				 fft_ubar, fft_vbar,
				 dtu0u, dtu0v,
                                 champubaru, champubarv,
				 A1, A2, matr, conjmatr,
                                 sigma, llambda, Ntaumm)

    champu5, champv5 = d2tduftau(t,
				 fft_ubar, fft_vbar,
				 dtu0u, dtu0v,
                                 A1, A2,
				 matr, conjmatr,
				 sigma, llambda, Ntaumm)

    champu6, champv6 = dtduftau(t, fft_ubar, fft_vbar,
                                dttu0baru, dttu0barv,
                                A1, A2,
				matr, conjmatr,
				sigma, llambda,
                                Ntaumm)

    champu7, champv7 = d2uftau(t,
			       fft_ubar, fft_vbar,
                               dttu0u, dttu0v,
			       champubaru, champubarv,
                               A1, A2,
			       matr, conjmatr,
                               sigma, llambda, Ntaumm)

    champu8, champv8 = dtduftau(t,
				fft_ubar, fft_vbar,
                                dttu0u, dttu0v,
				A1, A2,
				matr, conjmatr,
				sigma, llambda, Ntaumm)

    champu9, champv9 = duftau(t,
			      fft_ubar, fft_vbar,
                              dtttu0baru, dtttu0barv,
			      A1, A2,
			      matr, conjmatr,
			      sigma, llambda,
                              Ntaumm)

    champu10, champv10 = d3tftau(t,
				 fft_ubar, fft_vbar,
				 A1, A2,
                                 matr, conjmatr,
				 sigma, llambda, Ntaumm)

    champu11, champv11 = d2tduftau(t,
				   fft_ubar, fft_vbar,
                                   champubaru, champubarv,
				   A1, A2,
				   matr, conjmatr,
				   sigma, llambda, Ntaumm)

    champu = (champu1 .+ champu2 .+ 2 * champu3 .+ 2 * champu4
                     .+ 2 * champu5 .+ 2 * champu6 .+ champu7
		     .+ champu8 .+ champu9 .+ champu10 .+ champu11)

    champv = (champv1 .+ champv2 .+ 2 * champv3 .+ 2 * champv4
             .+ 2 * champv5 .+ 2 * champv6 .+ champv7 .+ champv8
	     .+ champv9 .+ champv10 .+ champv11)

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dttth1baru = epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dttth1barv = epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    champu1, champv1 = d2uftau(t,
			       fft_ubar .+ h1u, fft_vbar .+ h1v,
                               dtu1u .+ dth1u, dtu1v .+ dth1v,
                               champubaru .+ dth1baru,
			       champubarv .+ dth1barv,
                               A1, A2,
			       matr, conjmatr,
			       sigma, llambda, Ntaumm)

    champu2, champv2 = duftau(t,
			      fft_ubar .+ h1u,
			      fft_vbar .+ h1v,
                              dttu1baru .+ dtth1baru,
			      dttu1barv .+ dtth1barv,
                              A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu3, champv3 = dtduftau(t,
				fft_ubar .+ h1u,
				fft_vbar .+ h1v,
                                dtu1u .+ dth1u,
				dtu1v .+ dth1v,
				A1, A2,
				matr, conjmatr,
				sigma, llambda, Ntaumm)

    champu4, champv4 = dtduftau(t,
				fft_ubar .+ h1u, fft_vbar .+ h1v,
                                champubaru .+ dth1baru,
                                champubarv .+ dth1barv,
				A1, A2,
				matr, conjmatr,
				sigma, llambda,
                                Ntaumm)

    champu5, champv5 = d2tftau(t,
			       fft_ubar .+ h1u, fft_vbar .+ h1v,
			       A1, A2,
                               matr, conjmatr, sigma, llambda, Ntaumm)

    champu = champu1 .+ champu2 .+ champu3 .+ champu4 .+ champu5 .- dttth1baru
    champv = champv1 .+ champv2 .+ champv3 .+ champv4 .+ champv5 .- dttth1barv

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dtth2baru = epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dtth2barv = epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t,
			      fft_ubar .+ h1u,
			      fft_vbar .+ h1v,
                              champubaru .+ dth1baru,
			      champubarv .+ dth1barv,
                              A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t,
			      fft_ubar .+ h1u,
			      fft_vbar .+ h1v,
			      A1, A2,
                              matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu_fft = fft(champu1 .+ champu2 .- dtth1baru, dims=1)
    champv_fft = fft(champv1 .+ champv2 .- dtth1barv, dims=1)

    dttu0baru = champu_fft[1, :] / Ntaumm
    dttu0barv = champv_fft[1, :] / Ntaumm

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth2baru = epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dth2barv = epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t,
			      fft_ubar .+ h2u,
			      fft_vbar .+ h2v,
                              champubaru .+ dth2baru,
			      champubarv .+ dth2barv,
                              A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t,
			      fft_ubar .+ h2u,
			      fft_vbar .+ h2v,
			      A1, A2,
                              matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu = champu1 .+ champu2 .- dtth2baru
    champv = champv1 .+ champv2 .- dtth2barv

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dtC3u = champubaru .+ epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dtC3v = champubarv .+ epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    champgu = ffu .- dtauh3u .- dtC3u
    champgv = ffv .- dtauh3v .- dtC3v

    champgu_fft = fft(champgu, dims=1)
    champgv_fft = fft(champgv, dims=1)

    champmoyu = champgu_fft[1, :] / Ntaumm
    champmoyv = champgv_fft[1, :] / Ntaumm

    champgu_fft[1, :] .= 0
    champgv_fft[1, :] .= 0

    ichampgu = ifft(champgu_fft ./ (1im * Ktaubis), dims=1)
    ichampgv = ifft(champgv_fft ./ (1im * Ktaubis), dims=1)

    champubaru, champubarv, ichampgu, ichampgv, champmoyu, champmoyv

end


# Rotation utilities

export RotationLinOp, RotationParametericLinOp, RotationParametericDelayedEval, JacobianRotationEvaluated
export rotation_linop, rotation, Jacobian, ∂

## Linear operator

struct RotationLinOp{T}<:AbstractLinearOperator{T,3,3}
    R::AbstractArray{T,3}
    c::Union{Nothing,AbstractArray{T,2}}
    s::Union{Nothing,AbstractArray{T,2}}
end

rotation_linop(R::AbstractArray{T,3}; c::Union{Nothing,AbstractArray{T,2}}=nothing, s::Union{Nothing,AbstractArray{T,2}}=nothing) where {T<:Real} = RotationLinOp{T}(R, c, s)

function rotation_linop(φ::AbstractArray{T,2}; derivative::Bool=false) where {T<:Real}
    c = cos.(φ)
    s = sin.(φ)
    Rout = rotation_matrix(c, s; derivative=derivative)
    ~derivative ? (return RotationLinOp{T}(Rout, c, s)) : (return (RotationLinOp{T}(Rout[1], c, s), Rout[2]))
end

function rotation_matrix(c::AbstractArray{T,2}, s::AbstractArray{T,2}; derivative::Bool=false) where {T<:Real}
    c1 = c[:,1]; s1 = s[:,1]
    c2 = c[:,2]; s2 = s[:,2]
    c3 = c[:,3]; s3 = s[:,3]

    R = similar(c1, length(c1), 3, 3)
    R[:,1,1] .= c1.*c2;            R[:,1,2] .= -c2.*s1;           R[:,1,3] .= -s2;
    R[:,2,1] .= c3.*s1-c1.*s2.*s3; R[:,2,2] .= c1.*c3+s1.*s2.*s3; R[:,2,3] .= -c2.*s3;
    R[:,3,1] .= c1.*c3.*s2+s1.*s3; R[:,3,2] .= c1.*s3-c3.*s1.*s2; R[:,3,3] .= c2.*c3;
    R = permutedims(R, (1, 3, 2))

    if derivative
        ∂R = similar(c1, length(c1), 3, 3, 3)

        ∂R[:,1,1,1] .= -s1.*c2;            ∂R[:,1,2,1] .= -c2.*c1;            ∂R[:,1,3,1] .= 0;
        ∂R[:,2,1,1] .= c3.*c1+s1.*s2.*s3;  ∂R[:,2,2,1] .= -s1.*c3+c1.*s2.*s3; ∂R[:,2,3,1] .= 0;
        ∂R[:,3,1,1] .= -s1.*c3.*s2+c1.*s3; ∂R[:,3,2,1] .= -s1.*s3-c3.*c1.*s2; ∂R[:,3,3,1] .= 0;

        ∂R[:,1,1,2] .= -c1.*s2;     ∂R[:,1,2,2] .= s2.*s1;      ∂R[:,1,3,2] .= -c2;
        ∂R[:,2,1,2] .= -c1.*c2.*s3; ∂R[:,2,2,2] .= s1.*c2.*s3;  ∂R[:,2,3,2] .= s2.*s3;
        ∂R[:,3,1,2] .= c1.*c3.*c2;  ∂R[:,3,2,2] .= -c3.*s1.*c2; ∂R[:,3,3,2] .= -s2.*c3;

        ∂R[:,1,1,3] .= 0;                  ∂R[:,1,2,3] .= 0;                  ∂R[:,1,3,3] .= 0;
        ∂R[:,2,1,3] .= -s3.*s1-c1.*s2.*c3; ∂R[:,2,2,3] .= -c1.*s3+s1.*s2.*c3; ∂R[:,2,3,3] .= -c2.*c3;
        ∂R[:,3,1,3] .= -c1.*s3.*s2+s1.*c3; ∂R[:,3,2,3] .= c1.*c3+s3.*s1.*s2;  ∂R[:,3,3,3] .= -c2.*s3;

        ∂R = permutedims(∂R, (1, 3, 2, 4))
    end

    ~derivative ? (return R) : (return (R, ∂R))
end

AbstractLinearOperators.domain_size(R::RotationLinOp) = (size(R.R,1), "nk", 3)
AbstractLinearOperators.range_size(R::RotationLinOp) = domain_size(R)

function AbstractLinearOperators.matvecprod(R::RotationLinOp{T}, K::AbstractArray{T,3}) where {T<:Real}
    nt, nk, _ = size(K)
    return sum(reshape(R.R, nt, 1, 3, 3).*reshape(K, nt, nk, 1, 3); dims=4)[:,:,:,1]
end

function AbstractLinearOperators.matvecprod_adj(R::RotationLinOp{T}, K::AbstractArray{T,3}) where {T<:Real}
    nt, nk, _ = size(K)
    return sum(reshape(R.R, nt, 1, 3, 3).*reshape(K, nt, nk, 3, 1); dims=3)[:,:,1,:]
end

## Parameteric linear operator

struct RotationParametericLinOp end

rotation() = RotationParametericLinOp()

(::RotationParametericLinOp)(φ::AbstractArray{T,2}) where {T<:Real} = rotation_linop(φ)

(R::RotationParametericLinOp)() = R

## Parameteric delayed evaluation

struct RotationParametericDelayedEval{T}
    K::AbstractKSpaceSampling{T}
end

Base.:*(::RotationParametericLinOp, K::AbstractKSpaceSampling{T}) where {T<:Real} =  RotationParametericDelayedEval{T}(K)

(RK::RotationParametericDelayedEval{T})(φ::AbstractArray{T,2}) where {T<:Real} = rotation_linop(φ)*coord(RK.K)

## Jacobian of rotation evaluated

struct JacobianRotationEvaluated{T}<:AbstractLinearOperator{Complex{T},2,3}
    K::AbstractKSpaceSampling{T}
    ∂R::AbstractArray{T,4}
end

function Jacobian(RK::RotationParametericDelayedEval{T}, φ::AbstractArray{T,2}) where {T<:Real}
    Rφ, ∂Rφ = rotation_linop(φ; derivative=true)
    return Rφ*coord(RK.K), JacobianRotationEvaluated{T}(RK.K, ∂Rφ)
end

∂(RK::RotationParametericDelayedEval{T}, φ::AbstractArray{T,2}) where {T<:Real} = Jacobian(RK, φ)

AbstractLinearOperators.domain_size(∂RK::JacobianRotationEvaluated) = (size(∂RK.∂R,1),3)
AbstractLinearOperators.range_size(∂RK::JacobianRotationEvaluated) = (size(∂RK.K)...,3)

function AbstractLinearOperators.matvecprod(∂RK::JacobianRotationEvaluated{T}, Δφ::AbstractArray{CT,2}) where {T<:Real,CT<:RealOrComplex{T}}
    ΔR = sum(∂RK.∂R.*reshape(Δφ, :, 1, 1, 3); dims=4)[:,:,:,1]
    return rotation_linop(ΔR)*coord(∂RK.K)
end
Base.:*(∂RK::JacobianRotationEvaluated{T}, Δφ::AbstractArray{T,2}) where {T<:Real} = AbstractLinearOperators.matvecprod(∂RK, Δφ)

function AbstractLinearOperators.matvecprod_adj(∂RK::JacobianRotationEvaluated{T}, ΔRK::AbstractArray{CT,3}) where {T<:Real,CT<:RealOrComplex{T}}
    nt, nk, _ = size(ΔRK)
    KΔRK = sum(ΔRK.*reshape(coord(∂RK.K), nt, nk, 1, 3); dims=2)[:,1,:,:]
    return sum(∂RK.∂R.*reshape(KΔRK, nt, 3, 3, 1); dims=2:3)[:,1,1,:]
end

function LinearAlgebra.dot(∇u::AbstractArray{Complex{T},3}, ∂RK::JacobianRotationEvaluated{T}) where {T<:Real}
    nt, nk, _ = size(∇u)
    ∂R = ∂RK.∂R
    K = coord(∂RK.K)
    J = similar(∇u, nt, nk, 3)
    ∂jRK = similar(∇u, nt, nk, 3)
    @inbounds for j = 1:3
        ∂jRK .= sum(reshape(∂R[:,:,:,j], nt, 1, 3, 3).*reshape(K, nt, nk, 1, 3); dims=4)[:,:,:,1]
        J[:,:,j] .= sum(∇u.*∂jRK; dims=3)[:,:,1]
    end
    return J
end


PositiveFactorizations.floattype(::Type{Union{T,Missing}}) where {T} = float(T)

function newton_raphson(fgh!, x::AbstractArray{T}; ρ=one(T)/2, c=1e-4, tol=1e-4, max_iter=1000) where T
    n = length(x)
    R = Base.promote_op(/, T, T)
    grad = zeros(R, n)
    hes = zeros(R, n, n)
    y = zero(R)
    for i = 1:max_iter
        y = fgh!(x, grad, hes, true)
        search_dir = -(cholesky(Positive, hes) \ grad)
        norm(search_dir) > tol || return x, y, grad, hes
        step_size = one(T)
        while fgh!(x + search_dir * step_size, grad, hes, false) > y + c * step_size * grad'search_dir
            step_size *= ρ
            step_size > 1e-10 || error("Line search failed! Problematic Hessian or Gradient?")
        end
        x .+= search_dir * step_size
    end
    throw(ConvergenceException(max_iter))
end


using Parameters

struct DSGE_Model{data <: AbstractArray, alfa <: Float64}
    "Data"
    data::data
    "Alpha"
    alfa::alfa
end

function (problem::DSGE_Model)(pars)
    @unpack alfa, data = problem
    @unpack bet, epsilon, theta, sig, s2, phi, phi_pi, phi_y, rho_v = pars
    llh = log_like_dsge([alfa,bet,epsilon,theta,sig,s2,phi,phi_pi,phi_y,rho_v],data)
        return llh+logpdf(prior_bet,bet)+logpdf(prior_epsilon,epsilon)+logpdf(prior_theta,theta)+logpdf(prior_sig,sig)+logpdf(prior_s2,s2)+logpdf(prior_phi,phi)+logpdf(prior_phi_pi,phi_pi)+logpdf(prior_phi_y,phi_y)+logpdf(prior_rho_v,rho_v)
end

function problem_transform(p::DSGE_Model)
    as((bet = as_unit_interval, epsilon=as_positive_real,theta = as_unit_interval, sig = as_positive_real, s2 = as_positive_real, phi = as_positive_real, phi_pi = as_positive_real, phi_y = as_positive_real, rho_v = as_unit_interval))
end


module SplineFEEC

using CompactBases # Used for basis creation
using DomainSets
import DomainSets: (..), ×

using LinearAlgebra
using StaticArrays
using FastGaussQuadrature
using Kronecker
import Kronecker: ⊗
using FFTW

using Plots
using Cubature

export (..), ×, ⊗, ⊙, adjoint,
        PeriodicBSpline, DirichletBSpline, PeriodicBSplineTensorProductBasis, DirichletBSplineTensorProductBasis,
        mass_matrix, stiffness_matrix, L2_prod, col_mat, L2_proj, weighted_mass_matrix, interior_product_matrix,
        grad, curl, div, ∇, disc_diff, disc_grad, disc_curl, disc_div,
        Basis, TensorProductBasis, BSplineTensorProductBasis,
        DirichletMultipatchBSpline, mass_matrix_d, L2_prod_d, L2_proj_d, col_mat_d

include("bases.jl")
include("circulant.jl")
include("util.jl")
include("operator.jl")
include("multipatch.jl")
include("examples.jl")

end


function DSO_min_LoadShift(hours,
	nodes::Nodes,
	av_el_price,
	av_rd_price_dict,
	x_ag,
	x_re,
	x_ind,
	s_ag,
	s_re,
	s_ind,
	isworkhour
	)

# create subset of nodes
T = hours
J = nodes.id
Sbase = nodes.bmva

# calculate average nodal load
loadmean = Dict(i => sum(nodes.load[i][hours]/Sbase)/(length(hours)) for i in keys(nodes.load))

DSO_mod = JuMP.Model(with_optimizer(Gurobi.Optimizer))

	### VARIABLE DECLARATION
		### Variables for sectoral shift
			@variable(DSO_mod, ΔD_up_ag[T, J] >= 0);
			@variable(DSO_mod, ΔD_up_re[T, J] >= 0);
			@variable(DSO_mod, ΔD_up_ind[T, J] >= 0);
			@variable(DSO_mod, ΔD_dn_ag[T, J] >= 0);
			@variable(DSO_mod, ΔD_dn_re[T, J] >= 0);
			@variable(DSO_mod, ΔD_dn_ind[T, J] >= 0);
		### Variables for total shift
			@variable(DSO_mod, ΔD_up[T, J] >= 0);
			@variable(DSO_mod, ΔD_dn[T, J] >= 0);
		### Variable for new Loadprofile
			@variable(DSO_mod, Load[T, J] >= 0);

	### CONSTRAINT DECLARATION
		### Constraints for total shift
			@constraint(DSO_mod, ShiftUpTotal[j = J, t = T], ΔD_up[t, j] == ΔD_up_ag[t, j] + ΔD_up_re[t, j] + ΔD_up_ind[t, j]);
			@constraint(DSO_mod, ShiftDnTotal[j = J, t = T], ΔD_dn[t, j] == ΔD_dn_ag[t, j] + ΔD_dn_re[t, j] + ΔD_dn_ind[t, j]);
		### Constraints for agricultural, residential shift
			@constraint(DSO_mod, ShiftAgUpMax[j = J, t = T], ΔD_up_ag[t, j] <= x_ag*s_ag*loadmean[j]/Sbase);
			@constraint(DSO_mod, ShiftAgDnMax[j = J, t = T], ΔD_dn_ag[t, j] <= x_ag*s_ag*loadmean[j]/Sbase);
			@constraint(DSO_mod, ShiftReUpMax[j = J, t = T], ΔD_up_re[t, j] <= x_re*s_re*loadmean[j]/Sbase);
			@constraint(DSO_mod, ShiftReDnMax[j = J, t = T], ΔD_dn_re[t, j] <= x_re*s_re*loadmean[j]/Sbase);
		### Constraints for industrial shift
			# if mod(hours[24],168) < 120
			@constraint(DSO_mod, ShiftIndUpMax[j = J, t = T], ΔD_up_ind[t, j] <= x_ind*s_ind*loadmean[j]/Sbase*isworkhour[mod(t,24)]);
			@constraint(DSO_mod, ShiftIndDnMax[j = J, t = T], ΔD_dn_ind[t, j] <= x_ind*s_ind*loadmean[j]/Sbase*isworkhour[mod(t,24)]);
			#	else
			#@constraint(DSO_mod, ShiftIndUpMax[j = J, t = T], ΔD_up_ind[t, j] == 0);
			#@constraint(DSO_mod, ShiftIndDnMax[j = J, t = T], ΔD_dn_ind[t, j] == 0);
			# end
		### Constraints for ShiftBalances
			@constraint(DSO_mod, ShiftBalAg[j = J],
			sum(ΔD_up_ag[t, j] - ΔD_dn_ag[t, j] for t in hours) == 0);
			@constraint(DSO_mod, ShiftBalRe[j = J],
			sum(ΔD_up_re[t, j] - ΔD_dn_re[t, j] for t in hours) == 0);
			@constraint(DSO_mod, ShiftBalInd[j = J],
			sum(ΔD_up_ind[t, j] - ΔD_dn_ind[t, j] for t in hours) == 0);

		### Constraint for new Loadprofile
			@constraint(DSO_mod, ShiftedLoad[j = J, t = T],
			Load[t, j] == nodes.load[j][t]/Sbase + ΔD_up[t, j] - ΔD_dn[t, j]
			);

	### Objective Function
			@objective(DSO_mod, Min,
			sum(sum(av_el_price*Load[t, j] + av_rd_price_dict[j][t]*(ΔD_up[t, j]-ΔD_dn[t, j]) for t in T) for j in J)
				)

		# Initiate optimisation process
		JuMP.optimize!(DSO_mod)

ShiftUp = JuMP.value.(ΔD_up) * Sbase
ShiftDn = JuMP.value.(ΔD_dn) * Sbase
ShiftAgUp = JuMP.value.(ΔD_up_ag) * Sbase
ShiftAgDn = JuMP.value.(ΔD_dn_ag) * Sbase
ShiftReUp = JuMP.value.(ΔD_up_re) * Sbase
ShiftReDn = JuMP.value.(ΔD_dn_re) * Sbase
ShiftIndUp = JuMP.value.(ΔD_up_ind) * Sbase
ShiftIndDn = JuMP.value.(ΔD_dn_ind) * Sbase
Load_shifted_opt = JuMP.value.(Load) * Sbase

return(Load_shifted_opt,
	ShiftUp,
	ShiftDn,
)


#############################################################################
# modular.jl
# Handles modular functions of a given vector x.
#############################################################################

export modular
export lovasz
export sign, monotonicity, modularity, evaluate

mutable struct ModularAtom <: SubmodFunc   # TODO: extend to sets other than plain set varables
  head::Symbol
  id_hash::UInt64
  children::Tuple{AbstractExpr, CombiSet}
  size::Tuple{Int, Int}
  param::AbstractExpr
  setvariables::Array{CombiSet}

  function ModularAtom(w::AbstractExpr, S::CombiSet)
    if w.size != (S.cardinality, 1)
      error("Cannot define a modular function when the vector's size is different from the (card(baseset), 1).")
    else
      children = (w, S)
      setvariables = get_sv(S)
      return new(:modular, hash(children), children, (1, 1), w, setvariables)
    end
  end
end

modular(w::Variable, S::CombiSet) = ModularAtom(w, S)

modular(w::Values, S::CombiSet) = ModularAtom(Constant(w), S)

function lovasz(F::ModularAtom, x::Variable)
  if length(F.setvariables) == 1
    if evaluate(F, [])[1] != 0
      error("A combinatorial function should be 0 at the empty set to derive its Lovasz extension.")
    else
      if x.size[1] == F.setvariables[1].cardinality
        if isa(F.children[2], SetVariable)
          return vecdot(F.children[1], x)
        else
          return LovaszExtAtom(F, x)
        end
      else
        error("The size of the continuous variable should be the same as the baseset of the combinatorial variable of the combinatorial function.")
      end
    end
  else
    error("A combinatorial function should be sigle-variant to obtain its Lovasz extension.")
  end
end

function lovasz(F::ModularAtom, x::AbsAtom)
  if length(F.setvariables) == 1
    if length(x.children) == 1 && isa(x.children[1], Variable)
      if evaluate(F, [])[1] != 0
        error("A combinatorial function should be 0 at the empty set to derive its Lovasz extension.")
      else
        var = get_v(x)
        if length(var) != 1
          error("Functions with other than one variable are not supported.")
        else
          if var[1].size[1] == F.setvariables[1].cardinality
            if isa(F.children[2], SetVariable)
              return vecdot(F.children[1], x)
            else
              return LovaszExtAbsAtom(F, x)
            end
          else
            error("The size of the continuous variable should be the same as the baseset of the combinatorial variable of the combinatorial function.")
          end
        end
      end
    else
      error("Only able to define Lovasz extesions on absolute values of a variable.")
    end
  else
    error("A combinatorial function should be sigle-variant to obtain its Lovasz extension.")
  end
end

function sign(F::ModularAtom)
  return sign(F.param)
end

function monotonicity(F::ModularAtom)
  return (sign(F.children[2]) * Nondecreasing(), sign(F.children[1]) * Nondecreasing())
end

function modularity(F::ModularAtom)
  return Modularity()
end

function evaluate(F::ModularAtom) # TODO only allow scalar input
  summ = 0.0
  param = evaluate(F.param)
  elements = []
  elements = vcat(elements, get_elements(F.children[2]))
  elements = Set(elements)
  for i in elements
    summ += param[i]
  end
  # Compatibility with Convex.jl
  return summ
end


using Distributions

module GKDistribution

  import Distributions:
         ContinuousUnivariateDistribution,
         location,
         scale,
         skewness,
         kurtosis,
         params,
         minimum,
         maximum,
         insupport,
         quantile,
         logpdf,
         fit_mle,
         Normal,
         cdf,
         pdf

  import Optim:
         optimize
  import Roots:
         fzero

  include("gk.jl")
  export GK,
         logpdf,
         quantile,
         location,
         scale,
         skewness,
         kurtosis,
         asymmetry,
         logpdf,
         quantile,
         params,
         minimum,
         maximum,
         insupport,
         pdf

end # module


import IntertemporalChoiceHeuristics:
    load,
    fit,
    models,
    make_closures,
    LogLoss,
    invlogit,
    stderrs

inputs, weights = load()

loss = LogLoss()
λ = 0.0

for model in models
    res = fit(model, inputs, weights)

    nll, nll_gr! = make_closures(model, loss, λ, inputs, weights)

    mle = Optim.minimizer(res)

    se = stderrs(model, inputs, weights)

    K = length(mle)

    n_points = 256
    for j in 1:K
        parameters = copy(mle)
        grid = linspace(
            mle[j] - 6 * se[j],
            mle[j] + 6 * se[j],
            n_points,
        )
        for i in 1:n_points
            parameters[j] = grid[i]
            @printf(
                "%s\t%d\t%s\t%s\t%s\t%s\n",
                model,
                j,
                grid[i],
                nll(parameters),
                mle[j],
                se[j],
            )
        end
    end
end


include("RCF/kummer_extensions.jl")
include("RCF/class_fields.jl")
include("RCF/cyclo.jl")
include("RCF/rcf.jl")
include("RCF/autos.jl")
include("RCF/conductor.jl")
include("RCF/rcf_subfield.jl")


###
# Optimization problems for VoPP Upper and Lower Bounds
###
#Throughout we assume mu = 1, c = 0, and hbar s.t. E[hbar(Vc)] = 0.
#Methods with aleading underscore are not true bounds but approximations for debugging/configuration

module vopp

using JuMP, Gurobi, Ipopt, LinearAlgebra, QuadGK, Roots

include("LambertW.jl")
include("Helpers.jl")
include("closedFormUB.jl")
#include("debugging.jl")
include("upperBounds.jl")
include("UpperBoundsSymmetric.jl")
include("UpperBoundsUnimodal.jl")
include("lowerBoundsSymmetric.jl")
include("lowerBoundsUnimodal.jl")

end #ends module


using Test, DiscreteExteriorCalculus, DiscretePDEs
using UniqueVectors: UniqueVector
using AdmittanceModels: lossless_modes_dense, apply_transform

geo_code = GeoCode(gdspy_open(joinpath(@__DIR__, "cpw_resonator.gds"), "main"),
    layer_translation=Dict(23 => 50),
    layer_extrusion=Dict(21 => 10, 22 => 10, 23 => -100, 24 => 50),
    layer_volume_difference=[(24, 22, true), (23, 24, true), (23, 22, true)],
    layer_volume_groups=Dict("Superconductor" => [22], "Substrate" => [23], "Vacuum" => [24]),
    layer_mesh_size=Dict(21 => 10))
file_name = joinpath(@__DIR__, "cpw_resonator.geo")
geo_write!(file_name, geo_code, characteristic_length_factor=1.5, random_factor=1e-8,
    random_factor_3D=1e-8)

initialize!()
gmsh_open!(file_name)
N, K = 3, 4
mesh!(K)

node_tags, points, tcomp = get_triangulated_complex(N, K, 1e-6)
group_dict = get_physical_groups(node_tags, points)
comp = tcomp.complex
orient!(comp)
m = Metric(N)
mesh = Mesh(tcomp, circumcenter(m))

μ⁻_form = get_material(comp, 1/DiscretePDEs.μ₀, 3)
Λ⁻_form = get_material(comp, 0, 2)
σ_form = get_material(comp, 0, 2)
ϵr = 11.9 # silicon
ϵ_form = DiscretePDEs.ϵ₀ * (get_material(comp, 1, 2) +
    get_material(comp, group_dict["Substrate"], ϵr-1, 2))

pso, null_basis = electrodynamics_pso(m, mesh, Vector{Cell{N}}[],
    append!(subcomplex(comp, group_dict["Superconductor"]), boundary(comp)),
    μ⁻_form, Λ⁻_form, σ_form, ϵ_form)
constrained_pso = apply_transform(pso, null_basis)

λs, vs = lossless_modes_dense(constrained_pso, min_freq=1e9)
freqs = imag.(λs)/(2π)

hfss_answers = [47.0946, 140.140] * 1e9
@testset "cpw resonator" begin
    @test isapprox(freqs[1:2], hfss_answers, rtol=1.5e-1)
end

if false # set to true to display plots with Gmsh
    comp_points = UniqueVector([c.points[1] for c in comp.cells[1]])
    ordering = [findfirst(isequal(p), comp_points) for p in points]
    for i in 1:2
        v = null_basis * vs[:,i]
        v /= maximum(abs.(v))
        vec_A = sharp(m, comp, v)
        add_field!("Vector potential mode $i", node_tags, vec_A[ordering])
    end
    gui!()
end
rm(file_name)


using Wavelets

wf = waveletfilter(WT.db4)
N = 1024;
x0 = rand(N,N);
L = int(log2(N)-2)
tn = 10
#y=x0*0.0

println("iwt (N=",N,"x",N,"), ", L, " levels")
f(x0, L, wf) = for i = 1:tn; idwt(x0, wf, L); end
f(x0, L, wf);
@time f(x0, L, wf);


cd(@__DIR__); include("setups/grid23x22.jl")
gr(dpi = 200)

## build rNGWF
rNGWF, dic_l2x = rngwf_all_vectors(D, 𝚽; σ = 0.2 * maximum(D), thres = 0.15)
rNGWF_dual = (rNGWF * rNGWF') \ rNGWF
Γ = rngwf_lx(dic_l2x)
f = digit_img[:]

important_idx = sortperm((rNGWF' * f).^2; rev = true)
red_box_inds = [41, 43, 66, 69, 71, 73, 75, 84, 85, 86, 89, 99]
plot(layout = Plots.grid(10,10), size = (2300, 2200))
for i = 1:100
    plot!(clims = (-0.002, 0.02))
    grid_vector_plot!(important_idx[i], i, rNGWF)
    if i in red_box_inds
        plot_square!(Nx, Ny; subplot = i)
    end
end
plt = current()
savefig(plt, "../figs/Grid23x22_fdigit3_rngwf_top100.png")


mutable struct EquilibrationData{T <: Real, S} <: PreconditionerDataK2{T, S}
  P::Diagonal{T, S}
  C_equi::Diagonal{T, S}
end

function Equilibration(
  id::QM_IntData,
  fd::QM_FloatData{T},
  regu::Regularization{T},
  D::AbstractVector{T},
  K::Union{LinearOperator{T}, AbstractMatrix{T}},
) where {T <: Real}
  P = Diagonal(similar(D, id.nvar + id.ncon))
  P.diag .= one(T)
  C_equi = Diagonal(similar(D, id.nvar + id.ncon))
  return EquilibrationData(P, C_equi)
end

function update_preconditioner!(
  pdat::EquilibrationData{T},
  pad::PreallocatedData{T},
  itd::IterData{T},
  pt::Point{T},
  id::QM_IntData,
  fd::QM_FloatData{T},
  cnts::Counters,
) where {T <: Real}
  equilibrate_K2!(
    fd.Q.data,
    fd.A,
    pad.D,
    pad.regu.δ,
    id.nvar,
    id.ncon,
    pad.pdat.P,
    pad.pdat.C_equi,
    fd.uplo;
    ϵ = T(1.0e-4),
    max_iter = 100,
  )
  pdat.P.diag .= pdat.P.diag .^ 2
end


# optical_absorption.jl

"""
----------------------------------------------------------------------
Polaron absorption coefficient Γ(Ω).
----------------------------------------------------------------------
"""

"""
optical_absorption_ac(Ω::Float64, β::Float64, α::Float64, v::Float64, w::Float64)

    Calculate the absorption coefficient Γ(Ω) for the polaron at finite temperatures (equation (11a) in Devreese's et al.) for a given frequency Ω. β is the thermodynamic beta. v and w are the variational Polaron parameters that minimise the free energy, for the supplied α Frohlich coupling.
"""
function optical_absorption_ac(Ω, β, α, v, w)
    X = χ(Ω, β, α, v, w)
    Reχ = real(X)
    Imχ = imag(X)
    result = Ω * Imχ / (Ω^4 - 2 * Ω^2 * Reχ + Reχ^2 + Imχ^2)
    return result
end

"""
optical_absorption_ac(Ω::Float64, α::Float64, v::Float64, w::Float64)

    Calculate the absorption coefficient Γ(Ω) for the polaron at zero temperatures (equation (11a) in Devreese's et al.) for a given frequency Ω. β is the thermodynamic beta. v and w are the variational Polaron parameters that minimise the free energy, for the supplied α Frohlich coupling.
"""
function optical_absorption_ac(Ω, α, v, w)
    X = χ(Ω, α, v, w)
    Reχ = real(X)
    Imχ = imag(X)
    result = Ω * Imχ / (Ω^4 - 2 * Ω^2 * Reχ + Reχ^2 + Imχ^2)
    return result
end

"""
optical_absorption_dc(β::Float64, α::Float64, v::Float64, w::Float64)

    Calculate the absorption coefficient Γ(Ω) for the polaron at finite temperatures (equation (11a) in Devreese's et al.) at zero frequency. β is the thermodynamic beta. v and w are the variational Polaron parameters that minimise the free energy, for the supplied α Frohlich coupling.
"""
function optical_absorption_dc(β, α, v, w)
    X = χ_dc(β, α, v, w)
    Reχ = real(X)
    Imχ = imag(X)
    result = Ω * Imχ / (Ω^4 - 2 * Ω^2 * Reχ + Reχ^2 + Imχ^2)
    return result
end


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

# TODO: QuadGK is used to "hack" normalised Jacobi Polynomials.
#       to make sure they are orthonormal, not just orthogonal,
#       but this should be done properly...
#       or - indeed - move all this to a more general Pair potential
#       with arbitrary orthogonality - measure

using QuadGK

import JuLIP: evaluate,
              evaluate!,
              evaluate_d!,
              read_dict,
              write_dict

import JuLIP.MLIPs: alloc_B, alloc_dB, IPBasis

import Base.==

export Jacobi, Chebyshev

Chebyshev(N, args...; kwargs...) = Jacobi(0.5, 0.5, N, args...; kwargs...)

"""
`Jacobi{T} : ` represents the basis of Jacobi polynomials
parameterised by α, β up to some fixed maximum degree. Recall that Jacobi
polynomials are orthogonal on [-1,1] w.r.t. the weight
w(x) = (1-x)^α (1+x)^β.

### Constructor:
```
Jacobi(α, β, N)   # N = max degree
```

### Evaluate basis and derivatives:
```
x = 2*(rand() - 0.5)
P = zeros(N)
evaluate!(P, J, x)
dP = zeros(N)
evaluate_d!(P, dP, J, x)   # evaluates both P, dP
```

### Notes

`Jacobi(...)` precomputes the recursion coefficients using arbitrary
precision arithmetic, then stores them as `Vector{Float64}`. The recursion
is then given by
```
P_{n} = (A[n] * x + B[n]) * P_{n-1} + C[n] * P_{n-2}
```
"""
struct Jacobi{T} <: IPBasis
   α::T
   β::T
   A::Vector{T}
   B::Vector{T}
   C::Vector{T}
   nrm::Vector{T}
end

==(J1::Jacobi, J2::Jacobi) = (
      (J1.α == J2.α) && (J1.β == J2.β) && (length(J1) == length(J2))
   )

write_dict(J::Jacobi{T}) where {T} =
      Dict( "__id__" => "ACE_Jacobi",
            "T" => write_dict(T),
            "alpha" => J.alpha,
            "beta" => J.beta,
            "N" => length(J.A),
            "normalise" => !isempty(J.nrm) )

read_dict(::Val{SHIPs_Jacobi}, D::Dict) = read_dict(Val{ACE_Jacobi}(), D)

read_dict(::Val{:ACE_Jacobi}, D::Dict; T = read_dict(D["T"])) =
      Jacobi(T(D["alpha"]), T(D["beta"]), D["N"], T = T;
             normalise = D["normalise"])

function Jacobi(α, β, N, T=Float64; normalise=true)
   # precompute the recursion coefficients
   A = zeros(T, N)
   B = zeros(T, N)
   C = zeros(T, N)
   for n = 2:N
      c1 = big(2*n*(n+α+β)*(2*n+α+β-2))
      c2 = big(2*n+α+β-1)
      A[n] = T( big(2*n+α+β)*big(2*n+α+β-2)*c2 / c1 )
      B[n] = T( big(α^2 - β^2) * c2 / c1 )
      C[n] = T( big(-2*(n+α-1)*(n+β-1)*(2n+α+β)) / c1 )
   end
   J = Jacobi(T(α), T(β), A, B, C, T[])
   if normalise
      integrand = x -> evaluate(J, x).^2 * ((1-x)^α * (1+x)^β)
      nrm2 = quadgk(integrand, -1.0, 1.0)[1]
      J = Jacobi(T(α), T(β), A, B, C, nrm2.^(-0.5))
   end
   return J
end

Base.length(J::Jacobi) = maxdegree(J) + 1
maxdegree(J::Jacobi) = length(J.A)
alloc_B(J::Jacobi{T}, args...) where {T} = zeros(T, length(J))
alloc_dB(J::Jacobi{T}, args...) where {T} = zeros(T, length(J))

function evaluate!(P::AbstractVector, tmp, J::Jacobi, x)
   N = maxdegree(J)
   @assert length(P) >= N + 1
   @assert N+1 <= length(P)
   α, β = J.α, J.β
   P[1] = 1
   if N > 0
      #       (A * x + B) * P[1]
      #       A = 0.5 * (α+β+2)
      #       B = (α+1) - 0.5 * (α+β+2)
      P[2] = (α+1) + 0.5 * (α+β+2) * (x-1)
      if N > 1
         # 3-pt recursion
         @inbounds for n = 2:N
            P[n+1] = (J.A[n] * x + J.B[n]) * P[n] + J.C[n] * P[n-1]
         end
      end
   end
   # normalise
   if !isempty(J.nrm)  # if we want an orthonormal basis
      P .= P .* J.nrm
   end
   return P
end

function evaluate_d!(P::AbstractVector, dP::AbstractVector, tmp,
                    J::Jacobi, x::Number)
   N = maxdegree(J)
   @assert length(P) >= N+1
   @assert length(dP) >= N+1
   @assert N+1 <= min(length(P), length(dP))
   α, β = J.α, J.β
   P[1] = 1
   dP[1] = 0
   if N > 0
      P[2] = (α+1) + 0.5 * (α+β+2) * (x-1)
      dP[2] = 0.5 * (α+β+2)
      if N > 1
         @inbounds for n = 2:N
            c1 = J.A[n] * x + J.B[n]
            c2 = J.C[n]
            P[n+1] = c1 * P[n] + c2 * P[n-1]
            dP[n+1] = J.A[n] * P[n] + c1 * dP[n] + J.C[n] * dP[n-1]
         end
      end
   end
   if !isempty(J.nrm)  # if we want an orthonormal basis
      P .= P .* J.nrm
      dP .= dP .* J.nrm
   end
   # return P, dP
   return dP
end


abstract type MaternLikelihood{T<:Real} <: RegressionLikelihood{T} end

"""
    Matern3_2Likelihood(ρ::Real=1.0)

## Arguments
- `ρ::Real` : lengthscale

---

Matern 3/2 likelihood for regression:
see [wiki page](https://en.wikipedia.org/wiki/Mat%C3%A9rn_covariance_function)
---

For the analytical solution, it is augmented via:
```math
p(y|f,\\omega) = \\mathcal{N}(y|f,\\sigma^2\\omega)
```
Where ``\\omega \\sim \\mathcal{IG}(\\frac{\\nu}{2},\\frac{\\nu}{2})`` where ``\\mathcal{IG}`` is the inverse gamma distribution
See paper [Robust Gaussian Process Regression with a Student-t Likelihood](http://www.jmlr.org/papers/volume12/jylanki11a/jylanki11a.pdf)
"""
struct Matern3_2Likelihood{T<:Real,A<:AbstractVector{T}} <: MaternLikelihood{T}
    ρ::T
    c::A
    θ::A
    function Matern3_2Likelihood{T}(ρ::T) where {T<:Real}
        return new{T,Vector{T}}(ρ)
    end
    function Matern3_2Likelihood{T}(ρ::T, c²::A, θ::A) where {T<:Real,A<:AbstractVector{T}}
        return new{T,A}(ρ, c², θ)
    end
end

function Matern3_2Likelihood(ρ::T=1.0) where {T<:Real}
    return Matern3_2Likelihood{T}(ρ)
end

function implemented(
    ::Matern3_2Likelihood, ::Union{<:AnalyticVI,<:QuadratureVI,<:GibbsSampling}
)
    return true
end

function init_likelihood(
    likelihood::Matern3_2Likelihood{T},
    inference::AbstractInference{T},
    ::Int,
    nSamplesUsed::Int,
) where {T}
    if inference isa AnalyticVI || inference isa GibbsSampling
        Matern3_2Likelihood{T}(
            likelihood.ρ, abs2.(T.(rand(T, nSamplesUsed))), zeros(T, nSamplesUsed)
        )
    else
        Matern3_2Likelihood{T}(likelihood.ρ)
    end
end

function (l::Matern3_2Likelihood)(y::Real, f::Real)
    u = sqrt(3) * abs(y - f) / l.ρ
    return 4 * l.ρ / sqrt(3) * (one(T) + u) * exp(-u)
end

function Distributions.loglikelihood(l::Matern3_2Likelihood, y::Real, f::Real)
    u = sqrt(3) * abs(y - f) / l.ρ
    return log(4 * l.ρ / sqrt(3)) + log(one(T) + u) - u
end

function Base.show(io::IO, l::Matern3_2Likelihood{T}) where {T}
    return print(io, "Matern 3/2 likelihood (ρ = $(l.ρ))")
end

function compute_proba(
    l::Matern3_2Likelihood{T}, μ::AbstractVector{<:Real}, σ²::AbstractVector{<:Real}
) where {T<:Real}
    return μ, max.(σ², zero(T)) .+ 4 * l.ρ^2 / 3
end

## Local Updates ##
function local_updates!(
    l::Matern3_2Likelihood,
    y::AbstractVector,
    μ::AbstractVector{T},
    diag_cov::AbstractVector{T},
) where {T<:Real}
    l.c .= sqrt.(diag_cov + abs2.(μ - y))
    return l.θ .= 3.0 ./ (2.0 .* sqrt.(3) * l.c * l.ρ .+ 2 * l.ρ^2)
end

function sample_local!(l::Matern3_2Likelihood, y::AbstractVector, f::AbstractVector)
    l.c .= rand.(GeneralizedInverseGaussian.(3 / (2 * l.ρ^2), 2.0 .* abs2.(f - y), 1.5))
    set_ω!(l, l.c)
    return nothing
end

@inline ∇E_μ(l::Matern3_2Likelihood, ::AOptimizer, y::AbstractVector) = (2.0 * l.θ .* y,)
@inline ∇E_Σ(l::Matern3_2Likelihood, ::AOptimizer, ::AbstractVector) = (l.θ,)

## ELBO  ##

function expecLogLikelihood(
    l::Matern3_2Likelihood,
    i::AnalyticVI,
    y::AbstractVector,
    μ::AbstractVector,
    diag_cov::AbstractVector,
)
    return throw(error("Not implemented yet"))
end

function InverseGammaKL(l::Matern3_2Likelihood{T}) where {T}
    α_p = l.ν / 2
    β_p = α_p * l.σ^2
    return model.inference.ρ * sum(broadcast(InverseGammaKL, l.α, l.β, α_p, β_p))
end

## PDF and Log PDF Gradients ## (verified gradients)

@inline function ∇loglikehood(l::Matern3_2Likelihood, y::Real, f::Real)
    return 3.0 * (y - f) / (l.ρ * (abs(f - y) * sqrt(3) + l.ρ))
end

@inline function hessloglikelihood(l::Matern3_2Likelihood, y::Real, f::Real)
    return 3.0 / (l.ρ + sqrt(3) * abs(f - y))^2
end


"""
    KSparseLMO{T}(K::Int, right_hand_side::T)

LMO for the K-sparse polytope:
```
C = B_1(τK) ∩ B_∞(τ)
```
with `τ` the `right_hand_side` parameter.
The LMO results in a vector with the K largest absolute values
of direction, taking values `-τ sign(x_i)`.
"""
struct KSparseLMO{T} <: LinearMinimizationOracle
    K::Int
    right_hand_side::T
end

function compute_extreme_point(lmo::KSparseLMO{T}, direction; kwargs...) where {T}
    K = min(lmo.K, length(direction))
    K_indices = sortperm(direction[1:K], by=abs, rev=true)
    K_values = direction[K_indices]
    for idx in K+1:length(direction)
        new_val = direction[idx]
        # new greater value: shift everything right
        if abs(new_val) > abs(K_values[1])
            K_values[2:end] .= K_values[1:end-1]
            K_indices[2:end] .= K_indices[1:end-1]
            K_indices[1] = idx
            K_values[1] = new_val
            # new value in the interior
        elseif abs(new_val) > abs(K_values[K])
            # NOTE: not out of bound since unreachable with K=1
            j = K - 1
            while abs(new_val) > abs(K_values[j])
                j -= 1
            end
            K_values[j+1:end] .= K_values[j:end-1]
            K_indices[j+1:end] .= K_indices[j:end-1]
            K_values[j+1] = new_val
            K_indices[j+1] = idx
        end
    end
    v = spzeros(T, length(direction))
    for (idx, val) in zip(K_indices, K_values)
        v[idx] = -lmo.right_hand_side * sign(val)
    end
    return v
end

function convert_mathopt(
    lmo::KSparseLMO{T},
    optimizer::OT;
    dimension::Integer,
    kwargs...,
) where {T,OT}
    τ = lmo.right_hand_side
    n = dimension
    K = min(lmo.K, n)
    MOI.empty!(optimizer)
    x = MOI.add_variables(optimizer, n)
    tinf = MOI.add_variable(optimizer)
    MOI.add_constraint(optimizer, MOI.VectorOfVariables([tinf; x]), MOI.NormInfinityCone(n + 1))
    MOI.add_constraint(optimizer, tinf, MOI.LessThan(τ))
    t1 = MOI.add_variable(optimizer)
    MOI.add_constraint(optimizer, MOI.VectorOfVariables([t1; x]), MOI.NormOneCone(n + 1))
    MOI.add_constraint(optimizer, t1, MOI.LessThan(τ * K))
    return MathOptLMO(optimizer)
end

"""
    BirkhoffPolytopeLMO

The Birkhoff polytope encodes doubly stochastic matrices.
Its extreme vertices are all permutation matrices of side-dimension `dimension`.
"""
struct BirkhoffPolytopeLMO <: LinearMinimizationOracle end

function compute_extreme_point(
    ::BirkhoffPolytopeLMO,
    direction::AbstractMatrix{T};
    kwargs...,
) where {T}
    n = size(direction, 1)
    n == size(direction, 2) ||
        DimensionMismatch("direction should be square and matching BirkhoffPolytopeLMO dimension")
    res_mat = Hungarian.munkres(direction)
    m = spzeros(Bool, n, n)
    (rows, cols, vals) = SparseArrays.findnz(res_mat)
    @inbounds for i in eachindex(cols)
        m[rows[i], cols[i]] = vals[i] == 2
    end
    m = convert(SparseArrays.SparseMatrixCSC{Float64,Int64}, m)
    return m
end

function compute_extreme_point(
    lmo::BirkhoffPolytopeLMO,
    direction::AbstractVector{T};
    kwargs...,
) where {T}
    nsq = length(direction)
    n = isqrt(nsq)
    return compute_extreme_point(lmo, reshape(direction, n, n); kwargs...)
end

function convert_mathopt(
    ::BirkhoffPolytopeLMO,
    optimizer::OT;
    dimension::Integer,
    kwargs...,
) where {OT}
    n = dimension
    MOI.empty!(optimizer)
    (x, _) = MOI.add_constrained_variables(optimizer, fill(MOI.Interval(0.0, 1.0), n * n))
    xmat = reshape(x, n, n)
    for idx in 1:n
        # column constraint
        MOI.add_constraint(
            optimizer,
            MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(ones(n), xmat[:, idx]), 0.0),
            MOI.EqualTo(1.0),
        )
        # row constraint
        MOI.add_constraint(
            optimizer,
            MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(ones(n), xmat[idx, :]), 0.0),
            MOI.EqualTo(1.0),
        )
    end
    return MathOptLMO(optimizer)
end

"""
    ScaledBoundLInfNormBall(lower_bounds, upper_bounds)

Polytope similar to a L-inf-ball with shifted bounds or general box constraints.
Lower- and upper-bounds are passed on as abstract vectors, possibly of different types.
For the standard L-inf ball, all lower- and upper-bounds would be -1 and 1.
"""
struct ScaledBoundLInfNormBall{T, VT1 <: AbstractVector{T}, VT2 <: AbstractVector{T}} <: LinearMinimizationOracle
    lower_bounds::VT1
    upper_bounds::VT2
end

function compute_extreme_point(lmo::ScaledBoundLInfNormBall, direction; kwargs...)
    v = copy(lmo.lower_bounds)
    for i in eachindex(direction)
        if direction[i] * lmo.upper_bounds[i] < direction[i] * lmo.lower_bounds[i]
            v[i] = lmo.upper_bounds[i]
        end
    end
    return v
end

"""
    ScaledBoundL1NormBall(lower_bounds, upper_bounds)

Polytope similar to a L1-ball with shifted bounds.
It is the convex hull of two scaled and shifted unit vectors for each axis (shifted to the center of the polytope, i.e., the elementwise midpoint of the bounds).
Lower and upper bounds are passed on as abstract vectors, possibly of different types.
For the standard L1-ball, all lower and upper bounds would be -1 and 1.
"""
struct ScaledBoundL1NormBall{T, VT1 <: AbstractVector{T}, VT2 <: AbstractVector{T}} <: LinearMinimizationOracle
    lower_bounds::VT1
    upper_bounds::VT2
end

function compute_extreme_point(lmo::ScaledBoundL1NormBall, direction; kwargs...)
    idx = 0
    lower = false
    val = zero(eltype(direction))
    if length(direction) != length(lmo.upper_bounds)
        throw(DimensionMismatch())
    end
    @inbounds for i in eachindex(direction)
        scale_factor = lmo.upper_bounds[i] - lmo.lower_bounds[i]
        scaled_dir = direction[i] * scale_factor
        if scaled_dir > val
            val = scaled_dir
            idx = i
            lower = true
        elseif -scaled_dir > val
            val = -scaled_dir
            idx = i
            lower = false
        end
    end
    # compute midpoint for all coordinates, replace with extreme coordinate on one
    # TODO use smarter array type if bounds are FillArrays
    v = (lmo.lower_bounds + lmo.upper_bounds) / 2
    # handle zero direction
    idx = max(idx, 1)
    v[idx] = ifelse(lower, lmo.lower_bounds[idx], lmo.upper_bounds[idx])
    return v
end


"""
Module MSM implements Calvet and Fisher's Markov Switching Multifractal model.

# Example usge:
```jldocs
julia> m = MSMmodel()
MSMmodel
  k: Int64 4
  b: Float64 5.0
  m₀: Float64 1.7
  γₖ: Float64 0.9
  σ₀: Float64 0.01
# Fit model to data x
julia> fit!(m, x);
# Predict volatility of x
julia> s = predict(m);
julia> histogram(s)
# Simulate a model
julia> simulate(m, nsims = 10000)
```
"""
module MSM

export MSMmodel, fit!, fitglobal!, predict, simulate, tune

using Parameters, DataFrames, Statistics, Random, Optim, Kronecker, LinearAlgebra, Distributions, BlackBoxOptim

const PPI = 1.0 / sqrt(2π)

"""
    MSMmodel(k, b, m₀, γₖ, σ₀)

Constructor for Markov Switching Multifractal (MSM) model.

# Example:
```jldocs
julia> m = MSMmodel()
MSMmodel
  k: Int64 4
  b: Float64 5.0
  m₀: Float64 1.7
  γₖ: Float64 0.9
  σ₀: Float64 0.01
```
"""
@with_kw mutable struct MSMmodel
    @deftype Float64
    k::Int = 4
    @assert k >= 1
    b = 5.0
    @assert b > 1.0
    m₀ = 1.7
    @assert m₀ > 1.0 && m₀ < 2.0
    γₖ = 0.90
    @assert γₖ > 0.0 && γₖ < 1.0
    σ₀ = 0.01
    @assert σ₀ > 0.0
end

# transition probability matrix
function Φ(b, γₖ, k::Integer)
    γ = [1 - (1 - γₖ)^(b^(j - k)) for j ∈ 1:k] .* 0.5
    return Matrix(reduce(⊗, [[1.0-u u; u 1.0-u] for u ∈ γ]))
end

# state dependent volatility
Σ(σ₀, m₀, k::Integer, M::Vector{Int}) = [σ₀ * sqrt(((2.0 - m₀)^u) * (m₀^(k - u))) for u ∈ M]

Σ(σ₀, m₀, k::Integer, m::Int) = σ₀ * sqrt(((2.0 - m₀)^m) * (m₀^(k - m)))

"""
    fit!(model::MSMmodel, x::Vector{Float64}; time_limit=60., solver=NelderMead(), g_tol=1.0e-6)

Fit MSM model to a vector of real-valued, de-meaned data using Optim's optimization methods. Recommended method is the default method.

# Example:
```jldocs
julia> using MSM;
julia> m = MSMmodel(); fit!(m, x)
```
"""
function fit!(model::MSMmodel, x::Vector{Float64}; time_limit = 60.0, solver = NelderMead(), g_tol = 1.0e-6)
    @unpack k, b, m₀, γₖ, σ₀ = model
    p0 = [b, m₀, γₖ, σ₀]
    l = [1.001, 1.001, 0.001, 0.0001]
    h = [50.0, 1.999, 0.999, 5.0]
    M = sum.([digits(q, base = 2, pad = k) for q = 0:2^k-1])
    N = length(M)
    Π = fill(1.0 / N, N)
    A = Φ(b, γₖ, k)
    res = optimize(p -> nll(p, k, M, A, Π, x), l, h, p0, Fminbox(solver), Optim.Options(time_limit = time_limit, g_tol = g_tol))
    b, m₀, γₖ, σ₀ = res.minimizer
    @pack! model = b, m₀, γₖ, σ₀
    return res
end

"""
    fitglobal!(model::MSMmodel, x::Vector{Float64}; method = :adaptive_de_rand_1_bin_radiuslimited,  maxsteps = 10000, tracemode = :silent)

Fit MSM model to a vector of real-valued, de-meaned data using BlackBoxOptim's global optimization methods. Recommended method is the default method.

# Example:
```jldocs
julia> using MSM;
julia> m = MSMmodel(); fitglobal!(m, x)
```

tracemode = :silent suppresses optimizer trace messages.
"""
function fitglobal!(model::MSMmodel, x::Vector{Float64}; method = :adaptive_de_rand_1_bin_radiuslimited, maxsteps = 10000, tracemode = :silent)
    @unpack k, b, m₀, γₖ, σ₀ = model
    l = [1.001, 1.001, 0.001, 0.0001]
    h = [50.0, 1.999, 0.999, 5.0]
    M = sum.([digits(q, base = 2, pad = k) for q = 0:2^k-1])
    N = length(M)
    Π = fill(1.0 / N, N)
    A = Φ(b, γₖ, k)
    res = bboptimize(p -> nll(p, k, M, A, Π, x), SearchRange = collect(zip(l, h)), Method = :adaptive_de_rand_1_bin_radiuslimited, MaxSteps = maxsteps, TraceMode = tracemode)
    b, m₀, γₖ, σ₀ = best_candidate(res)
    @pack! model = b, m₀, γₖ, σ₀
    return res
end

"""
    nll(p::Vector{<:Real}, k::Integer, M::Vector{Int}, A::Matrix{Float64}, Π::Vector{Float64}, rt::Vector{Float64})

compute negative log-likelihood of MSM.
"""
function nll(p::Vector{<:Real}, k::Integer, M::Vector{Int}, A::Matrix{Float64}, Π::Vector{Float64}, rt::Vector{Float64})
    b, m₀, γₖ, σ₀ = p
    A .= Φ(b, γₖ, k)
    σ = Σ(σ₀, m₀, k, M)
    ll = 0.0
    for r in rt
        ωr = [PPI * exp(-0.5 * r * r / (s * s)) / s for s in σ]
        w = ωr .* (A * Π)
        sw = sum(w)
        ll += log(sw)
        Π = w ./ sw
    end
    return -ll
end

"""
    predict(model::MSMmodel, paths, pathlen)

Predict volatility implied by an MSM model through simulation. `pathlen` is the length of observations over which volatility (standard deviation) is computed and `paths` is the number of paths simulated. Returns a vector of length nsamples containing volatilities.
"""
function predict(model::MSMmodel, paths, pathlen)
    s = fill(0.0, paths)
    @inbounds for i = 1:paths
        r = simulate(model, pathlen)
        s[i] = std(r)
    end
    return s
end

"""
    simulate(model::MSMmodel, pathlen)

Simulate an MSM model's trajectory of length pathlen. Returns a vector.
"""
function simulate(model::MSMmodel, pathlen)
    @unpack k, b, m₀, γₖ, σ₀ = model
    A = Φ(b, γₖ, k)
    M = sum.([digits(q, base = 2, pad = k) for q = 0:2^k-1])
    N = length(M)
    u = rand(Categorical(N))
    r = Vector{Float64}(undef, pathlen)
    r[1] = Σ(σ₀, m₀, k, M[u]) * randn()
    @inbounds for j = 2:pathlen
        u = rand(Categorical(A[u, :]))
        r[j] = Σ(σ₀, m₀, k, M[u]) * randn()
    end
    return r
end

"""
    tune(x::Vector{<:Real}; ks=2:12)

Fit MSM models for different k's and output results. k > 10 may take more than a minute to fit. Returns a dataframe.

"""
function tune(x::Vector{<:Real}; ks = 2:12)
    res = DataFrame()
    for j in ks
        model = MSM(k = j, σ₀ = std(x))
        out = @timed fit!(model, x)
        @unpack k, b, m₀, γₖ, σ₀ = model
        push!(res, (; ([:exec_time, :loglik, :k, :b, :m₀, :γₖ, :σ₀] .=> [out.time, -out.value.minimum, k, b, m₀, γₖ, σ₀])...))
        println("k = $(k) took $(out.time) seconds.")
    end
    return res
end

end


module ParaxialBeamPropagation

using FFTW, Interpolations,
ClassicalOrthogonalPolynomials,SpecialFunctions, Plots,
BenchmarkTools, PlutoUI,Printf,LaTeXStrings

export get_fourier_transform,display_fourier_transform,
propagate_beam,display_beam,animate_beam,HG,LG,airy_beam,vertical_obstacle,horizontal_obstacle,
circular_obstacle,rectangular_obstacle

include("Propagation.jl")
include("InitialProfiles.jl")

end


using Pkg
Pkg.activate(@__DIR__)

using ProgressiveHedging
import JuMP
import Ipopt

function create_model(scenario_id::ScenarioID)

    model = JuMP.Model(()->Ipopt.Optimizer())
    JuMP.set_optimizer_attribute(model, "print_level", 0)
    JuMP.set_optimizer_attribute(model, "tol", 1e-12)

    c = [1.0, 10.0, 0.01]
    d = 7.0
    a = 16.0

    α = 1.0
    β = 1.0
    γ = 1.0
    δ = 1.0
    ϵ = 1.0

    s1 = 8.0
    s2 = 4.0
    s11 = 9.0
    s12 = 16.0
    s21 = 5.0
    s22 = 18.0

    stage1 = JuMP.@variable(model, x[1:3] >= 0.0)
    JuMP.@constraint(model, x[3] <= 1.0)
    obj = zero(JuMP.GenericQuadExpr{Float64,JuMP.VariableRef})
    JuMP.add_to_expression!(obj, sum(c.*x))

    # Second stage
    stage2 = Vector{JuMP.VariableRef}()
    if scenario_id < scid(2)
        vref = JuMP.@variable(model, y >= 0.0)
        JuMP.@constraint(model, α*sum(x) + β*y >= s1)
        JuMP.add_to_expression!(obj, d*y)
    else
        vref = JuMP.@variable(model, y >= 0.0)
        JuMP.@constraint(model, α*sum(x) + β*y >= s2)
        JuMP.add_to_expression!(obj, d*y)
    end
    push!(stage2, vref)

    # Third stage
    stage3 = Vector{JuMP.VariableRef}()
    if scenario_id == scid(0)
        vref = JuMP.@variable(model, z[1:2])
        JuMP.@constraint(model, ϵ*sum(x) + γ*y + δ*sum(z) == s11)
        JuMP.add_to_expression!(obj, a*sum(z[i]^2 for i in 1:2))

    elseif scenario_id == scid(1)
        vref = JuMP.@variable(model, z[1:2])
        JuMP.@constraint(model, ϵ*sum(x) + γ*y + δ*sum(z) == s12)
        JuMP.add_to_expression!(obj, a*sum(z[i]^2 for i in 1:2))

    elseif scenario_id == scid(2)
        vref = JuMP.@variable(model, z[1:2])
        JuMP.@constraint(model, ϵ*sum(x) + γ*y + δ*sum(z) == s21)
        JuMP.add_to_expression!(obj, a*sum(z[i]^2 for i in 1:2))

    else
        vref = JuMP.@variable(model, z[1:2])
        JuMP.@constraint(model, ϵ*sum(x) + γ*y + δ*sum(z) == s22)
        JuMP.add_to_expression!(obj, a*sum(z[i]^2 for i in 1:2))
    end
    append!(stage3, vref)

    JuMP.@objective(model, Min, obj)

    vdict = Dict{StageID, Vector{JuMP.VariableRef}}([stid(1) => stage1,
                                                        stid(2) => stage2,
                                                        stid(3) => stage3,
                                                        ])

    return JuMPSubproblem(model, scenario_id, vdict)
end

scen_tree = ScenarioTree()
branch_node_1 = add_node(scen_tree, root(scen_tree))
branch_node_2 = add_node(scen_tree, root(scen_tree))

add_leaf(scen_tree, branch_node_1, 0.375)
add_leaf(scen_tree, branch_node_1, 0.125)

add_leaf(scen_tree, branch_node_2, 0.375)
add_leaf(scen_tree, branch_node_2, 0.125)

(niter, abs_res, rel_res, obj, soln_df, phd) = solve(scen_tree,
                                                     create_model,
                                                     ScalarPenaltyParameter(25.0);
                                                     atol=1e-8, rtol=1e-12, max_iter=500)
@show niter
@show abs_res
@show rel_res
@show obj
@show soln_df


using IterativeSolvers
using Test
using Random
using LinearAlgebra

@testset ("BiCGStab(l)") begin

Random.seed!(1234321)
n = 20

@testset "Matrix{$T}" for T in (Float32, Float64, ComplexF32, ComplexF64)
    A = rand(T, n, n) + 15I
    x = ones(T, n)
    b = A * x

    tol = √eps(real(T))

    @testset "BiCGStab($l)" for l = (2, 4)
        # Solve without preconditioner
        x1, his1 = bicgstabl(A, b, l, max_mv_products = 100, log = true, tol = tol)
        @test isa(his1, ConvergenceHistory)
        @test norm(A * x1 - b) / norm(b) ≤ tol

        # With an initial guess
        x_guess = rand(T, n)
        x2, his2 = bicgstabl!(x_guess, A, b, l, max_mv_products = 100, log = true, tol = tol)
        @test isa(his2, ConvergenceHistory)
        @test x2 == x_guess
        @test norm(A * x2 - b) / norm(b) ≤ tol

        # Do an exact LU decomp of a nearby matrix
        F = lu(A + rand(T, n, n))
        x3, his3 = bicgstabl(A, b, Pl = F, l, max_mv_products = 100, log = true, tol = tol)
        @test norm(A * x3 - b) / norm(b) ≤ tol
    end
end
end


using Logging
using JuMP, Gurobi
using DecisionProgramming

const N = 4
const health = [3*k - 2 for k in 1:N]
const test = [3*k - 1 for k in 1:(N-1)]
const treat = [3*k for k in 1:(N-1)]
const cost = [(3*N - 2) + k for k in 1:(N-1)]
const price = [(3*N - 2) + N]
const health_states = ["ill", "healthy"]
const test_states = ["positive", "negative"]
const treat_states = ["treat", "pass"]

@info("Creating the influence diagram.")
S = States([
    (length(health_states), health),
    (length(test_states), test),
    (length(treat_states), treat),
])
C = Vector{ChanceNode}()
D = Vector{DecisionNode}()
V = Vector{ValueNode}()
X = Vector{Probabilities}()
Y = Vector{Consequences}()

for j in health[[1]]
    I_j = Vector{Node}()
    X_j = zeros(S[I_j]..., S[j])
    X_j[1] = 0.1
    X_j[2] = 1.0 - X_j[1]
    push!(C, ChanceNode(j, I_j))
    push!(X, Probabilities(j, X_j))
end

for (i, k, j) in zip(health[1:end-1], treat, health[2:end])
    I_j = [i, k]
    X_j = zeros(S[I_j]..., S[j])
    X_j[2, 2, 1] = 0.2
    X_j[2, 2, 2] = 1.0 - X_j[2, 2, 1]
    X_j[2, 1, 1] = 0.1
    X_j[2, 1, 2] = 1.0 - X_j[2, 1, 1]
    X_j[1, 2, 1] = 0.9
    X_j[1, 2, 2] = 1.0 - X_j[1, 2, 1]
    X_j[1, 1, 1] = 0.5
    X_j[1, 1, 2] = 1.0 - X_j[1, 1, 1]
    push!(C, ChanceNode(j, I_j))
    push!(X, Probabilities(j, X_j))
end

for (i, j) in zip(health, test)
    I_j = [i]
    X_j = zeros(S[I_j]..., S[j])
    X_j[1, 1] = 0.8
    X_j[1, 2] = 1.0 - X_j[1, 1]
    X_j[2, 2] = 0.9
    X_j[2, 1] = 1.0 - X_j[2, 2]
    push!(C, ChanceNode(j, I_j))
    push!(X, Probabilities(j, X_j))
end

for (i, j) in zip(test, treat)
    I_j = [i]
    push!(D, DecisionNode(j, I_j))
end

for (i, j) in zip(treat, cost)
    I_j = [i]
    Y_j = zeros(S[I_j]...)
    Y_j[1] = -100
    Y_j[2] = 0
    push!(V, ValueNode(j, I_j))
    push!(Y, Consequences(j, Y_j))
end

for (i, j) in zip(health[end], price)
    I_j = [i]
    Y_j = zeros(S[I_j]...)
    Y_j[1] = 300
    Y_j[2] = 1000
    push!(V, ValueNode(j, I_j))
    push!(Y, Consequences(j, Y_j))
end

validate_influence_diagram(S, C, D, V)
sort!.((C, D, V, X, Y), by = x -> x.j)

P = DefaultPathProbability(C, X)
U = DefaultPathUtility(V, Y)

@info("Creating the decision model.")
U⁺ = PositivePathUtility(S, U)
model = Model()
z = DecisionVariables(model, S, D)
x_s = PathCompatibilityVariables(model, z, S, P, probability_cut = false)
EV = expected_value(model, x_s, U⁺, P)
@objective(model, Max, EV)

@info("Starting the optimization process.")
optimizer = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "IntFeasTol"      => 1e-9,
)
set_optimizer(model, optimizer)
optimize!(model)

@info("Extracting results.")
Z = DecisionStrategy(z)

@info("Printing decision strategy:")
print_decision_strategy(S, Z)

@info("State probabilities:")
sprobs = StateProbabilities(S, P, Z)
print_state_probabilities(sprobs, health)
print_state_probabilities(sprobs, test)
print_state_probabilities(sprobs, treat)

@info("Conditional state probabilities")
node = 1
for state in 1:2
    sprobs2 = StateProbabilities(S, P, Z, node, state, sprobs)
    print_state_probabilities(sprobs2, health)
    print_state_probabilities(sprobs2, test)
    print_state_probabilities(sprobs2, treat)
end

@info("Computing utility distribution.")
udist = UtilityDistribution(S, P, U, Z)

@info("Printing utility distribution.")
print_utility_distribution(udist)

@info("Printing statistics")
print_statistics(udist)


function Statistics.mean(states::Vector{DoubleCovidState}, N::Int)
    n_states = length(states)
    sumS = 0
    sumI1 = zeros(Int,length(first(states).I1))
    sumI2 = zeros(Int,length(first(states).I2))
    sumT1 = zeros(Int,size(first(states).T1))
    sumT2 = zeros(Int,size(first(states).T2))
    for s in states
        sumS += s.S
        sumI1 .+= s.I1
        sumI2 .+= s.I2
        sumT1 .+= s.T1
        sumT2 .+= s.T2
    end
    avgS = floor(Int,sumS/n_states)
    avgI1 = floor.(Int, sumI1./n_states)
    avgI2 = floor.(Int, sumI2./n_states)
    avgR = N - (avgS + sum(avgI1) + sum(avgI2))
    @assert avgR ≥ 0
    avgT1 = floor.(Int, sumT1./n_states)
    avgT2 = floor.(Int, sumT2./n_states)

    return DoubleCovidState(
        avgS,
        avgI1,
        avgI2,
        avgR,
        avgT1,
        avgT2,
        CovidPOMDPs.mean_params(states),
        first(states).prev_action,
        )
end


# This package deals with complex spherical coordinates and complex spherical harmonics. Many of the identities used are not well known, so we pay special attention to tests for these complex functions.

using Test, LinearAlgebra

@testset "Tests for complex transmission direction" begin

    k_eff = rand(-1:0.11:1.0) + rand(-1:0.11:1.0) * im
    incident_wavevector = rand(-1:0.11:1.0,3) + rand(-1:0.11:1.0,3) .* im
    surface_normal = rand(-1:0.11:1.0,3)
    surface_normal = surface_normal / norm(surface_normal)

    # make sure this surface_normal is the outward normal
    if real(dot(surface_normal, incident_wavevector)) > 0
        surface_normal = - surface_normal
    end

    normal_eff = transmission_direction(k_eff, incident_wavevector, surface_normal)

    @test sum(normal_eff .^2) ≈ 1.0 + 0.0im
    @test imag(k_eff * dot(-surface_normal, normal_eff)) > 0

    k_eff = rand(-1:0.1:1.0) + rand(-1:0.1:1.0) * im
    incident_wavevector = rand(-1:0.1:1.0,2) + rand(-1:0.1:1.0,2) .* im
    surface_normal = rand(-1:0.1:1.0,2)
    surface_normal = surface_normal / norm(surface_normal)

    if real(dot(surface_normal, incident_wavevector)) > 0
        surface_normal = - surface_normal
    end

    normal_eff = transmission_direction(k_eff, incident_wavevector, surface_normal)

    @test sum(normal_eff .^2) ≈ 1.0 + 0.0im
    @test imag(k_eff * dot(-surface_normal, normal_eff)) > 0

    k_eff = rand(-1:0.1:1.0) + rand() * im
    k = rand()
    θin = rand(-pi/2.0:0.2:pi/2.0)
    θeff = transmission_angle_wiener(k, k_eff, θin)

    normal2 = [cos(θeff+pi/2.0),sin(θeff+pi/2.0)]

    surface_normal = [0.0,-1.0];
    incident_wavevector = k*[cos(θin+pi/2.0),sin(θin+pi/2.0)]
    normal_eff = transmission_direction(k_eff, incident_wavevector, surface_normal)

    @test norm(normal_eff - normal2) < 1e-10

end

@testset "Tests for complex transmission angle" begin

    surface_normal = [0.0,0.0,-1.0]
    xs = [rand(-1.01:0.1:1.0,3) + rand(-1.01:0.1:1.0,3)*im for i = 1:100];

    # make sure these are incident upon the surface
    xs = [real(dot(surface_normal, x)) > 0 ? -x : x for x in xs];

    @test all(cartesian_to_radial_coordinates(x)[2] == transmission_angle(x, surface_normal) for x in xs)

    surface_normal = [-1.0,0.0]
    xs = [rand(-1.01:0.1:1.0,2) + rand(-1.01:0.1:1.0,2)*im for i = 1:100];
    xs = [real(dot(surface_normal, x)) > 0 ? -x : x for x in xs];

    @test all(cartesian_to_radial_coordinates(x)[2] == transmission_angle(x, surface_normal) for x in xs)

    N = 100
    is = 1:N
    ns = [rand(-1.01:0.1:1.0,2) for i = is];
    ns = ns ./ norm.(ns);

    v_ins = [rand(-1.01:0.1:1.0,2) for i = is]
    v_ins = [real(dot(ns[i], v_ins[i])) > 0 ? -v_ins[i] : v_ins[i] for i in is]

    θ_ins = transmission_angle.(v_ins, ns)

    k_effs = rand(-1.01:0.1:1.0,N) + rand(-1.01:0.1:1.0,N) .* im
    v_effs = k_effs .* transmission_direction.(k_effs, v_ins, ns)
    θ_effs = transmission_angle.(v_effs, ns)

    # Check snells law, i.e. that the components orthogonal to the surface are the same for the two wavevectors
    @test maximum(norm(v_effs[i] - v_ins[i] - ns[i] .* dot(ns[i], v_effs[i] - v_ins[i])) for i in is) < 1e-14
    @test maximum(abs.((k_effs  .* sin.(θ_effs)) .^2 - (norm.(v_ins)  .* sin.(θ_ins)) .^2)) < 1e-10

    # Check that the effective wave attenuates as it moves into the material
    @test all(imag.(dot.(- ns,v_effs)) .> 0)

    # NOTE sum(v_effs[i] .^2) ≈ k_effs[i] ^2, however sqrt(sum(v_effs[i] .^2)) is not necessarily equal to k_effs[i], as the direction of v_effs[i] is chosen so that the wave attenuates when transmitted into the material.

    # When real part of dot(-n,v) is positive then the wave propagates into the material, in which case if real(k_effs) > 0 then we can recover k_effs from the v_effs.
    ns = [ -rand(0.2:0.1:1.0,2) for i = is]
    ns = ns ./ norm.(ns)

    v_ins = [rand(0.2:0.1:1.0,2) for i = is]
    θ_ins = transmission_angle.(v_ins, ns)

    @test all(-pi/2 .< real.(θ_ins) .< pi/2)

    # choose positive real part
    k_effs = rand(0.01:0.1:1.0,N) + rand(-1.1:0.1:1.0,N) .* im;
    v_effs = k_effs .* transmission_direction.(k_effs, v_ins, ns);
    θ_effs = transmission_angle.(v_effs, ns);

    @test maximum(abs.(k_effs  .* sin.(θ_effs) - norm.(v_ins)  .* sin.(θ_ins)) ) < 1e-10

    #angle from x-axis
    θ_ns = [transmission_angle(-n, [-1.0,0.0]) for n in ns];

    @test maximum(norm(v_ins[i] - norm(v_ins[i]) .* [cos(θ_ins[i] + θ_ns[i]), sin(θ_ins[i] + θ_ns[i])]) for i in is) < 1e-10
    @test maximum(norm(v_effs[i] - k_effs[i] .* [cos(θ_effs[i] + θ_ns[i]), sin(θ_effs[i] + θ_ns[i])]) for i in is) < 1e-10

# test for 3-dimensions

    ns = [rand(-1.01:0.1:1.0,3) for i = is];
    ns = ns[is] ./ norm.(ns[is]);

    v_ins = [rand(-1.01:0.1:1.0,3) for i = is];
    v_ins = [real(dot(ns[i], v_ins[i])) > 0 ? -v_ins[i] : v_ins[i] for i in is];

    θ_ins = transmission_angle.(v_ins, ns);

    v_effs = k_effs .* transmission_direction.(k_effs, v_ins, ns);
    θ_effs = transmission_angle.(v_effs, ns);

    # Check snells law, i.e. that the components orthogonal to the surface are the same for the two wavevectors
    @test maximum(norm(v_effs[i] - v_ins[i] - ns[i] .* dot(ns[i], v_effs[i] - v_ins[i])) for i in is) < 1e-14
    @test maximum(abs.((k_effs  .* sin.(θ_effs)) .^2 - (norm.(v_ins)  .* sin.(θ_ins)) .^2)) < 1e-10

    # Check that the effective wave attenuates as it moves into the material
    @test all(imag.(dot.(- ns,v_effs)) .>= 0.0)
end


function fGn_predict_cond(X::AbstractVector{<:Real}, α::AbstractMatrix, u::Integer=1; bias::Bool=true, kwargs...)
    # estimation of bias  β = μy - α * μx
    n, k = size(α)  # n: step of prediction, k: number of samples used for prediction

    gy = (t,n) -> t .+ collect(1:n)  # grid of prediction, starting from t
    # gx = (t,k,u) -> collect((t-k*u+1):u:t)  # grid of samples, ending by t
    gx = (t,k,u) -> reverse(collect(t:-u:(t-k*u+1)))  # grid of samples, ending by t

    β::AbstractVector = if bias
        βm = hcat([X[gy(t,n)] - α*X[gx(t,k,u)] for t=(k*u):(length(X)-n)]...)  # t=(k*u):(length(X)-n)
        mean(βm, dims=2)[:]
    else
        zeros(n)
    end

    # conditional mean and covariance
    μ::AbstractVector = α * X[gx(length(X),k,u)] + β
#     Σ::AbstractMatrix = Σyy - Σyx * iΣxx * Σyx'

    return (μ=μ, β=β)
end

function fGn_predict_replict(X::AbstractVector{<:Real}, hurst::Real, u::Integer=1, k::Integer=1; kwargs...)
    N = length(X)
    μ = sign(hurst-0.5) * mean(X[end:-u:(end-k*u+1)])
#     μ = sign(hurst-0.5) * mean(X[end-k+1:end])

    return (μ=μ,)
end

"""
MLE of fGn model and prediction by replication.

# Args
- s: sub window size
- l: length of decorrelation
- k: number of samples used for prediction
"""
function fGn_MLE_estim_predict_replict(X::AbstractVector{<:Real}, d::Integer, s::Integer, l::Integer, k::Integer=1; kwargs...)
    est = fGn_MLE_estim(X, d, s, l; kwargs...)
    μ = sign(est.hurst-0.5) * mean(X[end-k+1:end])

    return (hurst=est.hurst, σ=est.σ, μ=μ)
end

"""
MLE of fGn model and prediction by conditional statistics.

# Args
- s: sub window size
- l: length of decorrelation
- k: number of samples used for prediction
- u: downsampling factor for samples
- n: step of prediction
"""
function fGn_MLE_estim_predict_cond(X::AbstractVector{<:Real}, d::Integer, s::Integer, l::Integer, k::Integer, u::Integer=1, n::Integer=d; bias::Bool=false, kwargs...)
    @assert 0 < k && 0 < n
    # @assert n+k*u <= length(X)

    # estimation of hurst and volatility
    est = fGn_MLE_estim(X, d, s, l; kwargs...)

    P = FractionalGaussianNoise(est.hurst, d)  # fGn with estimated hurst at step d
    gy = t -> t .+ collect(1:n)  # grid of prediction, starting from t
    gx = t -> collect(reverse(t:-u:(t-k*u+1)))  # grid of samples, ending by t

    # covariance matrices
    Σyy = covmat(P, gy(0))
    Σyx = covmat(P, gy(0), gx(0))
    Σxx = covmat(P, gx(0))
    iΣxx = pinv(Matrix(Σxx))
    Σ = Σyy - Σyx * iΣxx * Σyx'

    α::AbstractMatrix = Σyx * iΣxx  # kernels of conditional mean
    # estimation of bias  β = μy - α * μx
    β::AbstractVector = if bias
        # βm = hcat([X[gy(t)] - α*X[gx(t)] for t=1:length(X) if gx(t)[1]>0 && gy(t)[end]<=length(X)]...)  # t=(k*u):(length(X)-n)
        βm = hcat([X[gy(t)] - α*X[gx(t)] for t=k*u:length(X)-n]...)  # t=(k*u):(length(X)-n)
        mean(βm, dims=2)[:]
    else
        zeros(n)
    end

    # conditional mean and covariance
    μ = α * X[gx(length(X))] + β

    return (hurst=est.hurst, σ=est.σ, μ=μ, Σ=Σ, α=α, β=β)
end

"""
# Args
- X: each row corresponds to a symbol
"""
function fGn_MLE_estim_predict_cond(X::AbstractMatrix{<:Real}, d::Integer, s::Integer, l::Integer, k::Integer, u::Integer=1, n::Integer=d; bias::Bool=false, kwargs...)
    @assert 0 < k && 0 < n
    # @assert n+k*u <= length(X)

    # estimation of hurst and volatility
    res = []  # results of estimation
    for r=1:size(X,1)
        push!(res, fGn_MLE_estim(X[r,:], d, s, l; kwargs...))
    end

    Hs = [x.hurst for x in res]  # vector of Hurst
    # Hμ = mean(Hs)
    # Hσ = std(Hs)
    Vs = [x.σ for x in res]  # vector of volatility
    # Vμ = mean(Vs)
    # Vσ = std(Vs)

    # hurst = mean(Hs)
    hurst = median(Hs)

    P = FractionalGaussianNoise(hurst, d)  # fGn with estimated hurst at step d
    gy = t -> t .+ collect(1:n)  # grid of prediction, starting from t
    gx = t -> collect(reverse(t:-u:(t-k*u+1)))

    # covariance matrices
    Σyy = covmat(P, gy(0))
    Σyx = covmat(P, gy(0), gx(0))
    Σxx = covmat(P, gx(0))
    iΣxx = pinv(Matrix(Σxx))
    Σ = Σyy - Σyx * iΣxx * Σyx'

    α::AbstractMatrix = Σyx * iΣxx  # kernels of conditional mean
    # estimation of bias  β = μy - α * μx
    β::AbstractMatrix = if bias
        βm = [X[:,gy(t)]' - α*X[:,gx(t)]' for t=k*u:size(X,2)-n]  # t=(k*u):(length(X)-n)
        mean(βm)
    else
        zeros(size(X,1),n)
    end

    # conditional mean and covariance
    μs = X[:,gx(size(X,2))] * α' + β
    # for r=1:size(X,1)
    #     push!(μs, α*X[r,gx(size(X,2))] + β)
    # end

    return (hurst=Hs, σ=Vs, μ=μs, Σ=Σ, α=α, β=β)
end

function fGn_MLE_estim_predict_old(X::AbstractVector{<:Real}, d::Integer, s::Integer, l::Integer, k::Integer, u::Integer, n::Integer; dmode::Symbol=:regular, kwargs...)
    N = length(X)  # number of samples

    # when sub window equals to rolling window then `rolling_vectorize()` has no effect
    est = fGn_MLE_estim(X, d, s, l; kwargs...)

    # convention of time arrow: from left to right
    sgrid::AbstractVector{<:Integer} = Int[]  # grid of historical samples
    S::AbstractVector{<:Real} = Real[]  # value of historical samples

    if dmode == :logscale  # logscale downsampling
        tidx = findall(reverse(logtrain(N, N÷u)))
        # select the last k samples
        sgrid = (k==0 || k>=length(tidx)) ? tidx : tidx[end-k+1:end]
        S = X[sgrid]
    else # regular downsampling
        Sm = vec2mat(X, u, keep=:tail)
        S0 = mean(Sm, dims=1)[:]
        # S0 = Sm[1,:]
        S = (k==0 || k>=length(S0)) ? S0 : S0[end-k+1:end]
        sgrid = u*(1:length(S))
    end

    # conditional mean and covariance
    pgrid = sgrid[end] .+ (1:n)  # grid of prediction
    prd = cond_mean_cov(FractionalGaussianNoise(est.hurst, d), pgrid, sgrid, S)
    C = prd.C[end,:]

    return (hurst=est.hurst, σ=est.σ, μ=prd.μ, C=C, Σ=prd.Σ)
    # return hurst, σ, μ, σ^2*Σ
end

"""
Power-law estimator and predictor.

# Args
- X: input matrix. The first row is the observation of fBm and the others are fGn
- lags: integer time lags (increment step) used to compute each component of `X` starting from the second row. Values in `lags` must be all distinct.
- s: time lag for prediction
- d: (average) downsampling factor

# Returns
- H, σ, (μr, Vr), (μy, Vy): estimation of hurst, volatility, conditional mean and covariance of fGn and fBm.

# Notes
- This function is intended to be used with `rolling_apply` where `X` is some observation on a time window. - The time arrow is on the second dimension (i.e. horizontal) from left to right.
- μr and

"""
function powlaw_estim_predict(X::AbstractMatrix{<:Real}, lags::AbstractVector{<:Integer}, s::Integer, d::Integer, k::Integer=0; kwargs...)
    # @assert s in lags
    # @assert "all values in lags must be distinct"

    N = size(X, 2)  # number of rows
    Y = view(X, 1, :)  # first row is an observation of fBm
    R = view(X, 2:size(X,1), :)  # the other rows are observations of fGn, corresponding to `lags`
    sidx = findall(lags.==s)[1]  # index of `s` in `lags`

    H, σ, opm = powlaw_estim(R, lags; kwargs...)

    # convention of time arrow: from left to right
    # method 1: logscale downsampling
    tidx0 = findall(reverse(logtrain(N, N÷d)))
    # @assert k <= length(tidx0)
    tidx = k>0 ? tidx0[end-k+1:end] : tidx0  # select the last k samples
    S = R[sidx,tidx]

    # # method 2: regular downsampling with averaging
    # S0 = mean(vec2mat(R[sidx,:], d, keep=:tail), dims=1)[:]
    # # S0 = transpose(vec2mat(R[sidx,:], d, keep=:tail))[1,:]
    # S = k>0 ? S0[end-k+1:end] : S0
    # tidx = d*(1:size(S,1))

    # conditional mean and covariance
    G = tidx[end]+1:tidx[end]+s  # grid of prediction
    # fGn: `S` is the observation of fGn corresponding to `s`
    μr, Σr = cond_mean_cov(FractionalGaussianNoise(H, s), G, tidx, S)
    # fBm: `Y` is the observation of fBm
    μy, Σy = cond_mean_cov(FractionalBrownianMotion(H), G, tidx, Y[tidx])

    return H, σ, (μr, σ^2*Σr), (μy, σ^2*Σy)
end

# """
# EXPERIMENTAL: prediction by recursive conditional mean
# No visible improvements over classical conditional mean for long range prediction.
# """
# function powlaw_estim_predict(X::AbstractMatrix{<:Real}, lags::AbstractVector{<:Integer}, s::Integer, k::Integer; kwargs...)
#     # @assert k>0 && d>0
#     H, σ, opm = FracFin.powlaw_estim(X, lags; kwargs...)
#     l = size(X,2)÷4
#     Cv = cond_mean_coeff(FractionalGaussianNoise(H, lags[s]), k, l)
#     μc = Cv * X[s, end-l+1:end]
#     return H, σ, μc
# end


using Distances
function  e1(StartPoint, GradientVolume, StepSize)
    #  E1 is a function which performs one step of the Euler ray tracing
    #
    #   EndPoint = E1(StartPoint, GradientVolume, StepSize);
    #
    #    inputs :
    #        StartPoint: 2D or 3D location in vectorfield
    #        GradientVolume: Vectorfield
    #        Stepsize : The stepsize
    #
    #   outputs :
    #        EndPoint : The new location (zero if outside image)
    #
    # Function is written by D.Kroon University of Twente (Oct 2010)

    # Linear interpolation of current location
    xBas=[0 0 1 1].+int(floor(StartPoint[1]))
    yBas=[0 1 0 1].+int(floor(StartPoint[2]))
    xCom=StartPoint[1].-int(floor(StartPoint[1]) )
    yCom=StartPoint[2].-int(floor(StartPoint[2]) )

    # Linear interpolation percentages.
    perc=[(1-xCom) * (1-yCom);
          (1-xCom) * yCom    ;
          xCom     * (1-yCom);
          xCom     * yCom    ]

    # Split in Gradient Volumes
    Gx=GradientVolume[:,:,1]
    Gy=GradientVolume[:,:,2]

    xBas=min(maximum(xBas,1),size(Gx,1))
    yBas=min(maximum(yBas,1),size(Gx,2))

    ind=sub2ind(size(Gx),xBas[:],yBas[:])
    gradient=[sum(Gx[ind].*perc) sum(Gy[ind].*perc)]
    gradient=gradient./(sqrt(sum(gradient.^2))+eps())

    # Set a step in the direction of the gradient.
    EndPoint=StartPoint[:]-StepSize*gradient[:]
    check = (EndPoint[1]<1)|(EndPoint[2]<1)|(EndPoint[1]>size(Gx,1))|(EndPoint[2]>size(Gx,2))
    if check
        EndPoint=[0,0]
        print("Finished")
    end
    return EndPoint
end

# else
#     %Linear interpolation of current location
#     xBas=[0 0 0 0 1 1 1 1]+floor(StartPoint(1));
#     yBas=[0 0 1 1 0 0 1 1]+floor(StartPoint(2));
#     zBas=[0 1 0 1 0 1 0 1]+floor(StartPoint(3));
#     xCom=StartPoint(1)-floor(StartPoint(1));
#     yCom=StartPoint(2)-floor(StartPoint(2));
#     zCom=StartPoint(3)-floor(StartPoint(3));

#     % Linear interpolation percentages.
#     perc=[(1-xCom) * (1-yCom) * (1-zCom);
#           (1-xCom) * (1-yCom) * zCom;
#           (1-xCom) * yCom     * (1-zCom);
#           (1-xCom) * yCom     * zCom;
#           xCom     * (1-yCom) * (1-zCom);
#           xCom     * (1-yCom) * zCom;
#           xCom     * yCom     * (1-zCom);
#           xCom     * yCom     * zCom;];

#     % Split in Gradient Volumes
#     Gx=GradientVolume(:,:,:,1);
#     Gy=GradientVolume(:,:,:,2);
#     Gz=GradientVolume(:,:,:,3);

#     xBas=min(max(xBas,1),size(Gx,1));
#     yBas=min(max(yBas,1),size(Gx,2));
#     zBas=min(max(zBas,1),size(Gx,3));

#     ind=sub2ind(size(Gx),xBas(:),yBas(:),zBas(:));
#     gradient=[sum(Gx(ind).*perc) sum(Gy(ind).*perc) sum(Gz(ind).*perc)];
#     gradient=gradient./(sqrt(sum(gradient.^2))+eps);

#     % Set a step in the direction of the gradient.
#     EndPoint=StartPoint-StepSize*gradient(:);
#     check=(EndPoint(1)<1)||(EndPoint(2)<1)||(EndPoint(3)<1)||(EndPoint(1)>size(Gx,1))||(EndPoint(2)>size(Gx,2))||(EndPoint(3)>size(Gx,3));
#     if(check), EndPoint=[0;0;0]; end
# end


#############################################################
#Data Import and Preprocesing

using Knet
include(Knet.dir("data", "cifar.jl"))
include(Knet.dir("data", "mnist.jl"))
function import_data(dataset::String)
	if dataset == "CIFAR10"
		smallimgs, labels, smallimgstest, labelstest = cifar10()
	    smallimgs = reshape(smallimgs, 32^2 * 3, 50000)
	    smallimgstest = reshape(smallimgstest, 32^2 * 3, 10000)
	elseif dataset == "MNIST"
		smallimgs, labels, smallimgstest, labelstest = mnist()
	    smallimgs = reshape(smallimgs, 28^2, 60000)
	    smallimgstest = reshape(smallimgstest, 28^2, 10000)
	end
	smallimgs, labels .- 1, smallimgstest, labelstest .- 1
end

function subtractmean!(data)
        m = mean(data, dims = 2)
        d, n = size(data)
        for i in 0:n-1
                BLAS.axpy!(-1., m, 1:d, data, i*d + 1: (i+1) * d)
        end
end
function subtractmean(data)
        m = deepcopy(data)
        subtractmean!(m)
        m
end

function normalizelinewise!(data)
    max = maximum(data, dims = 2)
    d, n = size(data)
    for i in 1:d
        (max[i] != 0.) && (data[i,:] ./= max[i])
    end
    subtractmean!(data)
end

function standardizelinewise!(data)
    means = mean(data, dims = 2)
    stds = std(data, dims = 2)
    d, n = size(data)
    for i in 1:d
        data[i,:] .-= means[i]
        (stds[i] != 0.) && (data[i,:] ./= stds[i])
    end
end

#scale data between [-1,1]
function rescaledata!(data)
        absmax = maximum(abs(data))
        scale!(data, 1. / absmax)
end

# data is dxn array of n d-dimensional random vectors with mean 0
function whiten(data; method = :f_ZCA)
        f = svdfact(data)
        eval(method)(f) * sqrt(size(data, 2) - 1)
end

f_ZCA(f::LinearAlgebra.SVD) = f[:U] * f[:Vt]
f_PCA(f::LinearAlgebra.SVD) = f[:Vt]


module ExtrapTests

using Base.Test, DualNumbers
using Interpolations

f(x) = sin((x-3)*2pi/9 - 1)
xmax = 10
A = Float64[f(x) for x in 1:xmax]

itpg = interpolate(A, BSpline(Linear()), OnGrid())

etpg = extrapolate(itpg, Flat())

@test etpg[-3] == etpg[-4.5] == etpg[0.9] == etpg[1.0] == A[1]
@test etpg[10.1] == etpg[11] == etpg[148.298452] == A[end]

etpf = @inferred(extrapolate(itpg, NaN))

@test @inferred(size(etpf)) == (xmax,)
@test isnan(@inferred(getindex(etpf, -2.5)))
@test isnan(etpf[0.999])
@test_approx_eq @inferred(getindex(etpf, 1)) f(1)
@test_approx_eq etpf[10] f(10)
@test isnan(@inferred(getindex(etpf,10.001)))

@test etpf[2.5,1] == etpf[2.5]   # for show method
@test_throws BoundsError etpf[2.5,2]
@test_throws ErrorException etpf[2.5,2,1]  # this will probably become a BoundsError someday

@test isa(@inferred(getindex(etpf, dual(-2.5,1))), Dual)

etpl = extrapolate(itpg, Linear())
k_lo = A[2] - A[1]
x_lo = -3.2
@test_approx_eq etpl[x_lo] A[1]+k_lo*(x_lo-1)
k_hi = A[end] - A[end-1]
x_hi = xmax + 5.7
@test_approx_eq etpl[x_hi] A[end]+k_hi*(x_hi-xmax)

xmax, ymax = 8,8
g(x, y) = (x^2 + 3x - 8) * (-2y^2 + y + 1)

itp2g = interpolate(Float64[g(x,y) for x in 1:xmax, y in 1:ymax], (BSpline(Quadratic(Free())), BSpline(Linear())), OnGrid())
etp2g = extrapolate(itp2g, (Linear(), Flat()))

@test_approx_eq @inferred(getindex(etp2g, -.5, 4)) itp2g[1,4]-1.5*epsilon(etp2g[dual(1,1),4])
@test_approx_eq @inferred(getindex(etp2g, 5, 100)) itp2g[5,ymax]

etp2ud = extrapolate(itp2g, ((Linear(), Flat()), Flat()))
@test_approx_eq @inferred(getindex(etp2ud, -.5, 4)) itp2g[1,4] - 1.5*epsilon(etp2g[dual(1,1),4])
@test @inferred(getindex(etp2ud, 5, -4)) == etp2ud[5,1]
@test @inferred(getindex(etp2ud, 100, 4)) == etp2ud[8,4]
@test @inferred(getindex(etp2ud, -.5, 100)) == itp2g[1,8] - 1.5 * epsilon(etp2g[dual(1,1),8])

etp2ll = extrapolate(itp2g, Linear())
@test_approx_eq @inferred(getindex(etp2ll, -.5, 100)) itp2g[1,8]-1.5*epsilon(etp2ll[dual(1,1),8]) + (100-8) * epsilon(etp2ll[1,dual(8,1)])

# Allow element types that don't support conversion to Int (#87):
etp87g = extrapolate(interpolate([1.0im, 2.0im, 3.0im], BSpline(Linear()), OnGrid()), 0.0im)
@test @inferred(getindex(etp87g, 1)) == 1.0im
@test @inferred(getindex(etp87g, 1.5)) == 1.5im
@test @inferred(getindex(etp87g, 0.75)) == 0.0im
@test @inferred(getindex(etp87g, 3.25)) == 0.0im

etp87c = extrapolate(interpolate([1.0im, 2.0im, 3.0im], BSpline(Linear()), OnCell()), 0.0im)
@test @inferred(getindex(etp87c, 1)) == 1.0im
@test @inferred(getindex(etp87c, 1.5)) == 1.5im
@test @inferred(getindex(etp87c, 0.75)) == 0.75im
@test @inferred(getindex(etp87c, 3.25)) == 3.25im
@test @inferred(getindex(etp87g, 0)) == 0.0im
@test @inferred(getindex(etp87g, 3.7)) == 0.0im

# Make sure it works with Gridded too
etp100g = extrapolate(interpolate(([10;20],),[100;110], Gridded(Linear())), Flat())
@test @inferred(getindex(etp100g, 5)) == 100
@test @inferred(getindex(etp100g, 15)) == 105
@test @inferred(getindex(etp100g, 25)) == 110
end

include("type-stability.jl")


# Data source for london tube https://www.whatdotheyknow.com/request/distance_between_adjacent_underg
using CSV, DataFrames

df = CSV.read("data/london/london.csv", DataFrame)

id = 1
stations = Dict()
edges = Dict()
for row in Tables.namedtupleiterator(df)
    station_a = getindex(row, Symbol("Station from (A)"))
    if !haskey(stations, station_a)
        push!(stations, station_a => id)
        global id += 1
    end
    station_b = getindex(row, Symbol("Station to (B)"))
    if !haskey(stations, station_b)
        push!(stations, station_b => id)
        global id += 1
    end

    (from,to) = (stations[station_a], stations[station_b])
    distance = getindex(row, Symbol("Distance (Kms)"))
    if !haskey(edges, (from,to)) && !haskey(edges, (to,from))
        push!(edges, (from,to) => round(Int, distance * 1000)) # in meters
    end
end

n, E, C = id - 1, collect(keys(edges)), collect(values(edges))

output = open("data/london/london.gr", "w")
write(output, "c London Tube Stations Graph with distances\n")
write(output, "c Source: https://www.whatdotheyknow.com/request/distance_between_adjacent_underg\n")
write(output, "c\n")
write(output, "p st ", string(n), " ", string(length(E)), "\n")
write(output, "c\n")
for (i, (a, b)) in enumerate(E)
    write(output, "a ", string(a), " ", string(b), " ", string(C[i]), "\n")
end
close(output)


using CayleyDickson
using Base.Test: @test, @test_throws

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigInt()
    l = Exo1SplitQuaternion(Exo1Real(a, b))
    r = Exo1SplitQuaternion(a, 0, 0, 0, b, 0, 0, 0)
    l == r
end

@test begin
    a = CayleyDickson.randomBigFloat()
    isreal(Exo1SplitQuaternion(a))
end

@test begin
    a = 1
    !isreal(Exo1SplitQuaternion(a, a, a, a))
end

@test begin
    z = random(Exo1SplitQuaternion{Int})
    z == +(z)
end

@test begin
    a = CayleyDickson.randomBigFloat()
    real(Exo1SplitQuaternion(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, Exo1SplitQuaternion(1.0, 2, 3, 4, 5, 6, 7, 8))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, s: 3.0, is: 4.0, a: 5.0, ia: 6.0, sa: 7.0, (is)a: 8.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1SplitQuaternion(1.0, 2, 3, 4, 5, 6, 7))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, s: 3.0, is: 4.0, a: 5.0, ia: 6.0, sa: 7.0, (is)a: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1SplitQuaternion(1.0, 2, 3, 4, 5, 6))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, s: 3.0, is: 4.0, a: 5.0, ia: 6.0, sa: 0.0, (is)a: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1SplitQuaternion(1.0, 2, 3, 4, 5))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, s: 3.0, is: 4.0, a: 5.0, ia: 0.0, sa: 0.0, (is)a: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1SplitQuaternion(1.0, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, s: 3.0, is: 4.0, a: 0.0, ia: 0.0, sa: 0.0, (is)a: 0.0]"
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    d = CayleyDickson.randomBigFloat()
    l = Exo1SplitQuaternion(SplitQuaternion(a, b, c, d))
    r = Exo1SplitQuaternion(a, b, c, d)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    l = Exo1SplitQuaternion(Binion(a, b))
    r = Exo1SplitQuaternion(a, b)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    l = Exo1SplitQuaternion(a, b, c)
    r = Exo1SplitQuaternion(a, b, c, 0)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigFloat})
    y = random(Exo1SplitQuaternion{BigFloat})
    x != y
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo1SplitQuaternion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = random(SplitQuaternion{BigInt})
    z = random(Exo1SplitQuaternion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo1SplitQuaternion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    a = random(SplitQuaternion{BigInt})
    z = random(Exo1SplitQuaternion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = x * y
    r = y * x
    l != r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    z = random(Exo1SplitQuaternion{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    z = random(Exo1SplitQuaternion{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l != r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = alternatorL(x, y)
    r = zero(Exo1SplitQuaternion{BigInt})
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = alternatorR(x, y)
    r = zero(Exo1SplitQuaternion{BigInt})
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    conj(conj(x)) == x
end

@test begin
    x = random(Exo1SplitQuaternion{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = conj(x * y)
    r = conj(y) * conj(x)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = conj(x + y)
    r = conj(x) + conj(y)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = conj(x - y)
    r = conj(x) - conj(y)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{Rational{BigInt}})
    y = random(Exo1SplitQuaternion{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{Rational{BigInt}})
    y = random(Exo1SplitQuaternion{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{Rational{BigInt}})
    y = random(Exo1SplitQuaternion{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    z = random(Exo1SplitQuaternion{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo1SplitQuaternion{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    a = random(SplitQuaternion{BigInt})
    x = random(Exo1SplitQuaternion{BigInt})
    l = commutator(a, x)
    r = zero(Exo1SplitQuaternion{BigInt})
    l != r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    z = random(Exo1SplitQuaternion{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    z = random(Exo1SplitQuaternion{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    z = random(Exo1SplitQuaternion{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    quadrance(x) > big(0) || quadrance(x) < big(0) || quadrance(x) == big(0)
end

@test begin
    x = random(Exo1SplitQuaternion{BigInt})
    y = random(Exo1SplitQuaternion{BigInt})
    l = quadrance(x * y)
    r = quadrance(x) * quadrance(y)
    l == r
end

@test_throws ErrorException begin
    inv(Exo1SplitQuaternion(0, 0.0))
end

@test_throws ErrorException begin
    random(Exo1SplitQuaternion{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(Exo1SplitQuaternion{Int})
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo1SplitQuaternion{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo1SplitQuaternion{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(Exo1SplitQuaternion{BigInt})
    z == z + zero(Exo1SplitQuaternion{BigInt})
end

@test begin
    z = random(Exo1SplitQuaternion{BigInt})
    z == z * one(Exo1SplitQuaternion{BigInt})
end

@test begin
    z = random(Exo1SplitQuaternion{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(Exo1SplitQuaternion{BigInt})
    z == z * one(z)
end

@test begin
    widen(Exo1SplitQuaternion{Int32}) == Exo1SplitQuaternion{Int64}
end

@test begin
    float(Exo1SplitQuaternion{Int64}) == Exo1SplitQuaternion{Float64}
end

@test begin
    float(Exo1SplitQuaternion{Float64}) == Exo1SplitQuaternion{Float64}
end


using NCDatasets
using Formatting

include("./SSM_z_res.jl")

fn_o = ARGS[1]

Nz = length(zs_SSM) - 1
N_zs = length(zs_SSM)

Dataset(fn_o, "c") do ds

    defDim(ds, "Nz", Nz)
    defDim(ds, "zs", N_zs)

    defVar(ds, "zs", zs_SSM, ("zs",))

end


using PyPlot
using PyCall
using DelimitedFiles
using Statistics
mpl = pyimport("tikzplotlib")

res = readdlm("result.txt")
d = Dict()
for i = 1:size(res,1)
    if haskey(d, res[i,1])
        push!(d[res[i,1]], res[i,2])
    else
        d[res[i,1]] = [res[i,2]]
    end
end

arr = [0.0,0.002,0.004,0.006,0.008,0.01,0.03,0.05,0.06, 0.07, 0.075,0.08,0.09,0.1]
s = zeros(length(arr))
m = zeros(length(arr))

for (ix,k) in enumerate(arr)
    s[ix] = std(d[k])
    m[ix] = mean(d[k])
end

semilogy(arr, m)
fill_between(arr, m - s, m+s, alpha=0.5)
xlabel("\$\\sigma_{\\mathrm{noise}}\$")
ylabel("Error")
mpl.save("poroerror.tex")