using Printf
using Statistics
using LinearAlgebra
using DelimitedFiles
using IterativeSolvers
using SparseArrays
using TickTock

# rearrange the edge sequence in the increasing order of layer id and starting node id
function edge_seq_rearrangement(edge_seq::Array{Int64,2})
    layer_list = setdiff(edge_seq[:,4],[])
    L = maximum(layer_list)
    if length(setdiff([i for i = 1:L],layer_list)) != 0
        println("please label layers in the order 1, 2, ..., L")
    end
    node_list = setdiff(edge_seq[:,1],[])
    N = maximum(edge_seq[:,1])
    if length(setdiff([i for i = 1:N],node_list)) != 0
        println("please label nodes in the order 1, 2, ..., N")
    end

    edge_seq_rearranged = zeros(Int64,size(edge_seq))
    id = 1
    for ell in layer_list
        edge_list_ell = findall(isequal(ell),edge_seq[:,4])
        M = sparse(edge_seq[edge_list_ell,1],edge_seq[edge_list_ell,2],edge_seq[edge_list_ell,3])
        inds = findall(!iszero,M)
        a = getindex.(inds, 1)
        b = getindex.(inds, 2)
        edge_seq_rearranged[id:id+length(a)-1,1] = b
        edge_seq_rearranged[id:id+length(a)-1,2] = a
        edge_seq_rearranged[id:id+length(a)-1,3] = M[inds]
        edge_seq_rearranged[id:id+length(a)-1,4] = ell*ones(Int64,length(a))
        id = id+length(a)
    end

    return edge_seq_rearranged
end

# solve a set of equations to obtain beta_i, beta_ij, and gamma_ij, for any number of layers
# this function is used for rhoA > (rhoA)^o in layer one under a given strategy configuration
# input---
# multilayer network in the form of an edge sequanece: edge_seq [node id, node id, edge weight, layer id], where edge weights here are set to be any positive integer
# initial strategy configuration: xi_seq [node id, strategy, layer id]
# output---
# the first column of "solution" gives beta_ij
# the other columns of "solution" gives gamma_ij, where the ll_th column is the strategy assortment between layer 1 and layter ll
function beta_gamma(edge_seq::Array{Int64,2}, xi_seq::Array{Int64,2})
    N = maximum(edge_seq[:,1])
    L = maximum(edge_seq[:,4])
    M = spzeros(Float64,N*L,N)
    for ell = 1:L
        edge_list_ell = findall(isequal(ell),edge_seq[:,4])
        M_ell = sparse(edge_seq[edge_list_ell,1],edge_seq[edge_list_ell,2],edge_seq[edge_list_ell,3],N,N)
        M[(ell-1)*N+1:ell*N,:] = M_ell
    end
    pi = spzeros(Float64,L*N)
    for ell = 1:L
        pi_2 = sum(M[(ell-1)*N+1:ell*N,:], dims = 2)
        pi_2 = pi_2[:,1]
        pi_transient = spzeros(Float64, N)
        presence = findall(!iszero,pi_2)
        pi_transient[presence] = pi_2[presence]/sum(pi_2)
        pi[(ell-1)*N+1:ell*N] = pi_transient

        M_transient = M[(ell-1)*N+1:ell*N,:]
        M_transient[presence,:] = M_transient[presence,:]./pi_2[presence]
        M[(ell-1)*N+1:ell*N,:] = M_transient
    end

    xi = spzeros(Float64,L*N)
    for ell = 1:L
        xi_transient = spzeros(Int64,N)
        presence = findall(isequal(ell),xi_seq[:,3])
        xi_transient[xi_seq[presence,1]] = xi_seq[presence,2]
        xi[(ell-1)*N+1:ell*N] = xi_transient
    end

    # obtain beta_i
    presence_list = findall(!iszero,pi[1:N])
    N1 = length(presence_list)
    absence_list = findall(iszero,pi[1:N])
    MatA_i = M[1:N,:]-sparse(Matrix(1.0I, N, N))
    pi_list = findall(!isequal(presence_list[1]),[i for i = 1:N])
    MatA_i[:,pi_list] = MatA_i[:,pi_list] - MatA_i[:,presence_list[1]]*transpose(pi[pi_list]/pi[presence_list[1]])
    MatA_i_reduced = MatA_i[pi_list,pi_list]

    xi_hat = sum(pi[1:N].*xi[1:N])
    MatB_i = (xi_hat*ones(Float64,N)-xi[1:N])*N1
    MatB_i_reduced = MatB_i[pi_list]

    beta_i_solution_reduced = idrs(MatA_i_reduced, MatB_i_reduced)
    beta_i_solution = zeros(Float64, N)
    beta_i_solution[pi_list] = beta_i_solution_reduced
    beta_i_solution[presence_list[1]] = -sum(pi[1:N].*beta_i_solution)/pi[presence_list[1]]

    # obtain beta_ij
    Mat1 = M[1:N,:]
    inds = findall(!iszero,Mat1)
    a = getindex.(inds, 1)
    b = getindex.(inds, 2)
    vec1 = [i for i = 1:N]
    vec1 = transpose(vec1)
    X1 = N*(a.-1)*ones(Int64,1,N)+ones(Int64,length(a))*vec1
    Y1 = N*(b.-1)*ones(Int64,1,N)+ones(Int64,length(b))*vec1
    W1 = Mat1[inds]*ones(Int64,1,N)

    vec2 = [(i-1)*N for i = 1:N]
    vec2 = transpose(vec2)
    X2 = a*ones(Int64,1,N)+ones(Int64,length(a))*vec2
    Y2 = b*ones(Int64,1,N)+ones(Int64,length(b))*vec2
    W2 = Mat1[inds]*ones(Int64,1,N)

    X11 = reshape(X1,:,1)
    Y11 = reshape(Y1,:,1)
    W11 = reshape(W1,:,1)
    X22 = reshape(X2,:,1)
    Y22 = reshape(Y2,:,1)
    W22 = reshape(W2,:,1)
    MatA_ij = sparse(X11[:,1],Y11[:,1],W11[:,1],N^2,N^2)/2+sparse(X22[:,1],Y22[:,1],W22[:,1],N^2,N^2)/2
    MatA_ij = MatA_ij - sparse(Matrix(1.0I, N*N, N*N))
    beta_obtained = [(i-1)*N+i for i = 1:N]
    beta_missed = setdiff([i for i = 1:N^2], beta_obtained)
    MatA_ij_reduced = MatA_ij[beta_missed, beta_missed]

    MatB_ij = -sum(MatA_ij[:,beta_obtained].*transpose(beta_i_solution), dims = 2)
    Mat = xi_hat*ones(Float64,N,N)*N1/2-xi[1:N]*transpose(xi[1:N])*N1/2
    Mat[absence_list,:] = spzeros(Float64,length(absence_list),N)
    Mat[:,absence_list] = spzeros(Float64,N,length(absence_list))
    MatB_ij = MatB_ij+reshape(transpose(Mat),:)
    MatB_ij_reduced = MatB_ij[beta_missed]

    beta_ij_solution_reduced = idrs(MatA_ij_reduced, MatB_ij_reduced)
    beta_ij_solution = zeros(Float64, N*N)
    beta_ij_solution[beta_missed] = beta_ij_solution_reduced
    beta_ij_solution[beta_obtained] = beta_i_solution

    solution = spzeros(Float64, N*N, L)
    solution[:,1] = beta_ij_solution
    # obtain gamma_ij
    if L > 1
        for ell = 2:L
            xi_ell = xi[(ell-1)*N+1:ell*N]
            pi_ell = pi[(ell-1)*N+1:ell*N]
            presence_ell = findall(!iszero, pi_ell)
            N_ell = length(presence_ell)
            GammaA_ij = sparse(X11[:,1],Y11[:,1],W11[:,1],N^2,N^2)*(N_ell-1)/(N1+N_ell-1)

            Mat_ell = M[(ell-1)*N+1:ell*N,:]
            inds_ell = findall(!iszero,Mat_ell)
            a_ell = getindex.(inds_ell, 1)
            b_ell = getindex.(inds_ell, 2)
            vec_ell = [(i-1)*N for i = 1:N]
            vec_ell = transpose(vec_ell)
            X_ell = a_ell*ones(Int64,1,N)+ones(Int64,length(a_ell))*vec_ell
            Y_ell = b_ell*ones(Int64,1,N)+ones(Int64,length(b_ell))*vec_ell
            W_ell = Mat_ell[inds_ell]*ones(Int64,1,N)
            X_ellell = reshape(X_ell,:,1)
            Y_ellell = reshape(Y_ell,:,1)
            W_ellell = reshape(W_ell,:,1)
            GammaA_ij = GammaA_ij+sparse(X_ellell[:,1],Y_ellell[:,1],W_ellell[:,1],N^2,N^2)*(N1-1)/(N1+N_ell-1)

            vec_gamma1 = reshape(ones(Int64, N)*vec1,:)
            vec_gamma2 = reshape(transpose(vec1)*ones(Int64,1,N),:)
            X_gamma1 = N*(a.-1)*ones(Int64,1,N*N)+ones(Int64,length(a))*transpose(vec_gamma1)
            Y_gamma1 = N*(b.-1)*ones(Int64,1,N*N)+ones(Int64,length(b))*transpose(vec_gamma2)
            W_gamma1 = Mat1[inds]*ones(Int64,1,N*N)

            vec_gamma3 = [(i-1)*N for i in vec_gamma1]
            vec_gamma4 = [(i-1)*N for i in vec_gamma2]
            X_gamma2 = a_ell*ones(Int64,1,N*N)+ones(Int64,length(a_ell))*transpose(vec_gamma3)
            Y_gamma2 = b_ell*ones(Int64,1,N*N)+ones(Int64,length(b_ell))*transpose(vec_gamma4)
            W_gamma2 = Mat_ell[inds_ell]*ones(Int64,1,N*N)

            X_gamma_11 = reshape(X_gamma1,:,1)
            Y_gamma_11 = reshape(Y_gamma1,:,1)
            W_gamma_11 = reshape(W_gamma1,:,1)
            X_gamma_22 = reshape(X_gamma2,:,1)
            Y_gamma_22 = reshape(Y_gamma2,:,1)
            W_gamma_22 = reshape(W_gamma2,:,1)
            GammaA_ij = GammaA_ij + sparse(X_gamma_11[:,1],Y_gamma_11[:,1],W_gamma_11[:,1],N^2,N^2).*sparse(X_gamma_22[:,1],Y_gamma_22[:,1],W_gamma_22[:,1],N^2,N^2)/(N1+N_ell-1)

            GammaA_ij = GammaA_ij - sparse(Matrix(1.0I, N*N, N*N))
            pi_transient = pi[1:N].*pi_ell
            presence_list = findall(!iszero,pi_transient)
            absence_list = findall(iszero,pi_transient)
            GammaA_delete = (presence_list.-1)*N+presence_list
            GammaA_ij[:,GammaA_delete] = GammaA_ij[:,GammaA_delete] - GammaA_ij[:,GammaA_delete[1]]*transpose(pi[presence_list])/pi[presence_list[1]]
            pi_list = findall(!isequal(GammaA_delete[1]),[i for i = 1:N*N])
            GammaA_ij_reduced = GammaA_ij[pi_list,pi_list]

            xi_ell_hat = sum(pi_ell.*xi_ell)
            GammaB_ij = (xi_hat*xi_ell_hat*ones(Float64,N,N) - xi[1:N]*transpose(xi_ell))*N1*N_ell/(N1+N_ell-1)
            presence1 = findall(iszero, pi[1:N])
            presence1 = presence1[:,1]
            presence_ell = findall(iszero, pi_ell)
            presence_ell = presence_ell[:,1]
            GammaB_ij[presence1,:] = spzeros(Float64, length(presence1), N)
            GammaB_ij[:,presence_ell] = spzeros(Float64, N, length(presence_ell))
            GammaB_ij  = reshape(transpose(GammaB_ij), :)
            GammaB_ij_reduced = GammaB_ij[pi_list]

            gamma_ij_solution_reduced = idrs(GammaA_ij_reduced, GammaB_ij_reduced)
            gamma_ij_solution = zeros(Float64, N*N)
            gamma_ij_solution[pi_list] = gamma_ij_solution_reduced
            gamma_ij_solution[GammaA_delete[1]] = -sum(pi[presence_list].*gamma_ij_solution[GammaA_delete])/pi[presence_list[1]]
            solution[:,ell] = gamma_ij_solution
        end
    end

    return solution
end

# provide the value of coefficient of Eq. (2), namely theta and phi
# input---
# multilayer network in the form of an edge sequanece: edge_seq [node id, node id, edge weight, layer id]
# initial strategy configuration: xi_seq [node id, strategy, layer id]
# for an edge between node i and j in layer ll with weight wij_ll, two terms should be added into the edge sequence, i.e. [i j wij_ll ll] and [j i wij_ll ll]
# the code here requires wij_ll to be a positive integer
# output---
# strategy assortment:  theta_phi[:,1] = [theta0, theta1, theta2, theta3], strategy correlation within layer 1
#                       theta_phi[:,ll] = [phi00, phi01, phi20, phi21], strategy correlation between layer 1 and layer ll
# if there are only two layers, in Eq. (2) in the main text, [theta0, theta1, theta2, theta3] = theta_phi[:,1] and [phi00, phi01, phi20, phi21] = theta_phi[:,2]
function bc_multilayer_DB(edge_seq::Array{Int64,2}, xi_seq::Array{Int64,2})
    edge_seq = edge_seq_rearrangement(edge_seq)
    solution = beta_gamma(edge_seq, xi_seq)
    N = maximum(edge_seq[:,1])
    L = maximum(edge_seq[:,4])
    M = spzeros(Float64,N*L,N)
    for ell = 1:L
        edge_list_ell = findall(isequal(ell),edge_seq[:,4])
        M_ell = sparse(edge_seq[edge_list_ell,1],edge_seq[edge_list_ell,2],edge_seq[edge_list_ell,3],N,N)
        M[(ell-1)*N+1:ell*N,:] = M_ell
    end
    pi = spzeros(Float64,L*N)
    for ell = 1:L
        pi_2 = sum(M[(ell-1)*N+1:ell*N,:], dims = 2)
        pi_2 = pi_2[:,1]
        pi_transient = spzeros(Float64, N)
        presence = findall(!iszero,pi_2)
        pi_transient[presence] = pi_2[presence]/sum(pi_2)
        pi[(ell-1)*N+1:ell*N] = pi_transient

        M_transient = M[(ell-1)*N+1:ell*N,:]
        M_transient[presence,:] = M_transient[presence,:]./pi_2[presence]
        M[(ell-1)*N+1:ell*N,:] = M_transient
    end

    M1 = M[1:N,:]
    beta = transpose(reshape(solution[:,1], N, :))
    theta1 = sum(pi[1:N].*M1.*beta)
    M2 = M1*M1
    theta2 = sum(pi[1:N].*M2.*beta)
    M3 = M2*M1
    theta3 = sum(pi[1:N].*M3.*beta)

    theta_phi = zeros(Float64, 4, L)
    theta_phi[:,1] = [0.0,theta1,theta2,theta3]

    if L > 1
        for ell = 2:L
            gamma = transpose(reshape(solution[:,ell], N, :))
            M_ell = M[(ell-1)*N+1:ell*N,:]
            phi01 = sum(pi[1:N].*M_ell.*gamma)
            phi20 = sum(pi[1:N].*M2.*gamma)
            M_ell21 = M2*M_ell
            phi21 = sum(pi[1:N].*M_ell21.*gamma)
            theta_phi[:,ell] = [0,phi01,phi20,phi21]
        end
    end

    return theta_phi
end

