using Test
using Distributions
using StatsBase
using Random

using EnsembleKalmanProcesses.ParameterDistributionStorage

@testset "ParameterDistributionStorage" begin
    @testset "ParameterDistributionType" begin
        # Tests for the ParameterDistributionType
        d = Parameterized(Gamma(2.0, 0.8))
        @test d.distribution == Gamma(2.0, 0.8)

        d = Samples([1 2 3; 4 5 6])
        @test d.distribution_samples == [1 2 3; 4 5 6]
        d = Samples([1 4; 2 5; 3 6]; params_are_columns = false)
        @test d.distribution_samples == [1 2 3; 4 5 6]
        d = Samples([1, 2, 3, 4, 5, 6])
        @test d.distribution_samples == [1 2 3 4 5 6]
        d = Samples([1, 2, 3, 4, 5, 6]; params_are_columns = false)
        @test d.distribution_samples == reshape([1, 2, 3, 4, 5, 6], :, 1)

    end
    @testset "ConstraintType" begin
        # Tests for the ConstraintType
        # The predefined transforms
        c1 = bounded_below(0.2)
        @test isapprox(c1.constrained_to_unconstrained(1.0) - (log(1.0 - 0.2)), 0)
        @test isapprox(c1.unconstrained_to_constrained(0.0) - (exp(0.0) + 0.2), 0)

        c2 = bounded_above(0.2)
        @test isapprox(c2.constrained_to_unconstrained(-1.0) - (log(0.2 - -1.0)), 0)
        @test isapprox(c2.unconstrained_to_constrained(10.0) - (0.2 - exp(10.0)), 0)

        c3 = bounded(-0.1, 0.2)
        @test isapprox(c3.constrained_to_unconstrained(0.0) - (log((0.0 - -0.1) / (0.2 - 0.0))), 0)
        @test isapprox(c3.unconstrained_to_constrained(1.0) - ((0.2 * exp(1.0) + -0.1) / (exp(1.0) + 1)), 0)
        @test_throws DomainError bounded(0.2, -0.1)

        #an example with user defined invertible transforms
        c_to_u = (x -> 3 * x + 14)
        u_to_c = (x -> (x - 14) / 3)

        c4 = Constraint(c_to_u, u_to_c)
        @test isapprox(c4.constrained_to_unconstrained(5.0) - c_to_u(5.0), 0)
        @test isapprox(c4.unconstrained_to_constrained(5.0) - u_to_c(5.0), 0)

    end

    @testset "ParameterDistribution" begin

        d = Parameterized(Normal(0, 1))
        c = no_constraint()
        name = "unconstrained_normal"
        u = ParameterDistribution(d, c, name)
        @test u.distributions == [d]
        @test u.constraints == [c]
        @test u.names == [name]

        # Tests for the ParameterDistribution
        d = Parameterized(MvNormal(4, 0.1))
        c = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]

        name = "constrained_mvnormal"
        u = ParameterDistribution(d, c, name)
        @test u.distributions == [d]
        @test u.constraints == c
        @test u.names == [name]
        @test_throws DimensionMismatch ParameterDistribution(d, c[1:3], name)
        @test_throws DimensionMismatch ParameterDistribution(d, c, [name, "extra_name"])

        # Tests for the ParameterDistribution
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1.0 3.0 5.0 7.0; 9.0 11.0 13.0 15.0])
        c2 = [bounded(10, 15), no_constraint()]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        u = ParameterDistribution([d1, d2], [c1, c2], [name1, name2])
        @test u.distributions == [d1, d2]
        @test u.constraints == cat([c1, c2]..., dims = 1)
        @test u.names == [name1, name2]
    end

    @testset "getter functions" begin
        # setup for the tests:
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1 2 3 4])
        c2 = [bounded(10, 15)]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        u = ParameterDistribution([d1, d2], [c1, c2], [name1, name2])

        # Test for get_dimension(s)
        @test get_total_dimension(u1) == 4
        @test get_total_dimension(u2) == 1
        @test get_total_dimension(u) == 5
        @test get_dimensions(u1) == [4]
        @test get_dimensions(u2) == [1]
        @test get_dimensions(u) == [4, 1]

        # Tests for get_name
        @test get_name(u1) == [name1]
        @test get_name(u) == [name1, name2]

        # Tests for get_n_samples
        @test typeof(get_n_samples(u)[name1]) <: String
        @test get_n_samples(u)[name2] == 4

        # Tests for get_distribution
        @test get_distribution(d1) == MvNormal(4, 0.1)
        @test get_distribution(u1)[name1] == MvNormal(4, 0.1)
        @test typeof(get_distribution(d2)) == Array{Int64, 2}
        @test typeof(get_distribution(u2)[name2]) == Array{Int64, 2}

        d = get_distribution(u)
        @test d[name1] == MvNormal(4, 0.1)
        @test typeof(d[name2]) == Array{Int64, 2}

        # Test for get_all_constraints
        @test get_all_constraints(u) == cat([c1, c2]..., dims = 1)
    end

    @testset "statistics functions" begin

        # setup for the tests:
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1 2 3 4])
        c2 = [bounded(10, 15)]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = Parameterized(Beta(2, 2))
        c3 = [no_constraint()]
        name3 = "unconstrained_beta"
        u3 = ParameterDistribution(d3, c3, name3)

        u = ParameterDistribution([d1, d2], [c1, c2], [name1, name2])

        d4 = Samples([1 2 3 4 5 6 7 8; 8 7 6 5 4 3 2 1])
        c4 = [no_constraint(), no_constraint()]
        name4 = "constrained_MVsampled"
        v = ParameterDistribution([d1, d2, d3, d4], [c1, c2, c3, c4], [name1, name2, name3, name4])

        # Tests for sample distribution
        seed = 2020
        Random.seed!(seed)
        s1 = rand(MvNormal(4, 0.1), 1)
        Random.seed!(seed)
        @test sample_distribution(u1) == s1

        Random.seed!(seed)
        s1 = rand(MvNormal(4, 0.1), 3)
        Random.seed!(seed)
        @test sample_distribution(u1, 3) == s1

        Random.seed!(seed)
        idx = StatsBase.sample(collect(1:size(d2.distribution_samples)[2]), 1)
        s2 = d2.distribution_samples[:, idx]
        Random.seed!(seed)
        @test sample_distribution(u2) == s2

        Random.seed!(seed)
        idx = StatsBase.sample(collect(1:size(d2.distribution_samples)[2]), 3)
        s2 = d2.distribution_samples[:, idx]
        Random.seed!(seed)
        @test sample_distribution(u2, 3) == s2

        Random.seed!(seed)
        s1 = sample_distribution(u1, 3)
        s2 = sample_distribution(u2, 3)
        Random.seed!(seed)
        s = sample_distribution(u, 3)
        @test s == cat([s1, s2]..., dims = 1)

        #Test for get_logpdf
        @test_throws ErrorException get_logpdf(u, zeros(get_total_dimension(u)))
        x_in_bd = [0.5]
        Random.seed!(seed)
        lpdf3 = logpdf.(Beta(2, 2), x_in_bd)[1] #throws deprecated warning without "."
        Random.seed!(seed)
        @test isapprox(get_logpdf(u3, x_in_bd) - lpdf3, 0.0; atol = 1e-6)
        @test_throws DimensionMismatch get_logpdf(u3, [0.5, 0.5])

        #Test for get_cov, get_var
        block_cov = cat([get_cov(d1), get_var(d2), get_var(d3), get_cov(d4)]..., dims = (1, 2))
        @test isapprox(get_cov(v) - block_cov, zeros(get_total_dimension(v), get_total_dimension(v)); atol = 1e-6)
        #Test for get_mean
        means = cat([get_mean(d1), get_mean(d2), get_mean(d3), get_mean(d4)]..., dims = 1)
        @test isapprox(get_mean(v) - means, zeros(get_total_dimension(v)); atol = 1e-6)

    end

    @testset "transform functions" begin
        #setup for the tests
        tol = 1e-8
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1.0 3.0 5.0 7.0; 9.0 11.0 13.0 15.0])
        c2 = [bounded(10, 15), no_constraint()]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        u = ParameterDistribution([d1, d2], [c1, c2], [name1, name2])

        x_unbd = rand(MvNormal(6, 3), 1000)  #6 x 1000
        # Tests for transforms
        x_real_constrained1 = mapslices(x -> transform_unconstrained_to_constrained(u1, x), x_unbd[1:4, :]; dims = 1)
        @test isapprox(
            x_unbd[1:4, :] -
            mapslices(x -> transform_constrained_to_unconstrained(u1, x), x_real_constrained1; dims = 1),
            zeros(size(x_unbd[1:4, :]));
            atol = tol,
        )

        x_real_constrained2 = mapslices(x -> transform_unconstrained_to_constrained(u2, x), x_unbd[5:6, :]; dims = 1)
        @test isapprox(
            x_unbd[5:6, :] -
            mapslices(x -> transform_constrained_to_unconstrained(u2, x), x_real_constrained2; dims = 1),
            zeros(size(x_unbd[5:6, :]));
            atol = tol,
        )

        x_real = mapslices(x -> transform_unconstrained_to_constrained(u, x), x_unbd; dims = 1)
        x_unbd_tmp = mapslices(x -> transform_constrained_to_unconstrained(u, x), x_real; dims = 1)
        @test isapprox(x_unbd - x_unbd_tmp, zeros(size(x_unbd)); atol = tol)

        # Tests transforms for other input structures
        @test isapprox(transform_unconstrained_to_constrained(u1, x_unbd[1:4, :]), x_real_constrained1; atol = tol)
        @test isapprox(
            x_unbd[1:4, :] - transform_constrained_to_unconstrained(u1, x_real_constrained1),
            zeros(size(x_unbd[1:4, :]));
            atol = tol,
        )
        @test isapprox(transform_unconstrained_to_constrained(u2, x_unbd[5:6, :]), x_real_constrained2; atol = tol)
        @test isapprox(
            x_unbd[5:6, :] - transform_constrained_to_unconstrained(u2, x_real_constrained2),
            zeros(size(x_unbd[5:6, :]));
            atol = tol,
        )
        @test isapprox(
            transform_unconstrained_to_constrained(u2, [x_unbd[5:6, :], x_unbd[5:6, :]]),
            [x_real_constrained2, x_real_constrained2];
            atol = tol,
        )
    end
end


#=
mdFnn.jl:

a Julia based alternative for running multi-dimensional recurrence analysis. Based on Dan Mønster's code: https://github.com/danm0nster/mdembedding

requires:
using Pkg
Pkg.add("LinearAlgebra") # diag caculation
Pkg.add("Distances") # matrix pairwise distances
Pkg.add("ElasticArrays") # efficent matrices

using Distances, ElasticArrays, LinearAlgebra, StatsBase;
=#

function mdFnn(data, tau, maxEmb, numSamples, Rtol, Atol)
    if typeof(data) <: Dataset
        data = data |> Matrix
    end

    dims = size(data,2);
    N = size(data,1);

    fnnPerc = [100.0]; # first FNN

    if dims == 1
        Ra = sum(cov(data));
    else
        Ra = sum(diag(cov(data))); # estimate of attractor size
    end

    if (N-tau*(maxEmb-1)) < numSamples # check whether enough data points exist for random sampling
        numSamples = N-tau*(maxEmb-1);
        samps = 1:1:numSamples;
    else
        #numSamples = N-tau*(maxEmb-1);
        samps = sortperm(rand(numSamples));
    end

    embData = ElasticArray{Float64}(undef, N-tau*(maxEmb-1), 0);

    for i = 1:maxEmb # embed data
        append!(embData,data[1+(i-1)*tau:end-(maxEmb-i)*tau, 1:dims]);
        dists =  pairwise(Euclidean(),embData.^2; dims=1);
        r2d1 = Array{Float64}(undef, 0);
        yRd1 = Array{Int}(undef, 0);
        for j = 1:numSamples # get nearest neighbors and distances
            temp = sort(dists[:,samps[j]]);
            coord = sortperm(dists[:,samps[j]]);
            push!(r2d1, temp[2]);
            push!(yRd1, coord[2] |> Int);
        end
        if i == 1
            global r2d = r2d1;
            global yRd = r2d1;
        else
            fnnTemp = Array{Float64}(undef, 0);
            for j = 1:length(r2d1)
                temp = dists[:,samps[j]];
                # check whether neighbors are false
                if sqrt(abs((temp[yRd1[j]] - r2d[j])/r2d[j])) > Rtol || abs(temp[yRd1[j]] - r2d[j])/Ra > Atol
                    append!(fnnTemp, 1);
                else
                    append!(fnnTemp, 0);
                end
            end
            push!(fnnPerc, 100*sum(fnnTemp)/length(fnnTemp)); # compute percentage of FNN
            r2d = r2d1;
            yRd = r2d1;
        end
    end

    return(fnnPerc)
end


#### Hansen`s model and modified model
### By Gilberto Boaretto and Daniel Coutinho
##### License: MIT

## To run this file, make sure you have packages Plots and LaTeXStrings in the system
## The Plots are better viewed with some IDE like Jupyter or Atom.
## Otherwise, go to the repl, press ] and type add <name-of-the-package>
## Make sure that you have the gensys file

include(string(pwd(),"/src/gensys.jl")) #Change this to where you saved gensys

## A word of warning about our gensys implementation: we return the matrices from the model with a unique solution in the case of multiple solutions. This is done because the procedure is not a 100 percent debbuged and in some cases with unique solution we return a warning of no unique solution (set gamma = 0, for example)

## Parameters for the original Hansen model, following McCandless book.

theta = 0.36
delta = 0.025
beta = 0.99
A = 1.72
h0 = 0.58
gamma = 0.95
sig = 0.0712

#Hansen's model with indivisible hours

B = A*log(1-h0)/h0

r_bar = 1/beta-(1-delta) #This is the interest rate in the steady state.

## The matrices Γ_0, Γ_1, Π and Ψ for Gensys:

G0 = zeros(7,7)
G1 = zeros(7,7)
Pi = zeros(7,2)
Psi = zeros(7)

G0[1,1] = 1
G0[1,2] = -r_bar*beta

G0[3,5] = 1

G0[7,7] = 1

###G1####

G1[1,1] = 1

G1[2,1] = 1
G1[2,3] = 1
G1[2,4] = -1

G1[3,1] = -(r_bar/theta-delta)
G1[3,4] = r_bar/theta
G1[3,5] = (1-delta)

