# In this tutorial, we will look at how to
# - Evaluate `CellFields` at arbitrary points
# - Interpolate finite element functions defined on different
# triangulations. We will consider examples for
#    - Lagrangian finite element spaces
#    - Raviart Thomas finite element spaces
#    - Vector-Valued Spaces
#    - Multifield finite element spaces

# ## Problem Statement
# Let $\mathcal{T}_1$ and $\mathcal{T}_2$ be two triangulations of a
# domain $\Omega$. Let $V_i$ be the finite element space defined on
# the triangulation $\mathcal{T}_i$ for $i=1,2$. Let $f_h \in V_1$. The
# interpolation problem is to find $g_h \in V_2$ such that
#
# ```math
# dof_k^{V_2}(g_h) = dof_k^{V_2}(f_h),\quad \forall k \in
# \{1,\dots,N_{dof}^{V_2}\}
# ```

# ## Setup
# For the purpose of this tutorial we require `Test`, `Gridap` along with the
# following submodules of `Gridap`

using Test
using Gridap
using Gridap.CellData
using Gridap.Visualization

# We now create a computational domain on the unit square $[0,1]^2$ consisting
# of 5 cells per direction

domain = (0,1,0,1)
partition = (5,5)
𝒯₁ = CartesianDiscreteModel(domain, partition)

# ## Background
# `Gridap` offers the feature to evaluate functions at arbitrary
# points in the domain. This will be shown in the next
# section. Interpolation then takes advantage of this feature to
# obtain the `FEFunction` in the new space from the old one by
# evaluating the appropriate degrees of freedom. Interpolation works
# using the composite type `Interpolable` to tell `Gridap` that the
# argument can be interpolated between triangulations.

# ## Interpolating between Lagrangian FE Spaces

# Let us define the infinite dimensional function

f(x) = x[1] + x[2]

# This function will be interpolated to the source `FESpace`
# $V_1$. The space can be built using

reffe₁ = ReferenceFE(lagrangian, Float64, 1)
V₁ = FESpace(𝒯₁, reffe₁)

# Finally to build the function $f_h$, we do

fₕ = interpolate_everywhere(f,V₁)

# To construct arbitrary points in the domain, we use `Random` package:

using Random
pt = Point(rand(2))
pts = [Point(rand(2)) for i in 1:3]

# The finite element function $f_h$ can be evaluated at arbitrary points (or
# array of points) by

fₕ(pt), fₕ.(pts)

# We can also check our results using

@test fₕ(pt) ≈ f(pt)
@test fₕ.(pts) ≈ f.(pts)

# Now let us define the new triangulation $\mathcal{T}_2$ of
# $\Omega$. We build the new triangulation using a partition of 20 cells per
# direction. The map can be passed as an argument to
# `CartesianDiscreteModel` to define the position of the vertices in
# the new mesh.

partition = (20,20)
𝒯₂ = CartesianDiscreteModel(domain,partition)

# As before, we define the new `FESpace` consisting of second order
# elements

reffe₂ = ReferenceFE(lagrangian, Float64, 2)
V₂ = FESpace(𝒯₂, reffe₂)

# Now we interpolate $f_h$ onto $V_2$ to obtain the new function
# $g_h$. The first step is to create the `Interpolable` version of
# $f_h$.

ifₕ = Interpolable(fₕ)

# Then to obtain $g_h$, we dispatch `ifₕ` and the new `FESpace` $V_2$
# to the `interpolate_everywhere` method of `Gridap`.

gₕ = interpolate_everywhere(ifₕ, V₂)

# We can also use
# `interpolate` if interpolating only on the free dofs or
# `interpolate_dirichlet` if interpolating the Dirichlet dofs of the
# `FESpace`.

ḡₕ = interpolate(ifₕ, V₂)

# The finite element function $\bar{g}_h$ is the same as $g_h$ in this
# example since all the dofs are free.

@test gₕ.cell_dof_values ==  ḡₕ.cell_dof_values

# Now we obtain a finite element function using `interpolate_dirichlet`

g̃ₕ = interpolate_dirichlet(ifₕ, V₂)

# Now $\tilde{g}_h$ will be equal to 0 since there are
# no Dirichlet nodes defined in the `FESpace`. We can check by running

g̃ₕ.cell_dof_values

# Like earlier we can check our results for `gₕ`:

@test fₕ(pt) ≈ gₕ(pt) ≈ f(pt)
@test fₕ.(pts) ≈ gₕ.(pts) ≈ f.(pts)

# We can visualize the results using Paraview

writevtk(get_triangulation(fₕ), "source", cellfields=["fₕ"=>fₕ])
writevtk(get_triangulation(gₕ), "target", cellfields=["gₕ"=>gₕ])

# which produces the following output

# ![Target](../assets/interpolation_fe/source_and_target.png)

# ## Interpolating between Raviart-Thomas FESpaces

# The procedure is identical to Lagrangian finite element spaces, as
# discussed in the previous section. The extra thing here is that
# functions in Raviart-Thomas spaces are vector-valued. The degrees of
# freedom of the RT spaces are fluxes of the function across the edge
# of the element. Refer to the
# [tutorial](https://gridap.github.io/Tutorials/dev/pages/t007_darcy/)
# on Darcy equation with RT for more information on the RT
# elements.

# Assuming a function

f(x) = VectorValue([x[1], x[2]])

# on the domain, we build the associated finite dimensional version
# $f_h \in V_1$.

reffe₁ = ReferenceFE(raviart_thomas, Float64, 1) # RT space of order 1
V₁ = FESpace(𝒯₁, reffe₁)
fₕ = interpolate_everywhere(f, V₁)

# As before, we can evaluate the RT function on any arbitrary point in
# the domain.

fₕ(pt), fₕ.(pts)

# Constructing the target RT space and building the `Interpolable`
# object,

reffe₂ = ReferenceFE(raviart_thomas, Float64, 1) # RT space of order 1
V₂ = FESpace(𝒯₂, reffe₂)
ifₕ = Interpolable(fₕ)

# we can construct the new `FEFunction` $g_h \in V_2$ from $f_h$

gₕ = interpolate_everywhere(ifₕ, V₂)

# Like earlier we can check our results

@test gₕ(pt) ≈ f(pt) ≈ fₕ(pt)

# ## Interpolating vector-valued functions

# We can also interpolate vector-valued functions across
# triangulations. First, we define a vector-valued function on a
# two-dimensional mesh.

f(x) = VectorValue([x[1], x[1]+x[2]])

# We then create a vector-valued reference element containing linear
# elements along with the source finite element space $V_1$.

reffe₁ = ReferenceFE(lagrangian, VectorValue{2,Float64}, 1)
V₁ = FESpace(𝒯₁, reffe₁)
fₕ = interpolate_everywhere(f, V₁)

# The target finite element space $V_2$ can be defined in a similar manner.

reffe₂ = ReferenceFE(lagrangian, VectorValue{2,Float64}, 2)
V₂ = FESpace(𝒯₂, reffe₂)

# The rest of the process is similar to the previous sections, i.e.,
# define the `Interpolable` version of $f_h$ and use
# `interpolate_everywhere` to find $g_h \in V₂$.

ifₕ = Interpolable(fₕ)
gₕ = interpolate_everywhere(ifₕ, V₂)

# We can then check the results

@test gₕ(pt) ≈ f(pt) ≈ fₕ(pt)

# ## Interpolating Multi-field Functions

# Similarly, it is possible to interpolate between multi-field finite element
# functions. First, we define the components $h_1(x), h_2(x)$ of a
# multi-field function $h(x)$ as follows.

h₁(x) = x[1]+x[2]
h₂(x) = x[1]

# Next we create a Lagrangian finite element space containing linear
# elements.

reffe₁ = ReferenceFE(lagrangian, Float64, 1)
V₁ = FESpace(𝒯₁, reffe₁)

# Next we create a `MultiFieldFESpace` $V_1 \times V_1$ and
# interpolate the function $h(x)$ to the source space $V_1$.

V₁xV₁ = MultiFieldFESpace([V₁,V₁])
fₕ = interpolate_everywhere([h₁, h₂], V₁xV₁)

# Similarly, the target multi-field finite element space is created
# using $\Omega_2$.

reffe₂ = ReferenceFE(lagrangian, Float64, 2)
V₂ = FESpace(𝒯₂, reffe₂)
V₂xV₂ = MultiFieldFESpace([V₂,V₂])

# Now, to find $g_h \in V_2 \times V_2$, we first extract the components of
# $f_h$ and obtain the `Interpolable` version of the components.

fₕ¹, fₕ² = fₕ
ifₕ¹ = Interpolable(fₕ¹)
ifₕ² = Interpolable(fₕ²)

# We can then use `interpolate_everywhere` on the `Interpolable`
# version of the components and obtain $g_h \in V_2 \times V_2$ as
# follows.

gₕ = interpolate_everywhere([ifₕ¹,ifₕ²], V₂xV₂)

# We can then check the results of the interpolation, component-wise.

gₕ¹, gₕ² = gₕ
@test fₕ¹(pt) ≈ gₕ¹(pt)
@test fₕ²(pt) ≈ gₕ²(pt)

# ## Acknowledgements

# Gridap contributors acknowledge support received from Google,
# Inc. through the Google Summer of Code 2021 project [A fast finite
# element interpolator in
# Gridap.jl](https://summerofcode.withgoogle.com/projects/#6175012823760896).