# an example about the using of above functions
# see Fig. 5A in the main text, where nodes are labelled and initial strategy configuration is presented in SFig. 2
# the edge sequence is given in the form [node id, node id, edge weight, layer id]
edge_seq = [1 6 1 1;
            2 6 1 1;
            3 6 1 1;
            4 6 1 1;
            5 6 1 1;
            6 1 1 1;
            6 2 1 1;
            6 3 1 1;
            6 4 1 1;
            6 5 1 1;
            1 2 1 2;
            1 3 1 2;
            1 4 1 2;
            1 5 1 2;
            1 6 1 2;
            2 1 1 2;
            3 1 1 2;
            4 1 1 2;
            5 1 1 2;
            6 1 1 2]
# the strategy configuration is given in the form [node id, strategy, layer id]
xi_seq = [1 1 1;
          2 0 1;
          3 0 1;
          4 0 1;
          5 0 1;
          6 0 1;
          1 0 2;
          2 0 2;
          3 0 2;
          4 0 2;
          5 0 2;
          6 1 2]
# see Eq.(2) in the main text
# theta_phi[:,1] = [theta0, theta1, theta2, theta3], strategy correlation within layer 1
# theta_phi[:,2] = [phi00, phi01, phi20, phi21], strategy correlation between layer 1 and layer 2
theta_phi = bc_multilayer_DB(edge_seq, xi_seq)


#=

   Sequence waiting times.

   From
   Gunnar Blom, Lars Holst, Dennis Sandell:
   "Problems and Snapshots from the World of Probability"
   Page 4ff, Problem 1.4 Patterns I, (Part Problem 1)

   Problem: Waiting time for a pattern (coin flipping)

   What is the waiting time for the patterns[0,1] in a coin
   flipping sequence? For the pattern [1,1]?
   Answer:
     * waiting time for [0,1] = 4
     * waiting time for [1,1] = 6

   This model verifies this result:
   - [0,1]: 4.0031
   - [1,1]: 6.0300

   Mean length of some other sequences:
   - [0,1,0]: 9.9668
   - [1,1,1]: 14.1199

   Cf ~/webppl/sequence_waiting_times_1.wppl

=#
using Turing, StatsPlots
include("jl_utils.jl")

@model function sequence_waiting_times_1(pattern)
    patternLen = length(pattern)

    function flipSequence(a)
        len = length(a)
        if len >= patternLen && a[len-patternLen+1:end] == pattern
            return a
        else
            return flipSequence(push!(a,rand(Bernoulli(0.5))*1))
        end
    end

    a = flipSequence([])
    len ~ Dirac(length(a))

end

function run_seqence_waiting_times_1(pattern)
    println("\npattern: $pattern")
    model = sequence_waiting_times_1(pattern)

    # chns = sample(model, Prior(), 10_000)
    # chns = sample(model, MH(), 10_000)
    chns = sample(model, PG(5), 10_000)
    # chns = sample(model, SMC(), 10_000)
    # chns = sample(model, IS(), 10_000)

    display(chns)
    # display(plot(chns))

    # show_var_dist_pct(chns,:len)
end

run_seqence_waiting_times_1([0,1])
run_seqence_waiting_times_1([1,1])
run_seqence_waiting_times_1([0,1,0])
run_seqence_waiting_times_1([1,1,1])


"""
    struct AbstractGegenbauerBasis{P} <: AbstractMultipleOrthogonalBasis{P}
        polynomials::Vector{P}
    end

Orthogonal polynomial with respect to the univariate weight function ``w(x) = (1 - x^2)^{\\alpha - 1/2}`` over the interval ``[-1, 1]``.
"""
abstract type AbstractGegenbauerBasis{P} <: AbstractMultipleOrthogonalBasis{P} end

even_odd_separated(::Type{<:AbstractGegenbauerBasis}) = true
reccurence_second_coef(::Type{<:AbstractGegenbauerBasis}, degree) = 0

"""
    struct LegendreBasis{P} <: AbstractGegenbauerBasis{P}
        polynomials::Vector{P}
    end

Orthogonal polynomial with respect to the univariate weight function ``w(x) = 1`` over the interval ``[-1, 1]``.
"""
struct LegendreBasis{P} <: AbstractGegenbauerBasis{P}
    polynomials::Vector{P}
end

polynomial_type(::Type{<:LegendreBasis}, V::Type) = MP.polynomialtype(V, Float64)

reccurence_first_coef(::Type{<:LegendreBasis}, degree) = (2degree - 1)
reccurence_third_coef(::Type{<:LegendreBasis}, degree) = -(degree - 1)
reccurence_deno_coef(::Type{<:LegendreBasis}, degree) = degree

degree_one_univariate_polynomial(::Type{<:LegendreBasis}, variable::MP.AbstractVariable) = MA.@rewrite(variable + 0)


immutable RandIntSampler  # for generating Int samples in [0, K-1]
    a::Int
    Ku::Uint
    U::Uint

    RandIntSampler(K::Int) = (Ku = uint(K); new(1, Ku, div(typemax(Uint), Ku) * Ku))
    RandIntSampler(a::Int, b::Int) = (Ku = uint(b-a+1); new(a, Ku, div(typemax(Uint), Ku) * Ku))
end

function rand(s::RandIntSampler)
    x = rand(Uint)
    while x >= s.U
        x = rand(Uint)
    end
    s.a + int(rem(x, s.Ku))
end

randi(K::Int) = rand(RandIntSampler(K))
randi(a::Int, b::Int) = rand(RandIntSampler(a, b))


using JLD2, Plots, Printf, LinearAlgebra
include(pwd() * "/analysis_scripts/" * "post_analysis.jl")
searchdir(path, key) = filter(x -> occursin(key, x), readdir(path))
mesoscale_dir = pwd()
checkpoints = searchdir(mesoscale_dir, "iteration")
file = jldopen( mesoscale_dir * "/" * checkpoints[1])
b = file["tracers"]["b"]["data"][2:end-1, 2:end-1, 2:end-1]
u = file["velocities"]["u"]["data"][2:end-1, 2:end-1, 2:end-1]
v = file["velocities"]["v"]["data"][2:end-1, 2:end-1, 2:end-1]
w = file["velocities"]["w"]["data"][2:end-1, 2:end-1, 2:end-1]
v = (v[:, 1:end-1, :] + v[:, 2:end, :]) .* 0.5 # put everything on cell centers
w = (w[:, :, 1:end-1] + w[:,:, 2:end]) .* 0.5  # put everything on cell centers
close(file)
n = length(b)
avg_b = sum(b) / n
tke = @. u^2 + v^2 + w^2
avg_tke = sum(tke) / n
println("the average buoyancy is " * string(avg_b))
println("the average tke is " * string(avg_tke))
##
const Nx = size(b)[1]
const Ny = size(b)[2]
const Nz = size(b)[3]
const Lz = 3000.0
const Lx = Ly = 10^6 * 1.0
##
@inline function zC(k)
    return - Lz + (k-0.5) / Nz * Lz
end
@inline function yC(j)
    return  (j-0.5) / Ny * Ly
end
@inline function xC(j)
    return  (j-0.5) / Nx * Lx
end

x = xC.( collect(1:Nx) )
y = yC.( collect(1:Ny) )
z = zC.( collect(1:Nz) )

x2 = (x[1:end-1] + x[2:end]) / 2
y2 = (y[1:end-1] + y[2:end]) / 2
z2 = (z[1:end-1] + z[2:end]) / 2
##
∂x = PartialDerivative((x, 1))
∂y = PartialDerivative((y, 2))
∂z = PartialDerivative((z, 3))
∇ = [∂x, ∂y, ∂z]
##
field_label = [(u, "u"), (v, "v"), (w, "w"), (b, "b"), ( w .* w, "ww")]
selection = 4
field = field_label[selection][1]
label = field_label[selection][2]
cmax = maximum(field)
cmin = minimum(field)
clims = (cmin, cmax)
p1 = contourf(y, z, field[ 1, :, :]',
    color = :thermometer, title = "Zonal Slice " * label,
    xlabel = "Meridional [m]", ylabel = "Depth [m]"
    , clims = clims, linewidth = 0)
##
field_label = [(u, "u"), (v, "v"), (w, "w"), (b, "b"), ( w .* w, "ww")]
selection = 4
field = field_label[selection][1]
label = field_label[selection][2]
cmax = maximum(field)
cmin = minimum(field)
clims = (cmin, cmax)
ϕ = sum(field, dims = 1)[1,:,:] ./ Nx
p1 = contourf(y, z, ϕ',
    color = :thermometer, title = "Zonal Average " * label,
    xlabel = "Meridional [m]", ylabel = "Depth [m]"
    , clims = clims, linewidth = 0)

##
field_label = [(u, "u"), (v, "v"), (w, "w"), (b, "b"), (w .* w, "ww"), (u .* u, "uu"), (v .* v, "vv"), (∂z( w .* w), "d(w .* w) / dz")]
selection = 1
field = sum(field_label[selection][1], dims = (1,2)) ./ (Nx * Ny)
label = field_label[selection][2]
cmax = maximum(field)
cmin = minimum(field)
clims = (cmin, cmax)
p1 = scatter(field[1,1,:],  z,
    color = :blue, xlabel = "Field", label = "Horizontally Averaged " * label,
    ylabel = "Depth [m]", xlims = (cmin, cmax)
    , legend = :bottomright)

##
# day_label = @sprintf("%.2f ", sim_day[i])
# surface values
field_label = [(u, "u"), (v, "v"), (w, "w"), (b, "b"), (u .* b, "ub")]
selection = 4
field = field_label[selection][1]
label = field_label[selection][2]
cmax = maximum(field)
cmin = minimum(field)
clims = (cmin, cmax)
p1 = contourf(x, y, field[ :, :, end]',
    color = :thermometer, title = "Surface " * label,
    xlabel = "Zonal [m]", ylabel = "Meridional [m]"
    , clims = clims, linewidth = 0)

## Ertel potential vorticity
U  = [u, v, w]
ω  = ∇ × U
∇b = [∂x(b), ∂y(b), ∂z(b)]
f  = -1e-4
β  = 1e-11
zero_ϕ = ω[1] .* 0
Ω = [zero_ϕ, zero_ϕ, zero_ϕ .+ -1e-4 .+ β .* reshape(y2, (1,191,1))]
@. ω +=  Ω
vec_pv = [ω[1] .* ∇b[1] ,  ω[2] .* ∇b[2] , ω[3] .* ∇b[3]]
pv = vec_pv[1] + vec_pv[2] + vec_pv[3]
magnitude = zeros(size(ω[1]))
@inbounds @simd for i in 1:length(ω[1])
    ω1  = sqrt(ω[1][i]^2  + ω[2][i]^2  + ω[3][i]^2)
    ∇b1 = sqrt(∇b[1][i]^2 + ∇b[2][i]^2 + ∇b[3][i]^2)
    magnitude[i] = ω1 * ∇b1
end
alignment = pv ./ magnitude

##
field = pv
label = "Alignment"
ϕ = field[ :, :, end]
cmax = maximum(ϕ)
cmin = minimum(ϕ)
clims = (cmin, cmax)
p1 = contourf(x2, y2, ϕ',
    color = :thermometer, title = "Potential Vorticity ",
    xlabel = "Zonal [m]", ylabel = "Meridional [m]"
    , clims = clims, linewidth = 0)

##
sf_escale = 32
sf = @. exp( -(y - Ly/2)^2 / (Ly^2 / sf_escale) ) - exp( -(Ly - Ly/2)^2 / (Ly^2 / sf_escale) )
plot(y, sf, ylims = (0,1))

nw_escale = 2*10^4
nw = @. exp( (x - Ly) / nw_escale )
plot!(y, nw, ylims = (0,1))
##

bd_escale = 1 / 200
bd = @. exp(- bd_escale * ( z + Lz))
plot(z, bd, ylims = (0,1))

##
# Checking relaxation profile
Δb = 10 * 2e-3
h = 1000
Nx = 192
Nz = 32
@inline function zC(k)
    return - Lz + (k-0.5) / Nz * Lz
end

@inline function yC(j)
    return  (j-0.5) / Nx * Ly
end

relaxation_profile(j) = Δb * (yC(j)/ Ly)
relaxation_profile_north_2(k) = Δb * ( exp(zC(k)/h) - exp(-Lz/h) ) / (1 - exp(-Lz/h))

tmp1 = relaxation_profile.( collect(1:192))
tmp2 = relaxation_profile_north_2.( collect(1:32))

zonal_average_b = sum(b, dims = (1, )) ./ Nx
tmp1 = reshape(tmp1, (1, 192, 1))
plot( b[1, : , end])
plot!(tmp1)
new_field = zonal_average_b .- tmp1
plot(y, new_field[1, :, end])
##
plot( b[2, end , :], legend = false)
for i in 1:10:192
    plot!(b[i, end, :])
end

plot!(tmp2)

##
# Check Geostrophy
∇pʰ = hydrostatic_pressure(b,x,y,z)
zero_ϕ = zeros(size(∇pʰ[1]))
f  = -1e-4
β  = 1e-11
Ω = [zero_ϕ, zero_ϕ, zero_ϕ .+ f .+ β .* reshape(y2, (1,length(y2),1))]
u_avg = (avg_other(u,1)[1:end-1,:,:] + avg_other(u,1)[2:end,:,:]) ./ 2
v_avg = (avg_other(v,1)[1:end-1,:,:] + avg_other(v,1)[2:end,:,:]) ./ 2
w_avg = (avg_other(w,1)[1:end-1,:,:] + avg_other(w,1)[2:end,:,:]) ./ 2
U = [u_avg, v_avg, w_avg]
coriolis_force = [-Ω[3] .* v_avg,  Ω[3] .* u_avg, zero_ϕ]

norm(coriolis_force[1] .+ ∇pʰ[1]) / norm(coriolis_force[1])
norm(coriolis_force[2] .+ ∇pʰ[2]) / norm(coriolis_force[2])

##
field_label = [(coriolis_force[1], "-Ω v"), (-∇pʰ[1], "-∂x(pʰ)"), (coriolis_force[2], "Ω u"), (-∇pʰ[2], "-∂y(pʰ)")]
selection = 2+2
field = field_label[selection][1]
label = field_label[selection][2]
cmax = maximum(field)
cmin = minimum(field)
clims = (cmin, cmax)
p1 = contourf(x2, y2, field[ :, :, end]',
    color = :thermometer, title = "100 [m] depth " * label,
    xlabel = "Zonal [m]", ylabel = "Meridional [m]"
    , clims = clims, linewidth = 0)

selection = 1+2
field = field_label[selection][1]
label = field_label[selection][2]
p2 = contourf(x2, y2, field[ :, :, end]',
    color = :thermometer, title = "100 [m] depth " * label,
    xlabel = "Zonal [m]", yaxis = false
    , clims = clims, linewidth = 0)
plot(p1,p2)
##
gr(size=(700,400))
p3 = plot(p1,p2)
savefig(p3, pwd() * "/geostrophic_1.pdf")
##
p = ∫dz(b, z)
h_p = sum(p, dims = (1,2)) ./ ( (Nx -1) * (Ny-1))
plot(h_p[:], z2)


"""
DNRM2 returns the euclidean norm of a vector via the function
name, so that

   DNRM2 := sqrt( x'*x )

-- This version written on 25-October-1982.
   Modified on 14-October-1993 to inline the call to DLASSQ.
   Sven Hammarling, Nag Ltd.
"""
function tron_dnrm2(n,x,incx)
    one = 1.0
    zero = 0.0

    if n < 1 || incx < 1
        xnorm = zero
    elseif n == 1
        xnorm = abs(x[1])
    else
        scale = zero
        ssq = one

        # The following loop is equivalent to this call to the LAPACK
        # auxiliary routine:
        # CALL DLASSQ( N, X, INCX, SCALE, SSQ )

        for ix=1:incx:1+(n-1)*incx
            if x[ix] != zero
                absxi = abs(x[ix])
                if scale < absxi
                    ssq = one + ssq*(scale/absxi)^2
                    scale = absxi
                else
                    ssq = ssq + (absxi/scale)^2
                end
            end
        end
        xnorm = scale*sqrt(ssq)
    end

    return xnorm
