using tele
using NCDatasets, Revise, Statistics, PyPlot
using CFTime, DSP, FFTW
using DataFrames, NaNMath
#analyzing data from https://psl.noaa.gov/data/gridded/data.noaa.ersst.v5.html

#FILE PATHS
data_dir = "/home/brynn/Documents/JP/860/tsa_data/"
ersst_p = data_dir * "sst.mnmean.nc"
# cesm_p2 = data_dir * "tos_Omon_CCSM4_past1000_r1i1p1_135001-185012.nc"
cesm_p2 = data_dir * "tos_Omon_CCSM4_historical_r1i1p1_185001-200512.nc"

#READ NETCDFS TO BOUNDING BOX
bbox = [71, 55, 290, 350]
(late, lone, timee, datae) = read_netcdf(ersst_p, "sst", bbox)
(latc, lonc, timec, datac) = read_netcdf(cesm_p2, "tos", bbox)

#FILTER CHARACTERISTICS
responsetype = Lowpass(0.1*10^(-7),fs = 3.80265176 * 10^(-7))
designmethod = Butterworth(4)

#Compute raw average time series over bounding box
datac = datac .- 273.15 #convert to celsius
tavge_raw = spatial_mean(datae)
tavgc_raw = spatial_mean(datac)

#ersst
dataep = zeros(length(timee), length(lone), length(late))
permutedims!(dataep, datae, [3,1,2])
filte = filt(digitalfilter(responsetype, designmethod), dataep)
permutedims!(datae, filte, [2,3,1])
anome = anomaly(datae, timee)
#filte holds filtered anomalized data

#cesm
timec = reinterpret(typeof(timee[begin]), timec) #convert from DatetimeNoLeap to DateTime

datacp = zeros(length(timec), length(lonc), length(latc))
permutedims!(datacp, datac, [3,1,2])
filtc = filt(digitalfilter(responsetype, designmethod), datacp)
permutedims!(datac, filtc, [2,3,1])
anomc = anomaly(datac, timec)
#filtc holds filtered, anomalized data

#filtered, anomalized time series
tavge_anom = spatial_mean(anome)
tavgc_anom = spatial_mean(anomc)

#Make contour plots
savetype = ".pdf"
contplot(latc,lonc,temp_mean(datac), "images/datac"*savetype)
contplot(late,lone,temp_mean(datae), "images/datae"*savetype)
contplot(latc,lonc,temp_mean(anomc), "images/anomc"*savetype)
contplot(late,lone,temp_mean(anome), "images/anome"*savetype)

#Make time series plots
filtseriesplot(tavgc_raw, tavgc_anom, timec, tavge_raw, tavge_anom, timee, "", "images/timeseries"*savetype)
# filtseriesplot(tavgc_anom, tavgc_anom_filt, timec, tavge_anom, tavge_anom_filt, timee, "", "images/anom"*savetype)

#COMPUTE FFTs
ffte_raw = fft(tavge_raw)
ffte_anom = fft(tavge_anom)
ffte_freq = FFTW.fftfreq(length(timee), 3.80265176 * 10^(-7))

fftc_raw = fft(tavgc_raw)
fftc_anom = fft(tavgc_anom)
fftc_freq = FFTW.fftfreq(length(timec), 3.80265176 * 10^(-7))

#PLot FFT
fftplot(ffte_freq[:], [ffte_raw, ffte_anom], "images/ffte"*savetype, "ERSST")
fftplot(fftc_freq[:], [fftc_raw, fftc_anom], "images/fftc"*savetype, "CESM")

#Plot FFT
figure()
ax = subplot()
halfway = 1006
plot(ffte_freq[begin+1:halfway], abs.(ffte_anom)[begin+1:halfway], alpha = 0.5)
halfway = 936
plot(fftc_freq[begin+1:halfway], abs.(fftc_anom)[begin+1:halfway], alpha = 0.5)
text(7 * 10 ^-10, 400, "f = 5.5 × 10^(-10) Hz")
text(7 * 10 ^-10, 380, "T = 47 years")
ax.set_xlim(ffte_freq[begin+1], 2 * 0.1*10^(-7))
ax.set_ylim(0, 500)
xlabel("Frequency [Hz]")
ylabel("Intensity")
legend(["ERSST", "CESM"])
savefig("images/fftec"*savetype)
close()

#histogram
#fig2b - boxplotGPCC Total # of Precips Obs
figure(figsize = (6, 8))
boxdict = boxplot([tavge_raw, tavgc_raw], notch = true, labels =  ["ERSST", "CESM (1850 - 2005)"])
grid(true)
ylabel("SST [°C]")
savefig("images/boxes_ersst"*savetype)
close()

#Extreme values analysis
df = DataFrame(min = [], max = [], mean = [], median = [], std = [])
push!(df, character(tavge_raw))
push!(df, character(tavgc_raw))

quant = quantile(tavgc_raw, [0.25, 0.5, 0.75])
thresh = (quant[3] - quant[1])* 1 + quant[3]
ind = findall(>(thresh), tavgc_raw)

extreme1 = temp_mean(datac[:,:, ind])
extreme2 = temp_mean(anomc[:,:, ind])

#Pot extreme values
contplot(latc,lonc,extreme1, "images/ext_raw"*savetype)
contplot(latc,lonc,extreme2, "images/ext_anom"*savetype)


include("../examples/lorenz.jl")
using LinearAlgebra
using JLD
using SharedArrays
using Distributed
function sens(s,nSteps)
    d = 3
    n_runup = 2000
    u = rand(d)
    for i = 1:n_runup
    u = next(u, s)
    end
    u_trj = step(u, s, nSteps)
    x, y, z = view(u_trj,1,:), view(u_trj,2,:), view(u_trj,3,:)
    du_trj = dstep(u_trj, s)
    ddu_trj = d2step(u_trj,s)
    f_trj = flow(u_trj,s)

    vs = zeros(d)
    q = rand(d)
    q /= norm(q)

    pp = pert(u_trj, s)
    dpp = dpert(u_trj, s)
    q1 = rand(d)
    vs1 = zeros(d)

    dJds_st = 0.
    dJds_ust = 0.
    dJds_c = 0.
    N = 2000
    nSteps = nSteps + 1
    g = zeros(nSteps)
    J = (z .- 28.0).^2.0
    #dJdu = [0., 0., 1.]
    # Df represents the derivative of f
    # on the
    # unstable manifold.
    a = zeros(nSteps)
    c = zeros(nSteps)
    beta = zeros(nSteps)
    dJf = zeros(nSteps)
    Da = zeros(nSteps)
    Dc = zeros(nSteps)
    Dq = zeros(d)
    Dq1 = zeros(d)
    Dvs = zeros(d)
    Dvs1 = zeros(d)
    d2varphi = zeros(d,d,d)
    d2q = zeros(d,d)
    rho = s[2]

    # nSteps large number.
    for i = 1:nSteps-1
        dui = du_trj[:,:,i] # for large systems, can't store Jacobian.
        ppi = pp[:,i]
        dppi = dpp[:,:,i]
    	beta[i+1] = norm(f_trj[:,i+1])
    	f1 = f_trj[:,i+1]/beta[i+1]

        q1 .= dui*q
        alpha = norm(q1)
        q1 ./= alpha
        alpha2 = alpha*alpha
    	th = dot(q1, f1)
        th2 = 1.0 - th*th
        r1 = dunit_flow([x[i+1],y[i+1],z[i+1]],s)*q1

        vs1 .= dui*vs + ppi
        a[i+1] = dot(vs1, q1 - th*f1)/th2
    	c[i+1] = dot(vs1, f1 - th*q1)/th2

        d2varphi .= permutedims(ddu_trj[:,:,:,i],
				 [1,3,2])
    	for j = 1:d
    		d2q[:,j] .= d2varphi[:,:,j]*q
    	end

        Dq1 .= d2q*q/alpha2 + dui*Dq/alpha2
        dalphadx = dot(Dq1, q1)*alpha
	Dq1 .= Dq1 .- dalphadx*q1/alpha

        Dvs1 .= d2q*vs/alpha + dui*Dvs/alpha + dppi*q1
    	Dth = dot(q1, r1) + dot(f1, Dq1)
    	constant = 2*th*Dth/th2/th2
    	Da[i+1] = constant*dot(vs1, q1 - th*f1)
    	Da[i+1] += 1/th2*(dot(vs1, Dq1 - Dth*f1 - th*r1) +
    				 dot(Dvs1, q1 - th*f1))

    	Dc1 = constant*dot(vs1, f1 - th*q1)
    	Dc1 += 1/th2*(dot(vs1, r1 - Dth*q1 - th*Dq1) +
    				 dot(Dvs1, f1 - th*q1))

        g[i+1] = g[i]/alpha - dalphadx/alpha
        vs1 .= vs1 .- c[i+1]*f1 .- a[i+1]*q1

    	Dvs1 .= Dvs1 .- a[i+1]*Dq1 .- Da[i+1]*q1 .- c[i+1]*r1 .- Dc1*f1

	dJdu = [0.,0., 2*(z[i] - 28.0)]
        dJds_st += dot(dJdu, vs)/nSteps
    	dJf[i+1] = dot(dJdu, f1)

        vs .= vs1
        q .= q1
        Dvs .= Dvs1
	Dq .= Dq1
    end

    for n = 1:N
        J_shift = J[n+1:end]
    	dJf_shift = dJf[n+1:end]
    	betan_shift = beta[n+1:end]

        nJ = length(J_shift)
        g_shift = g[2:nJ+1]
        Da_shift = Da[2:nJ+1]
        a_shift = a[2:nJ+1]
    	c_shift = c[2:nJ+1]
    	beta_shift = beta[2:nJ+1]

        dJds_ust -= dot(J_shift,Da_shift)/nJ
        dJds_ust -= dot(J_shift,g_shift.*a_shift)/nJ
    	dJds_c += dot(c_shift./beta_shift, dJf_shift.*betan_shift)/nJ
    end
    return dJds_st + dJds_ust + dJds_c

end
function get_sens(s)
    nSteps = 200000
    n_exps = size(s)[2]
    dJds = zeros(n_exps)
    var_dJds = zeros(n_exps)
    n_rep = 16
    dJds_proc = SharedArray{Float64}(n_rep)
    for k=1:n_exps
        sk = s[:,k]
        dJds_proc .= 0.
        t = @distributed for j=1:n_rep
            dJds_proc[j] = sens(sk,nSteps)/n_rep
        end
        wait(t)
        dJds[k] = sum(dJds_proc)
        var_dJds[k] = n_rep*sum(dJds_proc.*dJds_proc) - (dJds[k])^2
    end
    save("../data/sens/lorenz/dvarz_ds2_2000.jld", "rho",
         s[2,:], "dJds", dJds, "var_dJds", var_dJds)
end


function mstprimjarnik(G::DictGraph; source = 0)::Tuple{SortedDict{Int, Array{Int,1}},PriorityQueue{Int, Int}}
    𝐷 = SortedDict{Int, Array{Int,1}}(i => [typemax(Int), -1] for i in keys(G.relation))
    𝒬 = PriorityQueue{Int, Int}(i => typemax(Int) for i in vertices(G))

    𝐷[source][1] = 0
    𝒬[source] = 0

    while !isempty(𝒬)
        𝑢 = dequeue_pair!(𝒬)
        for ℒ in G[𝑢.first]                    #for each source adj to 𝑢
            if ℒ[1] ∈ keys(𝒬) && ℒ[2] < 𝒬[ℒ[1]]
                𝐷[ℒ[1]][2] = 𝑢.first;
                𝐷[ℒ[1]][1] = ℒ[2]
                𝒬[ℒ[1]] = ℒ[2]
            end
        end
    end

    #create a new digraph T where T is a spanning tree.
    𝑇 = DictGraph(directed = true)

    for p in 𝐷
        p.first != source && addedge!(𝑇, p.second[2], p.first; weight = p.second[1])
    end

    (𝐷,𝑇)
end


abstract type StateSpaceModel <: DynamicIterator
end

struct StateObs{T,S}
    x::T
    obs::S
end

function rand(rng::AbstractRNG, U::StateObs)
    x = rand(rng, U.x)
    y = rand(rng, U.obs(Gaussian(x, false*I)))
    x, y
end

"""
```
LinearStateSpaceModel <: StateSpaceModel

LinearStateSpaceModel(sys, obs)
```

Combines a linear system `sys` and an observations model `obs` and
to a linear statespace model in a modular way.

Evolves `StateObs` objects.
"""
struct LinearStateSpaceModel{Tsys,Tobs} <: StateSpaceModel
    sys::Tsys
    obs::Tobs
end

function evolve(SSM::LinearStateSpaceModel, (t, x)::Pair{<:Any, <:Gaussian})
    t, x = evolve(SSM.sys, t => x)
    t => StateObs(x, SSM.obs)
end

function evolve(SSM::LinearStateSpaceModel, (t, (x,y))::Pair{<:Any, <:Tuple})
    t, x = evolve(SSM.sys, t => x)
    t => StateObs(x, SSM.obs)
end

dims(SSM) = size(SSM.H)

llikelihood(::LinearStateSpaceModel, yres, S) = logpdf(Gaussian(zero(yres), symmetrize!(S)), yres)


include("utils.jl")

@testset "Oblique manifold" begin
    M = Oblique(3, 2)
    @testset "Oblique manifold Basics" begin
        @test Sphere(2)^3 === Oblique(3, 3)
        @test Sphere(2)^(3,) === PowerManifold(Sphere(2), 3)
        @test ^(Sphere(2), 2) === Oblique(3, 2)
        @test typeof(^(Sphere(2), 2)) == Oblique{3,2,ℝ,2}
        @test repr(M) == "Oblique(3,2; field = ℝ)"
        @test representation_size(M) == (3, 2)
        @test manifold_dimension(M) == 4
        p = [2, 0, 0]
        p2 = [p p]
        @test !is_point(M, p)
        @test_throws DomainError is_point(M, p, true)
        @test !is_point(M, p2)
        @test_throws CompositeManifoldError is_point(M, p2, true)
        @test !is_vector(M, p2, 0.0)
        @test_throws CompositeManifoldError{ComponentManifoldError{Int64,DomainError}} is_vector(
            M,
            p2,
            [0.0, 0.0, 0.0],
            true,
        )
        @test !is_vector(M, p2, [0.0, 0.0, 0.0])
        @test_throws DomainError is_vector(M, p, [0.0, 0.0, 0.0], true) # p wrong
        @test injectivity_radius(M) ≈ π
        x = [1.0 0.0 0.0; 1.0 0.0 0.0]'
        @test_throws DomainError is_vector(M, x, [0.0, 0.0, 0.0], true) # tangent wrong
        y = [1.0 0.0 0.0; 1/sqrt(2) 1/sqrt(2) 0.0]'
        z = [1/sqrt(2) 1/sqrt(2) 0.0; 1.0 0.0 0.0]'
        basis_types = (DefaultOrthonormalBasis(),)
        transports = [ParallelTransport(), PowerVectorTransport(ParallelTransport())]
        test_manifold(
            M,
            [x, y, z],
            test_forward_diff=false,
            test_reverse_diff=false,
            test_vector_spaces=true,
            test_project_tangent=false,
            test_musical_isomorphisms=true,
            test_default_vector_transport=true,
            vector_transport_methods=transports,
            basis_types_to_from=basis_types,
            exp_log_atol_multiplier=1,
        )
    end
