## Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra
using StaticArrays
using HybridArrays

# %%code cell

function finklet(j, l)
    # -------- set filters
    jrad = 7-j
    dθ = π/8        # 8 angular bins hardwired
    θ_l = dθ*l
    # -------- define coordinates
    nx = 256
    xbox = LinRange(-nx/2, nx/2-1 , nx)
    # make a 256x256 grid of X
    sx = xbox' .* ones(nx)
    sy = ones(nx)' .* xbox
    r  = sqrt.((sx).^2 + (sy).^2)
    θ  = mod.(atan.(sy, sx).+π .-θ_l,2*π)
    nozeros = r .> 0
    logr = log2.(r[nozeros])
    r[nozeros] = logr
    # -------- in Fourier plane, envelope of psi_j,l
    mask = (abs.(θ.-π).<= dθ) .& (abs.(r.-jrad) .<= 1)
    # -------- angular part
    ang = cos.((θ.-π).*4)
    # -------- radial part
    rad = cos.((r.-jrad).*π./2)
    psi = mask.*ang.*rad             #mask times angular part times radial part
    return psi
end

function fink_filter_bank(J,L)
    fink_filter = Array{Float64, 4}(undef, 256, 256, J, L)
    for l = 1:L
        for j = 1:J
            fink_filter[:,:,j,l]=fftshift(finklet(j-1,l-1))
        end
    end
    return fink_filter
end

function had!(A,B)
    m,n = size(A)
    @assert (m,n) == size(B)
    for j in 1:n
       for i in 1:m
         @inbounds A[i,j] *= B[i,j]
       end
    end
    return A
end

fink_filter_set = fink_filter_bank(8,8)

test_img = zeros(256,256)
copyto!(test_img,fink_filter_set[:,:,1,3])

function DHC(image, filter_set; norm_on = 1, coeff_12_on =1, coeff_20_on = 1)
    FFTW.set_num_threads(2)
    (Nx, Ny) = size(image)

    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []

    ## Coeff at (1,0)
    #takes care of image normalization
    S0 = zeros(2)
    if norm_on == 1
        S0[1] = mean(image)
        norm_im = image.-S0[1]
        S0[2] = sum(abs2.(norm_im))/(Nx*Ny)
        norm_im = norm_im./sqrt(Nx*Ny*S0[2])
    else
        norm_im = image
    end

    append!(out_coeff,S0)

    ##Coeff at (1,1)
    #store the fft of images and im*filters for later
    im_fd_0 = fft(norm_im)

    S1 = zeros((J,L))
    im_fd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    @views for l = 1:L
        for j = 1:J
            im_fd_0_1[:,:,j,l] .= im_fd_0
        end
    end

    if coeff_20_on == 1
        im_rd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    end

    @views for l = 1:L
        for j = 1:J
            had!(im_fd_0_1[:,:,j,l],filter_set[:,:,j,l]) #wavelet already in fft domain not shifted
            S1[j,l]+=sum(abs2.(im_fd_0_1[:,:,j,l])) #normalization choice arb to make order unity
            if coeff_20_on == 1
                im_rd_0_1[:,:,j,l] .= ifft(im_fd_0_1[:,:,j,l])
            end
        end
    end
    append!(out_coeff,S1)

    ##Coeff at (2,0)
    if coeff_20_on == 1
        S20 = zeros((J, L, J, L))

        Atmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,Nx,Ny))
        Btmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,Nx,Ny))
        interm = 0
        @views for l2 = 1:L
            for j2 = 1:J
                copyto!(Btmp, im_rd_0_1[:,:,j2,l2])
                for l1 = 1:L
                    for j1  = 1:J
                        copyto!(Atmp, im_rd_0_1[:,:,j1,l1])
                        S20[j1,l1,j2,l2] += sum(abs.(Atmp.*Btmp))
                    end
                end
            end
        end

        append!(out_coeff,S20)
    end

    ##Coeff at (1,2) level
    if coeff_12_on == 1
        S12 = Array{Float64,4}(undef, J, J, L, L)
        for j1 = 1:J
            for l1 = 1:L
                not_rot_im = fftshift(im_fd_0_1[j1,l1,:,:])
                for l2 = 1:L
                    for j2 = 1:J
                        S12[j1,j2,l1,l2] += sum(abs.(not_rot_im.*fftshift(im_fd_0_1[j2,l2,:,:])))
                    end
                end
            end
        end
    append!(out_coeff,S12)
    end
    return out_coeff
end

function dual_conj(A,B)
    @. abs(A.*B)
end

temp = DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

@time DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Profile.clear()
@profile DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.samples = 5
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 120
@benchmark DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

##Trying out static speed ups

function dual_conj_pre(A,B)
    Atmp = SMatrix{256,256}(A)
    Btmp = SMatrix{256,256}(B)
    @. abs(Atmp.*Btmp)
end

@time dual_conj_pre(test_img,test_img)

test_img_s = @SMatrix rand(4,4)

function mult_wrapper(A,B)
    sum(abs.(A.*B))
end

@time mult_wrapper(test_img_s,test_img_s)

test_img_2 = rand(4,4)

@time mult_wrapper(test_img_2,test_img_2)

A = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(randn(256,256))

mult_wrapper(A,A)

@benchmark mult_wrapper(A,A)

A = randn(256,256)

@benchmark mult_wrapper(A,A)

function mult_wrapper(A,B)
    A = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(randn(256,256))
    B = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(randn(256,256))
    sum(abs.(A.*B))
end

@benchmark mult_wrapper(A,A)

Btmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
copyto!(Btmp, fink_filter_set[:,:,5,5])

## Suppose I was willing to reshape in DHC_ISO

function structure_test(J,L)
    S20 = zeros((J, L, J, L))
    @views for l2 = 1:L
        for j2 = 1:J
            for l1 = 1:L
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = rand()
                end
            end
        end
    end
end

@benchmark structure_test(8,8)

function structure_test(J,L)
    S20 = zeros(J*L*J*L)
    @views for l2 = 1:L
        for j2 = 1:J
            for l1 = 1:L
                for j1  = 1:J
                    S20[(j1-1)*J+(l1-1)*L+(j2-1)*J+l2] = rand()
                end
            end
        end
    end
end

@benchmark structure_test(8,8)

function structure_test(J,L)
    S20 = zeros(J*L*J*L)
    Atmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
    Btmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
    @views for l2 = 1:L
        for j2 = 1:J
            copyto!(Btmp, fink_filter_set[:,:,j2,l2])
            for l1 = 1:L
                for j1  = 1:J
                    copyto!(Atmp, fink_filter_set[:,:,j1,l1])
                    S20[(j1-1)*J+(l1-1)*L+(j2-1)*J+l2] = sum(abs.(Atmp.*Btmp))
                end
            end
        end
    end
end

structure_test(8,8)

@time structure_test(8,8)

@benchmark structure_test(8,8)

function structure_test(J,L)
    S20 = zeros(J,L,J,L)
    Atmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
    Btmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
    @views for l2 = 1:L
        for j2 = 1:J
            copyto!(Btmp, fink_filter_set[:,:,j2,l2])
            for l1 = 1:L
                for j1  = 1:J
                    copyto!(Atmp, fink_filter_set[:,:,j1,l1])
                    S20[j1,l1,j2,l2] = sum(abs.(Atmp.*Btmp))
                end
            end
        end
    end
end

structure_test(8,8)

@benchmark structure_test(8,8)

## Using BLAS

function structure_test(J,L,test_img_1)
    S20 = zeros((J, L, J, L))
    @views for l2 = 1:L
        for j2 = 1:J
            for l1 = 1:L
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = sum(abs.(test_img_1))
                end
            end
        end
    end
end

structure_test(8,8,rand(256,256))

@benchmark structure_test(8,8,rand(256,256))

function structure_test(J,L,test_img_1)
    S20 = zeros((J, L, J, L))
    @views for l2 = 1:L
        for j2 = 1:J
            for l1 = 1:L
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = BLAS.asum(test_img_1)
                end
            end
        end
    end
end

structure_test(8,8,rand(256,256))

@benchmark structure_test(8,8,rand(256,256))

# %% codecell

function DHC(image, filter_set; norm_on = 1, coeff_12_on =1, coeff_20_on = 1)
    FFTW.set_num_threads(2)
    (Nx, Ny) = size(image)

    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []

    ## Coeff at (1,0)
    #takes care of image normalization
    S0 = zeros(2)
    if norm_on == 1
        S0[1] = mean(image)
        norm_im = image.-S0[1]
        S0[2] = sum(abs2.(norm_im))/(Nx*Ny)
        norm_im = norm_im./sqrt(Nx*Ny*S0[2])
    else
        norm_im = image
    end

    append!(out_coeff,S0)

    ##Coeff at (1,1)
    #store the fft of images and im*filters for later
    im_fd_0 = fft(norm_im)

    S1 = zeros((J,L))
    im_fd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    @views for l = 1:L
        for j = 1:J
            im_fd_0_1[:,:,j,l] .= im_fd_0
        end
    end

    if coeff_20_on == 1
        im_rd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    end

    @views for l = 1:L
        for j = 1:J
            had!(im_fd_0_1[:,:,j,l],filter_set[:,:,j,l]) #wavelet already in fft domain not shifted
            S1[j,l]+=sum(abs2.(im_fd_0_1[:,:,j,l])) #normalization choice arb to make order unity
            if coeff_20_on == 1
                im_rd_0_1[:,:,j,l] .= ifft(im_fd_0_1[:,:,j,l])
            end
        end
    end
    append!(out_coeff,S1)

    ##Coeff at (2,0)
    if coeff_20_on == 1
        S20 = zeros((J, L, J, L))

        Atmp = zeros(ComplexF64,Nx,Ny)
        Btmp = zeros(ComplexF64,Nx,Ny)
        interm = 0
        @views for l2 = 1:L
            for j2 = 1:J
                copyto!(Btmp, im_rd_0_1[:,:,j2,l2])
                for l1 = 1:L
                    for j1  = 1:J
                        copyto!(Atmp, im_rd_0_1[:,:,j1,l1])
                        S20[j1,l1,j2,l2] += BLAS.asum(Atmp.*Btmp)
                    end
                end
            end
        end

        append!(out_coeff,S20)
    end

    ##Coeff at (1,2) level
    if coeff_12_on == 1
        S12 = Array{Float64,4}(undef, J, J, L, L)
        for j1 = 1:J
            for l1 = 1:L
                not_rot_im = fftshift(im_fd_0_1[j1,l1,:,:])
                for l2 = 1:L
                    for j2 = 1:J
                        S12[j1,j2,l1,l2] += sum(abs.(not_rot_im.*fftshift(im_fd_0_1[j2,l2,:,:])))
                    end
                end
            end
        end
    append!(out_coeff,S12)
    end
    return out_coeff
end

temp = DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

@time DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Profile.clear()
@profile DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.samples = 5
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 120
@benchmark DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

## Factor of 2 from BLAS asum
# Need to do something for .*

BLAS.dot(test_img,test_img)

dot(test_img,test_img)

At = [1 2 3; 4 5 6; 7 8 9]
Bt = [1 2 3; 4 5 6; 7 8 9]
dot(At,Bt)

1^2+2^2+3^2+4^2+5^2+6^2+7^2+8^2+9^2

## Ok, so asum was a cool find, but dot(abs) is prob better

function structure_test1(J,L,test_img)
    S1 = zeros((J,L))
    @views for l = 1:L
        for j = 1:J
            S1[j,l]+=BLAS.dot(test_img,test_img) #normalization choice arb to make order unity
        end
    end
end

structure_test1(8,8,rand(256,256))

@time structure_test1(8,8,rand(256,256))

@benchmark structure_test1(8,8,rand(256,256))

function structure_test2(J,L,test_img)
    S1 = zeros((J,L))
    @views for l = 1:L
        for j = 1:J
            S1[j,l]+=sum(abs2.(test_img))#normalization choice arb to make order unity
        end
    end
end

structure_test2(8,8,rand(256,256))

@time structure_test2(8,8,rand(256,256))

@benchmark structure_test2(8,8,rand(256,256))

## Yeah, ok, so that is much better. Can we take abs as we allocate?

function structure_test(J,L,test_img_1)
    S20 = zeros((J, L, J, L))
    Atmp = zeros(Float64,256,256)
    Btmp = zeros(Float64,256,256)
    @views for l2 = 1:L
        copyto!(Btmp, abs.(test_img_1))
        for j2 = 1:J
            for l1  = 1:L
                copyto!(Atmp, abs.(test_img_1))
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = BLAS.dot(Atmp,Btmp)
                end
            end
        end
    end
end

structure_test(8,8,rand(256,256))

@time structure_test(8,8,rand(256,256))

@benchmark structure_test(8,8,rand(256,256))

function structure_test(J,L,test_img_1)
    S20 = zeros((J, L, J, L))
    Atmp = zeros(ComplexF64,256,256)
    Btmp = zeros(ComplexF64,256,256)
    @views for l2 = 1:L
        copyto!(Btmp, test_img_1)
        for j2 = 1:J
            for l1  = 1:L
                copyto!(Atmp, test_img_1)
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = BLAS.asum(Atmp.*Btmp)
                end
            end
        end
    end
end

structure_test(8,8,rand(256,256))

@time structure_test(8,8,rand(256,256))

@benchmark structure_test(8,8,rand(256,256))