@testset "Dijkstra" begin
    g4 = PathDiGraph(5)
    d1 = float([ 0 1 2 3 4; 5 0 6 7 8; 9 10 0 11 12; 13 14 15 0 16; 17 18 19 20 0])
    d2 = sparse(float([ 0 1 2 3 4; 5 0 6 7 8; 9 10 0 11 12; 13 14 15 0 16; 17 18 19 20 0]))

    for g in testdigraphs(g4)
      y = @inferred(dijkstra_shortest_paths(g, 2, d1))
      z = @inferred(dijkstra_shortest_paths(g, 2, d2))

      @test y.parents == z.parents == [0, 0, 2, 3, 4]
      @test y.dists == z.dists == [Inf, 0, 6, 17, 33]

      y = @inferred(dijkstra_shortest_paths(g, 2, d1; allpaths=true))
      z = @inferred(dijkstra_shortest_paths(g, 2, d2; allpaths=true))
      @test z.predecessors[3] == y.predecessors[3] == [2]

      @test @inferred(enumerate_paths(z)) == enumerate_paths(y)
      @test @inferred(enumerate_paths(z))[4] ==
          enumerate_paths(z,4) ==
          enumerate_paths(y,4) == [2,3,4]
    end

    gx = PathGraph(5)
    add_edge!(gx,2,4)
    d = ones(Int, 5,5)
    d[2,3] = 100
    for g in testgraphs(gx)
      z = @inferred(dijkstra_shortest_paths(g,1,d))
      @test z.dists == [0, 1, 3, 2, 3]
      @test z.parents == [0, 1, 4, 2, 4]
    end

    # small function to reconstruct the shortest path; I copied it from somewhere, can't find the original source to give the credits
    # @Beatzekatze on github
    spath(target, dijkstraStruct, source) = target == source ? target : [spath(dijkstraStruct.parents[target], dijkstraStruct, source) target]
    function spaths(ds, targets, source)
        shortest_paths = []
        for i in targets
            push!(shortest_paths,spath(i,ds,source))
        end
        return shortest_paths
    end

    G = LightGraphs.Graph()
    add_vertices!(G,4)
    add_edge!(G,2,1)
    add_edge!(G,2,3)
    add_edge!(G,1,4)
    add_edge!(G,3,4)
    add_edge!(G,2,2)
    w = [0. 3. 0. 1.;
        3. 0. 2. 0.;
        0. 2. 0. 3.;
        1. 0. 3. 0.]

    for g in testgraphs(G)
      ds = @inferred(dijkstra_shortest_paths(g,2,w))
      # this loop reconstructs the shortest path for vertices 1, 3 and 4
      @test spaths(ds, [1,3,4], 2) == Array[[2 1],
                                            [2 3],
                                            [2 1 4]]

    # here a selflink at source is introduced; it should not change the shortest paths
      w[2,2] = 10.0
      ds = @inferred(dijkstra_shortest_paths(g,2,w))
      shortest_paths = []
      # this loop reconstructs the shortest path for vertices 1, 3 and 4
      @test spaths(ds, [1,3,4], 2) == Array[[2 1],
                                            [2 3],
                                            [2 1 4]]
    end
end


# fashion mnist mlp
#
# reference: https://github.com/FluxML/model-zoo/blob/master/mnist/mlp.jl

using Flux
using Flux: onehotbatch, argmax, crossentropy, throttle, @epochs
using BSON: @save, @load
using Base.Iterators: repeated

using MLDatasets # FashionMNIST
using ColorTypes: N0f8, Gray

const Img = Matrix{Gray{N0f8}}

function prepare_train()
    # load full training set
    train_x, train_y = FashionMNIST.traindata() # 60_000

    trainrange = 1:6_000 # 1:60_000
    imgs = Img.([train_x[:,:,i] for i in trainrange])
    # Stack images into one large batch
    X = hcat(float.(reshape.(imgs, :))...) |> gpu
    # One-hot-encode the labels
    Y = onehotbatch(train_y[trainrange], 0:9) |> gpu
    X, Y
end

function prepare_test()
    # load full test set
    test_x,  test_y  = FashionMNIST.testdata() # 10_000

    testrange = 1:1_000 # 1:10_000
    test_imgs = Img.([test_x[:,:,i] for i in testrange])
    tX = hcat(float.(reshape.(test_imgs, :))...) |> gpu
    tY = onehotbatch(test_y[testrange], 0:9) |> gpu
    tX, tY
end

X, Y = prepare_train()
tX, tY = prepare_test()

m = Chain(
  Dense(28^2, 32, relu),
  Dense(32, 10),
  softmax) |> gpu

loss(x, y) = crossentropy(m(x), y)

accuracy(x, y) = mean(argmax(m(x)) .== argmax(y))

dataset = repeated((X, Y), 200)
evalcb = () -> @show(loss(X, Y))
opt = ADAM(params(m))

@epochs 5 Flux.train!(loss, dataset, opt, cb = throttle(evalcb, 2))

accuracy(X, Y)
# 0.983

accuracy(tX, tY)
# 0.83

# @save "fashion_mnist_mlp_model.bson" m
# weights = Tracker.data.(params(m)) ;
# @save "fashion_mnist_mlp_weights.bson" weights


using Test
using Mimi

include("../src/main_model.jl")
include("../src/compute_scc.jl")

m = getpage("RCP4.5 & SSP2", true, true, true)
run(m)
uds = [undiscount_scc(m, 2020), undiscount_scc(m, 2040), undiscount_scc(m, 2075)]
uds_compare = [0.8765, 0.5279, 0.2162] # From Excel PAGE-ICE
@test uds ≈ uds_compare rtol=1e-4

function checksccs(sccs, sccomp)
    sccobs = [quantile(sccs, [.05, .25, .5, .75, .95]); mean(sccs)]
    sccerror = abs.((sccobs - scccomp) ./ [10, 1, 1, 1, 10, 1])
    sccsd = std(sccs) / sqrt(1000)
    @test all(sccerror .< 3*sccsd) # average off by < 3x SD
end

m = getpage("RCP4.5 & SSP2", true, true, true)
sccs = compute_scc_mcs(m, 1000, year=2020)
checksccs(sccs, [37.031, 89.742, 161.868, 299.439, 847.020, 261.739])

# m = getpage("RCP8.5 & SSP5", true, true, true)
# sccs = compute_scc_mcs(m, 1000, year=2020)

m = getpage("RCP4.5 & SSP2", true, true, false)
sccs = compute_scc_mcs(m, 1000, year=2020)
checksccs(sccs, [53.903, 129.330, 217.337, 361.739, 886.998, 306.914])

# m = getpage("RCP8.5 & SSP5", true, true, false)
# sccs = compute_scc_mcs(m, 1000, year=2020)

end


# Computes samples of unitary invariant ensembles matrices of size nxn
#
# InvariantEnsemble(str,n) represents an n x n unitary invariant ensemble
# with distribution
#
#   exp(- Tr Q(M)) dM
#
# str specifies an ensemble with precomputed recurrence coefficients.
#  The currently include ensembles are
#
#       Quartic:            Q(M) = n M^4
#       Eight:              Q(M) = n M^8
#       HODecay:            Q(M) = n (M^4/20 - 4/15M^3 +M^2/5 + 8/5M)
#       CoshUnscaled:       Q(M) = cosh(M)
#       QuarticUnscaled:    Q(M) = M^4
#       EightUnscaled:      Q(M) = M^8
#
# References:
#    Olver, Rao & Trogdon 2014 arXiv:1404.007

module InvariantEnsembles
    using ApproxFun, RandomMatrices, Compat

export InvariantEnsemble

## Construct orthogonal polynomials as Funs from first moment and recurrence relationship
function orthonormalpolynomials(μ0,α::Vector,β::Vector,d)
    n=length(α)+1

    p=Array(Fun{Chebyshev,Float64},n)
    p[1] = Fun([μ0],d)
    p[2] = multiplybyx(p[1])./β[1] - p[1].*α[1]./β[1]
    for k = 3:n
       p[k] =  multiplybyx(p[k-1])./β[k-1] - p[k-1].*α[k-1]./β[k-1] - p[k-2].*β[k-2]./β[k-1]
    end
    p
end

function orthonormalpolynomials(μ0,α::Function,β::Function,d,n)
    orthonormalpolynomials(μ0,map(α,[1:n-1]),map(β,[1:n-1]),d)
end

function legendrepolynomials(n)
  orthonormalpolynomials(1./sqrt(2),k->0,k->sqrt(k^2./(4k^2-1)),[-1.,1.],n)
end
function scaledhermitepolynomials(n)
  orthonormalpolynomials(n^(1/4)/(1.0π)^(1/4),k->0,k->sqrt((k)/(2n)),[-3.,3.],n)
end

## Evaluate orthogonal polynomials at points x from first moment and recurrence relationship
function orthonormalpolynomialsvalues(μ0,α::Vector,β::Vector,x)
    n=length(α)+1
    m=length(x)

    p=Array(Float64,m,n)

    p[:,1] = μ0*ones(m)
    p[:,2] = x.*p[:,1]./β[1] - p[:,1].*α[1]./β[1]
    for k = 3:n
       p[:,k] =  x.*p[:,k-1]./β[k-1] - p[:,k-1].*α[k-1]./β[k-1] - p[:,k-2].*β[k-2]./β[k-1]
    end
    p
end

## Ensembles

type InvariantEnsemble{D<:Interval}
    basis::Array{Float64,2}         # the m x n array of the first n weighted OPs evaluated at m Chebyshev points
    domain::D                       # the sub domain  of the real line
end

InvariantEnsemble(basis,d::Vector)=InvariantEnsemble(basis,Interval(d))

## n, where the invariant ensemble is n x n
Base.size(ie::InvariantEnsemble)=size(ie.basis,2),size(ie.basis,2)
Base.size(ie::InvariantEnsemble,n)=size(ie.basis,2)