end

if isequal(BLAS_LIBRARY, :OpenBlas)
    dnrm2(n,x,incx) = BLAS.nrm2(n, x, incx)
    dnrm2(n, x::CuArray, incx) = CUBLAS.nrm2(n, x)
else
    dnrm2(n,x,incx) = tron_dnrm2(n, x, incx)
end

@inline function dnrm2(n::Int,x::CuDeviceArray{Float64,1},incx::Int)
    tx = threadIdx().x

    v = 0.0
    if tx <= n  # No check on ty so that each warp has v.
        @inbounds v = x[tx]*x[tx]
    end

    # shfl_down_sync() will automatically sync threads in a warp.

    offset = 16
    while offset > 0
        v += CUDA.shfl_down_sync(0xffffffff, v, offset)
        offset >>= 1
    end
    v = sqrt(v)
    v = CUDA.shfl_sync(0xffffffff, v, 1)

    return v
end


using Plots, ProgressMeter, LinearAlgebra
gr()

m, d  = 1000, 2
coef3 = 2 * rand(d,d,d) .- 1
coef2 = 2 * rand(d,d) .- 1
coef1 = 2 * rand(d,1)
coef0 = rand(1)

struct Stump

    j :: Int64
    theta :: Float64
    b :: Float64

    function Stump(j, theta, b)
        @assert j > 0
        new(j, theta, b)
    end

end

function (self :: Stump)( x )
    self.b * sign.(self.theta .- x[:,self.j])
end

function erreur(self :: Stump, x, y, D)
    sum( D .* (y .!= self(x)))
end

mutable struct Hypotheses

    weights :: Vector{Float64}
    stumps  :: Array{Stump,1}
    D       :: Vector{Vector{Float64}}

    function Hypotheses( )
        weights = Float64[]
        stumps  = Stump[]
        D       = Vector{Float64}[]
        new( weights, stumps, D)
    end

end

function (self::Hypotheses)(x)
    T = length(self.weights)
    m = size(x)[1]
    s = zeros(Float64, m )
    for t in 1:T
        s .+= self.weights[t] .* self.stumps[t](x)
    end
    sign.(s)
end

function erreur(self,x,y)
    m = length(y)
    sum(self(x) .!= y)/m
end

function plot_results( x, y, ytest)

    p    = plot(layout=(1,2); )
    blue = y .< 0
    red  = y .> 0
    scatter!(p[1,1],x[red,1],x[red,2],label=:Xb, title=:test)
    scatter!(p[1,1],x[blue,1],x[blue,2],label=:or)
    blue = ytest .< 0
    red  = ytest .> 0
    scatter!(p[1,2],x[red,1],x[red,2],label=:Xb, title=:test)
    scatter!(p[1,2],x[blue,1],x[blue,2],label=:or)
    p

end

function ERM_Stump( x :: Array{Float64,2},
                    y :: Vector{Float64},
                    D :: Vector{Float64})
    m, d       = size(x)
    Fstar      = 1e15
    F          = Fstar
    bstar      = 1
    thetastar  = 0.
    jstar      = 1
    indice_tri = hcat([sortperm(x[:,j]) for j in 1:d]...) #np.argsort(x,axis=0)
    xsort      = zeros(Float64,m+1)
    for j in 1:d
        xsort[1:m] = x[indice_tri,j][:,j]
        xsort[end] = xsort[m]+1
        ysort = y[indice_tri[:,j]]
        Dsort = D[indice_tri[:,j]]
        F = sum(Dsort .* (ysort .+ 1)/2)
        if F < Fstar
            bstar = 1
            Fstar = F
            thetastar = xsort[1] - 1
            jstar = j
        end
        for i in 1:m
            F = F - ysort[i] * Dsort[i]
            if F < Fstar && xsort[i] < xsort[i+1]
                bstar=1
                Fstar=F
                thetastar=(xsort[i]+xsort[i+1])/2
                jstar=j
            end
        end
        F = sum(Dsort .* (1 .- ysort)/2)
        if F < Fstar
            bstar = -1
            Fstar = F
            thetastar = xsort[1]-1
            jstar = j
        end
        for i in 1:m
            F = F + ysort[i] * Dsort[i]
            if F < Fstar && xsort[i]<xsort[i+1]
                bstar = -1
                Fstar = F
                thetastar = (xsort[i]+xsort[i+1])/2
                jstar = j
            end
        end
    end
    Stump(jstar, thetastar, bstar)
end

function generate_data(m, d)
    x = 2 * rand(m,d) .- 1
    y = zeros(Float64,m)
    for i in 1:d, j in 1:d, k in 1:d
        y .+= coef3[i,j,k] .* x[:,i] .* x[:,j] .* x[:,k]
    end
    y = sign.(y .+ diag((x * coef2) * transpose(x)) .+ x * coef1 .+ coef0);
    x, vec(y)
end

x, y = generate_data(m, d)

#d = 2
#m = 1000
#f = open("test1.txt")
#data = [ transpose(map(x->parse(Float64,x),split(line," "))) for line in eachline(f)]
#data = vcat(data...)
#
#x = data[:,1:2]
#y = data[:,3]

D = 1 / m .* ones(Float64, m)
h = ERM_Stump(x,y,D)
plot_results( x, y, h(x))

function adaboost( x, y , T)

    m = size(x)[1]
    D = 1 / m * ones(Float64,m)
    H = Hypotheses()
    epsilon = Float64[]

    @showprogress 1 for t in 1:T
        ht = ERM_Stump(x, y, D)
        hx = ht(x)
        e  = erreur(ht, x, y, D)
        w0 = 0.5 * log(1 / e-1)
        push!(epsilon, e)
        push!(H.weights, w0)
        push!(H.stumps, ht)
        #push!(H.D, D)
        D = D .* exp.(-w0 * y .* hx)
        D = D / sum(D)
    end
    H, epsilon
end

H, epsilon = adaboost( x, y, 100)
yresult = H(x)
plot_results( x, y, yresult)

a=100*(1-erreur(H,x,y))
println("Performance sur les données de test :",a[1],"%")

xtest, ytest = generate_data(m, d)
plot_results( xtest, ytest, H(xtest))

a=100*(1-erreur(H,xtest,ytest))
println("Performance sur les données de test :",a[1],"%")

err       = Float64[]
err0      = zeros(Float64,m)
err_test  = Float64[]
err_test0 = zeros(Float64,m)
for (weight, stump)  in zip(H.weights, H.stumps)
    err0 .+= weight .* stump(x)
    push!(err, sum( sign.(err0) .!= y) / m)
    err_test0 .+= weight .* stump(xtest)
    push!(err_test, sum(sign.(err_test0) .!= ytest) / m)
end

plot(err; marker=:o)
plot!(err_test; marker=:s)
xlabel!("Itérations")
ylabel!("Erreur en %")

scatter(epsilon)
xlabel!("steps")
ylabel!("Epsilon")
ylims!(minimum(epsilon), 0.5)

f = open("data_banknote_authentication.txt")
databrute = [ transpose(map(x->parse(Float64,x),split(line,","))) for line in eachline(f)]
data = vcat(databrute...);

@show n = size(data)[1]
print("Pourcentage de vrais billets :",(n-sum(data[:,5]))/n*100,"%")

using Random
taille       = n
taille_train = trunc(Int64,taille*0.8)
taille_test  = taille-taille_train
indice       = randperm(taille)
train        = data[indice[1:taille_train],:]
test         = data[indice[taille_train+1:end],:]
xtrain       = train[:,1:4]
ytrain       = 2 * train[:,5] .- 1
xtest        = test[:,1:4]
ytest        = 2 * test[:,5] .- 1
d            = 4
m            = taille_train

H, epsilon = adaboost(xtrain, ytrain, 50)

yresult     =  H(xtrain)
yresulttest =  H(xtest);

a=100*(1 .- erreur(H,xtrain,ytrain))
print("Performance sur les données d'entraînement :",a[1],"%")

a=100*(1-erreur(H,xtest,ytest))
print("Performance sur les données de test :",a[1],"%")

nt        = length(H.stumps)
err       = Float64[]
err0      = zeros(Float64,length(ytrain))
err_test  = Float64[]
err_test0 = zeros(Float64,length(ytest))

for i in 1:nt
    err0 .+= H.weights[i] .* H.stumps[i](xtrain)
    push!(err, sum(sign.(err0) .!= ytrain) / length(ytrain))
    err_test0 .+= H.weights[i] .* H.stumps[i](xtest)
    push!(err_test, sum(sign.(err_test0) .!= ytest)/length(ytest))
end

plot(err)
plot!(err_test)
xlabel!("Itérations ")
ylabel!("Erreur en %")

plot(epsilon)
xlabel!("Itérations")
ylabel!("Epsilon")
ylims!(minimum(epsilon), 0.5)

print("borne d'erreur théorique ",exp(-2*sum((1/2 .- epsilon).^2))*100,"%")

plot_results(xtrain, ytrain, yresult)

plot_results(xtest, ytest, yresulttest)

f = open("EEG_Eye_State.txt")

databrute = [ transpose(map(x->parse(Float64,x),split(line,","))) for line in eachline(f)]
data = vcat(databrute...)

n = size(databrute)[1]
println("Nombre de données : ", n)

count = sum( data[:,15] .== 0)
println("Pourcentage de classe 0 :",count/n*100,"%")

n_train  = trunc(Int64,n/3)
n_valid  = trunc(Int64,n/3)
n_test   = n-n_train-n_valid
indice   = randperm(n)
train    = data[indice[1:n_train],:]
valid    = data[indice[n_train+1:n_train+n_valid],:]
test     = data[indice[n_train+n_valid+1:end],:]
x        = train[:,1:14]
y        = 2 * train[:,15] .- 1
xvalid   = valid[:,1:14]
yvalid   = 2 * valid[:,15] .- 1
xtest    = test[:,1:14]
ytest    = 2 * test[:,15] .- 1
d        = 14
m        = n_train

(H,epsilon)=adaboost(x,y,2000)

yresult     = H(x)
yresulttest = H(xtest)

a = sum(yresult .== y) / m * 100
println("Performance sur les données d'entraînement :",a[1],"%")

a = sum(yresulttest .== ytest) / n_test*100
println("Performance sur les données de test :",a[1],"%")

err        = Float64[]
err0       = zeros(Float64,length(y))
err_valid  = Float64[]
err_valid0 = zeros(Float64,length(yvalid))
err_test   = []
err_test0  = zeros(Float64,length(ytest))
nt = length(H.stumps)
for i in 1:nt
    err0 .+= H.weights[i] .* H.stumps[i](x)
    push!(err, sum(sign.(err0) .!= y) / length(y))
    err_test0 .+= H.weights[i] .* H.stumps[i](xtest)
    push!(err_test, sum(sign.(err_test0) .!= ytest) / length(ytest))
    err_valid0 .+= H.weights[i]*H.stumps[i](xvalid)
    push!(err_valid,sum(sign.(err_valid0) .!= yvalid) / length(yvalid))
end

plot(err)
plot!(err_test)
plot!(err_valid)
xlabel!("Itérations")
ylabel!("Erreur en %")

plot(epsilon)
xlabel!("Itérations")
ylabel!("Epsilon")
ylims!(minimum(epsilon), 0.5)

println("borne d'erreur théorique : ",exp(-2*sum((1/2 .- epsilon).^2))*100,"%")

i = argmin(err_valid)
println(i)
println(err_test[i])
println(minimum(err_valid))


using Random, LinearAlgebra

"Returns a copy of `mat` with all coeficients on the diagonal equals to zero."
function nulldiag(mat::AbstractMatrix)
    @assert size(mat, 1) == size(mat, 2)
    mat .* .!I(size(mat, 1))
end

"Returns a copy of an array `mat` (containing objects on wich a norm is
defined) where all objects have a norm of 1."
function normalize(mat::Array)
    mat ./ norm.(mat)
end

"Returns a copy of `mat` where columns order is randomly shuffled."
function shufflecols(mat::AbstractMatrix)
    mat[:, randperm(size(mat, 2))]
end

"""
Returns an copy of `mat` where columns are offset to the left.

# Exemple
```julia-repl
julia> a = [1 2 3;
            4 5 6]
julia> offsetcols(a)
2 3 1
5 6 4
```
"""
function offsetcols(mat::AbstractMatrix)
    if size(mat, 2) == 1
        return mat
    end
    perms = collect(2:size(mat, 2))
        push!(perms, 1)
    mat[:, perms]
end

"""
Returns a square matrix of all differences of the elements provided in the `mat`
column matrix or vector.

# Exemple
```julia-repl
julia> diffs([a; b; c])
a-a b-a c-a
a-b b-b c-b
a-c b-c c-c
```
"""
function diffs(mat::AbstractMatrix)
    dupl = repeat(mat, 1, size(mat, 1))
    permutedims(dupl) - dupl
end


"""
    generate random particle-hole-symmetric positive-definite hamiltonian.
"""
function test_random_particle_hole_positive_definite_hamiltonian(N::Int)
    @assert mod(N,2)==0
    S = rand(N).+0.5
    q1,r1 = qr(rand(N,N))
    q2,r2 = qr(rand(N,N))
    S = diagm(vcat(S,S))
    q = [q1 q2;q2 q1]
    return transpose(q)*S*q
end

"""
    test if the transform V is canonical
"""
function test_canonicality(V::Matrix,ϵ::Float64)
    @assert size(V)[1]==size(V)[2]
    N = size(V)[1]÷2
    η = diagm(vcat([1.0 for i=1:N],[-1.0 for i=1:N]))
    @test norm(transpose(V)*η*V-η)<ϵ
end

"""
    test if the transform V is canonical with zeromodes
"""
function test_canonicality_zeromode(V::Matrix,ϵ::Float64)
    @assert size(V)[1]==size(V)[2]
    N = size(V)[1]÷2
    η = diagm(vcat([1.0 for i=1:N],[-1.0 for i=1:N]))
    test_η = transpose(V)*η*V
    for i = 1:2N
        if abs(test_η[i,i])<1E-10
            test_η[i,i] = i>N ? -1.0 : 1.0
        end
    end
    @test norm(test_η-η)<ϵ
end

"""
    test if the matrix is diagonalized
"""
function test_diag(s::Matrix,ϵ::Float64=1E-10)
    @test norm(s-diagm(diag(s)))<ϵ
end

"""
    test if S = transpose(V)*A*V
    this will hold if use normal and abnormal deal zeromodes
"""
function test_right(s::Matrix,V::Matrix,A::Matrix,ϵ::Float64)
    N = size(s)[1]÷2
    @test norm(transpose(V)*A*V-s)<ϵ
end

"""
    test if my_inverse(V^T)*S*my_inverse(V) = A
    this will not hold if use normal and abnormal deal zeromodes
"""
function test_full_right(s::Matrix,V::Matrix,A::Matrix,ϵ::Float64)
    N = size(s)[1]÷2
    @test norm(transpose(V)*A*V-s)<ϵ
    Vs = V[1:N,1:N]
    Us = V[1:N,1+N:2N]
    iV = [transpose(Vs) -transpose(Us);-transpose(Us) transpose(Vs)]
    @test norm(A-transpose(iV)*s*iV)<ϵ
end

"""
    OBC Schwinger Boson hamiltonian
"""