G1[4,3] = 1-theta
G1[4,4] = -1
G1[4,5] = theta
G1[4,7] = 1

G1[5,2] = -1
G1[5,4] = 1
G1[5,5] = -1

G1[6,3] = -1
G1[6,4] = 1
G1[6,6] = -1

G1[7,7] = gamma

Pi[1,1] = 1
Pi[1,2] = -beta*r_bar

Psi[7] = 1

sol1 = gensys(G0,G1,Psi,Pi)

irf1 = irf(sol1,100,0.01) #yes, this takes the whole model and knows where to look in the object of the type of the model which matrices to use. Neat!

using Plots
using LaTeXStrings

plot(irf1[:,1], w = 2, label = "C")
plot!(irf1[:,2], w = 2, label = "r")
plot!(irf1[:,3], w = 2, label = "h")
plot!(irf1[:,4], w = 2, label = "y")
plot!(irf1[:,5], w = 2, label = "k")
plot!(irf1[:,6], w = 2, label = "w", line = :dash, color = "red")
hline!([0], color = "black", w = 2, label = "0") #zero is kind of hard to see

###########################
##### Modified model######
#########################

theta = 0.36
delta = 0.025
beta = 0.99
A = 1.72
h0 = 0.58
gamma = 0.95
sig = 1#3.18 #CRRA coeficient
xi = 0 # 0.832 #habit formation

r_bar = 1/beta-(1-delta)
B = A*log(1-h0)/h0

rho = sig/((1-xi)*(1-beta*xi))
kappa = beta*(1-beta*xi)*(1-xi)/(sig*(1+beta*xi^2))*r_bar

G0 = zeros(9,9)
G1 = zeros(9,9)
Pi = zeros(9,3)
Psi = zeros(9)

###### Filling the matrices Γ_0, Γ_1, Π and Ψ for Gensys ############

########## Eq 1: Euler ##############

G0[1,1] = -beta*xi/(1+beta*xi^2)
G0[1,2] = 1+beta*xi/(1+beta*xi^2)
G0[1,4] = -kappa

G1[1,3] = -xi/(1+beta*xi^2)
G1[1,2] = 1+xi/(1+beta*xi^2)

Pi[1,1] = -beta*xi/(1+beta*xi^2)
Pi[1,2] = 1+beta*xi/(1+beta*xi^2)
Pi[1,3] = -kappa

####Eq 2 and 3: dummy equations#############

G0[2,2] = 1
G1[2,1] = 1

G0[3,3] = 1
G1[3,2] = 1

########## Eq 4: Labour supply ###############

G1[4,1] = -rho*beta*xi
G1[4,2] = rho*(1+beta*xi^2)
G1[4,3] = -rho*xi
G1[4,8] = -1

Pi[4,2] = rho*beta*xi

########  Eq 5: Production Function ###########

G1[5,5] = (1-theta)
G1[5,6] = -1
G1[5,7] = theta
G1[5,9] = 1

########## Eq 6: capital flux ####################

G0[6,7] = 1

G1[6,2] = -(r_bar/theta - delta)
G1[6,6] = r_bar/theta
G1[6,7] = (1-delta)

############### Eq 7: capital return ####################

G1[7,4] = -1
G1[7,6] = 1
G1[7,7] = -1

############## Eq 8: wage #########################

G1[8,5] = -1
G1[8,6] = 1
G1[8,8] = -1

############## Eq 9: autoregressive shock ################

G0[9,9] = 1

G1[9,9] = gamma

Psi[9] = 1

sol2 = gensys(G0,G1,Psi,Pi)

G0_aux = G0
G1_aux = G1
Psi_aux = Psi
Pi_aux = Pi

irf2 = irf(sol2,100,0.01)

plot(irf2[:,1], w = 2, label = "C")
plot!(irf2[:,2], w = 2, label = "r")
plot!(irf2[:,5], w = 2, label = "h")
plot!(irf2[:,6], w = 2, label = "y")
plot!(irf2[:,7], w = 2, label = "k")
plot!(irf2[:,8], w = 2, label = "w", line = :dash, color = "red")
hline!([0], color = "black", w = 2, label = "0")

# Sanity check: with sig = 1 and xi = 0, the model is exactly the same as before

plot(irf2[:,1], w = 2, label = "Modified Model")
plot!(irf1[:,1], w = 2, label = "Hansen's Model")
title!("c")

plot(irf2[:,4], w = 2, label = "Modified Model")
plot!(irf1[:,2], w = 2, label = "Hansen's Model")
title!("r")

plot(irf2[:,5], w = 2, label = "Modified Model")
plot!(irf1[:,3], w = 2, label = "Hansen's Model")
title!("h")

plot(irf2[:,6], w = 2, label = "Modified Model")
plot!(irf1[:,4], w = 2, label = "Hansen's Model")
title!("y")

plot(irf2[:,7], w = 2, label = "Modified Model")
plot!(irf1[:,5], w = 2, label = "Hansen's Model")
title!("k")

plot(irf2[:,8], w = 2, label = "Modified Model")
plot!(irf1[:,6], w = 2, label = "Hansen's Model")
title!("w")

plot(irf2[:,9], w = 2, label = "Modified Model")
plot!(irf1[:,7], w = 2, label = "Hansen's Model")
title!(L"\lambda")


mutable struct RegressionData
    X::Array{Float64,2}
    Y::Vector{Float64}
    w::Vector{Float64}
end

"""
    getdata(; Xdata::XData = XData(),
                         wdist::BenchmarkwDists = BinChoice(),
                         noisedist::BenchmarkNoiseDists = NoNoise(),
                         SNR::Float64 = 0.0,
                         n::Int = 100,
                         nfeatures::Int = 100,
                         sparsity::Int = 10
                )

Get regression data. Created to be stand-alone from the rest of the
benchmarking. Maybe will remove later, want user to only interface with
BenchmarkData object.
"""
function getdata(; Xdata::XData = XData(),
                         wdist::BenchmarkwDists = BinChoice(),
                         noisedist::BenchmarkNoiseDists = NoNoise(),
                         SNR::Float64 = 0.0,
                         n::Int = 100,
                         nfeatures::Int = 100,
                         sparsity::Int = 10
                        )
    bd = BenchmarkData(Xdata, wdist, noisedist, SNR, [n], nfeatures, sparsity)
    getdata(bd, 1)