#Takes in list of OPs, constructs phis
# can be unstable for large n
function InvariantEnsemble{F<:Fun}(p::Array{F},V::Function,d,n::Integer)
    error("Reimplement with matrix")

#     wsq=IFun(x->exp(-n/2.*V(x)),d)
#
#     #We now extend the lengths
#     m=length(wsq)+2n
#     p=map(f->pad(f,m),p)
#
#     # We want to multiply, but in value space to
#     # take advantage of wsq being very small
#
#     wsqv=exp(-n/2.*V(points(wsq.domain,m)))
#     q=map(f->IFun(chebyshevtransform(wsqv.*values(f)),f.domain),p)
#
#     InvariantEnsemble(q)

end

##Adaptively construct invariant ensemble from weight, first moment and recurrence relationship
function adaptiveie(w,μ0,α,β,d)
  for logn = 4:20
    m=2^logn + 1
    pts=points(Interval(d),m)
    pv=orthonormalpolynomialsvalues(μ0,α,β,pts)

    wv=w(pts)

    cfs=chebyshevtransform(pv[:,end].^2.*wv)

    if maximum(abs(cfs[end-8:end])) < 200*eps()
        ## chop to minimize oversample

        m=length(chop!(cfs,200*eps()))

        pts=points(Interval(d),m)
        pv=orthonormalpolynomialsvalues(μ0,α,β,pts)
        wv=w(pts)

        return InvariantEnsemble(diagm(sqrt(wv))*pv,d)
    end
  end

  error("Did not converge")
end

#Reads in recurrence relationship and constructs OPs
function InvariantEnsemble(str::AbstractString,V::Function,d,n::Integer)
    file = joinpath(dirname(@__FILE__), "../data/CoefficientDatabase/" * str * "/" * string(n))
    μ0=readdlm(file * "norm.csv")[1]
    A=readdlm(file * "rc.csv",',')
    a=reshape((A[1,1:end-1]),size(A)[2]-1)
    b=reshape(A[2,2:end],size(A)[2]-1)

    adaptiveie(x->exp(-n.*V(x)),μ0,a,sqrt(b),d)
end

# For constructing InvariantEnsembles that do not scale with n
function InvariantEnsembleUnscaled(str::AbstractString,V::Function,d,n::Integer)
    file = joinpath(dirname(@__FILE__), "../data/CoefficientDatabaseUnscaled/" * str * "/")
    μ0=readdlm(file * "norm.csv")[1]
    A=readdlm(file * "rc.csv",',')
    a=reshape((A[1,1:end-1]),size(A)[2]-1)
    b=reshape(A[2,2:end],size(A)[2]-1)

    adaptiveie(x->exp(-V(x)),μ0,a[1:n-1],sqrt(b[1:n-1]),d)
end

#Decides whether to use built in recurrence or read it in
# Also contains ensemble data
function InvariantEnsemble(str::AbstractString,n::Integer)
    if(str == "GUE")
        InvariantEnsemble(str,x->x.^2,[-3.,3.],n)
    elseif(str == "Quartic")
        InvariantEnsemble(str,x->x.^4,[-3.,3.],n)
    elseif(str == "Eight")
        InvariantEnsemble(str,x->x.^8,[-3.,3.],n)
    elseif(str == "HODecay")
        InvariantEnsemble(str,x->x.^4/20 -4/15*x.^3 + x.^2/5+8/5*x,[-4.,6.],n)
    elseif(str == "Legendre")
        error("Legendre not implemented")
        #InvariantEnsembles(map(f->pad(f,2n),legendrepolynomials(n)))
    elseif(str == "CoshUnscaled")
        InvariantEnsembleUnscaled("Cosh",x->cosh(x),[-2acosh(1.0n),2acosh(1.0n)],n)
    elseif(str == "EightUnscaled")
        InvariantEnsembleUnscaled("Eight",x->x.^8,[-3.0n^(1/8),3.0n^(1/8)],n)
    elseif(str == "QuarticUnscaled")
        InvariantEnsembleUnscaled("Quartic",x->x.^4,[-3.0n^(1/4),3.0n^(1/4)],n)
    end
end

##  Construct the invariant ensemble kernel K_n(x,x)
iekernel(q::Array{Float64,2},d)=iekernel(q,d,plan_chebyshevtransform(q[:,1]))
function iekernel(q::Array{Float64,2},d,plan::Function)
    n=size(q)[1]
    m=size(q)[2]
    ret=zeros(n)
    for i = 1:m
        for k = 1:n
            ret[k] += q[k,i].*q[k,i]
        end
    end

  Fun(chebyshevtransform(ret,plan),d)
end

samplespectra(str::AbstractString,n::Integer,m::Integer)=samplespectra(InvariantEnsemble(str,n),m)

# Sample eigenvalues of invariant ensemble, m times
function RandomMatrices.eigvalrand(p::InvariantEnsemble,m::Integer)
    q = p.basis
    plan = plan_chebyshevtransform(q[:,1])
    pts=chebyshevpoints(size(q,1))

    hcat([samplespectra(p.basis,p.domain,plan,pts) for i=1:m]...)'
end

RandomMatrices.eigvalrand(p::InvariantEnsemble)=[eigvalrand(p,1)[1,:]...]

# Sample invariant ensemble
function Base.rand(p::InvariantEnsemble)
    Q=rand(Haar(2),size(p,1))
    Q*diagm(eigvalrand(p))*Q'
end