end


abstract type Primitive{T} end

struct NormedPrimitive{P,T} <: Primitive{T}
  radius::T
  NormedPrimitive{P,T}(radius::T) where {P,T} = new{P,T}(radius)
  NormedPrimitive{P}(radius::T) where {P,T} = NormedPrimitive{P,T}(radius)
end

norm(p::Point, ::Type{<:NormedPrimitive{P}}) where {P} = norm(p, P)

(np::NormedPrimitive)(p) = norm(p, typeof(np)) - np.radius

(≈)(x::NormedPrimitive, y::NormedPrimitive) = typeof(x) == typeof(y) && x.radius ≈ y.radius

const HyperSphere{T} = NormedPrimitive{2,T}

norm(p::Point, ::Type{<:HyperSphere}) = hypot(p...)

const HyperCube{T} = NormedPrimitive{Inf,T}

"""
    Ellipsoid(semiaxes)

An ellipsoid with semi-axes `semiaxes`.
Is equivalent to a scaled [`HyperSphere`](@ref).
"""
const Ellipsoid{Dim,T} = Scaled{HyperSphere{T},Dim,T}
Ellipsoid(radius::T, transf::Scaling{Dim,T}) where {Dim,T} = Ellipsoid{Dim,T}(HyperSphere(radius), transf)

function Ellipsoid(semiaxes::AbstractVector)
  radius = norm(semiaxes, HyperSphere)
  Ellipsoid(radius, Scaling(semiaxes ./ radius))
end
Ellipsoid(semiaxes::Number...) = Ellipsoid(collect(semiaxes))

(≈)(x::Ellipsoid, y::Ellipsoid) = x.obj.radius .* x.transf.vec ≈ y.obj.radius .* y.transf.vec

Base.show(io::IO, elps::Ellipsoid{Dim,T}) where {Dim,T} = print(io, "Ellipsoid{$Dim, $T}($(elps.transf.vec .* elps.obj.radius))")

# A box around the origin. This is not the same as `Meshes.Box`, it is missing the translation part.
const Box{Dim,T} = Scaled{HyperCube{T},Dim,T}
Box(radius::T, transf::Scaling{Dim,T}) where {Dim,T} = Box{Dim,T}(HyperCube(radius), transf)

apply(t::Scaling, s::NormedPrimitive) = typeof(s)(norm(s, t.vec) * s.radius)
apply(t::Rotation, s::HyperSphere) = s

const Circle{T} = Projection{2,HyperSphere{T}}
const Square{T} = Projection{2,HyperCube{T}}


@testset "bigfloat initial convergence #720" begin
    f(x) = x[1]^2
    x0 = BigFloat[0]
    obj = OnceDifferentiable(f, x0; autodiff=:forward)
    for method in (GradientDescent, BFGS, LBFGS, AcceleratedGradientDescent, MomentumGradientDescent, ConjugateGradient, Newton, NewtonTrustRegion, SimulatedAnnealing)
        result = Optim.optimize(f, x0, method())
    end
end


# Preconditioned Crank–Nicolson algorithm tools

function generate_innovation(W0::Number,t,rng)
  dt = diff(t)
  Wnew = cumsum([zero(W0);[sqrt(dt[i]).*wiener_randn(rng,typeof(W0))
            for i in 1:length(t)-1]])
end

function generate_innovation(W0,t,rng)
  dt = diff(t)
  Wnew = cumsum([[zero(W0)];[sqrt(dt[i]).*wiener_randn(rng,W0)
            for i in 1:length(t)-1]])
end

"""
    pCN!(noise::AbstractNoiseProcess, ρ; reset=true,reverse=false,indx=nothing)

Create a new, but correlated noise process from `noise` and additional entropy with correlation ρ.
This update defines an autoregressive process in the space of Wiener (or noise process) trajectories which can be used as proposal distribution in Metropolis-Hastings algorithms (often called "preconditioned Crank–Nicolson scheme".)

External links
 * [Preconditioned Crank–Nicolson algorithm on Wikipedia](https://en.wikipedia.org/wiki/Preconditioned_Crank–Nicolson_algorithm)
"""
function pCN!(source::AbstractNoiseProcess{T,N,Vector{T2},inplace}, ρ;
              reset=true,reverse=false,indx=nothing) where {T,N,T2,inplace}

  # generate new Wiener process similar to the one in source
  Wnew = generate_innovation(source.W[1],source.t,source.rng)

  source.W = ρ * source.W + sqrt(one(ρ)-ρ^2) * Wnew
  source.u = ρ * source.u + sqrt(one(ρ)-ρ^2) * Wnew
  NoiseWrapper(source,reset=reset,reverse=reverse,indx=indx)
end

"""
    pCN(noise::AbstractNoiseProcess, ρ; reset=true,reverse=false,indx=nothing)

Create a new, but correlated noise process from `noise` and additional entropy with correlation ρ.
"""
function pCN(source::AbstractNoiseProcess{T,N,Vector{T2},inplace}, ρ;
              reset=true,reverse=false,indx=nothing) where {T,N,T2,inplace}

  source′ = deepcopy(source)

  # generate new Wiener process similar to the one in source
  Wnew = generate_innovation(source′.W[1],source′.t,source′.rng)

  source′.W = ρ * source′.W + sqrt(one(ρ)-ρ^2) * Wnew
  source′.u = ρ * source′.u + sqrt(one(ρ)-ρ^2) * Wnew
  NoiseWrapper(source′,reset=reset,reverse=reverse,indx=indx)
end

"""
    pCN(noise::NoiseGrid, ρ; reset=true, rng = Xorshifts.Xoroshiro128Plus(rand(UInt64)))

Create a new, but correlated noise process from `noise` and additional entropy with correlation ρ.
This update defines an autoregressive process in the space of Wiener (or noise process) trajectories which can be used as proposal distribution in Metropolis-Hastings algorithms (often called "preconditioned Crank–Nicolson scheme".)

External links
 * [Preconditioned Crank–Nicolson algorithm on Wikipedia](https://en.wikipedia.org/wiki/Preconditioned_Crank–Nicolson_algorithm)
"""
function pCN(source::NoiseGrid, ρ; reset=true, rng = Xorshifts.Xoroshiro128Plus(rand(UInt64)))

  # generate new Wiener process similar to the one in source
  t = source.t
  Wnew = generate_innovation(source.W[1],source.t,rng)

  W = ρ * source.W + sqrt(one(ρ)-ρ^2) * Wnew
  NoiseGrid(t,W,source.Z;reset=reset)
end


module TestParameter

using ModiaLang
using DifferentialEquations
@usingModiaPlot
include("$(ModiaLang.path)/models/Blocks.jl")

inputSignal(t) = sin(t)

FirstOrder = Model(
    T = parameter | Var(value=0.2),
    x = Var(init=0.3),
    equations = :[u = inputSignal(time/u"s"),
                  T * der(x) + x = u,
                  y = 2*x]
)

firstOrder = @instantiateModel(FirstOrder)

simulate!(firstOrder, Tsit5(), stopTime = 10, log=false, logParameters=false, logEvaluatedParameters=false, requiredFinalStates = [-0.3617373025974107])

plot(firstOrder, ["u", "x", "der(x)", "y"], figure=1)

# Second order system:
#     w = 1.0
#     D = 0.1
#     k = 2.0
#     der(x1) = x2
#     der(x2) = -w^2*x1 - 2*D*w*x2 + w^2*u
#          y  = k*x1
SecondOrder = Model(
    w = 20.0,
    D =  0.1,
    k =  2.0,
    sys = StateSpace | Map(A = parameter | :([  0        1;
                                              -w^2  -2*D*w]),
                           B = parameter | :([0; w^2]),
                           C = parameter | :([k 0]),
                           D = parameter | :(zeros(1)),
                           x = Var(init = zeros(2))),
    equations = :[sys.u = 1.0]
)

using Test
@test_skip begin
secondOrder = @instantiateModel(SecondOrder, unitless=true)

simulate!(secondOrder, merge=Map(D=0.3), logParameters=true, logEvaluatedParameters=true,
          requiredFinalStates = [0.9974089572681231, 0.011808652820321723])
plot(secondOrder, [("sys.u", "sys.y"); "sys.x"], figure = 2)

simulate!(secondOrder, merge=Map(w=100.0, D=0.1, k=1.1), logParameters=true, logEvaluatedParameters=true,
          requiredFinalStates = [0.9999806306477742, -0.00391695258236727])
plot(secondOrder, [("sys.u", "sys.y"); "sys.x"], figure = 3)
end
end


using BenchmarkTools

# Get the matrix dimensions N from the command line argument.
N = parse(Int, ARGS[1])

A = ones(N, N)
for i = 1:N
    A[i, i] = 6
end

println("--------------------------")
println(@sprintf "Matric square root: %d" N)
println("--------------------------")

# Take the square root of matrix A.
@btime begin
    B = sqrtm(A)
end


# This is part of Calendars.jl. See the copyright note there.
# ===================== European calendar ===================

# The Gregorian calendar is used for dates on and after
# CE 1582-10-15, otherwise the Julian calendar is used.

# Is the date a valid European date?
function isValidDateEuropean(cd::CDate, warn=true)
    cal, year, month, day = cd
    if cal != EC
        warn && @warn(Warning(cd))
        return false
    end

    dn = DayNumberGregorian(year, month, day)
    if dn < GregorysBreak
        return isValidDateJulian((JD, year, month, day), warn)
    end
    return isValidDateGregorian((CE, year, month, day), warn)
end

# Computes the day number from a date.
# All dates prior to CE 1582-10-15 are interpreted as Julian dates.
function DayNumberEuropean(year::DPart, month::DPart, day::DPart)
    dn = DayNumberGregorian(year, month, day)
    if dn < GregorysBreak
        return DayNumberJulian(year, month, day)
    end
    return dn
end

# Computes the day number from a valid date.
function DayNumberEuropean(cd::CDate)
    cal, year, month, day = cd
    return DayNumberEuropean(year, month, day)
end

# Computes the historic date from a day number.
# All dates prior to CE 1582-10-15 are returned as Julian dates.
function DateEuropean(dn::DPart)
    d = dn < GregorysBreak ? DateJulian(dn) : DateGregorian(dn)
    cal, year, month, day = d
    return (EC, year, month, day)
end

# Return the days this year so far.
function DayOfYearEuropean(year::DPart, month::DPart, day::DPart)
    dn = DayNumberGregorian(year, month, day)
    if dn < GregorysBreak
        return DayOfYearJulian(year, month, day)
    end
    return DayOfYearGregorian(year, month, day)
end

# Return the day number of the first day in the given European year.
function YearStartEuropean(year::DPart)
    if year <= 1582
        return YearStartJulian(year)
    end
    return YearStartGregorian(year)
end

# Return the day number of the last day in the given European year.
function YearEndEuropean(year::DPart)
    if year < 1582
        return YearEndJulian(year)
    end
    return YearEndGregorian(year)
end

function MonthsInYearEuropean(year::DPart)
    return 12
end

function DaysInYearEuropean(year::DPart)
    if year < 1582
        return DaysInYearJulian(year)
    end
    return DaysInYearGregorian(year)
end


using NonsmoothOptim
using Test

@testset "NonsmoothOptim.jl" begin
  @testset "Subgradient" begin
    @testset "1D: abs(x)" begin
      f = (x) -> abs(x[1])
      g = (x) -> [(x[1] < 0) ? -1.0 : 1.0]

      @testset "Minimise 1D" begin
        p = UnconstrainedNonSmoothProblem(f, g, 1, Minimise)
        a = SubgradientMethod(ConstantStepSize(1.0), 2)

        x0 = [10.0]
        x1 = solve(p, a, x0)
        @test length(x1) == 1
        @test x1[1] ≈ 8.0
      end

      @testset "Minimise 1D (start at optimum)" begin
        p = UnconstrainedNonSmoothProblem(f, g, 1, Minimise)
        a = SubgradientMethod(ConstantStepSize(1.0), 2)

        x0 = [0.0]
        x1 = solve(p, a, x0)
        @test length(x1) == 1
        @test x1[1] ≈ x0[1]
      end

      @testset "Maximise 1D" begin
        p = UnconstrainedNonSmoothProblem(f, g, 1, Maximise)
        a = SubgradientMethod(ConstantStepSize(1.0), 2)

        x0 = [10.0]
        x1 = solve(p, a, x0)
        @test length(x1) == 1
        @test x1[1] ≈ 12.0
      end

      @testset "Maximise 1D (2)" begin
        p = UnconstrainedNonSmoothProblem(f, g, 1, Maximise)
        a = SubgradientMethod(ConstantStepSize(1.0), 2)

        x0 = [-10.0]
        x1 = solve(p, a, x0)
        @test length(x1) == 1
        @test x1[1] ≈ -12.0
      end

      @testset "Info callback" begin
        cb_called = false
        cb = (k, f, x, g, f_best, x_best, δt) -> begin
          cb_called = true
        end
        p = UnconstrainedNonSmoothProblem(f, g, 1, Maximise)
        a = SubgradientMethod(ConstantStepSize(1.0), 2)
        solve(p, a, [0.0], info_callback=cb)

        @test cb_called
      end
    end
  end

  @testset "Projected subgradient" begin
    @testset "1D: abs(x)" begin
      f = (x) -> abs(x[1])
      g = (x) -> [(x[1] < 0) ? -1.0 : 1.0]
      proj = (x) -> [200.0]

      @testset "Minimise 1D" begin
        p = ProjectedConstrainedNonSmoothProblem(f, g, proj, 1, Minimise)
        a = ProjectedSubgradientMethod(ConstantStepSize(1.0), 2)

        x0 = [10.0] # Infeasible solution! Must be projected!
        x1 = solve(p, a, x0)
        @test length(x1) == 1
        @test x1[1] ≈ 200.0
      end
    end
  end

  @testset "Bundle" begin
  end
end


using Roots

function polynomialGenerator(a...)
    n = length(a)-1
    poly =  function(x)
                return sum([a[i+1]*x^i for i in 0:n])
            end
    return poly
end