function schwinger_boson_hamiltonian(N::Int,bc::String)
    A = 1
    λ = 1
    begin
        Q = zeros(N,N)
        for id in CartesianIndices((1:N,1:N))
            if id[1]-id[2]==1
                Q[id] = 1.0*A
            elseif id[2] - id[1] ==1
                Q[id] = -1.0*A
            end
        end

        if bc == "PBC"
            Q[N,1] = -1.0*A
            Q[1,N] =  1.0*A
        end
        Q = Q/2.0
        Q
    end
    ϵ = 5*1E-10
    Lambda = diagm([λ for i = 1:N])
    M = [Lambda Q;transpose(Q) Lambda]
    return M
end


function face_sampling(n::Int,dim::Int, tol::T) where T

	X1 = [collect(range(-tol,stop=tol, length=n)) for i in 1:dim]
	X = Iterators.product(X1 ...)
	sp = vec(collect.(X))
	sp = [sp[i] for i in 1:Int(length(sp)) if any(sp[i] .== tol) || any(sp[i] .== -tol) || sp[i] == zeros(dim)]
	return unique(sp)

end

function uniform_sampling(n::Int,dim::Int)
	X1 = [collect(range(-1,stop=1, length=n)) for i in 1:dim]
	X = Iterators.product(X1 ...)
	sp = vec(collect.(X))::Vector{Vector{Float64}}
	return unique(sp)
end

# Samples only the center point
function center_sampling(n::Int, dim::Int)
	return [zeros(dim)]
end

function sampleX(sampling_type::Symbol, ns::Int, dim::Int, T)

	if sampling_type == :uniform
		return uniform_sampling(ns, dim)
	end

	if sampling_type == :face
		return face_sampling(ns, dim, convert(T,0.98))
	end

	if sampling_type == :random
		return LHC_sampling(ns, dim)
	end

	if sampling_type == :center
		return center_sampling(ns, dim)
	end

	error("Unknown sampling method type. Please leave default for uniform sampling or select either uniform or face or random.")
end

# Auxiliary function for sampling inputs and disturbances
function sampleUW(n::Int, dim::Int, ub::Vector{T}, lb::Vector{T}) where T
	X1 = [collect(range(lb[i],stop=ub[i], length=n)) for i in 1:dim]
	X = Iterators.product(X1 ...)
	sp = vec(collect.(X))::Vector{Vector{T}}
	return unique(sp)
end

function sample_points!(y::Vector{Vector{T}}, bx::AbstractHyperrectangle, Xsp::Vector{Vector{T}}) where T
	@simd for i in 1:length(Xsp)
		@inbounds y[i] .= Xsp[i] .* radius_hyperrectangle(bx) .+ LS.center(bx) # [sp[i][1]*box.radius[1]+box.center[1], sp[i][2]*box.radius[2]+box.center[2]]
	end
end


# bootstrapped class of curve methods

# Forward curves

"""
    ForwardStarting(curve,forwardstart)

Rebase a `curve` so that `discount`/`accumulation`/etc. are re-based so that time zero from the new curves perspective is the given `forwardstart` time.

# Examples

```julia-repl
julia> zero = [5.0, 5.8, 6.4, 6.8] ./ 100
julia> maturity = [0.5, 1.0, 1.5, 2.0]
julia> curve = Yields.Zero(zero, maturity)
julia> fwd = Yields.ForwardStarting(curve, 1.0)

julia> discount(curve,1,2)
0.9275624570410582

julia> discount(fwd,1) # `curve` has effectively been reindexed to `1.0`
0.9275624570410582
```

# Extended Help

While `ForwardStarting` could be nested so that, e.g. the third period's curve is the one-period forward of the second period's curve, it will be more efficient to reuse the initial curve from a runtime and compiler perspective.

`ForwardStarting` is not used to construct a curve based on forward rates. See  [`Forward`](@ref) instead.
"""
struct ForwardStarting{T,U} <: AbstractYield
    curve::U
    forwardstart::T
end

function discount(c::ForwardStarting, to)
    discount(c.curve, c.forwardstart, to + c.forwardstart)
end

function Base.zero(c::ForwardStarting, to,cf::C) where {C<:CompoundingFrequency}
    z = forward(c.curve,c.forwardstart,to+c.forwardstart)
    return convert(cf,z)
end

"""
    zero(curve,time)
    zero(curve,time,CompoundingFrequency)

Return the zero rate for the curve at the given time. If not specified, will use `Periodic(1)` compounding.
"""
Base.zero(c::YC, time) where {YC<:AbstractYield} = zero(c, time, Periodic(1))
function Base.zero(c::YC, time, cf::C) where {YC<:AbstractYield,C<:CompoundingFrequency}
    return convert(cf, Continuous(c.zero(time))) # c.zero is a curve of continuous rates represented as floats. explicitly wrap in continuous before converting
end

"""
    Constant(rate::Real, cf::CompoundingFrequency=Periodic(1))
    Constant(r::Rate)

Construct a yield object where the spot rate is constant for all maturities. If `rate` is not a `Rate` type, will assume `Periodic(1)` for the compounding frequency

# Examples

```julia-repl
julia> y = Yields.Constant(0.05)
julia> discount(y,2)
0.9070294784580498     # 1 / (1.05) ^ 2
```
"""
struct Constant{T} <: AbstractYield
    rate::T
end

function Constant(rate::T, cf::C = Periodic(1)) where {T<:Real,C<:CompoundingFrequency}
    return Constant(Rate(rate, cf))
end

Base.zero(c::Constant, time) = c.rate
Base.zero(c::Constant, time, cf::CompoundingFrequency) = convert(cf, c.rate)
rate(c::Constant) = c.rate
rate(c::Constant, time) = c.rate
discount(r::Constant, time) = discount(r.rate, time)
discount(r::Constant, from, to) = discount(r.rate, to - from)
accumulation(r::Constant, time) = accumulation(r.rate, time)
accumulation(r::Constant, from, to) = accumulation(r.rate, to - from)

"""
    Step(rates,times)

Create a yield curve object where the applicable rate is the effective rate of interest applicable until corresponding time. If `rates` is not a `Vector{Rate}`, will assume `Periodic(1)` type.

# Examples

```julia-repl
julia>y = Yields.Step([0.02,0.05], [1,2])

julia>rate(y,0.5)
0.02

julia>rate(y,1.5)
0.05

julia>rate(y,2.5)
0.05
```
"""
struct Step <: AbstractYield
    rates
    times
end

Step(rates) = Step(rates, collect(1:length(rates)))

function rate(y::Step, time)
    i = findfirst(t -> time <= t, y.times)
    if isnothing(i)
        return y.rates[end]
    else
        return y.rates[i]
    end
end

function discount(y::Step, time)
    v = 1 / (1 + y.rates[1])^min(y.times[1], time)

    if y.times[1] >= time
        return v
    end

    for i = 2:length(y.times)

        if y.times[i] >= time
            # take partial discount and break
            v /= (1 + y.rates[i])^(time - y.times[i-1])
            break
        else
            # take full discount and continue
            v /= (1 + y.rates[i])^(y.times[i] - y.times[i-1])
        end

    end

    return v
end

"""
    Zero(rates, maturities; interpolation=CubicSpline())

Construct a yield curve with given zero-coupon spot `rates` at the given `maturities`. If `rates` is not a `Vector{Rate}`, will assume `Periodic(1)` type.

See [`bootstrap`](@ref) for more on the `interpolation` parameter, which is set to `CubicSpline()` by default.
"""
function Zero(rates::Vector{<:Rate}, maturities; interpolation=CubicSpline())
    # bump to a constant yield if only given one rate
    length(rates) == 1 && return Constant(first(rates))
    return _zero_inner(rates,maturities,interpolation)
end

# zero is different than the other boostrapped curves in that it doesn't actually need to bootstrap
# because the rate are already zero rates. Instead, we just cut straight to the
# appropriate interpolation function based on the type dispatch.
function _zero_inner(rates::Vector{<:Rate}, maturities, interp::CubicSpline)
    continuous_zeros = rate.(convert.(Continuous(), rates))
    return YieldCurve(
        rates,
        maturities,
        cubic_interp([0.0; maturities],[first(continuous_zeros); continuous_zeros])
    )
end

function _zero_inner(rates::Vector{<:Rate}, maturities, interp::LinearSpline)
    continuous_zeros = rate.(convert.(Continuous(), rates))
    return YieldCurve(
        rates,
        maturities,
        linear_interp([0.0; maturities],[first(continuous_zeros); continuous_zeros])
    )
end

# fallback for user provied interpolation function
function _zero_inner(rates::Vector{<:Rate}, maturities, interp)
    continuous_zeros = rate.(convert.(Continuous(), rates))
    return YieldCurve(
        rates,
        maturities,
        interp([0.0; maturities],[first(continuous_zeros); continuous_zeros])
    )
end

#fallback if `rates` aren't `Rate`s. Assume `Periodic(1)` per Zero docstring
function Zero(rates, maturities; interpolation=CubicSpline())
    Zero(Periodic.(rates, 1), maturities; interpolation)
end

function Zero(rates; interpolation=CubicSpline())
    # bump to a constant yield if only given one rate
    maturities = collect(1:length(rates))
    return Zero(rates, maturities; interpolation)
end

"""
    Par(rates, maturities; interpolation=CubicSpline())

Construct a curve given a set of bond equivalent yields and the corresponding maturities. Assumes that maturities <= 1 year do not pay coupons and that after one year, pays coupons with frequency equal to the CompoundingFrequency of the corresponding rate.

See [`bootstrap`](@ref) for more on the `interpolation` parameter, which is set to `CubicSpline()` by default.

# Examples

```julia-repl

julia> par = [6.,8.,9.5,10.5,11.0,11.25,11.38,11.44,11.48,11.5] ./ 100
julia> maturities = [t for t in 1:10]
julia> curve = Par(par,maturities);
julia> zero(curve,1)
Rate(0.06000000000000005, Periodic(1))

```
"""
function Par(rates::Vector{<:Rate}, maturities; interpolation=CubicSpline())
    # bump to a constant yield if only given one rate
    if length(rates) == 1
        return Constant(rate[1])
    end
    return YieldCurve(
        rates,
        maturities,
        # assume that maturities less than or equal to 12 months are settled once, otherwise semi-annual
        # per Hull 4.7
        bootstrap(rates, maturities, [m <= 1 ? nothing : 1 / r.compounding.frequency for (r, m) in zip(rates, maturities)], interpolation)
    )
end

function Par(rates::Vector{T}, maturities; interpolation=CubicSpline()) where {T<:Real}
    return Par(Rate.(rates), maturities; interpolation)
end

"""
    Forward(rate_vector,maturities)

Takes a vector of 1-period forward rates and constructs a discount curve. If rate_vector is not a vector of `Rates` (ie is just a vector of `Float64` values), then
the assumption is that each value is `Periodic` rate compounded once per period.

# Examples

```julia-repl
julia> Yields.Forward( [0.01,0.02,0.03] );

julia> Yields.Forward( Yields.Continuous.([0.01,0.02,0.03]) );

```
"""
function Forward(rates::Vector{<:Rate}, maturities)
    # convert to zeros and pass to Zero
    disc_v = Vector{Float64}(undef, length(rates))

    v = 1.0

    for i = 1:length(rates)
        Δt = maturities[i] - (i == 1 ? 0 : maturities[i-1])
        v *= discount(rates[i], Δt)
        disc_v[i] = v
    end

    z = (1.0 ./ disc_v) .^ (1 ./ maturities) .- 1 # convert disc_v to zero
    return Zero(z, maturities)
end

# if rates isn't a vector of Rates, then we assume periodic rates compounded once per period.
function Forward(rates, maturities)
    return Forward(Yields.Periodic.(rates, 1), maturities)
end

Forward(rates) = Forward(rates, collect(1:length(rates)))

"""
    Yields.CMT(rates, maturities; interpolation=CubicSpline())

Takes constant maturity (treasury) yields (bond equivalent), and assumes that instruments <= one year maturity pay no coupons and that the rest pay semi-annual.

See [`bootstrap`](@ref) for more on the `interpolation` parameter, which is set to `CubicSpline()` by default.

# Examples

```
# 2021-03-31 rates from Treasury.gov
rates =[0.01, 0.01, 0.03, 0.05, 0.07, 0.16, 0.35, 0.92, 1.40, 1.74, 2.31, 2.41] ./ 100
mats = [1/12, 2/12, 3/12, 6/12, 1, 2, 3, 5, 7, 10, 20, 30]

Yields.CMT(rates,mats)
```
"""
function CMT(rates::Vector{T}, maturities; interpolation=CubicSpline()) where {T<:Real}
    rs = map(zip(rates, maturities)) do (r, m)
        if m <= 1
            Rate(r, Periodic(1 / m))
        else
            Rate(r, Periodic(2))
        end
    end

    CMT(rs, maturities;interpolation)
end

function CMT(rates::Vector{<:Rate}, maturities; interpolation=CubicSpline())
    return YieldCurve(
        rates,
        maturities,
        # assume that maturities less than or equal to 12 months are settled once, otherwise semi-annual
        # per Hull 4.7
        bootstrap(rates, maturities, [m <= 1 ? nothing : 0.5 for m in maturities], interpolation)
    )
end

"""
    OIS(rates,maturities)
Takes Overnight Index Swap rates, and assumes that instruments <= one year maturity are settled once and other agreements are settled quarterly with a corresponding CompoundingFrequency.

See [`bootstrap`](@ref) for more on the `interpolation` parameter, which is set to `CubicSpline()` by default.

"""
function OIS(rates::Vector{T}, maturities; interpolation=CubicSpline()) where {T<:Real}
    rs = map(zip(rates, maturities)) do (r, m)
        if m <= 1
            Rate(r, Periodic(1 / m))
        else
            Rate(r, Periodic(4))
        end
    end

    return OIS(rs, maturities; interpolation)
end
function OIS(rates::Vector{<:Rate}, maturities ; interpolation=CubicSpline())
    return YieldCurve(
        rates,
        maturities,
        # assume that maturities less than or equal to 12 months are settled once, otherwise quarterly
        # per Hull 4.7
        bootstrap(rates, maturities, [m <= 1 ? nothing : 1 / 4 for m in maturities], interpolation)
    )
end


immutable IIRFilter
	filterobjPtr::Ptr{Void}
end

function Base.display(filt::IIRFilter)
	ccall((:iirfilt_rrrf_print,"libliquid"), Void, (Ptr{Void}, ), filt.filterobjPtr);
end

function createiirfilter(b::Array{Float32}, a::Array{Float32})
	ptr = ccall((:iirfilt_rrrf_create, "libliquid"), Ptr{Void}, (Ptr{Float32}, UInt, Ptr{Float32}, UInt),
		b,length(b), a, length(a));
	return IIRFilter(ptr);
end

@enum IIRFType BUTTER=0 CHEBY1 CHEBY2 ELLIP BESSEL
@enum IIRBType LPF=0 HPF BPF BSF

function getftype(input::AbstractString)
	if(input in ["lo", "low", "lowpass", "lpf"])
		return LPF
	elseif(input in ["hi","high", "hpf","highpass"])
		return HPF
	elseif(input in ["band", "bandpass", "bpf"])
		return BPF
	elseif(input in ["bandstop", "bsf"])
		return BSF
	else
		error("invalid band-type: " , input)
	end
end

function iir_filt(fc::Float64, f0::Float64, dtype::IIRFType,	band::AbstractString = "lpf",
	order::Int = 10,   Ap::Float64=1.0, As::Float64=60.0)
    btype = getftype(band);
	N = order;
	if(btype == BPF || btype == BSF)
		N = 2*order;
	end
	# r = N % 2;
	# L = (N-r) / 2;
	B = Array(Cfloat, Int(N+1));
	A = Array(Cfloat, Int(N+1));
	if(fc<=0.0 || fc>=0.5)
		error("cut-off frequency must be normalized and in range (0,0.5)")
	end
	if(f0<=0.0 || f0>=0.5)
		error("center frequency must be normalized and in range (0,0.5)")
	end
	ccall((:liquid_iirdes, "libliquid"), Void, (UInt, UInt, UInt, UInt, Cfloat,Cfloat,Cfloat,Cfloat,Ptr{Cfloat},Ptr{Cfloat}),
		dtype, btype, 1, N, Cfloat(fc), Cfloat(f0), Cfloat(Ap), Cfloat(As), B,A);
	return B,A
