using ParameterizedFunctions, OrdinaryDiffEq, DiffEqParamEstim
using BlackBoxOptim, NLopt, Plots,QuadDIRECT
gr(fmt=:png)

loc_bounds = Tuple{Float64,Float64}[(0, 1), (0, 1), (0, 1), (0, 1)]
glo_bounds = Tuple{Float64,Float64}[(0, 5), (0, 5), (0, 5), (0, 5)]
loc_init = [0.5,0.5,0.5,0.5]
glo_init = [2.5,2.5,2.5,2.5]

fitz = @ode_def FitzhughNagumo begin
  dv = v - v^3/3 -w + l
  dw = τinv*(v +  a - b*w)
end a b τinv l

p = [0.7,0.8,0.08,0.5]              # Parameters used to construct the dataset
r0 = [1.0; 1.0]                     # initial value
tspan = (0.0, 30.0)                 # sample of 3000 observations over the (0,30) timespan
prob = ODEProblem(fitz, r0, tspan,p)
tspan2 = (0.0, 3.0)                 # sample of 300 observations with a timestep of 0.01
prob_short = ODEProblem(fitz, r0, tspan2,p)

dt = 30.0/3000
tf = 30.0
tinterval = 0:dt:tf
t  = collect(tinterval)

h = 0.01
M = 300
tstart = 0.0
tstop = tstart + M * h
tinterval_short = 0:h:tstop
t_short = collect(tinterval_short)

#Generate Data
data_sol_short = solve(prob_short,Vern9(),saveat=t_short,reltol=1e-9,abstol=1e-9)
data_short = convert(Array, data_sol_short) # This operation produces column major dataset obs as columns, equations as rows
data_sol = solve(prob,Vern9(),saveat=t,reltol=1e-9,abstol=1e-9)
data = convert(Array, data_sol)

plot(data_sol_short)

plot(data_sol)

obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),tstops=t_short)
res1 = bboptimize(obj_short;SearchRange = glo_bounds, MaxSteps = 7e3)
# Lower tolerance could lead to smaller fitness (more accuracy)

obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),tstops=t_short,reltol=1e-9)
res1 = bboptimize(obj_short;SearchRange = glo_bounds, MaxSteps = 7e3)
# Change in tolerance makes it worse

obj_short = build_loss_objective(prob_short,Vern9(),L2Loss(t_short,data_short),tstops=t_short,reltol=1e-9,abstol=1e-9)
res1 = bboptimize(obj_short;SearchRange = glo_bounds, MaxSteps = 7e3)
# using the moe accurate Vern9() reduces the fitness marginally and leads to some increase in time taken

obj_short = build_loss_objective(prob_short,Vern9(),L2Loss(t_short,data_short),tstops=t_short,reltol=1e-9,abstol=1e-9)