polynomial = polynomialGenerator(1,3,-10)
zeroVals = find_zeros(polynomial,-10,10)
println("Zeros of the function f(x): ", zeroVals)


import NNlib: softmax, softmax!, ∇softmax, ∇softmax!,
              logsoftmax, logsoftmax!, ∇logsoftmax, ∇logsoftmax!

using CUDA.CUDNN: CUDNN_SOFTMAX_LOG, CUDNN_SOFTMAX_MODE_CHANNEL,
                CUDNN_SOFTMAX_FAST, CUDNN_SOFTMAX_ACCURATE, cudnnSoftmaxForward!,
                cudnnSoftmaxBackward

# Softmax

# @denizyuret: do not do inplace operations with softmax/logsoftmax when (1) cpu version is not, (2) one can use softmax!
function softmax(x::T; dims=1) where {T<:DenseCuArray}
    softmax!(similar(x), x; dims)
end

function ∇softmax(dy::T, x::T, y::T; dims=1) where {T<:DenseCuArray}
    ∇softmax!(similar(x), dy, x, y; dims)
end

function logsoftmax(x::T; dims=1) where {T<:DenseCuArray}
    logsoftmax!(similar(x), x; dims)
end

function ∇logsoftmax(dy::T, x::T, y::T; dims=1) where {T<:DenseCuArray}
    ∇logsoftmax!(similar(x), dy, x, y; dims)
end

# @denizyuret: backup implementations for unsupported/slow size/dims combinations:
function _softmax!(y::T, x::T; dims) where {T<:DenseCuArray}
    y .= exp.(x .- maximum(x; dims))
    y ./= sum(y; dims)
end

function _∇softmax!(dx::T, dy::T, x::T, y::T; dims) where {T<:DenseCuArray}
    dx .= y .* (dy .- sum(dy .* y; dims))
end

function _logsoftmax!(y::T, x::T; dims) where {T<:DenseCuArray}
    y .= x .- maximum(x; dims)
    y .-= log.(sum(exp.(y); dims))
end

function _∇logsoftmax!(dx::T, dy::T, x::T, y::T; dims) where {T<:DenseCuArray}
    dx .= dy .- sum(dy; dims) .* exp.(y)
end

# Trick by @norci to use cudnn for softmax dims args that are contiguous:
# If dims=(dmin:dmax) then CUDNN_SOFTMAX_MODE_CHANNEL does the trick with reshape
#    (1, prod(size(x)[1:dmin-1]), prod(size(x)[dmin:dmax]), :)
# softmaxdims returns nothing when the backup implementation should be used.

function softmaxdims(x, dims)
    dims === Colon() && return (1, 1, length(x), 1)
    mind,maxd = minimum(dims),maximum(dims)
    all(i in dims for i in mind:maxd) || return nothing # cannot handle if not contiguous
    stride = dimsize = 1
    for i in 1:(mind-1); stride *= size(x,i); end # Using size(x,i) assumes trailing dims = 1, robust to maxd > ndims(x)
    for i in mind:maxd; dimsize *= size(x,i); end
    batchsize = length(x)÷(stride*dimsize)
    # Here is a region where cudnn is slower, so we go with the backup:
    batchsize == 1 && 64 <= stride <= 4096 && 64 <= dimsize <= 4096 && return nothing
    return (1, stride, dimsize, batchsize)
end

# Determine softmax algo based on math_mode

softmaxalgo() = (CUDA.math_mode()===CUDA.FAST_MATH ? CUDNN_SOFTMAX_FAST : CUDNN_SOFTMAX_ACCURATE)

# Main implementations:

function softmax!(y::T, x::T = y; dims=1) where {T<:DenseCuArray}
    s = softmaxdims(x, dims)
    s === nothing && return _softmax!(y, x; dims)
    cudnnSoftmaxForward!(reshape(y,s), reshape(x,s); mode = CUDNN_SOFTMAX_MODE_CHANNEL, algo = softmaxalgo())
    return y
end

function ∇softmax!(dx::T, dy::T, x::T, y::T; dims=1) where {R,T<:DenseCuArray{R}}
    s = softmaxdims(x, dims)
    s === nothing && return _∇softmax!(dx, dy, x, y; dims)
    xDesc = cudnnTensorDescriptor(reshape(x,s))
    alpha, beta = scalingParameter(R,1), scalingParameter(R,0)
    cudnnSoftmaxBackward(handle(), softmaxalgo(), CUDNN_SOFTMAX_MODE_CHANNEL,
                         alpha, xDesc, y, xDesc, dy, beta, xDesc, dx)
    return dx
end

function logsoftmax!(y::T, x::T = y; dims=1) where {T<:DenseCuArray}
    s = softmaxdims(x, dims)
    s === nothing && return _logsoftmax!(y, x; dims)
    cudnnSoftmaxForward!(reshape(y,s), reshape(x,s); mode = CUDNN_SOFTMAX_MODE_CHANNEL, algo = CUDNN_SOFTMAX_LOG)
    return y
end

function ∇logsoftmax!(dx::T, dy::T, x::T, y::T; dims=1) where {R,T<:DenseCuArray{R}}
    s = softmaxdims(x, dims)
    s === nothing && return _∇logsoftmax!(dx, dy, x, y; dims)
    xDesc = cudnnTensorDescriptor(reshape(x,s))
    alpha, beta = scalingParameter(R,1), scalingParameter(R,0)
    cudnnSoftmaxBackward(handle(), CUDNN_SOFTMAX_LOG, CUDNN_SOFTMAX_MODE_CHANNEL,
                         alpha, xDesc, y, xDesc, dy, beta, xDesc, dx)
    return dx
end


function loglike(x)
    l = 0
    H .= H.^2
    Q .= Q.^2
    for t in 1:n
        v = y[t] - Z*a - d
        F = Z*P*Z' + Diagonal(H)
        a_t = a + P*Z'*F^(-1)*v
        a = T*a_t + c
        P_t = P - P*Z'*F^(-1)*Z*P
        P = T*P_t*T' + R*Diagonal(Q)*R'

        l += -1/2*logabsdet(F) - 1/2 v'*F*v
    end
    return l
end


using MCCGA

lower = [-32.768, -32.768, -32.768]
upper = [32.768, 32.768, 32.768]

result = MCCGA.mccga(
        lower = lower,
        upper = upper,
        costfunction = Functions.ackley,
        popsize = 100,
        maxsamples = 100000,
    )


using CounterfactualRegret: IIEMatrixGame, Kuhn

function CFRMatrixTest(sol_type, N::Int)
    ## Rock Paper Scissors
    game = IIEMatrixGame([
        (0,0) (-1,1) (1,-1);
        (1,-1) (0,0) (-1,1);
        (-1,1) (1,-1) (0,0)
    ])
    trainer = sol_type(game)
    train!(trainer, N)
    s1 = trainer.I[0].s
    s1 ./= sum(s1)
        @test ≈(s1, fill(1/3,3), atol=0.01)
    s2 = trainer.I[1].s
    s2 ./= sum(s2)
        @test ≈(s2, fill(1/3,3), atol=0.01)
    F_eval = evaluate(trainer)
        @test all( .≈(F_eval,(0,0), atol=0.01))
    MC_eval = MonteCarloEvaluate(trainer,1)
        @test all( .≈(MC_eval,(0,0), atol=0.01))

    ## Prisoners Dilemma
    game = IIEMatrixGame([
        (-1,-1) (-3,0);
        (0,-3) (-2,-2)
    ])
    trainer = sol_type(game)
    train!(trainer, N)
    s1 = trainer.I[0].s
    s1 ./= sum(s1)
        @test ≈(s1, [0,1], atol=0.01)
    s2 = trainer.I[1].s
    s2 ./= sum(s2)
        @test ≈(s2, [0,1], atol=0.01)
    F_eval = evaluate(trainer)
        @test all( .≈(F_eval,(-2,-2), atol=0.01))
    MC_eval = MonteCarloEvaluate(trainer,1)
        @test all( .≈(MC_eval,(-2,-2), atol=0.01))

    # https://sites.math.northwestern.edu/~clark/364/handouts/bimatrix-mixed.pdf
    game = CounterfactualRegret.IIEMatrixGame([
        (1,1) (0,0) (0,0);
        (0,0) (0,2) (3,0);
        (0,0) (2,0) (0,3);
    ])
    trainer = sol_type(game)
    train!(trainer, N)
    NEs = [[6/11,3/11,2/11], [0,3/5,2/5], [1,0,0]]
    s1 = trainer.I[0].s
    s1 ./= sum(s1)
        @test begin
            ≈(s1, NEs[1], atol=0.01) ||
            ≈(s1, NEs[2], atol=0.01) ||
            ≈(s1, NEs[3], atol=0.01)
        end
    s2 = trainer.I[1].s
    s2 ./= sum(s2)
        @test begin
            ≈(s2, NEs[1], atol=0.01) ||
            ≈(s2, NEs[2], atol=0.01) ||
            ≈(s2, NEs[3], atol=0.01)
        end
    F_eval = evaluate(trainer)
    @test begin
        all( .≈(F_eval,(6/5,6/5), atol=0.01)) ||
        all( .≈(F_eval,(6/11,6/11), atol=0.01)) ||
        all( .≈(F_eval,(1,1), atol=0.01))
    end
    MC_eval = MonteCarloEvaluate(trainer,1)
    @test begin
        all( .≈(F_eval,(6/5,6/5), atol=0.01)) ||
        all( .≈(F_eval,(6/11,6/11), atol=0.01)) ||
        all( .≈(F_eval,(1,1), atol=0.01))
    end

    sol = sol_type(game; debug=true)
    train!(sol, 100_000)

    @test CFR.infokeytype(sol) == Int
end

function CFRKuhnTest(sol_type, N::Int, atol::Float64)
    game = Kuhn()
    trainer = sol_type(game)
    train!(trainer, N)

    s11__ = trainer.I[(1,1,SA[-1,-1,-1])].s
    s11__ ./= sum(s11__)
    α = s11__[2]
    @test 0 ≤ α ≤ 1/3
    s1200 = trainer.I[(1,2,SA[-1,-1,-1])].s
    s1200 ./= sum(s1200)
    @test ≈(s1200, [1,0], atol=atol)
    s1300 = trainer.I[(1,3,SA[-1,-1,-1])].s
    s1300 ./= sum(s1300)
    @test ≈(s1300, [1-3α,3α], atol=atol)

    s210_ = trainer.I[(2,1,SA[0,-1,-1])].s
    s210_ ./= sum(s210_)
    @test ≈(s210_, [2/3,1/3], atol=atol)
    s220_ = trainer.I[(2,2,SA[0,-1,-1])].s
    s220_ ./= sum(s220_)
    @test ≈(s220_, [1,0], atol=atol)
    s230_ = trainer.I[(2,3,SA[0,-1,-1])].s
    s230_ ./= sum(s230_)
    @test ≈(s230_, [0,1], atol=atol)

    s211_ = trainer.I[(2,1,SA[1,-1,-1])].s
    s211_ ./= sum(s211_)
    @test ≈(s211_, [1,0], atol=atol)
    s221_ = trainer.I[(2,2,SA[1,-1,-1])].s
    s221_ ./= sum(s221_)
    @test ≈(s221_, [2/3,1/3], atol=atol)
    s231_ = trainer.I[(2,3,SA[1,-1,-1])].s
    s231_ ./= sum(s231_)
    @test ≈(s231_, [0,1], atol=atol)

    s1101 = trainer.I[(1,1,SA[0,1,-1])].s
    s1101 ./= sum(s1101)
    @test ≈(s1101, [1,0], atol=atol)
    s1201 = trainer.I[(1,2,SA[0,1,-1])].s
    s1201 ./= sum(s1201)
    @test ≈(s1201, [2/3-α,1/3+α], atol=atol)
    s1301 = trainer.I[(1,3,SA[0,1,-1])].s
    s1301 ./= sum(s1301)
    @test ≈(s1301, [0,1], atol=atol)
end

@testset "IIE Solvers" begin
    @testset "CFR Matrix" begin CFRMatrixTest(CFRSolver, 100_000) end
    @testset "CFR Kuhn" begin CFRKuhnTest(CFRSolver, 100_000, 0.03) end

    @testset "CSCFR Matrix" begin CFRMatrixTest(CSCFRSolver, 100_000) end
    @testset "CSCFR Kuhn" begin CFRKuhnTest(CSCFRSolver, 1_000_000, 0.03) end

    @testset "DCFR Matrix" begin CFRMatrixTest(DCFRSolver, 100_000) end
    @testset "DCFR Kuhn" begin CFRKuhnTest(DCFRSolver, 200_000, 0.03) end

    @testset "ESCFR Matrix" begin CFRMatrixTest(ESCFRSolver, 500_000) end
    @testset "ESCFR Kuhn" begin CFRKuhnTest(ESCFRSolver, 1_000_000, 0.03) end
end


function basic_linalg_problem(;Ns=1000,Nc=100)

  A = Diagonal(2*ones(Ns))
  C = Diagonal(ones(Nc))
  B2 = zeros(Nc,Ns)
  for j in 1:min(Nc,Ns)
      B2[j,j] = 1.0
  end
  B1t = B2'
  rhs1v = ones(Ns)
  rhs2v = 2*ones(Nc)

  Abig = [A B1t;B2 C]
  rhsbig = [rhs1v;rhs2v]
  solex = similar(rhsbig)
  solex .= Abig\rhsbig;

  return A,B2,B1t,C,rhs1v,rhs2v, solex

end