end

# liquid_iirdes(_ftype, _btype, _format, _n, _fc, _f0, _Ap, _As, *_B, *_A);
function butterworth(fc::Float64, f0::Float64, 	band::AbstractString = "lpf",
	order::Int = 10,   Ap::Float64=1.0, As::Float64=60.0)
	B,A = iir_filt(fc, f0, BUTTER, band, order, Ap, As )
end

function cheby1(fc::Float64, f0::Float64, 	band::AbstractString = "lpf",
	order::Int = 10,   Ap::Float64=1.0, As::Float64=60.0)
	B,A = iir_filt(fc, f0, CHEBY1, band, order, Ap, As )
end

function cheby2(fc::Float64, f0::Float64, 	band::AbstractString = "lpf",
	order::Int = 10,   Ap::Float64=1.0, As::Float64=60.0)
	B,A = iir_filt(fc, f0, CHEBY2, band, order, Ap, As )
end

function ellip(fc::Float64, f0::Float64, 	band::AbstractString = "lpf",
	order::Int = 10,   Ap::Float64=1.0, As::Float64=60.0)
	B,A = iir_filt(fc, f0, ELLIP, band, order, Ap, As )
end

function bessel(fc::Float64, f0::Float64, 	band::AbstractString = "lpf",
	order::Int = 10,   Ap::Float64=1.0, As::Float64=60.0)
	B,A = iir_filt(fc, f0, BESSEL, band, order, Ap, As )
end


# This file is part of the ValidatedNumerics.jl package; MIT licensed

doc"""An `IntervalBox` is an $N$-dimensional rectangular box, given
by a Cartesian product of $N$ `Interval`s.
"""

immutable IntervalBox{N, T} <: FixedVector{N, Interval{T}}  # uses FixedSizeArrays.jl package
    _ :: NTuple{N, Interval{T}}
end

IntervalBox(x::Interval) = IntervalBox( (x,) )  # single interval treated as tuple with one element

## arithmetic operations
# Note that standard arithmetic operations are implemented automatically by FixedSizeArrays.jl

mid(X::IntervalBox) = [mid(x) for x in X]

## set operations

⊆(X::IntervalBox, Y::IntervalBox) = all([x ⊆ y for (x,y) in zip(X, Y)])

∩(X::IntervalBox, Y::IntervalBox) = IntervalBox([x ∩ y for (x,y) in zip(X, Y)]...)

isempty(X::IntervalBox) = any(isempty, X)

diam(X::IntervalBox) = maximum([diam(x) for x in X])

emptyinterval(X::IntervalBox) = IntervalBox(map(emptyinterval, X))


@testset "Grid Search" begin
    @test Optim.grid_search(x -> (1.0 - x)^2, [-1:0.1:1.0;]) == 1
end
# Cartesian product over 2 dimensions.
#grid_search(x -> (1.0 - x[1])^2 + (2.0 - x[2])^2, product[-1:0.1:1.0, -1:0.1:1.0])


using OrdinaryDiffEq, Zygote, DiffEqSensitivity, Test

p_model = [1f0]

u0 = Float32.([0.0])

function dudt(du, u, p, t)
    du[1] = p[1]
end

prob = ODEProblem(dudt,u0,(0f0,99.9f0))

function predict_neuralode(p)
    _prob = remake(prob,p=p)
    Array(solve(_prob,Tsit5(), saveat=0.1))
end

loss(p) = sum(abs2,predict_neuralode(p))/length(p)

p_model_ini = copy(p_model)

@test !iszero(Zygote.gradient(loss,p_model_ini)[1])

## https://github.com/SciML/DiffEqSensitivity.jl/issues/675

u0 = Float32[2.0; 0.0] # Initial condition
p = [-0.1 2.0; -2.0 -0.1]

datasize = 30 # Number of data points
tspan = (0.0f0, 1.5f0) # Time range
# tsteps = range(tspan[1], tspan[2], length = datasize) # Split time range into equal steps for each data point
tsteps = (rand(datasize) .* (tspan[2] - tspan[1])  .+ tspan[1]) |> sort

function f(du,u,p,t)
    du .= p*u
end

function loss(p)
    prob = ODEProblem(f,u0,tspan,p)
    sol = solve(prob,Tsit5(),saveat=tsteps,sensealg=InterpolatingAdjoint())
    sum(sol)
end

Zygote.gradient(loss, p)[1]

@test !(Zygote.gradient(loss, p)[1] .|> iszero |> all)


# Regularize

"""
  Regularize(f::ProximableFunction, rho::Real, a::AbstractArray)

Given function `f`, returns `g(x) = f(x) + (rho/2)||x-a||²`.
"""

immutable Regularize{T <: ProximableFunction, S <: Real, A <: Union{Real, AbstractArray}} <: ProximableFunction
  f::T
  rho::S
  a::A
  function Regularize(f::T, rho::S, a::A)
    if rho <= 0.0
      error("parameter ρ must be positive")
    else
      new(f, rho, a)
    end
  end
end

Regularize{T <: ProximableFunction, S <: Real, A <: AbstractArray}(f::T, rho::S, a::A) = Regularize{T, S, A}(f, rho, a)

Regularize{T <: ProximableFunction, S <: Real}(f::T, rho::S=1.0, a::S=0.0) = Regularize{T, S, S}(f, rho, a)

function (g::Regularize){T <: RealOrComplex}(x::AbstractArray{T})
	return g.f(x) + g.rho/2*vecnorm(x-g.a)^2
end