# Sample invariant ensemble, m times
function Base.rand(p::InvariantEnsemble,m::Integer)
    ei = eigvalrand(p,m)
    Array{Complex{Float64},2}[(    Q=rand(Haar(2),size(p,1));
    Q*diagm([ei[k,:]...])*Q') for k=1:size(ei,1)]
end

## samplespectra is back end for eigvalrand
samplespectra(q::Array{Float64,2},d)=samplespectra(q,d,plan_chebyshevtransform(q[:,1]),chebyshevpoints(size(q,1)))

function samplespectra(q::Array{Float64,2},d,plan::Function,pts)
    n = size(q,2)
    r=Array(Float64,n)

    for k=1:n-1
        m = n - k + 1
        r[k] = samplecdf(normalizedcumsum!(iekernel(q,d,plan).coefficients/m))
        f=Float64[bary(q[:,j],pts,r[k]) for j = 1:m]
        r[k] = fromcanonical(d,r[k])
        Q=null(f')
        q=q*Q
    end

    r[n] = sample(iekernel(q,d,plan))

    r
end

function iekernel{F<:Fun}(p::Array{F})
    ret = 0.*p[1]
    for i = 1:length(p)
        ret += fasttimes(p[i],p[i])
    end

    ret
end

function samplespectra{F<:Fun}(p::Array{F})
    n = length(p)
    r=sample(iekernel(p)/n)
    if n==1
        r
    else
        f=map(q->q[r],p)
        Q=null(f')'
        [r ;samplespectra(Q*p)]
    end
end

iekernel(p::InvariantEnsemble)=iekernel(p.basis,p.domain)
samplespectra(p::InvariantEnsemble)=samplespectra(p.basis,p.domain)

## Spectra database supports building up a database of precomputed samples

function spectradatabase(str,n::Integer,m::Colon)
  file = joinpath(dirname(@__FILE__), "../data/SpectraDatabase/" * str * "/" * string(n) * ".csv");

  if(filesize(file) == 0)
    []
  else
    readdlm(file,',')
  end
end

function spectradatabase(str,n::Integer,m::Integer)
    if  filesize(joinpath(dirname(@__FILE__), "../data/CoefficientDatabase")) == 0 ||
        filesize(joinpath(dirname(@__FILE__), "../data/CoefficientDatabaseUnscaled")) == 0
        error("No coefficient database found.  InvariantEnsembles package corrupted.")
    end

    dir = joinpath(dirname(@__FILE__), "../data/SpectraDatabase")
    if filesize(dir) == 0
        warn("Creating SpectraDatabase folder " * dir)

        mkdir(dir)
    end

    dir = joinpath(dirname(@__FILE__), "../data/SpectraDatabase/" * str)

    if filesize(dir) == 0
        warn("Creating folder " * dir)

        mkdir(dir)
    end

  file = dir * "/" * string(n) * ".csv"

  if filesize(file) == 0
    warn("No existing samples found.  Creating sample database.");

    ie = InvariantEnsemble(str,n)

    A=eigvalrand(ie,m+1)
    writedlm(file,A,',')
  else
    A = readdlm(file,',')

    if ( m > size(A)[1])
      warn(string(m) * " is greater than current samples " * string(size(A)[1]) * ".  Generating more samples");

      ie = InvariantEnsemble(str,n)

      An=eigvalrand(ie,m-size(A)[1]+1)

      A = vcat(A,An)

      writedlm(file,A,',')
    end
  end

    return A[1:m,:]
end

end #module


filepath = joinpath(Pkg.dir("Survival", "examples"), "rossi.csv")
rossi = readtable(filepath)
rossi[:event] = Event.(rossi[:week],rossi[:arrest] .== 0)

outcome = coxph( @formula(event ~ fin+age+race+wexp+mar+paro+prio) ,rossi ; tol = 1e-8)

filepath_coefs = joinpath(Pkg.dir("Survival", "test"), "expected_coefmat.jld")
expected_coefmat = JLD.load(filepath_coefs, "expected_coefmat")
@test outcome.coefmat.cols[1:3] ≈ expected_coefmat.cols[1:3] atol = 1e-6
println("Cox regression is fine")
println("$(outcome.coefmat) \n ≈ \n $expected_coefmat")


# Simple propositional logic

module PropositionalLogic

using ..LogicBase

export PropositionalVariable, BinaryConnective, assignments, ⊢

abstract Formula

immutable PropositionalVariable <: Formula
    var::Symbol
end

immutable Negate <: Formula
    f::Formula
end

immutable BinaryConnective{T} <: Formula
    x::Formula
    y::Formula
end

function Base.writemime(io::IO, ::MIME"text/plain", pv::PropositionalVariable)
    print(io, pv.var)
end

function Base.writemime(io::IO, ::MIME"text/plain", n::Negate)
    print(io, '¬')
    writemime(io, "text/plain", n.f)
end

function Base.writemime{T}(io::IO, ::MIME"text/plain", bc::BinaryConnective{T})
    print(io, '(')
    writemime(io, "text/plain", bc.x)
    print(io, ' ', T, ' ')
    writemime(io, "text/plain", bc.y)
    print(io, ')')
end

immutable Assignment
    table::Dict{PropositionalVariable, Bool}
end

assign(a::Assignment, v::PropositionalVariable) = a.table[v]
assign(a::Assignment, n::Negate) = ¬assign(a, n.f)
assign(a::Assignment, c::BinaryConnective{:∧}) =
    assign(a, c.x) && assign(a, c.y)
assign(a::Assignment, c::BinaryConnective{:∨}) =
    assign(a, c.x) || assign(a, c.y)
assign(a::Assignment, c::BinaryConnective{:⇒}) = assign(a, c.x) ⇒ assign(a, c.y)
assign(a::Assignment, c::BinaryConnective{:⇔}) = assign(a, c.x) ⇔ assign(a, c.y)

immutable AssignmentIterator
    over::Vector{PropositionalVariable}
end

Base.start(ai::AssignmentIterator) = UInt(1) << length(ai.over)
function Base.next(ai::AssignmentIterator, state::UInt)
    state -= 1
    len = length(ai.over)
    assignmentdata = [
        v => state & (UInt(1) << (len - i)) ≠ 0
        for (i, v) in enumerate(ai.over)]
    Assignment(assignmentdata), state
end
Base.done(ai::AssignmentIterator, state::UInt) = state == UInt(0)

vars(v::PropositionalVariable) = Set([v])
vars(n::Negate) = vars(n.f)
vars(bc::BinaryConnective) = vars(bc.x) ∪ vars(bc.y)

assignments(vs::Vector{PropositionalVariable}) = AssignmentIterator(vs)
assignments(vs::Set{PropositionalVariable}) = assignments(collect(vs))
assignments(f::Formula) = assignments(vars(f))
assignments(f::Formula, g::Formula) = assignments(vars(f) ∪ vars(g))
assignments(vs::Vector{Formula}) = assignments(foldl(∪, Set(), vs))

#= Tautology & Entailment =#
⊢(::Tuple{}, f::Formula) = all(a -> assign(a, f), assignments(f))
⊢(f::Formula, g::Formula) =
    all(a -> assign(a, f) ⇒ assign(a, g), assignments(f, g))
⊢(fs, g::Formula) = all(a -> ¬assign(a, g) ∧ ¬⋀(map(f -> assign(a, f), fs)))

#= Convenience Functions =#
LogicBase. ¬(f::Formula) = Negate(f)
LogicBase. ∧(f::Formula, g::Formula) = BinaryConnective{:∧}(f, g)
LogicBase. ∨(f::Formula, g::Formula) = BinaryConnective{:∨}(f, g)
LogicBase. ⊻(f::Formula, g::Formula) = BinaryConnective{:⊻}(f, g)
LogicBase. ⇒(f::Formula, g::Formula) = BinaryConnective{:⇒}(f, g)
LogicBase. ⇔(f::Formula, g::Formula) = BinaryConnective{:⇔}(f, g)

end


module Dec1Part1
function main(file_path::String)
    lines = readlines(file_path)
    lines = parse.(Int32, lines)
    n_lines = length(lines)
    for i in 1:n_lines
        for j in i + 1:n_lines
            val1 = lines[i]
            val2 = lines[j]
            if val1 + val2 == 2020
                return val1 * val2
            end
        end
    end
end
end


using Distributed
addprocs(2) # add 2 workers

@everywhere using Pkg
@everywhere Pkg.activate(joinpath(@__DIR__, "..", "examples"))

@everywhere using ProgressiveHedging
@everywhere import JuMP
@everywhere import Ipopt

@everywhere function two_stage_model(scenario_id::ScenarioID)

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

    scen = value(scenario_id)

    ref = JuMP.@variable(model, x >= 0.0)
    stage1 = [ref]

    ref = JuMP.@variable(model, y >= 0.0)
    stage2 = [ref]

    b_s = scen == 0 ? 11.0 : 4.0
	c_s = scen == 0 ? 0.5 : 10.0

    JuMP.@constraint(model, x + y == b_s)

    JuMP.@objective(model, Min, 1.0*x + c_s*y)

    return JuMPSubproblem(model,
                          scenario_id,
                          Dict(stid(1) => stage1,
                               stid(2) => stage2)
                          )
end

scen_tree = two_stage_tree(2)

(niter, abs_res, rel_res, obj, soln_df, phd) = solve(scen_tree,
                                                     two_stage_model,
                                                     ScalarPenaltyParameter(1.0)
                                                     )
@show niter
@show abs_res
@show rel_res
@show obj
@show soln_df


### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("FVheat2d");

useLog("FVheat2dlog")

# Configuration setup
domain(2)
solverType(FV)
timeStepper(RK4)

# Mesh
n = 20 # number of elements
mesh(QUADMESH, elsperdim=n)

# Variables and BCs
u = variable("u", location=CELL)
boundary(u, 1, FLUX, "0")

# Time interval and initial condition
T = 0.1;
timeInterval(T)
initial(u, "(sin(pi*x)*sin(2*pi*y))^4")

# The flux and source terms of the conservation equation
# F and S in the following equation:
# Dt(int(u dx)/A) = int(S dx) - int(F.n ds)
coefficient("D", 0.1) # Diffusion rate
flux(u, "-D*dot(grad(u),normal())")

printLatex(u)

#exportCode("fvheat2dcode") # uncomment to export generated code to a file
#importCode("fvheat2dcode") # uncomment to import code from a file

solve(u)

# output result to file
output_values(u, "fvheat2d", format="vtk");

@finalize()

##### Uncomment below to plot
# xy = Finch.fv_info.cellCenters

# using Plots
# pyplot();
# display(plot(xy[1,:], xy[2,:], u.values[:], st=:surface))


module HWconstrained

greet() = print("Hello World!")

	using JuMP, NLopt, DataFrames, Ipopt
	using LinearAlgebra

	export data, table_NLopt, table_JuMP, obj, constr, max_NLopt, max_JuMP

	function data(a=0.5)
		n = 3
		p = [1, 1, 1]
		e = [2, 0, 0]
		ns = 4
		nss = ns^2
		z1 = [0.72, 0.92, 1.12, 1.32]
		z2 = [0.86, 0.96, 1.06, 1.16]
		za = [[1, i, j] for i in z1 for j in z2]
		z = vcat(za'...)
		pi = repeat([1/nss], nss)
		return Dict("a"=>a,"ns"=>ns,"nss"=>nss,"e"=>e,"p"=>p,"z"=>z,"pi"=>pi, "za"=>za)
	end

	function max_JuMP(a=0.5)
		d = data(a)
		m = Model(with_optimizer(Ipopt.Optimizer))
		@variable(m, c >= 0)
		@variable(m, x[1:3])
		@NLconstraint(m, c + sum(d["p"][i] * (x[i] - d["e"][i]) for i in 1:3) == 0)
		@NLobjective(m, Max, -exp(-a * c) + sum(-exp(-a * sum(d["za"][j][i] * x[i] for i in 1:3)) * d["pi"][j] for j in 1:16))
		JuMP.optimize!(m)
		return Dict("obj"=>objective_value(m),"c"=>value(c),"omegas"=>[value(x[i]) for i in 1:length(x)])
	end

	function table_JuMP()
		d = DataFrame(a=[0.5;1.0;5.0],c = zeros(3),omega1=zeros(3),omega2=zeros(3),omega3=zeros(3),fval=zeros(3))
		for i in 1:nrow(d)
			xx = max_JuMP(d[i,:a])
			d[i,:c] = xx["c"]
			d[i,:omega1] = xx["omegas"][1]
			d[i,:omega2] = xx["omegas"][2]
			d[i,:omega3] = xx["omegas"][3]
			d[i,:fval] = xx["obj"]
		end
		return d
	end

	function obj(x::Vector,grad::Vector,data::Dict)
		a = data["a"]
		z = data["z"]
		u(x) = -exp(-a * x)
		up(x) = a * exp(-a * x)
		if length(grad) > 0
			grad[1] = up(x[1])
			grad[2:end] = sum(data["pi"] .* z .* up.(z * x[2:end]), dims=1)
		end
		return u(x[1]) + sum(data["pi"] .* u.(z * x[2:end]))
	end

	function constr(x::Vector,grad::Vector,data::Dict)
		if length(grad) > 0
			grad[1] = 1
			grad[2:end] = data["p"]
		end
		return x[1] + sum(data["p"] .* (x[2:end] - data["e"]))
	end

	function max_NLopt(a=0.5)
		d = data(a)
		opt = Opt(:LD_MMA, 4)
		lower_bounds!(opt, [0., -Inf, -Inf, -Inf])
		max_objective!(opt, (x, g)->obj(x, g, d))
		inequality_constraint!(opt, (x, g)->constr(x, g, d))
		ftol_rel!(opt, 1e-9)
		NLopt.optimize(opt, vcat(0, d["e"]))
	end

	function table_NLopt()
		d = DataFrame(a=[0.5;1.0;5.0],c = zeros(3),omega1=zeros(3),omega2=zeros(3),omega3=zeros(3),fval=zeros(3))
		for i in 1:nrow(d)
			xx = max_NLopt(d[i,:a])
			for j in 2:ncol(d)-1
				d[i,j] = xx[2][j-1]
			end
			d[i,end] = xx[1]
		end
		return d
	end

end # module


using Plots

@fastmath function SI4R_fs(N::Int64,bet::Float64,gamm::Float64)
    Psi = (N+1)*(N+2)*(N+3)*(N+4)/24;
    final_size = zeros(N+1,1);
    p_vec = zeros(Psi,1);
    p_vec[2]=1;
    @inbounds for Z5 in 0:N
        w::Int64 = Psi - (N-Z5+1)*(N-Z5+2)*(N-Z5+3)*(N-Z5+4)/24 + 1;
        @inbounds for Z4 in Z5:N
            a5 = 4*gamm*(Z4-Z5);
            @inbounds for Z3 in Z4:N
                a4 = 4*gamm*(Z3-Z4);
                @inbounds for Z2 in Z3:N;
                    a3 = 4*gamm*(Z2-Z3);
                    @inbounds for Z1 in Z2:N
                    a1 = bet*(N-Z1)*(Z1-Z5);
                    a2 = 4*gamm*(Z1-Z2);
                    tot = a1+a2+a3+a4+a5;
                    if Z1-Z5 == 0
                        final_size[Z5+1] = p_vec[w];
                    end
                    if a1 > 0
                       p_vec[w+1] = p_vec[w+1]+ p_vec[w]*a1/tot;
                    end
                    if a2 > 0
                        p_vec[w+N-Z2] = p_vec[w+N-Z2]+ p_vec[w]*a2/tot;
                    end
                    if a3 > 0
                        place3::Int64 = (N-Z3)*(N-Z3+1)/2;
                        p_vec[w+place3] = p_vec[w+place3]+ p_vec[w]*a3/tot;
                    end
                    if a4 > 0
                        place4::Int64 = (N-Z4)*(N-Z4+1)*(N-Z4+2)/6;
                        p_vec[w+place4] = p_vec[w+place4] + p_vec[w]*a4/tot;
                    end
                    if a5 > 0
                        p_vec[w] = p_vec[w]*a5/tot;
                    end
                    w = w + 1;
                end
            end
        end
    end
end
    return final_size
end

final_size = SI4R_fs(N,bet,gamm);

@time final_size = SI4R_fs(N,bet,gamm);

bar(0:N,final_size)


using WeibullCount

begin
    x = 2
    c = 0.36787944117144233
    λ = 20.08553692

    d = WeibullCountModel(λ, c)

    @test pdf(d, x) ≥ 0
    @test all(pdf(d, [x]) .≥ 0)

    @test cdf(d, x) ≥ 0
    @test all(cdf(d, [x]) .≥ 0)
end


# SteinMaternWeightedKernel
#
# These represent a family of kernels of the form
#
# k(x, y) = Q(x,y) k_0(x,y)
#
# where Q(x,y) is such that k is indeed a valid kernel and
# k_0 is the Matern radial kernel. That is,
#
# k_0(x, y) = (1 + gamma * r) * exp{-gamma * r}
#
# where r = ||x-y||. This is a norm for the order
# 2 Sobolev space H_2. This is the Matern covariance
# associated with v = 3/2.
#
# See http://www.gaussianprocess.org/gpml/chapters/RW4.pdf
# or http://arxiv.org/pdf/1204.6448.pdf for more details.

abstract type SteinMaternWeightedKernel <: SteinKernel end

### METHODS TO IMPLEMENT

# the function that evaluates Q(x,y)
function Q(ker::SteinMaternWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    error("Must implement the method Q")
end

# the function that implements gradx Q
function gradxQ(ker::SteinMaternWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    error("Must implement the method gradxQ")
end

# the function that implements gradx \dot grady Q
function gradxyQ(ker::SteinMaternWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    error("Must implement the method gradxyQ")
end

### INDUCED METHODS

function _maternk(ker::SteinMaternWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    r = norm(x - y)
    gamma = ker.gamma
    (1 + gamma * r) * exp(-gamma * r)
end

function _materngradxk(ker::SteinMaternWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    r = norm(x - y)
    gamma = ker.gamma
    (gamma^2 * exp(-gamma * r)) .* (y - x)
end

function _materngradyk(ker::SteinMaternWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    _materngradxk(ker, y, x)
end

function _materngradxyk(ker::SteinMaternWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    r = norm(x - y)
    gamma = ker.gamma
    p = length(x)
    gamma^2 * (p - gamma * r) * exp(-gamma * r)
end

function gradyQ(ker::SteinMaternWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    gradxQ(ker, y, x)
end

# defintion of the base kernel
function k(ker::SteinMaternWeightedKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    Q(ker, x, y) * _maternk(ker, x, y)
end

# implements \grad_x k(x,y)
function gradxk(ker::SteinMaternWeightedKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    _maternk(ker, x, y) * gradxQ(ker, x, y) + Q(ker, x, y) * _materngradxk(ker, x, y)
end

# implements <\grad_y, \grad_x k(x,y)>
function gradxyk(ker::SteinMaternWeightedKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    order0 = gradxyQ(ker, x, y) * _maternk(ker, x, y)
    order1 = dot(gradxQ(ker, x, y), _materngradyk(ker, x, y)) +
        dot(gradyQ(ker, x, y), _materngradxk(ker, x, y))
    order2 = Q(ker, x, y) * _materngradxyk(ker, x, y)

    order0 + order1 + order2
end


# # Heat Equation + Slab Tutorial

# In this tutorial, we demonstrate simple sequential coupling
# of two PDE models using the `ClimaCore.jl` backends.

# # Model 1
# Model 1 represents a simplified atmosphere (atm)
# and solves the [heat
# equation](https://en.wikipedia.org/wiki/Heat_equation)
# in a one-column domain:

# ``
# \frac{∂ T}{∂ t} + ∇ ⋅ (-μ ∇T) = 0
# ``

#    with top and bottom boundary conditions set to fixed-temperature (non-zero Dirichlet) and fixed-flux (non-zero Neumann) conditions, respectively:

# ``
#    T_{top}  = 280 K, \,\,\,\,\,\,\, \frac{∂ T_{bottom}}{∂ t} = - ∇ F_{sfc}
# ``

# where
#  - `t` is time
#  - `μ` is the thermal diffusivity
#  - `T` is the temperature
#  - `F_sfc` is the thermal boundary flux (see below for calculation)

# # Model 2
# Model 2 assumes a simplified soil (lnd) domain to be a slab, represented by an ODE for its surface temperature:

# ``
#    \frac{dT_{sfc}}{dt} = - (F_{accumulated} + G ) / h_{lnd}
# ``

#    where

# ``
#    F_{accumulated} = {F_{integrated}} / Δt_{coupler}
# ``

# where
#  - `Δt_coupler` is the duration of the coupling cycle
#  - `T_sfc` is the temperature
#  - `h_lnd` is the slab thickness
#  - `F_integrated` and `F_accumulated` thermal boundary fluxes, respectively (see below for calculation)
#  - `G` represents soil physics, which we assume to be 0 for now

# # Coupling and Flux Calculation

# We use this Model 1 (usually this is done by the model with the shortest timestep)
# to calculate and accumulate the downward surface fluxes, `F_sfc`:

# ``
#    F_{sfc} = - λ (T_{sfc} - T1)
# ``

# ``
#    d(F_{integrated})/dt  = F_{sfc}
# ``

# where
#  - `T1` is the atm temperature near the surface (here assumed equal to the first model level)
#  - `λ` a constant relaxation timescale

# Note that in a more realistic setup the above equations would be weighted by their domains' densities and thermal heat capacities, so that the thermal flux would have the units of W m$^{-2}$.
# Here we assume these are unity for both domains.

# Sequential coupling has the following steps:
# 1) pre-Model 1: supply Model 1 with `T_sfc` for the `F_sfc` calculation; reset `F_integrated` to zero
# 2) run Model 1: step forward for all Model 1 timesteps within one coupling cycle using `F_sfc` as the bottom boundary condition; accumulate `F_integrated` at each (sub-)step
# 3) post-Model 1: pass `F_integrated` into coupler and convert to `F_accumulated` for the correct units.
# 4) pre-Model 2: supply Model 2 with `F_accumulated`
# 5) run Model 2: step forward for all Model 2 timesteps within one coupling cycle;
# 6) post-Model 2: state variable, `T_sfc` of Model 2 into coupler.
# 7) repeat steps 1-6 for all coupling timesteps.

# # Implementation

# ## [Loading Packages](@id Loading-code-coupler)

# First, we'll load our pre-requisites:
#  - load CliMA packages under development - you may need to add unregistered packages in Pkg, e.g.:
## import Pkg; Pkg.add(url="https://github.com/CliMA/ClimaCore.jl",rev="main")

#  - load external packages:
import LinearAlgebra, UnPack
import ClimaCore: Fields, Domains, Topologies, Meshes, DataLayouts, Operators, Geometry, Spaces

using Base: show_supertypes
using OrdinaryDiffEq: ODEProblem, solve, SSPRK33

using Logging: global_logger
using TerminalLoggers: TerminalLogger

using RecursiveArrayTools

using OrdinaryDiffEq

using Statistics

#src ## Setup Logging Information
global_logger(TerminalLogger()) #src
const CI = !isnothing(get(ENV, "CI", nothing)) #src

# ## Define Parameters
#  - Global Constants
const FT = Float64;

#  - Experiment-specific Parameters
parameters = (
    ## atmos parameters
    zmin_atm = FT(0.0), # height of atm stack bottom [m]
    zmax_atm = FT(1.0), # height of atm stack top [m]
    n = 15,  # number of elements in atm stack
    μ = FT(0.0001), # diffusion coefficient [m^2 / s]
    T_top = FT(280.0), # fixed temperature at the top of the domain_atm [K]
    T_atm_ini = FT(280.0), # initial condition of at temperature (isothermal) [K]
    ## slab parameters
    h_lnd = FT(0.5), # depth of slab layer [m]
    T_lnd_ini = FT(260.0), # initial condition of at temperature (isothermal) [K]
    ## coupling parameters
    λ = FT(1e-5), # transfer coefficient
)

# ## Define Model Functions

# - Model 1 (atm) Equations
"""
    ∑tendencies_atm!(du, u, (parameters, T_sfc), t)

Heat diffusion equation
    dT/dt =  ∇ μ ∇ T
    where
        T  = 280 K              at z = zmax_atm
        dT/dt = - ∇ F_sfc       at z = zmin_atm

We also use this model to calculate and accumulate the downward surface fluxes, F_sfc:
    F_sfc = - λ * (T_sfc - T1)
    d(F_integrated)/dt  = F_sfc
    where
        F_integrated is reset to 0 at the beginning of each coupling cycle
        T1 = atm temperature near the surface (here assumed equal to the first model level)
"""
function ∑tendencies_atm!(du, u, (parameters, T_sfc), t)
    T = u.x[1] # u.x = vector of prognostic variables from DifferentialEquations
    F_sfc = calculate_flux(T_sfc[1], parent(T)[1], parameters)
    ## set BCs
    bcs_bottom = Operators.SetValue(Geometry.Cartesian3Vector(F_sfc)) # F_sfc is converted to a Cartesian vector in direction 3 (vertical)
    bcs_top = Operators.SetValue(FT(parameters.T_top))

    gradc2f = Operators.GradientC2F(top = bcs_top) # Dirichlet BC (center-to-face)
    gradf2c = Operators.DivergenceF2C(bottom = bcs_bottom) # Neumann BC (face-to-center)

    ## tendency calculations
    @. du.x[1] = gradf2c(parameters.μ * gradc2f(T)) # dT/dt
    du.x[2] .= -F_sfc[1] # d(F_integrated)/dt
end

# - Model 2 (lnd) Equations
"""
    ∑tendencies_lnd!(dT_sfc, T_sfc, (parameters, F_accumulated), t)

Slab layer equation
    lnd d(T_sfc)/dt = - (F_accumulated + G) / h_lnd
    where
        F_accumulated = F_integrated / Δt_coupler
"""
function ∑tendencies_lnd!(dT_sfc, T_sfc, (parameters, F_accumulated), t)
    G = 0.0 # place holder for soil dynamics
    @. dT_sfc = (-F_accumulated + G) / parameters.h_lnd
end

# - Surface Flux Calculation (coarse bulk formula)
calculate_flux(T_sfc, T1, parameters) = -parameters.λ * (T_sfc - T1);

# - Coupler Communication Functions
# These functions export / import / transform variables
# These functions are now just place holders for coupler transformations (e.g. regridding, masking, etc)
coupler_get_(x) = x;
coupler_put_(x) = x;

# ## Model Initialization
# - initialize atm model domain and grid
domain_atm = Domains.IntervalDomain(
    Geometry.ZPoint{FT}(parameters.zmin_atm),
    Geometry.ZPoint{FT}(parameters.zmax_atm);
    boundary_tags = (:bottom, :top),
);
mesh_atm = Meshes.IntervalMesh(domain_atm, nelems = parameters.n); # struct, allocates face boundaries to 5,6: atmos
center_space_atm = Spaces.CenterFiniteDifferenceSpace(mesh_atm); # collection of the above, discretises space into FD and provides coords

# - initialize prognostic variables, either as ClimaCore's Field objects or as Arrays
T_atm_0 = Fields.ones(FT, center_space_atm) .* parameters.T_atm_ini; # initiates a spatially uniform atm progostic var
T_lnd_0 = [parameters.T_lnd_ini]; # initiates lnd progostic var
ics = (; atm = T_atm_0, lnd = T_lnd_0)

# - specify timestepping information
stepping = (;
    Δt_min = 0.02,
    timerange = (0.0, 6.0),
    Δt_coupler = 1.0,
    odesolver = SSPRK33(),
    nsteps_atm = 8, # number of timesteps of atm per coupling cycle
    nsteps_lnd = 1, # number of timesteps of lnd per coupling cycle
);

# ## Define the sequential coupling loop
function coupler_solve!(stepping, ics, parameters)
    t = 0.0
    Δt_min = stepping.Δt_min
    Δt_coupler = stepping.Δt_coupler
    t_start = stepping.timerange[1]
    t_end = stepping.timerange[2]

    ## init coupler fields
    coupler_F_sfc = [0.0]
    coupler_T_lnd = copy(ics.lnd)

    ## atmos copies of coupler variables
    atm_T_lnd = copy(coupler_T_lnd)
    atm_F_sfc = copy(coupler_F_sfc)

    ## SETUP ATMOS
    ## put all prognostic variable arrays into a vector and ensure that solve can partition them
    T_atm = ics.atm
    Y_atm = ArrayPartition((T_atm, atm_F_sfc))
    prob_atm = ODEProblem(∑tendencies_atm!, Y_atm, (t_start, t_end), (parameters, atm_T_lnd))
    integ_atm = init(prob_atm, stepping.odesolver, dt = Δt_min, saveat = 10 * Δt_min)

    ## land copies of coupler variables
    T_lnd = ics.lnd
    lnd_F_sfc = copy(coupler_F_sfc)

    ## SETUP LAND
    prob_lnd = ODEProblem(∑tendencies_lnd!, T_lnd, (t_start, t_end), (parameters, lnd_F_sfc))
    integ_lnd = init(prob_lnd, stepping.odesolver, dt = Δt_min, saveat = 10 * Δt_min)

    ## coupler stepping
    for t in (t_start:Δt_coupler:t_end)

        ## STEP ATMOS
        ## pre_atmos
        integ_atm.p[2] .= coupler_get_(coupler_T_lnd) # integ_atm.p is the parameter vector of an ODEProblem from DifferentialEquations
        integ_atm.u.x[2] .= [0.0] # surface flux to be accumulated

        ## run atmos
        ## NOTE: use (t - integ_atm.t) here instead of Δt_coupler to avoid accumulating roundoff error in our timestepping.
        OrdinaryDiffEq.step!(integ_atm, t - integ_atm.t, true)

        ## post_atmos
        coupler_F_sfc .= coupler_put_(integ_atm.u.x[2]) / Δt_coupler

        ## STEP LAND
        ## pre_land
        lnd_F_sfc .= coupler_get_(coupler_F_sfc)

        ## run land
        OrdinaryDiffEq.step!(integ_lnd, t - integ_lnd.t, true)

        ## post land
        coupler_T_lnd .= coupler_put_(integ_lnd.u) # update T_sfc
    end

    return integ_atm, integ_lnd
end;

# ## Run the Coupler Model Simulation
integ_atm, integ_lnd = coupler_solve!(stepping, ics, parameters);
sol_atm, sol_lnd = integ_atm.sol, integ_lnd.sol;

# ## Postprocessing and Visualization

# Each integrator output (`sol_atm`, `sol_lnd`), contains the DifferentialEquations variable `.u` (the name is hard coded).
# If `ArrayPartition` was used for combining multiple prognostic variables, `u` will include an additional variable `x` (also hard coded)
# `parent()` accesses the `Field` values.
# So, for example, the structure of `u` from Model 1 is:

# `parent(sol_atm.u[<time-index>].x[<ArrayPartition-index>])[<z-index>,<variable-index>]`

ENV["GKSwstype"] = "nul"
import Plots
Plots.GRBackend()

show_plots = isdefined(Main, :SHOWPLOTS) ? SHOWPLOTS : true

path = joinpath(dirname(@__FILE__), "images/")
mkpath(path);

# - Vertical profile at start and end
t0_ = parent(sol_atm.u[1].x[1])[:, 1];
tend_ = parent(sol_atm.u[end].x[1])[:, 1];
z_centers = parent(Fields.coordinate_field(center_space_atm))[:, 1];
show_plots ?
Plots.png(
    Plots.plot(
        [t0_ tend_],
        z_centers,
        title = "model 1: atm",
        labels = ["t=0" "t=end"],
        xlabel = "T (K)",
        ylabel = "z (m)",
    ),
    joinpath(path, "tc1_f1.png"),
) : nothing
# ![](images/tc1_f1.png)

# - Conservation: absolute "energy" of both models with time
# convert to the same units (analogous to energy conservation, assuming that is both domains density=1 and thermal capacity=1)
lnd_sfc_u_t = [u[1] for u in sol_lnd.u] .* parameters.h_lnd;
atm_sum_u_t = [sum(parent(u.x[1])[:]) for u in sol_atm.u] .* (parameters.zmax_atm - parameters.zmin_atm) ./ parameters.n;
v1 = lnd_sfc_u_t .- lnd_sfc_u_t[1];
v2 = atm_sum_u_t .- atm_sum_u_t[1];
show_plots ?
Plots.png(
    Plots.plot(
        sol_lnd.t,
        [v1 v2 v1 + v2],
        labels = ["lnd" "atm" "tot"],
        xlabel = "time (s)",
        ylabel = "pseudo-energy (J / m2)",
    ),
    joinpath(path, "tc1_f2.png"),
) : nothing
# ![](images/tc1_f2.png)

# - Conservation: relative error with time
total = atm_sum_u_t + lnd_sfc_u_t;
rel_error = (total .- total[1]) / mean(total);
show_plots ?
Plots.png(
    Plots.plot(sol_lnd.t, rel_error, labels = ["tot"], xlabel = "time (s)", ylabel = "relative error"),
    joinpath(path, "tc1_f3.png"),
) : nothing
# ![](images/tc1_f3.png)

#src # - Animation
#src anim = Plots.@animate for u in sol_atm.u
#src     Plots.plot(u.x[1], xlim=(220,280))
#src end
#src Plots.mp4(anim, joinpath(path, "heat.mp4"), fps = 10)
#src
#src function linkfig(figpath, alt = "")
#src     # buildkite-agent upload figpath
#src     # link figure in logs if we are running on CI
#src     if get(ENV, "BUILDKITE", "") == "true"
#src         artifact_url = "artifact://$figpath"
#src         print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
#src     end
#src end
#src
#src dirname = "heat"
#src linkfig("output/$(dirname)/heat_end.png", "Heat End Simulation")


using Test
using Random
using LinearAlgebra
using SparseArrays
using JLD2, FileIO
using TraceEstimation

# Diagonal Approximation method as described by P. Fika works only for SPD matrix with low condition number (< 500)
@testset "Diagonal Approximation (P.Fika)" begin
    Random.seed!(1234323)
    @testset "Dense SPD Hermitian Matrices" begin
        @testset "a[i,j] = exp(-2 * abS(i - j)) (small size)" begin
            println("Executing Test 01: Dense SPD Small Size")
            A = rand(610, 610)
            for i in 1:610
                for j in 1:610
                    A[i, j] = exp(-2 * abs(i - j))
                end
            end
            obv = diagapp(A)
            acv = tr(inv(A))
            @test isapprox(obv, acv, rtol=10)
        end
        @testset "a[i,j] = exp(-2 * abS(i - j)) (large size)" begin
            println("Executing Test 02: Dense SPD Large Size")
            A = rand(5100, 5100)
            for i in 1:5100
                for j in 1:5100
                    A[i, j] = exp(-2 * abs(i - j))
                end
            end
            obv = diagapp(A)
            acv = tr(inv(A))
            @test isapprox(obv, acv, rtol=10)
        end
        @testset "Random generated SPD matrix (small size) (N=30)" begin
            println("Executing Test 03: Random SPD Small Size (N = 30)")
            A = rand(810, 810)
            A = A + A' + 30I
            while isposdef(A) == false
                A = rand(810, 810)
                A = A + A' + 30I
            end
            obv = diagapp(A)
            acv = tr(inv(A))
            @test isapprox(obv, acv, rtol=10)
        end
        @testset "Random generated SPD matrix (small size) (N=60)" begin
            println("Executing Test 04: Random SPD Small Size (N=60)")
            A = rand(810, 810)
            A = A + A' + 30I
            while isposdef(A) == false
                A = rand(810, 810)
                A = A + A' + 30I
            end
            obv = diagapp(A)
            acv = tr(inv(A))
            @test isapprox(obv, acv, rtol=10)
        end
        @testset "Random generated SPD matrix (large size) (N=30)" begin
            println("Executing Test 05: Random SPD Large Size")
            A = rand(5100, 5100)
            A = A + A' + 300I
            while isposdef(A) == false
                A = rand(5100, 5100)
                A = A + A' + 300I
            end
            obv = diagapp(A)
            acv = tr(inv(A))
            @test isapprox(obv, acv, rtol=10)
        end
    end
    @testset "Random generated SPD matrix (large size, large condition number)" begin
            println("Executing Test 05.5: Random SPD Large Size Larger Condition")
            A = Symmetric(rand(5005, 5005))
            A = A + 10I
            while isposdef(A) == false
                A = A + 10I
            end
            obv = diagapp(A)
            acv = tr(inv(A))
            @test isapprox(obv, acv, rtol=10)
        end
    @testset "Sparse SPD Hermitian Matrices" begin
        @testset "Random generated Sparse SPD (small size)" begin
            println("Executing Test 06: Sparse Random SPD Small Size")
            A = Symmetric(sprand(1000, 1000, 0.7))
            A = A+50*I
            while isposdef(A) == false
                A = Symmetric(sprand(1000, 1000, 0.7))
                A = A+50*I
            end
            obv = diagapp(A)
            A = Matrix(A)
            acv = tr(inv(A))
            @test isapprox(obv, acv, rtol=10)
        end
        #= Condition number too large (> 1000)
        @testset "TopOpt Test (Large Condition Number)" begin
            println("Executing Test 07: TopOpt Large Condition Number")
            s = (40, 10) # increase to increase the matrix size
            xmin = 0.9 # decrease to increase the condition number
            problem = HalfMBB(Val{:Linear}, s, (1.0, 1.0), 1.0, 0.3, 1.0)
            solver = FEASolver(Displacement, Direct, problem, xmin = xmin)
            n = length(solver.vars)
            solver.vars[rand(1:n, n÷2)] .= 0
            solver()
            K = solver.globalinfo.K
            obv = diagapp(K)
            M = Matrix(K)
            acv = tr(inv(M))
            @test isapprox(obv, acv, rtol=10)
        end
        =#
        @testset "TopOpt Test (Modified Condition Number)" begin
            println("Executing Test 08: TopOpt Modified Condition Number")
            println("Loading JLD2 file")
            #=
            s = (40, 10) # increase to increase the matrix size
            xmin = 0.9 # decrease to increase the condition number
            problem = HalfMBB(Val{:Linear}, s, (1.0, 1.0), 1.0, 0.3, 1.0)
            solver = FEASolver(Displacement, Direct, problem, xmin = xmin)
            n = length(solver.vars)
            solver.vars[rand(1:n, n÷2)] .= 0
            solver()
            K = solver.globalinfo.K
            K = K+1*I
            =#
            @load "topoptfile.jld2" K
            obv = diagapp(K)
            M = Matrix(K)
            acv = tr(inv(M))
            @test isapprox(obv, acv, rtol=10)
        end
    end
    #=
    @testset "CuArrays Test" begin
        @testset "Dense Random CuArray" begin
            println("Executing Test 09: CuArray Small Size")
            A = cu(rand(400,400))
            A = A+A'+40*I
            function mycurand(type::Type, size::Int64)
                cu(rand(type, size))
            end
            function mycurand(range::StepRange{<:Any, <:Any}, size::Int64)
                cu(rand(range, size))
            end
            obv = diagapp(A, randfunc=mycurand)
            M = Matrix(A)
            acv = tr(inv(M))
            @test isapprox(obv, acv, rtol=10)
        end
    end =#
end


function mass(u, q, w, v, x, t, mu, eta)
    m = [1.0, 1.0, 1.0, 1.0];
    return m;
end
function flux(u, q, w, v, x, t, mu, eta)
    gam = mu[1];
    gam1 = gam - 1.0;
    Re = mu[2];
    Pr = mu[3];
    Minf = mu[4];
    Re1 = 1/Re;
    M2 = Minf^2;
    c23 = 2.0/3.0;
    fc = 1/(gam1*M2*Re*Pr);
    r = u[1];
    ru = u[2];
    rv = u[3];
    rE = u[4];
    rx = q[1];
    rux = q[2];
    rvx = q[3];
    rEx = q[4];
    ry = q[5];
    ruy = q[6];
    rvy = q[7];
    rEy = q[8];
    r1 = 1/r;
    uv = ru*r1;
    vv = rv*r1;
    E = rE*r1;
    ke = 0.5*(uv*uv+vv*vv);
    p = gam1*(rE-r*ke);
    h = E+p*r1;
    fi = [ru, ru*uv+p, rv*uv, ru*h, rv, ru*vv, rv*vv+p, rv*h];
    ux = (rux - rx*uv)*r1;
    vx = (rvx - rx*vv)*r1;
    kex = uv*ux + vv*vx;
    px = gam1*(rEx - rx*ke - r*kex);
    Tx = gam*M2*(px*r - p*rx)*r1^2;
    uy = (ruy - ry*uv)*r1;
    vy = (rvy - ry*vv)*r1;
    key = uv*uy + vv*vy;
    py = gam1*(rEy - ry*ke - r*key);
    Ty = gam*M2*(py*r - p*ry)*r1^2;
    txx = Re1*c23*(2*ux - vy);
    txy = Re1*(uy + vx);
    tyy = Re1*c23*(2*vy - ux);
    fv = [0, txx, txy, uv*txx + vv*txy + fc*Tx, 0, txy, tyy, uv*txy + vv*tyy + fc*Ty];
    f = fi+fv;
    f = reshape(f,(4,2));

    return f;
end
function source(u, q, w, v, x, t, mu, eta)
    s = [0.0, 0.0, 0.0, 0.0];
    return s;
end
function ubou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    gam = mu[1];
    gam1 = gam - 1.0;
    r = u[1];
    ru = u[2];
    rv = u[3];
    rE = u[4];
    nx = n[1];
    ny = n[2];

    r1 = 1/r;
    uv = ru*r1;
    vv = rv*r1;
    E = rE*r1;
    p = gam1*(rE-r*0.5*(uv*uv+vv*vv));
    h = E+p*r1;
    a = sqrt(gam*p*r1);

    run = ru*nx + rv*ny;
    rut = -ru*ny + rv*nx;
    un = run/r;
    ut = rut/r;

    K = [ 1  1  0  1 ; un-a  un  0  un+a ; ut  ut  1  ut ; h - un*a  (1/2)*(un^2 + ut^2)  ut  h+un*a];
    Kinv = (gam1/(2*a^2))*[ h + (a/gam1)*(un-a)  -(un+a/gam1)  -ut  1; -2*h + (4/gam1)*a^2  2*un  2*ut  -2 ; -2*(ut*a^2)/gam1  0  2*(a^2)/gam1  0 ; h - a*(un+a)/gam1  -un+a/gam1  -ut  1];
    T = [ 1 0 0 0; 0 nx ny 0; 0 -ny nx 0; 0 0 0 1];
    Tinv = [ 1 0 0 0; 0 nx -ny 0; 0 ny nx 0; 0 0 0 1];
    Lambda = [ tanh(1e3*(un-a)) 0 0 0 ; 0 tanh(1e3*(un)) 0 0 ; 0 0 tanh(1e3*(un)) 0 ; 0 0 0 tanh(1e3*(un+a)) ];
    # E = simplify(K * Lambda * Kinv);
    # An = simplify(Tinv * E * T);
    E = (K * Lambda * Kinv);
    An = (Tinv * E * T);

    # freestream boundary condition
    uinf = mu[5:8]; # freestream flow
    uinf = uinf[:];
    u = u[:];          # state variables
    ui = 0.5*((u.+uinf) + An*(u.-uinf));  # Riemann solution
    uw = copy(u);
    uw[2] = 0; # zero velocity
    uw[3] = 0; # zero velocity

    ub = hcat(uw, ui); # wall and freestream boundary conditions
    return ub;
end
function fbou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    f = flux(u, q, w, v, x, t, mu, eta);
    fi = f[:,1]*n[1] + f[:,2]*n[2] + tau[1]*(u-uhat); # numerical flux at freestream boundary
    fw = copy(fi);
    fw[1] = 0.0;   # zero velocity
    fw[end] = 0.0; # adiabatic wall -> zero heat flux
    fb = hcat(fw[:], fi[:]); # wall and freestream boundary conditions

    return fb;
end
function initu(x, mu, eta)

    u0 = mu[5:8]; # freestream flow

    return u0;
end


@doc raw"""
LevStehnoexp{T} <: Kernel
The kernel for of the levensteindistance
`` k(x, x^\prime) = \exp(-levensthein(x_i,x_j))``

"""

struct LevStehnoexp{T} <: Kernel
    s::T
end

"""
levensthein(xᵢ,xⱼ)
dynamic calculations of the model the Levenstein distance between two vectors of modules

"""
function levenshtein(xᵢ, xⱼ)
    n, m = length(xᵢ), length(xⱼ)
    lev = zeros(Int, n+1, m+1)
    lev[:,1] = 0:n
    lev[1,:] = 0:m
    for j in 1:m
        for i in 1:n
            lev[i+1, j+1] = min(
                lev[i, j+1] + 1,
                lev[i+1, j] + 1,
                xᵢ[i]==xⱼ[j] ? lev[i, j] : lev[i, j] + 1
            )
        end
    end
    return last(lev)
end
ew(k::LevStehnoexp, x::AbstractVector{N} where N) = exp.(-k.s .*zeros(length(x)))
ew(k::LevStehnoexp, x::AbstractVector{N} where N, x′::AbstractVector{N} where N) = [exp(-k.s*levenshtein(xᵢ,xⱼ)) for (xᵢ,xⱼ) in zip(x, x′)]

pw(k::LevStehnoexp, x::AbstractVector{N} where N) = reshape([exp(-k.s*levenshtein(xᵢ,xⱼ)) for xᵢ in x for xⱼ in x] ,(length(x),length(x)))
pw(k::LevStehnoexp, x::AbstractVector{N} where N, x′::AbstractVector{N} where N) = reshape([exp(-k.s*levenshtein(xᵢ,xⱼ)) for xⱼ in x′ for xᵢ in x] ,(length(x),length(x′)))

####
#Cosine kernel
####

####
#make vectors first
####

function myword2vec(word,letter2index)
    word2vec = letter2index  |> t -> length(t) |> l-> zeros(l)
    for letter in word
        word2vec[letter2index[letter]] += 1
    end
    return word2vec
end

function cossim(xi,xj)
    letter2index = Dict( letter => index for (index,letter) in enumerate(Set([xi...,xj...])))
    xi_v = myword2vec(xi,letter2index)
    xj_v = myword2vec(xj,letter2index)
    return dot(xi_v,xj_v)/(norm(xi_v) * norm(xj_v))
end

struct CosStehno <: Kernel end

ew(::CosStehno, x::AbstractVector{N} where N) = ones(length(x))
ew(::CosStehno, x::AbstractVector{N} where N, x′::AbstractVector{N} where N) = [cossim(xᵢ,xⱼ) for (xᵢ,xⱼ) in zip(x, x′)]

pw(k:: CosStehno, x::AbstractVector{N} where N) = reshape([cossim(xᵢ,xⱼ) for xᵢ in x for xⱼ in x] ,(length(x),length(x)))
pw(k:: CosStehno, x::AbstractVector{N} where N, x′::AbstractVector{N} where N) = reshape([cossim(xᵢ,xⱼ) for xⱼ in x′ for xᵢ in x] ,(length(x),length(x′)))


function CRT_SIA(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T, tol)
  # boundary conditions
  BC = Array{Any}(N_comp)
  for i in 1:N_comp
      BC[i] = createBC(m)
      # Dirichlet BC on the left side
      BC[i].left.a[:] = 0.0
      BC[i].left.b[:] = 1.0
      BC[i].left.c[:] = C_bc[i]
  end

  # initial condition
  C_old = Array{Any}(N_comp)
  for i in 1:N_comp
      C_old[i] = createCellVariable(m, C_0[i], BC[i]) # initial value of concentration
  end

  # Advective term
  u_vec = createFaceVariable(m, u)
  M_adv = convectionUpwindTerm(u_vec) # matrix of coef for the advection term

  # Diffusive term
  M_diff = Array{Any}(N_comp)
  for i in 1:N_comp
    D_face = createFaceVariable(m, D[i]*ϕ)
    M_diff[i] = diffusionTerm(D_face)      # matrix of coef for the diffusion term
  end

  # Boundary terms
  M_bc = Array{Any}(N_comp)
  RHS_bc  = Array{Any}(N_comp)
  for i in 1:N_comp
      M_bc[i], RHS_bc[i] = boundaryConditionTerm(BC[i]) # matrix and RHS for the BC
  end

  # Constructing stationary part of each system matrix
  M_stat = Array{Any}(N_comp)
  for i in 1:N_comp
    M_stat[i] = M_adv - M_diff[i] + M_bc[i]
  end

  # Preparing transient terms
  M_t = Array{Any}(N_comp)
  RHS_t = Array{Any}(N_comp)

  maxit = 5000

  # Time-stepping
  for t in dt:dt:T
      # transport equation
      for i in 1:N_comp
          M_t[i], RHS_t[i] = transientTerm(C_old[i], dt, ϕ)
          M = M_t[i]+M_stat[i]
          RHS = RHS_t[i]+RHS_bc[i]
          C_new = solveLinearPDE(m, M, RHS)
          C_old[i] = copyCell(C_new)
      end

      # Iterative loop
      iter = 0
      error = 1000
      eps0 = 0
      eps1 = 0

      while (error > tol && iter < maxit)
        # reaction equation
        eps1 = epsilon.(internalCells(C_old[1]), internalCells(C_old[2]),
                internalCells(C_old[3]), nu_A, nu_B, nu_C, K_eq)
        C_old[1] = createCellVariable(m, internalCells(C_old[1])-nu_A*eps1, BC[1])
        C_old[2] = createCellVariable(m, internalCells(C_old[2])-nu_B*eps1, BC[2])
        C_old[3] = createCellVariable(m, internalCells(C_old[3])+nu_C*eps1, BC[3])

        error = sum(abs(eps1-eps0))
        eps0 = eps1
        iter += 1
    end

    #print(iter, ", ", t, ", ", error, "\n")
    #print(sum(internalCells(C_old[3])-K_eq*internalCells(C_old[1]).*internalCells(C_old[2])), "\n") #mass balance
  end

  return C_old
end


push!(LOAD_PATH, string(pwd(), "/src/"))
using WRIdual
using LinearAlgebra, SparseArrays
using PyPlot
using JLD2

## Set slowness model and discretization step

n = (201, 201)
h = (10e0, 10e0)
val = 0.5e0
x = reshape((collect(1:n[1]).-1f0).*h[1], :, 1)
z = reshape((collect(1:n[2]).-1f0).*h[2], 1, :)
v = 1000e0.*(2e0.-val*exp.(-((x./1000e0.-1e0)/0.25e0).^2e0.-((z./1000e0.-1e0)/0.5e0).^2e0))
m = 1e0./v.^2e0

## Create data

# Set frequencies
freqs = [6e0]

# Set source/receiver geometry
ix_src = 3:4:199
Ns = length(ix_src)
iz_src = 2
x_src = (Float64.(ix_src).-1e0)*h[1]
z_src = (Float64.(iz_src).-1e0)*h[2]
ix_rcv = 3:199
iz_rcv = 200
x_rcv = (Float64.(ix_rcv).-1e0)*h[1]
z_rcv = (Float64.(iz_rcv).-1e0)*h[2]

# Setting source wavelet
src_enc = C.(Matrix(sparse(I, length(x_src), length(x_src))))

# Absorbing boundary condition padding
nabc = (33, 33, 33, 33)

# Solving PDE for each frequency
dat = gendata(m, h, freqs, x_rcv, z_rcv, x_src, z_src, src_enc; nabc = nabc, verbose = true)

## Save data & related parameters

@save "./data/Gausslens/Gausslens_data.jld" h m freqs x_rcv z_rcv x_src z_src src_enc nabc dat


import Pkg
Pkg.instantiate() # downloads all dependencies for the current project

using DataFrames, CSV, StatsBase

df = CSV.read("../externals/Core.Math.Data/data/Pejcic_318.csv", copycols = true)

println(corspearman(df.ATV, df.ATT))