function CRT_GIA_kin_fb(m, Nx, ϕ, u, D, k_f, k_b, nu_A, nu_B, nu_C, a, b, c, C_bc, C_0, N_comp, dt, T, tol)
  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)

  # Preparing Jacobian terms
  J_eps = Array{Any}(N_comp)
  RHS_eps = Array{Any}(N_comp)

  # Preparing system matrices and RHS
  M_system = Array{Any}(N_comp)
  RHS_system = Array{Any}(N_comp)
  C_new = Array{Any}(N_comp)
  C_temp = Array{Any}(N_comp)

  # Defining time-stepping parameters
  maxit = 5000 # maximum number of iterations

  # Time-stepping
  for t in dt:dt:T
    iter = 0
    error = 1000 # forcing at least one Newton-iteration

    #Assigning C_old as the new start guess for the Newton-Iteration and creating the transport related system matrix
    for i in 1:N_comp
      C_temp[i] = copyCell(C_old[i])
      M_t[i], RHS_t[i] = transientTerm(C_old[i], dt, ϕ)
      M_system[i] = M_t[i] + M_stat[i]
    end

    #Newton iterations loop
    while (error > tol && iter < maxit)

      # Constructing Jacobian terms
      J_eps[1] = linearSourceTerm(createCellVariable(m, a*k_f.*(internalCells(C_temp[1]).^(a-1)).*(internalCells(C_temp[2]).^b), BC[1]))
      J_eps[2] = linearSourceTerm(createCellVariable(m, b*k_f.*(internalCells(C_temp[1]).^a).*(internalCells(C_temp[2]).^(b-1)), BC[2]))
      J_eps[3] = linearSourceTerm(createCellVariable(m, -c*k_b.*(internalCells(C_temp[3]).^(c-1)), BC[3]))

      RHS_reacA = constantSourceTerm(createCellVariable(m, -k_f.*(internalCells(C_temp[1]).^a).*(internalCells(C_temp[2]).^b).*(1-a-b)+k_b*(internalCells(C_temp[3]).^c)*(1-c), BC[1]))
      RHS_reacB = constantSourceTerm(createCellVariable(m, -k_f.*(internalCells(C_temp[1]).^a).*(internalCells(C_temp[2]).^b).*(1-a-b)+k_b*(internalCells(C_temp[3]).^c)*(1-c), BC[2]))
      RHS_reacC = constantSourceTerm(createCellVariable(m, -k_f.*(internalCells(C_temp[1]).^a).*(internalCells(C_temp[2]).^b).*(1-a-b)+k_b*(internalCells(C_temp[3]).^c)*(1-c), BC[3]))

      # Constructing the RHS's
      RHS_system[1] = RHS_t[1] + RHS_bc[1] - nu_A*RHS_reacA
      RHS_system[2] = RHS_t[2] + RHS_bc[2] - nu_B*RHS_reacB
      RHS_system[3] = RHS_t[3] + RHS_bc[3] + nu_C*RHS_reacC

      RHS = [RHS_system[1];
            RHS_system[2];
            RHS_system[3]]

      # Constructing block system matrix
      M = [M_system[1]-(nu_A*J_eps[1]) -(nu_A*J_eps[2]) -(nu_A*J_eps[3]);
          -(nu_B*J_eps[1]) M_system[2]-(nu_B*J_eps[2]) -(nu_B*J_eps[3]);
          (nu_C*J_eps[1]) (nu_C*J_eps[2]) M_system[3]+(nu_C*J_eps[3])]

      # Solving the system
      sol = M\RHS

      # Putting solution into JFVM structure
      for i in 1:N_comp
        C_new[i] = copyCell(C_temp[i])
        C_new[i].value[:] = sol[((i-1)*(Nx+2)+1):(i*(Nx+2))]
      end

      # Calculating convergence criteria
      error = 0
      for i in 1:N_comp
        error = error + sum(abs(C_new[i].value-C_temp[i].value))
        C_temp[i] = copyCell(C_new[i])
      end

      iter = iter + 1
    end

    #if (t == dt) break end

    for i in 1:N_comp
      C_old[i] = copyCell(C_new[i])
    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


# Dynamically keep track of the eigenvalues system of a hermitian, positive
# definite matrix that is growing diagonally.
#   * We need to keep track of its inverse, which is not going to be very stable numerically, but it should do the trick
#   of tracking eigenvalues using the Cauchy Interlace Theorem.
#   * We also want to keep track of it's eigensystem, which is going to serve for the Lanczos Algorithm with selective
#   re-orthogonalizations!
#   * It starts keeping track of the eigensystem when it's been asked to, and it will start doing it.

mutable struct DynamicSymTridiagonal{T<:AbstractFloat}
    alphas::Vector{T}       # Diaognal
    betas::Vector{T}        # Lower & upper Diaognal
    L::Vector{T}            # The lower diagonal of the unit-bidiagonal matrix L
    U::Vector{T}            # Lower diagonal of the upper bi-diagonal matrix U
    k::Int64                # Size of the matrix.

    V::Matrix{T}            # Eigenvectors
    thetas::Vector{T}       # Eigen values.

    function DynamicSymTridiagonal{T}(alpha::T) where {T<:Float64}
        this = new{T}()
        this.alphas = Vector{T}()
        push!(this.alphas, alpha)
        this.betas = Vector{T}()
        this.L = Vector{T}()
        this.U = Vector{T}()
        push!(this.U, alpha)
        this.k = 1
        return this
    end

    function DynamicSymTridiagonal(alpha::AbstractFloat)
        T = typeof(alpha)
        return DynamicSymTridiagonal{T}(alpha)
    end
end

"""
    append an alpha and a beta (diagonal and sub & super diagonal) elemnt to the current sym diagonal matrix.
"""
function (this::DynamicSymTridiagonal{T})(alpha::T, beta::T) where {T<:AbstractFloat}
    push!(this.alphas, alpha)
    push!(this.betas, beta)
    push!(this.U, alpha - beta^2/this.U[end])
    push!(this.L, beta/this.U[end - 1])
    this.k += 1
    return this
end

"""
    Using the Lapack library to get the eigen system of the curent T, and then after
    this, it will keep track of the eigensystem as elements are added to the dynamic matrix.
"""
function InitializeEigenSystem(this::DynamicSymTridiagonal{T}, max_size::Int64) where {T<:AbstractFloat}
    if max_size > 32728 || max_size < 2
        error("Eigen system estimated max size can't be: $(max_size)")
    end
    Trid = GetT(this)
    theta, V = eigen(Trid)
    Vbigger = zeros(T, max_size, max_size)
    Vbigger[1:this.k, 1:this.k] = V
    this.V = Vbigger
    Theta = zeros(T, max_size)
    Theta[1:this.k] = theta
    this.thetas = Theta
    # Sort by absolute values.
    SortedIdx = sortperm(abs.(this.thetas))
    this.thetas = this.thetas[SortedIdx]
    this.V = this.V[:, SortedIdx]
    return this.thetas, this.V
end

"""
    Update the eigensystem of the bigger matrix using the preivous eigen system. Use Power Iterations
    iteratively.
"""
function UpdateEigenSystemPowItr(this::DynamicSymTridiagonal{T}) where {T <: AbstractFloat}
    # if !(isdefined(Base, :this.V))  # cheks if defined first.
    #     error("Must initialize the eigensystem with an estimated size before calling this function. ")
    # end
    k = this.k
    V = this.V
    Θ = this.thetas
    A = this
    if k > size(this.V, 1)
        error("Hasn't implemented this part yet, k exceed the maximal size of eigen system.")
    end
    Avⱼ = zeros(T, k)
    for j in 1: k
        vⱼ = view(V, 1:k, j)
        vⱼ .+= 1e-10*rand(T, k)
        if j >= 2   # Components on to previous ortho eigen vectors removed.
            vⱼ .-= view(V, 1:k, 1:j - 1)*view(V, 1:k, 1:j - 1)'*vⱼ
        end
        for Itr in 1: k^2 + 1e3      # Power iterations
            vⱼᵀvⱼ = dot(vⱼ, vⱼ)
            Apply!(A, vⱼ, Avⱼ)
            if j >= 2  # Components on to previous ortho eigen vectors removed.
                Avⱼ .-= view(V, 1:k, 1:j - 1)*view(V, 1:k, 1:j - 1)'*Avⱼ
            end
            λ̃ = dot(vⱼ, Avⱼ)/vⱼᵀvⱼ  # Rayleigh Quotient.
            ∇r = 2(Avⱼ - λ̃*vⱼ)      # Gradient of Rayleigh Quotient.
            if norm(∇r, Inf) <= 1e-8
                V[1:k, j] .= vⱼ/sqrt(vⱼᵀvⱼ)
                Θ[j] = λ̃
                break               # go to Next Eigenvector.
            end
            vⱼ .= Avⱼ/norm(Avⱼ)
        end
    end
    return
end

"""
    Get the L matrix for this instance.
"""
function GetL(this::DynamicSymTridiagonal)
    return Bidiagonal(fill(1, this.k), convert(Vector{Float64}, this.L), :L)
end

"""
    Get the U matrix for this instance.
"""
function GetU(this::DynamicSymTridiagonal)
    return Bidiagonal(convert(Vector{Float64}, this.U), this.betas, :U)
end

"""
    Get the Tridiagonal Matrix
"""
function GetT(this::DynamicSymTridiagonal)
    return SymTridiagonal(this.alphas, this.betas)
end

"""
    Solve this matrix against another matrix or vector.
"""
function Base.:\(
    this::DynamicSymTridiagonal{T},
    b::Union{Matrix{T}, Vector{T}}
    ) where {T<:AbstractFloat}

    return
end

"""
    Multiply a vector on the left hand size of this dynamically growing matrix
    T.
"""
function Base.:*(this::DynamicSymTridiagonal{T}, b::AbstractArray{T}) where {T <: AbstractFloat}
    @assert length(b) == this.k "Dimension of vector b does match the number of rows in k. Expect $(this.k) but get $(length(b))"
    β = this.betas
    α = this.alphas
    v = similar(b)
    v[1] = α[1]*b[1] + β[1]*b[2]
    for Idx in 2: this.k - 1
        v[Idx] = β[Idx - 1]*b[Idx - 1] + α[Idx]*b[Idx] + β[Idx]*b[Idx + 1]
    end
    v[this.k] = β[this.k - 1]*b[end - 1] + α[this.k]*b[end]
    return v
end

"""
    Memory friendly version of multiplications.
"""
function Apply!(this::DynamicSymTridiagonal{T}, b::AbstractArray{T}, v::Vector{T}) where {T <: AbstractFloat}
    @assert length(b) == this.k "Dimension of vector b does match the number of rows in k. Expect $(this.k) but get $(length(b))"
    @assert length(v) == this.k "Dimension of mutable vector v deosn't match the number of rows in k. Expect $(this.k) but get $(length(v))"
    β = this.betas
    α = this.alphas
    v[1] = α[1]*b[1] + β[1]*b[2]
    for Idx in 2: this.k - 1
        v[Idx] = β[Idx - 1]*b[Idx - 1] + α[Idx]*b[Idx] + β[Idx]*b[Idx + 1]
    end
    v[this.k] = β[this.k - 1]*b[end - 1] + α[this.k]*b[end]
    return v
end

# BASIC TESTING ----------------------------------------------------------------
using LinearAlgebra, Logging, Plots
@info "Basic Testing"
n = 64
A = SymTridiagonal(fill(-2.0, n), fill(1.0,n - 1))
T_dy = DynamicSymTridiagonal(A[1, 1])
InitializeEigenSystem(T_dy, n)
for Idx in 2: n
    T_dy(A[Idx, Idx], A[Idx - 1, Idx])
    @time UpdateEigenSystemPowItr(T_dy)
    # display(T_dy.V)
end
@info "Dynamic T Eigen System"
display(T_dy.V)
L = GetL(T_dy)
U = GetU(T_dy)
b = rand(n)
@assert norm(GetT(T_dy) - L*U) < 1e-10
@assert norm(T_dy*b - GetT(T_dy)*b) < 1e-10


function solve_modelARCDiag(H, g, nlp_stop, PData::PDataFact, α::T) where {T}
    # Solve the ARC subproblem once diagonalized into Δ
    # printstyled("On est dans solve_modelARCDiag ⇊ \n", color = :red)
    # Setup the problem
    M = ones(T, size(PData.Δ))

    ϵ1 = sqrt(eps(T)) ./ T(100.0)
    ϵ = T.(ϵ1 .* (1.0 .+ PData.λ))
    λ = max(ϵ, PData.λ .+ ϵ)

    d̃ = -(PData.Δ .+ λ .* M) .\ PData.g̃
    normd̃ = sqrt(d̃ ⋅ d̃)
    seuil = normd̃ ./ α

    # Solve the subproblem (Δ + λ I) d̃ = -g̃ such that λ = ||d̃||/α
    d̃, λ = solve_diag(λ, PData.Δ, PData.g̃, seuil, α, ϵ)

    PData.λ = λ
    # Transform back d̃ into d

    d = AInv(PData, d̃)

    return d, λ
end


#!/usr/bin/env julia

Springpot =  RheoModelClass(
        # Model name
        name="springpot",
        # Model parameters,
        p = [:cᵦ, :β],
        # Relaxation modulus
        G = quote
              cᵦ*t^(-β)/gamma(1 - β)
            end,
        # Creep modulus
        J = quote
              (t^β)/(cᵦ*gamma(1 + β))
            end,
        # Storage modulus
        Gp = quote
                cᵦ*(ω^β)*cos(π*β/2)
             end,
        # Loss modulus
        Gpp = quote
                cᵦ*(ω^β)*sin(π*β/2)
              end,
        # Constraints
        constraint = quote
                (0<β<1)
                end,
        # Network
        info= "
                ____ ╱╲ ____
                     ╲╱  cᵦ, β
                "
        )

Spring =  RheoModelClass(
        # Model name
        name="spring",
        # Model parameters,
        p = [:k],
        # Relaxation modulus
        G = quote
              k
            end,
        # Creep modulus
        J = quote
              1/k
            end,
        # Storage modulus
        Gp = quote
                k
             end,
        # Loss modulus
        Gpp = quote
                0.0
              end,
        # Network
        info= "
                ___╱╲  ╱╲  ╱╲  ________
                     ╲╱  ╲╱  ╲╱  k
                "
        )

Dashpot =  RheoModelClass(
        # Model name
        name="dashpot",
        # Model parameters,
        p = [:η],
        # Creep modulus
        J = quote
              t/η
            end,
        # Storage modulus
        Gp = quote
                0.0
             end,
        # Loss modulus
        Gpp = quote
                η*ω
              end,
        # Network
        info= "
                 ___
             _____| |_____
                 _|_|
                     η
                "
        )


@define_scenario simple_scenario = begin
    t_xi::Int64
    s_xi::Int64
    @zero begin
        return simple_scenario(0, 0)
    end
    @expectation begin
        t_xi = Int(floor(sum([probability(s)*s.t_xi for s in scenarios])))
        s_xi = Int(floor(sum([probability(s)*s.s_xi for s in scenarios])))
        return simple_scenario(t_xi, s_xi)
    end
end

@sampler simple_sampler = begin
    timesteps::UnitRange{Int64}
    #n::Int64
    simple_sampler(timesteps::UnitRange{Int64}) = new(timesteps)

    @sample simple_scenario begin
        @parameters timesteps
        return simple_scenario(rand(timesteps[1:end-2]), rand([-1, 1]))
    end
end
## Try adding F to the scenario
@define_scenario F_scenario = begin
    t_xi::Int64
    s_xi::Int64
    F_xi::Float64
    @zero begin
        return F_scenario(0, 0, 0.)
    end
    @expectation begin
        t_xi = Int(floor(sum([probability(s)*s.t_xi for s in scenarios])))
        s_xi = Int(floor(sum([probability(s)*s.s_xi for s in scenarios])))
        F_xi = mean([probability(s)*s.F_xi for s in scenarios])
        return F_scenario(t_xi, s_xi, F_xi)
    end
end