opt = Opt(:GN_ORIG_DIRECT_L, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

opt = Opt(:GN_CRS2_LM, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

opt = Opt(:GN_ISRES, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

opt = Opt(:GN_ESCH, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

opt = Opt(:LN_BOBYQA, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LN_NELDERMEAD, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LD_SLSQP, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LN_COBYLA, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LN_NEWUOA_BOUND, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LN_PRAXIS, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LN_SBPLX, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LD_MMA, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

obj = build_loss_objective(prob,Vern9(),L2Loss(t,data),tstops=t,reltol=1e-9,abstol=1e-9)
res1 = bboptimize(obj;SearchRange = glo_bounds, MaxSteps = 4e3)

opt = Opt(:GN_ORIG_DIRECT_L, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

opt = Opt(:GN_CRS2_LM, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 20000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

opt = Opt(:GN_ISRES, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 50000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

opt = Opt(:GN_ESCH, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 20000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

opt = Opt(:LN_BOBYQA, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LN_NELDERMEAD, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-9)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

opt = Opt(:LD_SLSQP, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),tstops=t_short)
lower = [0,0,0,0]
upper = [1,1,1,1]
splits = ([0,0.3,0.7],[0,0.3,0.7],[0,0.3,0.7],[0,0.3,0.7])
@time root, x0 = analyze(obj_short,splits,lower,upper)

minimum(root)

obj = build_loss_objective(prob,Vern9(),L2Loss(t,data),tstops=t,reltol=1e-9,abstol=1e-9)
lower = [0,0,0,0]
upper = [5,5,5,5]
splits = ([0,0.5,1],[0,0.5,1],[0,0.5,1],[0,0.5,1])
@time root, x0 = analyze(obj_short,splits,lower,upper)

minimum(root)


using Stats
using Base.Test

## moments

wv = weights(ones(5) * 2.0)

@test_approx_eq skewness(1:5) 0.0
@test_approx_eq skewness([1, 2, 3, 4, 5]) 0.0
@test_approx_eq skewness([1, 2, 2, 2, 5]) 1.1731251294063556
@test_approx_eq skewness([1, 4, 4, 4, 5]) -1.1731251294063556

@test_approx_eq skewness([1, 2, 2, 2, 5], wv) 1.1731251294063556

@test_approx_eq kurtosis(1:5) -1.3
@test_approx_eq kurtosis([1, 2, 3, 4, 5]) -1.3
@test_approx_eq kurtosis([1, 2, 3, 3, 2]) -1.1530612244897953

@test_approx_eq kurtosis([1, 2, 3, 4, 5], wv) -1.3

## variability

@test_approx_eq variation([1:5]) 0.527046276694730

@test_approx_eq sem([1:5]) 0.707106781186548

@test_approx_eq mad([1:5], 3) 1.4826
@test_approx_eq mad([1:5]) 1.4826
@test_approx_eq mad(1:5) 1.4826

# min/max related

@test minmax(1:5) == (1,5)
@test minmax([3, 2, 7, 1, 4]) == (1,7)
@test minmax([3.0, NaN, 1.0, NaN, 2.0]) == (1.0, 3.0)

@test_approx_eq midrange([1, 2, 3, 4, 5]) 3.0
@test_approx_eq midrange([1, 2, 2, 2, 5]) 3.0
@test_approx_eq midrange([1, 4, 4, 4, 5]) 3.0
@test_approx_eq midrange([1:5]) 3.0

@test_approx_eq range([1, 2, 3, 4, 5]) 4.0
@test_approx_eq range([1, 2, 2, 2, 5]) 4.0
@test_approx_eq range([1, 4, 4, 4, 5]) 4.0
@test_approx_eq range([1:15]) 14.0

# quantile & friends

@test_approx_eq quantile(1:5) [1:5]
@test_approx_eq nquantile(1:5, 2) [1, 3, 5]
@test_approx_eq nquantile(1:5, 4) [1:5]

@test_approx_eq percentile([1:5], 25) 2.0
@test_approx_eq percentile([1:5], [25, 50, 75]) [2.0, 3.0, 4.0]
@test_approx_eq iqr(1:5) 2.0

# mode & modes

@test mode([1, 2, 3, 3, 2, 2, 1], 1:3) == 2
@test modes([1, 2, 3, 3, 2, 2, 1], 1:3) == [2]
@test modes([1, 3, 2, 3, 3, 2, 2, 1], 1:3) == [2, 3]

@test mode([1, 2, 3, 3, 2, 2, 1]) == 2
@test modes([1, 2, 3, 3, 2, 2, 1]) == [2]
@test sort(modes([1, 3, 2, 3, 3, 2, 2, 1])) == [2, 3]

# summarystats

s = summarystats(1:5)
@test isa(s, Stats.SummaryStats)
@test s.min == 1.0
@test s.max == 5.0
@test_approx_eq s.mean 3.0
@test_approx_eq s.median 3.0
@test_approx_eq s.q25 2.0
@test_approx_eq s.q75 4.0


# Copyright (c) 2021
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

using LinearAlgebra;

"""
    L0EM(A::Matrix{Float64}, b::Vector{Float64}; maxiter=100, epsilon=0.01)

Find the solution to Ax=b using an efficient EM algoritm that directly solves
the L0 optimization problem.

### Input

- `A`       -- Matrix: Ax=b
- `b`       -- Vector: Ax=b
- `maxiter` -- (optional) number of optmization iterations
- `epsilon` -- (optional) threshold to stop optimizing
- `lambda`  -- (optional) regularization

### Output

Solution to Ax=b (Vector{Float64})

### Example

julia> A = randn(10, 100);
julia> b = randn(10);
julia> x = L0EM(A, b, maxiter=5, epsilon=.01);
julia> println(x)

### Algorithm

This function implements Liu and Li's L0EM algorithms in https://arxiv.org/pdf/1407.7508v1.pdf.

"""
function L0EM(A::Matrix{Float64},
              b::Vector{Float64};
              lambda=.001,
              epsilon=.001,
              maxiter=50)
    local n, p, theta, A_eta, eta;
    eps_stop = .01;
    eps_zero = .01;
    n, p = size(A);
    eye = Matrix{Float64}(I, p, p);

    # get the initial solution
    theta = inv(A'*A + lambda*eye)*A'*b;

    # continue to optimize theta
    for i = 1:maxiter
        eta = theta;

        # copy the squared eta terms into a matrix so we can use the hammard product.
        A_eta = repeat(eta.^2, 1, n)';
        A_eta = A_eta.*A;

        # update theta
        theta = inv(A_eta'*A + lambda*eye)*A_eta'*b;

        # check to break the loop due to a small difference in norm
        if norm(theta-eta, 2) <= epsilon
            break;
        end
    end
    x = theta;
    i = abs.(x) .< epsilon;
    x[i] = zeros(sum(i));
    return x;
end

# A = randn(10, 50); b = randn(10);
# x = L0EM(A, b);
# println(x)


# setting working directory
cd("/Volumes/SSD Hans/Github/MacroFall2020/AdvMacro")

#Defining vars
a = 1/3
b = 0.6
k_star = (a*b)^(1/(1-a))
r = 1/b

# Initializing array for loop
m, n = 100,2
K = fill(0.0, (m, n))

#defining policy fn
g = x -> a*b*(x^a)

# loop for capital
k = 0.8*k_star
for i in 1:m
    K[i, 1] = k
    K[i, 2] = g(K[i,1])
    global k = K[i, 2]
end

# defining fn's for the rest of the variables
w = x -> (1-a)x^a
y = x -> x^a
c = (p,q) -> y(p)-q
r = x -> a*(x^(a-1))
int = r.(K[1:100,2])
wage = w.(K[1:100,1])
output = y.(K[1:100,1])
cons = c.(K[1:100,1],K[1:100,2])

# importing packages for plotting
import Pkg; Pkg.add("Plots")
using Plots
gr()

# plotting capital
plot(1:10,fill(k_star,10), label="capital SS")
plot!(1:10,K[1:10,1], label="capital path", linestyle = :dot, legend = :bottomright)
xlabel!("t")
title!("Capital")
savefig("graphs/capital")

# plotting output

plot(1:10,fill(y(k_star),10), label="output SS")
plot!(1:10,output[1:10], label="output path", linestyle = :dot, legend = :bottomright)
xlabel!("t")
title!("Output")
savefig("graphs/output")

# plotting consumption

plot(1:10,fill(c(k_star,k_star),10), label="consumption SS")
plot!(1:10,cons[1:10], label="consumption path", linestyle = :dot, legend = :bottomright)
xlabel!("t")
title!("Consumption")
savefig("graphs/consumption")

# plotting wage

plot(1:10,fill(w(k_star),10), label="wage SS")
plot!(1:10,wage[1:10], label="wage path", linestyle = :dot, legend = :bottomright)
xlabel!("t")
title!("Wage")
savefig("graphs/wage")

# plotting capital rental rate
plot(1:10,fill(r(k_star),10), label="capital rental rate SS")
plot!(1:10,int[1:10], label="rate path", linestyle = :dot, legend = :topright)
xlabel!("t")
title!("Capital rental rate")
savefig("graphs/capitalr")

# Increase in productivity permanently by 5%

k_starp = (a*b*1.5)^(1/(1-a))
#defining policy fn with productivity increases
g1 = x -> a*b*1.05*(x^a)

# Initializing array for loop
K1 = fill(0.0, (m, n))

# loop for capital
k1 = k_starp
for i in 1:m
    K1[i, 1] = k1
    K1[i, 2] = g(K[i,1])
    global k1 = K[i, 2]
end

# defining fn's for the rest of the variables
w1 = x -> (1-a)*(x^a) #does not change
y1 = x -> 1.5*(x^a)
c1 = (p,q) -> y(p)-q #does not change
r1 = x -> a*(1.5*x^(a-1))
int_p = r1.(K1[1:100,2])
wage_p = w1.(K1[1:100,1])
output_p = y1.(K1[1:100,1])
cons_p = c1.(K1[1:100,1],K1[1:100,2])

# plotting capital
plot(1:10,fill(k_star,10), label="capital SS")
plot!(1:10,K1[1:10,1], label="capital path", linestyle = :dot, legend = :bottomright)
xlabel!("t")
title!("Capital (productivity increase)")
savefig("graphs/capitalp")

# plotting output

plot(1:10,fill(y1(k_star),10), label="output SS")
plot!(1:10,output_p[1:10], label="output path", linestyle = :dot, legend = :bottomright)
xlabel!("t")
title!("Output (productivity increase)")
savefig("graphs/outputp")

# plotting consumption

plot(1:10,fill(c1(k_star,k_star),10), label="consumption SS")
plot!(1:10,cons_p[1:10], label="consumption path", linestyle = :dot, legend = :bottomright)
xlabel!("t")
title!("Consumption (productivity increase)")
savefig("graphs/consumptionp")

# plotting wage

plot(1:10,fill(w1(k_star),10), label="wage SS")
plot!(1:10,wage_p[1:10], label="wage path", linestyle = :dot, legend = :bottomright)
xlabel!("t")
title!("Wage (productivity increase)")
savefig("graphs/wagep")

# plotting capital rental rate
plot(1:10,fill(r1(k_star),10), label="capital rental rate SS")
plot!(1:10,int_p[1:10], label="rate path", linestyle = :dot, legend = :topright)
xlabel!("t")
title!("Capital rental rate (productivity increase)")
savefig("graphs/capitalrp")


#=

  Euler #34 in Julia.

  Problem 34
  """
  145 is a curious number, as 1! + 4! + 5! = 1 + 24 + 120 = 145.

  Find the sum of all numbers which are equal to the sum of the
  factorial of their digits.

  Note: as 1! = 1 and 2! = 2 are not sums they are not included.
  """

  This Julia program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

include("Euler.jl")

function factorial_sum(n)
    # return (split(string(n),"").|>i->factorial(parse(Int,i)))|>sum
    return (digits(n).|>i->factorial(i))|>sum

end

# 0.01974533s
function euler34a()
    s = 0;
    for n in 10:100000
        if n == factorial_sum(n)
            s += n
        end
    end
    return s
end

# 0.02070016s
function euler34b()
    return filter(n->n == factorial_sum(n),10:100000)|>sum
end

# 0.01983716s
function euler34c()
    return [n for n in 10:100000 if n == factorial_sum(n)]|>sum
end

# run_euler(euler34a)
# run_euler(euler34b)
run_euler(euler34c)


using RobustShortestPath

data = [
 1   4  79   31  66  28;
 1   2  59   97  41  93;
 2   4  31   21  50  40;
 2   3  90   52  95  38;
 2   5   9   23  95  59;
 2   6  32   57  73   7;
 3   9  89  100  38  21;
 3   8  66   13   4  72;
 3   6  68   95  58  58;
 3   7  47   12  56  20;
 4   3  14   19  36  84;
 4   9  95   65  88  42;
 4   8  88   13  62  54;
 5   3  44    8  62  53;
 5   6  83   66  30  19;
 6   7  33    3   7   8;
 6   8  37   99  29  46;
 7  11  79   54  23   3;
 7  12  10   37  35  43;
 8   7  95   71  85  56;
 8  10   0   95  16  64;
 8  12  30   38  16   3;
 9  10   5   69  51  71;
 9  11  44   60  60  17;
10  13  79   78  16  59;
10  14  91   59  64  61;
11  14  53   38  84  77;
11  15  80   85  78   6;
11  13  56   23  26  85;
12  15  75   80  31  38;
12  14   1  100  18  40;
13  14  48   28  45  33;
14  15  25   71  33  56;
]

start_node = data[:,1] #first column of data
end_node = data[:,2] #second column of data
p = data[:,3] #third
q = data[:,4] #fourth
c = data[:,5] #fifth
d = data[:,6] #sixth

# Link length vectors for single-uncertain-coefficient robust shortest paths
cc = p.*c
dd = (p+q).*(c+d) - p.*c

# Setting origin and destination nodes
origin = 1
destination = 15
println("Origin=$origin, Destination=$destination")

println("----------------------------------------------------")
println("Single Coefficient Case")
# For each Gamma from 0 to 6, obtain the robust shortest path
known_cost_solution_one = Dict(0 => 6060, 1=> 15024, 2=> 20864, 3=> 26604, 4=> 31293, 5=> 32291)
for Gamma=0:5
	robust_path, robust_x, worst_case_cost = get_robust_path(start_node, end_node, cc, dd, Gamma, origin, destination)
	println("Gamma=$Gamma: Robust Path is $(robust_path') and the worst-case cost is $worst_case_cost.")

    @assert worst_case_cost == known_cost_solution_one[Gamma]
end

println("----------------------------------------------------")
println("Two Coefficient Case")

for Gamma_u=1:5
    for Gamma_v=1:5
        robust_path, robust_x, worst_case_cost = get_robust_path_two(start_node, end_node, p, q, c, d, Gamma_u, Gamma_v, origin, destination)
        println("(Gamma_u,Gamma_v)=($Gamma_u,$Gamma_v): Robust Path is $(robust_path') and the worst-case cost is $worst_case_cost.")
    end
end

test_Gamma_u = 5
test_Gamma_v = 5
robust_path, robust_x, worst_case_cost = get_robust_path_two(start_node, end_node, p, q, c, d, test_Gamma_u, test_Gamma_v, origin, destination)

println(worst_case_cost)
@assert worst_case_cost==32291.0


using KrylovMethods
using Base.Test
include("getDivGrad.jl")

println("=== Testing SSOR for real matrix ===")
# small full system

# CG: test sparse Laplacian
A = getDivGrad(32,32,32)
rhs = randn(size(A,1))
tol = 1e-2

# tests with A being matrix
x1,flag1,relres1,iter1,resvec1       = ssor(A,rhs,tol=tol,maxIter=100)
x2,flag2,relres2,iter2,resvec2       = ssor(A,rhs,tol=tol,maxIter=100,omega=2/3,out=1)
x3,flag3,relres3,iter3,resvec3     = ssor(A,rhs,tol=tol,maxIter=100,omega=1/2,out=1,storeInterm=true)

# test relative residuals
@test norm(A*x1-rhs)/norm(rhs) <= tol
@test norm(A*x2-rhs)/norm(rhs) <= tol
@test norm(A*x3[:,end]-rhs)/norm(rhs) <= tol

println("=== Testing SSOR as Preconditioner ===")
omega = 1.2;
d = omega./diag(A);
x = zeros(length(rhs)) # pre allocation for the preconditioner result.
PC = r -> (x[:]=0.0; return ssorPrecTrans!(A,x,r,d));
y = KrylovMethods.cg(A,rhs,tol=1e-12,maxIter=200,M=PC,out=1)[1]
@test norm(A*y-rhs)/norm(rhs) <= 1e-12

println("=== Testing SSOR for complex matrix ===")
# small full system

# CG: test sparse Laplacian
A = getDivGrad(32,32,32)
A += 1im*speye(size(A,1))
rhs = randn(size(A,1)) + 1im * randn(size(A,1))
tol = 1e-2

#
x0,flag0,relres0,iter0,resvec0       = ssor(A,rhs,tol=tol,maxIter=3,out=2)
@test flag0==-1
@test iter0==3
@test length(resvec0)==3

# tests with A being matrix
x1,flag1,relres1,iter1,resvec1       = ssor(A,rhs,tol=tol,maxIter=100)
x2,flag2,relres2,iter2,resvec2       = ssor(A,rhs,tol=tol,maxIter=100,omega=2/3,out=1)
x3,flag3,relres3,iter3,resvec3       = ssor(A,rhs,tol=tol,maxIter=100,omega=1/2,out=1)

# test relative residuals
@test norm(A*x1-rhs)/norm(rhs) <= tol
@test norm(A*x2-rhs)/norm(rhs) <= tol
@test all(diff(resvec1).<0)
@test all(diff(resvec2).<0)
@test all(diff(resvec3).<0)

println("=== SSOR : All tests passed. ====")


using DSGE
using Base.Test

# Test `hessizero` in context of Rosenbrock function
function rosenbrock(x::Vector)
    a = 100
    b = 1
    return a*(x[2]-x[1]^2.0)^2.0 + b*(1-x[1])^2.0
end

function rosenbrock_hessian(x::Vector)
    H = zeros(eltype(x), 2, 2)
    H[1,1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2.0
    H[1,2] = -400.0 * x[1]
    H[2,1] = -400.0 * x[1]
    H[2,2] = 200.0
    return H
end

# At the min, ensure no negatives in diagonal
x0 = [1.0, 1.0]
hessian_expected = rosenbrock_hessian(x0)
hessian, _ = DSGE.hessizero(rosenbrock, x0; check_neg_diag=true)
@testset "Check valid hessian at the minimum" begin
    @test_matrix_approx_eq hessian_expected hessian
end

# Not at the min (indeed, we are at max), ensure throws error
x1 = [1.0, 1.0]
rosenbrock_neg(x) = -rosenbrock(x)
@testset "Throw error for hessian calculation not at the min" begin
    @test_throws Exception hessian, _ = DSGE.hessizero(rosenbrock_neg, x1; check_neg_diag=true)
end

# Not at the min, check matches closed form
x1 = Vector[[0.5, 1.5], [-1.0, -1.0]]
@testset "Check closed-form of the hessian" begin
    for x in x1
        hessian_expected = rosenbrock_hessian(x)
        hessian, _ = DSGE.hessizero(rosenbrock, x)
        @test @test_matrix_approx_eq hessian_expected hessian
    end
end


"""
    conductance(h::Hypergraph, subset::Set{Int})::Float64

Calculate unweighted hypergraph conductance of `subset`.
note: ∅ ⊊ `subset` ⊊ `1:nhv(h)`

For more information see `1. Introduction` at:
Spectral Properties of Hypergraph Laplacian and Approximation Algorithms
auhtors: T-H. Hubert Chan, Anand Louis, Zhihao Gavin Tang, and Chenzi Zhang
"""
function conductance(h::Hypergraph, subset::Set{Int})::Float64
  if isempty(subset) error("`subset` is not allowed empty.") end
  if subset == Set(1:nhv(h)) error("`subset` is not allowed true subset of `h`.") end
  subset2 = setdiff(Set(1:nhv(h)), subset)
  volS = sum([sum(values(gethyperedges(h, node))) for node in subset])
  volV_S = sum([sum(values(gethyperedges(h, node))) for node in subset2])
  cutS = 0
  for he in 1:nhe(h)
    he_vertices = keys(getvertices(h, he))
    if isempty(intersect(he_vertices, subset)) continue end
    if isempty(intersect(he_vertices, subset2)) continue end
    cutS += 1
  end
  return cutS / min(volS, volV_S)
end


"""
    smooth(x_curr::Gaussian, x_next_smoothed::Gaussian, Ah::AbstractMatrix, Qh::AbstractMatrix)

SMOOTH step of the (extended) Kalman smoother, or (extended) Rauch-Tung-Striebel smoother.
It is implemented in Joseph Form:
```math
m_{n+1}^P = A(h)*m_n
P_{n+1}^P = A(h)*P_n*A(h) + Q(h)

G = P_n * A(h)^T * (P_{n+1}^P)^{-1}
m_n^S = m_n + G * (m_{n+1}^S - m_{n+1}^P)
P_n^S = (I - G*A(h)) P_n (I - G*A(h))^T + G * Q(h) * G + G * P_{n+1}^S * G
```
"""
function smooth(
    x_curr::Gaussian,
    x_next_smoothed::Gaussian,
    Ah::AbstractMatrix,
    Qh::AbstractMatrix,
)
    x_pred = predict(x_curr, Ah, Qh)

    P_p = x_pred.Σ
    P_p_inv = inv(P_p)

    G = x_curr.Σ * Ah' * P_p_inv

    smoothed_mean = x_curr.μ + G * (x_next_smoothed.μ - x_pred.μ)
    smoothed_cov =
        (X_A_Xt(x_curr.Σ, (I - G * Ah)) + X_A_Xt(Qh, G) + X_A_Xt(x_next_smoothed.Σ, G))
    x_curr_smoothed = Gaussian(smoothed_mean, smoothed_cov)
    return x_curr_smoothed, G
end
function smooth(
    x_curr::SRGaussian,
    x_next_smoothed::SRGaussian,
    Ah::AbstractMatrix,
    Qh::SRMatrix,
)
    x_pred = predict(x_curr, Ah, Qh)

    P_p = x_pred.Σ
    P_p_inv = inv(P_p)

    G = x_curr.Σ * Ah' * P_p_inv

    smoothed_mean = x_curr.μ + G * (x_next_smoothed.μ - x_pred.μ)

    _R = [
        x_curr.Σ.squareroot' * (I - G * Ah)'
        Qh.squareroot' * G'
        x_next_smoothed.Σ.squareroot' * G'
    ]
    _, P_s_R = qr(_R)
    smoothed_cov = SRMatrix(P_s_R')

    x_curr_smoothed = Gaussian(smoothed_mean, smoothed_cov)
    return x_curr_smoothed, G
end

function smooth!(x_curr, x_next, Ah, Qh, integ, diffusion=1)
    # x_curr is the state at time t_n (filter estimate) that we want to smooth
    # x_next is the state at time t_{n+1}, already smoothed, which we use for smoothing
    @unpack d, q = integ.cache
    @unpack x_pred = integ.cache
    @unpack C1, G1, G2, C2 = integ.cache

    # Prediction: t -> t+1
    predict_mean!(x_pred, x_curr, Ah)
    predict_cov!(x_pred, x_curr, Ah, Qh, C1, diffusion)

    # Smoothing
    # G = x_curr.Σ * Ah' * P_p_inv
    P_p_chol = Cholesky(x_pred.Σ.squareroot, :L, 0)
    G = rdiv!(_matmul!(G1, x_curr.Σ.mat, Ah'), P_p_chol)

    x_curr.μ .+= G * (x_next.μ .- x_pred.μ)

    # Joseph-Form:
    M, L = C2.mat, C2.squareroot
    D = length(x_pred.μ)

    _matmul!(G2, G, Ah)
    copy!(view(L, 1:D, 1:D), x_curr.Σ.squareroot)
    _matmul!(view(L, 1:D, 1:D), G2, x_curr.Σ.squareroot, -1.0, 1.0)

    _matmul!(view(L, 1:D, D+1:2D), _matmul!(G2, G, sqrt.(diffusion)), Qh.squareroot)
    _matmul!(view(L, 1:D, 2D+1:3D), G, x_next.Σ.squareroot)

    # _matmul!(M, L, L')
    # chol = cholesky!(Symmetric(M), check=false)

    QL =
        false && issuccess(chol) ? Matrix(chol.U)' :
        eltype(L) <: Union{Float16,Float32,Float64} ? lq!(L).L : qr(L').R'
    copy!(x_curr.Σ.squareroot, QL)
    _matmul!(x_curr.Σ.mat, QL, QL')

    return nothing
end


# https://optimization.mccormick.northwestern.edu/index.php/Disjunctive_inequalities
using JuMP
using DisjunctiveProgramming

m = Model()
@variable(m, -5 ≤ x[1:2] ≤ 10)
@disjunction(
    m,
    begin
        con1[i=1:2], 0 ≤ x[i] ≤ [3,4][i]
    end,
    begin
        con2[i=1:2], [5,4][i] ≤ x[i] ≤ [9,6][i]
    end,
    reformulation = :big_m,
    name = :y
)
print(m)

# ┌ Warning: [con1[1] : x[1] in [0.0, 3.0], con1[2] : x[2] in [0.0, 4.0]] uses the `MOI.Interval` set. Each instance of the interval set has been split into two constraints, one for each bound.
# ┌ Warning: [con2[1] : x[1] in [5.0, 9.0], con2[2] : x[2] in [4.0, 6.0]] uses the `MOI.Interval` set. Each instance of the interval set has been split into two constraints, one for each bound.
# Feasibility
# Subject to
#  XOR(disj_y) : y[1] + y[2] == 1.0         <- XOR constraint
#  con1[1,lb] : -x[1] + 5 y[1] <= 5.0      <- left-side of con1[1]
#  con1[1,ub] : x[1] + 7 y[1] <= 10.0      <- right-side of con1[1]
#  con1[2,lb] : -x[2] + 5 y[1] <= 5.0      <- left-side of con1[2]
#  con1[2,ub] : x[2] + 6 y[1] <= 10.0      <- right-side of con1[2]
#  con2[1,lb] : -x[1] + 10 y[2] <= 5.0     <- left-side of con2[1]
#  con2[1,ub] : x[1] + y[2] <= 10.0        <- right-side of con2[1]
#  con2[2,lb] : -x[2] + 9 y[2] <= 5.0      <- left-side of con2[2]
#  con2[2,ub] : x[2] + 4 y[2] <= 10.0      <- right-side of con2[2]
#  x[1] >= -5.0                             <- varaible bounds
#  x[2] >= -5.0                             <- variable bounds
#  x[1] <= 10.0                             <- variable bounds
#  x[2] <= 10.0                             <- variable bounds
#  y[1] >= 0.0                              <- lower bound on binary
#  y[2] >= 0.0                              <- lower bound on binary
#  y[1] <= 1.0                              <- upper bound on binary
#  y[2] <= 1.0                              <- upper bound on binary
#  y[1] binary                              <- indicator variable (1st disjunct) is binary
#  y[2] binary                              <- indicator variable (2nd disjunct) is binary


# Script to generate Figures 16 and 17
using TestImages, Plots, SparseArrays, LinearAlgebra, Wavelets, MultiscaleGraphSignalTransforms
include("auxilaries.jl")

# Set up the resolution and display size
gr(dpi=200, size=(800,600))

# Load the cameraman image and subsample it to make it as 128x128 image
img = testimage("camera");
matrix = convert(Array{Float64,2}, img)[1:4:512,1:4:512]
m, n = size(matrix)

# Fig. 16a
display(heatmap(matrix, ratio=1, yaxis =:flip, showaxis = false, ticks = false,
    color = :grays, clim = (0,1), colorbar = false))
savefig("cameraman.pdf")
savefig("cameraman.png")

# Now, let's compute the weight matrix based on the Gaussian affinity
# of the small windows (or radius r) around pixels and the pixel locations.
# Set up the key parameters
r = 5; # specify radius of neighbors
σ = 0.007
# Do the weight matrix computation
W007 = image_Gaussian_affinity(matrix, r, σ)

# Preprocess to generate G (GraphSig struct) and GP (GraphPart struct)
# Note that GraphSig requires a matrix data even if it is just a one vector, i.e.,
# f = matrix[:] does not work!
G007 = GraphSig(W007, f = reshape(matrix, (length(matrix), 1)))
GP007 = partition_tree_fiedler(G007)
dmatrix007 = ghwt_analysis!(G007, GP=GP007)

# Construct or search the specific basis
# Haar
BS_haar007 = bs_haar(GP007)
dvec_haar007 = dmatrix2dvec(dmatrix007, GP007, BS_haar007)

# eGHWT
dvec_eghwt007, BS_eghwt007 = eghwt_bestbasis(dmatrix007, GP007)

# Generate Fig. 17a
top_vectors_plot2(dvec_eghwt007, m, n, BS_eghwt007, GP007)
savefig("cameraman_eghwt09_sigma007.pdf")
savefig("cameraman_eghwt09_sigma007.png")

# Now change the σ value
σ = 0.07
W07 = image_Gaussian_affinity(matrix, r, σ)
# Preprocess to generate G (GraphSig struct) and GP (GraphPart struct)
G07 = GraphSig(W07, f = reshape(matrix, (length(matrix), 1)))
GP07 = partition_tree_fiedler(G07)
dmatrix07 = ghwt_analysis!(G07, GP=GP07)

# Construct or search the specific basis
# Haar
BS_haar07 = bs_haar(GP07)
dvec_haar07 = dmatrix2dvec(dmatrix07, GP07, BS_haar07)
# eGHWT
dvec_eghwt07, BS_eghwt07 = eghwt_bestbasis(dmatrix07, GP07)
# Generate Fig. 17b
top_vectors_plot2(dvec_eghwt07, m, n, BS_eghwt07, GP07)
savefig("cameraman_eghwt09_sigma07.pdf")
savefig("cameraman_eghwt09_sigma07.png")
# Finally, do the classical Haar transform in Wavelets.jl
dvec_classichaar = dwt(matrix, wavelet(WT.haar))

# Fig. 16b (Approximation error plot)
DVEC = [ dvec_classichaar[:], dvec_haar07[:], dvec_eghwt07[:],
         dvec_haar007[:], dvec_eghwt007[:] ]
T = [ "Classical Haar", "Graph Haar (σ = 0.07)",
      "eGHWT (σ = 0.07)", "Graph Haar (σ = 0.007)",
      "eGHWT (σ = 0.007)"]
L = [ (:dashdot,:orange), (:dashdot, :red), (:dashdot, :black),
      (:solid, :red), (:solid, :black) ]
approx_error2(DVEC, T, L, 0.5)
savefig("cameraman_approx_error.pdf")
savefig("cameraman_approx_error.png")


module CompactBasisFunctions

    using ContinuumArrays
    using StaticArrays

    import Base: *
    import ContinuumArrays: Mul, QMul2, QuasiAdjoint, (..), @simplify

    export Derivative, ℵ₁

    include("vandermonde_matrix.jl")

    export Basis, basis, nbasis, grid, degree, order

    include("basis.jl")

    export Bernstein, BernsteinDerivative,
           Chebyshev, ChebyshevDerivative,
           ChebyshevT, ChebyshevTDerivative,
           ChebyshevU, ChebyshevUDerivative,
           Lagrange, LagrangeDerivative,
           LagrangeGauß, LagrangeLobatto,
           Legendre, LegendreDerivative

    include("bernstein.jl")
    include("chebyshev.jl")
    include("lagrange.jl")
    include("legendre.jl")

end


struct SpatialGrid{T <: AbstractFloat }
    x          :: OffsetVector{T}
    global_DoF :: Vector{Vector{Int64}}
    nfree      :: Int64
    nfixed     :: Int64
    max_r      :: Int64
    pt         :: Vector{T} # Quadrature points on (-1,1)
    wt         :: Vector{T} # Quadrature weights
end

"""
   Φ = shape_funcs(r, ξ)

Returns shape functions `Φ[i,j] = Φ_(i-1)(ξ[j])`.
"""
function shape_funcs(r::Int64, ξ::AbstractVector{T}
                    ) where { T <: AbstractFloat }
    J = length(ξ)
    Φ = zeros(T, r, J)
    for j = 1:J
        Φ[1,j] = ( 1 - ξ[j] ) / 2
        Φ[2,j] = ( 1 + ξ[j] ) / 2
    end
    P = Array{T}(undef, r, J)
    legendre_polys!(P, ξ) # P[i,j] = P_(i-1)(ξ[j])
    for j = 1:J, i = 3:r
        a = convert(T, 2(2i-3))
        Φ[i,j] = ( P[i,j] - P[i-2,j] ) / sqrt(a)
    end
    return Φ
end

function SpatialGrid(x::OffsetVector{T}, r::Vector{Int64},
                     bc::Tuple{String,String},
                     num_x_Gauss_pts::Int64) where { T <: AbstractFloat }
    P = length(x) - 1
    @argcheck length(r) == P DimensionMismatch
    global_DoF = Vector{Vector{Integer}}(undef, P)
    total =  0
    max_r = 0
    for p = 1:P
        @argcheck r[p] ≥ 2
        global_DoF[p] = Array{Integer}(undef, r[p])
        total += r[p]
        if r[p] ≥ max_r
            max_r = r[p]
        end
    end
    total -= P-1
    if bc == ("essential", "essential")
        nfixed = 2
        nfree = total - nfixed
        global_DoF[1][1] = nfree + 1
        m = 0
        if P > 1
            for j = 2:r[1]
                m += 1
                global_DoF[1][j] = m
            end
            for p = 2:P-1
                global_DoF[p][1] = global_DoF[p-1][2]
                for j = 2:r[p]
                    m += 1
                    global_DoF[p][j] = m
                end
            end
            global_DoF[P][1] = global_DoF[P-1][2]
            global_DoF[P][2] = nfree + 2
            for j = 3:r[P]
                m += 1
                global_DoF[P][j] = m
            end
        else
            global_DoF[1][2] = nfree + 2
            for j = 3:r[1]
                global_DoF[1][j] = j-2
            end
        end
    elseif bc == ("essential", "natural")
        nfixed = 1
        nfree = total - nfixed
        global_DoF[1][1] = nfree + 1
        m = 0
        for j = 2:r[1]
            m += 1
            global_DoF[1][j] = m
        end
        for p = 2:P
            global_DoF[p][1] = global_DoF[p-1][2]
            for j = 2:r[p]
                m += 1
                global_DoF[p][j] = m
            end
        end
    elseif bc == ("natural", "essential")
        nfixed = 1
        nfree = total - nfixed
        m = 0
        for j = 1:r[1]
            m += 1
            global_DoF[1][j] = m
        end
        for p = 2:P-1
            global_DoF[p][1] = global_DoF[p-1][2]
            for j = 2:r[p]
                m += 1
                global_DoF[p][j] = m
            end
        end
        global_DoF[P][1] = global_DoF[P-1][2]
        global_DoF[P][2] = nfree + 1
        for j = 3:r[P]
            m += 1
            global_DoF[P][j] = m
        end
    elseif bc == ("natural", "natural")
        nfixed = 0
        nfree = total - nfixed
        m = 0
        for j = 1:r[1]
            m += 1
            global_DoF[1][j] = m
        end
        for p = 2:P
            global_DoF[p][1] = global_DoF[p-1][2]
            for j = 2:r[p]
                m += 1
                global_DoF[p][j] = m
            end
        end
    else
        throw(ArgumentError("unrecognised boundary conditions"))
    end
    pt, wt = GaussQuadrature.legendre(T, num_x_Gauss_pts)
    return SpatialGrid{T}(x, global_DoF, nfree, nfixed, max_r, pt, wt)
end

function SpatialGrid(x::OffsetVector{T}, r::Int64,
                     bc::Tuple{String,String},
                     num_Gauss_pts) where { T <: AbstractFloat }
    P = length(x) - 1
    vecr = fill(r, P)
    return SpatialGrid(x, vecr, bc, num_Gauss_pts)
end

function ref_element_matrices!(M::Matrix{T}, S::Matrix{T},
                               κ::T) where { T <: AbstractFloat }
    sz = size(M, 1)
    @argcheck size(M, 2) == sz
    @argcheck size(S) == (sz, sz)
    @argcheck sz >= 2
    fill!(M, zero(T))
    fill!(S, zero(T))
    two = convert(T, 2)
    three = convert(T, 3)
    M[1, 1] = M[2, 2] =  2 / three
    M[2, 1] = M[1, 2] =  1 / three
    S[1, 1] = S[2, 2] =  κ / 2
    S[2, 1] = S[1, 2] = -κ / 2
    if sz == 2
        return
    end
    six = convert(T, 6)
    if sz >= 3
        M[1, 3] = M[2, 3] = -1 /sqrt(six)
        M[3, 1] = M[1, 3]
        M[3, 2] = M[2, 3]
    end
    if sz >= 4
        ten = convert(T, 10)
        M[1, 4] = 1 / ( 3 * sqrt(ten) )
        M[2, 4] = -M[1, 4]
        M[4, 1] = M[1, 4]
        M[4, 2] = M[2, 4]
    end
    for j = 3:sz
        M[j, j] = two / ((2j-5)*(2j-1))
        S[j, j] = κ
    end
    for j = 3:sz-2
        a = convert(T, (2j-3)*(2j+1))
        M[j, j+2] = M[j+2, j] = -1 / ( (2j-1) * sqrt(a) )
    end
end

function assembled_matrices(grid::SpatialGrid,
                            κ::T) where { T <: AbstractFloat }
    x, gdof = grid.x, grid.global_DoF
    nfree, nfixed, max_r = grid.nfree, grid.nfixed, grid.max_r
    P = length(gdof)
    Mref = zeros(T, max_r, max_r)
    Sref = zeros(T, max_r, max_r)
    ref_element_matrices!(Mref, Sref, κ)
    IM, JM, VM = Int64[], Int64[], Float64[]
    IS, JS, VS = Int64[], Int64[], Float64[]
    for p = 1:P
        hp = x[p] - x[p-1]
        rp = length(gdof[p])
        for i = 1:rp
            gi = gdof[p][i]
            if gi ≤ nfree
                for j = 1:rp
                    gj = gdof[p][j]
                    push!(IM, gi)
                    push!(JM, gj)
                    push!(VM, (hp/2)*Mref[i,j])
                    push!(IS, gi)
                    push!(JS, gj)
                    push!(VS, (2/hp)*Sref[i,j])
                end
            end
        end
    end
    M = SparseArrays.sparse(IM, JM, VM, nfree, nfree+nfixed)
    S = SparseArrays.sparse(IS, JS, VS, nfree, nfree+nfixed)
    return M, S
end

function element_load_vector!(Fp::AbstractArray{T},
                              elt::Tuple{T,T}, f::Function,
                              Φ::Matrix{T}, pt::Vector{T}, wt::Vector{T}
                             ) where { T <: AbstractFloat }
    r = length(Fp)
    J = length(pt)
    @argcheck size(Φ,1) ≤ r  DimensionMismatch
    @argcheck size(Φ,2) == J DimensionMismatch
    Jacobian = ( elt[2] - elt[1] ) / 2
    fill!(Fp, zero(T))
    for j = 1:J
        xj = ( ( 1 - pt[j] ) * elt[1] + ( 1 + pt[j] ) * elt[2] ) / 2
        for i = 1:r
            Fp[i] += wt[j] * f(xj) * Φ[i,j] * Jacobian
        end
    end
end

function assembled_load_vector!(F::Vector{T}, grid::SpatialGrid{T},
                                f::Function) where { T <: AbstractFloat }
    x, gdof, pt, wt = grid.x, grid.global_DoF, grid.pt, grid.wt
    nfree, nfixed, max_r = grid.nfree, grid.nfixed, grid.max_r
    @argcheck length(F) == nfree
    P = length(x) - 1
    Φ = shape_funcs(max_r, pt)
    Fp = zeros(T, max_r)
    fill!(F, zero(T))
    for p = 1:P
        rp = length(gdof[p])
        element_load_vector!(view(Fp, 1:rp), (x[p-1],x[p]), f, Φ, pt, wt)
        for i = 1:length(gdof[p])
            gi = gdof[p][i]
            if gi ≤ nfree
                F[gi] += Fp[i]
            end
        end
    end
end

function assembled_load_vector(grid::SpatialGrid{T},
                               f::Function) where { T <: AbstractFloat }
    nfree = grid.nfree
    F = zeros(T, nfree)
    assembled_load_vector!(F, grid, f)
    return F
end

function spatial_points!(xvals::Vector{T}, grid::SpatialGrid{T},
                         ξ::AbstractVector{T}) where { T <: AbstractFloat }
    x, gdof = grid.x, grid.global_DoF
    Nx = length(xvals)
    P = length(x) - 1
    pts_per_interval = length(ξ)
    n = 0
    if Nx == 1 + P * ( pts_per_interval - 1 )
        for p = 1:P
            for j = 1:pts_per_interval - 1
                xpj = ( ( 1 - ξ[j] ) * x[p-1] + (1 + ξ[j] ) * x[p] ) / 2
                n += 1
                xvals[n] = xpj
            end
        end
        n += 1
        xvals[n] = x[P]
    elseif Nx == P * pts_per_interval
        for p = 1:P
            for j = 1:pts_per_interval
                xpj = ( ( 1 - ξ[j] ) * x[p-1] + (1 + ξ[j] ) * x[p] ) / 2
                n += 1
                xvals[n] = xpj
            end
        end
    else
        throw(ArgumentError("xvals length does not match grid"))
    end
end

function evaluate_fem1d_soln!(Uvals::Vector{T}, U::Vector{T},
                              grid::SpatialGrid{T},
                              ξ::AbstractVector{T}
                             ) where { T <: AbstractFloat }
    x, gdof = grid.x, grid.global_DoF
    nfree, nfixed, max_r = grid.nfree, grid.nfixed, grid.max_r
    @argcheck length(U) == nfree DimensionMismatch
    Nx = length(Uvals)
    P = length(x) - 1
    pts_per_interval = length(ξ)
    Φ = shape_funcs(max_r, ξ)
    n = 0
    if Nx == 1 + P * ( pts_per_interval - 1 ) # ξ[1] = -1, ξ[end] = +1
        for p = 1:P
            rp = length(gdof[p])
            for j = 1:pts_per_interval-1
                n += 1
                Uvals[n] = zero(T)
                for i = 1:rp
                    gpi = gdof[p][i]
                    if gpi ≤ nfree
                        Uvals[n] += U[gpi] * Φ[i,j]
                    end
                end
            end
        end
        n += 1
        gpi = gdof[P][2]
        if gpi ≤ nfree
            Uvals[n] = U[gpi]
        end
    elseif Nx == P * pts_per_interval  # ξ[1] > -1 or ξ[end] < +1
        for p = 1:P
            rp = length(gdof[p])
            for j = 1:pts_per_interval
                n += 1
                Uvals[n] = zero(T)
                for i = 1:rp
                    gpi = gdof[p][i]
                    if gpi ≤ nfree
                        Uvals[n] += U[gpi] * Φ[i,j]
                    end
                end
            end
        end
    else
        throw(ArgumentError("Uvals length does not match grid"))
    end
end

function evaluate_fem1d_soln(U::Vector{T},
                             grid::SpatialGrid{T},
                             pts_per_interval::Int64
                            ) where { T <: AbstractFloat }
    x, nfree = grid.x, grid.nfree
    @argcheck length(U) == nfree DimensionMismatch
    P = ubound(x, 1)
    Nx = 1 + P * ( pts_per_interval - 1 )
    ξ = collect( range(-one(T), stop=one(T), length=pts_per_interval) )
    xvals = zeros(T, Nx)
    spatial_points!(xvals, grid, ξ)
    Uvals = zeros(T, Nx)
    evaluate_fem1d_soln!(Uvals, U, grid, ξ)
    return xvals, Uvals
end

function evaluate_pcwise_poly(U::Array{T}, t::OffsetVector{T},
                              grid::SpatialGrid{T},
                              pts_per_time_interval::Int64,
                              pts_per_space_interval::Int64
                             ) where { T <: AbstractFloat }
    τ = range(-one(T), stop=one(T), length=pts_per_time_interval)
    ξ = range(-one(T), stop=one(T), length=pts_per_space_interval)
    xvals, pcwise_t, pcwise_U = evaluate_pcwise_poly(U, t, grid, τ, ξ)
    return xvals, pcwise_t, pcwise_U
end

function evaluate_pcwise_poly(U::Array{T}, t::OffsetVector{T},
                              grid::SpatialGrid{T},
                              τ::AbstractVector{T},
                              ξ::AbstractVector{T}
                             ) where { T <: AbstractFloat }
    x = grid.x
    P = length(x) - 1
    nfree, rt, N = size(U)
    pts_per_space_interval = length(ξ)
    pts_per_time_interval = length(τ)
    @argcheck axes(t, 1) == 0:N
    if ξ[1] > -1.0 || ξ[end] < 1.0
        Nx = P * pts_per_space_interval
    else
        Nx = 1 + P * ( pts_per_space_interval - 1 )
    end
    xvals = zeros(T, Nx)
    spatial_points!(xvals, grid, ξ)
    Ψ = Array{Float64}(undef, rt, pts_per_time_interval)
    legendre_polys!(Ψ, τ)
    pcwise_t = zeros(T, pts_per_time_interval, N)
    pcwise_U = zeros(T, Nx, pts_per_time_interval, N)
    Uvals = zeros(T, Nx)
    for n = 1:N
        for i = 1:pts_per_time_interval
            pcwise_t[i,n] = ( (1-τ[i])*t[n-1] + (1+τ[i])*t[n] ) / 2
            for k = 1:rt
                evaluate_fem1d_soln!(Uvals, U[:,k,n], grid, ξ)
                pcwise_U[:,i,n] .+= Uvals * Ψ[k,i]
            end
        end
    end
    return xvals, pcwise_t, pcwise_U
end

function evaluate_pcwise_error(pcwise_U::Array{T}, u::Function,
                               pcwise_t::Array{T}, xvals::Vector{T}
                              ) where { T <: AbstractFloat }
    Nx = length(xvals)
    N = size(pcwise_t, 2)
    pts_per_time_interval = size(pcwise_t, 1)
    pcwise_err = zeros(Nx, pts_per_time_interval, N)
    for n = 1:N
        for i = 1:pts_per_time_interval
            for p = 1:Nx
                pcwise_err[p,i,n] = ( pcwise_U[p,i,n]
                                     - u(xvals[p], pcwise_t[i,n]) )
            end
        end
    end
    return pcwise_err
end

function dynamic_load_vector!(Fn::Matrix{T}, In::Vector{Float64},
                              grid::SpatialGrid{T}, f::Function,
                              τ::Vector{T}, w::Vector{T},
                              Ψ::AbstractMatrix{T}) where { T <: AbstractFloat }
    nfree, rt = size(Fn)
    @argcheck grid.nfree == nfree
    @argcheck length(In) == 2
    fill!(Fn, 0.0)
    Fni = zeros(nfree)
    Δt = In[2] - In[1]
    for i = 1:length(τ)
        ti = ( ( 1 - τ[i] ) * In[1] + ( 1 + τ[i] ) * In[2] ) / 2
        assembled_load_vector!(Fni, grid, x -> f(x,ti))
        for j = 1:rt
            Fn[:,j] += (Δt/2) * w[i] * Fni * Ψ[j,i]
        end
    end
end

function PDEDG(κ::Float64, f::Function, U0::Vector{Float64},
               grid::SpatialGrid{Float64}, t::OffsetVector{Float64},
               rt::Int64, Mt::Integer) where T <: AbstractFloat
    x, gdof = grid.x, grid.global_DoF
    nfree, nfixed = grid.nfree, grid.nfixed

    G = coef_G(Float64, rt)
    K = coef_K(Float64, rt, rt)
    H = zeros(rt, rt)
    for j = 1:rt
        H[j,j] = 1 / (2j-1)
    end

    N = length(t) - 1
    U = zeros(nfree*rt, N)
    Fn = zeros(nfree, rt)
    M, S = assembled_matrices(grid, κ)
    M = M[1:nfree,1:nfree]
    S = S[1:nfree,1:nfree]
    Ψ_at_minus1 = Vector{Float64}(undef, rt)
    legendre_polys!(Ψ_at_minus1, -1.0)
    τ, w = GaussQuadrature.legendre(Mt)
    Ψ = Array{Float64}(undef, rt, Mt)
    legendre_polys!(Ψ, τ)
    dynamic_load_vector!(Fn, t[0:1], grid, f, τ, w, Ψ)
    b = Fn[1:nfree*rt] + kron(Ψ_at_minus1, M*U0)
    Δt1 = t[1] - t[0]
    G_kron_M = kron(G, M)
    H_kron_S = kron(H, S)
    A = G_kron_M + Δt1 * H_kron_S
    U[:,1] = A \ b
    K_kron_M = kron(K, M)
    for n = 2:N
        Δtn = t[n] - t[n-1]
        dynamic_load_vector!(Fn, t[n-1:n], grid, f, τ, w, Ψ)
        b = Fn[1:nfree*rt] + K_kron_M * U[:,n-1]
        A = G_kron_M + Δtn * H_kron_S
        U[:,n] = A \ b
    end
    return reshape(U, nfree, rt, N)
end

"""
    t, U = PDEDG(κ, f, U0, grid, N, rt, num_Gauss_pts)

Faster version for uniform time steps.
"""
function PDEDG(κ::Float64, f::Function, U0::Vector{Float64},
               grid::SpatialGrid{Float64}, max_t::Float64, N::Int64,
               rt::Int64, Mt::Integer)
    x, gdof = grid.x, grid.global_DoF
    nfree, nfixed = grid.nfree, grid.nfixed

    G = coef_G(Float64, rt)
    K = coef_K(Float64, rt, rt)
    H = zeros(rt, rt)
    for j = 1:rt
        H[j,j] = 1 / (2j-1)
    end

    t_ = collect(range(0, max_t, length=N+1))
    t = OffsetVector(t_, 0:N)
    U = zeros(nfree*rt, N)
    Fn = zeros(nfree, rt)
    M, S = assembled_matrices(grid, κ)
    M = M[1:nfree,1:nfree]
    S = S[1:nfree,1:nfree]
    Ψ_at_minus1 = Vector{Float64}(undef, rt)
    legendre_polys!(Ψ_at_minus1, -1.0)
    τ, w = GaussQuadrature.legendre(Mt)
    Ψ = Array{Float64}(undef, rt, Mt)
    legendre_polys!(Ψ, τ)
    dynamic_load_vector!(Fn, t[0:1], grid, f, τ, w, Ψ)
    b = Fn[1:nfree*rt] + kron(Ψ_at_minus1, M*U0)
    Δt = max_t / N
    G_kron_M = kron(G, M)
    H_kron_S = kron(H, S)
    A = G_kron_M + Δt * H_kron_S
    LU = SparseArrays.lu(A)
    U[:,1] = LU \ b
    K_kron_M = kron(K, M)
    for n = 2:N
        dynamic_load_vector!(Fn, t[n-1:n], grid, f, τ, w, Ψ)
        b = Fn[1:nfree*rt] + K_kron_M * U[:,n-1]
        U[:,n] = LU \ b
    end
    return t, reshape(U, nfree, rt, N)
end

function spatial_L2_projection(f::Function, grid::SpatialGrid{Float64})
    M, S = assembled_matrices(grid, 0.0)
    nfree = grid.nfree
    Pf = M[1:nfree,1:nfree] \ assembled_load_vector(grid, f)
    return Pf
end

"""
    err_τ = maxerr(τ, U, t, u, ξ, grid, cutoff=0.0)

Returns err_τ[i] = max |U(x,t)-u(x,t)| over x values in each spatial
interval corresponding to ξ the values in [-1,1], and over the t value
in each time subinterval corresponding to τ[i] in [-1,1], except for time
intervals with t[n] ≤ cutoff.
"""
function maxerr(τ:: AbstractVector{Float64}, U::Array{Float64},
                t::OffsetVector{Float64}, u::Function,
                ξ::AbstractVector{Float64}, grid::SpatialGrid,
                cutoff=0.0)
    xvals, pcwise_t, pcwise_U = evaluate_pcwise_poly(U, t, grid, τ, ξ)
    pcwise_err = evaluate_pcwise_error(pcwise_U, u, pcwise_t, xvals)
    N = size(pcwise_t, 2)
    pts_per_time_interval = size(pcwise_t, 1)
    Nx = size(pcwise_err, 1)
    err_τ = zeros(pts_per_time_interval)
    for n = 1:N
        if t[n] ≤ cutoff
            continue
        end
        for i = 1:pts_per_time_interval
            for p = 1:Nx
                err_τ[i] = max(abs(pcwise_err[p,i,n]), err_τ[i])
            end
        end
    end
    return err_τ
end

"""
    U = FPDEDG(κ, f, U0, grid, t, rt, num_t_Gauss_pts)

DG solver for fractional PDE.
"""
function FPDEDG(κ::Float64, f::Function, U0::Vector{Float64},
               grid::SpatialGrid{Float64}, t::OffsetVector{Float64},
               rt::Int64, Mt::Int64, store::Store{Float64})
    α = store.α
    x, gdof = grid.x, grid.global_DoF
    nfree, nfixed = grid.nfree, grid.nfixed

    G = coef_G(Float64, rt)
    K = coef_K(Float64, rt, rt)
    Nt = length(t) - 1
    Hn = OffsetArray{Matrix{Float64}}(undef, 0:Nt-1)
    for n = 0:Nt-1
        Hn[n] = Array{Float64}(undef, rt, rt)
    end
    τ, wτ = rule(store.legendre[Mt])
    Ψ = view(store.Ψ, 1:rt, 1:Mt)
    legendre_polys!(Ψ, τ)

    U = zeros(nfree*rt, Nt)
    Fn = zeros(nfree, rt)
    M, S = assembled_matrices(grid, κ)
    M = M[1:nfree,1:nfree]
    S = S[1:nfree,1:nfree]
    Ψ_at_minus1 = Array{Float64}(undef, rt)
    legendre_polys!(Ψ_at_minus1, -1.0)
    dynamic_load_vector!(Fn, t[0:1], grid, f, τ, wτ, Ψ)
    b = Fn[1:nfree*rt] + kron(Ψ_at_minus1, M*U0)
    coef_Hn!(Hn, 1, 0, t, Mt, store)
    G_kron_M = kron(G, M)
    Hnn_kron_S = kron(Hn[0], S)
    A = G_kron_M + Hnn_kron_S
    U[:,1] = A \ b
    K_kron_M = kron(K, M)
    for n = 2:Nt
        coef_Hn!(Hn, n, n-1, t, Mt, store)
        Hnn_kron_S = kron(Hn[0], S)
        A = G_kron_M + Hnn_kron_S
        fill!(b, 0.0)
        for l = 1:n-1
            Hnl_kron_S = kron(Hn[n-l], S)
            b .= b .+ Hnl_kron_S * U[:,l]
        end
        dynamic_load_vector!(Fn, t[n-1:n], grid, f, τ, wτ, Ψ)
        b .= Fn[1:nfree*rt] .+ K_kron_M * U[:,n-1] .- b
        U[:,n] = A \ b
    end
    return reshape(U, nfree, rt, Nt)
end

"""
    t, U = FPDEDG(κ, f, U0, grid, max_t, N, rt, num_t_Gauss_pts)

Specialised version for uniform time steps.
"""
function FPDEDG(κ::Float64, f::Function, U0::Vector{Float64},
                grid::SpatialGrid{Float64}, max_t::Float64,
                N::Int64, rt::Int64, Mt::Int64, store::Store{Float64})
    α = store.α
    x, gdof = grid.x, grid.global_DoF
    nfree, nfixed = grid.nfree, grid.nfixed

    Δt = max_t / N
    G = coef_G(Float64, rt)
    K = coef_K(Float64, rt, rt)
    H = coef_H_uniform!(rt, N, Mt, store)

    t_ = collect(range(0, max_t, length=N+1))
    t = OffsetVector(t_, 0:N)

    U = zeros(nfree*rt, N)
    Fn = zeros(nfree, rt)
    M, S = assembled_matrices(grid, κ)
    M = M[1:nfree,1:nfree]
    S = S[1:nfree,1:nfree]
    Ψ_at_minus1 = Vector{Float64}(undef, rt)
    legendre_polys!(Ψ_at_minus1, -1.0)
    τ, w = GaussQuadrature.legendre(Mt)
    Ψ = Array{Float64}(undef, rt, Mt)
    legendre_polys!(Ψ, τ)
    dynamic_load_vector!(Fn, t[0:1], grid, f, τ, w, Ψ)
    b = Fn[1:nfree*rt] + kron(Ψ_at_minus1, M*U0)
    G_kron_M = kron(G, M)
    Hnn_kron_S = kron(H[0], S)
    A = G_kron_M + Δt^α * Hnn_kron_S
    Fact = SparseArrays.lu(A)
    U[:,1] = Fact \ b
    K_kron_M = kron(K, M)
    for n = 2:N
        fill!(b, 0.0)
        for l = 1:n-1
            Hnl_kron_S = kron(H[n-l], S)
            b .= b .+ Δt^α * Hnl_kron_S * U[:,l]
        end
        dynamic_load_vector!(Fn, t[n-1:n], grid, f, τ, w, Ψ)
        b .= Fn[1:nfree*rt] .+ K_kron_M * U[:,n-1] .- b
        U[:,n] = Fact \ b
    end
    return t, reshape(U, nfree, rt, N)
end

function functional(pcwise_U::Array{Float64}, pcwise_t::Array{Float64},
                    func::Function, grid::SpatialGrid,
                    ξ::Vector{Float64}, w::Vector{Float64})
    x = grid.x
    pts_per_time_interval = size(pcwise_t,1)
    N = size(pcwise_t, 2)
    P = ubound(x, 1)
    pts_per_space_interval = length(ξ)
    Fnctl = zeros(pts_per_time_interval, N)

end


######## Stan diagnose example  ###########

using Stan

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

bernoulli = "
data {
  int<lower=0> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
    y ~ bernoulli(theta);
}
"

bernoullidata = [
  Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1]),
  Dict("N" => 10, "y" => [0, 1, 0, 0, 0, 0, 1, 0, 0, 1]),
  Dict("N" => 10, "y" => [0, 0, 0, 0, 0, 0, 1, 0, 1, 1]),
  Dict("N" => 10, "y" => [0, 0, 0, 1, 0, 0, 0, 1, 0, 1])
]

stanmodel = Stanmodel(Diagnose(Gradient(epsilon=1e-6)), name="bernoulli", model=bernoulli);

diags = stan(stanmodel, bernoullidata, ProjDir, CmdStanDir=CMDSTAN_HOME);
diags[1]["diagnose"] |> display

end # cd


# [ include("../src/"*s) for s in readdir("../src") ]
using PolyChaos
# using LinearAlgebra
# import FFTW
# import SpecialFunctions
# numerically compute recurrence coefficients for (almost) Gaussian density w(t)
N = 10
w(t) = exp(-t^2)
lb, ub = -Inf, Inf
@time α, β = rm_compute(w,lb,ub;Nquad=200,Npoly=N)
# analytical solution
α_ana = zeros(N)
β_ana = [ √π; [0.5*k for k=1:N-1] ]
# compare
display("Deviation for α: $(α-α_ana)")
display("Deviation for β: $(β-β_ana)")
## do the same for Chebyshev polynomials #4
v(t) = sqrt(1-t)/sqrt(1+t)
lb, ub = -1+1e-8, 1-1e-8
@time α, β = rm_compute(v,lb,ub;Nquad=2000,Npoly=N)
# analytical solution
α_ana = [-0.5; zeros(N-1) ]
β_ana = [ π; [0.25 for k=1:N-1] ]
# compare
display("Deviation for α: $(α-α_ana)")
display("Deviation for β: $(β-β_ana)")


# Import libraries.
using Turing, Flux, Plots, Random, ReverseDiff

# Hide sampling progress.
Turing.setprogress!(false);

# Use reverse_diff due to the number of parameters in neural networks.
Turing.setadbackend(:reversediff)

# Number of points to generate.
N = 80
M = round(Int, N / 4)
Random.seed!(1234)

# Generate artificial data.
x1s = rand(M) * 4.5; x2s = rand(M) * 4.5;
xt1s = Array([[x1s[i] + 0.5; x2s[i] + 0.5] for i = 1:M])
x1s = rand(M) * 4.5; x2s = rand(M) * 4.5;
append!(xt1s, Array([[x1s[i] - 5; x2s[i] - 5] for i = 1:M]))

x1s = rand(M) * 4.5; x2s = rand(M) * 4.5;
xt0s = Array([[x1s[i] + 0.5; x2s[i] - 5] for i = 1:M])
x1s = rand(M) * 4.5; x2s = rand(M) * 4.5;
append!(xt0s, Array([[x1s[i] - 5; x2s[i] + 0.5] for i = 1:M]))

# Store all the data for later.
xs = [xt1s; xt0s]
ts = [ones(2*M); zeros(2*M)]

# Plot data points.
function plot_data()
    x1 = map(e -> e[1], xt1s)
    y1 = map(e -> e[2], xt1s)
    x2 = map(e -> e[1], xt0s)
    y2 = map(e -> e[2], xt0s)

    Plots.scatter(x1,y1, color="red", clim = (0,1))
    Plots.scatter!(x2, y2, color="blue", clim = (0,1))
end

plot_data()

# Turn a vector into a set of weights and biases.
function unpack(nn_params::AbstractVector)
    W₁ = reshape(nn_params[1:6], 3, 2);
    b₁ = reshape(nn_params[7:9], 3)

    W₂ = reshape(nn_params[10:15], 2, 3);
    b₂ = reshape(nn_params[16:17], 2)

    Wₒ = reshape(nn_params[18:19], 1, 2);
    bₒ = reshape(nn_params[20:20], 1)
    return W₁, b₁, W₂, b₂, Wₒ, bₒ
end

# Construct a neural network using Flux and return a predicted value.
function nn_forward(xs, nn_params::AbstractVector)
    W₁, b₁, W₂, b₂, Wₒ, bₒ = unpack(nn_params)
    nn = Chain(Dense(W₁, b₁, tanh),
               Dense(W₂, b₂, tanh),
               Dense(Wₒ, bₒ, σ))
    return nn(xs)
end;

# Create a regularization term and a Gaussain prior variance term.
alpha = 0.09
sig = sqrt(1.0 / alpha)

# Specify the probabalistic model.
@model function bayes_nn(xs, ts)
    # Create the weight and bias vector.
    nn_params ~ MvNormal(zeros(20), sig .* ones(20))

    # Calculate predictions for the inputs given the weights
    # and biases in theta.
    preds = nn_forward(xs, nn_params)

    # Observe each prediction.
    for i = 1:length(ts)
        ts[i] ~ Bernoulli(preds[i])
    end
end;

# Perform inference.
N = 5000
ch = sample(bayes_nn(hcat(xs...), ts), HMC(0.05, 4), N);

# Extract all weight and bias parameters.
theta = MCMCChains.group(ch, :nn_params).value;

# Plot the data we have.
plot_data()

# Find the index that provided the highest log posterior in the chain.
_, i = findmax(ch[:lp])

# Extract the max row value from i.
i = i.I[1]

# Plot the posterior distribution with a contour plot.
x_range = collect(range(-6,stop=6,length=25))
y_range = collect(range(-6,stop=6,length=25))
Z = [nn_forward([x, y], theta[i, :])[1] for x=x_range, y=y_range]
contour!(x_range, y_range, Z)

# Return the average predicted value across
# multiple weights.
function nn_predict(x, theta, num)
    mean([nn_forward(x, theta[i,:])[1] for i in 1:10:num])
end;

# Plot the average prediction.
plot_data()

n_end = 1500
x_range = collect(range(-6,stop=6,length=25))
y_range = collect(range(-6,stop=6,length=25))
Z = [nn_predict([x, y], theta, n_end)[1] for x=x_range, y=y_range]
contour!(x_range, y_range, Z)

# Number of iterations to plot.
n_end = 500

anim = @gif for i=1:n_end
    plot_data()
    Z = [nn_forward([x, y], theta[i,:])[1] for x=x_range, y=y_range]
    contour!(x_range, y_range, Z, title="Iteration $i", clim = (0,1))
end every 5

using Bijectors
using Turing: Variational
using AdvancedVI

m = bayes_nn(hcat(xs...), ts);

q = Variational.meanfield(m)

μ = randn(length(q))
ω = -1 .* ones(length(q))

q = AdvancedVI.update(q, μ, exp.(ω));

advi = ADVI(10, 5_000)
q_hat = vi(m, advi, q);

samples = transpose(rand(q_hat, 5000))
ch_vi = Chains(reshape(samples, size(samples)..., 1), string.(MCMCChains.namesingroup(ch, :nn_params)));

# Extract all weight and bias parameters.
theta = MCMCChains.group(ch_vi, :nn_params).value;

# Plot the average prediction.
plot_data()

n_end = 1500
x_range = collect(range(-6,stop=6,length=25))
y_range = collect(range(-6,stop=6,length=25))
Z = [nn_predict([x, y], theta, n_end)[1] for x=x_range, y=y_range]
contour!(x_range, y_range, Z)

# Specify the network architecture.
network_shape = [
    (3,2, :tanh),
    (2,3, :tanh),
    (1,2, :σ)]

# Regularization, parameter variance, and total number of
# parameters.
alpha = 0.09
sig = sqrt(1.0 / alpha)
num_params = sum([i * o + i for (i, o, _) in network_shape])

# This modification of the unpack function generates a series of vectors
# given a network shape.
function unpack(θ::AbstractVector, network_shape::AbstractVector)
    index = 1
    weights = []
    biases = []
    for layer in network_shape
        rows, cols, _ = layer
        size = rows * cols
        last_index_w = size + index - 1
        last_index_b = last_index_w + rows
        push!(weights, reshape(θ[index:last_index_w], rows, cols))
        push!(biases, reshape(θ[last_index_w+1:last_index_b], rows))
        index = last_index_b + 1
    end
    return weights, biases
end

# Generate an abstract neural network given a shape,
# and return a prediction.
function nn_forward(x, θ::AbstractVector, network_shape::AbstractVector)
    weights, biases = unpack(θ, network_shape)
    layers = []
    for i in eachindex(network_shape)
        push!(layers, Dense(weights[i],
            biases[i],
            eval(network_shape[i][3])))
    end
    nn = Chain(layers...)
    return nn(x)
end

# General Turing specification for a BNN model.
@model bayes_nn_general(xs, ts, network_shape, num_params) = begin
    θ ~ MvNormal(zeros(num_params), sig .* ones(num_params))
    preds = nn_forward(xs, θ, network_shape)
    for i = 1:length(ts)
        ts[i] ~ Bernoulli(preds[i])
    end
end

# Perform inference.
num_samples = 500
ch2 = sample(bayes_nn_general(hcat(xs...), ts, network_shape, num_params), NUTS(0.65), num_samples);

# This function makes predictions based on network shape.
function nn_predict(x, theta, num, network_shape)
    mean([nn_forward(x, theta[i,:], network_shape)[1] for i in 1:10:num])
end;

# Extract the θ parameters from the sampled chain.
params2 = MCMCChains.group(ch2, :θ).value

plot_data()

x_range = collect(range(-6,stop=6,length=25))
y_range = collect(range(-6,stop=6,length=25))
Z = [nn_predict([x, y], params2, length(ch2), network_shape)[1] for x=x_range, y=y_range]
contour!(x_range, y_range, Z)

if isdefined(Main, :TuringTutorials)
    Main.TuringTutorials.tutorial_footer(WEAVE_ARGS[:folder], WEAVE_ARGS[:file])
end


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

@testset "Clebsch-Gordan" begin

@info("Testset Clebsch-Gordan")

#---

using PyCall, Test, ACE, ACE.SphericalHarmonics, JuLIP.Testing, StaticArrays
using JuLIP: evaluate
using ACE.SphericalHarmonics: index_y
using ACE.RPI.Rotations3D: ClebschGordan

#---

try
	sympy = pyimport("sympy")
	spin = pyimport("sympy.physics.quantum.spin")

	pycg(j1, m1, j2, m2, j3, m3, T=Float64) =
      		spin.CG(j1, m1, j2, m2, j3, m3).doit().evalf().__float__()

	cg = ClebschGordan()

	@info("compare implementation against `sympy`")
	let ntest = 0
   		while ntest < 200
      			j1 = rand(0:10)
      			j2 = rand(0:10)
      			J = rand(abs(j1-j2):min(10,(j1+j2)))
      			m1 = rand(-j1:j1)
      			for m2 = -j2:j2
         			M = m1+m2
         			if abs(M) <= J
            				ntest += 1
            				print_tf(@test cg(j1,m1, j2,m2, J,M) ≈ pycg(j1,m1, j2,m2, J,M))
            				# print_tf(@test clebschgordan(j1,m1, j2,m2, J,M) ≈ )
         			end
      			end
   		end
	end
	println()

	##

	@info("Checking the SphH expansion in terms of CG coeffs")
	# expansion coefficients of a product of two spherical harmonics in terms a
	# single spherical harmonic
	# see e.g. https://en.wikipedia.org/wiki/Clebsch–Gordan_coefficients
	# this is the magic formula that we need, on which everything else is based
	for ntest = 1:200
   		# two random Ylm  ...
   		l1, l2 = rand(1:10), rand(1:10)
   		m1, m2 = rand(-l1:l1), rand(-l2:l2)
   		# ... evaluated at random spherical coordinates
   		θ = rand() * π
   		φ = (rand()-0.5) * 2*π
   		R = SVector( cos(φ)*sin(θ), sin(φ)*sin(θ), cos(θ) )
   		# evaluate all relevant Ylms (up to l1 + l2)
   		Ylm = evaluate(SHBasis(l1+l2), R)
   		# evaluate the product p = Y_l1_m1 * Y_l2_m2
   		p = Ylm[index_y(l1,  m1)] * Ylm[index_y(l2,m2)]
   		# and its expansion in terms of CG coeffs
   		p2 = 0.0
   		M = m1 + m2  # all other coeffs are zero

   		for L = abs(M):(l1+l2)
      			p2 += sqrt( (2*l1+1)*(2*l2+1) / (4 * π * (2*L+1)) ) *
            		cg(l1,  0, l2,  0, L, 0) *
            		cg(l1, m1, l2, m2, L, M) *
            		Ylm[index_y(L, M)]
   		end
   		print_tf((@test (p ≈ p2) || (abs(p-p2) < 1e-15)))
	end
	println()

catch e   # try importing some python packages
	@warn "sympy didn't import, No Clebsch-Gordan tests are running"
end

#---

end # @testset


using EngEconomics

# Given
aInit = 1700
bInit = 2100
cInit = 3750
aAnnualBenefit = 1000
bAnnualBenefit = 1000
cAnnualBenefit = 1000
aLife = 2
bLife = 3
cLife = 6
i = 0.08

# Find: Present Worth Value of each Item then pick the best option
# Start by determining the number of periods for each product
years = 6
aPeriods = years / aLife
bPeriods = years / bLife
cPeriods = years / cLife

# Next, determine the present worth of product A
aInitPW = -aInit * (1 + presentWorthFactor(i, aLife) + presentWorthFactor(i, 2 * aLife))
aAnnualBenefitPW = (aAnnualBenefit * seriesPresentAmountFactor(i, aLife)) * (1 + presentWorthFactor(i, aLife) + presentWorthFactor(i, 2 * aLife))
aPW = aInitPW + aAnnualBenefitPW

# Determine PW of Product B
bInitPW = -bInit * (1 + presentWorthFactor(i, bLife))
bAnnualBenefitPW = (bAnnualBenefit * seriesPresentAmountFactor(i, bLife)) * (1 + presentWorthFactor(i, bLife))
bPW = bInitPW + bAnnualBenefitPW

# Determine PW of Product C
cInitPW = -cInit * (1)
cAnnualBenefitPW = (cAnnualBenefit * seriesPresentAmountFactor(i, cLife)) * (1)
cPW = cInitPW + cAnnualBenefitPW

# Compare Alternatives, pick the greatest one
maxPW = max(aPW, bPW, cPW)

if maxPW == aPW
	println("Pick A")
elseif maxPW == bPW
	println("Pick B")
else
	println("Pick C")
end


"""
    BiweightMidcovariance(; c=9.0, modify_sample_size=false)

The biweight midcovariance is a covariance estimator that is resilient to outliers.

The technique derives originally from astrophysics [1], and is implemented in the Python
module [Astropy](https://www.astropy.org/) [2], as well as in NIST's Dataplot [3].

Consider two random variables ``x`` and ``y``, for which we have ``n`` observations
``\\{(x_i, y_i)\\}``. The biweight midcovariance is then defined to be:
```math
n_s\\cdot\\frac{
    \\sum_{|u_i|<1,|v_i|<1}(x_i - M_x)(1 - u_i^2)^2(y_i - M_y)(1 - v_i^2)^2
}{
    \\left(\\sum_{|u_i|<1}(1 - u_i^2)(1-5u_i^2)\\right)
    \\left(\\sum_{|v_i|<1}(1 - v_i^2)(1-5v_i^2)\\right)
}
```
where ``n_s`` is the sample size, ``M_x`` and ``M_y`` are the medians of ``\\{x_i\\}`` and
``\\{y_i\\}`` respectively, and
```math
\\begin{aligned}
u_i &= \\frac{x_i - M_x}{c \\cdot \\mathrm{MAD}_x} \\\\
v_i &= \\frac{y_i - M_y}{c \\cdot \\mathrm{MAD}_y},
\\end{aligned}
```
where ``\\mathrm{MAD}`` represents the median absolute deviation,
```math
\\begin{aligned}
\\mathrm{MAD}_x &= \\mathrm{median}(\\left\\{\\left|x_i - M_x\\right|\\right\\}) \\\\
\\mathrm{MAD}_y &= \\mathrm{median}(\\left\\{\\left|y_i - M_y\\right|\\right\\}).
\\end{aligned}
```

If either ``\\mathrm{MAD}_x = 0`` or ``\\mathrm{MAD}_y = 0``, the pairwise covariance is
defined to be zero.

The parameter ``c`` is a tuning constant, for which the default is ``9.0``.
Larger values will reduce the number of outliers that are removed — i.e. reducing
robustness, but increasing sample efficiency.

# Fields
- `c::Float64`: The tuning constant corresponding to ``c`` above.
- `modify_sample_size::Bool`: If `false`, then we use a sample size ``n_s`` equal to the
    total number of observations ``n``. This is consistent with the standard definition of
    biweight midcovariance in the literature. Otherwise, we count only those elements which
    are not rejected as outliers in the numerator, i.e. those for which ``|u_i|<1``
    and ``|v_i|<1``.
    This follows the implementation in astropy [2].

# Complexity
- Space: ``O(p^2)``
- Time: ``O(np^2)``

# References
[1] Beers, Flynn, and Gebhardt (1990; AJ 100, 32)
"Measures of Location and Scale for Velocities in Clusters of Galaxies -- A Robust Approach"

[2] [Astropy biweight_midcovariance](https://docs.astropy.org/en/stable/api/astropy.stats.biweight_midcovariance.html)

[3] [NIST Dataplot biweight midcovariance](https://www.itl.nist.gov/div898/software/dataplot/refman2/auxillar/biwmidc.htm)
"""
struct BiweightMidcovariance <: CovarianceEstimator
    c::Float64
    modify_sample_size::Bool
end
function BiweightMidcovariance(; c::Real=9.0, modify_sample_size::Bool=false)
    c > 0 || throw(ArgumentError("c must be positive, got $c"))
    return BiweightMidcovariance(Float64(c), modify_sample_size)
end

function covzm(ce::BiweightMidcovariance, x::AbstractVector{<:Real})
    MADx = median!(abs.(x))
    numerator = zero(eltype(x))

    # If MADx is zero, return zero.
    iszero(MADx) && return zero(one(eltype(x)) / 1)

    denominator = zero(eltype(x))
    count = 0
    for xi in x
        ui² = (xi / (ce.c * MADx))^2
        # Benchmarking suggests that having this branch is similar or faster to a form
        # using `ifelse`, where we always compute the updates but avoid the branch.
        ui² >= 1 && continue
        count += 1
        numerator += xi^2 * (1 - ui²)^4
        denominator += (1 - ui²) * (1 - 5 * ui²)
    end
    n = ifelse(ce.modify_sample_size, count, length(x))
    return n * numerator / (denominator^2)
end

function cov(
    ce::BiweightMidcovariance,
    x::AbstractVector{<:Real};
    mean::Union{Nothing,<:Real}=nothing,
)
    Mx = mean === nothing ? median(x) : mean
    return covzm(ce, x .- Mx)
end

function covzm(
    ce::BiweightMidcovariance, x::AbstractVector{<:Real}, y::AbstractVector{<:Real}
)
    MADx = median!(abs.(x))
    MADy = median!(abs.(y))

    # Promote types between x & y for numerator
    numerator = zero(promote_type(eltype(x), eltype(y)))

    # If either of the MADs are zero, return zero.
    #   `numerator` is always zero here, and of the correct type.
    iszero(MADx) && return numerator
    iszero(MADy) && return numerator

    denominator_x = zero(eltype(x))
    denominator_y = zero(eltype(y))
    count = 0
    for (xi, yi) in zip(x, y)
        ui² = (xi / (ce.c * MADx))^2
        vi² = (yi / (ce.c * MADy))^2

        # Compute the updates, and increment if desired.
        # This way should avoid branches, and benchmarks suggest it's 1-2% faster than
        # using conditionals.
        ax = (1 - ui²) * (1 - 5 * ui²)
        ay = (1 - vi²) * (1 - 5 * vi²)
        anum = xi * (1 - ui²)^2 * yi * (1 - vi²)^2

        denominator_x += ifelse(ui² < 1, ax, zero(denominator_x))
        denominator_y += ifelse(vi² < 1, ay, zero(denominator_y))
        i_count, i_numerator = ifelse(ui² < 1 && vi² < 1, (1, anum), (0, zero(numerator)))
        count += i_count
        numerator += i_numerator
    end

    n = ifelse(ce.modify_sample_size, count, length(x))
    return n * numerator / (denominator_x * denominator_y)
end

function cov(
    ce::BiweightMidcovariance, x::AbstractVector{<:Real}, y::AbstractVector{<:Real}
)
    if size(x) != size(y)
        throw(ArgumentError("x & y have different sizes, $(size(x)) & $(size(y))"))
    end

    return covzm(ce, x .- median(x), y .- median(y))
end

function cov(
    ce::BiweightMidcovariance,
    X::AbstractMatrix{<:Real};
    mean::Union{Nothing,AbstractVector{<:Real},AbstractMatrix{<:Real}}=nothing,
    dims::Int=1,
)
    dims ∈ (1, 2) || throw(ArgumentError("Argument dims can only be 1 or 2 (given: $dims)"))

    # Ascertain the number of observations `n` and number of variables `p`.
    n, p = dims == 1 ? size(X) : (size(X, 2), size(X, 1))

    # Temporary storage to use when computing medians.
    temp = Vector{eltype(X)}(undef, n)

    # If the `mean` argument isn't provided, then we centralise with the median.
    # After this statement, Xc is standardised to be of shape (n, p)
    Xc = if mean === nothing
        # The following is equivalent to
        #   X .- median(X; dims=dims)
        # However, this implementation significantly reduces the runtime of this function.
        # This is primarily achieved by reducing the number of allocations -- we allocate a
        # single temporary buffer which we allow `median!` to mutate, and then copy the next
        # section of data into it.
        X = dims == 2 ? transpose(X) : X
        Xc = copy(X)
        @inbounds for i in 1:p
            # Avoiding copyto!(, ... @view(X[:, i])), because this can allocate a view if
            # the call to copyto! isn't inlined.
            for j in 1:n
                temp[j] = X[j, i]
            end
            Xc[:, i] .-= median!(temp)
        end
        Xc
    else
        Xc = X .- mean
        # Standardise the orientation of Xc.
        dims == 2 ? transpose(Xc) : Xc
    end

    # Compute all the median absolute deviations.
    # To avoid extra allocations, use a temporary buffer `temp` which will be populated with
    # the absolute values of the observations of each variable, and which we then allow
    # median! to mutate.
    T = typeof(one(eltype(X)) / 1)  # This is the output type of median.
    MAD = Matrix{T}(undef, 1, p)
    @inbounds for i in 1:p
        # Avoiding map!(abs, temp, @view(Xc[:, i])), because this can allocate a view if
        # the call to map! isn't inlined.
        for j in 1:n
            temp[j] = abs(Xc[j, i])
        end
        MAD[1, i] = median!(temp)
    end

    u² = @. (Xc / (ce.c * MAD))^2

    onesubu² = max.((1 .- u²), 0)
    a = @. Xc * onesubu²^2
    numerator = a' * a

    a[:, :] .= @. onesubu² * (1 - 5 * u²)  # Re-use `a` memory to avoid an allocation.
    b = sum(a; dims=1)
    denominator = b' * b

    n = if ce.modify_sample_size
        mask = onesubu² .> 0
        mask' * mask
    else
        n
    end

    # Re-use the numerator memory for the result.
    result = numerator
    result .*= n ./ denominator

    # Whereever the MAD is zero, we should set the result to zero.
    mad_zero = dropdims(iszero.(MAD); dims=1)
    result[mad_zero, :] .= zero(eltype(result))
    result[:, mad_zero] .= zero(eltype(result))

    return result
end


remap(x, from, to) = (x - from[1]) * (to[2] - to[1]) / (from[2] - from[1]) + to[1]


module TestShapeDiagrams
using Test

using Catlab.Theories, Catlab.CategoricalAlgebra.ShapeDiagrams

A, B, C = Ob(FreeCategory, :A, :B, :C)

# Spans.
f, g = Hom(:f, C, A), Hom(:g, C, B)
span = Span(f,g)
@test apex(span) == C
@test left(span) == f
@test right(span) == g
f = Hom(:f, A, B)
@test_throws Exception Span(f,g)

# Cones.
f, g = Hom(:f, C, A), Hom(:g, C, B)
cone = Cone(C,[f,g])
@test apex(cone) == C
@test leg(cone,1) == f
@test leg(cone,2) == g
f = Hom(:f, A, B)
@test_throws Exception Cone(C,[f,g])

# Cospans.
f, g = Hom(:f, A, C), Hom(:g, B, C)
cospan = Cospan(f,g)
@test base(cospan) == C
@test left(cospan) == f
@test right(cospan) == g

# Cocones.
f, g = Hom(:f, A, C), Hom(:g, B, C)
cocone = Cocone(C,[f,g])
@test base(cocone) == C
@test leg(cocone,1) == f
@test leg(cocone,2) == g
f = Hom(:f, A, B)
@test_throws Exception Cocone(C,[f,g])

# Diagrams.

f, g, h = Hom(:f, A, C), Hom(:g, B, C), Hom(:h, A, B)
diag = Diagram([A,B,C],[(1,3,f),(2,3,g),(1,2,h)])
@test_throws Exception Diagram([A,B,C],[(1,2,f),(2,3,g),(1,2,h)])

end


# p14.jl - solve nonlinear BVP u_xx = exp(u), u(-1)=u(1)=0
#         (compare p13.jl)

N = 16;
(D,x) = cheb(N); D2 = D^2; D2 = D2[2:N,2:N];
u = zeros(N-1);
change = 1; it = 0;
while change > 1e-15                   # fixed-point iteration
    unew = D2\exp.(u);
    change = norm(unew-u,Inf);
    u = unew; it = it+1;
end
u = [0;u;0];
clf(); axes([.1,.4,.8,.5]);
plot(x,u,".",markersize=6);
xx = -1:.01:1;
uu = polyval(polyfit(x,u),xx);
plot(xx,uu), grid(true);
title("no. steps = $it      u(0) =$(u[Int(N/2+1)])")


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, StatsPlots, LaTeXStrings
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.9"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	dist = TriangularDist(0,2,1)
	xGrid = 0:0.01:2
	uGrid = 0:0.01:1
end

# ╔═╡ a715f0ac-2034-11eb-029c-9be9db193aed
begin
	p1 = plot( xGrid, pdf.(dist,xGrid), c=:blue,
			xlims=(0,2), ylims=(0,1.1),
			xlabel="x", ylabel="f(x)")
	p2 = plot( xGrid, cdf.(dist,xGrid), c=:blue,
			xlims=(0,2), ylims=(0,1),
			xlabel="x", ylabel="F(x)")
	p3 = plot( uGrid,quantile.(dist,uGrid), c=:blue,
			xlims=(0,1), ylims=(0,2),
			xlabel="u", ylabel=(L"F^{-1}(u)"))

	plot(p1, p2, p3, legend=false, layout=(1,3), size=(1200, 400))
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.9"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═a715f0ac-2034-11eb-029c-9be9db193aed
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


"""
Uses UK Ordnance Survey "A Guide to Coordinate Systems in Great Britain"
https://www.ordnancesurvey.co.uk/documents/resources/guide-coordinate-systems-great-britain.pdf

Worked example (Airy1830)
E = 651409.903 # meters
N = 313177.270 # meters

"""
function OSENtoLLA(E,N)
    # Airy 1830
    a = 6377563.396
    b = 6356256.909
    e2 = (a^2 - b^2)/a^2
    E0 = 400000
    N0 = -100000
    F0 = 0.9996012717
    ϕ0 = 49 * (π/180)
    λ0 = -2 * (π/180)

    ϕ1 = (N - N0)/a*F0 + ϕ0
    n=(a-b)/(a+b)
    M_function(ϕ) = b*F0*( (1 + n + 5/4*n^2 + 5/4*n^3) * (ϕ-ϕ0)
        - (3*n + 3*n^2 + 21/8*n^3) * sin(ϕ-ϕ0) * cos(ϕ+ϕ0)
        + (15/8*n^2 + 15/8*n^3) * sin(2*(ϕ-ϕ0)) * cos(2*(ϕ+ϕ0))
        - 35/24*n^3 * sin(3*(ϕ-ϕ0))*cos(3*(ϕ+ϕ0)))
    residual(M) = (N-N0-M)
    M = M_function(ϕ1)

    while residual(M) > 0.00000001
        ϕ1 = (residual(M)/a*F0) + ϕ1
        M = M_function(ϕ1)
    end

    ϕ1 / (π/180)

    ν = a * F0 * abs(1-e2*sin(ϕ1)^2)^(-0.5) # \nu
    ρ = a * F0 * (1-e2) * (1-e2*sin(ϕ1)^2)^(-1.5)
    η2 = ν/ρ-1

    VII = tan(ϕ1)/(2*ρ*ν)
    VIII = tan(ϕ1)/(24*ρ*ν^3) * (5 + 3*tan(ϕ1)^2 + η2 - 9*(tan(ϕ1)^2)*η2)
    IX = tan(ϕ1)/(720*ρ*ν^5) * (61 + 90*tan(ϕ1)^2 + 45*(tan(ϕ1)^4))
    X = sec(ϕ1)/ν
    XI = sec(ϕ1)/(6*ν^3) * (ν/ρ + 2*tan(ϕ1)^2)
    XII = sec(ϕ1)/(120*ν^5) * (5 + 28*tan(ϕ1)^2 + 24*tan(ϕ1)^4)
    XIIA = sec(ϕ1)/(5040*ν^7) * (61 + 662*tan(ϕ1)^2 + 1320*tan(ϕ1)^4 + 720*tan(ϕ1)^6)

    ϕ = ϕ1 - VII*(E-E0)^2 + VIII*(E-E0)^4 - IX*(E-E0)^6
    λ = λ0 + X*(E-E0) - XI*(E-E0)^3 + XII*(E-E0)^5 - XIIA*(E-E0)^7

    return ϕ/(π/180), λ/(π/180)
end


using NonLinearReactionAdvectionDiffusionWithFrontData;

@testset "Дельта-функция                                   " begin
    N = 20;
    X = [i/N for i in 0:N];

    d = [ delta(x, X, 0.5) for x in X ]

    @testset "формула прямоугольников" begin
        I = 0.0;
        for n in 1:N
            δx = X[n+1] - X[n]
            I += d[n] * δx
        end
        @test isapprox(I, 1.0)
    end

    @testset "формула трапеций" begin
        I = 0.0;
        for n in 1:N
            δx = X[n+1] - X[n]
            I += ((d[n] + d[n+1]) / 2.0) * δx
        end
        @test isapprox(I, 1.0)
    end

    # Возвращает только один ненулевой элемент
    @test length( filter( x -> x != 0, d) ) == 1

    @testset "Проверка области определения дельта-функции" begin
        @test_throws DomainError delta(-2, X, 0.5)
        @test_throws DomainError delta(2, X, 0.5)
        @test_throws DomainError delta(0.5, X, -2)
        @test_throws DomainError delta(0.5, X, 2)
        @test_throws DomainError delta(0.5, X, 1.0)
    end

end


############################# abstract QR type #################################
abstract type AbstractQR{T} <: Factorization{T} end
# AbstractQR assumes existence of n, m, Q, R fields
Base.size(F::AbstractQR) = (F.n, F.m)
Base.size(F::AbstractQR, i::Int) = i > 2 ? 1 : size(F)[i]
Base.eltype(F::AbstractQR{T}) where T = T

Base.:\(F::AbstractQR, x::AbstractVecOrMat) = ldiv!(F, copy(x))

# some helpers
"""
```
    number_of_rotations(n::Int, m::Int)
```
Computes the number of Givens rotations that are necessary to compute the QR
factorization of a general matrix of size n by m.
"""
number_of_rotations(n::Int, m::Int) = n*m - (m*(m+1)) ÷ 2
number_of_rotations(A::AbstractMatOrFac) = number_of_rotations(size(A)...)
# min_nm = min(m, m) # IDEA: for overdetermined systems

# to append k columns
number_of_rotations_to_append_column(n::Int, m::Int, k::Int = 1) = k*(n-m) - (k*(k+1)) ÷ 2
# to remove a single column at index k in R
number_of_rotations_to_remove_column(m::Int, k::Int) = m-k

"""
```
    allocate_rotations(T::DataType, n::Int, m::Int)
```
Allocates a vector of Givens rotations of a length that is necessary to compute
the QR factorization of a general matrix of size n by m.
"""
function allocate_rotations(T::DataType, n::Int, m::Int)
    Vector{Givens{T}}(undef, number_of_rotations(n, m))
end
function allocate_rotations(A::AbstractMatOrFac)
    allocate_rotations(eltype(A), size(A)...)
end


using ModelingToolkit, MethodOfLines, LinearAlgebra, OrdinaryDiffEq
using DomainSets

@test_broken begin #@testset "Inviscid Burgers equation, 1D, u(0, x) ~ x" begin
    @parameters x t
    @variables u(..)
    Dx = Differential(x)
    Dt = Differential(t)
    x_min = 0.0
    x_max = 1.0
    t_min = 0.0
    t_max = 1.0

    analytic_u(t, x) = x / (t + 1)

    eq = Dt(u(t, x)) + u(t, x) * Dx(u(t, x)) ~ 0

    bcs = [u(0, x) ~ x,
        u(t, x_min) ~ analytic_u(t, x_min),
        u(t, x_max) ~ analytic_u(t, x_max)]

    domains = [t ∈ Interval(t_min, t_max),
        x ∈ Interval(x_min, x_max)]

    dx = 0.05

    @named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])

    disc = MOLFiniteDifference([x => dx], t, upwind_order=1)

    prob = discretize(pdesys, disc)

    sol = solve(prob, Tsit5())

    grid = get_discrete(pdesys, disc)
    x_disc = grid[x]
    solu = [map(d -> sol[d][i], grid[u(t, x)]) for i in 1:length(sol[t])]

    for (i, t) in enumerate(sol.t[5:end])
        u_analytic = analytic_u.([t], x_disc)
        u_disc = solu[i+4]
        @test_broken all(isapprox.(u_analytic, u_disc, atol=1e-3))
    end
end

@test_broken begin #@testset "Inviscid Burgers equation, 1D, u(0, x) ~ x, Non-Uniform" begin
    @parameters x t
    @variables u(..)
    Dx = Differential(x)
    Dt = Differential(t)
    x_min = 0.0
    x_max = 1.0
    t_min = 0.0
    t_max = 1.0

    analytic_u(t, x) = x / (t + 1)

    eq = Dt(u(t, x)) + u(t, x) * Dx(u(t, x)) ~ 0

    bcs = [u(0, x) ~ x,
        u(t, x_min) ~ analytic_u(t, x_min),
        u(t, x_max) ~ analytic_u(t, x_max)]

    domains = [t ∈ Interval(t_min, t_max),
        x ∈ Interval(x_min, x_max)]

    @named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])

    dx = 0:0.05:1
    dx = collect(dx)
    dx[2:end-1] .= dx[2:end-1] .+ rand([0.001, -0.001], length(dx[2:end-1]))

    disc = MOLFiniteDifference([x => dx], t, upwind_order=1)

    prob = discretize(pdesys, disc)

    sol = solve(prob, Tsit5())

    grid = get_discrete(pdesys, disc)
    x_disc = grid[x]
    solu = [map(d -> sol[d][i], grid[u(t, x)]) for i in 1:length(sol[t])]

    for (i, t) in enumerate(sol.t[5:end])
        u_analytic = analytic_u.([t], x_disc)
        u_disc = solu[i+4]
        @test_broken all(isapprox.(u_analytic, u_disc, atol=1e-3))
    end
end

# # Exact solutions from: https://www.sciencedirect.com/science/article/pii/S0898122110003883
# @test_broken begin# @testset "Test 01: Burger's Equation 2D" begin
#        @parameters x y t
#        @variables u(..) v(..)
#        Dt = Differential(t)
#        Dx = Differential(x)
#        Dy = Differential(y)
#        Dxx = Differential(x)^2
#        Dyy = Differential(y)^2

#        R = 80.0

#        x_min = y_min = t_min = 0.0
#        x_max = y_max = 1.0
#        t_max = 0.05

#        u_exact(x,y,t) = 3/4 - 1/(4*(1+exp(R*(-t-4x+4y)/32)))
#        v_exact(x,y,t) = 3/4 + 1/(4*(1+exp(R*(-t-4x+4y)/32)))

#        eq = [Dt(u(x,y,t)) + u(x,y,t)*Dx(u(x,y,t)) + v(x,y,t)*Dy(u(x,y,t)) ~ 1/R *(Dxx(u(x,y,t)) + Dyy(u(x,y,t))),
#              Dt(v(x,y,t)) + u(x,y,t)*Dx(v(x,y,t)) + v(x,y,t)*Dy(v(x,y,t)) ~ 1/R *(Dxx(v(x,y,t)) + Dyy(v(x,y,t)))]

#        domains = [x ∈ Interval(x_min, x_max),
#                   y ∈ Interval(y_min, y_max),
#                   t ∈ Interval(t_min, t_max)]

#        bcs = [u(x,y,0) ~ u_exact(x,y,0),
#               u(0,y,t) ~ u_exact(0,y,t),
#               u(x,0,t) ~ u_exact(x,0,t),
#               u(1,y,t) ~ u_exact(1,y,t),
#               u(x,1,t) ~ u_exact(x,1,t),

#               v(x,y,0) ~ v_exact(x,y,0),
#               v(0,y,t) ~ v_exact(0,y,t),
#               v(x,0,t) ~ v_exact(x,0,t),
#               v(1,y,t) ~ v_exact(1,y,t),
#               v(x,1,t) ~ v_exact(x,1,t)]

#        @named pdesys = PDESystem(eq,bcs,domains,[t,x,y],[u(x,y,t),v(x,y,t)])

#        # Method of lines discretization
#        dx = 0.05
#        dy = 0.05

#        order = 4

#        discretization = MOLFiniteDifference([x=>dx, y=>dy], t, approx_order=order)

#        # Convert the PDE problem into an ODE problem
#        prob = discretize(pdesys,discretization)

#        sol = solve(prob,Tsit5())

#        Nx = floor(Int64, (x_max - x_min) / dx) + 1
#        Ny = floor(Int64, (y_max - y_min) / dy) + 1

#        @variables u[1:Nx,1:Ny](t)
#        @variables v[1:Nx,1:Ny](t)
#        t = sol[t]
#        # anim = @animate for k in 1:length(t)
#        #        solu′ = reshape([sol[u[(i-1)*Ny+j]][k] for i in 1:Nx for j in 1:Ny],(Nx,Ny))
#        #        solv′ = reshape([sol[v[(i-1)*Ny+j]][k] for i in 1:Nx for j in 1:Ny],(Nx,Ny))
#        #        heatmap(solu′)
#        # end
#        # gif(anim, "plots/Burgers2Dsol.gif", fps = 5)

#        solu′ = reshape([sol[u[(i-1)*Ny+j]][end] for i in 1:Nx for j in 1:Ny],(Nx,Ny))
#        solv′ = reshape([sol[v[(i-1)*Ny+j]][end] for i in 1:Nx for j in 1:Ny],(Nx,Ny))

#        r_space_x = x_min:dx:x_max
#        r_space_y = y_min:dy:y_max

#        asfu = reshape([u_exact(t_max,r_space_x[i],r_space_y[j]) for j in 1:Ny for i in 1:Nx],(Nx,Ny))
#        asfv = reshape([v_exact(t_max,r_space_x[i],r_space_y[j]) for j in 1:Ny for i in 1:Nx],(Nx,Ny))

#        asfu[1,1] = asfu[1, end] = asfu[end, 1] = asfu[end, end] = 0.
#        asfv[1,1] = asfv[1, end] = asfv[end, 1] = asfv[end, end] = 0.

#        # anim = @animate for T in t
#        #        asfu = reshape([u_exact(T,r_space_x[i],r_space_y[j]) for j in 1:Ny for i in 1:Nx],(Nx,Ny))
#        #        asfv = reshape([v_exact(T,r_space_x[i],r_space_y[j]) for j in 1:Ny for i in 1:Nx],(Nx,Ny))

#        #        heatmap(asfu)
#        # end
#        # gif(anim, "plots/Burgers2Dexact.gif", fps = 5)

#        mu = max(asfu...)
#        mv = max(asfv...)
#        @test_broken asfu / mu ≈ solu′ / mu  atol=0.2
#        @test_broken asfv / mv ≈ solv′ / mv  atol=0.2
# end


export AbstractNLS # TODO: do not export
export parameter_size, residue_size
export eval_r, eval_r_J
export eval_nls_fobj, eval_nls_∇fobj, eval_nls_∇∇fobj

using LinearAlgebra: dot, mul!
using LinearAlgebra.BLAS: BlasFloat, syrk!, gemv!

@doc raw"""
```julia
abstract type AbstractNLS end
```

Defines an abstract non-linear least squares problem (NLS). In our
context such problem is essentially a differentiable function ``\mathbf{r}``:

```math
\mathbf{r}: \theta\in\mathbb{R}^{n_θ}\mapsto \mathbf{r}(\mathbf{\theta})\in\mathbb{R}^{n_S}
```
where:
- ``\mathbf{r}(\mathbf{θ})∈\mathbb{R}^{n_S}`` is the residue vector,
- ``\mathbf{θ}∈\mathbb{R}^{n_θ}`` is the parameter vector to be optimized

The objective function to minimize is:

```math
f(θ)=\frac{1}{2}\| \mathbf{r}(θ) \|^2
```

The classical approach uses a linear approximation of ``\mathbf{r}``:
```math
\mathbf{r}(\mathbf{θ}+δ\mathbf{θ})\approx \mathbf{r}(\mathbf{θ}) + \mathbf{J}(\mathbf{θ})\cdot δ\mathbf{θ}
```
where ``\mathbf{J}`` is the Jacobian:
```math
\mathbf{J}_{i,j}=\partial_j r^i(\mathbf{θ}),\ i\in[1,n_S],\ j\in[1,n_θ]
```
This leads to
```math
f(\mathbf{θ}+δ\mathbf{θ})\approx f(\mathbf{θ}) + \langle \nabla f, δ\mathbf{θ} \rangle + \frac{1}{2}  \langle \nabla^2 f \cdot δ\mathbf{θ},  δ\mathbf{θ} \rangle
```

Where the gradient ``\nabla f`` is ``\mathbf{J}^t \mathbf{r}`` and the
(approximate) Hessian ``\nabla^2 f`` is ``\mathbf{J}^t \mathbf{J}``.

To implement such model, you must define the following functions:
- [`parameter_size`](@ref) : returns ``n_θ``
- [`residue_size`](@ref) : returns ``n_S``
- [`eval_r`](@ref) : compute ``\mathbf{r}``
- [`eval_r_J`](@ref) : compute ``(\mathbf{r}, \mathbf{J})``
"""
abstract type AbstractNLS end

# ================================================================
# Interface...
# ================================================================
#
@doc raw"""
    parameter_size(nls::AbstractNLS)

Return the dimension ``n_θ`` of the parameter vector ``θ``.
"""
parameter_size(nls::AbstractNLS) = @assert(false,"To implement")

@doc raw"""
    sample_size(nls::AbstractNLS)

Return the dimension ``n_S`` of the residue vector ``r``.
"""
residue_size(nls::AbstractNLS) = @assert(false,"To implement")

@doc raw"""
```julia
eval_r(nls::AbstractNLS,
        θ::AbstractVector) -> r
```

Compte the residual vector ``\mathbf{r}``
"""
eval_r(nls::AbstractNLS,
       θ::AbstractVector) = @assert(false,"To implement")

@doc raw"""
```julia
eval_r_J(nls::AbstractNLS,θ::AbstractVector) -> (r,J)
```

Compute the residual the vector ``\mathbf{r}`` and its Jacobian ``\mathbf{J}``
"""
eval_r_J(nls::AbstractNLS,
         θ::AbstractVector) = @assert(false,"To implement")

# ================================================================
# Convenience functions...
# ================================================================
#

@doc raw"""
```julia
eval_nls_fobj(r::AbstractVector{T}) -> f(θ)
```

Compute ``f(θ)=\frac{1}{2}\| \mathbf{r}(\mathbf{θ}) \|^2``
"""
eval_nls_fobj(r::AbstractVector) = dot(r,r)/2

@doc raw"""
```julia
eval_nls_∇fobj(r,J) -> ∇fobj
```

Compute the gradient: ``\nabla f(\mathbf{θ}) = \mathbf{J}^t\mathbf{r}``
"""
function eval_nls_∇fobj(r::AbstractVector, J::AbstractMatrix)
    J'*r
end

@doc raw"""
```julia
eval_nls_∇∇fobj(J) -> ∇∇fobj
```

Compute the (approximate) Hessian: ``\nabla^2 f(\mathbf{θ}) = \mathbf{J}^t\mathbf{J}``
"""
function eval_nls_∇∇fobj(J::AbstractMatrix)
    Symmetric(J'*J)
end


# retrieve apm.jl from GitHub or APMonitor.com
include("../apm.jl")

s = "http://byu.apmonitor.com"
a = "minlp"

apm(s,a,"clear all")

#  minimize    f(x)
#  subject to  h(x) > 0
#              g(x) = 0
#
# In this case:
#  f(x) = x1*x4*(x1+x2+x3) + x3
#  h(x) = x1*x2*x3*x4 - 25
#  g(x) = x1^2 + x2^2 + x3^2 + x4^2 - 40
#
# x1 and x2 are continuous variables
# int_x3 and int_x4 are integer variables

if isfile("minlp.apm")
  # Load model file (if it exists)
  apm_load(s,a,"minlp.apm")
else
  # Declare variables and equations here
  # Declare variables
  apm(s,a,"Variables")
  apm(s,a," x1 = 1, >=1, <=5")
  apm(s,a," x2 = 5, >=1, <=5")
  apm(s,a," int_x3 = 5, >=1, <=5")
  apm(s,a," int_x4 = 1, >=1, <=5")

  # Declare objective and equations
  apm(s,a,"Equations")
  apm(s,a," minimize x1*int_x4*(x1+x2+int_x3) + int_x3")
  apm(s,a," x1*x2*int_x3*int_x4 > 25")
  apm(s,a," x1^2 + x2^2 + int_x3^2 + int_x4^2 = 40")
end

# Solve optimization problem
output = apm(s,a,"solve")
println(output)

# Retrieve solution
y = apm_sol(s,a)

println("Solution Results")
println("  x1 = " * string(y["x1"]))
println("  x2 = " * string(y["x2"]))
println("  int_x3 = " * string(y["int_x3"]))
println("  int_x4 = " * string(y["int_x4"]))


# The sea floor is getting steeper. Maybe the sleigh keys got carried this way?
#
# A massive school of glowing lanternfish swims past. They must spawn quickly
# to reach such large numbers - maybe exponentially quickly? You should model
# their growth rate to be sure.
#
# Although you know nothing about this specific species of lanternfish, you
# make some guesses about their attributes. Surely, each lanternfish creates a
# new lanternfish once every 7 days.
#
# However, this process isn't necessarily synchronized between every
# lanternfish - one lanternfish might have 2 days left until it creates another
# lanternfish, while another might have 4. So, you can model each fish as a
# single number that represents the number of days until it creates a new
# lanternfish.
#
# Furthermore, you reason, a new lanternfish would surely need slightly longer
# before it's capable of producing more lanternfish: two more days for its
# first cycle.
#
# So, suppose you have a lanternfish with an internal timer value of 3:
#
# - After one day, its internal timer would become 2.
# - After another day, its internal timer would become 1.
# - After another day, its internal timer would become 0.
# - After another day, its internal timer would reset to 6, and it would create
#   a new lanternfish with an internal timer of 8.
# - After another day, the first lanternfish would have an internal timer of 5,
#   and the second lanternfish would have an internal timer of 7.
#
# A lanternfish that creates a new fish resets its timer to 6, not 7 (because 0
# is included as a valid timer value). The new lanternfish starts with an
# internal timer of 8 and does not start counting down until the next day.
#
# Realizing what you're trying to do, the submarine automatically produces a
# list of the ages of several hundred nearby lanternfish (your puzzle input).
# For example, suppose you were given the following list:
#
# 3,4,3,1,2
#
# This list means that the first fish has an internal timer of 3, the second
# fish has an internal timer of 4, and so on until the fifth fish, which has an
# internal timer of 2. Simulating these fish over several days would proceed as
# follows:
#
# Initial state: 3,4,3,1,2
# After  1 day:  2,3,2,0,1
# After  2 days: 1,2,1,6,0,8
# After  3 days: 0,1,0,5,6,7,8
# After  4 days: 6,0,6,4,5,6,7,8,8
# After  5 days: 5,6,5,3,4,5,6,7,7,8
# After  6 days: 4,5,4,2,3,4,5,6,6,7
# After  7 days: 3,4,3,1,2,3,4,5,5,6
# After  8 days: 2,3,2,0,1,2,3,4,4,5
# After  9 days: 1,2,1,6,0,1,2,3,3,4,8
# After 10 days: 0,1,0,5,6,0,1,2,2,3,7,8
# After 11 days: 6,0,6,4,5,6,0,1,1,2,6,7,8,8,8
# After 12 days: 5,6,5,3,4,5,6,0,0,1,5,6,7,7,7,8,8
# After 13 days: 4,5,4,2,3,4,5,6,6,0,4,5,6,6,6,7,7,8,8
# After 14 days: 3,4,3,1,2,3,4,5,5,6,3,4,5,5,5,6,6,7,7,8
# After 15 days: 2,3,2,0,1,2,3,4,4,5,2,3,4,4,4,5,5,6,6,7
# After 16 days: 1,2,1,6,0,1,2,3,3,4,1,2,3,3,3,4,4,5,5,6,8
# After 17 days: 0,1,0,5,6,0,1,2,2,3,0,1,2,2,2,3,3,4,4,5,7,8
# After 18 days: 6,0,6,4,5,6,0,1,1,2,6,0,1,1,1,2,2,3,3,4,6,7,8,8,8,8
#
# Each day, a 0 becomes a 6 and adds a new 8 to the end of the list, while each
# other number decreases by 1 if it was present at the start of the day.
#
# In this example, after 18 days, there are a total of 26 fish. After 80 days,
# there would be a total of 5934.
#
# Find a way to simulate lanternfish. How many lanternfish would there be after
# 80 days?

# read input
fish = parse.(Int8, split(readline("input.txt"), ','))

for _ in 1:80
  fish .-= 1
  replicating = findall(isequal(-1), fish)
  if length(replicating) > 0
    fish[replicating] .= 6
    append!(fish, repeat([8], length(replicating)))
  end
end

println("Number of fish: ", length(fish))


function construct_inpTheta(X_test, X_domain, X_boundary, set_sigma)
    N_test = size(X_test)[1]
    N_domain = size(X_domain)[1]
    N_boundary = size(X_boundary)[1]
    inpTheta = zeros((N_test, 2*N_domain+N_boundary))
    for iter_i = 1:N_test
        for iter_j = 1:N_domain
            inpTheta[iter_i,iter_j] = Delta_x_kappa(X_test[iter_i,1],X_test[iter_i,2],X_domain[iter_j,1],X_domain[iter_j,2],set_sigma)
        end
        for iter_j = N_domain+1:2*N_domain
            inpTheta[iter_i,iter_j] = kappa(X_test[iter_i,1],X_test[iter_i,2],X_domain[iter_j-N_domain,1],X_domain[iter_j-N_domain,2],set_sigma)
        end
        for iter_j = 2*N_domain+1:2*N_domain+N_boundary
            inpTheta[iter_i,iter_j] = kappa(X_test[iter_i,1],X_test[iter_i,2],X_boundary[iter_j-2*N_domain,1],X_boundary[iter_j-2*N_domain,2],set_sigma)
        end
    end
    return inpTheta
end

function extend_solution(X_test, X_domain, X_boundary, L, v, set_sigma)
    inp_Theta = construct_inpTheta(X_test, X_domain, X_boundary, set_sigma)
    return inp_Theta*(L'\(L\v))
end

function get_extended_solution(num_pts, X_domain, X_boundary, L, sol, set_sigma)
    x = range(0,1,length=num_pts)
    y = range(0,1,length=num_pts)
    xx = x' .* ones(num_pts)
    yy = ones(num_pts)' .* y

    xxv = Matrix{Float64}(reshape(xx,(num_pts^2,1)))
    yyv = Matrix{Float64}(reshape(yy,(num_pts^2,1)))

    X_test = hcat(xxv,yyv)

    N_domain = size(X_domain)[1]
    N_boundary = size(X_boundary)[1]
    rhs_f = zeros((N_domain,1))
    bdy_g = zeros((N_boundary,1))

    for iter_i = 1:N_domain
        rhs_f[iter_i] = f(X_domain[iter_i,1],X_domain[iter_i,2])
    end
    for iter_i = 1:N_boundary
        bdy_g[iter_i] = g(X_boundary[iter_i,1],X_boundary[iter_i,2])
    end

#     X_test=hcat([[x[i];y[j]] for i = 1 : num_pts for j = 1 : num_pts])
    temp_vec = [alpha*sol.^m-rhs_f; sol; bdy_g]
    u_extended = Matrix{Float64}(reshape(extend_solution(X_test, X_domain, X_boundary, L, temp_vec, set_sigma),(num_pts,num_pts)))
    return x, y, u_extended
end