##
function DHC(image, filter_set; norm_on = 1, coeff_12_on =1, coeff_20_on = 1)
    FFTW.set_num_threads(2)
    (Nx, Ny) = size(image)

    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []

    ## Coeff at (1,0)
    #takes care of image normalization
    S0 = zeros(2)
    if norm_on == 1
        S0[1] = mean(image)
        norm_im = image.-S0[1]
        S0[2] = BLAS.dot(norm_im,norm_im)/(Nx*Ny)
        norm_im = norm_im./sqrt(Nx*Ny*S0[2])
    else
        norm_im = image
    end

    append!(out_coeff,S0)

    ##Coeff at (1,1)
    #store the fft of images and im*filters for later
    im_fd_0 = fft(norm_im)

    S1 = zeros((J,L))
    im_fd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    @views for l = 1:L
        for j = 1:J
            im_fd_0_1[:,:,j,l] .= im_fd_0
        end
    end

    if coeff_20_on == 1
        im_rd_0_1 = Array{Float64,4}(undef, Nx, Ny, J, L)
    end

    Atmp = zeros(ComplexF64,Nx,Ny)
    Btmp = zeros(Float64,Nx,Ny)
    Ctmp = zeros(Float64,Nx,Ny)
    @views for l = 1:L
        for j = 1:J
            had!(im_fd_0_1[:,:,j,l],filter_set[:,:,j,l]) #wavelet already in fft domain not shifted
            Btmp .= abs.(im_fd_0_1[:,:,j,l])
            S1[j,l]+=BLAS.dot(Btmp,Btmp) #normalization choice arb to make order unity
            if coeff_20_on == 1
                Atmp .= ifft(im_fd_0_1[:,:,j,l])
                im_rd_0_1[:,:,j,l] .= abs.(Atmp)
            end
        end
    end
    append!(out_coeff,S1)

    ##Coeff at (2,0)
    if coeff_20_on == 1
        S20 = zeros(Float64,J, L, J, L)
        @views for l2 = 1:L
            for j2 = 1:J
                copyto!(Btmp, im_rd_0_1[:,:,j2,l2])
                for l1 = 1:L
                    for j1  = 1:J
                        copyto!(Ctmp, im_rd_0_1[:,:,j1,l1])
                        S20[j1,l1,j2,l2] += BLAS.dot(Btmp,Ctmp)
                    end
                end
            end
        end

        append!(out_coeff,S20)
    end

    ##Coeff at (1,2) level
    if coeff_12_on == 1
        S12 = zeros(Float64,J, L, J, L)
        for l2 = 1:L
            for j2 = 1:J
                not_rot_im = fftshift(im_fd_0_1[:,:,j1,l1])
                for l1 = 1:L
                    for j1 = 1:J
                        S12[j1,l1,j2,l2] += sum(abs.(not_rot_im.*fftshift(im_fd_0_1[:,:,j2,l2])))
                    end
                end
            end
        end
    append!(out_coeff,S12)
    end
    return out_coeff
end

temp = DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

@time DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Profile.clear()
@profile DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.samples = 5
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 120
@benchmark DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

## Starting to test the S12 coeff

function fftshift_wrap()
    A = rand(256,256)
    fftshift(A)
end

@time fftshift_wrap()

@benchmark fftshift_wrap()

function fftshift_wrap()
    A = rand(256,256)
end

@time fftshift_wrap()

@benchmark fftshift_wrap()


@test Id(2) == Symp([1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1])


@doc raw"""
    ExponentialSmoothing(
        y::Vector{Fl};
        trend::Bool = false,
        damped_trend::Bool = false,
        seasonal::Int = 0
    ) where Fl

Linear exponential smoothing models. These models are also known as ETS in the literature.
This model is estimated using the Kalman filter for linear Gaussian state space models, for this
reason the possible models are the following ETS with additive errors:
 * ETS(A, N, N)
 * ETS(A, A, N)
 * ETS(A, Ad, N)
 * ETS(A, N, A)
 * ETS(A, A, A)
 * ETS(A, Ad, A)

Other softwares have use the augmented least squares approach and have all the possible ETS
combinations. The Kalman filter approach might be slower than others but have the advantages of
filtering the components.

# References
 * Hyndman, Rob, Anne B. Koehler, J. Keith Ord, and Ralph D. Snyder.
   Forecasting with exponential smoothing: the state space approach.
   Springer Science & Business Media, 2008.
 * Hyndman, Robin John; Athanasopoulos, George.
   Forecasting: Principles and Practice.
   2nd ed. OTexts, 2018.
"""
mutable struct ExponentialSmoothing <: StateSpaceModel
    hyperparameters::HyperParameters
    trend::Bool
    damped_trend::Bool
    seasonal::Int
    system::LinearUnivariateTimeInvariant
    results::Results

    function ExponentialSmoothing(y::Vector{Fl};
                                  trend::Bool = false,
                                  damped_trend::Bool = false,
                                  seasonal::Int = 0
                                  ) where Fl
        @assert seasonal != 1 "seasonal must be different than 1"
        if damped_trend
            @assert trend
        end
        # Define system matrices
        Z = build_Z(Fl, trend, seasonal)
        T = build_T(Fl, trend, seasonal)
        R = zeros(Fl, size(T, 1), 1)
        R[1] = 1
        d = zero(Fl)
        c = zeros(Fl, size(T, 1))
        H = zero(Fl)
        Q = ones(Fl, 1, 1)

        system = LinearUnivariateTimeInvariant{Fl}(y, Z, T, R, d, c, H, Q)

        # Define hyperparameters names
        names = build_names(trend, damped_trend, seasonal)
        hyperparameters = HyperParameters{Fl}(names)

        return new(hyperparameters, trend, damped_trend,
                   seasonal, system, Results{Fl}())
    end
end

num_states(trend::Bool, seasonal::Int) = 2 + trend + seasonal

function build_Z(Fl::DataType, trend::Bool, seasonal::Int)
    n_states = num_states(trend, seasonal)
    Z = zeros(Fl, n_states)
    Z[1:2] .= 1
    if seasonal > 0
        Z[3 + trend] = 1
    end
    return Z
end

function build_T(Fl::DataType, trend::Bool, seasonal::Int)
    n_states = num_states(trend, seasonal)
    T = zeros(Fl, n_states, n_states)
    T[2, 2] = 1
    if trend
        T[2:3, 3] .= 1
    end
    if seasonal > 0
        T[3 + trend, end] = 1
        for i in 3+trend+1:n_states
            T[i, i-1] = 1
        end
    end
    return T
end

function build_names(trend::Bool, damped_trend::Bool, seasonal::Int)
    names = String["sigma2", "smoothing_level"]
    if trend
        push!(names, "smoothing_trend")
        if damped_trend
            push!(names, "damping_trend")
        end
    end
    if seasonal > 0
        push!(names, "smoothing_seasonal")
    end
    push!(names, "initial_level")
    if trend
        push!(names, "initial_trend")
    end
    if seasonal > 0
        for i in 1:seasonal - 1
            push!(names, "initial_seasonal_$i")
        end
    end
    return names
end

function diff_es(y::Vector{Fl}, s::Int) where Fl
    # Seasonal differencing
    y = y[s:end] - y[1:end-s+1]
    return y
end

function default_filter(model::ExponentialSmoothing)
    Fl = typeof_model_elements(model)
    n_states = num_states(model.trend, model.seasonal)
    steadystate_tol = Fl(1e-5)
    a1 = zeros(Fl, n_states)
    P1 = Fl(1e6) .* Matrix{Fl}(I, n_states, n_states)
    skip_llk = 0
    return UnivariateKalmanFilter(a1, P1, skip_llk, steadystate_tol)
end

function initial_hyperparameters!(model::ExponentialSmoothing)
    Fl = typeof_model_elements(model)
    observations = model.system.y[findall(!isnan, model.system.y)]
    observed_variance = var(observations)
    initial_hyperparameters = Dict{String,Fl}(
        "sigma2" => observed_variance,
        "smoothing_level" => Fl(0.1),
        "initial_level" => observations[1]
    )
    if model.trend
        initial_hyperparameters["smoothing_trend"] = Fl(0.01)
        initial_hyperparameters["initial_trend"] = observations[2] - observations[1]
        if model.damped_trend
            initial_hyperparameters["damping_trend"] = Fl(0.97)
        end
    end
    if model.seasonal > 0
        # If the model is seasonal we update the initial_level and initial_trend
        initial_hyperparameters["smoothing_seasonal"] = Fl(0.01)
        # TODO this might be wrong for missing observations in the first s observations
        initial_hyperparameters["initial_level"] = mean(observations[1:model.seasonal - 1])
        if model.trend
            obs = diff_es(observations, model.seasonal)
            initial_hyperparameters["initial_trend"] = obs[2] - obs[1]
        end
        for i in 1:model.seasonal - 1
            initial_hyperparameters["initial_seasonal_$i"] = observations[i] -
                                                             mean(observations[1:model.seasonal - 1])
        end
    end
    set_initial_hyperparameters!(model, initial_hyperparameters)
    return model
end

function constrain_hyperparameters!(model::ExponentialSmoothing)
    Fl = typeof_model_elements(model)
    constrain_variance!(model, "sigma2")
    constrain_box!(model, "smoothing_level", Fl(0.0001), Fl(0.9999))
    constrain_identity!(model, "initial_level")
    if model.trend
        ub = min(Fl(0.9999), get_constrained_value(model, "smoothing_level"))
        constrain_box!(model, "smoothing_trend", Fl(0.0001), ub)
        constrain_identity!(model, "initial_trend")
        if model.damped_trend
            constrain_box!(model, "damping_trend", Fl(0.8), Fl(0.98))
        end
    end
    if model.seasonal > 0
        ub = min(Fl(0.9999), 1 - get_constrained_value(model, "smoothing_level"))
        constrain_box!(model, "smoothing_seasonal", Fl(0.0001), ub)
        for i in 1:model.seasonal - 1
            constrain_identity!(model, "initial_seasonal_$i")
        end
    end
    return model
end

function unconstrain_hyperparameters!(model::ExponentialSmoothing)
    Fl = typeof_model_elements(model)
    unconstrain_variance!(model, "sigma2")
    unconstrain_box!(model, "smoothing_level", Fl(0.0001), Fl(0.9999))
    unconstrain_identity!(model, "initial_level")
    if model.trend
        ub = min(Fl(0.9999), get_constrained_value(model, "smoothing_level"))
        unconstrain_box!(model, "smoothing_trend", Fl(0.0001), ub)
        unconstrain_identity!(model, "initial_trend")
        if model.damped_trend
            unconstrain_box!(model, "damping_trend", Fl(0.8), Fl(0.98))
        end
    end
    if model.seasonal > 0
        ub = min(Fl(0.9999), 1 - get_constrained_value(model, "smoothing_level"))
        unconstrain_box!(model, "smoothing_seasonal", Fl(0.0001), ub)
        for i in 1:model.seasonal - 1
            unconstrain_identity!(model, "initial_seasonal_$i")
        end
    end
    return model
end

function fill_model_system!(model::ExponentialSmoothing)
    model.system.Q[1] = get_constrained_value(model, "sigma2")
    i = 1
    model.system.R[i] = 1 - get_constrained_value(model, "smoothing_level")
    i += 1
    model.system.R[i] = get_constrained_value(model, "smoothing_level")
    i += 1
    if model.trend
        model.system.R[i] = get_constrained_value(model, "smoothing_trend")
        i += 1
        if model.damped_trend
            model.system.T[2:3, 3] .= get_constrained_value(model, "damping_trend")
        end
    end
    if model.seasonal > 0
        model.system.R[1] -= get_constrained_value(model, "smoothing_seasonal")
        model.system.R[i] = get_constrained_value(model, "smoothing_seasonal")
    end
    return model
end

function fill_model_filter!(filter::KalmanFilter, model::ExponentialSmoothing)
    Fl = typeof_model_elements(model)
    initial_level = get_constrained_value(model, "initial_level")
    initial_state = [0; initial_level]
    if model.trend
        initial_trend = get_constrained_value(model, "initial_trend")
        initial_state = vcat(initial_state, initial_trend)
    end
    if model.seasonal > 0
        initial_seasonal = Fl[]
        for i in 1:model.seasonal - 1
            push!(initial_seasonal, get_constrained_value(model, "initial_seasonal_$i"))
        end
        last_initial_seasonal = -sum(initial_seasonal)
        push!(initial_seasonal, last_initial_seasonal)
        initial_state = vcat(initial_state, initial_seasonal)
    end
    filter.kalman_state.a = model.system.T * initial_state
    filter.kalman_state.P = model.system.R * model.system.Q * model.system.R'
    return nothing
end

has_exogenous(model::ExponentialSmoothing) = false

function reinstantiate(model::ExponentialSmoothing, y::Vector{Fl}) where Fl
    return ExponentialSmoothing(y;
                trend = model.trend,
                damped_trend = model.damped_trend,
                seasonal = model.seasonal
            )
end

function model_name(model::ExponentialSmoothing)
    E = "A"
    T = model.trend ?
        model.damped_trend ? "Ad" : "A" :
        "N"
    S = model.seasonal > 0 ? "A" : "N"
    return "ETS($E,$T,$S)"
end

function dict_components(model::ExponentialSmoothing)
    dict_components = OrderedDict{String, Int}()
    i = 2
    dict_components["Trend"] = i
    i += 1
    if model.trend
        dict_components["Slope"] = i
        i += 1
    end
    if model.seasonal > 0
        dict_components["Seasonal"] = i
    end
    return dict_components
end

@doc raw"""
    auto_ets(y::Vector{Fl}; seasonal::Int = 0) where Fl

Automatically fits the best [`ExponentialSmoothing`](@ref) model according to the best AIC
between the models:
 * ETS(A, N, N)
 * ETS(A, A, N)
 * ETS(A, Ad, N)

If the user provides the time series seasonality it will search between the models
 * ETS(A, N, A)
 * ETS(A, A, A)
 * ETS(A, Ad, A)

# References
 * Hyndman, Robin John; Athanasopoulos, George.
 Forecasting: Principles and Practice.
 2nd ed. OTexts, 2018.
"""
function auto_ets(y::Vector{Fl}; seasonal::Int = 0) where Fl
    models = StateSpaceModel[]
    models_aic = Fl[]
    @assert seasonal != 1 "seasonal must be different than 1"
    m1 = ExponentialSmoothing(y; trend = false, damped_trend = false, seasonal = seasonal)
    fit!(m1)
    push!(models, m1)
    push!(models_aic, m1.results.aic)
    m2 = ExponentialSmoothing(y; trend = true, damped_trend = false, seasonal = seasonal)
    fit!(m2)
    push!(models, m2)
    push!(models_aic, m2.results.aic)
    m3 = ExponentialSmoothing(y; trend = true, damped_trend = true, seasonal = seasonal)
    fit!(m3)
    push!(models, m3)
    push!(models_aic, m3.results.aic)
    best_aic_idx = findmin(models_aic)[2] # index of the best BIC
    return models[best_aic_idx]
end


mutable struct Polynomial <: AbstractComponent
    coeff::Vector{Parameter}

    function Polynomial(args...)
        coeff = [Parameter(arg) for arg in args]
        return new(coeff)
    end
end

struct Polynomial_cdata <: AbstractComponentData; end
cdata(comp::Polynomial, domain::AbstractDomain) = Polynomial_cdata()

function evaluate!(cdata::Polynomial_cdata, output::AbstractArray{Float64}, domain::AbstractDomain, coeffs...)
    output .= coeffs[1]
    x = domain[1]
    for deg in 1:length(coeffs)-1
        output .+= x.^deg .* coeffs[deg+1]
    end
    return output