@sampler F_sampler = begin
    timesteps::UnitRange{Int64}
    F_range::Tuple{Float64, Float64}
    #n::Int64
    F_sampler(timesteps::UnitRange{Int64}, F_range::Tuple{Float64, Float64}) = new(timesteps, F_range)

    @sample F_scenario begin
        @parameters timesteps F_range
        return F_scenario(rand(timesteps[1:end-2]), rand([-1, 1]), F_range[1]+rand()*F_range[2])
    end
end


using NiLang, NiLang.AD
using Test, Random
using ForwardDiff
using NiLangCore: default_constructor

@testset "basic instructions, ULogarithmic" begin
    x, y = default_constructor(ULogarithmic{Int}, 1),
    default_constructor(ULogarithmic{Int}, 2)
	@instr x *= y
    @test x == default_constructor(ULogarithmic{Int}, 3)
    @test y == default_constructor(ULogarithmic{Int}, 2)

	@test PlusEq(gaussian_log)(1.0, 2.0) == (1.0+log(1+exp(2.0)), 2.0)

    x, y, z = default_constructor(ULogarithmic{Float64}, 7.0),
    default_constructor(ULogarithmic{Float64}, 2.0),
    default_constructor(ULogarithmic{Float64}, 3.0)
	@instr x *= y + z
	@test check_inv(MulEq(+), (x, y, z))
	@test x.log ≈ log(exp(7.0) * (exp(2.0) + exp(3.0)))
    x, y, z = default_constructor(ULogarithmic{Float64}, 7.0),
    default_constructor(ULogarithmic{Float64}, 5.0),
    default_constructor(ULogarithmic{Float64}, 3.0)
	@instr x *= y - z
	@test x.log ≈ log(exp(7.0) * (exp(5.0) - exp(3.0)))

    x, y, z = default_constructor(ULogarithmic{Float64}, 7.0),
    default_constructor(ULogarithmic{Float64}, 5.0),
    default_constructor(ULogarithmic{Float64}, 3.0)
	@instr x *= y^3.4
	@test x.log ≈ log(exp(5.0)^3.4 * exp(7.0))

    x, y, z = default_constructor(ULogarithmic{Float64}, 7.0),
    default_constructor(ULogarithmic{Float64}, 5.0),
    default_constructor(ULogarithmic{Float64}, 3.0)
	@instr x *= 3
	@test x.log ≈ log(exp(7.0) * 3)
end

@testset "error on += and -=" begin
    @i function f(x::ULogarithmic)
        x += ULogarithmic(3.0)
    end

    @test_throws MethodError f(ULogarithmic(2.0))
    @test_throws MethodError (~f)(ULogarithmic(2.0))
end


###############################################
#FEniCS tutorial demo program: Convection-diffusion-reaction for a system
#describing the concentration of three species A, B, C undergoing a simple
#first-order reaction A + B --> C with first-order decay of C. The velocity
#is given by the flow field w from the demo navier_stokes_cylinder.py.
#  u_1' + w . nabla(u_1) - div(eps*grad(u_1)) = f_1 - K*u_1*u_2
#  u_2' + w . nabla(u_2) - div(eps*grad(u_2)) = f_2 - K*u_1*u_2
#  u_3' + w . nabla(u_3) - div(eps*grad(u_3)) = f_3 + K*u_1*u_2 - K*u_3
###############################################

module ft09

using FenicsPy

T = 5.0            # final time
num_steps = 500    # number of time steps
Δt = T / num_steps # time step size
eps = 0.01         # diffusion coefficient
K = 10.0           # reaction rate

# Read mesh from file
mesh = Mesh("navier_stokes_cylinder/cylinder.xml.gz")

# Define function space for velocity
W = VectorFunctionSpace(mesh, "P", 2)

# Define function space for system of concentrations
P1 = FiniteElement("P", triangle, 1)
element = MixedElement([P1, P1, P1])
V = FunctionSpace(mesh, element)

# Define test functions
v_1, v_2, v_3 = TestFunctions(V)

# Define functions for velocity and concentrations
w = FeFunction(W)
u = FeFunction(V)
u_n = FeFunction(V)

# Split system functions to access components
u_1, u_2, u_3 = split(u)
u_n1, u_n2, u_n3 = split(u_n)

# Define source terms
f_1 = Expression("pow(x[0]-0.1,2)+pow(x[1]-0.1,2)<0.05*0.05 ? 0.1 : 0",
                 degree=1)
f_2 = Expression("pow(x[0]-0.1,2)+pow(x[1]-0.3,2)<0.05*0.05 ? 0.1 : 0",
                 degree=1)
f_3 = Constant(0)

# Define expressions used in variational forms
K = Constant(K)
eps = Constant(eps)

# Define variational problem
F = ((u_1 - u_n1) / Δt)*v_1*dx + dot(w, grad(u_1))*v_1*dx +
  eps*dot(grad(u_1), grad(v_1))*dx + K*u_1*u_2*v_1*dx  +
  ((u_2 - u_n2) / Δt)*v_2*dx + dot(w, grad(u_2))*v_2*dx +
  eps*dot(grad(u_2), grad(v_2))*dx + K*u_1*u_2*v_2*dx  +
  ((u_3 - u_n3) / Δt)*v_3*dx + dot(w, grad(u_3))*v_3*dx +
  eps*dot(grad(u_3), grad(v_3))*dx - K*u_1*u_2*v_3*dx + K*u_3*v_3*dx -
  f_1*v_1*dx - f_2*v_2*dx - f_3*v_3*dx

# Create time series for reading velocity data
timeseries_w = TimeSeries("navier_stokes_cylinder/velocity_series")

# Create VTK files for visualization output
vtkfile_u_1 = File("reaction_system/u_1.pvd")
vtkfile_u_2 = File("reaction_system/u_2.pvd")
vtkfile_u_3 = File("reaction_system/u_3.pvd")

# Create progress bar
#progress = Progress("Time-stepping")
#set_log_level(PROGRESS)

# Time-stepping
t = 0
for i = 1:num_steps

    # Update current time
    global t += Δt

    # Read velocity from file
    timeseries_w.retrieve(w.vector(), t)

    # Solve variational problem for time step
    solve(F == 0, u)

    # Save solution to file (VTK)
    _u_1, _u_2, _u_3 = u.split()
    vtkfile_u_1 << (_u_1, t)
    vtkfile_u_2 << (_u_2, t)
    vtkfile_u_3 << (_u_3, t)

    # Update previous solution
    u_n.assign(u)

    # Update progress bar
    #progress.update(t / T)
end

end # module ft09


function init_observable_mappings!(m::SmetsWouters)

    observables = OrderedDict{Symbol,Observable}()
    population_mnemonic = get(get_setting(m, :population_mnemonic))

    ############################################################################
    ## 1. GDP
    ############################################################################
    gdp_fwd_transform = function (levels)
        # FROM: Level of nominal GDP (FRED :GDP series)
        # TO:   Quarter-to-quarter percent change of real, per-capita GDP, adjusted for population smoothing

        levels[!,:temp] = percapita(m, :GDP, levels)
        gdp = 1000 * nominal_to_real(:temp, levels)
        oneqtrpctchange(gdp)
    end

    gdp_rev_transform = loggrowthtopct_annualized_percapita

    observables[:obs_gdp] = Observable(:obs_gdp, [:GDP__FRED, population_mnemonic, :GDPDEF__FRED],
                                       gdp_fwd_transform, gdp_rev_transform,
                                       "Real GDP Growth", "Real GDP Growth Per Capita")

    ############################################################################
    ## 2. Hours per-capita
    ############################################################################

    hrs_fwd_transform = function (levels)
        # FROM: Average weekly hours (AWHNONAG) & civilian employment (CE16OV)
        # TO:   log (3 * per-capita weekly hours / 100)
        # Note: Not sure why the 3 is there.

        levels[!,:temp] = levels[!,:AWHNONAG] .* levels[!,:CE16OV]
        weeklyhours = percapita(m, :temp, levels)
        100*log.(3 * weeklyhours / 100)
    end

    hrs_rev_transform = logleveltopct_annualized_percapita

    observables[:obs_hours] = Observable(:obs_hours, [:AWHNONAG__FRED, :CE16OV__FRED],
                                         hrs_fwd_transform, hrs_rev_transform,
                                         "Hours Per Capita", "Log Hours Per Capita")

    ############################################################################
    ## 3. Wages
    ############################################################################

    wages_fwd_transform = function (levels)
        # FROM: Nominal compensation per hour (:COMPNFB from FRED)
        # TO: quarter to quarter percent change of real compensation (using GDP deflator)

        oneqtrpctchange(nominal_to_real(:COMPNFB, levels))
    end

    wages_rev_transform = loggrowthtopct_annualized

    observables[:obs_wages] = Observable(:obs_wages, [:COMPNFB__FRED, :GDPDEF__FRED],
                                         wages_fwd_transform, wages_rev_transform,
                                         "Percent Change in Wages",
                                         "Q-to-Q Percent Change of Real Compensation (using GDP deflator)")

    ############################################################################
    ## 4. GDP Deflator
    ############################################################################

    gdpdeflator_fwd_transform = function (levels)
        # FROM: GDP deflator (index)
        # TO:   Approximate quarter-to-quarter percent change of gdp deflator,
        #       i.e.  quarterly gdp deflator inflation

        oneqtrpctchange(levels[!,:GDPDEF])
    end

    gdpdeflator_rev_transform = loggrowthtopct_annualized

    observables[:obs_gdpdeflator] = Observable(:obs_gdpdeflator, [:GDPDEF__FRED],
                                               gdpdeflator_fwd_transform, gdpdeflator_rev_transform,
                                               "GDP Deflator",
                                               "Q-to-Q Percent Change of GDP Deflator")

    ############################################################################
    ## 5. Nominal short-term interest rate (3 months)
    ############################################################################

    nominalrate_fwd_transform = function (levels)
        # FROM: Nominal effective federal funds rate (aggregate daily data at a
        #       quarterly frequency at an annual rate)
        # TO:   Nominal effective fed funds rate, at a quarterly rate

        annualtoquarter(levels[!,:DFF])
    end

    nominalrate_rev_transform = quartertoannual

    observables[:obs_nominalrate] = Observable(:obs_nominalrate, [:DFF__FRED],
                                               nominalrate_fwd_transform, nominalrate_rev_transform,
                                               "Nominal FFR",
                                               "Nominal Effective Fed Funds Rate")

    ############################################################################
    ## 6. Consumption
    ############################################################################

    consumption_fwd_transform = function (levels)
        # FROM: Nominal consumption
        # TO:   Real consumption, approximate quarter-to-quarter percent change,
        #       per capita, adjusted for population filtering

        levels[!,:temp] = percapita(m, :PCE, levels)
        cons = 1000 * nominal_to_real(:temp, levels)
        oneqtrpctchange(cons)
    end

    consumption_rev_transform = loggrowthtopct_annualized_percapita

    observables[:obs_consumption] = Observable(:obs_consumption, [:PCE__FRED, population_mnemonic],
                                               consumption_fwd_transform, consumption_rev_transform,
                                               "Consumption growth per capita",
                                               "Consumption growth adjusted for population filtering")

    ############################################################################
    ## 7. Investment growth per capita
    ############################################################################

    investment_fwd_transform = function (levels)
        # FROM: Nominal investment
        # INTO: Real investment, approximate quarter-to-quarter percent change,
        #       per capita, adjusted for population filtering

        levels[!,:temp] = percapita(m, :FPI, levels)
        inv = 10000 * nominal_to_real(:temp, levels)
        oneqtrpctchange(inv)
    end

    investment_rev_transform  = loggrowthtopct_annualized_percapita

    observables[:obs_investment] = Observable(:obs_investment, [:FPI__FRED, population_mnemonic],
                                              investment_fwd_transform, investment_rev_transform,
                                              "Real Investment per capita",
                                              "Real investment per capita, adjusted for population filtering")

    ############################################################################
    # 8 - (8 + n_anticipated_shocks)
    ############################################################################

    for i = 1:n_anticipated_shocks(m)
        # FROM: OIS expectations of $i-period-ahead interest rates at a quarterly rate
        # TO:   Same

        ant_fwd_transform = function (levels)
            levels[!, Symbol("ant$i")]
        end

        ant_rev_transform = quartertoannual

        observables[Symbol("obs_nominalrate$i")] = Observable(Symbol("obs_ant$i"), [Symbol("ant$(i)__OIS")],
                                                      ant_fwd_transform, ant_rev_transform,
                                                      "Anticipated Shock $i",
                                                      "$i-period ahead anticipated monetary policy shock")
    end

    m.observable_mappings = observables
end


using Cubature
using BasicInterpolators: CubicSplineInterpolator, NoBoundaries
using Base.Threads: @threads
using Test

using ClearSky
using ClearSky: 𝓇stream, dIdω, P2ω

const R = 8.31446262

##

totalopticaldepth(σ, g, μ, Pₛ, Pₜ=0) = dτdP(σ, g, μ)*(Pₛ - Pₜ)

function analyticaloutgoing(σ, g, μ, cₚ, Pₛ, Tₛ)
    #whole atmosphere optical depth, to a pressure of 1 Pa
    τ∞ = totalopticaldepth(σ, g, μ, Pₛ)
    #adiabat exponent
    γ = R/(μ*cₚ)
    #blackbody power
    P = stefanboltzmann(Tₛ)
    #analytical OLR solution (Principles of Planetary Climate equation 4.32)
    P*(exp(-τ∞) + τ∞^(-4γ)*hquadrature(τ′->exp(-τ′)*τ′^(4γ), 0, τ∞)[1])
end

function numericaloutgoing(σ, g, μ, cₚ, Pₛ, Tₛ, Pₜ=1e-3, tol=1e-9)
    #nice big, wide wavenumber sample
    ν = [logrange(1e-6, 1e5, 10000, 4); 1e6]
    #absorber wrapper
    U = UnifiedAbsorber(GrayGas(σ, ν))
    #dry adiabat, whole atmosphere
    Γ = DryAdiabat(Tₛ, Pₛ, cₚ, μ)
    # 1/cos(θ)
    m = 1.0
    #mean molar mass
    fμ(T,P) = μ
    #transformed coords
    ω₁, ω₂ = P2ω(Pₛ), P2ω(Pₜ)
    #compute OLR
    olr = zeros(length(ν))
    @threads for i ∈ eachindex(ν)
        I₀ = planck(ν[i], Tₛ)
        olr[i] = π*stream(dIdω, I₀, ω₁, ω₂, U, i, g, m, Γ, fμ, tol)
    end
    #make an interpolator for smooth integration
    ϕ = CubicSplineInterpolator([0.0; ν], [0.0; olr], NoBoundaries())
    #integrate it
    hquadrature(ν->ϕ(ν), 0.0, maximum(ν))[1]