end
"""
    getdata(bd::BenchmarkData, i::Int)

Get regression data, the i^th value in the range for n will be used.
"""
function getdata(bd::BenchmarkData, i::Int)
    w = getw(bd)
    X = getX(bd, i)
    # In case we are using a multivariate distribution for x
    if size(X, 2) != (bd.nfeatures)
        error("X distribution $(bd.Xdata.dist) doesn't match the number of
                features requested, $(bd.nfeatures).")
    end
    Y = X * w
    noise = getnoise(bd, Y)
    Y .+= noise
    RegressionData(X, Y, w)
end


using Test
using Latexify
using ModelingToolkit

### Tips for generating latex tests:
### Latexify has an unexported macro:
###
### Latexify.@generate_test
###
### which generates a test using a given latexify function.
### For example:
###
### Latexify.@generate_test latexify([1, 2, 3], [4, 5, 6]; env=:mdtable)
###
### This puts a ready-made test in your clipboard which you can paste into the
### test file.
###
### Just be sure to remove all such macros before you commit a change since it
### will cause issues with Travis.

@parameters t σ ρ β
@variables x(t) y(t) z(t)
@derivatives D'~t

eqs = [D(x) ~ σ*(y-x)*D(x-y)/D(z),
       0 ~ σ*x*(ρ-z)/10-y,
       D(z) ~ x*y^(2//3) - β*z]

@test latexify(eqs) ==
raw"\begin{align}
\mathrm{derivative}\left( \mathrm{x}\left( t \right), t \right) =& \frac{\sigma \cdot \left( \mathrm{y}\left( t \right) - \mathrm{x}\left( t \right) \right) \cdot \mathrm{derivative}\left( \mathrm{x}\left( t \right) - \mathrm{y}\left( t \right), t \right)}{\mathrm{derivative}\left( \mathrm{z}\left( t \right), t \right)} \\
0 =& \frac{\sigma \cdot \mathrm{x}\left( t \right) \cdot \left( \rho - \mathrm{z}\left( t \right) \right)}{10} - \mathrm{y}\left( t \right) \\
\mathrm{derivative}\left( \mathrm{z}\left( t \right), t \right) =& \mathrm{x}\left( t \right) \cdot \left( \mathrm{y}\left( t \right) \right)^{\frac{2}{3}} - \beta \cdot \mathrm{z}\left( t \right)
\end{align}
"


#!/usr/bin/env julia
# del2z <delta.z@aliyun.com>

using Flux

using Flux.Tracker

W = param(2) # 2.0 (tracked)
b = param(3) # 3.0 (tracked)

f(x) = W * x + b

grads = Tracker.gradient(() -> f(4), params(W, b))

grads[W] # 4.0
grads[b] # 1.0


# Seed random generator to avoid random effects in testing
srand(1)

X = [1 1; 2 2;  1 -1]
y = [-1; -1; 1]

@test_approx_eq_eps optimize(RidgeReg(X, y), optimizer=:closed_form).theta [-0.0666666666666667,-0.23333333333333334] 1e-5
@test_approx_eq_eps optimize(RidgeReg(X, y, λ=1.0), optimizer=:closed_form).theta [0.06060606060606055,-0.606060606060606] 1e-5
@test_approx_eq_eps optimize(RidgeReg(X, y, λ=10.0), optimizer=:closed_form).theta [0.1791607732201792,-0.7732201791607733] 1e-5
@test_approx_eq_eps optimize(RidgeReg(X, y), optimizer=:sgd).theta [-0.0666666666666667,-0.23333333333333334] 5e-2
@test_approx_eq_eps optimize(RidgeReg(X, y, λ=1.0), optimizer=:sgd).theta [0.06060606060606055,-0.606060606060606] 5e-2
@test_approx_eq_eps optimize(RidgeReg(X, y, λ=10.0), optimizer=:sgd).theta [0.1791607732201792,-0.7732201791607733] 5e-2
@test_approx_eq_eps optimize(RidgeReg(X, y), optimizer=:l_bfgs).theta [-0.0666666666666667,-0.23333333333333334] 1e-5
@test_approx_eq_eps optimize(RidgeReg(X, y, λ=1.0), optimizer=:l_bfgs).theta [0.06060606060606055,-0.606060606060606] 1e-5
@test_approx_eq_eps optimize(RidgeReg(X, y, λ=10.0), optimizer=:l_bfgs).theta [0.1791607732201792,-0.7732201791607733] 1e-5
show(IOBuffer(), RidgeReg(X, y))

@test_throws DimensionMismatch RidgeReg(X', y)
@test_throws ArgumentError optimize(RidgeReg(X, y), optimizer=:blubs)
@test_throws ArgumentError optimize(RidgeReg(X, y, λ=0.0))


# System of chemical reactions similar in structure to ModelingToolkit.jl ReactionSystem
# allow stochastic variables (sampled from the geometric distribution) as stoichometry matrix coefficients (hence Mod)
# Note that this construction is very limited compared to ReactionSystem and is not integrated within the broader framework of ModelingToolkit.jl
# parts of code adapted from
# https://github.com/SciML/ModelingToolkit.jl/blob/master/src/systems/reaction/reactionsystem.jl
"""
$(TYPEDEF)

A system of chemical reactions. The structure is heavily based on
[`ModelingToolkit.ReactionSystem`](https://mtk.sciml.ai/stable/systems/ReactionSystem)
and the relevant fields are left identical for the sake of familiar users.

# Fields
$(FIELDS)

# Notes
- The molecule numbers of each chemical species must be explicitly defined as
  [`Symbolics.@variables`](@ref) and reaction constants as [`ModelingToolkit.@parameters`](@ref).
  The time, ``t``, must also be initialised as a `ModelingToolkit.parameter`
  and be used to indicate the time-dependence of species' variables.
- To specify that a certain reaction product is geometrically distributed, the
  corresponding stoichiometric matrix element must be a `ModelingToolkit.parameter`
  denoting the mean of the distribution.

# Example (also see [here](@ref geometric-and-conditional))
```julia
using MomentClosure

@parameters t, c₁, c₂, c₃, c₄, Ω
@variables X(t), Y(t)

# stoichiometric matrix
S_mat = [ 1 -1  1 -1;
         -1  1  0  0]

# propensity functions
a = [c₁*X*Y*(X-1)/Ω^2, c₂*X, c₃*Ω, c₄*X]

rn = ReactionSystemMod(t, [X, Y], [c₁, c₂, c₃, c₄, Ω], a, S_mat)
```
"""
struct ReactionSystemMod
    """Independent variable (usually time)"""
    iv::Any
    """Dependent (state) variables representing amount of each species."""
    states::Vector
    """Parameter variables."""
    ps::Vector
    """Propensity functions of all reactions"""
    a::Vector
    """Stoichiometric matrix"""
    S::Matrix

    function ReactionSystemMod(iv, states, ps, a, S)
        if size(S)[1] != size(states)[1]
            error("Inconsistent stoichiometric matrix dimensions and number of species")
        elseif size(S)[2] != size(a)[1]
            error("Inconsistent stoichiometric matrix dimensions and number of reactions")
        else
            if (typeof(iv) != Num)
                error("the indepent variable is not Num type (possibly breaking update of Catalyst/ModelingToolkit)")
            else
            # initially using @parameters and @variables macros we have defined Num type symbols
            # which have to be converted to SymbolicUtils supported types for further manipulations
                new(value(iv), value.(states), value.(ps), simplify.(value.(a)), value.(S))
            end
        end
    end

end

# Implementing a number of basic functions allowing easy access to reaction network properties.
# The functions are identical to the Catalyst ones (full credit to the developers) in
# order to keep the APIs consistent
# TODO: use getters

"""
    species(rn::ReactionSystemMod)
Given a [`ReactionSystemMod`](@ref), return a vector of species variables expressed as
[`Term{Real}`](https://mtk.sciml.ai/stable/IR/#Types-1). Extension of Catalyst's
[`species`](https://catalyst.sciml.ai/stable/api/catalyst_api/#Catalyst.species) used with a
[`ReactionSystem`](https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem).
"""
function species(rn::ReactionSystemMod)
    rn.states
end

"""
    params(rn::ReactionSystemMod)
Given a [`ReactionSystemMod`](@ref), return a vector of parameter variables
expressed as `Sym{ModelingToolkit.Parameter{Real}`. Extension of  Catalyst's
[`params`](https://catalyst.sciml.ai/stable/api/catalyst_api/#Catalyst.params) used with a
[`ReactionSystem`](https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem).
"""
function params(rn::ReactionSystemMod)
    rn.ps
end

"""
    speciesmap(rn::ReactionSystemMod)
Given a [`ReactionSystemMod`](@ref), return a Dictionary mapping from species to
species indices. Extension of Catalyst's
[`speciesmap`](https://catalyst.sciml.ai/stable/api/catalyst_api/#Catalyst.speciesmap) used with a
[`ReactionSystem`](https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem).
"""
function speciesmap(rn::ReactionSystemMod)
    Dict(S => i for (i,S) in enumerate(species(rn)))
end

"""
    paramsmap(rn::ReactionSystemMod)
Given a [`ReactionSystemMod`](@ref), return a Dictionary mapping from parameters to
parameter indices.  Extension of Catalyst's
[`paramsmap`](https://catalyst.sciml.ai/stable/api/catalyst_api/#Catalyst.paramsmap) used with a
[`ReactionSystem`](https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem).
"""
function paramsmap(rn::ReactionSystemMod)
    Dict(p => i for (i,p) in enumerate(params(rn)))
end

"""
    numspecies(rn::ReactionSystemMod)
Return the number of species within the given [`ReactionSystemMod`](@ref). Extension of Catalyst's
[`numspecies`](https://catalyst.sciml.ai/stable/api/catalyst_api/#Catalyst.numspecies) used with a
[`ReactionSystem`](https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem).
"""
function numspecies(rn::ReactionSystemMod)
    length(rn.states)
end

"""
    numreactions(rn::ReactionSystemMod)
Return the number of reactions within the given [`ReactionSystemMod`](@ref). Extension of Catalyst's
[`numreactions`](https://catalyst.sciml.ai/stable/api/catalyst_api/#Catalyst.numreactions) used with a
[`ReactionSystem`](https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem).
"""
function numreactions(rn::ReactionSystemMod)
    length(rn.a)
end

"""
    numparams(rn::ReactionSystemMod)
Return the number of parameters within the given [`ReactionSystemMod`](@ref). Extension of Catalyst's
[`numparams`](https://catalyst.sciml.ai/stable/api/catalyst_api/#Catalyst.numparams) used with a
[`ReactionSystem`](https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem).
"""
function numparams(rn::ReactionSystemMod)
    length(rn.ps)
end

"""
    get_S_mat(rn::Union{ReactionSystem, ReactionSystemMod})
Return the (net) stoichiometric matrix of the given [`ReactionSystem`]
(https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem)
or [`ReactionSystemMod`](@ref). Note that in case of a Catalyst's [`ReactionSystem`]([`ReactionSystem`]
(https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem),
the transpose of [`netstoichmat`]((https://catalyst.sciml.ai/stable/api/catalyst_api/#Catalyst.netstoichmat)
is returned.
"""
function get_S_mat(rn::Union{ReactionSystem, ReactionSystemMod}; smap=speciesmap(rn))
    if rn isa ReactionSystem
        netstoichmat(rn; smap)'
    else
        ordering = [smap[s] for s in species(rn)]
        view(rn.S, ordering, :)
    end
end

"""
    propensities(rn::Union{ReactionSystem, ReactionSystemMod}; combinatoric_ratelaw=true)
Return a vector of propensity functions of all reactions in the given [`ReactionSystem`]
(https://catalyst.sciml.ai/stable/api/catalyst_api/#ModelingToolkit.ReactionSystem)
or [`ReactionSystemMod`](@ref).

Notes:
- `combinatoric_ratelaw=true` uses binomials in calculating the propensity functions
  of a `ReactionSystem`, see the notes for [`ModelingToolkit.jumpratelaw`]
  (https://mtk.sciml.ai/stable/systems/ReactionSystem/#ModelingToolkit.jumpratelaw).
  *Note* that this field is irrelevant using `ReactionSystemMod` as then the
  propensities are defined directly by the user.
"""
function propensities(rn::Union{ReactionSystem, ReactionSystemMod}; combinatoric_ratelaw=true)
    if typeof(rn) == ReactionSystem
        simplify.(jumpratelaw.(reactions(rn), combinatoric_ratelaw=combinatoric_ratelaw))
    else
        rn.a
    end
end


"""
Struct for Brownian Motion

		bmProcess=BrownianMotionVec(σ::num1,μ::num2) where {num1,num2 <: Number}

Where:\n
		σ	=	volatility of the process.
		μ	=	drift of the process.
"""
mutable struct BrownianMotionVec{num <: Number, num1 <: Number, num4 <: Number, numtype <: Number} <: AbstractMonteCarloEngine{numtype}
    σ::num
    μ::Curve{num1, num4}
    function BrownianMotionVec(σ::num, μ::Curve{num1, num4}) where {num <: Number, num1 <: Number, num4 <: Number}
        if σ <= 0
            error("Volatility must be positive")
        else
            zero_typed = zero(num) + zero(num1) + zero(num4)
            return new{num, num1, num4, typeof(zero_typed)}(σ, μ)
        end
    end
end

function BrownianMotion(σ::num, μ::Curve{num1, num4}) where {num <: Number, num1 <: Number, num4 <: Number}
    return BrownianMotionVec(σ, μ)
end

function simulate!(X, mcProcess::BrownianMotionVec, mcBaseData::MonteCarloConfiguration{type1, type2, type3, SerialMode, type5}, T::numb) where {numb <: Number, type1 <: Number, type2 <: Number, type3 <: StandardMC, type5 <: Random.AbstractRNG}
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    σ = mcProcess.σ
    μ = mcProcess.μ
    @assert T > 0
    dt = T / Nstep
    stddev_bm = σ * sqrt(dt)
    zero_drift = μ(dt * 0, dt)
    isDualZero = stddev_bm * 0 * zero_drift
    view(X, :, 1) .= isDualZero
    @inbounds for j = 1:Nstep
        tmp_ = μ((j - 1) * dt, dt)
        @inbounds for i = 1:Nsim
            x_i_j = @views X[i, j]
            @views X[i, j+1] = x_i_j + tmp_ + stddev_bm * randn(mcBaseData.rng)
        end
    end

    nothing
end

function simulate!(X, mcProcess::BrownianMotionVec, mcBaseData::MonteCarloConfiguration{type1, type2, type3, SerialMode, type5}, T::numb) where {numb <: Number, type1 <: Number, type2 <: Number, type3 <: AntitheticMC, type5 <: Random.AbstractRNG}
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    σ = mcProcess.σ
    μ = mcProcess.μ
    @assert T > 0
    dt = T / Nstep
    stddev_bm = σ * sqrt(dt)
    zero_drift = μ(dt * 0, dt)
    isDualZero = stddev_bm * 0 * zero_drift
    view(X, :, 1) .= isDualZero
    Nsim_2 = div(Nsim, 2)

    @inbounds for j = 1:Nstep
        tmp_ = μ((j - 1) * dt, dt)
        @inbounds for i = 1:Nsim_2
            Z = stddev_bm * randn(mcBaseData.rng)
            X[2*i-1, j+1] = X[2*i-1, j] + tmp_ + Z
            X[2*i, j+1] = X[2*i, j] + tmp_ - Z
        end
    end

    nothing
end


import FireflyModule.Firefly

function light_intensity(source_node::Firefly, destination_node::Firefly, λ::Float)::Float
    """Calculates the light intensity for the source firefly.
    """
    I_0 = source_node.cost
    r, _ = hamming_distance(source_node, destination_node)
    return I_0 * ℯ ^ (-λ * (r^2))
end


using GaussianMarkovRandomFields

using Random
using BenchmarkTools
using SparseArrays, LinearAlgebra
using LDLFactorizations
using Distributions, PDMats

using ForwardDiff
using ReverseDiff
# using Zygote # broken
using FiniteDiff
using Turing

using Test

function make_precision(ρ, k)
    return spdiagm(-1 => -ρ*ones(k-1), 0 => ones(k), 1 => -ρ*ones(k-1))
end

function test_variables()
    k = 50
    ρ = 0.4
    μ = zeros(k)
    Q = make_precision(ρ, k)
    return k, ρ, μ, Q
end

@testset "Constructors" begin
    k = 10
    T = Float32
    Q = sprandn(k, k, 0.1)
    Q = Q'Q + I
    μ = zeros(T, k)
    L, P = cholesky_ldl(Q)
    F = ldl(Q)

    g1 = GMRF(Q)
    g2 = GMRF(μ, Q)
    g3 = GMRF(μ, Q, L)
    g4 = GMRF(μ, F)
    g5 = GMRF(F)
    @test g1 == g2
    @test g1 == g3
    @test mean(g1) == μ
    @test prec(g1) == Q
    @test eltype(g2) == T
    for gf in [g4, g5]
        @test all(mean(gf) .≈ mean(g1))
        @test all(prec(gf) .≈ prec(g1))
    end
    #  @test all(var(g1) .≈ 1 ./ diag(Q))
    #  Tests for variance below
    @test_throws ErrorException cov(g1)
end

@testset "Linear Algebra" begin
    k, ρ, μ, Q = test_variables()

    L, P = cholesky_ldl(Q)
    L1 = sparse(cholesky(Q).L)
    @test all(Matrix(L * L') .≈ Matrix(Q)[P, P])
    @test prod(diag(L).^2) .≈ det(Q)
    @test all(Matrix(L) .≈ Matrix(sparse(cholesky(Q).L)))
    @test logdet_ldl(Q) .≈ log(det(Matrix(Q)))
    @test logdet_ldl(Q) .≈ logdet(Matrix(Q))
end

@testset "Likelihood" begin
    k, ρ, μ, Q = test_variables()
    Qd = Matrix(Q)

    d = GMRF(μ, Q)
    dd = MvNormalCanon(PDMat(Qd))
    dd1 = MvNormal(zeros(k), inv(Symmetric(Matrix(Q))))
    Random.seed!(1)
    x = rand(d)
    @test logpdf(d, x) ≈ logpdf(dd, x)
    @test logpdf(d, x) ≈ logpdf(dd1, x)
end

@testset "Performance" begin
    Random.seed!(1)
    k, ρ, μ, Q = test_variables()
    Qd = Matrix(Q)

    t_logdet_ldl = @benchmark logdet_ldl($Q)
    t_logdet_ldl_d = @benchmark logdet_ldl($Qd)
    t_logdet_d = @benchmark logdet($Qd)
    println(" ")
    println("logdet_ldl (sparse): $(minimum(t_logdet_ldl))")
    println("logdet_ldl (dense): $(minimum(t_logdet_ldl_d))")
    println("logdet (dense): $(minimum(t_logdet_d))")

    @test minimum(t_logdet_ldl.times) .< minimum(t_logdet_ldl_d.times)
    @test minimum(t_logdet_ldl.times) .< minimum(t_logdet_d.times)

    d = GMRF(μ, Q)
    d1 = MvNormalCanon(PDSparseMat(Q))
    dd = MvNormalCanon(PDMat(Qd))
    dd1 = MvNormal(zeros(k), inv(Symmetric(Matrix(Q))))
    x = rand(d)

    t_logpdf_s = @benchmark logpdf($d, $x)
    t_logpdf_s1 = @benchmark logpdf($d1, $x)
    t_logpdf_d = @benchmark logpdf($dd, $x)
    t_logpdf_d1 = @benchmark logpdf($dd1, $x)
    println(" ")
    println("logpdf (sparse, LDL): $(minimum(t_logpdf_s))")
    println("logpdf (sparse, MvNormalCanon): $(minimum(t_logpdf_s1))")
    println("logpdf (dense, MvNormalCanon): $(minimum(t_logpdf_d))")
    println("logdet (dense, MvNormal): $(minimum(t_logpdf_d1))")

    @test_broken minimum(t_logpdf_s.times) .< minimum(t_logpdf_s1.times)
    @test minimum(t_logpdf_s.times) .< minimum(t_logpdf_d.times)
    # @test minimum(t_logpdf_s.times) .< minimum(t_logpdf_d1.times)
end

@testset "AD" begin
    Random.seed!(1)
    k, ρ, μ, Q = test_variables()
    d = GMRF(μ, Q)
    x = rand(d)

    logistic(x) = 1 / (1 + exp(-x))

    function test_loglik(θ)
        σ = exp(θ[1])
        ρ = σ/2 * logistic(θ[2])
        μ = θ[3:end]
        Q = make_precision(ρ, length(μ))
        d = GMRF(μ, Q)
        return logpdf(d, x)
    end

    θ = randn(k + 2)
    Δfinite = FiniteDiff.finite_difference_gradient(test_loglik, θ)
    @test all(ForwardDiff.gradient(test_loglik, θ) .≈ Δfinite)
    @test all(ReverseDiff.gradient(test_loglik, θ) .≈ Δfinite)
    # Zygote broken for some reason
    # @test all(Zygote.gradient(test_loglik, θ) .≈ Δfinite)
end

@testset "Turing" begin
    Random.seed!(1)
    k, ρ, μ, Q = test_variables()
    d = GMRF(μ, Q)
    x = rand(d)

    @model function GMRFTest(x)
        n = length(x)
        ρ ~ Uniform(0, 0.5)
        Q = make_precision(ρ, n)
        x ~ GMRF(Q)
    end

    m = GMRFTest(x)
    Turing.setadbackend(:forwarddiff)
    @test typeof(sample(m, NUTS(), 50)) <: Turing.Chains
    Turing.setadbackend(:reversediff)
    @test typeof(sample(m, NUTS(), 50)) <: Turing.Chains
end

@testset "Marginal variance" begin
    ### Setup AR1 example
    function ar1_precision(n::I, ρ::T) where {I<:Integer,T}
        offdiagvec = -ρ * ones(T, n - 1)
        diagvec = T[one(ρ); ones(T, n - 2) .+ ρ^2; one(ρ)]
        return  spdiagm(-1 => offdiagvec, 0 => diagvec, 1 => offdiagvec)
    end
    function ar1_precision(n, ρ, σ)
        ρ, σ = promote(ρ, σ)
        return 1 / σ^2 * ar1_precision(n, ρ)
    end
    #  function ar1_precision(n, ρ::T, σ::AbstractVector{T}) where T
        #  scalemat = Diagonal(σ)
        #  return scalemat * ar1_precision(n, ρ) * scalemat
    #  end

    k, ρ, μ, Q = test_variables()

    Q1 = ar1_precision(k, ρ)
    Q2 = ar1_precision(k, ρ, 2.0)
    # Qns = ar1_precision(k, ρ, range(0.1, 1.9, length = k))

    n = 10_000   # number of measurements in time series
    σ = 2.0      # innovation variance
    ρ = 0.4      # autocorrelation coefficient

    Qbig = ar1_precision(n, ρ, σ)
    dbig = GMRF(Qbig)

    # Compare to the analytic marginal variance for an AR(1) process and ensure
    # that we don't overflow the stack on large precision matrices.
    @test all(var(dbig) .≈ σ^2 / (1 - ρ^2))

    # Test that all elements are computed accurately
    Q_small = ar1_precision(k, ρ)
    # Convert to generic sparse matrix
    L_small = dropzeros(sparse(cholesky(Matrix(Q_small)).L))
    # Invert directly to get full covariance matrix for comparison
    Σ_small = inv(Matrix(Q_small))

    Σ_elem = similar(Σ_small)
    for j in k:-1:1, i in k:-1:1
        Σ_elem[i, j] = GaussianMarkovRandomFields.cov_element(L_small, i, j)
    end

    @test all(Σ_elem .≈ Σ_small)
end


using DifferentialEquations
using SimpleDiffEq
using DiffEqSensitivity
using Random
using Distributions
using NestedSamplers
using StatsBase: sample, Weights
using MCMCChains: Chains, describe
using StatsPlots

function sir_ode!(du,u,p,t)
    (S,I,R,C) = u
    (β,c,γ) = p
    N = S+I+R
    infection = β*c*I/N*S
    recovery = γ*I
    @inbounds begin
        du[1] = -infection
        du[2] = infection - recovery
        du[3] = recovery
        du[4] = infection
    end
    nothing
end;

δt = 1.0
tmax = 40.0
tspan = (0.0,tmax)
obstimes = 1.0:1.0:tmax;

u0 = [990.0,10.0,0.0,0.0]; # S,I.R,Y

p = [0.05,10.0,0.25]; # β,c,γ

prob_ode = ODEProblem(sir_ode!,u0,tspan,p)
sol_ode = solve(prob_ode,Tsit5(),saveat=δt);

out = Array(sol_ode)
C = out[4,:];

X = C[2:end] .- C[1:(end-1)];

Random.seed!(1234);

Y = rand.(Poisson.(X));

function ll(x)
    (i0,β) = x
    I = i0*1000.0
    prob = remake(prob_ode,u0=[1000.0-I,I,0.0,0.0],p=[β,10.0,0.25])
    sol = solve(prob,Tsit5(),saveat=δt)
    out = Array(sol)
    C = out[4,:]
    X = C[2:end] .- C[1:(end-1)]
    nonpos = sum(X .<= 0)
    if nonpos > 0
        return Inf
    end
    sum(logpdf.(Poisson.(X),Y))
end;

priors = [
    Uniform(0, 0.1),
    Uniform(0, 0.1)
]

model = NestedModel(ll, priors);

spl = Nested(2, 10000, bounds=Bounds.MultiEllipsoid)

chain = sample(model, spl;
               param_names=["i0", "β"],
               chain_type=Chains)

describe(chain)

plot(chain)

include(joinpath(@__DIR__,"tutorials","appendix.jl"))
appendix()


"""
Determine the null and alternate hypotheses under z-tests given:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
"""
function hypotheses(::Type{OneSampleZTest}, n::Real, δ::Real, σ::Real)
    se = sqrt(1 / n) * σ
    null = Normal(0, se)
    alt = Normal(δ, se)
    return null, alt
end

"""
Determine the null and alternate hypotheses under t-tests given:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
"""
function hypotheses(::Type{EqualVarianceZTest}, n::Real, δ::Real, σ::Real)
    se = sqrt(2 / n) * σ
    null = Normal(0, se)
    alt = Normal(δ, se)
    return null, alt
end

# @doc """
# Determine the values of z that lead to overestimates of the magnitude of δ.
# """ ->
thresholds(alt::Normal) = -mean(alt), mean(alt)


module TestConnectance
    using Test
    using EcologicalNetworks
    using LinearAlgebra

    # Generate some data
    N = BipartiteProbabilisticNetwork([0.0 0.1 0.0; 0.2 0.0 0.2; 0.4 0.5 0.0])

    @test links(N) ≈ 1.4
    @test links_var(N) ≈ 0.9
    @test connectance(N) ≈ 1.4 / 9.0
    @test connectance_var(N) ≈ 0.011111111111111111111

    # Once more with a deterministic network
    N = BipartiteNetwork([false true false; false false true; true true true])

    @test links(N) ≈ 5
    @test connectance(N) ≈ 5/9
    @test linkage_density(N) ≈ 5/6

    # Once more with a deterministic unipartite network
    N = UnipartiteNetwork([false true false; false false true; true true true])

    @test links(N) ≈ 5
    @test connectance(N) ≈ 5/9
    @test linkage_density(N) ≈ 5/3

    # Links with quantitative networks
    K = BipartiteQuantitativeNetwork(rand(Float64, (3,3)))
    @test links(K) == 9

    # Convert to adjacency
    @test connectance(UnipartiteNetwork(Matrix(I, (10,10)))) ≈ connectance(UnipartiteQuantitativeNetwork(Matrix{Int64}(I, (10,10)).*2))

end


module ParameterDistributionStorage

## Imports
using Distributions
using Statistics
using StatsBase
using Random

## Exports

#types

export ParameterDistributionType
#objects
export Parameterized, Samples
export ParameterDistribution
export Constraint

#functions
export get_name, get_distribution, get_total_dimension, get_dimensions, get_all_constraints, get_n_samples
export sample_distribution
export no_constraint, bounded_below, bounded_above, bounded
export transform_constrained_to_unconstrained, transform_unconstrained_to_constrained
export get_logpdf, get_cov, get_var, get_mean, batch

## Objects
# for the Distribution
abstract type ParameterDistributionType end

"""
    struct Parameterized <: ParameterDistributionType

A distribution constructed from a parametrized formula (e.g Julia Distributions.jl)
"""
struct Parameterized <: ParameterDistributionType
    distribution::Distribution
end

"""
    struct Samples{FT<:Real} <: ParameterDistributionType

A distribution comprised of only samples, stored as columns of parameters
"""
struct Samples{FT<:Real} <: ParameterDistributionType
    distribution_samples::Array{FT,2} #parameters are columns
    Samples(distribution_samples::Array{FT,2}; params_are_columns=true) where {FT <: Real} = params_are_columns ? new{FT}(distribution_samples) : new{FT}(permutedims(distribution_samples,(2,1)))
    #Distinguish 1 sample of an ND parameter or N samples of 1D parameter, and store as 2D array
    Samples(distribution_samples::Array{FT,1}; params_are_columns=true) where {FT <: Real} = params_are_columns ? new{FT}(reshape(distribution_samples,1,:)) : new{FT}(reshape(distribution_samples,:,1))

end

# For the transforms
abstract type ConstraintType end

"""
    struct Constraint <: ConstraintType

Contains two functions to map between constrained and unconstrained spaces.
"""
struct Constraint <: ConstraintType
    constrained_to_unconstrained::Function
    unconstrained_to_constrained::Function
end

"""
    function no_constraint()

Constructs a Constraint with no constraints, enforced by maps x -> x and x -> x.
"""
function no_constraint()
    c_to_u = (x -> x)
    u_to_c = (x -> x)
    return Constraint(c_to_u,u_to_c)
end

"""
    function bounded_below(lower_bound::FT) where {FT <: Real}

Constructs a Constraint with provided lower bound, enforced by maps x -> log(x - lower_bound) and x -> exp(x) + lower_bound.
"""
function bounded_below(lower_bound::FT) where {FT <: Real}
    c_to_u = ( x -> log(x - lower_bound) )
    u_to_c = ( x -> exp(x) + lower_bound )
    return Constraint(c_to_u, u_to_c)
end

"""
    function bounded_above(upper_bound::FT) where {FT <: Real}

Constructs a Constraint with provided upper bound, enforced by maps x -> log(upper_bound - x) and x -> upper_bound - exp(x).
"""
function bounded_above(upper_bound::FT) where {FT<:Real}
    c_to_u = ( x -> log(upper_bound - x) )
    u_to_c = ( x -> upper_bound - exp(x) )
    return Constraint(c_to_u, u_to_c)
end

"""
    function bounded(lower_bound::FT, upper_bound::FT) where {FT <: Real}

Constructs a Constraint with provided upper and lower bounds, enforced by maps
x -> log((x - lower_bound) / (upper_bound - x))
and
x -> (upper_bound * exp(x) + lower_bound) / (exp(x) + 1)

"""
function bounded(lower_bound::FT, upper_bound::FT) where {FT <: Real}
    if (upper_bound <= lower_bound)
        throw(DomainError("upper bound must be greater than lower bound"))
    end
    c_to_u = ( x -> log( (x - lower_bound) / (upper_bound - x)) )
    u_to_c = ( x -> (upper_bound * exp(x) + lower_bound) / (exp(x) + 1) )
    return Constraint(c_to_u, u_to_c)
end

"""
    function len(c::Array{CType})

The number of constraints, each constraint has length 1.
"""
function len(c::CType) where {CType <: ConstraintType}
    return 1
end

function len(carray::Array{CType}) where {CType <: ConstraintType}
    return size(carray)[1]
end

"""
    function get_dimensions(d<:ParametrizedDistributionType)

The number of dimensions of the parameter space
"""
function dimension(d::Parameterized)
    return length(d.distribution)
end

function dimension(d::Samples)
    return size(d.distribution_samples)[1]
end

"""
    function n_samples(d::Samples)

The number of samples in the array
"""
function n_samples(d::Samples)
    return size(d.distribution_samples)[2]
end
function n_samples(d::Parameterized)
    return "Distribution stored in Parameterized form, draw samples using `sample_distribution` function"
end

"""
    struct ParameterDistribution

Structure to hold a parameter distribution, always stored as an array of distributions
"""
struct ParameterDistribution{PDType <: ParameterDistributionType, CType <: ConstraintType, ST <: AbstractString}
    distributions::Array{PDType}
    constraints::Array{CType}
    names::Array{ST}

    function ParameterDistribution(parameter_distributions::Union{PDType,Array{PDType}},
                                   constraints::Union{CType, Array{CType}, Array},
                                   names::Union{ST, Array{ST}}) where {PDType <: ParameterDistributionType,
                                                                       CType <: ConstraintType,
                                                                       ST <: AbstractString}

        parameter_distributions = isa(parameter_distributions, PDType) ? [parameter_distributions] : parameter_distributions
        n_parameters_per_dist = [dimension(pd) for pd in parameter_distributions]
        constraints = isa(constraints, Union{<:ConstraintType,Array{<:ConstraintType}}) ? [constraints] : constraints #to calc n_constraints_per_dist
        names = isa(names, ST) ? [names] : names

        n_constraints_per_dist = [len(c) for c in constraints]
        n_dists = length(parameter_distributions)
        n_names = length(names)
        if !(n_parameters_per_dist == n_constraints_per_dist)
            throw(DimensionMismatch("There must be one constraint per parameter in a distribution, use NoConstraint() type if no constraint is required"))
        elseif !(n_dists == n_names)
            throw(DimensionMismatch("There must be one name per parameter distribution"))
        else
            constraints = cat(constraints...,dims=1)

            new{PDType,ConstraintType,ST}(parameter_distributions, constraints, names)
        end
    end

end

## Functions

"""
    function get_name(pd::ParameterDistribution)

Returns a list of ParameterDistribution names
"""
function get_name(pd::ParameterDistribution)
    return pd.names
end

"""
    function get_dimensions(pd::ParameterDistribution)

The number of dimensions of the parameter space
"""
function get_dimensions(pd::ParameterDistribution)
    return [dimension(d) for d in pd.distributions]
end
function get_total_dimension(pd::ParameterDistribution)
    return sum(dimension(d) for d in pd.distributions)
end

"""
    function get_n_samples(pd::ParameterDistribution)

The number of samples in a Samples distribution
"""
function get_n_samples(pd::ParameterDistribution)
    return Dict{String,Any}(pd.names[i] => n_samples(d) for (i,d) in enumerate(pd.distributions))
end

"""
    function get_all_constraints(pd::ParameterDistribution)

returns the (flattened) array of constraints of the parameter distribution
"""
function get_all_constraints(pd::ParameterDistribution)
    return pd.constraints
end

"""
    function batch(pd:ParameterDistribution)

Returns a list of contiguous [collect(1:i), collect(i+1:j),... ] used to split parameter arrays by distribution dimensions
"""
function batch(pd::ParameterDistribution)
    #chunk xarray to give to the different distributions.
    d_dim = get_dimensions(pd) #e.g [4,1,2]
    d_dim_tmp = Array{Int64}(undef,size(d_dim)[1]+1)
    d_dim_tmp[1] = 0
    for i = 2:size(d_dim)[1]+1
        d_dim_tmp[i] = sum(d_dim[1:i-1]) # e.g [0,4,5,7]
    end

    return [collect(d_dim_tmp[i]+1:d_dim_tmp[i+1]) for i = 1:size(d_dim)[1]] # e.g [1:4, 5:5, 6:7]
end

"""
    function get_distribution(pd::ParameterDistribution)

Returns a `Dict` of `ParameterDistribution` distributions, with the parameter names as dictionary keys. For parameters represented by `Samples`, the samples are returned as a 2D (parameter_dimension x n_samples) array
"""
function get_distribution(pd::ParameterDistribution)
    return Dict{String,Any}(pd.names[i] => get_distribution(d) for (i,d) in enumerate(pd.distributions))
end

function get_distribution(d::Samples)
    return d.distribution_samples
end
function get_distribution(d::Parameterized)
    return d.distribution
end

"""
    function sample_distribution(pd::ParameterDistribution)

Draws samples from the parameter distributions returns an array, with parameters as columns
"""
function sample_distribution(pd::ParameterDistribution)
    return sample_distribution(pd,1)
end

function sample_distribution(pd::ParameterDistribution, n_draws::IT) where {IT <: Integer}
    return cat([sample_distribution(d,n_draws) for d in pd.distributions]...,dims=1)
end

function sample_distribution(d::Samples, n_draws::IT)  where {IT <: Integer}
    n_stored_samples = n_samples(d)
    samples_idx = StatsBase.sample(collect(1:n_stored_samples), n_draws)
    if dimension(d) == 1
        return reshape(d.distribution_samples[:,samples_idx],:,n_draws) #columns are parameters
    else
        return d.distribution_samples[:,samples_idx]
    end
end

function sample_distribution(d::Parameterized, n_draws::IT) where {IT <: Integer}
    if dimension(d) == 1
        return reshape(rand(d.distribution,n_draws), :, n_draws) #columns are parameters
    else
        return rand(d.distribution, n_draws)
    end
end

"""
    function logpdf(pd::ParameterDistribution, xarray::Array{<:Real,1})

Obtains the independent logpdfs of the parameter distributions at xarray (non-Samples Distributions only), and returns their sum.
"""
function get_logpdf(d::Parameterized, xarray::Array{FT,1}) where {FT <: Real}
    return logpdf.(d.distribution, xarray)
end

function get_logpdf(pd::ParameterDistribution, xarray::Array{FT,1}) where {FT <: Real}
    #first check we don't have sampled distribution
    for d in pd.distributions
        if typeof(d) <: Samples
            throw(ErrorException("Cannot compute get_logpdf of Samples distribution. Consider using a Parameterized type for your prior."))
        end
    end
    #assert xarray correct dim/length
    if size(xarray)[1] != get_total_dimension(pd)
        throw(DimensionMismatch("xarray must have dimension equal to the parameter space"))
    end

    # get the index of xarray chunks to give to the different distributions.
    batches = batch(pd)

    # perform the logpdf of each of the distributions, and returns their sum
    return sum(cat([get_logpdf(d, xarray[batches[i]]) for (i,d) in enumerate(pd.distributions)]...,dims=1))
end

"""
    function get_cov(pd::ParameterDistribution)

returns a blocked covariance of the distributions
"""
function get_cov(d::Parameterized)
    return cov(d.distribution)
end

function get_cov(d::Samples)
    return cov(d.distribution_samples, dims=2) #parameters are columns
end

function get_var(d::Parameterized)
    return var(d.distribution)
end

function get_var(d::Samples)
    return var(d.distribution_samples)
end

function get_cov(pd::ParameterDistribution)
    #first check we don't have sampled distribution

    d_dims = get_dimensions(pd)

    # create each block (co)variance
    block_cov = Array{Any}(undef,size(d_dims)[1])
    for (i,dimension) in enumerate(d_dims)
        if dimension == 1
            block_cov[i] = get_var(pd.distributions[i])
        else
            block_cov[i] = get_cov(pd.distributions[i])
        end
    end

    return cat(block_cov...,dims=(1,2)) #build the block diagonal (dense) matrix

end

"""
    get_mean(pd::Parameterized)

returns a mean of parameterized distribution
"""
function get_mean(d::Parameterized)
    return mean(d.distribution)
end

"""
    get_mean(pd::Samples)

returns a mean of the samples
"""
function get_mean(d::Samples)
    return mean(d.distribution_samples, dims=2) #parameters are columns
end

"""
    function get_mean(pd::ParameterDistribution)

returns a mean of the distributions
"""
function get_mean(pd::ParameterDistribution)
    return cat([get_mean(d) for d in pd.distributions]...,dims=1)
end

#apply transforms

"""
    function transform_constrained_to_unconstrained(pd::ParameterDistribution, x::Array{<:Real})

Apply the transformation to map (possibly constrained) parameters `xarray` into the unconstrained space
"""
function transform_constrained_to_unconstrained(pd::ParameterDistribution, xarray::Array{FT,1}) where {FT <: Real}
    return cat([c.constrained_to_unconstrained(xarray[i]) for (i,c) in enumerate(pd.constraints)]...,dims=1)
end

"""
    function transform_constrained_to_unconstrained(pd::ParameterDistribution, xarray::Array{FT,2})

Apply the transformation to map (possibly constrained) parameter samples `xarray` into the unconstrained space.
Here, `xarray` contains parameters as columns and samples as rows.
"""
function transform_constrained_to_unconstrained(pd::ParameterDistribution, xarray::Array{FT,2}) where {FT <: Real}
    return Array(hcat([c.constrained_to_unconstrained.(xarray[i,:]) for (i,c) in enumerate(pd.constraints)]...)')
end

"""
    function transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{<:Real,1})

Apply the transformation to map parameters `xarray` from the unconstrained space into (possibly constrained) space
"""
function transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{FT,1}) where {FT <: Real}
    return cat([c.unconstrained_to_constrained(xarray[i]) for (i,c) in enumerate(pd.constraints)]...,dims=1)
end

"""
    function transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{FT,2})

Apply the transformation to map parameter samples `xarray` from the unconstrained space into (possibly constrained) space.
Here, `xarray` contains parameters as columns and samples as rows.
"""
function transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{FT,2}) where {FT <: Real}
    return Array(hcat([c.unconstrained_to_constrained.(xarray[i,:]) for (i,c) in enumerate(pd.constraints)]...)')
end

"""
    function transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{Array{FT,2},1})

Apply the transformation to map parameter sample ensembles `xarray` from the unconstrained space into (possibly constrained) space.
Here, `xarray` contains parameters sample ensembles for different EKP iterations.
"""
function transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{Array{FT,2},1}) where {FT <: Real}
    transf_xarray = []
    for elem in xarray
           push!(transf_xarray, transform_unconstrained_to_constrained(pd, elem))
       end
    return transf_xarray
end

end # of module


using DataFrames, CSV, Dates, StatsBase, Plots, TimeSeries; pyplot()

data = CSV.read("../data/temperatures.csv",copycols = true)
brisbane = data.Brisbane
dates = [Date(
            Year(data.Year[i]),
            Month(data.Month[i]),
            Day(data.Day[i])
        ) for i in 1:nrow(data)]

window1, window2 = 7, 14
d1 = values(moving(mean,TimeArray(dates,brisbane),window1))
d2 = values(moving(mean,TimeArray(dates,brisbane),window2))

grid = (2pi:-2pi/365:0) .+ pi/2
monthsNames = Dates.monthname.(dates[1:31:365])

plot(grid, d1[1:365],
    c=:blue, proj=:polar, label="Brisbane weekly average temp.")
plot!(grid, d2[1:365],
    xticks=([mod.((11pi/6:-pi/6:0) .+ pi/2,2pi) ;], monthsNames),
    c=:red, proj=:polar,
    label="Brisbane fortnightly average temp.", legend=:outerbottom)


__precompile__()

"""
Module for identifying low-rank subsets and subgraphs
"""
module LowRank
using Combinatorics, LinearAlgebra, StatsBase, Parameters
import Base.show
import Base.findmin, Base.findmax

include("generate_datasets.jl")
include("utils.jl")
include("trajectory.jl")
include("low_rank_subset.jl")

export find_low_rank_subset_checkall, find_low_rank_subset_sample_rep, FastLowRank

end # module

# TODO:
# allow for maximization of ∑ σᵢ² for i ≦ k instead of ∑ σᵢ² for i ≧ k+1
# allow for normalization in FastLowRank
# write a test that introduces Gaussian noise to each entry
# write spatial methods
# methods for objectives with normalization?
# write a method to find many low-rank subsets, as opposed to just one


# In this tutorial, we will learn
#  - How to solve a simple PDE with a DG method
#  - How to compute jumps and averages of quantities on the mesh skeleton
#  - How to implement the method of manufactured solutions
#  - How to integrate error norms
#  - How to generate Cartesian meshes in arbitrary dimensions

#
# ## Problem statement
#
# The goal of this tutorial is to solve a PDE using a DG formulation. For
# simplicity, we take the Poisson equation on the unit cube $\Omega \doteq
# (0,1)^3$ as the model problem, namely

# ```math
# \left\lbrace
# \begin{aligned}
# -\Delta u = f  \ &\text{in} \ \Omega,\\
# u = g \ &\text{on}\ \Gamma \doteq \partial\Omega,\\
# \end{aligned}
# \right.
# ```
# where $f$ is the source term and $g$ is the prescribed Dirichlet boundary
# function. In this tutorial, we follow the method of manufactured solutions
# since we want to illustrate how to compute discretization errors. We take
# $u(x) = 3 x_1 + x_2^2 + 2 x_3^3 + x_1 x_2 x_3 $ as the exact solution of the
# problem, for which $f(x)= -2 - 12x_3 $ and $g(x) = u(x)$. The selected
# manufactured solution $u$ is a third order multi-va riate polynomial, which
# can be represented exactly by the FE discretization that we are going to
# define below. In this scenario, the discretization error has to be close to
# the machine precision. We will use this result to validate the proposed
# implementation.
#
# ## Numerical Scheme
#
# We consider a DG formulation to approximate the problem. In particular, we
# consider the symmetric interior penalty method (see, e.g. [1], for specific
# details). For this formulation, the approximation space is made of
# discontinuous piece-wise polynomials, namely
#
# ```math
# V \doteq \{ v\in L^2(\Omega):\ v|_{T}\in Q_p(T) \text{ for all } T\in\mathcal{T}  \},
# ```
# where $\mathcal{T}$ is the set of all cells $T$ of the FE mesh, and $Q_p(T)$
# is a polynomial space of degree $p$ defined on a generic cell $T$. For
# simplicity, we consider Cartesian meshes in this tutorial. In this case, the
# space $Q_p(T)$ is made of multi-variate polynomials up to degree $p$ in each
# spatial coordinate.
#
# In order to write the weak form of the problem, we need to introduce some
# notation. The sets of interior and boundary facets associated with the FE
# mesh $\mathcal{T}$ are denoted here as $\mathcal{F}_\Lambda$ and
# $\mathcal{F}_{\Gamma}$ respectively. In addition, for a given
# function $v\in V$ restricted to the interior facets $\mathcal{F}_\Lambda$, we
# introduce the well known jump and mean value operators,
# ```math
# \begin{aligned}
# \lbrack\!\lbrack v\ n \rbrack\!\rbrack &\doteq v^+\ n^+ + v^- n^-,\\
# \{\! \!\{\nabla v\}\! \!\} &\doteq \dfrac{ \nabla v^+ + \nabla v^-}{2},
# \end{aligned}
# ```
# with $v^+$, and $v^-$ being the restrictions of $v\in V$ to the cells $T^+$,
# $T^-$ that share a generic interior facet in $\mathcal{F}_\Lambda$, and $n^+$,
# and $n^-$ are the facet outward unit normals from either the perspective of
# $T^+$ and $T^-$ respectively.
#
# With this notation, the weak form associated with the interior penalty
# formulation of our problem reads: find $u\in V$ such that $a(u,v) = l(v)$ for
# all $v\in V$. The bilinear and linear forms $a(\cdot,\cdot)$ and $l(\cdot)$
# have contributions associated with the bulk of $\Omega$, the boundary facets
# $\mathcal{F}_{\Gamma}$, and the interior facets $\mathcal{F}_\Lambda$,
# namely
#  ```math
# \begin{aligned}
# a(u,v) &= a_{\Omega}(u,v) + a_{\Gamma}(u,v) + a_{\Lambda}(u,v),\\
# l(v) &= l_{\Omega}(v) + l_{\Gamma}(v).
# \end{aligned}
# ```
# These contributions are defined as
# ```math
# \begin{aligned}
# a_{\Omega}(u,v) &\doteq \sum_{T\in\mathcal{T}} \int_{T} \nabla v \cdot \nabla u \ {\rm d}T,
# \\
# l_{\Omega}(v) &\doteq \int_{\Omega} v\ f \ {\rm d}\Omega,
# \end{aligned}
# ```
# for the volume,
# ```math
# \begin{aligned}
# a_{\Gamma}(u,v)
#    \doteq
#    & - \sum_{F\in\mathcal{F}_{\Gamma}} \int_{F} v\ (\nabla u \cdot n)  \ {\rm d}F
# \\ & - \sum_{F\in\mathcal{F}_{\Gamma}} \int_{F} (\nabla v \cdot n)\ u  \ {\rm d}F
# \\ & + \sum_{F\in\mathcal{F}_{\Gamma}} \dfrac{\gamma}{|F|} \int_{F} v\ u \ {\rm d}F,
# \\
# l_{\Gamma}(v)
# \doteq
#    & - \sum_{F\in\mathcal{F}_{\Gamma}} \int_{F} (\nabla v \cdot n)\ g  \ {\rm d}F
# \\ & + \sum_{F\in\mathcal{F}_{\Gamma}} \dfrac{\gamma}{|F|} \int_{F} v\ g \ {\rm d}F,
# \end{aligned}
# ```
# for the boundary facets and,
# ```math
# \begin{aligned}
# a_{\Lambda}(u,v)
# \doteq
#    & - \sum_{F\in\mathcal{F}_{\Lambda}} \int_{F} \lbrack\!\lbrack v\ n \rbrack\!\rbrack\cdot \{\! \!\{\nabla u\}\! \!\} \ {\rm d}F
# \\ & - \sum_{F\in\mathcal{F}_{\Lambda}} \int_{F} \{\! \!\{\nabla v\}\! \!\}\cdot \lbrack\!\lbrack u\ n \rbrack\!\rbrack \ {\rm d}F
# \\ & + \sum_{F\in\mathcal{F}_{\Lambda}} \dfrac{\gamma}{|F|} \int_{F} \lbrack\!\lbrack v\ n \rbrack\!\rbrack\cdot \lbrack\!\lbrack u\ n \rbrack\!\rbrack \ {\rm d}F,
# \end{aligned}
# ```
#  for the interior facets. In previous expressions, $|F|$ denotes the diameter
#  of the face $F$ (in our Cartesian grid, this is equivalent to the
#  characteristic mesh size $h$), and $\gamma$ is a stabilization parameter that
#  should be chosen large enough such that the bilinear form $a(\cdot,\cdot)$ is
#  stable and continuous. Here, we take $\gamma = p\ (p+1)$ as done in the
#  numerical experiments in reference [2].
#
# ## Manufactured solution
#
# We start by loading the Gridap library and defining the manufactured solution
# $u$ and the associated source term $f$ and Dirichlet function $g$.

using Gridap
u(x) = 3*x[1] + x[2]^2 + 2*x[3]^3 + x[1]*x[2]*x[3]
f(x) = -2 - 12*x[3]
g(x) = u(x)
# We also need to define the gradient of $u$ since we will compute the $H^1$
# error norm later. In that case, the gradient is simply defined as

∇u(x) = VectorValue(3        + x[2]*x[3],
                    2*x[2]   + x[1]*x[3],
                    6*x[3]^2 + x[1]*x[2])

# In addition, we need to tell the Gridap library that the gradient of the
# function `u` is available in the function `∇u` (at this moment `u` and `∇u`
# are two standard Julia functions without any connection between them). This
# is done by adding an extra method to the function `gradient` (aka `∇`)
# defined in Gridap:

import Gridap: ∇
∇(::typeof(u)) = ∇u

#  Now, it is possible to recover function `∇u` from function `u` as `∇(u)`. You
#  can check that the following expression evaluates to `true`.

∇(u) === ∇u

# ## Cartesian mesh generation
#  In order to discretize the geometry of the unit cube, we use the Cartesian mesh generator available in Gridap.

L = 1.0
domain = (0.0, L, 0.0, L, 0.0, L)
n = 4
partition = (n,n,n)
model = CartesianDiscreteModel(domain,partition)

# The type `CartesianDiscreteModel` is a concrete type that inherits from
# `DiscreteModel`, which is specifically designed for building Cartesian
# meshes. The `CartesianDiscreteModel` constructor takes a tuple containing
# limits of the box we want to discretize plus a tuple with the number of cells
# to be generated in each direction (here $4\times4\times4$ cells). You can
# write the model in vtk format to visualize it (see next figure).

 writevtk(model,"model")

# ![](../assets/dg_discretization/model.png)
#
#  Note that the `CaresianDiscreteModel` is implemented for arbitrary
#  dimensions. For instance, the following lines build a
#  `CartesianDiscreteModel` for the unit square $(0,1)^2$ with 4 cells per
#  direction

domain2D = (0.0, L, 0.0, L)
partition2D = (n,n)
model2D = CartesianDiscreteModel(domain2D,partition2D)

# You could also generate a mesh for the unit tesseract $(0,1)^4$ (i.e., the
# unit cube in 4D). Look how the 2D and 3D models are built and just follow the
# sequence.
#
# ## FE spaces
#
# On top of the discrete model, we create the discontinuous space $V$ as
# follows

order = 3
V = TestFESpace(model,
                ReferenceFE(lagrangian,Float64,order),
                conformity=:L2)

# We have select a Lagrangian, scalar-valued interpolation of order $3$ within
# the cells of the discrete model. Since the cells are hexahedra, the resulting
# Lagrangian shape functions are tri-cubic polynomials. In contrast to previous
# tutorials, where we have constructed $H^1$-conforming (i.e., continuous) FE
# spaces, here we construct a $L^2$-conforming (i.e., discontinuous) FE space.
# That is, we do not impose any type of continuity of the shape function on the
# cell boundaries, which leads to the discontinuous FE space $V$ of the DG
# formulation. Note also that we do not pass any information about the
# Dirichlet boundary to the `TestFESpace` constructor since the Dirichlet
# boundary conditions are not imposed strongly in this example.
#
# From the `V` object we have constructed in previous code snippet, we build
# the trial FE space as usual.

U = TrialFESpace(V)

# Note that we do not pass any Dirichlet function to the `TrialFESpace`
# constructor since we do not impose Dirichlet boundary conditions strongly
# here.
#
# ## Numerical integration
#
# Once the FE spaces are ready, the next step is to set up the numerical
# integration. In this example, we need to integrate in three different
# domains: the volume covered by the cells $\mathcal{T}$ (i.e., the
# computational domain $\Omega$), the surface covered by the boundary facets
# $\mathcal{F}_{\Gamma}$ (i.e., the boundary $\Gamma = \partial \Omega$), and
# the surface covered by the interior facets $\mathcal{F}_{\Lambda}$ (i.e. the
# so-called mesh skeleton). In order to integrate in $\Omega$ and on its
# boundary $\Gamma$, we use `Triangulation` and `BoundaryTriangulation` objects
# as already discussed in previous tutorials.

Ω = Triangulation(model)
Γ = BoundaryTriangulation(model)

# Here, we do not pass any boundary identifier to the `BoundaryTriangulation`
# constructor. In this case, an integration mesh for the entire boundary
# $\Gamma$ is constructed by default (which is just what we need in this
# example).
#
# In order to generate an integration mesh for the interior facets
# $\mathcal{F}_{\Lambda}$, we use a new type of `Triangulation` referred to as
# `SkeletonTriangulation`. It can be constructed from a `DiscreteModel` object
# as follows:

Λ = SkeletonTriangulation(model)

# As any other type of `Triangulation`, an `SkeletonTriangulation` can be
# written into a vtk file for its visualization (see next figure, where the
# interior facets $\mathcal{F}_\Lambda$ are clearly observed).

writevtk(Λ,"strian")

# ![](../assets/dg_discretization/skeleton_trian.png)
#
# Once we have constructed the triangulations needed in this example, we define
# the corresponding quadrature rules by passing the triangualtions
# together with the desired degree to the `Measure` function.

degree = 2*order

dΩ = Measure(Ω,degree)
dΓ = Measure(Γ,degree)
dΛ = Measure(Λ,degree)

# We still need a way to represent the unit outward normal vector to the
# boundary $\Gamma$, and the unit normal vector on the interior faces
# $\mathcal{F}_\Lambda$. This is done with the `get_normal_vector` getter.

n_Γ = get_normal_vector(Γ)
n_Λ = get_normal_vector(Λ)

# The `get_normal_vector` getter takes either a boundary or a skeleton
# triangulation and returns an object representing the normal vector to the
# corresponding surface. For boundary triangulations, the returned normal
# vector is the unit outwards one, whereas for skeleton triangulations the
# orientation of the returned normal is arbitrary. In the current
# implementation (Gridap v0.5.0), the unit normal is outwards to the cell with
# smaller id among the two cells that share an interior facet in
# $\mathcal{F}_\Lambda$.
#
# ## Weak form
#
# With these ingredients we can define the different terms in the weak form.
# First, we start with the terms $a_\Omega(\cdot,\cdot)$ , and
# $l_\Omega(\cdot)$ associated with integrals in the volume $\Omega$. This is
# done as in the tutorial for the Poisson equation.

a_Ω(u,v) = ∫( ∇(v)⊙∇(u) )dΩ
l_Ω(v) = ∫( v*f )dΩ

# The terms $a_{\Gamma}(\cdot,\cdot)$ and $l_{\Gamma}(\cdot)$ associated with
# integrals on the boundary $\Gamma$ are defined using an analogous approach:

h = L / n
γ = order*(order+1)
a_Γ(u,v) = ∫( - v*(∇(u)⋅n_Γ) - (∇(v)⋅n_Γ)*u + (γ/h)*v*u )dΓ
l_Γ(v)   = ∫(                - (∇(v)⋅n_Γ)*g + (γ/h)*v*g )dΓ

# Note that in the definition of the functions `a_Γ` and `b_Γ`, we have used
# the object `n_Γ` representing the outward unit normal to the boundary
# $\Gamma$. The code definition of `a_Γ` and `b_Γ` is indeed very close to
# the mathematical definition of the forms $a_{\Gamma}(\cdot,\cdot)$ and
# $b_{\Gamma}(\cdot)$.
#
# Finally, we need to define the term $a_\Lambda(\cdot,\cdot)$ integrated on
# the interior facets $\mathcal{F}_\Lambda$,

a_Λ(u,v) = ∫( - jump(v*n_Λ)⊙mean(∇(u))
              - mean(∇(v))⊙jump(u*n_Λ)
              + (γ/h)*jump(v*n_Λ)⊙jump(u*n_Λ) )dΛ

# Note that the arguments `v`, `u` of function `a_Λ` represent a test and trial
# function *restricted* to the interior facets $\mathcal{F}_\Lambda$. As
# mentioned before in the presentation of the DG formulation, the restriction
# of a function $v\in V$ to the interior faces leads to two different values
# $v^+$ and $v^-$ . In order to compute jumps and averages of the quantities
# $v^+$ and $v^-$, we use the functions `jump` and `mean`, which represent the
# jump and mean value operators $\lbrack\!\lbrack \cdot \rbrack\!\rbrack$ and
# $\{\! \!\{\cdot\}\! \!\}$ respectively. Note also that we have used the
# object `n_Λ` representing the unit normal vector on the interior facets. As a
# result, the notation used to define function `a_Λ` is very close to the
# mathematical definition of the terms in the bilinear form
# $a_\Lambda(\cdot,\cdot)$.
#
# Once the different terms of the weak form have been defined, we build and solve the FE problem.
a(u,v) = a_Ω(u,v) + a_Γ(u,v) + a_Λ(u,v)
l(v) = l_Ω(v) + l_Γ(v)

op = AffineFEOperator(a, l, U, V)
uh = solve(op)

# ## Discretization error
#
# We end this tutorial by quantifying the discretization error associated with
# the computed numerical solution `uh`. In DG methods a simple error indicator
# is the jump of the computed (discontinuous) approximation on the interior
# faces. We compute and visualize the jump of these values as follows (see next
# figure):

writevtk(Λ,"jumps",cellfields=["jump_u"=>jump(uh)])

# Note that the jump of the numerical solution is very small, close to the
# machine precision (as expected in this example with manufactured solution).
# ![](../assets/dg_discretization/jump_u.png)
#
#  A more rigorous way of quantifying the error is to measure it with a norm.
#  Here, we use the $L^2$ and $H^1$ norms, namely
#  ```math
# \begin{aligned}
#  \| w \|_{L^2}^2 & \doteq \int_{\Omega} w^2 \ \text{d}\Omega, \\
#  \| w \|_{H^1}^2 & \doteq \int_{\Omega} w^2 + \nabla w \cdot \nabla w \ \text{d}\Omega.
# \end{aligned}
# ```
#
# The discretization error can be computed in this example as the difference of
# the manufactured and numerical solutions.

e = u - uh

# We compute the error norms as follows. First, we implement the integrands of
# the norms we want to compute.

l2(u) = sqrt(sum( ∫( u⊙u )*dΩ ))
h1(u) = sqrt(sum( ∫( u⊙u + ∇(u)⊙∇(u) )*dΩ ))

# Then, we compute the corresponding integrals with the `integrate` function.
el2 = l2(e)
eh1 = h1(e)

# The `integrate` function returns a lazy object representing the contribution
# to the integral of each cell in the underlying triangulation. To end up with
# the desired error norms, one has to sum these contributions and take the
# square root. You can check that the computed error norms are close to machine
# precision (as one would expect).

tol = 1.e-10
@assert el2 < tol
@assert eh1 < tol

# ## References
#
# [1] D. N. Arnold, F. Brezzi, B. Cockburn, and L. Donatella Marini. Unified analysis of discontinuous Galerkin methods for elliptic problems. *SIAM Journal on Numerical Analysis*, 39 (5):1749–1779, 2001. doi:[10.1137/S0036142901384162](http://dx.doi.org/10.1137/S0036142901384162).
#
# [2] B. Cockburn, G. Kanschat, and D. Schötzau. An equal-order DG method for the incompressible Navier-Stokes equations. *Journal of Scientific Computing*, 40(1-3):188–210, 2009. doi:[10.1007/s10915-008-9261-1](http://dx.doi.org/10.1007/s10915-008-9261-1).
#


function ar1_sample_mean(N, beta, alpha, s)
    sm = 0.0
    x = beta / (1 - alpha)
    for i in 1:N
        sm += x
        x = beta + alpha * x + s * randn()
    end
    return sm / N
end

N = 1e7
beta = 1.0
alpha = 0.9
s = 1.0
tic()
result = ar1_sample_mean(N, beta, alpha, s)
println("mean = $result")
toc()


module significance_testing_module

using Distributions
using Statistics
using RCall

"""
# ____________________________
# Effects distribution fitting

`best_fitting_distribution(data::Array{Float64,1})`

Determine the best fitting distribution to model the data and heuristically estimate the p-values.

# Distributions
- [Bernoulli](https://en.wikipedia.org/wiki/Bernoulli_distribution)
- [Beta](https://en.wikipedia.org/wiki/Beta_distribution)
- [Binomial](https://en.wikipedia.org/wiki/Binomial_distribution)
- [Categorical](https://en.wikipedia.org/wiki/Categorical_distribution)
- [Dicscrete Uniform](https://en.wikipedia.org/wiki/Discrete_uniform_distribution)
- [Exponential](https://en.wikipedia.org/wiki/Exponential_distribution)
- [Gaussian](https://en.wikipedia.org/wiki/Normal_distribution)
- [Gamma](https://en.wikipedia.org/wiki/Gamma_distribution)
- [Geometric](https://en.wikipedia.org/wiki/Geometric_distribution)
- [Laplace](https://en.wikipedia.org/wiki/Laplace_distribution)
- [Pareto](https://en.wikipedia.org/wiki/Pareto_distribution)
- [Poisson](https://en.wikipedia.org/wiki/Poisson_distribution)
- [Inverse Gaussian or Wald](https://en.wikipedia.org/wiki/Inverse_Gaussian_distribution)
- [Uniform](https://en.wikipedia.org/wiki/Uniform_distribution_(continuous))

# Examples
```
using Distributions
using UnicodePlots
x1 = Distributions.rand(Distributions.Normal(0, 1), 100)
UnicodePlots.histogram(x1)
significance_testing_module.best_fitting_distribution(x1)
x2 = Distributions.rand(Distributions.Laplace(0, 1), 100)
UnicodePlots.histogram(x2)
significance_testing_module.best_fitting_distribution(x2)
```
"""
function best_fitting_distribution(data::Array{Float64,1})
    DIST_NAMES =   [Distributions.Bernoulli, Distributions.Beta, Distributions.Binomial, Distributions.Categorical,
                    Distributions.DiscreteUniform, Distributions.Exponential, Distributions.Normal, Distributions.Gamma,
                    Distributions.Geometric, Distributions.Laplace, Distributions.Pareto, Distributions.Poisson,
                    Distributions.InverseGaussian, Distributions.Uniform]
    DIST_INSTANCES = [try Distributions.fit_mle(D, data); catch nothing; end for D in DIST_NAMES]
    NEG_LOGLIK = [try -sum(Distributions.logpdf.(D, data)); catch nothing; end for D in DIST_INSTANCES]
    # NEG_LOGLIK = [try sum(Distributions.pdf.(D, data)); catch nothing; end for D in DIST_INSTANCES]
    DISTRIBUTIONS_DF = hcat((DIST_NAMES[NEG_LOGLIK .!= nothing],
                            DIST_INSTANCES[NEG_LOGLIK .!= nothing],
                            NEG_LOGLIK[NEG_LOGLIK .!= nothing])...)
    D = try
        (DISTRIBUTIONS_DF[argmin(DISTRIBUTIONS_DF[:,3]), 2], DISTRIBUTIONS_DF[argmin(DISTRIBUTIONS_DF[:,3]), 1])
    catch
        (nothing, "Failure to fit into any of the distributions tested.")
    end
    return(D)
end

"""
# ________________________________________________________________________________________
# Heuristic estimates of p-values and -log10(p-values) using the best-fitting distribution

`estimate_pval_lod(data::Array{Float64,1})`

# Output
1. p-values
2. -log10(p-values)

# Examples
```
using Distributions
using UnicodePlots
x1 = Distributions.rand(Distributions.Normal(0, 1), 100)
UnicodePlots.histogram(x1)
PVAL, LOD = significance_testing_module.estimate_pval_lod(x1)
UnicodePlots.scatterplot(LOD)
x2 = Distributions.rand(Distributions.Laplace(0, 1), 100)
UnicodePlots.histogram(x2)
PVAL, LOD = significance_testing_module.estimate_pval_lod(x2)
UnicodePlots.scatterplot(LOD)
```
"""
function estimate_pval_lod(data::Array{Float64,1})
    D, D_name = best_fitting_distribution(data)
    println(string("Distribution used: ", D_name))
    if (D == nothing) | (std(data) == 0)
        PVAL = repeat([1.0], inner=length(data))
        LOD = PVAL .- 1.0
    else
        PVAL = [Distributions.cdf(D, x) <= Distributions.ccdf(D, x) ? 2*Distributions.cdf(D, x) : 2*Distributions.ccdf(D, x) for x in data]
        LOD = -log.(10, PVAL)
        # UnicodePlots.scatterplot(LOD)
    end
    return(PVAL, LOD)
end

"""
# _______________________________________________________________________________________________
# Plot -log10(p-values) with a given Bonferroni threshold at user-defined false positive rate (α)

`plot_manhattan(;chrom::Array{Any,1}, pos::Array{Int64,1}, pval::Array{Float64,1}, lod::Array{Float64,1}, fpr::Float64=0.01, png_fname::String=string("Manhattan_plot-", hash(rand()), ".jpeg"))`

# Inputs
1. *chrom* [Array{Any,1}]: vector of chromosome names
2. *pos* [Array{Int64,1}]: vector of allele/SNP positions in base-pairs
3. *pval* [Array{Float64,1}]: vector of p-values
4. *lod* [Array{Float64,1}]: vector of -log10(p-values)
5. *fpr* [Float64=0.01]: false positive rate threshold or the alpha (default=0.01)
6. *png_fname* [String]: output filename of the Manhattan plot in portable network format (png) (default=string("Manhattan_plot-", hash(rand()), ".jpeg"))

# Output
1. Manhattan plot in portable network graphics (png) format

# Examples
```
using Distributions
chrom = repeat([1,2,3,"X", "Z"], inner=20)
pos = repeat(collect(1:20), outer=5)
eff = Distributions.rand(Distributions.Normal(0, 1), 100)
eff[[10,30,60,90]] = [10.0,20.0,30.0,40.0]
pval, lod = significance_testing_module.estimate_pval_lod(eff)
significance_testing_module.plot_manhattan(chrom=chrom, pos=pos, pval=pval, lod=lod)
```
"""
function plot_manhattan(;chrom::Array{Any,1}, pos::Array{Int64,1}, pval::Array{Float64,1}, lod::Array{Float64,1}, fpr::Float64=0.01, png_fname::String=string("Manhattan_plot-", hash(rand()), ".png"))
	### plot in R
    @rput chrom;
    @rput pos;
    @rput pval;
	@rput lod;
	@rput fpr;
	@rput png_fname;
	### remove missing values and convert infinite LOD into the maximum non-infinite LOD
    R"gwas_out = data.frame(CHROM=as.factor(unlist(chrom)), POS=as.numeric(unlist(pos)), PVALUES=as.numeric(unlist(pval)), LOD=as.numeric(unlist(lod)))";
	R"gwas_out = droplevels(gwas_out[gwas_out$CHROM != 'Intercept', ])";
	R"gwas_out$LOD[is.na(gwas_out$LOD)] = 0.0";
	R"gwas_out$LOD[is.infinite(gwas_out$LOD)] = max(gwas_out$LOD[!is.infinite(gwas_out$LOD)])";
    ### arranging chromosomes or scaffold end to end
    R"gwas_out$POS_CONTINUOUS = gwas_out$POS";
    R"contigs = levels(gwas_out$CHROM)";
    R"min_pos = aggregate(POS ~ CHROM, data=gwas_out, FUN=min)"
    R"max_pos = aggregate(POS ~ CHROM, data=gwas_out, FUN=max)"
    R"idx_chrom = max_pos$CHROM[2:nrow(max_pos)]"
    R"add_us = cumsum(max_pos$POS)[1:nrow(max_pos)-1]"
    R"for (i in 1:length(idx_chrom)){
            chr = idx_chrom[i]
            add = add_us[i]
            gwas_out$POS_CONTINUOUS[gwas_out$CHROM == chr] = gwas_out$POS_CONTINUOUS[gwas_out$CHROM == chr] + add
        }"
	### coloring
	R"colours = RColorBrewer::brewer.pal(6, 'Paired')";
	R"colours_points = rep(colours[c(1,3)], times=ceiling(length(contigs)/2))";
	R"colours_labels = rep(colours[c(2,4)], times=ceiling(length(contigs)/2))";
	R"bonferroni_threshold = -log10(fpr/nrow(gwas_out))";
	### plot -log10(pvalue) points
	R"png(png_fname, width=2000, height=800)";
	R"layout(matrix(c(1,1,1,2), nrow=1, byrow=TRUE))";
	R"par(cex=1.5)";
	R"plot(x=gwas_out$POS_CONTINUOUS, y=gwas_out$LOD, col=colours_points[as.numeric(gwas_out$CHROM)], ylim=c(0, max(c(gwas_out$LOD, bonferroni_threshold))), pch=20, main=paste0(' Bonferroni Threshold at ', round(bonferroni_threshold, 2)), xaxt='n', xlab='Chromosome or Scaffold', ylab=expression(log[10]~(pvalue)), las=1)";
	R"xaxis_pos = aggregate(gwas_out$POS_CONTINUOUS ~ gwas_out$CHROM, FUN=median)";
	R"mtext(side=1, at=xaxis_pos[,2], text=xaxis_pos[,1], cex=1.5)";
	R"abline(h=bonferroni_threshold, lty=2, lwd=2, col=colours[5])";
	R"grid()";
	# QQ plot (p-values are assumed to be uniformly distributed)
	# observed_pval = runif(100); observed_pval = observed_pval[order(observed_pval, decreasing=FALSE)] ### TEST
	# sort the observed p-values
	R"observed_pval = gwas_out$PVALUES[order(gwas_out$PVALUES, decreasing=FALSE)]";
	# calculate the pdf of the observed p-values, i.e. the probability density for a p-value interval which corresponds to each observed p-value
	R"observed_pval_density = density(observed_pval, n=length(observed_pval), from=0, to=1) ### calculate the density of the observed p-values between 0 and 1 because we are dealing with p-values which are probabilities!";
	# calculate the cummulative probabilities of the observed p-values based on the pdf: where Prob(pvalue) = Density(pvalue | pvalue_interval) * pvalue_interval
	R"observed_pval_cumprob = cumsum(observed_pval_density$y * (observed_pval_density$x[2]-observed_pval_density$x[1]))";
	# calculate the expected quantiles based on the cummulative probabilities of the observed p-values
	R"expected_pval = qunif(p=observed_pval_cumprob, min=0, max=1) ### calculate the expected quantiles based on the observed cummulative probabilities";
	# transform into -log10 scale
	R"observed_lod = -log10(observed_pval + 1e-200)";
	R"expected_lod = -log10(expected_pval + 1e-200)";
	# plot
	R"plot(x=c(min(observed_lod), max(observed_lod)), y=c(min(observed_lod), max(observed_lod)), type='n', , main='QQ Plot', xlab=expression(Expected~~log[10]~(pvalue)), ylab=expression(Observed~~log[10]~(pvalue)))";
	R"points(x=expected_lod, y=observed_lod, type='p', pch=20, col=colours[5])";
	R"lines(x=c(0, max(observed_lod)), y=c(0, max(observed_lod)), lty=2, lwd=2, col='gray')";
	R"grid()";
	R"dev.off()";
	return(0)
end

end # end of significance_testing_module


using ADCME
using AdFem

include("chip_unstructured_solver.jl")
include("chip_unstructured_geometry.jl")
include("plot_inverse_one_iter.jl")

trialnum = 1

k_mold = 0.014531
k_chip_ref = 2.60475
k_air = 0.64357

function k_exact(x, y)
    k_mold + 1000 * k_chip_ref * (x-0.49)^2 / (1 + x^2)
end

nu = 0.47893  # equal to 1/Re
power_source = 82.46295  #82.46295 = 1.0e6 divide by air rho cp   #0.0619 = 1.0e6 divide by chip die rho cp
buoyance_coef = 299102.83

u_std = 0.001
p_std = 0.000001225
T_infty = 300

NT = 14    # number of iterations for Newton's method, 8 is good for m=400

k_chip = Variable(ones(length(chip_fem_idx))) .* k_mold

heat_source_fem = zeros(ndof)
heat_source_fem[chip_fem_top_idx] .= power_source
heat_source_gauss = dof_to_gauss_points(heat_source_fem, mesh)
heat_source = compute_fem_source_term1(constant(heat_source_gauss), mesh)

B = constant(compute_interaction_matrix(mesh)) # function not exist
Laplace = constant(compute_fem_laplace_matrix1(nu * constant(ones(ngauss)), mesh))

# apply Dirichlet to velocity and temperature; set left bottom two points to zero to fix rank deficient problem for pressure
bd = [bd; bd .+ ndof;
      fvm_bd .+ 2*ndof;
     bd .+ (2*ndof+nelem)]

# add solid region into boundary condition for u, v, p, i.e. exclude solid when solving Navier Stokes
bd = [bd; solid_fem_idx; solid_fem_idx .+ ndof; solid_fvm_idx .+ 2*ndof]

S0 = constant(zeros(nelem+3*ndof))
S = solve_navier_stokes(S0, NT, k_chip)
S_computed = S[end, :]
S_data = matread("fn$trialnum/xz_chip_unstructured_data.mat")["V"]

loss =  mean((S_computed .- S_data)^2)
loss = loss * 1e10
# g = gradients(loss, θ)

# ---------------------------------------------------
_loss = Float64[]
cb = (vs, iter, loss)->begin
    global _loss
    push!(_loss, loss)
    printstyled("[#iter $iter] loss=$loss\n", color=:green)
    if mod(iter, 10)==1
        close("all")
        pix_plot_velo_pres_temp_cond(iter)
        matwrite("fn$trialnum/pix_loss$iter.mat", Dict("iter"=>iter,"loss"=>_loss, "k_chip"=>vs[1]))
    end
end

sess = Session(); init(sess)
BFGS!(sess, loss, vars = [k_chip], callback = cb)


#module Compologicircuits

using Catlab.WiringDiagrams
using Catlab.CategoricalAlgebra
using Catlab.Theories
using Catlab.Syntax
using Catlab.GAT
using Catlab.Present
using Catlab.Programs
using Catlab.Graphics
import Catlab.Graphics: Graphviz
import Catlab.Theories: dom, codom, id, compose, otimes, braid, munit, mcopy, delete, pair, proj1, proj2, ⊗, ⋅, σ

# Helper function to show graphviz diagrams
show_diagram(d::WiringDiagram) = to_graphviz(d,
    orientation=LeftToRight,
    labels=false, label_attr=:xlabel,
    node_attrs=Graphviz.Attributes(
        :fontname => "Courier",
    ),
    edge_attrs=Graphviz.Attributes(
        :fontname => "Courier",
    )
)

# The skeleton of n-dimensional bool spaces which are the input/output types of logic circuits
struct CircuitDom
    n::Int
end

# Implement a logic circuit as a function from Vector{Bool} to Vector{Bool}
struct Circuit
    dom::CircuitDom
    codom::CircuitDom
    impl::Function # Vector{Bool} of size dom -> Vector{Bool} of size codom
end

# Convenience constructor
Circuit(nins::Int, nouts::Int, impl) = Circuit(CircuitDom(nins), CircuitDom(nouts), impl)

# Apply a circuit to the specified input vector
apply(circuit::Circuit, input::Vector{Bool})::Vector{Bool} = circuit.impl(input)

# Basic logic gate implementations
iNOT = Circuit(1, 1, x->map(b->!b, x))

iAND = Circuit(2, 1, x->[x[1] && x[2]])

iOR = Circuit(2, 1, x->[x[1] || x[2]])

#iδ = Circuit(1, 2, x->[x[1], x[1]])

# Make circuits compositional by implementing them as an SMC instance
@instance CartesianCategory{CircuitDom, Circuit} begin
    id(A::CircuitDom) = Circuit(A,A, x->x)
    dom(f::Circuit) = f.dom
    codom(f::Circuit) = f.codom

    # Circuit composition is just function composition
    compose(f::Circuit, g::Circuit) = begin
        @assert(f.codom == g.dom)
        return Circuit(f.dom, g.codom, x->g.impl(f.impl(x)))
    end

    otimes(A::CircuitDom, B::CircuitDom) = CircuitDom(A.n + B.n)
    # Monoidal product of circuits runs both circuits (TODO: run them in parallel to improve performance) and concatenates the results
    otimes(f::Circuit, g::Circuit) = begin
        impl = x -> vcat(f.impl(x[1:f.dom.n]), g.impl(x[f.dom.n + 1:end]))
        return Circuit(f.dom.n + g.dom.n, f.codom.n + g.codom.n, impl)
    end

    # A symmetric braiding just swaps the A and B parts of the input vector
    braid(A::CircuitDom, B::CircuitDom) = begin
        impl = x -> vcat(x[A.n+1:end], x[1:A.n])
        n = A.n + B.n
        Circuit(n, n, impl)
    end

    # Monoidal unit is a 0-dimensional bool space (i.e. a point or empty list)
    munit(::Type{CircuitDom}) = CircuitDom(0)

    # Stuff for Cartesian Category
    mcopy(A::CircuitDom) = Circuit(A.n, A.n+A.n, x->vcat(x,x))
    delete(A::CircuitDom) = Circuit(A.n, 0, x->Bool[])

    pair(f::Circuit, g::Circuit) = begin
        @assert(f.dom == g.dom)
        return Circuit(f.dom.n, f.codom.n+g.codom.n, x->vcat(f.impl(x), g.impl(x)))
    end
    proj1(A::CircuitDom, B::CircuitDom) = Circuit(A.n+B.n, A.n, x->x[1:A.n])
    proj2(A::CircuitDom, B::CircuitDom) = Circuit(A.n+B.n, B.n, x->x[A.n+1:end])
end

# Syntactic representation of logic circuits

@present Circuits(FreeCartesianCategory) begin
    B::Ob
    NOT::Hom(B,B)
    AND::Hom(B⊗B,B)
    OR::Hom(B⊗B,B)
    XOR::Hom(B⊗B,B)
    FA::Hom(B⊗B⊗B,B⊗B)
end

# Functor from circuit diagrams to circuit implementations
gens = Dict(
    Circuits[:B]=>CircuitDom(1),
    Circuits[:NOT]=>iNOT,
    Circuits[:AND]=>iAND,
    Circuits[:OR]=>iOR
)
Impl(expr) = functor((CircuitDom, Circuit), expr, generators=gens)

# Make some circuits

XOR = @program Circuits (x::B, y::B) begin
    xnandy = NOT(AND(x,y))
    xory = OR(x,y)
    return AND(xnandy,xory)
end

XOR_expr = to_hom_expr(FreeCartesianCategory, XOR)

iXOR = Impl(XOR_expr)

# Print XOR truth table

print_xor() = for i in Iterators.product([false, true], [false, true])
    input = collect(i)
    output = apply(iXOR, input)
    println("XOR on $input = $output")
end

full_adder = @program Circuits (x::B, y::B, cin::B) begin
    z = XOR(x,y)
    s = XOR(z,cin)
    cout = OR(AND(z,cin),AND(x,y))
    return (s, cout)
end

AND = to_wiring_diagram(Circuits[:AND])
NOT = to_wiring_diagram(Circuits[:NOT])
OR = to_wiring_diagram(Circuits[:OR])
total_full_adder = ocompose(full_adder, [XOR, XOR, AND, AND, OR])
iFA = Impl(to_hom_expr(FreeCartesianCategory, total_full_adder))

# Print Full Adder truth table

print_fa() = for i in Iterators.product([false, true], [false, true], [false, true])
    input = collect(i)
    output = apply(iFA, input)
    println("FA on $input = $output")
end

four_bit_adder = @program Circuits (cin::B, x1::B, x2::B, x3::B, x4::B, y1::B, y2::B, y3::B, y4::B) begin
    s1, cout1 = FA(x4,y4,cin)
    s2, cout2 = FA(x3,y3,cout1)
    s3, cout3 = FA(x2,y2,cout2)
    s4, cout = FA(x1,y1,cout3)
    return (cout, s4, s3, s2, s1)
end

total_four_bit_adder = ocompose(four_bit_adder, [total_full_adder, total_full_adder, total_full_adder, total_full_adder])
tfba_expr = to_hom_expr(FreeCartesianCategory, total_four_bit_adder)
i4bitadder = Impl(tfba_expr)

sum = apply(i4bitadder,
    Bool[0,      #carry in
        0,1,1,1, #x
        0,1,1,0] #y
)

### Perform sub-circuit matching.
# The theory of directed wiring diagrams without outer ports.
@present TheoryNotOpenWiringDiagram(FreeSchema) begin
  Box::Ob
  (InPort, OutPort)::Ob
  Wire::Ob
  #(InWire, OutWire)::Ob
  PortValue::AttrType
  WireValue::AttrType
  BoxValue::AttrType
  BoxType::AttrType

  # To which in/out port is a wire attached?
  src::Hom(Wire, OutPort)
  tgt::Hom(Wire, InPort)

  in_port_type::Attr(InPort, PortValue)
  out_port_type::Attr(OutPort, PortValue)

  value::Attr(Box, BoxValue)
  box_type::Attr(Box, BoxType)
  wire_value::Attr(Wire, WireValue)
  #in_wire_value::Attr(InWire, WireValue)
  #out_wire_value::Attr(OutWire, WireValue)

  #in_tgt::Hom(InWire, InPort)
  #out_src::Hom(OutWire, OutPort)

  # To which box is an in/out port attached?
  in_port_box::Hom(InPort, Box)
  out_port_box::Hom(OutPort, Box)
end
F = @finfunctor TheoryNotOpenWiringDiagram DirectedWiringDiagrams.TheoryAttributedWiringDiagram begin
  Box => Box
  InPort => InPort
  OutPort => OutPort
  Wire => Wire
  #InWire => InWire
  #OutWire => OutWire

  PortValue => PortValue
  WireValue => WireValue
  BoxValue => BoxValue
  BoxType => BoxType

  src => src
  tgt => tgt

  in_port_type => in_port_type
  out_port_type => out_port_type

  value => value
  box_type => box_type
  wire_value => wire_value
  #in_wire_value => in_wire_value
  #out_wire_value => out_wire_value

  in_tgt => in_tgt
  out_src => out_src

  in_port_box => in_port_box
  out_port_box => out_port_box
end

@acset_type NOWDACSet(TheoryNotOpenWiringDiagram)

deltaF = DeltaMigration(F, NOWDACSet{Any, Any, Any, DataType})

cAND = deltaF(AND.diagram)
cOR = deltaF(OR.diagram)
cNOT = deltaF(NOT.diagram)
cXOR = deltaF(XOR.diagram)

## Pick out AND's in AND.
println("There are ", length(homomorphisms(cAND, cAND, monic=true)), " monic homomorphisms from AND to AND.")

## Pick out OR's in AND.
println("There are ", length(homomorphisms(cOR, cAND, monic=true)), " monic homomorphisms from OR to AND.")

# Pick out AND's in XOR.
println("There are ", length(homomorphisms(cAND, cXOR, monic=true)), " monic homomorphisms from AND to XOR.")

# Pick out circuits in total_four_bit_adder that can be replaced by FA's.
println("There are ", length(homomorphisms(deltaF(total_full_adder.diagram), deltaF(total_four_bit_adder.diagram), monic=true)), " monic homomorphisms from total_full_adder to total_four_bit_adder.")

#end


using DiscontinuousGherkins
using DiscontinuousGherkins: jacobi_gauss_lobatto, vandermonde_1d, Dmatrix1D

NODETOL = 1e-10

N = 8  # Order of polymomials used for approximation
K = 10 # Number of elements

domain = (0, 2)
grid = DiscontinuousGalerkin1DGrid(domain, K)

Np = N+1
Nfp = 1
Nfaces = 2

r = jacobi_gauss_lobatto(0, 0, N)

V = vandermonde_1d(N, r)
V⁻¹ = inv(V)

Dr = Dmatrix1D(N, r, V)

Emat = zeros(Np, Nfaces * Nfp)
Emat[1,  1] = 1
Emat[Np, 2] = 1
lift = V * (V' * Emat)

# build coordinates of all the nodes
VX = grid.node_coordinates
EToV = grid.element_to_node_connectivity
va = EToV[:, 1]'
vb = EToV[:, 2]'
x = ones(N+1, 1) * VX[Int.(va)] .+ 0.5 .* (r .+ 1) * (VX[Int.(vb)] - VX[Int.(va)])

# Compute the metric elements for the local mappings of the 1D elements.
xr = Dr * x
J = xr
rx = 1 ./ J

# Compute masks for edge nodes
Fmask = BitArray(undef, size(x)...)
Fmask .= 0
Fmask[1, 1] = 1
Fmask[end, end] = 1
Fx = x[Fmask]

# % Build surface normals and inverse metric at surface
nx = zeros(Nfp * Nfaces, K)
nx[1, :] .= -1
nx[2, :] .= +1

Fscale = 1 ./ J[Fmask]

## Build global connectivity arrays for 1D grid based on standard EToV input array from grid generator

Nfaces = 2

# Find number of elements and vertices
TotalFaces = Nfaces * K
Nv = K+1

# List of local face to local vertex connections
vn = Int[1, 2]

# Build global face to node sparse array
SpFToV = zeros(TotalFaces, Nv)
sk = 1;
for k in 1:K, face in 1:Nfaces
    SpFToV[sk, EToV[k, vn[face]] |> Int] = 1
end

# Build global face to global face sparse array
SpFToF = SpFToV * SpFToV' - I(TotalFaces)

% Find complete face to face connections
[faces1, faces2] = find(SpFToF==1);

% Convert face global number to element and face numbers
element1 = floor( (faces1-1)/Nfaces ) + 1;
face1 = mod( (faces1-1), Nfaces ) + 1;
element2 = floor( (faces2-1)/Nfaces ) + 1;
face2 = mod( (faces2-1), Nfaces ) + 1;

% Rearrange into Nelements x Nfaces sized arrays
ind = sub2ind([K, Nfaces], element1, face1);
EToE = (1:K)’*ones(1,Nfaces);
EToF = ones(K,1)*(1:Nfaces);
EToE(ind) = element2; EToF(ind) = face2;
return


function getCoefficientsRadIIA2(T=Float64)
    a = [[5//12  -1//12]
         [3//4    1//4 ]]
    b = [3//4, 1//4]
    c = [1//3, 1//1]

    CoefficientsRK(T, :RadIIA2, 3, a, b, c)
end

function getCoefficientsRadIIA3(T=Float64)
    a = @dec128 [
          [11/45 -  7*√6/360    37/225-169*√6/1800   -2/225+√6/75]
          [37/225+169*√6/1800   11/45 +  7*√6/360    -2/225-√6/75]
          [ 4/9  -    √6/36      4/9  +    √6/36      1/9        ]
        ]
    b = @dec128 [4/9-√6/36,   4/9+√6/36,  1/9 ]
    c = @dec128 [2/5-√6/10,   2/5+√6/10,  1   ]

    CoefficientsRK(T, :RadIIA3, 5, a, b, c)
end


using Distributions

function multiple_NB_dgp(RNG, N)
    log_sq_foot = Vector{Float64}(undef, N)  # Use zeros?
    live_in_super = Vector{Float64}(undef, N)  # Use zeros?
    traps = Vector{Int64}(undef, N)  # Use zeros?
    complaints = Vector{Int64}(undef, N)  # Use zeros?

    α = rand(RNG, Normal(log(4), .1))
    β = rand(RNG, Normal(-0.25, .1))
    β_super = rand(RNG, Normal(-0.5, 1))
    inv_ϕ = abs(rand(RNG, Normal(0, 1))) #TODO: fabs == abs?

    for i in 1:N
        log_sq_foot[i] = rand(RNG, Normal(1.5, 0.1))
        live_in_super[i] = rand(RNG, Bernoulli(0.5))
        traps[i] = rand(RNG, Poisson(8))
        r = α + log_sq_foot[i] + β * traps[i] + β_super * live_in_super[i]
        p = inv(inv_ϕ)

        # TODO: Log parametrization
        # complaints[i] = rand(RNG, NegativeBinomial(r, p))
    end

    return (α = α,
            β = β,
            β_super = β_super,
            log_sq_foot = log_sq_foot,
            live_in_super = live_in_super,
            traps = traps,
            complaints = complaints)
end


################################
# Complex Operator Definitions #
################################
@present ExtendedOperators <: ExtendedFreeExtCalc2D begin
  X::Space
  neg::Hom(DualForm1(X), DualForm1(X)) # negative
  neg₁::Hom(Form1(X), Form1(X)) # negates 1-forms
  dneg₁::Hom(DualForm1(X), DualForm1(X)) # negates 1-forms
  L₀::Hom(Form1(X)⊗DualForm2(X), DualForm2(X))
  L₁::Hom(Form1(X)⊗DualForm1(X), DualForm1(X))
  i₀::Hom(Form1(X)⊗DualForm2(X), DualForm1(X))
  i₁::Hom(Form1(X)⊗DualForm1(X), DualForm0(X))

  L₁′::Hom(Form1(X)⊗Form1(X), Form1(X))
  i₀′::Hom(Form1(X)⊗Form2(X), Form1(X))
  i₁′::Hom(Form1(X)⊗Form1(X), Form0(X))
  ∧₁₁′::Hom(Form1(X)⊗DualForm1(X), DualForm2(X))
  ∧₁₀′::Hom(Form1(X)⊗DualForm0(X), Form1(X))
end

rules = gen_dec_rules()

# Lie derivative between two primal 1-forms
Lie1Imp′ = @decapode ExtendedOperators begin
  (F1, F1′, Ḟ1)::Form1{X}
  Ḟ1 == i₀′(F1, d₁{X}(F1′)) + d₀{X}(i₁′(F1, F1′))
end
lie1_imp′ = diag2dwd(Lie1Imp′, in_vars = [:F1, :F1′], out_vars = [:Ḟ1])
rules[:L₁′] = lie1_imp′

# Internal product between a primal 1-form and a primal 2-form
I0Imp′ = @decapode ExtendedOperators begin
  (F1, Ḟ1)::Form1{X}
  F2::Form2{X}
  Ḟ1 == neg₁(∧₁₀′(F1, ⋆₂{X}(F2)))
end
i0_imp′ = diag2dwd(I0Imp′, in_vars = [:F1, :F2], out_vars = [:Ḟ1])
rules[:i₀′] = i0_imp′

# Internal product between two primal 1-forms
I1Imp′ = @decapode ExtendedOperators begin
  (F1, F1′)::Form1{X}
  F0::Form0{X}
  F0 == ⋆₀⁻¹{X}(∧₁₁′(F1, ⋆₁{X}(F1′)))
end
i1_imp′ = diag2dwd(I1Imp′, in_vars = [:F1, :F1′], out_vars = [:F0])
rules[:i₁′] = i1_imp′