end


## Vector of fun routines

function coefficientmatrix{N,F}(::Type{N},f::Vector{F},o...)
    if isempty(f)
        return Array(N,0,0)
    end

    n=mapreduce(ncoefficients,max,f)
    m=length(f)
    R=zeros(N,n,m)
    for k=1:m
        R[1:ncoefficients(f[k]),k]=coefficients(f[k],o...)
    end
    R
end

scalarorfuntype{S,T<:Number}(::Fun{S,T})=T
scalarorfuntype{T<:Number}(::T)=T
scalarorfuntype{T<:Number}(b::Vector{T})=T
scalarorfuntype(b::Vector{Any})=promote_type(map(scalarorfuntype,b)...)
scalarorfuntype{F<:Fun}(b::Vector{F})=promote_type(map(scalarorfuntype,b)...)

coefficientmatrix{F<:Fun}(Q::Vector{F},o...)=coefficientmatrix(scalarorfuntype(Q),Q,o...)
coefficientmatrix(Q::Vector{Any})=(@assert isempty(Q); zeros(0,0))

function values{D,N}(f::Vector{Fun{D,N}})
    n=mapreduce(ncoefficients,max,f)
    m=length(f)
    R=zeros(N,n,m)
    for k=1:m
        R[:,k] = values(pad(f[k],n))
    end
    R
end

function values{D,T}(p::Array{Fun{D,T},2})
    @assert size(p)[1] == 1

   values(vec(p))
end

## evaluation

#TODO: fix for complex
evaluate{T<:Fun}(A::AbstractArray{T},x::Number)=typeof(first(A)(x))[Akj(x) for Akj in A]

function evaluate{T<:Fun}(A::AbstractVector{T},x::AbstractVector)
    n=length(x)
    ret=Array(promote_type(eltype(x),mapreduce(eltype,promote_type,A)),length(A),n)

    for k=1:length(A)
        bkr=evaluate(A[k],x)

        for j=1:n
            ret[k,j]=bkr[j]
        end
    end

    ret
end

## Algebra