end

##

#define atmospheric parameters
g = 10 #gravity [m/s^2]
μ = 0.01 #molar mass [kg/molecule]
cₚ = 1e3 #heat capacity [J/kg/K]
Pₛ = 1e5 #surface pressure [Pa]
Tₛ = 300 #surface temperature [K]
σ = 10 .^ range(-29, -23, length=10) #gray gas absorption coefs [cm^2/molecule]

#test accuracy with different optical depths
abserr = similar(σ)
relerr = similar(σ)
τ = similar(σ)
for (i,σᵢ) ∈ enumerate(σ)
    τ[i] = totalopticaldepth(σᵢ, g, μ, Pₛ)
    OLRₐ = analyticaloutgoing(σᵢ, g, μ, cₚ, Pₛ, Tₛ)
    OLRₙ = numericaloutgoing(σᵢ, g, μ, cₚ, Pₛ, Tₛ, 1e-6)
    err = OLRₙ - OLRₐ
    abserr[i] = abs(err)
    relerr[i] = abserr[i] / OLRₐ
    @test relerr[i] < 0.01
end

##

"""using Plots

p = plot(τ, relerr,
    xaxis=:log,
    yaxis=:log,
    xlabel="Whole Atmosphere Optical Depth",
    ylabel="OLR [W/m^2] Relative Error",
    legend=false)"""


P = [ 0.5000000   0.0000000   0.0000000   0.0000000   0.0000000   0.0000000 0.0000000;
      0.0000000   1.0593543  -0.2507533   0.0000000   0.0000000   0.0000000 0.0000000;
      0.0000000  -0.2507533   1.0593543   0.0000000   0.0000000   0.0000000 0.0000000;
      0.0000000   0.0000000   0.0000000   0.6666667   0.0000000   0.0000000 0.0000000;
      0.0000000   0.0000000   0.0000000   0.0000000   0.6666667   0.0000000 0.0000000;
      0.0000000   0.0000000   0.0000000   0.0000000   0.0000000   0.6666667 0.0000000;
      0.0000000   0.0000000   0.0000000   0.0000000   0.0000000   0.0000000 0.5000000 ]

basSet = readBasisSetTX93("STO-3G.tx93")
geo = readGeometryXYZ("h2o.xyz")
bas = computeBasis(basSet,geo)

function computeMatrixCoulomb(basis::GaussianBasis, density::Matrix)
  N = length(basis.contractedBFs)

  J = zeros(N,N)
  μIndex = 0
  for (μ in basis.contractedBFs)
    μIndex+=1
    for (νIndex in 1:μIndex)
      ν = basis.contractedBFs[νIndex]
      ERImatrix = [computeElectronRepulsionIntegral(μ,ν,λ,σ) for λ in basis.contractedBFs, σ in basis.contractedBFs]
      J[μIndex,νIndex] = trace(ERImatrix*density)
      J[νIndex,μIndex] = J[μIndex,νIndex]
    end
  end
  return J
end

function computeMatrixExchange(basis::GaussianBasis, density::Matrix)
  N = length(basis.contractedBFs)

  K = zeros(N,N)
  μIndex = 0
  for (μ in basis.contractedBFs)
    μIndex+=1
    for (νIndex in 1:μIndex)
      ν = basis.contractedBFs[νIndex]
      ERImatrix = [computeElectronRepulsionIntegral(μ,λ,ν,σ) for λ in basis.contractedBFs, σ in basis.contractedBFs]
      K[μIndex,νIndex] = trace(ERImatrix*density)
      K[νIndex,μIndex] = K[μIndex,νIndex]
    end
  end
  return K
end

function computeEnergyInteratomicRepulsion(
  atom1::Atom,
  atom2::Atom)
  q1 = atom1.element.atomicNumber
  q2 = atom2.element.atomicNumber
  A = atom1.position
  B = atom2.position
  return q1*q2/(distance(A,B)^2)
end

function computeEnergyInteratomicRepulsion(geo::Geometry)
  result = 0.::Float64
  for (A in 1:length(geo.atoms))
    for (B in 1:A-1)
      result += computeEnergyInteratomicRepulsion(geo.atoms[A],geo.atoms[B])
    end
  end
  return result
end

function computeEnergyHartreeFock(
  basis::GaussianBasis,
  geometry::Geometry,
  density::Matrix)
  P = density
  T = computeMatrixKinetic(basis)
  V = computeMatrixNuclearAttraction(basis,geometry)
  J = computeMatrixCoulomb(basis,density)
  K = computeMatrixExchange(basis,density)
  G = 2J-K
  return trace((T+V+1/2*G)*P)
end

function computeMatrixFock(
  basis::GaussianBasis,
  geometry::Geometry,
  density::Matrix)
  T = computeMatrixKinetic(basis)
  V = computeMatrixNuclearAttraction(basis,geometry)
  J = computeMatrixCoulomb(basis,density)
  K = computeMatrixExchange(basis,density)
  return T+V+2J-K
end

function evaluateSCFStep(
  basis::GaussianBasis,
  geometry::Geometry,
  initialGuessDensity::Matrix,
  overlap::Matrix,
  electronNumber::Integer)

  N = electronNumber
  S = overlap
  F = computeMatrixFock(basis,geometry,initialGuessDensity)
  eig = eigfact(F,S)
  C = eig.vectors
  moEnergies = eig.values
  Cocc = eig.vectors[:,1:N]
  P = Cocc * Cocc.'

  return (moEnergies,P)
end

function evaluateSCF(
  basis::GaussianBasis,
  geometry::Geometry,
  initialGuessDensity::Matrix,
  electronNumber::Integer)

  N = electronNumber
  normalize!(basis)
  S = computeMatrixOverlap(basis)
  P = initialGuessDensity

  for (i in 1:30)
    energies,P = evaluateSCFStep(basis,geometry,P,S,N)
    println("E[$i]: $(computeEnergyHartreeFock(basis,geometry,P))")
  end
end

h2o = Geometry([Atom(Element("O"),Position(0.0, 0.0, 0.04851804)),
		Atom(Element("H"),Position(0.75300223,0.00000000,-0.51923377)),
		Atom(Element("H"),Position(-0.75300223,0.00000000,-0.51923377))])
angstrom2bohr!(h2o)

bfs = computeBasis(basSet,h2o)
normalize!(bfs)
bfs.contractedBFs[3],bfs.contractedBFs[5] = bfs.contractedBFs[5],bfs.contractedBFs[3]

function prettyprint(cgbf::ContractedGaussianBasisFunction,indent="")
  print(indent); dump(typeof(cgbf))
  for idx in 1:length(cgbf.coefficients)
    print(indent * "  $(cgbf.coefficients[idx]) × "); dump(typeof(cgbf.primitiveBFs[idx]))
    print(indent * "    center:   "); println(cgbf.primitiveBFs[idx].center)
    print(indent * "    exponent: "); println(cgbf.primitiveBFs[idx].exponent)
    print(indent * "    mqn:      "); println(cgbf.primitiveBFs[idx].mqn)
  end
end

function prettyprint(basis::GaussianBasis)
  dump(typeof(basis))
  indent = "  "

  println(indent * "contractedBFs: ")
  indent = indent * "  "

  for cgbf in basis.contractedBFs
    prettyprint(cgbf,indent)
  end

end


@testset "Array quantizer" begin

m = 2
k = 2
Q = QuantizedArrays.OrthogonalQuantization
for T in [UInt8, UInt16, UInt32, UInt64, Int8, Int16, Int32, Int64,
          Float16, Float32, Float64]
    for mᵢ in 1:2
        X = rand(T[0,1,2,3], m, 10)  # values in 0-3 to keep distances small
        aqm = build_quantizer(X, k=k, m=mᵢ, method=:sample)
        @test aqm isa ArrayQuantizer{Q, UInt8, Distances.SqEuclidean, T, 2}
        @test length(aqm.codebooks) == mᵢ
        qm = QuantizedArrays.quantize_data(aqm, X)
        @test qm isa Matrix{UInt8}
        @test size(qm) == (mᵢ, size(X,2))

        Xᵥ= rand(T[0,1,2,3,4], 10)
        aqv = build_quantizer(Xᵥ, k=k, m=mᵢ, method=:sample)
        @test aqv isa ArrayQuantizer{Q, UInt8, Distances.SqEuclidean, T, 1}
        @test length(aqv.codebooks) == 1
        qv = QuantizedArrays.quantize_data(aqv, Xᵥ)
        @test qv isa Matrix{UInt8}
        @test length(qv) == length(Xᵥ)
    end
end

T = Float32
Twrong = Float16
X = rand(T, m, 100)
m = 2
k = 2
aq = build_quantizer(X, k=k, m=m, method=:sample)
@test QuantizedArrays.codebooks(aq) == aq.codebooks

@test_throws AssertionError QuantizedArrays.quantize_data(aq, rand(T, m+1, 100))
@test_throws MethodError QuantizedArrays.quantize_data(aq, rand(Twrong, m, 100))

# The ArrayQuantizer works with any `m`:
X2 = rand(T, m+1, 100)
aq2 = ArrayQuantizer(X2, k=k, m=m, method=:sample)
@test aq2 isa ArrayQuantizer{Q, UInt8, Distances.SqEuclidean, T, 2}
@test length(aq2.codebooks) == m

aq3 = ArrayQuantizer(X, k=k, m=m+100, method=:sample)
@test aq3 isa ArrayQuantizer{Q, UInt8, Distances.SqEuclidean, T, 2}
@test length(aq3.codebooks) == m

end


date = Dates.DateTime(2020, 2, 20, 11, 11, 0)
lon, lat = [FT(80.0), FT(20.0)]

sza, azi, d = instantaneous_zenith_angle(
    date,
    lon,
    lat,
    param_set,
    eot_correction=false,
    milankovitch=false)

F = insolation(sza, d, param_set)
@test typeof(sza) == FT
@test typeof(azi) == FT
@test typeof(d) == FT
@test typeof(F) == FT

sza, azi, d = instantaneous_zenith_angle(
    date,
    lon,
    lat,
    param_set,
    eot_correction=false,
    milankovitch=true)

F = insolation(sza, d, param_set)
@test typeof(sza) == FT
@test typeof(azi) == FT
@test typeof(d) == FT
@test typeof(F) == FT

sza, azi, d = instantaneous_zenith_angle(
    date,
    lon,
    lat,
    param_set,
    eot_correction=true,
    milankovitch=false)

F = insolation(sza, d, param_set)
@test typeof(sza) == FT
@test typeof(azi) == FT
@test typeof(d) == FT
@test typeof(F) == FT

sza, azi, d = instantaneous_zenith_angle(
    date,
    lon,
    lat,
    param_set,
    eot_correction=true,
    milankovitch=true)

F = insolation(sza, d, param_set)
@test typeof(sza) == FT
@test typeof(azi) == FT
@test typeof(d) == FT
@test typeof(F) == FT


# For the conversion to go through ForwardDiff, I need
#Base.convert(FDT::Type{ForwardDiff.Dual{T,V,N}}, p::AIBECS.Parameters) where {T,V,N} = AIBECS.Parameters(convert(Vector{FDT}, vec(p))...)

@testset "Derivatives" begin
    nt = length(T_all)
    n = nt * nb
    @unpack xgeo = p
    x = xgeo * ones(n)
    testp = AIBECS.reconstruct(typeof(p), 2vec(p)) # other value for testing p
    @testset "∇ₓF" begin
        @test ForwardDiff.jacobian(x -> F(x, p), x) ≈ ∇ₓF(x, p) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> F(x, p), 2x) ≈ ∇ₓF(2x, p) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> F(x, p), -x) ≈ ∇ₓF(-x, p) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> F(x, testp), x) ≈ ∇ₓF(x, testp) rtol = 1e-14
    end
    @testset "∇ₚf" begin
        fAD = generate_objective(ωs, μx, σ²x, v, ωp, typeof(p))
        @test fAD(x,p) ≈ f(x,p)
        @test ForwardDiff.jacobian(p -> [fAD(x, p)], p) ≈ ∇ₚf(x, p) rtol = 1e-14
        @test ForwardDiff.jacobian(p -> [fAD(x, p)], testp) ≈ ∇ₚf(x, testp) rtol = 1e-14
        @test ForwardDiff.jacobian(p -> [fAD(2x, p)], p) ≈ ∇ₚf(2x, p) rtol = 1e-14
    end
    @testset "∇ₓf" begin
        @test ForwardDiff.jacobian(x -> [f(x, p)], x) ≈ ∇ₓf(x, p) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> [f(x, testp)], x) ≈ ∇ₓf(x, testp) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> [f(x, p)], 2x) ≈ ∇ₓf(2x, p) rtol = 1e-14
    end
end


using QMTK
using Compat.Test
using Compat.Distributed

@testset "Metropolis-Hasting" begin

space = RealSpace(min=-4, max=4)
sampler = MHSampler(space)

@everywhere function normal(x::T) where T <: Real
    factor = 1 / sqrt(2 * pi)
    return factor * exp(- x^2/2)
end

plan = SamplePlan(normal, sampler)
data = sample!(plan)

@test length(data) == 1000

end


##### Generic methods #####

"""
    size(d::MatrixDistribution)

Return the size of each sample from distribution `d`.
"""
Base.size(d::MatrixDistribution)

size(d::MatrixDistribution, i) = size(d)[i]

"""
    length(d::MatrixDistribution)

The length (*i.e* number of elements) of each sample from the distribution `d`.
"""
Base.length(d::MatrixDistribution)

"""
    rank(d::MatrixDistribution)

The rank of each sample from the distribution `d`.
"""
LinearAlgebra.rank(d::MatrixDistribution)

"""
    inv(d::MatrixDistribution)

If known, returns a `MatrixDistribution` instance representing the
distribution of inv(X), where X is a random matrix with distribution `d`.
"""
Base.inv(d::MatrixDistribution)

"""
    mean(d::MatrixDistribution)

Return the mean matrix of `d`.
"""
mean(d::MatrixDistribution)

"""
    var(d::MatrixDistribution)

Compute the matrix of element-wise variances for distribution `d`.
"""
var(d::MatrixDistribution) = ((n, p) = size(d); [var(d, i, j) for i in 1:n, j in 1:p])

"""
    cov(d::MatrixDistribution)

Compute the covariance matrix for `vec(X)`, where `X` is a random matrix with distribution `d`.
"""
function cov(d::MatrixDistribution)
    M = length(d)
    V = zeros(partype(d), M, M)
    iter = CartesianIndices(size(d))
    for el1 = 1:M
        for el2 = 1:el1
            i, j = Tuple(iter[el1])
            k, l = Tuple(iter[el2])
            V[el1, el2] = cov(d, i, j, k, l)
        end
    end
    return V + tril(V, -1)'