function prox!{T <: RealOrComplex}(g::Regularize, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  gr= g.rho*gamma
  gr2= 1/(1+gr)
  v = prox!(g.f, gr2.*(x+gr.*g.a), y, gr2*gamma)
  return v + g.rho/2*vecnorm(y-g.a)^2
end

function prox_naive{T <: RealOrComplex}(g::Regularize, x::AbstractArray{T}, gamma::Real=1.0)
  y, v = prox_naive(g.f, x./(1+gamma*g.rho)+g.a./(1/(gamma*g.rho)+1), gamma/(1+gamma*g.rho))
  return y, v + g.rho/2*vecnorm(y-g.a)^2
end

fun_name(f::Regularize) = string("Regularized ", fun_name(f.f))
fun_dom(f::Regularize) = fun_dom(f.f)
fun_expr(f::Regularize) = string(fun_expr(f.f),"+(ρ/2)||x-a||²")
fun_params(f::Regularize) = "ρ = $(f.rho), λ = $(f.f.lambda), a = $( typeof(f.a)<:Real ? f.a :typeof(f.a) )"


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 54fc6142-ff38-11ea-2884-ed7abed903ca
using Pkg, DrWatson

# ╔═╡ 54fc9980-ff38-11ea-134f-cffec3a049e6
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanSample
	using StatisticalRethinking
end

# ╔═╡ 550fc4ee-ff38-11ea-0bc0-8db9a84657b2
include(projectdir("models", "06", "m6.7s.jl"))

# ╔═╡ d6c8e622-ff39-11ea-38bb-513629d16e91
include(projectdir("models", "06", "m6.8s.jl"))

# ╔═╡ 8377e266-ff34-11ea-2963-b7dce70656f5
md"## Clip-06-20s.jl"

# ╔═╡ 54fd1328-ff38-11ea-22ee-03b3cbfb8b64
begin
	N = 1000
	df = DataFrame(
	  :h0 => rand(Normal(10,2 ), N),
	  :treatment => vcat(zeros(Int, Int(N/2)), ones(Int, Int(N/2))),
	  :M => rand(Bernoulli(), N)
	);

	d(i) = Binomial(1, 0.5 - 0.4 * df[i, :treatment] + 0.4 * df[i, :M])
	df.fungus = [rand(d(i), 1)[1] for i in 1:N]
	df.h1 = [df[i, :h0] + rand(Normal(5 + 3 * df[i, :M]), 1)[1] for i in 1:N]
end

# ╔═╡ 5509d90a-ff38-11ea-0563-a9cfe86b4e1e
md"##### Execute m6.7s & m6.8s."

# ╔═╡ 55103e56-ff38-11ea-0825-8b25859b9e32
begin
	(s, p) = plot_model_coef([m6_7s, m6_8s], [:a, :bt, :bf])
	p
end

# ╔═╡ 551932c2-ff38-11ea-28cd-5bfa4ad04d57
s

# ╔═╡ 5519a194-ff38-11ea-26c3-0140d22578e4
md"## End of clip-06-20s.jl"

# ╔═╡ Cell order:
# ╟─8377e266-ff34-11ea-2963-b7dce70656f5
# ╠═54fc6142-ff38-11ea-2884-ed7abed903ca
# ╠═54fc9980-ff38-11ea-134f-cffec3a049e6
# ╠═54fd1328-ff38-11ea-22ee-03b3cbfb8b64
# ╟─5509d90a-ff38-11ea-0563-a9cfe86b4e1e
# ╠═550fc4ee-ff38-11ea-0bc0-8db9a84657b2
# ╠═d6c8e622-ff39-11ea-38bb-513629d16e91
# ╠═55103e56-ff38-11ea-0825-8b25859b9e32
# ╠═551932c2-ff38-11ea-28cd-5bfa4ad04d57
# ╟─5519a194-ff38-11ea-26c3-0140d22578e4


function numericsummary(INPUT::Union{Array,DataFrame},VarNames::Symbol)  #数値型の基本統計量の算出
  SummaryVar=StatsBase.var(INPUT)
  SummaryStd=StatsBase.std(INPUT)
  SummaryQuartile=StatsBase.quantile( INPUT , [0.00, 0.25, 0.50, 0.75, 1.00])
  SummaryMedian=StatsBase.median(INPUT)
  SummaryMean=StatsBase.mean(INPUT)

  Output=DataFrame(colname=VarNames,Var=SummaryVar,Std=SummaryStd,Mean=SummaryMean, Min=SummaryQuartile[1] ,Quartile1st=SummaryQuartile[2] ,Median=SummaryMedian,Quartile3rd=SummaryQuartile[4],Max=SummaryQuartile[5])

  return Output
end

function nonnumericsummary(INPUT::Union{Array,DataFrame},VarNames::Symbol)  #非数値型の基本統計量の算出
  Output=combine(groupby(INPUT, VarNames, sort=false, skipmissing=false), VarNames=>df -> size(df, 1))
  rename!(Output,[VarNames,:count])

  return Output
end

function summarymerge(INPUT::DataFrame, Summary::DataFrame ,FirstMergeFlag::Bool)  #データの集約
  if FirstMergeFlag ==true
    Summary=INPUT
    FirstMergeFlag=false
  else
    Summary=vcat(Summary,INPUT)
  end

  return Summary,FirstMergeFlag
end

function univariate(INPUT::DataFrame;graphplot::Bool=false)
  VarNames=propertynames(INPUT)
  Vartype=eltype.(eachcol(INPUT))
  SummaryNonNum=DataFrame(colname="",hist=0)
  SummaryNum=DataFrame(colname="",Var=0.0,Std=0.0,Mean=0.0, Min=0.0 ,Quartile1st=0.0 ,Median=0.0,Quartile3rd=0.0,Max=0.0)
  FirstMergeFlagNum=true
  FirstMergeFlagNonNum=true
  typearray=0
  typearrayflag=true

  for  i = 1 : size(INPUT,2)
    if Vartype[i] <: Real
      #数値型の基本統計量の算出
      work=numericsummary(INPUT[:,i],VarNames[i])

      #基本統計量の集約
      SummaryNum,FirstMergeFlagNum=summarymerge(work,SummaryNum,FirstMergeFlagNum)
      #列がReal型かどうか判定
      if typearrayflag
        typearray=i
        typearrayflag=false
      else
        typearray=hcat(typearray,i)
      end
    else
      #非数値型の基本統計量の算出
      work=nonnumericsummary(INPUT,VarNames[i])

      #基本統計量の集約
      SummaryNonNum,FirstMergeFlagNonNum=summarymerge(work,SummaryNonNum,FirstMergeFlagNonNum)
    end
  end

  if(graphplot)
   # gr()
   # @df INPUT corrplot(cols(typearray),grid=true)
  end

  return SummaryNum,SummaryNonNum
end


struct IterativeInitializationCollisionModel{T <: Real, UT <: AbstractArray{T, 3}} <:
       CollisionModel
    τ::T

    # The incompressible equilibrium funciton can be split in a linear and a nonlinear component
    # the linear component depends on the local density while the nonlinear component depends
    # on the local velocity.
    # Since this collision operator keeps a constant velocity for each lattice node we can
    # precompute the nonlinear term for some performance improvements
    nonlinear_term::UT
end
function IterativeInitializationCollisionModel(
    q::Quadrature,
    τ::T,
    problem,
) where {T <: Real}
    nx, ny, nf = problem.NX, problem.NY, length(q.weights)
    nonlinear_term = Array{T}(undef, nx, ny, nf)
    ρ_0 = one(T)

    # Initialize f with the given velocity field
    x_range, y_range = range(problem)
    for x_idx in 1:nx, y_idx in 1:ny
        u = lattice_velocity(q, problem, x_range[x_idx], y_range[y_idx])
        u_squared = u[1]^2 + u[2]^2

        for f_idx in 1:nf
            u_dot_xi = dot(q.abscissae[:, f_idx], u)

            cs = q.speed_of_sound_squared
            a_H_1 = cs * u_dot_xi
            a_H_2 = cs^2 * (u_dot_xi * u_dot_xi) - cs * u_squared

            nonlinear_term[x_idx, y_idx, f_idx] =
                q.weights[f_idx] * ρ_0 * (a_H_1 + a_H_2 / 2)
        end
    end

    return IterativeInitializationCollisionModel(τ, nonlinear_term)
end
function collide!(
    collision_model::IterativeInitializationCollisionModel,
    q::Quadrature,
    f_in,
    f_out;
    time = 0.0,
)
    τ = collision_model.τ
    nx, ny, nf = size(f_out)
    for x_idx in 1:nx, y_idx in 1:ny
        ρ = density(q, f_in[x_idx, y_idx, :])
        @inbounds for f_idx in 1:nf
            feq = q.weights[f_idx] * ρ + collision_model.nonlinear_term[x_idx, y_idx, f_idx]

            f_out[x_idx, y_idx, f_idx] =
                (1 - 1 / τ) * f_in[x_idx, y_idx, f_idx] + (1 / τ) * feq
        end
    end
    return
end


using Random
using LinearAlgebra
using Distances

"""
Gaussian/RBF/Squared Exponential correlation function
"""
function rbf(x, y, θ=1.0)
    1/sqrt(2*pi)*exp.(-θ*0.5*(norm(x .- y))^2)
end

"""
Gaussian/RBF/Squared Exponential correlation function
"""
function rbf_single(dist, θ=1.0)
    1/sqrt(2*pi)*exp.(-θ*0.5*(dist)^2)
end

"""
Derivative of Gaussian/RBF/Squared Exponential correlation function
with respect to hyperparameter theta
"""
function rbf_prime(x, y, θ=1.0)
    - 1/sqrt(2*pi) * 0.5 * norm(x.-y)^2 * exp.(-θ*0.5*(norm(x .- y))^2)
end

"""
Derivative of Gaussian/RBF/Squared Exponential correlation function
with respect to hyperparameter theta
"""
function rbf_prime_single(dist, θ=1.0)
    - 1/sqrt(2*pi) * 0.5 * dist^2 * exp.(-θ*0.5*(dist)^2)
end

"""
Second derivative of Gaussian/RBF/Squared Exponential correlation function
with respect to hyperparameter theta
"""
function rbf_prime2(x, y, θ=1.0)
    1/sqrt(2*pi) * 0.25 * norm(x.-y)^4 * exp.(-θ*0.5*(norm(x .- y))^2)
end

"""
Second derivative of Gaussian/RBF/Squared Exponential correlation function
with respect to hyperparameter theta
"""
function rbf_prime2_single(dist, θ=1.0)
    1/sqrt(2*pi) * 0.25 * dist^4 * exp.(-θ*0.5*(dist)^2)
end

"""
Builds cross-covariance matrix using corr kernel function.
Ex. RBF, exponential, Matern

Inputs are assumed to be 1D arrays or 2D column vectors.
Returns matrix of size len(s1) times len(s2)
"""
function K(s1, s2, θ, corr=rbf)
    K = zeros(size(s1, 1), size(s2, 1))
    if s1 != s2
        for i = 1:size(s1, 1)
            for j = 1:size(s2, 1)
                cur = corr(s1[i, :], s2[j, :], θ)
                K[i, j] = isa(cur, Array) ? cur[1] : cur
            end
        end
    else #s1==s2 can compute half the matrix and fill in other half
        for i = 1:size(s1, 1)
            for j = 1:i
                cur = corr(s1[i, :], s2[j, :], θ)
                K[i, j] = K[j, i] = isa(cur, Array) ? cur[1] : cur
            end
        end
    end
    return size(s1, 1)==size(s2, 1) ? K+UniformScaling(1e-8) : K
end

#rows of s1 and s2 contain the coordinates of each location
function fastK(s1, s2, θ, corr=rbf_single)
    corrtheta = dist -> corr(dist, θ)
    res = corrtheta.(pairwise(Euclidean(), s1, s2, dims=1))
    return size(s1, 1)==size(s2, 1) ? res + UniformScaling(1e-8) : res
end

"""
Samples from a Normal distribution with mean 0 and covariance matrix K
Scales result by 40. Works in 1D.
"""
function sample(K)
    c = cholesky(K)
    v = 40*randn(size(K, 1), 1)
    return c.L*(v)
end


##########################################################
#           Implementación de Backpropagation
# Basado en el algoritmo 4 y 5, sección 5.3 de la memoria
##########################################################
export backpropagation!

using Random

VectorOrMatrix = Union{Matrix,Vector}
function descent_gradient(
    neural_network::AbstractOneLayerNeuralNetwork,
    X_train :: VectorOrMatrix,
    Y_train :: VectorOrMatrix,
    activation_function,
    derivative_activation_funcion,
    batch_size :: Int = 32,
    )
    len, _ = size(X_train)
    if len < batch_size
        error("batch_size should be equal or smaller than the size of train dataset,
        but $batch_size  > $len")
    end

    # Derivadas parciales a calcular
    n_1,d = size(neural_network.W1)
    s,n = size(neural_network.W2)

    partial_W1 = zeros(Float64, (n_1,d))
    partial_W2 = zeros(Float64, (s,n))
    derivative_B = zeros(Float64, (s,n))
    # Variables auxiliares para reducir número de operaciones
    sensibilities_img = zeros(Float64, n)

    index = first(randperm(len),batch_size)
    # Para cada muestra del entrenamiento
    for train_index in index
        # 1. Forward propagation almacenando resultado relevantes
        𝛅 = neural_network.W1 * push!(copy(X_train[train_index,:]),1)
        sensibilities_img = map(activation_function, 𝛅)
        derivative_sensibilities_img= map(derivative_activation_funcion, 𝛅)

        forward_propagation_x = neural_network.W2 * sensibilities_img
        error = forward_propagation_x .- Y_train[train_index,:]

        # 2. Parcial de B (W_2)
        for v in 1:n
            for w in 1:s
                partial_W2[w,v] = partial_W2[w,v] + error[w]*sensibilities_img[v]
            end
        end

        # 3. Parcial W_1
        for i in 1:n
            for k in 1:s
            derivative_B[k,i] =
                neural_network.W2[k,i]*derivative_sensibilities_img[i]
            end
        end

        for v in 1:n
            for k in 1:s
                difference_times_derivative = error[k]*derivative_B[k,s]
                partial_W2[k, v] += error[k]*derivative_sensibilities_img[v]
                partial_W1[v,d] += difference_times_derivative

                for u in 1:(d-1)
                    partial_W1[v,u] += difference_times_derivative * X_train[train_index, u]
                end
            end
        end

     end
     return partial_W1, partial_W2
end

"""
    function backpropagation!(
        neural_network::AbstractOneLayerNeuralNetwork,
        X_train :: Matrix,
        Y_train :: Vector,
        activation_function,
        derivative_activation_funcion,
        batch_size :: Int = 32,
        η :: Float64 =  0.1
    )

Compute backpropagation
"""
function backpropagation!(
    neural_network::AbstractOneLayerNeuralNetwork,
    X_train :: Matrix,
    Y_train :: Vector,
    activation_function,
    derivative_activation_funcion,
    batch_size :: Int = 32,
    η :: Float64 =  0.1
)
    ∂w1, ∂w2 = descent_gradient(neural_network,
         X_train, Y_train,
         activation_function, derivative_activation_funcion,
        batch_size)
        neural_network.W1 -= η*∂w1
        neural_network.W2 -= η*∂w2
    return neural_network
end


#=
# Steady frictional flow through a constant-area duct
This notebook demonstrates the use of tools for computing flow through a constant-area
 duct with frictional effects, also known as **Fanno flow**
=#

# ### Set up the module
using Gasdynamics1D
#-
using Plots
using LaTeXStrings

#=
### Some general aspects of Fanno flow
Most calculations with frictional flow involve the relationship between Mach
number and the reference length $L^*$, representing the length required from that
point to bring the flow to Mach number equal to 1. The relationship is based on
the dimensionless parameter

$$\dfrac{f L^*}{D}$$

where $f$ is the Darcy friction factor and $D$ is the duct diameter. Let's plot
 this relationship for air:
=#
Mrange = range(0.1,8,length=801)
fLDarray = []
for M in Mrange
    push!(fLDarray,value(FLStarOverD(MachNumber(M),FannoFlow,gas=Air)))
end
#-
plot(Mrange,fLDarray,xlim=(0,8),ylim=(0,10),xticks=0:1:8,xlabel="Mach number",ylabel=L"fL^*/D",legend=false)

#=
Notice that the reference length goes to zero at Mach number 1, as it should, by
its definition. It should also be noted that there are a few values of $fL^*/D$ for
which there are two possible Mach numbers---one subsonic and one supersonic. This
is analogous to the situation in isentropic flow with area changes.
=#
#=
### Example 1
Flow enters a duct of diameter 2 cm at Mach number equal to 0.1 and leaves at
Mach number equal to 0.5. Assume that the friction factor is equal to 0.024 throughout.

(a) What is the length of the duct?

(b) How much longer would the duct have to be to reach Mach number equal to 1 at the exit?
=#
#=
First, we use $M_1 = 0.1$ to calculate $f L_1^*/D$. We will do the same with $M_2$
 to calculate $f L_2^*/D$. The actual length $L$ is given by the difference between
 $L_1^*$ and $L_2^*$.
=#
f = FrictionFactor(0.024)
D = Diameter(2u"cm")
#-
M1 = MachNumber(0.1)
fL1star_over_D = FLStarOverD(M1,FannoFlow)
#-
M2 = MachNumber(0.5)
fL2star_over_D = FLStarOverD(M2,FannoFlow)
# Now calculate $L_1^*$ and $L_2^*$, and $L$ from their difference:
L1star = Length(fL1star_over_D*D/f)
L2star = Length(fL2star_over_D*D/f)
L = Length(L1star-L2star)
# The additional length needed to bring the flow to Mach number 1 is, by definition, $L_2^*$:
L2star
# So we only need 0.89 m more to bring the flow to sonic (choked) conditions.

#=
###  Example 2
Consider a flow of air through a circular pipe of diameter 3 cm, starting at
stagnation pressure 200 kPa, stagnation temperature 500 K, and velocity 100 m/s.
The friction factor in the pipe can be assumed to be 0.02 throughout.

(a) What length of pipe would be needed to drive the flow to Mach number 1 at the exit?

(b) What is the mass flow rate through the pipe at those conditions?

(c) Suppose the pipe were 30 m long. How would the mass flow rate change?
=#

# For (a), we are seeking $L_1^*$. First set the known conditions:
D = Diameter(3u"cm")
p01 = StagnationPressure(200u"kPa")
T01 = StagnationTemperature(500)
u1 = Velocity(100)
f = FrictionFactor(0.02)

#=
Now calculate the Mach number at the entrance. For this, we need $c_1$, the speed
of sound. We get this by using the known stagnation temperature $T_{01} = 500$ K
and velocity $u_1 = 100$ m/s. We will calculate the temperature $T_1$ from these
(using $h_1 = h_{01} - u_1^2/2$), and $c_1$ from that:
=#
h01 = StagnationEnthalpy(T01)
h1 = Enthalpy(h01,u1) # this computes h1 = h01 - u1^2/2
T1 = Temperature(h1)
c1 = SoundSpeed(T1)
M1 = MachNumber(u1/c1)
# Now we can compute $fL_1^*/D$ and then $L_1^*$. This is what we seek.
fL1starD = FLStarOverD(M1,FannoFlow)
L1star = Length(fL1starD*D/f)

#=
so we need a pipe of 16.6 m to reach Mach number 1.

Now let's compute the mass flow rate. We need the area and density. For the density,
we will first get the stagnation density (from $p_{01}$ and $T_{01}$) and then get
 density from the isentropic relation. (We are not using isentropic relations
 between two different points in the duct; we are using it to relate *local*
 stagnation conditions to *local* conditions.)
=#
A = Area(D)
ρ01 = StagnationDensity(p01,T01)
ρ1 = Density(ρ01,M1,Isentropic)
# and the mass flow rate is:
ṁ = MassFlowRate(ρ1*u1*A)

#=
(c) Now suppose the pipe is 30 m long. This becomes our new $L_1^*$, and all of
the conditions at the entrance must adjust accordingly. First find the Mach number
=#
Lstar = L = Length(30u"m")
fLstar_over_D = FLOverD(f*Lstar/D)
M1 = SubsonicMachNumber(fLstar_over_D,FannoFlow)

#=
Note that the entrance Mach number is lower. The flow slows down due to the longer pipe.
Now find the new values of $\rho_1$ and $u_1$ (by finding $c_1$):
=#
T1 = Temperature(T01,M1,Isentropic)
ρ1 = Density(ρ01,M1,Isentropic)
c1 = SoundSpeed(T1)
u1 = Velocity(c1,M1)
ṁ = MassFlowRate(ρ1*u1*A)

#=
So the mass flow rate is smaller than it was with the shorter pipe! For the same
reservoir conditions, less mass flow rate is delivered to a longer (choked) pipe.
=#


using CompScienceMeshes
using Test

ignd(x,y) = x^2 * cos(x*y)

Iref = 0.082739596439277

p1 = [0.0, 0.0]
p2 = [1.0, 0.0]
p3 = [0.0, 1.0]

N = length(CompScienceMeshes.triangleGaussW)
J = zeros(N)
for _n in 1 : N
    u, w = trgauss(_n)
    for _g in 1 : length(w)
        _p = u[1,_g]*p1 + u[2,_g]*p2 + (1-u[1,_g]-u[2,_g])*p3
        x, y = _p[1], _p[2]
        J[_n] += w[_g] * ignd(x,y)
    end
end

E = [0.3326898477547438,0.00128868370286643,0.0003505830254662883,0.0002990068871080405,0.00026769940415425264,1.4413762136992496e-7,1.8717613253691006e-7,1.8109645370453297e-12,3.4565487801639488e-12,3.0571872387445795e-12,3.3689945069086314e-12,3.3050899511418194e-12,3.3156568461898747e-12]

for _q in 1 : N
    @test abs(J[_q]-Iref)/abs(Iref) <=  E[_q]*2.0
end


#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 01-2018
#--------------------------------------------------------------------

function sconv(bnd1::Function, bnd2::Function, frhs::Function, asol::Function,
                                                                   Nx::Int, Ny::Int, M::Int)::Float64
    dbase = fdistribute(bnd1, bnd2, frhs, Nx, Ny, M)
    sumL2 = 0.0
    for m in 1:M, n in 1:M
        sPatch = interpolatePatch(dbase[[m,n]], 4Nx, 4Ny).value
        fPatch = Float64[asol(x,y) for x in chebgrid(4Nx, M, m), y in chebgrid(4Ny, M, n)]
        sumL2 += (L2norm(fPatch, sPatch, chebweights(4Nx)/M, chebweights(4Ny)/M))^2
    end
    return sqrt(sumL2)
end

function showconv(bnd1::Function, bnd2::Function, frhs::Function, analyticsol::Function,
                                                               maxmodes::Int, maxlevels::Int, h::Int)
    if maxlevels > 1
        L2 = 1
        p  = maxmodes
        println("----------------------------------------------------------------------------------")
        println("==> h-convergence")
        println("----------------------------------------------------------------------------------")
        for level in 0:maxlevels
            L2error = sconv(bnd1, bnd2, frhs, analyticsol, p, p, h^level)
            level == 0 ? R = 0 : R =  L2/L2error
            @printf("p = %s | np = %s | h^(p+1) = %s | R = %e | L2 error = %e \n",
                    lpad(p,4," "), lpad(h^level,4," "),  lpad(h^(p+1),4," "), R, L2error)
            L2 = L2error
        end
    else
        println("----------------------------------------------------------------------------------")
        println("==> p-convergence")
        println("----------------------------------------------------------------------------------")
        L2 = 1
        h  = 1
        level = 0
        for p in 1:2:maxmodes
            L2error = sconv(bnd1, bnd2, frhs, analyticsol, p, p, h^level)
            p == 1 ? R = 0 : R =  L2/L2error
            @printf("p = %s | np = %s | h^(p+1) = %s | R = %e | L2 error = %e \n",
                    lpad(p,4," "), lpad(h^level,4," "),  lpad(h^(p+1),4," "), R, L2error)
            L2 = L2error
        end
    end
end


## Binomial Example ####
## Note: Adapted from the Rate_4 example in Bayesian Cognitive Modeling
##  https://github.com/stan-dev/example-models/tree/master/Bayesian_Cognitive_Modeling

using Mamba, Stan

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  binomialstanmodel = "
  // Inferring a Rate
  data {
    int<lower=1> n;
    int<lower=0> k;
  }
  parameters {
    real<lower=0,upper=1> theta;
    real<lower=0,upper=1> thetaprior;
  }
  model {
    // Prior Distribution for Rate Theta
    theta ~ beta(1, 1);
    thetaprior ~ beta(1, 1);

    // Observed Counts
    k ~ binomial(n, theta);
  }
  generated quantities {
    int<lower=0> postpredk;
    int<lower=0> priorpredk;

    postpredk <- binomial_rng(n, theta);
    priorpredk <- binomial_rng(n, thetaprior);
  }
  "

  global stanmodel, rc, sim1
  stanmodel = Stanmodel(name="binomial", model=binomialstanmodel)

  binomialdata = [
    Dict("n" => 10, "k" => 5)
  ]

  rc, sim1 = stan(stanmodel, binomialdata, ProjDir, diagnostics=false,
    CmdStanDir=CMDSTAN_HOME)

  if rc == 0
    theta_sim = sim1[1:size(sim1,1), ["theta", "thetaprior"], :]
    predk_sim = sim1[1:size(sim1,1), ["postpredk", "priorpredk"], :]

    p = plot(theta_sim, [:trace, :mean, :density, :autocor], legend=false);
    draw(p, ncol=4, nrow=2, filename="$(stanmodel.name)-thetas", fmt=:svg)

    p = plot(predk_sim, [:bar], legend=false)
    draw(p, ncol=2, nrow=2, filename="$(stanmodel.name)-predictiv", fmt=:svg)
  end # rc == 0
end # cd


include("eliminacao01.jl");
include("sr00.jl");

function resolve00(A,b)

	(m,n) = size(A);

	U = eliminacao01([A b]);
	b = U[:,n+1];
	U = U[:,1:n];

	x = sr00(U,b);

	return x;
end


model = Dynamics.pendulum[1]

# Objective
n,m = model.n, model.m
x0 = [0; 0.]
xf = [pi; 0] # (ie, swing up)
Q = 1e-3*Diagonal(I,n)
Qf = 100.0*Diagonal(I,n)
R = 1e-2*Diagonal(I,m)
tf = 5.
obj_uncon = LQRObjective(Q, R, Qf, tf, x0, xf)

N = 41
U0 = ones(m,N)

# Test original problem
solver = Solver(model, obj_uncon, N=N)
res, stats = solve(solver, U0)
@test norm(res.X[N]-xf) < 1e-3

# Convert LQRObjective to Generic Objective
ell(x,u) = 0.5*x'Q*x + 0.5*u'R*u - xf'Q*x
ell(x) = 0.5*x'Qf*x - xf'Qf*x

costfun = GenericCost(ell,ell,n,m)
obj_generic = UnconstrainedObjective(costfun,tf,x0,xf)
solver_g = Solver(model, obj_generic, N=N)
res_g, stats_g = solve(solver_g, U0)
@test norm(res_g.X - res.X) == 0

# Try a non-linear cost function
ell2(x,u) = cos(x[1]) + u'R*u + x[2]^2*1e-3
ell2(x) = cos(x[1]) + x[2]^2*1e-1

costfun2 = GenericCost(ell2,ell2,n,m)
obj2 = UnconstrainedObjective(costfun2,tf,x0,xf)
solver_g2 = Solver(model, obj2, N=N)
res_g2, stats_g2 = solve(solver_g2, U0*0)
@test norm(res_g2.X[N] - xf) < 1e-6

# Constrained (non-zero initialization)
x_min = [-2*pi,-Inf]
x_max = [2*pi,Inf]
obj2_con = ConstrainedObjective(obj2,x_min=x_min,x_max=x_max)
solver_g2 = Solver(model, obj2_con, N=N)
res_g2_con, stats_g2_con = solve(solver_g2, U0)
@test norm(res_g2_con.X[N] - xf) < 1e-3


using IterativeSolvers

export bicgstabl, bicgstabl!

bicgstabl(A, b, Pl=1; kwargs...) =  bicgstabl!(IterativeSolvers.zerox(A,b), A, b, Pl; kwargs...)

function bicgstabl!(x, A, b, Pl=1; tol::Real=size(A,2)*eps(), maxiter::Int=size(A,2))
    K = KrylovSubspace(A, length(b), 1, Vector{IterativeSolvers.Adivtype(A,b)}[])
    IterativeSolvers.init!(K, x)
    bicgstabell(K; tol=tol, l=1)
end

function bicgstabell(K::KrylovSubspace; tol::Real=eps(), l::Int=1)
    k=-l
    r = zeros(3)
    IterativeSolvers.initrand!(K)
    #XXX Choose r̃[0]
    r[0]=b-IterativeSolvers.nextvec(K)
    u[-1]=0
    x[0]=K.v0
    ρ₀=1
    α=0
    ω=1

    while norm(r[k+l]) >= tol
        k+=l
        û[0]=u[k-1]; r̂[0]=r[k]; x̂[0]=x[k]
        ρ₀*=-ω
        for j=0:l-1 #BiCG part
            ρ₁=r̂[j]⋅r̃[0]; β=β[k+j]=α*ρ1/ρ₀; ρ₀=ρ₁
            for i=0:j
                û[i]=r̂[i]-β*û[i]
            end
            û[j+1]=K.A*û[j]
            γ=û[j+1]⋅r̃[0]; α=α[k+j]=ρ0/γ
            for i=0:j
                r̂[i]-=α*û[i+1]
            end
            r̂[j+1]=K.A*r̂[j]; x̂0+=α*û[0]
        end
        for j=1:l #Mod G-S; MR part
            for i=1:j-1
                τ[i,j]=r̂[j]⋅r̂[i]/σ[i]
                r̂[j]-=τ[i,j]*r̂[i]
            end
            σ[j]=r̂[j]⋅r̂[j];γ′[j]=r̂[0]⋅r̂[j]/σ[j]
        end
        ω=γ[l]=γ′[l]
        for j=l-1:-1:1 #γ=T\γ′
            γ[j]=γ′[j]-sum([τ[j,i]*γ[i] for i=j+1:l])
        end
        for j=1:l-1 #γ″=TSγ
            γ″[j]=γ[j+1]-sum([τ[j,i]*γ[i+1] for i=j+1:l-1])
        end
        x̂[0]+=γ[1]*r̂[0]; r̂[0]-=γ′[l]*r̂[l]; û[0]-=γ[l]*û[l] #update
        for j=1:l-1
            û[0]-=γ[j]*û[j]
            x̂[0]+=γ″[j]*r̂[j]
            r̂[0]-=γ′[j]*r̂[j]
        end
        u[k+l-1]=û[0]; r[k+l]=r̂[0]; x[k+l]=x̂[0]
    end
    #TODO What to return?
end


# prepare the sets N_{ij} and M_{ij} from the article
# for further details take a look at our article
function sets_m_n(airdist::Array{Float64,2}, hex::Int64)
    N = Array{Bool,3}(undef,hex,hex,hex) .= 0
    M = Array{Bool,3}(undef,hex,hex,hex) .= 0
    maxdist = Array{Float64,2}(undef,hex,hex) .= 0
    for j = 1:hex
        for v = 1:hex
            if adjacent[j,v] == 1
                for i = 1:hex
                    if airdist[i,v] < airdist[i,j]
                        N[i,j,v] = 1
                    end
                    maxdist[i,j] = max(maxdist[i,j], airdist[i,v])
                end
            end
        end
    end
    for j = 1:hex
        for v = 1:hex
            if adjacent[j,v] == 1
                for i = 1:hex
                    if airdist[i,v] < maxdist[i,j]
                        M[i,j,v] = 1
                    end
                end
            end
        end
    end
    card_n = sum(N, dims=3)[:,:,1]
    card_m = sum(M, dims=3)[:,:,1]
    return  N::Array{Bool,3},
            M::Array{Bool,3},
            card_n::Array{Int64,2},
            card_m::Array{Int64,2}
end


using ShunnHamQuadrature

using FastGaussQuadrature
using StaticArrays
using Combinatorics
using Plots

function _legendre(n,a,b)
    x, w = FastGaussQuadrature.gausslegendre(n)
    w .*= (b-a)/2
    x = (x.+1)/2*(b-a).+a
    collect(zip(x,w))
end

function duffy2D( rule)

    qps = _legendre(rule, 0.0, 1.0)

    QL = rule

    QP = QL*QL
    uP = zeros(2,QP)
    wP = zeros(QP)
    n = 1
    uP = Vector{SVector{2,Float64}}()

    for (x1,w1) in qps
        for (x2,w2) in qps
            push!(uP,[x1,x2*x1])

            wP[n]   = w1*w2*x1
            n += 1
        end
    end

    return uP,wP
end

function duffy3D(rule)

    qps = _legendre(rule, 0.0, 1.0)

    QL = rule

    QP = QL*QL*QL

    wP = zeros(QP)
    n = 1
    uP = Vector{SVector{3,Float64}}()

    for (x1,w1) in qps
        for (x2,w2) in qps
            for (x3,w3) in qps
                push!(uP,[x1,x2*x1,x3*x2*x1])

                wP[n]   = w1*w2*w3*x1^2*x2
                n += 1
            end
        end
    end

    return uP,wP
end

function duffy4D(rule)

    qps = _legendre(rule, 0.0, 1.0)

    QL = rule

    QP = QL*QL*QL*QL

    wP = zeros(QP)
    n = 1
    uP = Vector{SVector{4,Float64}}()

    for (x1,w1) in qps
        for (x2,w2) in qps
            for (x3,w3) in qps
                for (x4,w4) in qps
                    push!(uP,[x1,x2*x1,x3*x2*x1,x4*x3*x2*x1])

                    wP[n]   = w1*w2*w3*w4*x1^3*x2^2*x3
                    n += 1
                end
            end
        end
    end

    return uP,wP
end

function duffy5D(rule)

    qps = _legendre(rule, 0.0, 1.0)

    QL = rule

    QP = QL*QL*QL*QL*QL

    wP = zeros(QP)
    n = 1
    uP = Vector{SVector{5,Float64}}()

    for (x1,w1) in qps
        for (x2,w2) in qps
            for (x3,w3) in qps
                for (x4,w4) in qps
                    for (x5,w5) in qps
                        push!(uP,[x1,x2*x1,x3*x2*x1,x4*x3*x2*x1,x5*x4*x3*x2*x1])

                        wP[n]   = w1*w2*w3*w4*w5*x1^4*x2^3*x3^2*x4
                        n += 1
                    end
                end
            end
        end
    end

    return uP,wP
end

function test_one()
    return 1
end

function test_poly_3d(x,m)
    sum = 0.0
    denom = 0.0
    for i in 0:m
        for j in 0:m-i
            for k in 0:m-i-j
                sum += (i+1)*(j+1)*(k+1)*x[1]^i*x[2]^j*x[3]^k
            end
        end
        for j in 0:i+1
            denom += j
        end
    end
    return sum/denom
end

#evaluate polynom
function eval_poly(x,e)
    p = 1.0

    for (i,y) in enumerate(x)

        p*=y^e[i]
    end
    return p
end

function test_poly_2d(x,m)
    sum = 0.0
    denom = (m+1)*(m+2)
    for i in 0:m
        for j in 0:m-i
            sum += 2*(i+1)*(j+1)*x[1]^i*x[2]^j
        end
    end
    return sum/denom
end

#Test Order
exact=0.5

#vertices = [[0 0], [1 0], [1 1]]
#vertices = [[0 0 0], [1 0 0], [1 1 0], [1 1 1]]
#vertices = [[0 0 0 0], [1 0 0 0], [1 1 0 0], [1 1 1 0], [1 1 1 1]]
vertices = [[0 0 0 0 0], [1 0 0 0 0], [1 1 0 0 0], [1 1 1 0 0], [1 1 1 1 0], [1 1 1 1 1]]

(x,w) = duffy5D(12)
qps_high = zip(x,w)
println(sum(w))

order = []
err_ref =[]
err_new =[]

plot(size=(800,500),yaxis=:log)
for k in 1:6

    #=
    sh_ref = shunnham2D_ref(k)
    (x,w) = get_pts_wts(sh_ref,vertices)
    qps_ref = zip(x,w)
    println(sum(w))
    =#

    sh_new = shunnham5D(k)
    (x,w) = get_pts_wts(sh_new,vertices)
    qps_new = zip(x,w)

    println(sum(w))

    order = []
    err_ref =[]
    err_new =[]

    for i in 0:10

        exponents = collect(multiexponents(5,i))
        println(length(exponents))
        for e in exponents
            function integrand(x)
                eval_poly(x,e)
            end

            #Refrence
            exact  =  sum(w*integrand(x) for (x,w) in qps_high)

            #int_ref = sum(w*integrand(x) for (x,w) in qps_ref)
            int_new = sum(w*integrand(x) for (x,w) in qps_new)

            println(i," ",e)
            println(exact)
            #println(int_ref)
            println(int_new)

            push!(order,i)
            #push!(err_ref,abs(int_ref-exact))
            push!(err_new,abs(int_new-exact))
        end

    end

    scatter!(order,err_new,label=string(length(x))*"-Point Rule",markershape=Plots.supported_markers()[k+2])

end

plot!(xlims=(-1,13),xticks=0:1:12,ylims=(1e-17,10),ylabel="Max. rel. error",xlabel="Max. degree of polynom",legend=:outerright, title="Order of 5D Shunn-Ham Quadratures")


struct Quaternion{T<:Number}
    v::Vector3{T}
    w::T

    Quaternion{T}(v, w) where T = new{T}(v, w)

    Quaternion{T}() where T = new{T}(Vector3{T}(), one(T))
end

Quaternion(v::Vector3{U}, w::T) where {T,U} = Quaternion{promote_type(T,U)}(v, w)


module MultiFieldSparseMatrixAssemblersTests

using Gridap.Arrays
using Gridap.Algebra
using Gridap.Geometry
using Gridap.FESpaces
using Gridap.Fields
using Gridap.Integration
using SparseArrays
using Gridap.CellData
using Gridap.MultiField

using Test

order = 2

domain = (0,1,0,1)
partition = (3,3)
model = CartesianDiscreteModel(domain,partition)

degree = order
trian = get_triangulation(model)
quad = CellQuadrature(trian,degree)

trian_Γ = SkeletonTriangulation(model)
quad_Γ = CellQuadrature(trian_Γ,degree)

V = TestFESpace(model=model,order=order,reffe=:Lagrangian,conformity=:H1,valuetype=Float64)
Q = TestFESpace(model=model,order=order-1,reffe=:Lagrangian,conformity=:L2,valuetype=Float64)

U = TrialFESpace(V)
P = TrialFESpace(Q)

Y = MultiFieldFESpace([V,Q])
X = MultiFieldFESpace([U,P])

free_values = rand(num_free_dofs(X))
xh = FEFunction(X,free_values)
uh, ph = xh

dy = get_cell_basis(Y)
dx = get_cell_basis(X)
dv, dq = dy
du, dp = dx

cellmat = integrate(dv*du,trian,quad)
cellvec = integrate(dv*2,trian,quad)
cellids = get_cell_id(trian)
cellmatvec = pair_arrays(cellmat,cellvec)

dv_Γ, dq_Γ = restrict(get_cell_basis(Y), trian_Γ)
du_Γ, dp_Γ = restrict(get_cell_basis(X), trian_Γ)
cellmat_Γ = integrate(  jump(dv_Γ)*dp_Γ.⁺ + mean(dq_Γ)*jump(dp_Γ), trian_Γ,quad_Γ)
cellvec_Γ = integrate(  jump(dv_Γ) + mean(dq_Γ), trian_Γ,quad_Γ)
cellmatvec_Γ = pair_arrays(cellmat_Γ,cellvec_Γ)
cellids_Γ = get_cell_id(trian_Γ)

assem = SparseMatrixAssembler(SparseMatrixCSR{0},X,Y)

matvecdata = ([cellmatvec,cellmatvec_Γ],[cellids,cellids_Γ],[cellids,cellids_Γ])
matdata = ([cellmat,cellmat_Γ],[cellids,cellids_Γ],[cellids,cellids_Γ])
vecdata = ([cellvec,cellvec_Γ],[cellids,cellids_Γ])
data = (matvecdata,matdata,vecdata)

test_assembler(assem,matdata,vecdata,data)

A = assemble_matrix(assem,matdata)

A = allocate_matrix(assem,matdata)

assem = SparseMatrixAssembler(X,Y)
test_assembler(assem,matdata,vecdata,data)

end # module


module TestOTSIN
    using Test
    using EcologicalNetworks

    M = [1 2 1 1 1;
         0 5 0 0 0;
         3 1 1 0 0;
         1 1 1 2 1]

    a = [0.25, 0.25, 0.25, 0.25]
    b = [0.3, 0.1, 0.1, 0.25, 0.25]

    Q = optimaltransportation(M; a, b)
    @test Q isa AbstractMatrix
    @test sum(Q) ≈ 1.0
    @test sum(Q, dims=2)[:] ≈ a
    @test sum(Q, dims=1)[:] ≈ b
    @test optimaltransportation(M; a, b, λ=0.0) ≈ a * b'

    Q = optimaltransportation(M; a)
    @test sum(Q) ≈ 1.0
    @test sum(Q, dims=2)[:] ≈ a

    Q = optimaltransportation(M; b)
    @test sum(Q) ≈ 1.0
    @test sum(Q, dims=1)[:] ≈ b

    Q = optimaltransportation(M)
    @test sum(Q) ≈ 1.0

    N = BipartiteNetwork(M.>0)
    Q = optimaltransportation(N; a, b)

    @test Q isa BipartiteQuantitativeNetwork

end


module F1Method

#======================================================================
This package (the code below) implements the F-1 method as described
in the work of Pasquier et al. (2019). The numbers in parentheses
refer to the Equation numbers in the above manuscript. A bibtex
citation file is available in the GitHub repository.
======================================================================#

using LinearAlgebra, ForwardDiff, SciMLBase

"""
    Mem

Memory cache to store reusable objects.
Contains
- `s`   the steady-state solution, 𝒔(𝒑)
- `A`   the factors of 𝐀 = ∇ₓ𝑭(𝒔,𝒑)
- `∇s`  the derivative ∇𝒔(𝒑)
- `∇ₓf` the derivative ∇ₓ𝑓(𝒔,𝒑)
- `p`   the parameters 𝒑
The `Mem`-type object should be initialized with `initialize_mem`.
"""
mutable struct Mem{Ts, TA, T∇s, T∇ₓf, Tp}
    s::Ts     # 𝒔(𝒑)
    A::TA     # factors of 𝐀 = ∇ₓ𝑭(𝒔,𝒑)
    ∇s::T∇s   # ∇𝒔(𝒑)
    ∇ₓf::T∇ₓf # ∇ₓ𝑓(𝒔,𝒑)
    p::Tp     # 𝒑 that matches all but 𝒔(𝒑)
    psol::Tp  # 𝒑 that matches 𝒔(𝒑)
end

function update_mem!(f, F, ∇ₓf, mem, p, alg; options...)
    if p ≠ mem.p                      # only update mem if 𝒑 has changed
        update_solution!(F, mem, p, alg; options...)
        ∇ₚF = ForwardDiff.jacobian(p -> F(mem.s, p), p)
        mem.A = factorize(F.jac(mem.s, p)) # update factors of ∇ₓ𝑭(𝒔,𝒑)
        mem.∇s .= mem.A \ -∇ₚF           # update ∇𝒔
        mem.∇ₓf .= ∇ₓf(mem.s, p)         # update ∇ₓ𝑓(𝒔,𝒑)
        mem.p .= p                  # update 𝒑 for the variables above
    end
end

function update_solution!(F, mem, p, alg; options...)
    if p ≠ mem.psol
        prob = SteadyStateProblem(F, mem.s, p) # define problem
        mem.s .= solve(prob, alg; options...).u      # update 𝒔
        mem.psol .= p                          # update 𝒑 for 𝒔
    end
end

"""
    objective(f, F, mem, p, alg; options...)

Returns `f(x,p)` such that `F(x,p)=0` using the F-1 method.

Specifically, `objective(f, F, mem, p, alg; options...)`
evaluates the objective function defined by `f̂(p) = f(s(p),p)`, where
`s(p)`, which is the steady-state solution (i.e., such that `F(s(p),p)=0`)
is computed by the iterative Newton-type solver `alg`.
The memory cache `mem` must be supplied.
"""
function objective(f, F, mem, p, alg; options...)
    update_solution!(F, mem, p, alg; options...)
    return f(mem.s, p)
end

"""
    gradient(f, F, ∇ₓf, mem, p, alg; options...)

Returns the gradient of the `objective` function using the F-1 method.
"""
function gradient(f, F, ∇ₓf, mem, p, alg; options...)
    update_mem!(f, F, ∇ₓf, mem, p, alg; options...)
    s, ∇s, m = mem.s, mem.∇s, length(p)
    ∇ₚf = ForwardDiff.jacobian(p -> [f(s,p)], p)
    return mem.∇ₓf * ∇s + ∇ₚf
end

"""
    hessian(f, F, ∇ₓf, mem, p, alg; options...)

Returns the Hessian of the `objective` function using the F-1 method.
"""
function hessian(f, F, ∇ₓf, mem, p, alg; options...)
    update_mem!(f, F, ∇ₓf, mem, p, alg; options...)
    s, A, ∇s, m = mem.s, mem.A, mem.∇s, length(p)
    A⁻ᵀ∇ₓfᵀ = vec(A' \ mem.∇ₓf') # independent of (𝑗,𝑘)
    H(λ) = f(s+∇s*λ, p+λ) - F(s+∇s*λ, p+λ)' * A⁻ᵀ∇ₓfᵀ
    return ForwardDiff.hessian(H, zeros(m))
end

"""
    initialize_mem(x, p)

Initializes the memory cache for the F-1 method.
"""
function initialize_mem(F, ∇ₓf, x, p, alg; options...)
    x = copy(x)
    p = copy(p)
    psol = copy(p)
    prob = SteadyStateProblem(F, x, p)
    s = solve(prob, alg; options...).u
    A = factorize(F.jac(s, p))
    ∇ₚF = ForwardDiff.jacobian(p -> F(s, p), p)
    return Mem(s, A, A \ -∇ₚF, ∇ₓf(s, p), p, psol)
end

end


using Distributions
using DynamicHMC
using LogDensityProblems
using TransformVariables
using Random

struct Schaefer{F}
    r::F
    K::F
end
(s::Schaefer)(p) = p + s.r * p * (1 - p / s.K)

struct SchaeferProblem{T}
    p0::Vector{T}
    p1::Vector{T}
    σ::T
end
function (sp::SchaeferProblem)(θ)
    ll = zero(θ.r)
    ll += logpdf(Normal(100.0, 50.0), θ.K)
    ll += logpdf(LogNormal(log(0.1) - 0.1^2 / 2, 0.1), θ.r)
    schaefer = Schaefer(θ.r, θ.K)
    errdist = LogNormal(-sp.σ^2 / 2, sp.σ)
    for idx in eachindex(sp.p0)
        res = schaefer(sp.p0[idx]) / sp.p1[idx]
        ll += logpdf(errdist, res)
    end
    ll
end

schaefer = Schaefer(0.1, 100.0)
σ = 0.1
errdist = LogNormal(σ^2 / 2, σ)
p0 = 100rand(30)
p1 = zeros(100)
for idx in eachindex(p0)
    p1[idx] = schaefer(p0[idx])
end

sprob = SchaeferProblem(p0, p1, σ)

sprob((r = 0.1, K = 100.0))
t = as((r = asℝ₊, K = asℝ₊))
ℓ = TransformedLogDensity(t, sprob)

∇ℓ = ADgradient(:ForwardDiff, ℓ)

## Fails for non-finiteness somewhere?
post = mcmc_with_warmup(Random.GLOBAL_RNG, ∇ℓ, 100_000;
                        initialization = (q = [0.1, 100.0],),
                        reporter = NoProgressReport())

histogram2d(first.(post.chain), last.(post.chain), bins = 100)

using Optim

optimize(θ -> -sprob((r = θ[1], K = θ[2])), [0.1, 100.0])

struct Rosenbrock{T}
    a::T
    b::T
end
(rosen::Rosenbrock)(θ) = -(rosen.a - θ.x)^2 - rosen.b * (θ.y - θ.x^2)^2

rosen = Rosenbrock(1.0, 10.0)
rosen_t = as((x = asℝ, y = asℝ))
rosen_ll = TransformedLogDensity(rosen_t, rosen)
∇rosen_ll = ADgradient(:ForwardDiff, rosen_ll)
rosen_post = mcmc_with_warmup(Random.GLOBAL_RNG, ∇rosen_ll, 10_000;
                              reporter = NoProgressReport())

histogram2d(first.(rosen_post.chain), last.(rosen_post.chain), bins = 100)

RNG = Random.GLOBAL_RNG

results = DynamicHMC.mcmc_keep_warmup(RNG, ∇rosen_ll, 0; reporter = NoProgressReport())
steps = DynamicHMC.mcmc_steps(results.sampling_logdensity, results.final_warmup_state)
Q = results.final_warmup_state.Q

# a single update step
Q, tree_stats = DynamicHMC.mcmc_next_step(steps, Q)

grid = range(-1.5, 3.5; length = 101)

# extract the position
Q.q

DynamicHMC.Diagnostics.leapfrog_trajectory(∇rosen_ll, Q.q, -5:5)


module Quadratic_Loss
    using LoopVectorization

    function func(output_matrix::Array{Float32}, sample_matrix::Array{Float32})
        loss_matrix = zeros(Float32, size(output_matrix))
        @avxt for i in axes(loss_matrix, 1), j in axes(loss_matrix, 2)
            loss_matrix[i,j] = (output_matrix[i,j]-sample_matrix[i,j])^2
        end
        return loss_matrix
    end

    function prop!(δ::Array{Float32}, output_matrix::Array{Float32}, sample_matrix::Array{Float32})
        @avxt for i in eachindex(i)
            δ[i] = 2*(output_matrix[i]-sample_matrix[i])
        end
    end
end


module BesselFunctions

using AndExport
# using GSL
# using SpecialFunctions

# @xport sphj(n, x) = GSL.sf_bessel_jl(n, x)
# @xport sphn(n, x) = GSL.sf_bessel_yl(n, x)
# @xport sphh₊(n, x) = sphj(n, x) + im*sphn(n, x)
# @xport sphh₋(n, x) = sphj(n, x) - im*sphn(n, x)
# # @xport sphj(n, x) = sqrt(π/2x) * besselj(n+0.5, x)
# # @xport sphn(n, x) = sqrt(π/2x) * bessely(n+0.5, x)
# # @xport sphh₊(n, x) = sqrt(π/2x) * besselh1(n+0.5, x)
# # @xport sphh₋(n, x) = sqrt(π/2x) * besselh2(n+0.5, x)

# * Changed to use the complex versions in SpecialFunctions
using SpecialFunctions

@xport sphj(n,x) = sqrt(π/2x) * besselj(n+1/2, x)
@xport sphn(n,x) = sqrt(π/2x) * bessely(n+1/2, x)
@xport sphh₊(n, x) = sqrt(π/2x) * hankelh1(n+1/2, x)
@xport sphh₋(n, x) = sqrt(π/2x) * hankelh2(n+1/2, x)

#####

@xport dsphjdx(n, x) = n/x * sphj(n, x) - sphj(n+1,x)
@xport dsphndx(n, x) = n/x * sphn(n, x) - sphn(n+1,x)
@xport dsphh₊dx(n, x) = n/x * sphh₊(n, x) - sphh₊(n+1,x)
@xport dsphh₋dx(n, x) = n/x * sphh₋(n, x) - sphh₋(n+1,x)

@xport dsphjdr(n,k,r) = dsphjdx(n,k*r) * k
@xport dsphndr(n,k,r) = dsphndx(n,k*r) * k

@xport dsphh₊dr(n,k,r) = dsphh₊dx(n,k*r) * k
@xport dsphh₋dr(n,k,r) = dsphh₋dx(n,k*r) * k

@xport ricj(n,x) = x*sphj(n,x)
# Note: this -ve makes ricn(0,x) = cos(x)
@xport ricn(n,x) = -x*sphn(n,x)
# Note: these negatives are to recover the proper hankel behaviour
# (which is to obtain i^(-n-1) exp(iz) and i^(n+1) exp(-iz)
@xport rich₊(n,x) = ricj(n,x) - im*ricn(n,x)
@xport rich₋(n,x) = ricj(n,x) + im*ricn(n,x)

@xport dricjdx(n,x) = sphj(n,x) + x*dsphjdx(n,x)
@xport dricndx(n,x) = -sphn(n,x) + -x*dsphndx(n,x)
@xport drich₊dx(n,x) = dricjdx(n,x) - im*dricndx(n,x)
@xport drich₋dx(n,x) = dricjdx(n,x) + im*dricndx(n,x)

@xport dricjdr(n,k,r) = k * dricjdx(n,k*r)
@xport dricndr(n,k,r) = k * dricndx(n,k*r)
@xport drich₊dr(n,k,r) = k * drich₊dx(n,k*r)
@xport drich₋dr(n,k,r) = k * drich₋dx(n,k*r)

end


using Distributions
using DifferentialEquations
using DataFrames
using CSV
using PyPlot

"""
	function mcmc(data,cols,p,Population,total::Int)

Metropolis-Hastings MCMC of the SIR ODE model assuming a Poisson likelihood
for the cumulative number of cases and deaths, and a noninformative prior
data = DataFrame of
cols = [active cases, deaths, recovered cases]
p[1:end-1] = array of parameters used in the model
p[end] = day of first infection relative to day 0 of data (Jan 22, 2020)
Population = population of region in data
total = number of MCMC steps
output = (loglikelihoods, parameters, maximum loglikelihood, max loglikelihood parameters)

"""
function mcmc(data::DataFrame,cols::Array,p::Array,Population::Float64,total::Int)

	chi = ll(data,cols,p,Population)
	accept = 0
	pt = copy(p)
	chiml = chi
	pml = copy(p)

	chiout = Array{Float64,1}(undef,total)
	pout = Array{Float64,2}(undef,total,length(p))

	for step in 1:total
		sampleparams!(pt,p)
		chit = ll(data,cols,pt,Population)
		if rand() < exp(chi - chit)
			chi = chit
			p = copy(pt)
			accept += 1
			if chi < chiml
				chiml = chi
				pml = copy(p)
			end
		end
		chiout[step] = chi
		pout[step,:] = p
	end
	return chiout, pout, accept/total, chiml, pml
end

# generate MCMC step guess from current parameter value r using logNormal distirbution for all
# parameters except start day with scale parameter manually adjusted so approximately a third of guesses are accepted
function sampleparams!(pt,p)
	for i in 1:length(p) - 1
		d = Distributions.LogNormal(log(p[i]),.002)
		pt[i] = rand(d)
	end
	pt[end] = p[end] + .05*randn()
end

# loglikelihood of data given parameters at r using a Poisson likelihood function
function ll(data,cols,p,Population)
	total = length(data[!,cols[1]])-1
	prediction = model(p,total,Population)
	chi = 0
	for set in 1:3
		for i in 1:total
			d = Distributions.Poisson(max(prediction[i,set],1.))
			chi -= loglikelihood(d,[data[i,cols[set]]])
		end
	end
	return chi
end

# produce arrays of cases and deaths
function model(pin,total,Population)
	u0 = [0.,0.,0.,1/Population,0.,0.]
	p = pin[1:end-1]
	if pin[end] < total - 1
		start = ceil(Int,pin[end])
		tspan = (pin[end],total-1)
		prob = ODEProblem(sirc!,u0,tspan,p)
		sol = solve(prob,saveat=collect(start:total-1))
		# align solutions to time 0 of data and reorder rows and cols
		prediction = zeros(total,3)
		for set in 1:3
			prediction[max(0,start)+1:end,set] .= sol[set,max(0,-start)+1:end]
		end
		return prediction*Population
	else
		return zeros(total,3)
	end
end

# SIRC ODE model using Differential Equations
function sirc!(du,u,p,t)
	C,D,R,Y,Z,U = u
	beta,sigmac,sigmar,gammad,gammar = p
	du[1] = dC = sigmac*Y - gammar*C - gammad*C
	du[2] = dD = gammad*C
	du[3] = dR = gammar*C
	du[4] = dY = beta*Y*(1-Z) - sigmac*Y - sigmar*Y
	du[5] = dZ = beta*Y*(1-Z)
	du[6] = dU = sigmar*Y
end