## scalar fun times vector

 for op in (:*,:(Base.Ac_mul_B),:(Base.At_mul_B))
     @eval begin
         function ($op){T<:Number,V<:Number,D}(A::Array{T,2}, p::Vector{Fun{D,V}})
             cfs=$op(A,coefficientmatrix(p).')
             ret = Array(Fun{D,promote_type(T,V)},size(cfs,1))
             for i = 1:size(cfs,1)
                 ret[i] = chop!(Fun(vec(cfs[i,:]),first(p).space),eps())
             end
             ret
         end

         function ($op){T<:Number,D}(p::Vector{Fun{D,T}},A::Array{T,2})
             cfs=$op(A,coefficientmatrix(p).')
             ret = Array(Fun{D,T},size(cfs,1))
             for i = 1:size(cfs,1)
                 ret[i] = chop!(Fun(vec(cfs[i,:]),first(p).space),eps())
             end
             ret
         end
     end
 end

#Allow vecfun + constvec, etc
#can't just promote constant vector to a vector-valued fun because don't know the domain.
for op = (:+,:-,:.*,:./)
    @eval begin
        ($op){T<:Number,S,V}(f::Fun{S,V},c::AbstractArray{T})=devec($op(vec(f),c))
        ($op){T<:Number,S,V}(c::AbstractArray{T},f::Fun{S,V})=devec($op(c,vec(f)))
    end
end


using Test, MPI
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.Mesh.Geometry
using StaticArrays

MPI.Initialized() || MPI.Init()

@testset "LocalGeometry" begin
  FT = Float64
  ArrayType = Array

  xmin      = 0
  ymin      = 0
  zmin      = 0
  xmax      = 2000
  ymax      = 400
  zmax      = 2000

  Ne        = (20,2,20)

  polynomialorder = 4

  brickrange = (range(FT(xmin); length=Ne[1]+1, stop=xmax),
                range(FT(ymin); length=Ne[2]+1, stop=ymax),
                range(FT(zmin); length=Ne[3]+1, stop=zmax))
  topl = StackedBrickTopology(MPI.COMM_SELF, brickrange, periodicity = (true, true, false))

  grid = DiscontinuousSpectralElementGrid(topl,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = polynomialorder
                                          )

  S = ((xmax-xmin)/(polynomialorder*Ne[1]), (ymax-ymin)/(polynomialorder*Ne[2]), (zmax-zmin)/(polynomialorder*Ne[3]))
  Savg = cbrt(prod(S))
  M = SDiagonal(S.^-2)

  for e in 1:size(grid.vgeo,3)
    for n in 1:size(grid.vgeo,1)
      g = LocalGeometry(Val(polynomialorder), grid.vgeo, n, e)
      @test lengthscale(g) ≈ Savg
      @test Geometry.resolutionmetric(g) ≈ M
    end
  end
end


import Pkg
Pkg.activate(".")

using DistancePlan
using StaticArrays
import Makie

scene = Makie.Scene()

DistancePlan.draw!(Sphere(SVector{2}([0. 0.]), 1.))

Makie.save("unit_circle.png", scene, resolution=(1000, 1000))


@testset "Blackbody T=$T" for T in [2000, 4000, 6000]
    wave = range(1e3, 5e4, length = 1000)
    b = 2.897771955185172e7

    bb = @inferred blackbody(wave, T)
    @test typeof(bb) <: Spectra.Spectrum
    @test bb.T == T
    @test bb.wave[argmax(bb)] ≈ b / T rtol = 0.01

    wave *= u"angstrom"
    T *= u"K"
    bb = @inferred blackbody(wave, T)
    @test typeof(bb) <: Spectra.Spectrum
    @test unit(bb)[2] == u"W/m^2/angstrom"
    @test bb.T == T
    @test bb.wave[argmax(bb)] ≈ b * u"angstrom*K" / T rtol = 0.01
end


function odeint(y0::Vector{Float64},t0::Float64,tf::Float64,sparams::WK3.SolverParams,
    mparams::WK3.ModelParams,k::Float64)
    # parameters for dependent variable scaling
    tiny = 1e-30;

    # integration options
    saveflag = true; # save at intervals ≈ dtsav

    # output variables
    tout = Float64[];
    yout = Vector{Float64}[];

    # initialization
    y = zeros(sparams.nvar);
    ys = zeros(sparams.nvar);
    dy = zeros(sparams.nvar);
    hdid = 0.;
    hnext = 0.;
    t = t0;
    y[:] = y0;
    nok = 0;
    nbad = 0;

    # make sure we're stepping in the right direction (forwards vs. backwards)
    h = (tf-t0) >= 0 ? (sparams.h0 >= 0 ? sparams.h0: -sparams.h0) : (sparams.h0 >= 0 ? -sparams.h0 : sparams.h0);

    # assure storage of first step
    if saveflag
        tsav = t-sparams.dtsav*2;
    end

    while true
        # evaluate derivatives
        # println("Derivative evaluation:")
        wk3odes(t,y,dy,mparams,k);
        # scaling to monitor accuracy
        # println("Scaling:")
        ys .= abs.(y) .+ abs.(dy.*h).+tiny;
        # store intermediate results
        if saveflag && abs.(t-tsav) > abs(sparams.dtsav)
            push!(tout,t)
            push!(yout,y[:])
            tsav = t;
        end
        # decrease stepsize if overshooting end of integration window
        if (t+h-tf)*(t+h-t0) > 0
            h = tf-t;
        end
        # attempt an integration step
        # println("RKQS:")
        t,hdid,hnext=rkqs(y,dy,t,h,hdid,hnext,ys,sparams,mparams,k);
        # check valve state, set ventricular flow accordingly
        E,~ = elastancefn(t*ts,mparams,k);
        Pv = E.*(y[2].*Vs.-mparams.V0)./Ps;
        if Pv < y[1] && y[3] <= 0.
            y[3] = 0.;
        end
        # track times step size was changed
        if hdid == h
            nok+=1;
        else
            nbad+=1;
        end
        # normal exit if finished (saving final step)
        if (t-tf)*(tf-t0) >= 0
            push!(yout,y)
            push!(tout,t)
            return tout,yout
        end
        if abs.(hnext) <= sparams.hmin
            push!(yout,y)
            push!(tout,t)
            println("Step size below minimum allowable in odeint. Aborting. hnext = $hnext. t = $t.")
            return tout,yout
        end
    end
end


function matgen(n)
  tmp = 1.0 / n / n
  [ tmp * (i - j) * (i + j - 2) for i=1:n, j=1:n ]
end

function main()
  n = 100
  if length(ARGS) >= 1
    n = int(ARGS[1])
  end
  t = time()
  n = int(n / 2 * 2)
  a = matgen(n)
  b = matgen(n)
  c = a * b
  v = int(n/2) + 1
  println(c[v, v])
  println(time() - t)
end

main()


"""
CTMC

Continous Time Markov Chain
"""
mutable struct CTMC
    τ::AbstractFloat
    transitiontimes::Vector{AbstractFloat}
    states::Vector{Int}
    CTMC() = new()
end

function CTMC(q::Array, τ::AbstractFloat,state1::Int)

# create empty CTMC structure
    ctmc = CTMC()
    ctmc.τ=τ

    lastchange = 0.0
    currentstate = state1

    ctmc.states = [state1]
    ctmc.transitiontimes = [lastchange]

    sidx = Vector((1:size(q, 1)))

    while lastchange < τ

# get time for state change
        k_tot = sum(q[currentstate,:])
        Δt = rand(Exponential(1 / k_tot))

# get the new state
        ps = q[currentstate,:]
        ps = ps ./ sum(ps)
        deleteat!(ps,currentstate)
        xs = sidx[:]
        deleteat!(xs,currentstate)
        newstate = rand(DiscreteNonParametric(xs, ps))

# update CTMC
        lastchange +=Δt
        push!(ctmc.states,newstate)
        push!(ctmc.transitiontimes,lastchange)
        currentstate = newstate
    end
    return ctmc
end

"""
    getstate(ctmc::CTMC,t::AbstractFloat)

 return the state at time t
"""
function getstate(ctmc::CTMC,t::AbstractFloat)
    nstates=length(ctmc.states)
    for nn=2:nstates
        if t<ctmc.transitiontimes[nn]
            return ctmc.states[nn-1]
        end
    end
end

"""
    getnext(ctmc::CTMC,t::AbstractFloat)

 return the time and state of next transision
"""
function getnext(ctmc::CTMC,t::AbstractFloat)
    nstates=length(ctmc.states)
    for nn=1:nstates
        if t<ctmc.transitiontimes[nn]
            return ctmc.states[nn],ctmc.transitiontimes[nn]
        end
    end
end

"""
    intensitytrace(f::GenericFluor, nframes::Int, framerate::AbstractFloat;state1=1)

Calculate an intensity trace.

"""
function intensitytrace(f::GenericFluor, nframes::Int, framerate::Real;state1=1)

    endtime = (nframes) / framerate

    # generate CTMC
    ctmc=SMLMSim.CTMC(f.q,endtime,state1)

    # generate integrated photons
    exptime=1/framerate
    photons=zeros(nframes)
    for nn=1:nframes
        t=(nn-1)*exptime
        frameend=(nn*exptime)
        while t<frameend
            currentstate=SMLMSim.getstate(ctmc,t)
            (~,nexttime)=SMLMSim.getnext(ctmc,t)
            tend=min(frameend,nexttime)
            if currentstate==1 # the fluorescent state
                photons[nn]+=f.γ*(tend-t)
            end
            t=tend
        end
    end
    return photons
end

# function genroistack(psf::MicroscopePSFs.PSF, xsize, ysize, nframes, framerate)
# end

"""
    function kineticmodel(smd_true::SMLMData.SMLD2D,f::Molecule,nframes::Int,framerate::AbstractFloat;ndatasets::Int=1,minphotons=50.0)

generate noise-free blinking model from smd_true
"""
function kineticmodel(smd_true::SMLMData.SMLD2D,f::Molecule,nframes::Int,framerate::Real;ndatasets::Int=1,minphotons=50.0)

    state1=2;

    smd=SMLMData.SMLD2D(0)
    smd.ndatasets=ndatasets
    smd.nframes=nframes
    smd.datasize=deepcopy(smd_true.datasize)
    for dd=1:ndatasets, ll=1:length(smd_true.x)
        photons=SMLMSim.intensitytrace(f,nframes,framerate;state1=state1);
        framenum=findall(photons.>minphotons)
        n=length(framenum)
        push!(smd.photons,photons[framenum]...)
        push!(smd.framenum,framenum...)
        for nn = 1:n
            push!(smd.x, smd_true.x[ll])
            push!(smd.y, smd_true.y[ll])
            push!(smd.connectID, smd_true.connectID[ll])
            push!(smd.datasetnum, dd)
        end
    end
    return smd
end

function noise(smd_model::SMLMData.SMLD2D,σ_psf::AbstractFloat)
    n=length(smd_model.x)
    smd=deepcopy(smd_model)
    smd.σ_x=zeros(n)
    smd.σ_y=zeros(n)

    for nn=1:n
        σ=σ_psf/sqrt(smd_model.photons[nn])
        smd.x[nn]=smd_model.x[nn]+randn()*σ
        smd.y[nn]=smd_model.y[nn]+randn()*σ
        smd.σ_x[nn]=σ
        smd.σ_y[nn]=σ
    end
    return smd
end


# PerceptronLoss

using LaTeXStrings, LossFunctions, Measures, Latexify

Plots.plot(PerceptronLoss(), c=:olive, linewidth=2, annotation=(1, 1.6, L"L\left( a \right) = max\{0, - a \}", 18))


#####
##### Electrostatic Energy
#####

function electrostatic_energy(
    qa::Float64,
    qb::Float64,
    d::Float64,
    )
    return qa * qb / d
end

function electrostatic_energy(
    qa::Float64,
    ra::Vec3{Float64},
    qb::Float64,
    rb::Vec3{Float64},
    )
    d = norm(ra - rb)
    return electrostatic_energy(qa, qb, d)
end

function electrostatic_energy(
    a::AbstractParticle,
    b::AbstractParticle,
    )
    return electrostatic_energy(a.charge, a.position, b.charge, b.position)
end

function electrostatic_energy(
    g::AbstractParticleGroup,
    )
    ret = 0.0
    for i=1:length(g)-1, j=i+1:length(g)
        @inbounds ret += electrostatic_energy(g[i], g[j])
    end
    return ret
end

function electrostatic_energy(
    ga::AbstractParticleGroup,
    gb::AbstractParticleGroup,
    )
    ret = 0.0
    for i=1:length(ga), j=1:length(gb)
        @inbounds ret += electrostatic_energy(ga[i], gb[j])
    end
    return ret
end

#####
##### Electrostatic Potential
#####

function electrostatic_potential(
    q::Float64,
    d::Float64,
    )
    return q / d
end

function electrostatic_potential(
    q::Float64,
    ra::Vec3{Float64},
    rb::Vec3{Float64},
    )
    d = norm(ra - rb)
    return electrostatic_potential(q, d)
end

function electrostatic_potential(
    q::Float64,
    σ::Float64,
    d::Float64,
    )
    return q / d * erf(d / (sqrt(2) * σ))
end

function electrostatic_potential(
    q::Float64,
    σ::Float64,
    ra::Vec3{Float64},
    rb::Vec3{Float64},
    )
    d = norm(ra - rb)
    return electrostatic_potential(q, σ, d)
end

function electrostatic_potential(
    particle::AbstractParticle,
    point::Vec3{Float64},
    )
    return electrostatic_potential(particle.charge, particle.position, point)
end

function electrostatic_potential(
    particle::GaussianCharge,
    point::Vec3{Float64},
    )
    return electrostatic_potential(particle.charge, particle.sigma, particle.position, point)
end

function electrostatic_potential(
    particles::AbstractParticleGroup,
    point::Vec3{Float64},
    )
    ret = 0.0
    for i=1:length(particles)
        @inbounds ret += electrostatic_potential(particles[i], point)
    end
    return ret
end

function electrostatic_potential(
    scf::SCF,
    point::Vec3{Float64},
    )
    Vnuc = electrostatic_potential(scf.mole.nuclei, point)
    intor = Intor(scf.mole)
    Vel = intor("int1e_rinv", point) ⋅ scf.P * 2
    return Vnuc - Vel
end

function electrostatic_potential(
    obj::Union{AbstractParticle,AbstractParticleGroup,SCF},
    points::AbstractVector{Vec3{Float64}},
    )
    ret = zeros(length(points))
    for (i, point) in enumerate(points)
        ret[i] = electrostatic_potential(obj, point)
    end
    return ret
end


using TomoForward
using Images
using Plots
using XfromProjections
using StaticArrays
using PyCall
using Logging
using Random
using Printf

#include("../tv_primaldual_flow.jl")
#include("../optical_flow.jl")
include("../../phantoms/sperm_phantom.jl")

H,W = 609, 609
detmin, detmax = -38.0, 38.0
function radon_operator(height, width, detcount, angles)
    proj_geom = ProjGeom(0.5, detcount, angles)
    A = fp_op_parallel2d_line(proj_geom, height, width, detmin,detmax, detmin,detmax)
    return A
end
grid = collect(detmin:0.1:detmax)
r(s) = 1.0

images, tracks = get_sperm_phantom(301,r,grid)
frames = images[:,:,collect(1:31)]
@info size(frames)

w_tv = 0.1
w_flow  = 0.05
c=1.1

detcount = Int(floor(H*1.4))

cwd = @__DIR__
path = normpath(joinpath(@__DIR__, "result"))
!isdir(path) && mkdir(path)
cd(path)

#4 angles evenly spaced between 0 and pi
#angles = collect(range(0.0,π,length=5))[1:end-1]#angles = rand(0.0:0.001:π, 4)#collect(range(π/2,3*π/2,length=ang_nr+1))[1:end-1]
#@info angles
nangles = 3
a =rand(0.0:0.001:π, nangles*size(frames)[3]) #range(0.0,4*pi+pi/4, length=nangles*size(frames)[3])#
angles = map(t -> a[(t-1)*nangles+1:(t-1)*nangles+nangles], 1:size(frames)[3])
@info angles
#match size of input image (generating data)
As = map(t -> radon_operator(size(frames[:,:,1],1),size(frames[:,:,1],2),detcount, angles[t]),1:size(frames)[3])
bsex = zeros(size(As[1])[1],size(frames)[3])

map(t -> bsex[:,t] = As[t]*vec(frames[:,:,t]), 1:size(frames)[3])
@info "adding gaussion noise at level 0.01"
rho = 0.01
e = randn(size(bsex));
e = rho*norm(bsex)*e/norm(e);
bs = bsex + e;

u0s = zeros(H,W,size(frames)[3])
mask = ones(size(u0s)...)

As  = map(t -> radon_operator(H,W, detcount, angles[t]),1:size(u0s)[3])

@info "Reconstruction using tv regularization frame by frame"
niter=450
us_tv = zeros(H,W,size(frames)[3])
for t = 1:size(frames)[3]
    A = As[t]
    p = bs[:,t]
    u0 = u0s[:,:,t]
    us_tv[:,:,t] .= recon2d_tv_primaldual!(us_tv[:,:,t], A, p, niter, w_tv, c)
end
p1 = heatmap(grid, grid, Gray.(us_tv[:,:,1]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
plot!(tracks[271][:,1], tracks[271][:,2], aspect_ratio=:equal, linewidth=5)
p2 = heatmap(grid, grid, Gray.(us_tv[:,:,2]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
plot!(tracks[272][:,1], tracks[272][:,2], aspect_ratio=:equal, linewidth=5)
p3 = heatmap(grid, grid, Gray.(us_tv[:,:,3]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
plot!(tracks[273][:,1], tracks[273][:,2], aspect_ratio=:equal, linewidth=5)
p4 = heatmap(grid, grid, Gray.(us_tv[:,:,4]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
plot!(tracks[274][:,1], tracks[274][:,2], aspect_ratio=:equal, linewidth=5)
l = @layout [a b c d]
plot(p1, p2, p3, p4, layout = l, size=(2000,600), linewidth=5)
savefig(@sprintf "result_all_tv_2_%d" length(angles))

@info "Reconstructing using joint motion estimation and reconstruction"
niter1=100
niter2=500
u0s = deepcopy(us_tv)
us_flow = recon2d_tv_primaldual_flow(As, bs, u0s, niter1, niter2, w_tv, w_flow, c, mask, 0.012)

p1 = heatmap(grid, grid, Gray.(us_flow[:,:,1]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
plot!(tracks[271][:,1], tracks[271][:,2], aspect_ratio=:equal, linewidth=5)
p2 = heatmap(grid, grid, Gray.(us_flow[:,:,2]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
plot!(tracks[272][:,1], tracks[272][:,2], aspect_ratio=:equal, linewidth=5)
p3 = heatmap(grid, grid, Gray.(us_flow[:,:,3]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
plot!(tracks[273][:,1], tracks[273][:,2], aspect_ratio=:equal, linewidth=5)
p4 = heatmap(grid, grid, Gray.(us_flow[:,:,4]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
plot!(tracks[274][:,1], tracks[274][:,2], aspect_ratio=:equal, linewidth=5)
l = @layout [a b c d]
plot(p1, p2, p3, p4, layout = l, size=(2000,600), linewidth=5)
savefig(@sprintf "result_all_flow_2_%d" length(angles))

@info "Preparing results in human readable format"
anim = @animate for t=1:size(frames)[3]
    l = @layout [a b c]
    p1 = heatmap(grid, grid, Gray.(frames[:,:,t]),yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false, title="Ground Truth")
    p2 = heatmap(grid, grid, Gray.(us_flow[:,:,t]),yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false, title="Flow")
    plot!(tracks[t+270][:,1], tracks[t+270][:,2], aspect_ratio=:equal, linewidth=1)
    p3 = heatmap(grid, grid, Gray.(us_tv[:,:,t]),yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false, title="TV")
    plot!(tracks[t+270][:,1], tracks[t+270][:,2], aspect_ratio=:equal, linewidth=1)
    plot(p1, p2, p3, layout = l)
end

cwd = @__DIR__
path = normpath(joinpath(@__DIR__, "result"))
!isdir(path) && mkdir(path)
cd(path)
gif(anim, "reconstruction_flow_tv_2.gif", fps = 1)
cd(cwd)


#
# Computes the total potential energy of the system using r2
#

function utotal2(p,data)
    first_atom = zeros(Int64, data.nc, data.nc)
    next_atom = zeros(Int64, data.N)
    makell!(first_atom, next_atom, p, data)
    ut = 0.
    for ip = 1:data.N
        icell = trunc(Int64, p[ip][1]/data.cutoff) + 1
        jcell = trunc(Int64, p[ip][2]/data.cutoff) + 1
        for ic = icell-1:icell+1
            iw = wrap_cell(data.nc, ic)
            for jc = jcell-1:jcell+1
                jw = wrap_cell(data.nc, jc)
                jp = first_atom[iw,jw]
                while jp > 0
                    if ip < jp
                        r2 = r2pbc(p[ip], p[jp], data.bdim)
                        if r2 ≤ data.cutoff2
                            ut += upair2(r2, data)
                        end
                    end
                    jp = next_atom[jp]
                end
            end
        end
    end
    return ut
end


# The controller node
# This file is part of the openBuildNet framework developed at EPFL.
# (C) 2016 LA-EPFL
# Author: Truong X. Nghiem (truong.nghiem@gmail.com, xuan.nghiem@epfl.ch)

# MPC node
# See: http://users.isy.liu.se/johanl/yalmip/pmwiki.php?n=Examples.StandardMPC

using openbuildnet

# Model data
const A = [2. -1.;1. 0.2]
const B = [1. 0.]'
const x0 = [3., 1.];  # Initial state

# MPC data
nx = size(A,2); # Number of states
nu = size(B,2); # Number of inputs

# MPC data
Q = eye(nx);
R = 2.;
N = 7;

# Prepare the optimization model
using GLPKMathProgInterface
using Convex

function create_model()
  local u = Variable(nu, N);
  local x = Variable(nx, N);
  local x0v = Variable(nx);

  local constraints = [-1 <= u, u <= 1, -5 <= x, x <= 5];
  constraints += x[:,1] == A*x0v + B*u[:,1];
  objective = norm(Q*x0v,1) + norm(R*u[:,1],1);
  for k = 2:N
    objective = objective + norm(Q*x[:,k-1],1) + norm(R*u[:,k],1);
    constraints += x[:,k] == A*x[:,k-1] + B*u[:,k];
  end

  controller = minimize(objective, constraints);

  # returns a tuple of the controller, the parameter, and the solution
  (controller, x0v, u)
end

tic()
controller, x0v, u = create_model()
toc()

# Create the node
node = OBNNode("controller", "simplempc")

# Create the ports
input_x = create_input(node, "x", Vector{Float64})  # feedback
output_u = create_output(node, "u", Vector{Float64}) # control values

# Create Callbacks
on_init(node) do
  println("INIT")
end

on_term(node) do
  println("TERMINATE")
end

const MAINBLOCK = 0

US = Vector{Vector{Float64}}()  # To store the control values

on_block_output(node, MAINBLOCK) do
  global controller, x0v, u
  fix!(x0v, get(input_x))
  tic()
  solve!(controller, GLPKSolverLP())
  toc()
  uk = u.value[:,1]
  push!(US, uk)
  set(output_u, uk)
end

println("Ready to start simulation...")
status = runsafe(node, 20)
println("Final status: ", status)

# Plot the results
if status == 2
  using PyPlot
  USmat = hcat(US...)' # convert to matrix
  plot(USmat[:,1])
  show()
end


export air_density, absolute_pressure, vapor_pressure_over_water

"""
    air_density(P, T, ϕ)

Calculate density of humid air [kg/m^3] given absolute pressure `P` [Pa], temperature `T` [K] and relative humidity `ϕ` [1].
"""
function air_density(P, T, ϕ)
    R_d = 287.058
    R_v = 461.495

    P_sat = vapor_pressure_over_water(T - 273.15)
    P_v = ϕ * P_sat
    P_d = P - P_v

    (P_d / R_d + P_v / R_v) / T
end

"""
    absolute_pressure(p0, h)

Calculate the absolute pressure [Pa] at altitude `h` [m] for a given sea level static pressure [Pa].
"""
function absolute_pressure(P0, h)
    R = 8.3144598 # universal gas constant
    M = 0.0289644 # molar mass of Earth's air
    g0 = 9.80665 # gravitational acceleration

    if h < 11000
        Tb = 288.15
        Pb = 101325.00
        hb = 0.0
        Lb = -0.0065 # standard temperature lapse rate (K/m)
    else
        error("not implemented for altitude $(h) m")
    end

    P0 * (Tb / (Tb + Lb * (h - hb)))^(g0 * M / R / Lb)
end

"""
    vapor_pressure_over_water(T)

Compute the saturation vapor pressure over liquid water in Pa given temperature `T` in °C.

# References
- Arden L Buck (1981) - New Equations for Computing Vapor Pressure and Enhancement Factor
- Arden L Buck (1996)
"""
function vapor_pressure_over_water(Tc)
    1000 * 0.61121 * exp((18.678 - Tc / 234.5) * (Tc / (257.14 + Tc)))
end


module QRDowndate
# copyright (c) 2017 Ralph Smith
# SPDX-License-Identifier: MIT-Expat
# License-Filename: LICENSE.md

export qrdowndate, qrdowndate!
using LinearAlgebra

# classical Gramm-Schmidt
function gs_step(U,::Type{Val{:CGS}},v)
    g = U' * v
    r = v - U * g
    β = norm(r,2)
    rmul!(r,1/β)
    r, g, β
end

# modified Gramm-Schmidt
function gs_step(U,::Type{Val{:MGS}},v)
    n = size(U,2)
    r = copy(v)
    g = zeros(eltype(U),n)
    for k=1:n
        # paper has a temp vector q; we just use r here
        ζ = dot(view(U,:,k),r)
        r = r - ζ * view(U,:,k)
        g[k] = ζ
    end
    β = norm(r,2)
    rmul!(r,1/β)
    r, g, β
end

# backwards modified Gramm-Schmidt
function gs_step(U,::Type{Val{:BMGS}},v)
    n = size(U,2)
    r = copy(v)
    g = zeros(eltype(U),n)
    for k=n:-1:1
        ζ = dot(view(U,:,k),r)
        r = r - ζ * view(U,:,k)
        g[k] = ζ
    end
    β = norm(r,2)
    rmul!(r,1/β)
    r, g, β
end

"""
Given `U[m,n]`, find unit vector `u0[m]` and vector `d[n+1]` such that
`U' * u0 = 0` and `hcat(u0,U) * d = e1`.

Algorithm from [^Barlow2005]

[Barlow2005]: Barlow et al., BIT 45, 259 (2005)
"""
function getorth(U::AbstractMatrix{T},M1,M2) where T
    m,n = size(U)
    @assert m >= n
    e1 = zeros(T,m); e1[1] = 1
    v1,g1,β1 = gs_step(U,M1,e1)
    βmin = sqrt(1/2)
    if β1 == 0
        β2 = α = zero(T)
        f = copy(g1)
        d = vcat(zeros(eltype(g1),1),g1)
        branch = 10
    else
        v2,g2,β2 = gs_step(U,M2,v1)
        if β1 >= βmin
            f = copy(g1)
            α = β1
            # β2 = zero(T) # early guess since paper fails to specify
            u0 = copy(v1)
            d = vcat([α],g1)
            branch = 20
        else
            f = g1 + β1 * g2
            α = β1 * β2
            if β2 >= βmin
                u0 = v2
                d = vcat([α],f)
            end
            branch = 30
        end
    end
    if β2 < βmin
        # FIXME: inelegant
        cval = Inf
        k = 0
        for j=1:m
            ek = zeros(m); ek[j] = 1
            tval = norm(U' * ek)
            if tval < cval
                k = j
                cval = tval
            end
        end

        ek = zeros(m); ek[k] = 1
        w1,k1,γ1 = gs_step(U,M1,ek)
        u0,k2,γ2 = gs_step(U,M2,w1)
        if α > 0
            d1 = vcat(zeros(T,1),f)
            Uhat = hcat(u0, U)
            # confusingly this is named $v_2$ in the paper,
            # but is NOT the same as the above v2.
            v2b = -U * f; v2b[1] += one(T)
            rmul!(v2,1/α)
            w2,d2,γ3 = gs_step(Uhat,M1,v2b)
            d = d1 + α * d2
            branch += 1
        else
            branch += 2
        end
    end

    return u0, d, branch
end

"""
    qrdowndate(Q,R) -> Qbar, Rbar

downdate the QR decomposition of a matrix (by the top row).

Given a QR decomposition `X=Q*R`, computes `Qbar, Rbar` such that
`Xbar = Qbar * Rbar` where `Xbar == X[2:end,:]`.
"""
function qrdowndate(U,R)
    Unew,Rnew = copy(U),copy(R)
    qrdowndate!(Unew,Rnew)
    Unew[2:end,:],Rnew
end

"""
    qrdowndate!(Q,R)

downdate the QR decomposition of a matrix.

Given a QR decomposition `X=Q*R`, computes `Qbar, Rbar` such that
`Xbar = Qbar * Rbar` where `Xbar == X[2:end,:]`, overwriting
the bottom of `Q` and `R`.
"""
function qrdowndate!(U::AbstractMatrix{T},R::AbstractMatrix{T}) where T
    m,n = size(U)
    (m > n) || throw(ArgumentError("U must be tall"))
    # check size of R?
    u0, d, b = getorth(U,Val{:MGS},Val{:BMGS})

    ρ = d[1]
    r = d[2:end]
    u = zeros(T,1,n)
    # find V st
    #     V is a product of Givens rotations
    #     V' * d = norm(d)*e1
    #     V' * vcat(zeros(1,n),R) = vcat(z',Rbar) for Up.Tri. Rhat
    # Ubar = hcat(u0,U) * V
    for k=n:-1:1
        G,ρ = givens(ρ,r[k],1,2)
        tmp = vcat(Transpose(u[k:n]),Transpose(R[k,k:n]))
        lmul!(G,tmp)
        u[k:n] = tmp[1,:]
        R[k,k:n] = tmp[2,:]
        tmp2 = hcat(u0,U[:,k])
        rmul!(tmp2,adjoint(G))
        u0 = tmp2[:,1]
        U[:,k] = tmp2[:,2]
    end
    # unceremoniously return indicator of branch used in orthogonalizer
    b
end

end


function solve()
    a, b, w = [parse(Int, x) for x in split(readline())]
    w = 1000w

    can = []
    for i=1:1000000
        if a*i ≤ w ≤ b*i
            push!(can, i)
        end
    end
    if length(can) == 0
        println("UNSATISFIABLE")
    else
        println(can[1], " " , can[end])
    end
end

solve()


## Residual Constrained Alternating Minimization
# Function Definitions

function AltMinBackslash(b,r,K)
# Performs RCAM on a single matricised monochromatic slice
# Calls: minFrob

    # Generate L0
    m,n = size(b)
    L = randn(m,r)

    local R

    # Begin Alterations
    for k = 1:K

        # Solve for R
        R = minFrob(L,b)

        # Solve for L
        L = minFrob(R,b')

    end

    return L, R
end
export AltMinBackslash

function minFrob(L,b)
# Solves for each row of input
# Calls: minFrobRsub

    m = size(b,2)
    r = size(L,2)

    # Initialize R
    R = zeros(m,r)
    if ~isreal(b)
        R = complex(R)
    end

    # Solve for each row of R
    for k=1:m

        R[k,:] = minFrobRsub(L,b,k)

    end

    R = conj(R)

    return R
end

function minFrobRsub(L,b,k)
# Solves a single row

    # Solve for the kth row of R
    bsub = b[:,k]
    ind = find(x -> x != 0, bsub)
    bsub = bsub[ind]
    L = L[ind,:]

    v = L\bsub

    return v

end


using Base.Test
using SparseRF
using ScikitLearnBase

srand(2)
n,m = 10^3, 5 ;
features = rand(n,m);
weights = rand(-1:1,m);
labels = features * weights;

model = fit!(DecisionTreeRegressor(maxlabels=5, pruning_purity_threshold=0.1), features, labels)
@test R2(labels, predict(model, features)) > 0.8

model = fit!(RandomForestRegressor(ntrees=10, maxlabels=5, nsubfeatures=2), features, labels)
@test R2(labels, predict(model, features)) > 0.8

srand(2)
N = 3000
X = randn(N, 10)
y = randn(N)
maxdepth = 5
model = fit!(DecisionTreeRegressor(maxdepth=maxdepth), X, y)
@test depth(model) == maxdepth

## Test that the RNG arguments work as expected
srand(2)
X = randn(100, 10)
y = randn(100)
@test fit_predict!(RandomForestRegressor(; rng=10), X, y) ==
    fit_predict!(RandomForestRegressor(; rng=10), X, y)
@test fit_predict!(RandomForestRegressor(; rng=10), X, y) !=
    fit_predict!(RandomForestRegressor(; rng=22), X, y)


"""
Constructs SeqModel object from config dict.
"""
function construct_model(config::Dict,
                         max_time::Float64,
                         num_neurons::Int64)

    # Prior on sequence type proportions / relative frequencies.
    seq_type_proportions = SymmetricDirichlet(
        config[:seq_type_conc_param],
        config[:num_sequence_types]
    )

    # Prior on expected number of spikes induces by a sequence events.
    seq_event_amplitude = specify_gamma(
        config[:mean_event_amplitude],    # mean of gamma; α / β
        config[:var_event_amplitude]      # variance of gamma; α / β²
    )

    # Prior on relative response amplitudes per neuron to each sequence type.
    neuron_response_proportions = SymmetricDirichlet(
        config[:neuron_response_conc_param],
        num_neurons
    )

    # Prior on the response offsets and widths for each neuron.
    neuron_response_profile = NormalInvChisq(
        config[:neuron_offset_pseudo_obs],
        0.0, # prior mean
        config[:neuron_width_pseudo_obs],
        config[:neuron_width_prior],
    )

    # Prior on expected number of background spikes in a unit time interval.
    bkgd_amplitude = specify_gamma(
        config[:mean_bkgd_spike_rate],    # mean of gamma; α / β
        config[:var_bkgd_spike_rate]      # variance of gamma; α / β²
    )

    # Prior on relative background firing rates across neurons.
    bkgd_proportions = SymmetricDirichlet(
        config[:bkgd_spikes_conc_param],
        num_neurons
    )

    model = SeqModel(
        # constants
        max_time,
        config[:max_sequence_length],

        # warp parameters
        config[:num_warp_values],
        config[:max_warp],
        config[:warp_variance],

        # priors
        config[:seq_event_rate],
        seq_type_proportions,
        seq_event_amplitude,
        neuron_response_proportions,
        neuron_response_profile,
        bkgd_amplitude,
        bkgd_proportions
    )

    if (:num_threads in keys(config)) && config[:num_threads] > 0
        return DistributedSeqModel(model, config[:num_threads])
    else
        return model
    end
end


"""
Caluclates and returns the admittance laplacian of the powergrid pg.
The notation was choosen in accordance to: https://en.wikipedia.org/wiki/Nodal_admittance_matrix
Inputs:
        pg: Powergrid, a graph contaning nodes and lines
"""
function admittance_laplacian(pg::PowerGrid)
    Y =  zeros(Complex, length(pg.nodes), length(pg.nodes))
    y_sum = zeros(Complex, length(pg.nodes))
    for i in 1:length(pg.lines)
        From = pg.lines[i].from
        To = pg.lines[i].to
        Y[From,To] = -1 * pg.lines[i].Y
        y_sum[From] += pg.lines[i].Y
    end
    for j in 1:length(pg.nodes)
        Y[j,j] = pg.nodes[j].Y_n + y_sum[j]
    end
    return Y
end

"""
A Simple Power Flow Calculation for PowerDynamics.
The calculation is only correct for nodes where the phase θ,
voltage magnitude u are know. For powergrids contaning eg. SlackAlgebraic the
calculation can not be performed because the phases θ are unknown.
A Power-flow study would need to be performed before hand.
Inputs:
    pg: PowerGrid, a graph containg nodes and lines
    op: State, Operationpoint (or any other valid point) of the powergrid pg
Outputs:
    P: Vector, Active Power at the nodes
    Q: Vector, Reactive Power at the nodes
"""
function SimplePowerFlow(pg::PowerGrid, op::State)
    #if SlackAlgebraic ∈ pg.nodes .|> typeof
    #    throw(OperationPointError("Found a SlackAlgebraic in the powergrid."))
    #end
    Y = admittance_laplacian(pg)
    G = real(Y)
    B = imag(Y)

    P = zeros(Complex, length(pg.nodes))
    Q = zeros(Complex, length(pg.nodes))

    θ_idx = findall(:θ .∈ symbolsof.(pg.nodes))
    for i in θ_idx
        Vi = op[i, :v] # Voltage Magnitude of node i
        for k in θ_idx
            Vk = op[k,:v] # Volatge Magnitude of node k
            θik = op[i, :θ] - op[k, :θ] # Phase difference between node i and k
            P[i] += Vk * (G[i,k] * sin(θik) + B[i,k] * sin(θik))
            Q[i] += Vk * (G[i,k] * cos(θik) - B[i,k] * cos(θik))
        end
        P[i] *= Vi
        Q[i] *= Vi
    end
    return P, Q
end


module HypergeoMat

export hypergeomPQ
export lmvgamma
export mvgamma
export BesselA

module HypergeomPQ

export hypergeomPQ
import LinearAlgebra

function DictParts(m::Integer, n::Integer)
  D = Dict{Int64,Int64}()
  Last = hcat([0, m, m])
  fin = 0
  for i = 1:n
    NewLast = Array{Int64}(undef, 3, 0)
    for record in eachcol(Last)
      @inbounds manque = record[2]
      @inbounds l = min(manque, record[3])
      if l > 0
        @inbounds D[record[1]+1] = fin + 1
        x = Array{Int64}(undef, 3, l)
        for j = 1:l
          @inbounds x[:, j] = [fin + j, manque - j, j]
        end
        NewLast = hcat(NewLast, x)
        fin += l
      end
    end
    Last = NewLast
  end
  return (D, fin)
end

function _N(dico::Dict{I,I}, lambda::Vector{I}) where {I<:Integer}
  lambda = filter(x -> x > 0, lambda)
  if isempty(lambda)
    return 1
  end
  @inbounds dico[_N(dico, lambda[1:end-1])] + last(lambda)
end

function _T(
  alpha::Real,
  a::Vector{<:Number},
  b::Vector{<:Number},
  kappa::Vector{<:Integer}
)
  if isempty(kappa) || kappa[1] == 0
    return 1
  end
  i = lastindex(kappa)
  @inbounds c = kappa[i] - 1 - (i - 1) / alpha
  @inbounds d = alpha * kappa[i] - i
  @inbounds s = collect(1:(kappa[i]-1))
  e = d .- alpha .* s + map(j -> count(>=(j), kappa), s)
  g = e .+ 1
  ss = 1:(i-1)
  @inbounds f = alpha .* kappa[ss] - collect(ss) .- d
  h = f .+ alpha
  l = h .* f
  prod1 = prod(a .+ c)
  prod2 = prod(b .+ c)
  prod3 = prod((g .- alpha) .* e ./ (g .* (e .+ alpha)))
  prod4 = prod((l - f) ./ (l + h))
  out = prod1 / prod2 * prod3 * prod4
  return isinf(out) || isnan(out) ? 0 : out
end

function betaratio(
  kappa::Vector{I},
  mu::Vector{I},
  k::I,
  alpha::T,
) where {I<:Integer,T<:Real}
  @inbounds muk = mu[k]
  t = k - alpha * muk
  @inbounds u = map(i -> t + 1 - i + alpha * kappa[i], 1:k)
  @inbounds v = map(i -> t - i + alpha * mu[i], 1:(k-1))
  muPrime = dualPartition(mu)
  @inbounds w = map(i -> muPrime[i] - t - alpha * i, 1:(muk-1))
  prod1 = prod(u ./ (u .+ alpha .- 1))
  prod2 = prod((v .+ alpha) ./ v)
  prod3 = prod((w .+ alpha) ./ w)
  return alpha * prod1 * prod2 * prod3
end

function dualPartition(lambda::Vector{T}) where {T<:Integer}
  out = T[]
  if !isempty(lambda)
    @inbounds for i in 1:lambda[1]
      push!(out, count(>=(i), lambda))
    end
  end
  return out
end

# ------------------------------------------------------------------------------
function hypergeomI(
  m::I,
  alpha::Real,
  a::Vector{<:Number},
  b::Vector{<:Number},
  n::I,
  x::Number,
)::Number where {I<:Integer}
  function summation(
    i::Int64,
    z::Number,
    j::I,
    kappa::Vector{I},
  )
    function go(kappai::I, zz::Number, s::Number)
      @inbounds if i == 0 && kappai > j || i > 0 && kappai > min(kappa[i], j)
        return s
      else
        kappap = vcat(kappa, [kappai])
        t = _T(alpha, a, b, filter(x -> x > 0, kappap))
        zp = zz * x * (n - i + alpha * (kappai - 1)) * t
        sp = j > kappai && i <= n ? s + summation(i + 1, zp, j - kappai, kappap) : s
        spp = sp + zp
        go(kappai + 1, zp, spp)
      end
    end # end go ---------------------------------------------------------------
    go(I(1), z, 0)
  end # end summation ----------------------------------------------------------
  return 1 + summation(0, 1, m, I[])
end

# ------------------------------------------------------------------------------
"""
    hypergeomPQ(m, a, b, x[, alpha])

Compute the truncated hypergeometric function of a matrix argument given the
  eigen values of the matrix.

# Arguments
- `m`: truncation weight of the summation, a positive integer
- `a`: the "upper" parameters, a real or complex vector, possibly empty
- `b`: the "lower" parameters, a real or complex vector, possibly empty
- `x`: real or complex vector, the eigen values
- `alpha`: the alpha parameter, a positive number; if missing, `alpha=2` is used
"""
function hypergeomPQ(
  m::Integer,
  a::Vector{<:Number},
  b::Vector{<:Number},
  x::Vector{<:Number},
  alpha::Union{Nothing,Real} = nothing,
)
  if isnothing(alpha)
    alpha = 2
  end
  n = length(x)
  @inbounds if all(x .== x[1])
    @inbounds return hypergeomI(m, alpha, a, b, n, x[1])
  end
  function jack(
    k::I,
    beta::Number,
    c::I,
    t::I,
    mu::Vector{I},
    jarray::Array{<:Number,2},
    kappa::Vector{I},
    nkappa::I,
  ) where {I<:Integer}
    lmu = length(mu)
    for i in ((max(k, 1):count(>(0), mu)))
      @inbounds u = mu[i]
      @inbounds if lmu == i || u > mu[i+1]
        gamma = beta * betaratio(kappa, mu, i, alpha)
        mup = copy(mu)
        @inbounds mup[i] = u - 1
        filter!(>(0), mup)
        if length(mup) >= i && u > 1
          jarray = jack(i, gamma, c + 1, t, mup, jarray, kappa, nkappa)
        else
          if nkappa > 1
            if !isempty(mup)
              @inbounds jarray[nkappa, t] +=
                gamma * jarray[_N(dico, mup)-1, t-1] * x[t]^(c + 1)
            else
              @inbounds jarray[nkappa, t] += gamma * x[t]^(c + 1)
            end
          end
        end
      end
    end
    if k == 0
      if nkappa > 1
        @inbounds jarray[nkappa, t] += jarray[nkappa, t-1]
      end
    else
      @inbounds jarray[nkappa, t] += beta * x[t]^c * jarray[_N(dico, mu)-1, t-1]
    end
    return jarray
  end # end jack ---------------------------------------------------------------
  function summation(
    i::I,
    z::Number,
    j::I,
    kappa::Vector{I},
    jarray::Array{<:Number,2},
  ) where {I<:Integer}
    function go(kappai::I, zp::Number, s::Number)
      if i == n || i == 0 && kappai > j || i > 0 && kappai > min(last(kappa), j)
        return s
      end
      kappap = vcat(kappa, [kappai])
      nkappa = _N(dico, kappap) - 1
      zpp = zp * _T(alpha, a, b, kappap)
      @inbounds if nkappa > 1 && (length(kappap) == 1 || kappap[2] == 0)
        @inbounds jarray[nkappa, 1] =
          x[1] * (1 + alpha * (kappap[1] - 1)) * jarray[nkappa-1, 1]
      end
      for t in 2:n
        jarray = jack(I(0), 1, I(0), t, kappap, jarray, kappap, nkappa)
      end
      @inbounds sp = s + zpp * jarray[nkappa, n]
      if j > kappai && i <= n
        spp = summation(i + 1, zpp, j - kappai, kappap, jarray)
        go(kappai + 1, zpp, sp + spp)
      else
        go(kappai + 1, zpp, sp)
      end
    end # end go ---------------------------------------------------------------
    go(I(1), z, 0)
  end # end summation ----------------------------------------------------------
  (dico, Pmn) = DictParts(m, n)
  T = typeof(x[1])
  J = zeros(T, Pmn, n)
  @inbounds J[1, :] = cumsum(x)
  return 1 + summation(0, T(1), m, Int64[], J)
end

# univariate -------------------------------------------------------------------
"""
    hypergeomPQ(m, a, b, x)

Compute the truncated hypergeometric function of a scalar argument.

# Arguments
- `m`: truncation weight of the summation, a positive integer
- `a`: the "upper" parameters, a real or complex vector, possibly empty
- `b`: the "lower" parameters, a real or complex vector, possibly empty
- `x`: scalar, real or complex
"""
function hypergeomPQ(
  m::Integer,
  a::Vector{<:Number},
  b::Vector{<:Number},
  x::Number
)
  return hypergeomI(m, 2, a, b, 1, x)
end

# matrix argument --------------------------------------------------------------
"""
    hypergeomPQ(m, a, b, X[, alpha])

Compute the truncated hypergeometric function of a matrix argument. The
  hypergeometric function is usually defined for a symmetric real matrix
  or a Hermitian complex matrix but arbitrary square matrices are allowed.

# Arguments
- `m`: truncation weight of the summation, a positive integer
- `a`: the "upper" parameters, a real or complex vector, possibly empty
- `b`: the "lower" parameters, a real or complex vector, possibly empty
- `X`: a square matrix, real or complex
- `alpha`: the alpha parameter, a positive number; if missing, `alpha=2` is used
"""
function hypergeomPQ(
  m::Integer,
  a::Vector{<:Number},
  b::Vector{<:Number},
  X::Matrix{<:Number},
  alpha::Union{Nothing,Real} = nothing
)
  x = LinearAlgebra.eigvals(X)
  return hypergeomPQ(m, a, b, x, alpha)
end

end # end module HypergeomPQ

module Mvgamma

export lmvgamma
export mvgamma
import GSL

"""
    lmvgamma(z, p)

Compute the logarithm of the multivariate Gamma function.

# Arguments
- `z`: real or complex number
- `p`: positive integer, the dimension
"""
function lmvgamma(
  z::Union{R,T},
  p::Integer
) where {R<:Real,T<:Complex{R}}
  if real(z) <= 0
    throw(DomainError(z, "The real part of `z` is not positive."))
  end
  C = p*(p-1)/4.0 * log(pi)
  isComplex = eltype(z) <: Complex
  if isComplex
    z_re = real(z)
    z_im = imag(z)
    (a, b) = GSL.sf_lngamma_complex_e(z_re, z_im)
    S = complex(a.val, b.val)
    for i in 2:p
      (a, b) = GSL.sf_lngamma_complex_e(z_re + (1 - i)/2.0, z_im)
      S = S + complex(a.val, b.val)
    end
  else
    S = GSL.sf_lngamma(z)
    for i in 2:p
      S = S + GSL.sf_lngamma(z + (1 - i)/2.0)
    end
  end
  return C + S
end

"""
    mvgamma(z, p)

Compute the multivariate Gamma function.

# Arguments
- `z`: real or complex number
- `p`: positive integer, the dimension
"""
function mvgamma(
  z::Union{R,T},
  p::Integer
) where {R<:Real,T<:Complex{R}}
  if imag(z) == 0 && real(z) <= 0 && trunc(z) == z
    throw(DomainError(z, "The argument `z` is a non-positive integer."))
  end
  if real(z) > 0.0
    result = exp(lmvgamma(z, p))
  else
    n = 1 + floor(-real(z))
    pochhammer = z
    for i in 2:n
      pochhammer = pochhammer * (z + i - 1.0)
    end
    result = exp(lmvgamma(z + n, p)) / pochhammer
  end
  return result
end

end # end module Mvgamma

module Bessel

export BesselA

"""
    BesselA(m, X, nu)

Compute the truncated Herz's type one Bessel function of a matrix argument. It
  is usually defined for a symmetric real matrix or a Hermitian complex matrix
  but arbitrary square matrices are allowed.

# Arguments
- `m`: truncation weight of the hypergeometric function, a positive integer
- `X`: a square matrix, real or complex
- `nu`: the order parameter, real or complex number with `real(nu)>-1`
"""
function BesselA(
  m::Integer,
  X::Matrix{Number},
  nu::Number
)
  if real(nu) <= -1
    throw(DomainError(nu, "The real part of `nu` is smaller than -1."))
  end
  p = size(X)[1]
  b = nu + (p+1) / 2
  return hypergeomPQ(m, Int64[], [b], -X, 2) / mvgamma(b, p)
end

"""
    BesselA(m, x, nu)

Compute the truncated Herz's type one Bessel function of a matrix argument
  given the eigen values of the matrix.

# Arguments
- `m`: truncation weight of the hypergeometric function, a positive integer
- `x`: the eigen values, a vector of real or complex numbers
- `nu`: the order parameter, real or complex number with `real(nu)>-1`
"""
function BesselA(
  m::Integer,
  x::Vector{<:Number},
  nu::Number
)
  if real(nu) <= -1
    throw(DomainError(nu, "The real part of `nu` is smaller than -1."))
  end
  p = length(x)
  b = nu + (p+1) / 2
  return hypergeomPQ(m, Int64[], [b], -x, 2) / mvgamma(b, p)
end

using ..HypergeomPQ
using ..Mvgamma

end # end module Bessel

using .HypergeomPQ
using .Mvgamma

end


function acqr_kalman(y,m::acqr;s::Bool=true)

	# parameters
	a = m.a
	c = m.c
	q = m.q
	r = m.r
	m1 = m.m1
	P1 = m.P1

    if s
        # stationary
        (xtt,Ptt,xtt1,Ptt1,et,St,Kt,loglik) = acqr_kfilter_s(y,a,c,q,r,m1)
		(xtN,PtN,Pt1tN) = acqr_ksmoother_s(a,xtt,Ptt,xtt1,Ptt1)
	else
		(xtt,Ptt,xtt1,Ptt1,et,St,Kt,loglik) = acqr_kfilter(y,a,c,q,r,m1,P1)
		(xtN,PtN,Pt1tN) = acqr_ksmoother(a,xtt,Ptt,xtt1,Ptt1)
	end

	salida = acqrKal(y,m,xtt1,Ptt1,xtt,Ptt,xtN,PtN,s)

	return salida
end


using JuBrain
using Base.Test

model=NetworkModel()

model.parameters="""
τᵥ = 20*ms;  τₑ = 5*ms;   τᵢ = 10*ms
Vt = -50*mV; Vᵣ = -60*mV; El = -49*mV"""
model.equations="dv/dt=(gₑ+gᵢ-(v-El))/τᵥ;dgₑ/dt=-gₑ/τₑ;dgᵢ/dt=-gᵢ/τᵢ"
model.spike="v.>Vt"
model.reset="v=Vᵣ"
model.synapse="gₑ+=wₑ;gᵢ+=wᵢ"""

model.groups=makeGroups(groupName=["e","i"],groupSize=[3200,800])
initialize(model,groups="all",expr="v=Vᵣ+rand()*(Vt-Vᵣ)")
spConnect(model,pre="e",post="all",expr="wₑ=(60*0.27/10)*mV",p=0.02)
spConnect(model,pre="i",post="all",expr="wᵢ=(-20*4.5/10)*mV",p=0.02)

buildNetwork(name="temp",model=model,duration=1000)


using Revise
using SoapySDR
using GLMakie
using LinearAlgebra
using FFTW
using DSP

# Don't forget to add/import a device-specific plugin package!
# using xtrx_jll
# using SoapyLMS7_jll
using SoapyRTLSDR_jll
# using SoapyPlutoSDR_jll
# using SoapyUHD_jll

# Here we want to test the behavior in a loop, to ensure
# that we can block on buffer overflow conditions, and
# handle partial reads, measure latnecy, etc

freq = 104.3e6u"Hz"

function makie_viz(freq=freq)
    dev = open(Devices()[1])
    rx_chan = dev.rx[1]
    @show rx_chan
    rx_chan.gain_mode = true
    rx_chan.frequency = freq
    rx_stream = SoapySDR.Stream(ComplexF32, [rx_chan])
    @show rx_stream.mtu
    buf = Vector{ComplexF32}(undef, rx_stream.mtu)
    @show typeof(fft)
    sum = zeros(Float32, length(buf))
    norms = Observable(Vector{Float32}(undef,length(buf)))

    println("Makie Setup...")
    fig = Figure(); display(fig)
    ax = Axis(fig[1,1])
    xlims!(ax, 0, 100)
    ylims!(ax, -500, 500)
    lines!(ax, norms)

    h = Float32.(DSP.hamming(length(buf)))

    println("Starting fft loop..")
    Base.atexit(()->SoapySDR.deactivate!(rx_stream))

    SoapySDR.activate!(rx_stream)
    k = 1
    while true
        read!(rx_stream, (buf,))
        buf = fftshift(FFTW.fft!(buf.*h))
        norms[] .= (norm.(buf))
        #norms[] .= sum ./ k
        norms[] = norms[]
        println("fft") # makie needs a slight delay here
        k += 1
    end

end


using TensorTrains
using Test
using Random

function ttcheckerr(Y1, Y2)
    return norm(add(1.0, Y1, -1.0, Y2)) < length(Y1) * eps(Float64)*max(norm(Y1),norm(Y2))
end

@testset "TensorTrains" begin
    @testset "dot prod" begin
        T1 = ttones(3); T2 = ttones(3)
        @test dot(T1,T2) ≈ 8.0
        T1 = ttrandn(5); T2 = ttrandn(5)
        @test dot(T1,T2) ≈ vec(decompress(T1))'*vec(decompress(T2))
        @test dot(T1,T2) ≈ vec(unfold(T1))'*vec(unfold(T2))
    # @test dot(ttones)
    end;
    @testset "norm" begin
        T1 = ttones(3);
        @test norm(T1) ≈ √8
        T1 = ttrandn(5);
        @test norm(T1)  ≈ √dot(T1,T1)
        @test norm!(T1) ≈ √dot(T1,T1)
    end;
    @testset "add" begin
        T1 = ttones(3);
        add!(T1, ttones(3));
        @test norm(T1) ≈ 2*√8;
        add!(T1, -1.0, ttones(3));
        @test norm(T1) ≈ √8;
        T2 = add(2.0, T1, -1.0, T1);
        @test norm(T2) ≈ norm(T1);
    end
    @testset "orthogonality" begin
        T1 = ttones(4);
        add!(T1, ttones(4));
        add!(T1, ttones(4));
        rightOrth!(T1);
        @test norm(T1[end]) ≈ norm(T1);
        leftOrth!(T1);
        @test norm(T1[1]) ≈ norm(T1);
    end
    @testset "svd" begin
        T1 = ttones(4);
        add!(T1, ttones(4));
        add!(T1, ttones(4));
        svdtrunc!(T1, svd!(T1), 1e-16);
        @test ranks(T1) == [1,1,1];
        Random.seed!(1234);
        T1 = ttrandn(8);
        add!(T1, ttrandn(8));
        add!(T1, ttrandn(8));
        T2 = deepcopy(T1);
        svdtrunc!(T1, svd!(T1), 1e-1);
        @test norm(add(1.0, T1, -1.0, T2)) < 2e-1;
        T1 = deepcopy(T2);
        svdtrunc!(T1, svd!(T1), 2e-1);
        @test norm(add(1.0, T1, -1.0, T2)) < 4e-1;
    end
    @testset "matvec" begin
        Random.seed!(1234);
        d = 8;
        Id = tteye(d);
        T1 = ttrandn(d);
        @test norm(T1) ≈ norm(matvec(Id,T1))
        Y1 = matvec(add(1.0,Id,1.0,Id),T1);
        Y2 = matvec(scale(Id,2.0),T1);
        @test ttcheckerr(Y1,Y2);
        A1 = deepcopy(Id);
        add!(A1, 2.5, Id);
        Y1 = matvec(A1, T1);
        A2 = deepcopy(Id);
        scale!(A2, 3.5);
        Y2 = matvec(A2, T1);
        @test ttcheckerr(Y1,Y2);
    end
    @testset "matmat" begin
        Random.seed!(1234);
        d = 8;
        Id = tteye(d);
        A1 = deepcopy(Id);
        scale!(A1,Float64(π));
        T1 = ttrandn(d);
        A2 = matmat(Id,A1);
        Y1 = matvec(A1,T1);
        Y2 = matvec(A2,T1);
        @test ttcheckerr(Y1,Y2);
        A2 = matmat(A2,Id);
        Y2 = matvec(A2,T1);
        @test ttcheckerr(Y1,Y2);
    end
    @testset "hadamard" begin
        Random.seed!(1234);
        T1 = ttones(6);
        T2 = ttrandn(6);
        @test norm(hadamard(T1,T2)) ≈ norm(T2)
    end
    @testset "kron" begin
        Random.seed!(1234);
        A1 = tteye(8);
        A2 = tteye(8);
        T1 = ttrandn(8);
        T2 = ttrandn(8);
        scale!(A1, 0.5);
        scale!(A2, Float64(π));
        Y1 = matvec(ttkron(A1,A2),ttkron(T1,T2));
        Y2 = ttkron(matvec(A1,T1),matvec(A2,T2));
        @test ttcheckerr(Y1,Y2);
    end
end;


s₁ = Scenario(q₁ = -24.0, q₂ = -28.0, d₁ = 500.0, d₂ = 100.0, probability = 0.4)
s₂ = Scenario(q₁ = -28.0, q₂ = -32.0, d₁ = 300.0, d₂ = 300.0, probability = 0.6)
deferred = StochasticProgram([s₁,s₂], solver=GLPKSolverLP())

@first_stage deferred = begin
    @variable(model, x₁ >= 40)
    @variable(model, x₂ >= 20)
    @objective(model, Min, 100*x₁ + 150*x₂)
    @constraint(model, x₁+x₂ <= 120)
end defer

@second_stage deferred = begin
    @decision x₁ x₂
    @uncertain q₁ q₂ d₁ d₂
    @variable(model, 0 <= y₁ <= d₁)
    @variable(model, 0 <= y₂ <= d₂)
    @objective(model, Min, q₁*y₁ + q₂*y₂)
    @constraint(model, 6*y₁ + 10*y₂ <= 60*x₁)
    @constraint(model, 8*y₁ + 5*y₂ <= 80*x₂)
end defer


"""
    diffusion_time_sta(f)

Get the short term approximation of the diffusion time of the time profile `f`.
"""
function diffusion_time_sta end

function diffusion_time_sta(f::PGSE)
    (; δ, Δ) = f
    out =
        4 / 35 * ((Δ + δ)^(7 / 2) + (Δ - δ)^(7 / 2) - 2Δ^(7 / 2) - 2δ^(7 / 2)) /
        (δ^2 * (Δ - δ / 3))
    out^2
end

function diffusion_time_sta(f::DoublePGSE)
    (; δ, Δ, tpause) = f
    tm = δ + tpause
    out =
        2 / 35 * (
            (2Δ + tm + δ)^(7 / 2) +
            (2Δ + tm - δ)^(7 / 2) +
            (tm + δ)^(7 / 2) +
            (tm - δ)^(7 / 2) - 2(2Δ + tm)^(7 / 2) - 2(Δ + tm + δ)^(7 / 2) -
            2(Δ + tm - δ)^(7 / 2) +
            2(Δ + δ)^(7 / 2) +
            2(Δ - δ)^(7 / 2) - 2tm^(7 / 2) + 4(Δ + tm)^(7 / 2) - 4Δ^(7 / 2) - 4δ^(7 / 2)
        ) / (δ^2 * (Δ - δ / 3))
    out^2
end

function diffusion_time_sta(f::CosOGSE)
    (; δ, Δ) = f
    n = f.nperiod
    S = sin(2π * n * Δ / δ)
    C = cos(2π * n * Δ / δ)
    FS1 = fresnels(2√n)
    FS2 = fresnels(2√(Δ * n / δ))
    FS3 = fresnels(2√((Δ + δ) * n / δ))
    FS4 = fresnels(2√((Δ - δ) * n / δ))
    FC1 = fresnelc(2√n)
    FC2 = fresnelc(2√(Δ * n / δ))
    FC3 = fresnelc(2√((Δ + δ) * n / δ))
    FC4 = fresnelc(2√((Δ - δ) * n / δ))
    out =
        3 / 8 / √(n * δ) * (
            2Δ * FC2 * C + δ * FC4 * C - δ * FC3 * C - Δ * FC3 * C - Δ * FC4 * C +
            2Δ * FS2 * S +
            δ * FS4 * S - δ * FS3 * S - Δ * FS3 * S - Δ * FS4 * S + 2δ * FC1
        ) +
        9√δ / 32π / n^(3 / 2) *
          (+2FS2 * C - FS3 * C - FS4 * C - 2FC2 * S + FC3 * S + FC4 * S + 2FS1)
    out^2
end

function diffusion_time_sta(f::SinOGSE)
    (; δ, Δ) = f
    n = f.nperiod
    S = sin(2π * n * Δ / δ)
    C = cos(2π * n * Δ / δ)
    FS1 = fresnels(2√n)
    FS2 = fresnels(2√(Δ * n / δ))
    FS3 = fresnels(2√((Δ + δ) * n / δ))
    FS4 = fresnels(2√((Δ - δ) * n / δ))
    FC1 = fresnelc(2√n)
    FC2 = fresnelc(2√(Δ * n / δ))
    FC3 = fresnelc(2√((Δ + δ) * n / δ))
    FC4 = fresnelc(2√((Δ - δ) * n / δ))
    out =
        (
            64π * n^(3 / 2) * Δ^(3 / 2) +
            64π * n^(3 / 2) * δ^(3 / 2) +
            42δ^(3 / 2) * FS2 * C +
            42δ^(3 / 2) * FS1 - 42δ^(3 / 2) * FC2 * S + 32π * n^(3 / 2) * √(Δ - δ) * δ -
            32π * n^(3 / 2) * √(Δ + δ) * Δ - 32π * n^(3 / 2) * √(Δ + δ) * δ -
            32π * n^(3 / 2) * √(Δ - δ) * Δ +
            24π * n * Δ * √δ * FC2 * C +
            24π * n * Δ * √δ * FS2 * S +
            24π * n * δ^(3 / 2) * FC1 +
            21δ^(3 / 2) * FC3 * S +
            21δ^(3 / 2) * FC4 * S - 21δ^(3 / 2) * FS3 * C - 21δ^(3 / 2) * FS4 * C +
            12π * n * δ^(3 / 2) * FS4 * S +
            12π * n * δ^(3 / 2) * FC4 * C - 12π * n * δ^(3 / 2) * FC3 * C -
            12π * n * δ^(3 / 2) * FS3 * S - 12π * n * Δ * √δ * FC4 * C -
            12π * n * Δ * √δ * FC3 * C - 12π * n * Δ * √δ * FS3 * S -
            12π * n * Δ * √δ * FS4 * S
        ) / (96π * n^(3 / 2) * δ)
    out^2
end


function FVort2VecDSS(v1,v2,CG,Param)
OP=CG.OrdPoly+1;
NF=Param.Grid.NumFaces;
nz=Param.Grid.nz;
dXdxIC = Param.cache.dXdxIC
v1CG=reshape(v1[reshape(CG.Glob,OP*OP*NF,1),:]
  ,OP,OP,NF,nz);
v2CG=reshape(v2[reshape(CG.Glob,OP*OP*NF,1),:]
  ,OP,OP,NF,nz);
vC1=reshape(
  CG.DS*reshape(dXdxIC[:,:,:,:,1,1].*v2CG -
  dXdxIC[:,:,:,:,1,2].*v1CG
  ,OP,OP*NF*nz)
  ,OP,OP,NF,nz) -
  permute(
  reshape(
  CG.DS*reshape(
  permute(-v2CG.*dXdxIC[:,:,:,:,2,1] +
  v1CG.*dXdxIC[:,:,:,:,2,2]
  ,[2 1 3 4])
  ,OP,OP*NF*nz)
  ,OP,OP,NF,nz)
  ,[2 1 3 4]);

Vort=zeros(CG.NumG,nz);
for iM=1:size(CG.FaceGlob,1)
  i = CG.FaceGlob[iM].Ind
  arr = reshape(CG.Glob[:,i,:],OP*OP*size(i,1))
  Vort[arr,:] = Vort[arr,:] .+
    reshape(vC1[:,:,i,:]
    ,OP*OP*size(i,1),nz);
end
Vort=Vort./CG.M;
return Vort
end


# using Elliptic
"""
This function calculates the inductance for an N array of coaxial wire loops

The input is an N x 2 array of coordinates [R₁ Θ₁;...;Rₙ Θₙ]
The output is the overall inductance including any mutual inductances

For strange arrangements of loops, this may be useful, but for simple solenoids, other methods are desireable
Empirical formulas (``SimpleInduct'') will be more accurate and less cumbersome in those cases.
"""
function CalcInductance(Coords)
    #Thus function takes in two vectors listing the centroid [in polar coords]
    #of series-wound strands of wire to form a coil. It also takes in the index
    #(of the vector) that different coil "segments" are wound from. An example
    #being if you wanted to break up a Helmholtz pair into the two halves to
    #look at the self-inductances & mutual inductances of each independently
    #you would mark the middle index value in Segment Index List

    for i = 1:(length(Coords[:, 1])-1)
        println(i)
        for j = (i+1):length(Coords[:, 1])
            if Coords[j, :] == Coords[i, :]
                Coords[j, :] = Coords[j, :] .+ 1e-8
            end
        end
    end

    R_List = Coords[:, 1]
    Z_List = Coords[:, 2]
    L = length(R_List)

    μ₀ = 4 * π * 1e-7

    RMesh, AMesh = MPI_Tools.meshgrid(R_List,dims=2)

    KVal = similar(RMesh)
    EValue = similar(RMesh)

    KVal, EValue = (
        [ellipke.(fFunk(R_List, Z_List) .^ 2)[i, j][1] for i = 1:L, j = 1:L],
        [ellipke.(fFunk(R_List, Z_List) .^ 2)[i, j][2] for i = 1:L, j = 1:L],
    )

    M =
        μ₀ .* sqrt.(RMesh .* AMesh) .* (
            (2 .* (KVal .- EValue)) ./ fFunk(R_List, Z_List) .-
            fFunk(R_List, Z_List) .* KVal
        ) #https://thompsonrd.com/OSEE-inductance.pdf
    TotalL = sum(M)

    return TotalL
end

function fFunk(r,z)

    RMesh, AMesh = MPI_Tools.meshgrid(r,dims=2)
    ZMesh_A, ZMesh_B = MPI_Tools.meshgrid(z,dims=2)
    ZDelta = ZMesh_A .- ZMesh_B
    out = sqrt.( (4 .* RMesh.*AMesh ) ./ (ZDelta.^2 .+ (RMesh .+ AMesh).^2) )
    replace!(x -> x>=1 ? 1-eps() : x,out)

    return out
end

"""
SIMPLE INDUCTANCE FORMULAS FOR RADIO COILS*
BY HAROLD A. WHEELER
This is an empirical formula for inductance
    Takes in turns(N), radius (a), and length (b)
    the formulation was designed for inches, but by default
    this function assumes meters.
"""
function SimpleInduct(N,a,b;units = "meters")
    if units=="meters"
        a = a/25.4e-3
        b = b/25.4e-3
    elseif units == "inches"
    else
    error("unknown units")
    end

    L = (a^2*N^2 / (9*a + 10*b))*1e-6
end


# This file defines the reflection and transmission coefficients of a multilayer

abstract type Polarization end
struct te <: Polarization end
struct tm <: Polarization end

abstract type Structure end

struct Bulk{T <:OptProp,U <:OptProp} <: Structure
    ep1 :: T
    ep2 :: U
end
# Outer constructor for the Bulk type when one medium is vacuum
Bulk(ep2) = Bulk(Cst(1.0+im*0.0),ep2)

struct Layer <: Structure
    material  :: OptProp
    thickness :: Float64 # in meter. For semi-infinite media thickness = 0
end

Layer(material)  = Layer(material,0.0)

struct Film <: Structure
    film :: Layer
    substrate :: Layer
end

# new Vector type
const MultiLayer = Vector{Layer}

# Union types
const LayerOrMultiLayer = Union{Layer,MultiLayer}

# Accessor functions
substrate(structure :: Bulk) = Layer(structure.ep2)
substrate(structure :: MultiLayer) = structure[end]
substrate(structure :: Film) = structure.substrate

top_medium(structure :: Bulk) = Layer(structure.ep1)
top_medium(structure :: MultiLayer) = structure[1]
top_medium(structure :: Film) = Layer(Cst())

OpticalProperties.permittivity(layer::Layer,w) = permittivity(layer.material,w)

# Normal functions

function compute_kz(kx,eps,w)
    sqrt(eps*(w/c0)^2 - kx*kx)
end

function compute_kx(theta_0,eps,w)
    sqrt(eps)*sin(theta_0)*w/c0
end

function compute_ang(kx,eps,w)
    asin(real(kx)/real(sqrt(eps))/(w/c0))
end

function outer_media(b,kx,w)
    eps1 = permittivity(top_medium(b),w)
    k1z  = compute_kz(kx,eps1,w)

    eps2 = permittivity(substrate(b),w)
    k2z  = compute_kz(kx,eps2,w)
    return eps1,k1z,eps2,k2z
end

# Fresnel coefficient of a semi-infinite medium : perpendicular
function rt(pol :: te, eps1,eps2, k0z,k2z,w)
  # k0z = incident wavevector from medium 1
    r = (k0z - k2z)/(k0z + k2z)
    t = (2.0+im*0.0)*k0z/(k0z + k2z)
    return r,t
end

# Fresnel coefficient of a semi-infinite medium : parallel
function rt(pol :: tm, eps1,eps2, k0z,k2z,w)
  # k0z = incident wavevector from medium 1
    r = (eps2*k0z-eps1*k2z)/(eps2*k0z+eps1*k2z)
    t = (2.0+im*0.0)*k0z/(eps2*k0z+eps1*k2z)*sqrt(eps1)*sqrt(eps2)
    return r,t
end

function power_t(pol::te,eps1,eps2, k0z,k2z,t)
    real(k2z)/real(k0z)*abs(t)^2
end

function power_t(pol::tm,eps1,eps2, k0z,k2z,t)
    real(k2z/eps2)/real(k0z/eps1)*abs(eps2/eps1)*abs(t)^2
end

# Fresnel coefficient of a semi-infinite medium : as a function of incidence angle and frequency
function rt(structure :: Bulk, pol :: Polarization, kx ,w)
    eps1 = permittivity(structure.ep1,w)
    eps2 = permittivity(structure.ep2,w)

    k0z   = compute_kz(kx,eps1,w)
    k2z   = compute_kz(kx,eps2,w)

    r,t=rt(pol,eps1,eps2,k0z,k2z,w)
    return r, t
end

# Fresnel coefficient of a multilayered semi-infinite medium : as a function of incidence angle and frequency

function rt(structure :: MultiLayer, pol :: Polarization, kx ,w)
      S   =[1.0+0.0*im  0.0+0.0*im ;
              0.0+0.0*im  1.0+0.0*im ]
    if real(kx)<w/c0 || real(kx)>w/c0
        scattering_matrix!(S,structure,pol,kx,w)
        return S[2,1], S[1,1]
    else
        return 1.0+0.0*im, 0.0+0.0*im
  end
end

function scattering_matrix!(S,structure :: MultiLayer, pol :: Polarization, kx, w)
        for i=2:length(structure)
            eps1 :: Complex{Float64} = permittivity(structure[i-1].material,w)
            eps2 :: Complex{Float64} = permittivity(structure[i].material,w)
            k0z   = compute_kz(kx,eps1,w)
            k2z   = compute_kz(kx,eps2,w)
            r,t = rt(pol, eps1,eps2, k0z,k2z,w)

             S[1,1] =(S[1,1]*t*exp(im*structure[i-1].thickness*k0z))/(1.0+0.0*im - S[1,2]*r*exp(2.0*im*structure[i-1].thickness*k0z))
             S[1,2] =(S[1,2]*exp(2.0*im*structure[i-1].thickness*k0z)-r)/(1.0+0.0*im - S[1,2]*r*exp(2.0*im*structure[i-1].thickness*k0z))
             S[2,1] =(S[1,1]*S[2,2]*r*exp(im*structure[i-1].thickness*k0z))/t+S[2,1]
             S[2,2] =(S[2,2]*exp(im*structure[i-1].thickness*k0z)*(r*S[1,2] + 1.0+0.0*im))/t
        end
    return nothing
end

function rt(structure :: Film , pol :: Polarization , kx, w)
    filmmaterial      = structure.film.material
    substratematerial = structure.substrate.material
    th                = structure.film.thickness
    eps2 = permittivity(Cst(),w)
    eps1 = permittivity(filmmaterial,w)
    eps0 = permittivity(substratematerial,w)

    k0z = compute_kz(kx,eps0,w)
    k1z = compute_kz(kx,eps1,w)
    k2z = compute_kz(kx,eps2,w)

    r10 ,t10 = rt(pol, eps1,eps0, k1z,k0z,w)
    r21 ,t21 = rt(pol, eps2,eps1, k2z,k1z,w)

    r = (r21 + r10*exp(2*im*k1z*th))/(1 + r21*r10*exp(2*im*k1z*th))
    t = (t21*t10*exp(im*k1z*th))/(1 + r21*r10*exp(2*im*k1z*th))

    return r,t
end

function power_rt(structure, pol :: Polarization, kx ,w)

    eps1,k1z,eps2,k2z = outer_media(structure,kx,w)

    r,t=rt(structure,pol,kx,w)
    R = abs(r)^2
    T = power_t(pol,eps1,eps2, k1z,k2z,t)

    return R, T
end


using Smile
using Base.Test

function test_vec_approx_equal{T<:Real, S<:Real}(a::Vector{T}, b::Vector{S})
	@test length(a) == length(b)
	for i = 1 : length(a)
		@test isapprox(a[i], b[i])
	end
	true
end
function test_vec_approx_equal{T<:Real, S<:Real}(a::Vector{T}, b::Vector{S}, tol::Float64)
	@test length(a) == length(b)
	for i = 1 : length(a)
		@test abs(a[i] - b[i]) < tol
	end
	true
end

# write your own tests here
include("test_intarray.jl")
include("test_pattern.jl")
include("test_dataset.jl")
include("test_learning.jl")
include("test_sampling.jl")
include("tut1_create_a_bayesian_network.jl")
include("high_level_learning.jl")


"""
intck(x,y) computes the difference between two date objects and return number of day. Any of x and y can be missing value.
"""
intck(x::Missing,y::Missing) = missing
intck(x::Missing,y) = missing
intck(x,y::Missing) = missing
intck(x::Date,y::Date) = (x-y).value

"""
rescale(x,minx,maxx,minval,maxval) rescales x to run from minval and maxval, given x originaly runs from minx to maxx.
"""
function rescale(x,minx,maxx,minval,maxval)
    -(-maxx*minval+minx*maxval)/(maxx-minx)+(-minval+maxval)*x/(maxx-minx)
end
rescale(x::Missing,minx,maxx,minval,maxval) = missing
rescale(x::Vector,minx,maxx,minval,maxval) = rescale.(x,minx,maxx,minval,maxval)
rescale(x,minx,maxx) = rescale(x,minx,maxx,0.0,1.0)

"""
stdze(x) standardizes an array. It return missing for missing data points.
"""
function stdze(x)
    all(ismissing,x) && return x
    meandata = mean(skipmissing(x))
    vardata = var(skipmissing(x))
    (x .- meandata) ./ sqrt(vardata)
end

"""
lag(x,k) Creates a lag-k of the provided array x. The output will be an array the
same size as x (the input array), and the its type will be Union{Missing, T} where T is the type of input.
"""
function lag(x,k)
    res=zeros(Union{typeof(x[1]),Missing},length(x))
    for i in 1:k
        res[i] = missing
    end
    for i in k+1:length(x)
        res[i] = x[i-k]
    end
    res
end

lag(x)=lag(x,1)

"""
lead(x,k) Creates a lead-k of the provided array x. The output will be an array the
same size as x (the input array), and the its type will be Union{Missing, T} where T is the type of input.
"""
function lead(x,k)
    res=zeros(Union{typeof(x[1]),Missing},length(x))
    for i in 1:length(x)-k
        res[i] = x[i+k]
    end
    for i in length(x)-k+1:length(x)
        res[i] = missing
    end
    res
end
lead(x)=lead(x,1)

"""
dttodate(x) converts SAS or STATA dates (which is the number of days after 01-01-1960) to a Julia Date object.
dttodate(DataFrame,cols) converts the given columns to Date object.
"""
dttodate(x::Missing) = missing
dttodate(x) = Date(1960,1,1)+Day(x)
dttodate(x::Date) = x
function dttodate!(df::DataFrame,cols)
    for i in cols
        df[!,i] = dttodate.(df[!,i])
    end
end

# modifying some Base functions to suit for working with data with missing values

Base.maximum(x::AbstractArray{Missing,1})=missing
function Base.maximum(x::AbstractArray{Union{T,Missing},1}) where {T <: Number}
    res=typemin(T)
    cnt_miss=0
    for i in 1:length(x)
        if !ismissing(x[i])
            if x[i]>res
                res=x[i]
            end
        else
            cnt_miss+=1
        end
    end
    if cnt_miss==length(x)
        res=missing
    end
    res
end
Base.maximum(x::AbstractArray{Union{Date,Missing},1})=maximum(skipmissing(x))

Base.minimum(x::AbstractArray{Missing,1})=missing
function Base.minimum(x::AbstractArray{Union{T,Missing},1}) where {T <: Number}
    res=typemax(T)
    cnt_miss=0
    for i in 1:length(x)
        if !ismissing(x[i])
            if res>x[i]
                res=x[i]
            end
        else
            cnt_miss+=1
        end
    end
    if cnt_miss==length(x)
        res=missing
    end
    res
end
Base.minimum(x::AbstractArray{Union{Date,Missing},1})=minimum(skipmissing(x))

Base.sum(x::AbstractArray{Missing,1})=missing
function Base.sum(x::AbstractArray{Union{T,Missing},1}) where {T <: Number}
    res=zero(T)
    cnt_miss=0
    for i in 1:length(x)
        if !ismissing(x[i])
            res+=x[i]
        else
            cnt_miss+=1
        end
    end
    if cnt_miss==length(x)
        res=missing
    end
    res
end

Statistics.mean(x::AbstractArray{Missing,1})=missing
function Statistics.mean(x::AbstractArray{Union{T,Missing},1}) where {T <: Number}
    res=0.0
    cnt_nonmiss=0
    for i in 1:length(x)
        if !ismissing(x[i])
            res+=x[i]
            cnt_nonmiss+=1
        end
    end
    if cnt_nonmiss>0
        res=res/cnt_nonmiss
    else
        res=missing
    end
    res
end

Statistics.var(x::AbstractArray{Missing,1},df=true)=missing
function Statistics.var(x::AbstractArray{Union{T,Missing},1},df=true) where {T <: Number}
    res=0.0
    ss=0.0
    sval=0.0
    cnt_nonmiss=0
    for i in 1:length(x)
        if !ismissing(x[i])
            sval+=x[i]
            ss+=x[i] * x[i]
            cnt_nonmiss+=1
        end
    end
    if cnt_nonmiss>1
        res=ss/cnt_nonmiss - (sval/cnt_nonmiss)*(sval/cnt_nonmiss)
        if df
            res=(cnt_nonmiss/(cnt_nonmiss-1))*res
        end
    elseif cnt_nonmiss==1
        res=0.0
    else
        res=missing
    end
    res
end
Statistics.std(x::AbstractArray{Missing,1},df=true)=missing
function Statistics.std(x::AbstractArray{Union{T,Missing},1},df=true) where {T <: Number}
    sqrt(var(x,df))
end

function Statistics.median(v::AbstractArray{T,1}) where T
    isempty(v) && throw(ArgumentError("median of an empty array is undefined, $(repr(v))"))
    eltype(v)>:Missing && all(ismissing, v) && return missing
    (eltype(v)<:AbstractFloat || eltype(v)>:AbstractFloat) && any(isnan, v) && return convert(eltype(v), NaN)
    any(ismissing,v) ? v2=collect(skipmissing(v)) : v2=v
    inds = axes(v2, 1)
    n = length(inds)
    mid = div(first(inds)+last(inds),2)
    if isodd(n)
        return middle(partialsort!(v2,mid))
    else
        m = partialsort!(v2, mid:mid+1)
        return middle(m[1], m[2])
    end
end

function Statistics.quantile(x::Array{T,1},v) where T
    all(ismissing,x) && return missing
    quantile(skipmissing(x),v)
end


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

# %%
using Plots
default(size=(400, 250))

Base.:^(f::Function, n::Integer) = n == 0 ? identity : foldr(∘, (f for _ in 1:n))

# %%
f(x) = sin(x)
plot(f, 0, 2; label="") |> display

plot([n^(1/2)*(f^n)(1.0) for n in 0:10:200]; label="")
hline!([√3]; label="")

# %%
f(x) = x - x^3/10
plot(f, 0, 2; label="") |> display

plot([n^(1/2)*(f^n)(1.0) for n in 0:10:200]; label="")

# %%
f(x) = x - x^2/4
plot(f, 0, 2; label="") |> display

plot([n*(f^n)(1.0) for n in 0:25:500]; label="")

# %%
f(x) = x - x^(3/2)/2
plot(f, 0, 2; label="") |> display

plot([n^2*(f^n)(1.0) for n in 0:25:500]; label="")

# %%


X = vec([+1 -1 -1 -1 -1 +1;
         -1 +1 -1 -1 +1 -1;
         -1 -1 +1 +1 -1 -1;
         -1 -1 +1 +1 -1 -1;
         -1 -1 +1 +1 -1 -1;
         -1 +1 -1 -1 +1 -1;
         +1 -1 -1 -1 -1 +1.])

O = vec([+1 +1 +1 +1 +1 +1;
         +1 -1 -1 -1 -1 +1;
         +1 -1 -1 -1 -1 +1;
         +1 -1 -1 -1 -1 +1;
         +1 -1 -1 -1 -1 +1;
         +1 -1 -1 -1 -1 +1;
         +1 +1 +1 +1 +1 +1.])

F1 = vec([-1 -1 -1 -1 -1 -1;
          -1 +1 +1 +1 +1 -1;
          -1 +1 -1 -1 -1 -1;
          -1 +1 +1 +1 +1 -1;
          -1 +1 -1 -1 -1 -1;
          -1 +1 -1 -1 -1 -1;
          -1 -1 -1 -1 -1 -1.])

F2 = vec([+1 +1 +1 +1 +1 -1;
          +1 -1 -1 -1 -1 -1;
          +1 -1 -1 -1 -1 -1;
          +1 +1 +1 +1 +1 -1;
          +1 -1 -1 -1 -1 -1;
          +1 -1 -1 -1 -1 -1;
          +1 -1 -1 -1 -1 -1.])