end
cov(d::MatrixDistribution, ::Val{true}) = cov(d)

"""
    cov(d::MatrixDistribution, flattened = Val(false))

Compute the 4-dimensional array whose `(i, j, k, l)` element is `cov(X[i,j], X[k, l])`.
"""
function cov(d::MatrixDistribution, ::Val{false})
    n, p = size(d)
    [cov(d, i, j, k, l) for i in 1:n, j in 1:p, k in 1:n, l in 1:p]
end

# pdf & logpdf

# TODO: Remove or restrict - this causes many ambiguity errors...
_logpdf(d::MatrixDistribution, X::AbstractMatrix{<:Real}) = logkernel(d, X) + d.logc0

#  for testing
is_univariate(d::MatrixDistribution) = size(d) == (1, 1)
check_univariate(d::MatrixDistribution) = is_univariate(d) || throw(ArgumentError("not 1 x 1"))

##### Specific distributions #####

for fname in ["wishart.jl", "inversewishart.jl", "matrixnormal.jl",
              "matrixtdist.jl", "matrixbeta.jl", "matrixfdist.jl", "lkj.jl"]
    include(joinpath("matrix", fname))
end


# Script to create the test wavefunction
# Dynamics not required: imprint an exact GPE vortex core onto a Thomas-Fermi state

using SpecialFunctions, Plots, Plots.Measures, JLD2, LaTeXStrings, QuadGK, VortexDistributions
using QuantumFluidSpectra
import VortexDistributions: Λ, ψa
gr()

function new_plot(;size=(400,200))
    fs,tfs,lw = 12,8,1.5
    p = plot(legend=:topright,lw=lw,size=size,frame=:box,
    foreground_color_legend = nothing,
    xtickfontsize=tfs,ytickfontsize=tfs,xguidefontsize=fs,
    tickfont=font(fs, "Times"),
    yguidefontsize=fs,legendfontsize=tfs)
    return p
end

## square domain
L = 22
N = 512
μ = 30.0
g = 0.1
w = 1

## useful methods
V(x,y) = 0.5*w^2*(x^2 + y^2)
ψ0(x,y,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y)/μ,0.0)+im*0.0)
healinglength(x,y,μ,g) = 1/sqrt(g*abs2(ψ0(x,y,μ,g)))
R(w) = sqrt(2*μ/w^2)
Rtf = R(1)
ξ = healinglength(0,0.,μ,g)
x = LinRange(-L/2,L/2,N); y = x
X = (x,y)
n0 = abs2.(ψ0(0,0,μ,g))
ek_unit = pi*n0*Rtf^3/(2μ)

##
ψ = ψ0.(x,y',μ,g)
# heatmap(x,y,abs.(ψ),aspect_ratio=1.0)

## Imprint charge 1 vortex at the origin
pv = PointVortex(0,0.,1)
# v1 = ScalarVortex(ξ,pv)

# check ansatz core agrees exactly with TFv
v1 = ScalarVortex(Ansatz(ψa,ξ,Λ),pv)
psi = Torus(copy(ψ),x,y)
vortex!(psi,v1)
gr()
heatmap(x,y,abs.(psi.ψ),aspect_ratio=1.0)

## k grid for spectra, and reference values
kx,ky = kvecs((L,L),(N,N))
K = (kx,ky)
kR = 2*pi/Rtf
kxi = 2*pi/ξ

kmin = 0.1kR
kmax = 1.2kxi
Np = 1000
k = LinRange(kmin,kmax,Np)

εki = incompressible_spectrum(k,Psi(psi.ψ,X,K))
# calc spec and save
# ψv = psi.ψ
# psiv = Psi(ψv,X,K)
# εki = incompressible_spectrum(k,psiv)
# @save joinpath(@__DIR__,"test_psi.jld2") ψv x y εki

#@load joinpath(@__DIR__,"test_psi.jld2") ψv x y εki

## Fig 3 (a) ui power spectrum plot
pgfplotsx()
ep3a = new_plot()

# Analytic form homog. [PRX]
f(x) = x*(besseli(0,x)*besselk(1,x)-besseli(1,x)*besselk(0,x))
FΛ(x) = f(x/2/Λ)^2/x
plot!(k*ξ,FΛ.(k*ξ),line=(1,:blue,0.8),label=false)

# analytic spectra: vortex in trap
Tint(x,a,b)= x*sqrt((1-x^2)/(x^2+b^2))*besselj1(a*x)
Tv(a,b) = quadgk(x->Tint(x,a,b),0.,1.)[1]
Mv(x,y) = x*abs2(Tv(x,y))
ekin_v_a(k) = Mv(k*Rtf,ξ/Λ/Rtf)*ek_unit
plot!(k*ξ,ekin_v_a.(k)*ξ/μ,line=(4,:pink),label=false)

plot!(k*ξ,εki*ξ/μ,scale=:log10,line=(1,:red),legend=:topright,label=false,minorticks=true,grid=false)
xlims!(extrema(k*ξ)[1],ξ*kxi*1.15)
ylims!(1e-3,1e2)
xlabel!(L"k \xi")
ylabel!(L"\varepsilon^i_k(k)\xi/\mu",grid=false)
vline!([1.0],label=false,line=(1,:gray))
vline!([2π],label=false,line=(1,:gray))
vline!([kR*ξ],label=false,line=(1,:gray))
ya = 6e-2
annotate!(0.115,ya,text(L"\frac{2\pi}{R}",10))
annotate!(.83,ya,text(L"\frac{1}{\xi}",10))
annotate!(4.9,ya,text(L"\frac{2\pi}{\xi}",10))
annotate!(5,35,text(L"(a)",10))
# joinpath(@__DIR__,"test_data/central_vortex_εik.pdf") |>  savefig
# gr()
plot!()

##
# εki
# ekin_v_a.(k)
# maximum relative error as function of grid size
# @256 => 1.65979
# @512 => 0.07953
# @1024 => 0.01276
# @2048 => 0.00337
maximum(abs.((εki .- ekin_v_a.(k))./ekin_v_a.(k)))

errk = abs.((εki .- ekin_v_a.(k))./ekin_v_a.(k))
plot(k*ξ, errk,yscale=:log10)

## Two point velocity correlation [check units!]
# pgfplotsx()
ep3b = new_plot()
r = LinRange(0,2*Rtf,1000)
gi = gv(r,k,ekin_v_a.(k))
plot!(r/ξ,gi,label=false,grid=false)
vline!([Rtf/ξ],line=(1,:gray),label=false)
ylabel!(L"g_k^i(r)")
plot!(yticks=0:0.5:1.0)
xlabel!(L"r/\xi")
xlims!(extrema(r/ξ)...)
annotate!(45,.75,text(L"R",10))
annotate!(82,0.92,text(L"(b)",10))

## combined plot
l = @layout [a; b]
pc = plot(ep3a,ep3b,layout=l,size=(430,350))
# "2d_trapvtf_combined.pdf" |> savefig
plot!()


# # Tracking transient deformation of a cantilever beam: centered difference

# Source code: [`bending_wave_Ray_expl_cd_tut.jl`](bending_wave_Ray_expl_cd_tut.jl)

# ## Description

# A cantilever beam is given an initial velocity and then at time 0.0 it is
# suddenly stopped by fixing one of its ends. This sends a wave down the beam.

# The beam is modeled as a solid. Trapezoidal rule is used to integrate the
# equations of motion in time. No damping is present.

# ## Goals

# - Show how to create the discrete model for explicit dynamics.
# - Demonstrate  centered difference time stepping.

##
# ## Definitions

# Basic imports.
using LinearAlgebra
using Arpack

# This is the finite element toolkit itself.
using FinEtools

# The linear stress analysis application is implemented in this package.
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule

# Input parameters
E = 205000*phun("MPa");# Young's modulus
nu = 0.3;# Poisson ratio
rho = 7850*phun("KG*M^-3");# mass density
loss_tangent = 0.0001;
frequency = 1/0.0058;
Rayleigh_mass = 2*loss_tangent*(2*pi*frequency);
L = 200*phun("mm");
W = 4*phun("mm");
H = 8*phun("mm");
tolerance = W/500;
vmag = 0.1*phun("m")/phun("SEC");
tend = 0.013*phun("SEC");

##
# ## Create the discrete model

MR = DeforModelRed3D
fens,fes  = H8block(L,W,H, 50,1,4)

geom = NodalField(fens.xyz)
u = NodalField(zeros(size(fens.xyz,1),3)) # displacement field

nl = selectnode(fens, box=[L L -Inf Inf -Inf Inf], inflate=tolerance)
setebc!(u, nl, true, 1)
setebc!(u, nl, true, 2)
setebc!(u, nl, true, 3)
applyebc!(u)
numberdofs!(u)

corner = selectnode(fens, nearestto=[0 0 0])
cornerzdof = u.dofnums[corner[1], 3]

material = MatDeforElastIso(MR, rho, E, nu, 0.0)

femm = FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(3,2)), material)
femm = associategeometry!(femm, geom)
K = stiffness(femm, geom, u)
# Assemble the mass matrix as diagonal. The HRZ lumping technique is
# applied through the assembler of the sparse matrix.
hrzass = SysmatAssemblerSparseHRZLumpingSymm(0.0)
femm = FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(3,3)), material)
M = mass(femm, hrzass, geom, u)

# Form the damping matrix.
C = Rayleigh_mass * M

# Figure out the highest frequency in the model, and use a time step that is
# considerably larger than the period of that highest frequency.
evals, evecs = eigs(K, M; nev=1, which=:LM);
@show dt = 0.99 * 2/sqrt(evals[1]);

# The time stepping loop is protected by `let end` to avoid unpleasant surprises
# with variables getting clobbered by globals.
ts, corneruzs = let dt = dt
    # Initial displacement, velocity, and acceleration.
    U0 = gathersysvec(u)
    v = deepcopy(u)
    v.values[:, 3] .= vmag
    V0 = gathersysvec(v)
    F1 = fill(0.0, length(V0))
    U1 = fill(0.0, length(V0))
    V1 = fill(0.0, length(V0))
    A0 = fill(0.0, length(V0))
    A1 = fill(0.0, length(V0))

    # The times and displacements of the corner will be collected into two vectors
    ts = Float64[]
    corneruzs = Float64[]
    # Let us begin the time integration loop:
    t = 0.0;
    step = 0;
    while t < tend
        push!(ts, t)
        push!(corneruzs, U0[cornerzdof])
        t = t+dt;
        step = step + 1;
        (mod(step,1000)==0) && println("Step $(step): $(t)")
        # Zero out the load
        fill!(F1, 0.0);
        # Initial acceleration
        if step == 1
            A0 = M \ (F1)
        end
        # Update displacement.
        @. U1 = U0 + dt*V0 + (dt^2/2)*A0;
        # Compute updated acceleration.
        A1 .= M \ (-K*U1 + F1)
        # Update the velocities.
        @. V1 = V0 + (dt/2)*(A0 + A1)
        # Switch the temporary vectors for the next step.
        U0, U1 = U1, U0;
        V0, V1 = V1, V0;
        A0, A1 = A1, A0;
        if (t == tend) # Are we done yet?
            break;
        end
        if (t+dt > tend) # Adjust the last time step so that we exactly reach tend
            dt = tend-t;
        end
    end
    ts, corneruzs # return the collected results
end

##
# ## Plot the results

using Gnuplot

@gp "set terminal windows 4 " :-
@gp  :- ts corneruzs./phun("mm") "lw 2 lc rgb 'red' with lines title 'Displacement of the corner' "
@gp  :- "set xlabel 'Time [s]'"
@gp  :- "set ylabel 'Displacement [mm]'"

# The end.
true


# Apply variable selection using power likelihood to the CPP (collaborative perinatal project data)
include("varsel.jl")

module RunCPP
using VarSel
using PyPlot
draw() = (pause(0.001); get_current_fig_manager()[:window][:raise_]())

# Settings
n_total = 10^4  # total number of MCMC samples
n_keep = n_total  # number of MCMC samples to record
alphas = [100,500,1000,2000,Inf]  # values of alpha to use
nburn = int(n_keep/10)  # burn-in

# Variables to use
target_name = "V_BWGT"  # name of target variable
predictor_names = vec(readdlm("datasets/CPP/predictors.txt"))  # names of the predictor variables / covariates

# randomly permute order to check for possible MCMC mixing issues
#predictor_names = predictor_names[randperm(length(predictor_names))]

# Load data
data = readdlm("datasets/CPP/finaldde.csv",',')
names = vec(data[1,:])
values = data[2:end,:]
predictor_subset = [findfirst(names.==name)::Int64 for name in predictor_names]
@assert(minimum(predictor_subset)>0)
target_index = findfirst(names.==target_name)
for j in predictor_subset
    missing = vec(values[:,j].==".")
    values[missing,j] = mean(values[!missing,j])
end
println("Missing values in covariates have been replaced by the average of non-missing entries.")
missing_y = vec(values[:,target_index].==".")
y = convert(Array{Float64,1}, values[!missing_y,target_index])
x = convert(Array{Float64,2}, values[!missing_y,predictor_subset])
n = length(y)
println("Removed $(sum(missing_y)) records which were missing the target variable.")

# Preprocess data
y = (y-mean(y))./std(y)
x = (x.-mean(x,1))./std(x,1)
@assert(all(!isnan(y)))
@assert(all(!isnan(x)))
x = [ones(n) x] # append constant covariate for offset
predictor_names = ["constant",predictor_names]
p = size(x,2)

function interval(x,P)
    N = length(x)
    xsort = sort(x)
    ai = (1-P)/2
    l,u = xsort[int(floor(ai*N))], xsort[int(ceil((1-ai)*N))]
    return l,u
end

# Run sampler
nalphas = length(alphas)
use = nburn+1:n_keep
n_use = n_keep - nburn
k_posteriors = zeros(p+1,nalphas)
P_posteriors = zeros(p,nalphas)
beta_means = zeros(p,nalphas)
for (i_a,alpha) in enumerate(alphas)

    # Run sampler
    zeta = (1/n)/(1/n + 1/alpha)
    beta_r,lambda_r,keepers = VarSel.run(x,y,n_total,n_keep,zeta)
    k_r = vec(sum(beta_r.!=0, 1))
    P_nonzero = vec(mean(beta_r[:,use].!=0, 2))
    P_posteriors[:,i_a] = P_nonzero
    beta_means[:,i_a] = vec(mean(beta_r[:,use],2))

    # Compute posterior on k
    edges,counts = hist(k_r[use],-0.5:1:p+0.5)
    k_posteriors[:,i_a] = counts/n_use

    # Show traceplot of number of nonzero coefficients
    figure(1); clf(); hold(true)
    plot(k_r + 0.5*rand(length(k_r))-0.25,"k.",markersize=2)
    draw()
    savefig("cpp-traceplot-alpha=$alpha.png",dpi=150)

    # Show posterior of lambda
    figure(2); clf(); hold(true)
    plt.hist(lambda_r[use], 50)
    draw()
    savefig("cpp-lambda-alpha=$alpha.png",dpi=150)
    println("E(lambda|data) = ", mean(lambda_r[use]))

    # List of most probable nonzero coefs
    order = sortperm(P_nonzero,rev=true)

    # Show traceplots of top coefficients
    figure(3,figsize=(8,8)); clf(); hold(true)
    subplots_adjust(bottom=0.1,hspace=0.4)
    for ji = 1:8
        j = order[ji]
        subplot(8,1,ji)
        betas = vec(beta_r[j,:])
        plot(betas,"k.",markersize=2)
        xm,xM = xlim(); ym,yM = ylim()
        text((xM-xm)*0.05+xm,(yM-ym)*0.5+ym,predictor_names[j])
    end
    draw()
    savefig("cpp-traceplot-beta-alpha=$alpha.png",dpi=150)

    # Display c.d.f.s of coefficients
    figure(4,figsize=(8,8)); clf(); hold(true)
    subplots_adjust(bottom=0.1,hspace=0.4)
    for ji = 1:15
        j = order[ji]
        subplot(5,3,ji)
        betas = vec(beta_r[j,use])
        plot(sort([betas,betas]),sort([0:n_use-1,1:n_use])/n_use,"b-",markersize=0,linewidth=2)
        ylim(0,1)
        mb = max(0.01, maximum(abs(betas)))
        xlim(-mb,mb)
        xm,xM = xlim(); ym,yM = ylim()
        text((xM-xm)*0.05+xm,(yM-ym)*0.5+ym,predictor_names[j])
        l,u = interval(betas,0.95)
        plot([l,l],ylim(),"r-",markersize=0)
        plot([u,u],ylim(),"r-",markersize=0)
    end
    draw()
    savefig("cpp-beta-alpha=$alpha.png",dpi=150)

end

# Plot posteriors on k
figure(5,figsize=(8,3.2)); clf(); hold(true)
subplots_adjust(bottom=0.2)
colors = "bgyrm"
shapes = "ds^v*"
plot(0:p,VarSel.prior_k(0:p,p),"ko-",label="prior",markersize=8,linewidth=2)
for (i_a,alpha) in enumerate(alphas)
    lab = "\$\\alpha = " * (alpha<Inf ? repr(int(alpha)) : "\\infty") * "\$"
    plot(0:p,vec(k_posteriors[:,i_a]),"$(colors[i_a])$(shapes[i_a])-",label=lab,markersize=8,linewidth=2)
end
xticks(0:p)
xlim(0,16)
ylim(0,1)
xlabel("\$k \\,\\,\\mathrm{(\\#\\,of\\,nonzero\\,coefficients)}\$",fontsize=17)
ylabel("\$\\Pi(k|\\mathrm{data})\$",fontsize=15)
legend(numpoints=1)
draw()
savefig("cpp-k_posteriors.png",dpi=150)

# Sort coefs by overall probability of inclusion
order = sortperm(vec(mean(P_posteriors,2)),rev=true)
names_order = predictor_names[order]
for i in 1:16
    println(i,": ",names_order[i])
end

# Plot posterior probability of inclusion
figure(6,figsize=(8,3.2)); clf(); hold(true)
subplots_adjust(bottom=0.2)
colors = "bgyrm"
shapes = "ds^v*"
for (i_a,alpha) in enumerate(alphas)
    lab = "\$\\alpha = " * (alpha<Inf ? repr(int(alpha)) : "\\infty") * "\$"
    plot(1:p,vec(P_posteriors[order,i_a]),"$(colors[i_a])$(shapes[i_a])-",label=lab,markersize=8,linewidth=2)
end
d_show = 16
xticks(1:d_show)
xlim(0,d_show)
yticks(linspace(0,1,6))
ylim(0,1.1)
xlabel("\$j\\,\\,\\mathrm{(coefficient\\,index)}\$",fontsize=17)
ylabel("\$\\Pi(\\beta_j \\neq 0\\,|\\,\\mathrm{data})\$",fontsize=16)
legend(numpoints=1)
draw()
savefig("cpp-P_posteriors.png",dpi=150)

# Plot posterior means
figure(7,figsize=(8,3.2)); clf(); hold(true)
subplots_adjust(bottom=0.2)
colors = "bgyrm"
shapes = "ds^v*"
for (i_a,alpha) in enumerate(alphas)
    lab = "\$\\alpha = " * (alpha<Inf ? repr(int(alpha)) : "\\infty") * "\$"
    plot(1:p,vec(beta_means[order,i_a]),"$(colors[i_a])$(shapes[i_a])-",label=lab,markersize=8,linewidth=2)
end
d_show = 16
xticks(1:d_show)
xlim(0,d_show)
xlabel("\$j\\,\\,\\mathrm{(coefficient\\,index)}\$",fontsize=17)
ylabel("\$E(\\beta_j |\\mathrm{data})\$",fontsize=16)
legend(numpoints=1)
draw()
savefig("cpp-beta_means.png",dpi=150)

end # module


struct Spline{order,T,P<:AbstractVector} <: Basis{T}
    points::P
end
Spline{o,T}(pts::P) where {o,T,P} = Spline{o,T,P}(pts)

const LinearSpline = Spline{1}
const HeavisideSpline = Spline{0}

Spline{o}(pts::AbstractVector{T}) where {o,T} = Spline{o,float(T)}(pts)
Spline{o}(S::Spline) where {o} = Spline{o}(S.points)

for Typ in (:LinearSpline, :HeavisideSpline)
    STyp = string(Typ)
    @eval function summary(io::IO, L::$Typ)
        print(io, "$($STyp)(")
        print(IOContext(io, :limit=>true), L.points)
        print(io,")")
    end
end

axes(B::Spline{o}) where o =
    (Inclusion(first(B.points)..last(B.points)), OneTo(length(B.points)+o-1))
==(A::Spline{o}, B::Spline{o}) where o = A.points == B.points

function getindex(B::LinearSpline{T}, x::Number, k::Int) where T
    x ∈ axes(B,1) && 1 ≤ k ≤ size(B,2)|| throw(BoundsError())

    p = B.points
    n = length(p)

    k > 1 && x ≤ p[k-1] && return zero(T)
    k < n && x ≥ p[k+1] && return zero(T)
    x == p[k] && return one(T)
    x < p[k] && return (x-p[k-1])/(p[k]-p[k-1])
    return (x-p[k+1])/(p[k]-p[k+1]) # x ≥ p[k]
end

function getindex(B::HeavisideSpline{T}, x::Number, k::Int) where T
    x ∈ axes(B,1) && 1 ≤ k ≤ size(B,2)|| throw(BoundsError())

    p = B.points
    n = length(p)

    p[k] < x < p[k+1] && return one(T)
    p[k] == x && return one(T)/2
    p[k+1] == x && return one(T)/2
    return zero(T)
end

grid(L::HeavisideSpline) = L.points[1:end-1] .+ diff(L.points)/2
grid(L::LinearSpline) = L.points

## Sub-bases

## Mass matrix
function similar(AB::QMul2{<:QuasiAdjoint{<:Any,<:LinearSpline},<:LinearSpline}, ::Type{T}) where T
    n = size(AB,1)
    SymTridiagonal(Vector{T}(undef, n), Vector{T}(undef, n-1))
end
#
@simplify function *(Ac::QuasiAdjoint{<:Any,<:LinearSpline}, B::LinearSpline)
    M = Mul(Ac, B)
    copyto!(similar(M, eltype(M)), M)
end

function copyto!(dest::SymTridiagonal,
                 AB::QMul2{<:QuasiAdjoint{<:Any,<:LinearSpline},<:LinearSpline}) where T
    Ac,B = AB.A,AB.B
    A = parent(Ac)
    A.points == B.points || throw(ArgumentError())
    dv,ev = dest.dv,dest.ev
    x = A.points; n = length(x)
    length(dv) == n || throw(DimensionMismatch())

    dv[1] = (x[2]-x[1])/3
    @inbounds for k = 2:n-1
        dv[k] = (x[k+1]-x[k-1])/3
    end
    dv[n] = (x[n] - x[n-1])/3

    @inbounds for k = 1:n-1
        ev[k] = (x[k+1]-x[k])/6
    end

    dest
end

@simplify function *(Ac::QuasiAdjoint{<:Any,<:HeavisideSpline}, B::HeavisideSpline)
    A = parent(Ac)
    A.points == B.points || throw(ArgumentError("Cannot multiply incompatible splines"))
    Diagonal(diff(A.points))
end

## Derivative
function copyto!(dest::MulQuasiMatrix{<:Any,<:Tuple{<:HeavisideSpline,<:Any}},
                 M::QMul2{<:Derivative,<:LinearSpline})
    D, L = M.A, M.B
    H, A = dest.args
    x = H.points

    axes(dest) == axes(M) || throw(DimensionMismatch("axes must be same"))
    x == L.points || throw(ArgumentError("Cannot multiply incompatible splines"))
    bandwidths(A) == (0,1) || throw(ArgumentError("Not implemented"))

    d = diff(x)
    A[band(0)] .= inv.((-).(d))
    A[band(1)] .= inv.(d)

    dest
end

function similar(M::QMul2{<:Derivative,<:LinearSpline}, ::Type{T}) where T
    D, B = M.A, M.B
    n = size(B,2)
    ApplyQuasiMatrix(*, HeavisideSpline{T}(B.points),
        BandedMatrix{T}(undef, (n-1,n), (0,1)))
end

@simplify function *(D::Derivative, L::LinearSpline)
    M = Mul(D, L)
    copyto!(similar(M, eltype(M)), M)
end

##
# sum
##

function _sum(A::HeavisideSpline, dims)
    @assert dims == 1
    permutedims(diff(A.points))
end

function _sum(P::LinearSpline, dims)
    d = diff(P.points)
    ret = Array{float(eltype(d))}(undef, length(d)+1)
    ret[1] = d[1]/2
    for k = 2:length(d)
        ret[k] = (d[k-1] + d[k])/2
    end
    ret[end] = d[end]/2
    permutedims(ret)
end

_cumsum(H::HeavisideSpline{T}, dims) where T = LinearSpline(H.points) * tril(Ones{T}(length(H.points),length(H.points)-1) .* diff(H.points)',-1)


m13.1 <- ulam(
    alist(
        S ~ dbinom( N , p ) ,
        logit(p) <- a[tank] ,
        a[tank] ~ dnorm( 0 , 1.5 )
    ), data=dat , chains=4 , log_lik=TRUE )

## R code 13.3
m13.2 <- ulam(
    alist(
        S ~ dbinom( N , p ) ,
        logit(p) <- a[tank] ,
        a[tank] ~ dnorm( a_bar , sigma ) ,
        a_bar ~ dnorm( 0 , 1.5 ) ,
        sigma ~ dexp( 1 )
    ), data=dat , chains=4 , log_lik=TRUE )

## R code 13.13
dat <- list( Si=dsim$Si , Ni=dsim$Ni , pond=dsim$pond )
m13.3 <- ulam(
    alist(
        Si ~ dbinom( Ni , p ),
        logit(p) <- a_pond[pond],
        a_pond[pond] ~ dnorm( a_bar , sigma ),
        a_bar ~ dnorm( 0 , 1.5 ),
        sigma ~ dexp( 1 )
    ), data=dat , chains=4 )

## R code 13.21
library(rethinking)
data(chimpanzees)
d <- chimpanzees
d$treatment <- 1 + d$prosoc_left + 2*d$condition

dat_list <- list(
    pulled_left = d$pulled_left,
    actor = d$actor,
    block_id = d$block,
    treatment = as.integer(d$treatment) )

set.seed(13)
m13.4 <- ulam(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a[actor] + g[block_id] + b[treatment] ,
        b[treatment] ~ dnorm( 0 , 0.5 ),
        # adaptive priors
        a[actor] ~ dnorm( a_bar , sigma_a ),
        g[block_id] ~ dnorm( 0 , sigma_g ),
        # hyper-priors
        a_bar ~ dnorm( 0 , 1.5 ),
        sigma_a ~ dexp(1),
        sigma_g ~ dexp(1)
    ) , data=dat_list , chains=4 , cores=4 , log_lik=TRUE )

## R code 13.23
m13.5 <- ulam(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a[actor] + b[treatment] ,
        b[treatment] ~ dnorm( 0 , 0.5 ),
        a[actor] ~ dnorm( a_bar , sigma_a ),
        a_bar ~ dnorm( 0 , 1.5 ),
        sigma_a ~ dexp(1)
    ) , data=dat_list , chains=4 , cores=4 , log_lik=TRUE )

## R code 13.25
m13.6 <- ulam(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a[actor] + g[block_id] + b[treatment] ,
        b[treatment] ~ dnorm( 0 , sigma_b ),
        a[actor] ~ dnorm( a_bar , sigma_a ),
        g[block_id] ~ dnorm( 0 , sigma_g ),
        a_bar ~ dnorm( 0 , 1.5 ),
        sigma_a ~ dexp(1),
        sigma_g ~ dexp(1),
        sigma_b ~ dexp(1)
    ) , data=dat_list , chains=4 , cores=4 , log_lik=TRUE )

## R code 13.26
m13x <- ulam(
    alist(
        v ~ normal(0,3),
        x ~ normal(0,exp(v))
    ), data=list(N=1) , chains=4 )
precis(m13x)

## R code 13.27
m13y <- ulam(
    alist(
        v ~ normal(0,3),
        z ~ normal(0,1),
        gq> real[1]:x <<- z*exp(v)
    ), data=list(N=1) , chains=4 )
precis(m13y)

## R code 13.28
m13.4b <- ulam( m13.4 , chains=4 , cores=4 , control=list(adapt_delta=0.99) )
divergent(m13.4b)

## R code 13.29
m13.4nc <- ulam(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a_bar + z[actor]*sigma_a + # actor intercepts
                    x[block_id]*sigma_g +      # block intercepts
                    b[treatment] ,
        b[treatment] ~ dnorm( 0 , 0.5 ),
        z[actor] ~ dnorm( 0 , 1 ),
        x[block_id] ~ dnorm( 0 , 1 ),
        a_bar ~ dnorm( 0 , 1.5 ),
        sigma_a ~ dexp(1),
        sigma_g ~ dexp(1)
    ) , data=dat_list , chains=4 , cores=4 )


@testset "Contour" begin
    @polyvar x y
    p1 = 1 + x + y + x * y + y^2 + x^2 * y

    C = PolynomialAmoebas.contour(p1, res=(100, 100))
    @test length(C.ps) != 0
end