module Test_psutils

using PeriodicSystems
using DescriptorSystems
using Symbolics
using Test
using LinearAlgebra
using LinearAlgebra: BlasInt
using ApproxFun
#using BenchmarkTools

println("Test_psutils")

@testset "test_psutils" begin

# MB03BD example
A1 = Matrix{Float64}(I,3,3); A2 = [   1.0   2.0   0.0; 4.0  -1.0   3.0; 0.0   3.0   1.0]; A3 = Matrix{Float64}(I,3,3);
E1 =  [2.0   0.0   1.0; 0.0  -2.0  -1.0; 0.0   0.0   3.0]; E2 = Matrix{Float64}(I,3,3);
E3 = [ 1.0   0.0   1.0; 0.0   4.0  -1.0; 0.0   0.0  -2.0];
ev = eigvals(inv(E1)*A2*inv(E3))

# using the SLICOT wrapper
A = reshape([E1 A2 E3],3,3,3);
KMAX = 3
NMAX = 3
LDA1 = NMAX
LDA2 = NMAX
LDQ1 = NMAX
LDQ2 = NMAX
LDWORK = KMAX + max( 2*NMAX, 8*KMAX )
LIWORK = 2*KMAX + NMAX
QIND = Array{BlasInt,1}(undef, KMAX)
S = [-1,1,-1];
Q = Array{Float64,3}(undef, LDQ1, LDQ2, KMAX)
ALPHAR = Array{Float64,1}(undef, NMAX)
ALPHAI = Array{Float64,1}(undef, NMAX)
BETA = Array{Float64,1}(undef, NMAX)
SCAL = Array{BlasInt,1}(undef, NMAX)
IWORK = Array{BlasInt,1}(undef, LIWORK)
DWORK = Array{Float64,1}(undef, LDWORK)

mb03bd!('T','C','I',QIND,3,3,2,1,3,S,A,Q,ALPHAR, ALPHAI, BETA, SCAL, LIWORK, LDWORK)

poles = (ALPHAR+im*ALPHAI) ./ BETA .* (2. .^SCAL)

@test sort(real(poles)) ≈ sort(real(ev)) &&
      sort(imag(poles)) ≈ sort(imag(ev))

A = reshape([A2 inv(E3) inv(E1)],3,3,3);
mb03wd!('S','I',3,3,1,3,1,3,A,Q,ALPHAR, ALPHAI, LDWORK)

poles = (ALPHAR+im*ALPHAI)

@test sort(real(poles)) ≈ sort(real(ev)) &&
      sort(imag(poles)) ≈ sort(imag(ev))

# MB03VD example
A1 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.];
A2 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.];
n = 4; K = 2;
A = reshape([A1 A2], n, n, K);
@time H, Z, ihess = phess(A; rev = false);
@test check_psim(A,Z,H; rev = false) && istriu(H[:,:,ihess],-1) && ihess == 1

@time H1, Z1, ihess1 = phess(A; rev = false, withZ = false);
@test H[1] == H1[1] && H[2] == H1[2]

@time H, Z, ihess = phess(A, hind = 2, rev = false);
@test check_psim(A,Z,H; rev = false) && istriu(H[:,:,ihess],-1) && ihess == 2

@time H1, Z1, ihess1 = phess(A, hind = 2, rev = false, withZ = false);
@test H[1] == H1[1] && H[2] == H1[2]

@time H, Z, ihess = phess(A; rev = true);
@test check_psim(A,Z,H; rev = true) && istriu(H[:,:,ihess],-1) && ihess == 1

@time H1, Z1, ihess1 = phess(A; rev = true, withZ = false);
@test H[1] == H1[1] && H[2] == H1[2]

@time H, Z, ihess = phess(A; hind = 2, rev = true);
@test check_psim(A,Z,H; rev = true) && istriu(H[:,:,ihess],-1) && ihess == 2

@time H1, Z1, ihess1 = phess(A; hind = 2, rev = true, withZ = false);
@test H[1] == H1[1] && H[2] == H1[2]

ev = eigvals(A[:,:,1]*A[:,:,2])
@time S, Z, eigs, ischur, α, γ = pschur(A; rev = false);
@test check_psim(A,Z,S; rev = false) && istriu(S[:,:,ischur],-1) && ischur == 1
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

@time S1, Z1, eigs1, ischur1, α1, γ1 = pschur(A; rev = false, withZ = false);
@test S[1] == S1[1] && S[2] == S1[2] && eigs == eigs1

@time S, Z, eigs, ischur, α, γ = pschur(A, sind = 2, rev = false);
@test check_psim(A,Z,S; rev = false) && istriu(S[:,:,ischur],-1) && ischur == 2
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

ev = eigvals(A[:,:,2]*A[:,:,1])
@time S, Z, eigs, ischur, α, γ = pschur(A; rev = true);
@test check_psim(A,Z,S; rev = true) && istriu(S[:,:,ischur],-1) && ischur == 1
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

@time S, Z, eigs, ischur, α, γ = pschur(A; sind = 2, rev = true);
@test check_psim(A,Z,S; rev = true) && istriu(S[:,:,ischur],-1) && ischur == 2
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

A1 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.];
n = 4; K = 19; hind = 10; rev = false;
A = reshape(hcat([rand()*A1 for i in 1:K]...), n, n, K);
@time H, Z, ihess = phess(A; hind, rev);
@test check_psim(A,Z,H; rev) && istriu(H[:,:,ihess],-1) && ihess == hind

@time H, Z, ihess = phess1(A; hind, rev);
@test check_psim(A,Z,H; rev) && istriu(H[:,:,ihess],-1) && ihess == hind

@time S, Z, eigs, ischur, = pschur(A; sind = hind, rev);
@test check_psim(A,Z,S; rev)

@time S, Z, eigs, ischur, = pschur1(A; sind = hind, rev);
@test check_psim(A,Z,S; rev)

# this would fail
# n =4; K = 19; hind = 10; rev = false;
# A = reshape(hcat([rand()*A1 for i in 1:K]...), n, n, K);
# @time S, Z, eigs, ischur, = pschur2(A; sind = hind, rev);
# @test check_psim(A,Z,S; rev)

n =4; K = 15; hind = 10; rev = false;
A = reshape(hcat([rand()*A1 for i in 1:K]...), n, n, K);
@time S, Z, eigs, ischur, = pschur2(A; sind = hind, rev);
@test check_psim(A,Z,S; rev)

hind = 14; rev = true;
@time H, Z, ihess = phess(A; hind, rev);
@test check_psim(A,Z,H; rev) && istriu(H[:,:,ihess],-1) && ihess == hind

@time H, Z, ihess = phess1(A; hind, rev);
@test check_psim(A,Z,H; rev) && istriu(H[:,:,ihess],-1) && ihess == hind

@time S, Z, eigs, ischur, = pschur(A; sind = hind, rev);
@test check_psim(A,Z,S; rev)

@time S, Z, eigs, ischur, = pschur1(A; sind = hind, rev);
@test check_psim(A,Z,S; rev)

# this would fail
# n =4; K = 19; hind = 10; rev = false;
# A = reshape(hcat([rand()*A1 for i in 1:K]...), n, n, K);
# @time S, Z, eigs, ischur, = pschur2(A; sind = hind, rev);
# @test check_psim(A,Z,S; rev)

n =4; K = 15; hind = 10; rev = false;
A = reshape(hcat([rand()*A1 for i in 1:K]...), n, n, K);
@time S, Z, eigs, ischur, = pschur2(A; sind = hind, rev);
@test check_psim(A,Z,S; rev)

A1 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3];
A2 = [1.5 -.7 3.5; 1.  0.  2.; 1.5 -.7 2.5 ; 1.  0.  2.];
Ai=[A1,A2];
A = [Ai[j] for j = 1:2,  i in 1:9][:]

rev = false;
@time S, Z, eigs, ischur, = pschur(A; rev);
@test check_psim(A,Z,S; rev)

rev = true;
@time S, Z, eigs, ischur, = pschur(A; rev);
@test check_psim(A,Z,S; rev)

# @time S1, Z1, eigs1, ischur1 = pschurw(A; rev = false);
# @test check_psim(A,Z1,S1; rev = false)

A1 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.];
ev1 = eigvals(A1)
n = 4; K = 18;

k = 1; N = div(K,k)
Ap = PeriodicArray(reshape(hcat([A1 for i in 1:N]...), n, n, N), K; nperiod = k)
@time ev = pseig(Ap)
@test sort(real(ev)) ≈ sort(real(ev1.^K)) && sort(imag(ev)) ≈ sort(imag(ev1.^K))

k = 9; N = div(K,k)
Ap = PeriodicArray(reshape(hcat([A1 for i in 1:N]...), n, n, N),K; nperiod = k)
@time ev = pseig(Ap)
@test sort(real(ev)) ≈ sort(real(ev1.^K)) && sort(imag(ev)) ≈ sort(imag(ev1.^K))

k = 18; N = div(K,k)
Ap = PeriodicArray(reshape(hcat([A1 for i in 1:N]...), n, n, N),K; nperiod = k)
@time ev = pseig(Ap)
@test sort(real(ev)) ≈ sort(real(ev1.^K)) && sort(imag(ev)) ≈ sort(imag(ev1.^K))

# modified MB03VD example
A1 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.];
A2 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.] .+ 1;
A3 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.] .- 1;

n = 4; K = 3; sind = 2;
A = reshape([A1 A2 A3], n, n, K);
#AH, Z = phess(A; h);
@time S, Z, eigs, ischur = pschur(A; sind);
@test check_psim(A, Z, S) && istriu(S[:,:,ischur],-1)

ev = eigvals(A[:,:,1]*A[:,:,2]*A[:,:,3])
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

i1 = reverse(1:K)
@time a, e = psreduc_reg(view(A,:,:,i1)); eigs = eigvals(a,e);
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

# modified MB03VD example
A1 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.];
A2 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.] .+ 1;
A3 = [1.5 -.7 3.5 -.7; 1.  0.  2.  3.; 1.5 -.7 2.5 -.3; 1.  0.  2.  1.] .- 1;
n = 4; K = 3;
A = reshape([A1 A2 A3], n, n, K);
M = PeriodicArray(A,K);
MA = PeriodicMatrix([A1,A2,A3],K);

@time eigs = pseig(A); ev = pseig(A,fast = true);
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      isapprox(sort(imag(eigs)),sort(imag(ev)), atol = 1.e-7)
@test sort(eigs, by = real) ≈ sort(ev, by = real)
@test sort(eigs, by = abs) ≈ sort(ev, by = abs)
@test sort(psceig(M),by = real) ≈ sort(psceig(MA),by = real)

@time eigs = pseig(A); ev = pseig(M,fast = true);
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      isapprox(sort(imag(eigs)),sort(imag(ev)), atol = 1.e-7)
@test sort(eigs, by = real) ≈ sort(ev, by = real)
@test sort(eigs, by = abs) ≈ sort(ev, by = abs)

@time AH, Z, ihess = phess(A; rev = false);
@test check_psim(A,Z,AH; rev = false) && istriu(AH[:,:,ihess],-1) && ihess == 1

@time AH, Z, ihess = phess(A, hind = 2, rev = false);
@test check_psim(A,Z,AH; rev = false) && istriu(AH[:,:,ihess],-1) && ihess == 2

@time AH, Z, ihess = phess(A; rev = true);
@test check_psim(A,Z,AH; rev = true) && istriu(AH[:,:,ihess],-1) && ihess == 1

@time AH, Z, ihess = phess(A; hind = 2, rev = true);
@test check_psim(A,Z,AH; rev = true) && istriu(AH[:,:,ihess],-1) && ihess == 2

ev = eigvals(A[:,:,1]*A[:,:,2]*A[:,:,3])
@time S, Z, eigs, ischur, α, γ = pschur(A; rev = false);
@test check_psim(A,Z,S; rev = false) && istriu(S[:,:,ischur],-1) && ischur == 1
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

@time S, Z, eigs, ischur, α, γ = pschur(A, sind = 2, rev = false);
@test check_psim(A,Z,S; rev = false) && istriu(S[:,:,ischur],-1) && ischur == 2
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

ev = eigvals(A[:,:,3]*A[:,:,2]*A[:,:,1])
@time S, Z, eigs, ischur, α, γ = pschur(A; rev = true);
@test check_psim(A,Z,S; rev = true) && istriu(S[:,:,ischur],-1) && ischur == 1
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

@time S, Z, eigs, ischur, α, γ = pschur(A; sind = 2, rev = true);
@test check_psim(A,Z,S; rev = true) && istriu(S[:,:,ischur],-1) && ischur == 2
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

ev = eigvals(A[:,:,3]*A[:,:,2]*A[:,:,1])
a, e = psreduc_reg(A); eigs = eigvals(a,e);
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      sort(imag(eigs)) ≈ sort(imag(ev))

A = [A1, A2, A3];
@time eigs = pseig(A); ev = pseig(A,fast = true);
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      isapprox(sort(imag(eigs)),sort(imag(ev)), atol = 1.e-7)

@time eigs = pseig(A; rev = false); ev = pseig(A; fast = true, rev = false);
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      isapprox(sort(imag(eigs)),sort(imag(ev)), atol = 1.e-7)

@time eigs = pseig(A, 2; rev = false);
@time ev = pseig(A, 2; fast = true, rev = false);
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      isapprox(sort(imag(eigs)),sort(imag(ev)), atol = 1.e-7)

A3 = rand(5,4); A2 = rand(4,3); A1 = rand(3,5);
#A = [A1, A2, A3];
K = 3;
A = PeriodicMatrix([A1, A2, A3],K);
@time eigs = pseig(A); ev = pseig(A,fast = true);
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      isapprox(sort(imag(eigs)),sort(imag(ev)), atol = 1.e-7)
@test sort(eigs, by = real) ≈ sort(ev, by = real)
@test sort(eigs, by = abs) ≈ sort(ev, by = abs)

ev = eigvals(prod(A.M[K:-1:1])); nmin = minimum(size.(A.M,1))
@time S, Z, eigs, ischur, α, γ = pschur(A.M; rev = true);
@test check_psim(A.M,Z,S; rev = true) && istriu(S[ischur][1:nmin,1:nmin],-1) && eigs == α.*γ
@test isapprox(sort(real(eigs)),sort(real(ev)), atol = 1.e-7) &&
      isapprox(sort(imag(eigs)),sort(imag(ev)), atol = 1.e-7)

A = [A3, A2, A1];
ev = eigvals(prod(A)); nmin = minimum(size.(A,1))

@time S, Z, eigs, ischur, α, γ = pschur(A; rev = false);
@test check_psim(A,Z,S; rev = false) && istriu(S[ischur][1:nmin,1:nmin],-1) && eigs == α.*γ
@test sort(real(eigs)) ≈ sort(real(ev)) &&
      isapprox(sort(imag(eigs)),sort(imag(ev)), atol = 1.e-7)

# periodic time-varying matrices from IFAC2005 paper
@variables t
A = PeriodicSymbolicMatrix([cos(t) 1; 1 1-sin(t)],2*pi);
println("A(t) = $(A.F)")
B = HarmonicArray([0;1],[[1;0]],[[1;-1]],2*pi);
println("B(t) = $(convert(PeriodicSymbolicMatrix,B).F)")
C = PeriodicFunctionMatrix(t-> [sin(2*t)+cos(2*t) 1],pi);
println("C(t) = $(C.f(t))")
D = PeriodicFunctionMatrix(zeros(1,1),pi)
K = 100
Δ = 2*pi/K
ts = (0:K-1)*Δ

psys = ps(A,B,C,D);
psys = ps(A,B,C);
ps(A);
ps(B);
ps(C);
ps(D);
psys = ps(A,B,C,D,4*pi);
psys = ps(A,B,C,4*pi);
psys = ps(HarmonicArray,A,B,C,D,4*pi);
psys = ps(PeriodicSymbolicMatrix, A,B,C,4*pi);

psys = ps(rss(4,3,2),10)
@test islti(psys)
psys=ps(rand(2,2),rand(2),rand(3,2),rand(3),10)
psys=ps(rand(2,2),rand(2),rand(3,2),rand(3),10;Ts = 1)
psys=ps(rand(2,2),rand(2),rand(3,2),10;Ts = 1)

psys = PeriodicStateSpace(convert(PeriodicFunctionMatrix,A),
                          convert(PeriodicFunctionMatrix,B),
                          convert(PeriodicFunctionMatrix,C),
                          convert(PeriodicFunctionMatrix,D));
print(psys);
convert(PeriodicStateSpace{PeriodicSymbolicMatrix},psys);
convert(PeriodicStateSpace{HarmonicArray},psys);
convert(PeriodicStateSpace{PeriodicTimeSeriesMatrix},psys);
convert(PeriodicStateSpace{FourierFunctionMatrix},psys);

psys = PeriodicStateSpace(convert(PeriodicFunctionMatrix{:c,BigFloat},A),
                          convert(PeriodicFunctionMatrix{:c,BigFloat},B),
                          convert(PeriodicFunctionMatrix{:c,BigFloat},C),
                          convert(PeriodicFunctionMatrix{:c,BigFloat},D));

psys = ps(convert(PeriodicFunctionMatrix,A),
          convert(PeriodicFunctionMatrix,B),
          convert(PeriodicFunctionMatrix,C),
          convert(PeriodicFunctionMatrix,D));

psys = ps(convert(PeriodicFunctionMatrix,A),
          convert(PeriodicFunctionMatrix,B),
          convert(PeriodicFunctionMatrix,C));

psys = PeriodicStateSpace(convert(PeriodicSymbolicMatrix,A),
                          convert(PeriodicSymbolicMatrix,B),
                          convert(PeriodicSymbolicMatrix,C),
                          convert(PeriodicSymbolicMatrix,D));
print(psys);
convert(PeriodicStateSpace{PeriodicFunctionMatrix},psys);
convert(PeriodicStateSpace{HarmonicArray},psys);
convert(PeriodicStateSpace{PeriodicTimeSeriesMatrix},psys);
convert(PeriodicStateSpace{FourierFunctionMatrix},psys);

psys = ps(convert(PeriodicSymbolicMatrix,A),
          convert(PeriodicSymbolicMatrix,B),
          convert(PeriodicSymbolicMatrix,C),
          convert(PeriodicSymbolicMatrix,D));

psys = ps(convert(PeriodicSymbolicMatrix,A),
          convert(PeriodicSymbolicMatrix,B),
          convert(PeriodicSymbolicMatrix,C));

psys = PeriodicStateSpace(convert(HarmonicArray,A),
                          convert(HarmonicArray,B),
                          convert(HarmonicArray,C),
                          convert(HarmonicArray,D));
print(psys);
convert(PeriodicStateSpace{PeriodicSymbolicMatrix},psys);
convert(PeriodicStateSpace{PeriodicTimeSeriesMatrix},psys);
convert(PeriodicStateSpace{PeriodicFunctionMatrix},psys);
convert(PeriodicStateSpace{FourierFunctionMatrix},psys);

psys = ps(convert(HarmonicArray,A),
          convert(HarmonicArray,B),
          convert(HarmonicArray,C),
          convert(HarmonicArray,D));

psys = ps(convert(HarmonicArray,A),
          convert(HarmonicArray,B),
          convert(HarmonicArray,C));

At = PeriodicTimeSeriesMatrix(convert(PeriodicFunctionMatrix,A).f.(ts),A.period);
Bt = PeriodicTimeSeriesMatrix(convert(PeriodicFunctionMatrix,B).f.(ts),B.period);
Ct = PeriodicTimeSeriesMatrix(convert(PeriodicFunctionMatrix,C).f.(ts),C.period);
Dt = PeriodicTimeSeriesMatrix(convert(PeriodicFunctionMatrix,D).f.(ts),D.period);
psys = PeriodicStateSpace(At, Bt, Ct, Dt);
print(psys);
convert(PeriodicStateSpace{PeriodicFunctionMatrix},psys);
convert(PeriodicStateSpace{PeriodicSymbolicMatrix},psys);
convert(PeriodicStateSpace{HarmonicArray},psys);
convert(PeriodicStateSpace{FourierFunctionMatrix},psys);

psys = ps(At, Bt, Ct, Dt);
psys = ps(At, Bt, Ct);

psys = PeriodicStateSpace(convert(FourierFunctionMatrix,A),
                          convert(FourierFunctionMatrix,B),
                          convert(FourierFunctionMatrix,C),
                          convert(FourierFunctionMatrix,D));
print(psys);
convert(PeriodicStateSpace{PeriodicFunctionMatrix},psys);
convert(PeriodicStateSpace{HarmonicArray},psys);
convert(PeriodicStateSpace{PeriodicTimeSeriesMatrix},psys);
convert(PeriodicStateSpace{PeriodicSymbolicMatrix},psys);

psys = ps(convert(FourierFunctionMatrix,A),
          convert(FourierFunctionMatrix,B),
          convert(FourierFunctionMatrix,C),
          convert(FourierFunctionMatrix,D));
psys = ps(convert(FourierFunctionMatrix,A),
          convert(FourierFunctionMatrix,B),
          convert(FourierFunctionMatrix,C));

psys = ps(FourierFunctionMatrix,A,B,C,D);

# #new tests
# AF = ffm2hr(psys.A)
# convert(PeriodicSymbolicMatrix,AF)

# constant dimensions
Ad = PeriodicMatrix([[1. 0; 0 0], [1 1;1 1], [0 1; 1 0]], 6, nperiod = 2);
Bd = PeriodicMatrix( [[ 1; 0 ], [ 1; 1]] ,2);
Cd = PeriodicMatrix( [[ 1 1 ], [ 1 0]] ,2);
Dd = PeriodicMatrix( [[ 1 ]], 1);
psys = PeriodicStateSpace(Ad,Bd,Cd,Dd);
convert(PeriodicStateSpace{PeriodicArray},psys);
convert(PeriodicStateSpace{PeriodicTimeSeriesMatrix},psys);

psys = ps(Ad,Bd,Cd,Dd);
psys = ps(Ad,Bd,Cd);

Ad = PeriodicMatrix([[1. 0], [1;1]],2);
Bd = PeriodicMatrix( [[ 1 ], [ 1; 1]] ,2);
Cd = PeriodicMatrix( [[ 1 1 ], [ 1 ]] ,2);
Dd = PeriodicMatrix( [[ 1 ]], 1);
psys = PeriodicStateSpace(Ad,Bd,Cd,Dd);
print(psys);
convert(PeriodicStateSpace{PeriodicArray},psys);
convert(PeriodicStateSpace{PeriodicTimeSeriesMatrix},psys);

psys = ps(Ad,Bd,Cd,Dd);
psys = ps(Ad,Bd,Cd);

Ad = PeriodicArray(rand(Float32,2,2,10),10);
Bd = PeriodicArray(rand(2,1,2),2);
Cd = PeriodicArray(rand(1,2,3),3);
Dd = PeriodicArray(rand(1,1,1), 1);
psys = PeriodicStateSpace(Ad,Bd,Cd,Dd);
print(psys);
convert(PeriodicStateSpace{PeriodicMatrix},psys)
convert(PeriodicStateSpace{PeriodicTimeSeriesMatrix},psys)

psys = ps(Ad,Bd,Cd,Dd);
psys = ps(Ad,Bd,Cd);
psys = ps(convert(PeriodicMatrix,Ad), Bd, convert(PeriodicMatrix,Cd), Dd);
psys = ps(convert(PeriodicMatrix,Ad), Bd, convert(PeriodicMatrix,Cd));

# symbolic periodic
@variables t
A = [cos(t) 1; 1 1-sin(t)];
B = [cos(t)+sin(t); 1-sin(t)];
C = [sin(t)+cos(2*t) 1];
Ap = PeriodicSymbolicMatrix(A,2*pi);
Bp = PeriodicSymbolicMatrix(B,2*pi);
Cp = PeriodicSymbolicMatrix(C,2*pi);

# functional expressions
tA(t::Real) = [cos(t) 1; 1 1-sin(t)];
tB(t::Real) = [cos(t)+sin(t); 1-sin(t)];
tC(t::Real) = [sin(t)+cos(2*t) 1];
# store snapshots as 3d arrays
N = 200;
tg = collect((0:N-1)*2*pi/N);
time = (0:N-1)*2*pi/N;
# At = reshape(hcat(A.(t)...),2,2,N);
# Bt = reshape(hcat(B.(t)...),2,1,N);
# Ct = reshape(hcat(C.(t)...),1,2,N);

# time series expressions
At = PeriodicTimeSeriesMatrix(tA.(time),2*pi);
Bt = PeriodicTimeSeriesMatrix(tB.(time),2*pi);
Ct = PeriodicTimeSeriesMatrix(tC.(time),2*pi);

# harmonic expressions
@time Ahr = ts2hr(At);
@test Ahr.values[:,:,1] ≈ [0. 1; 1 1] && Ahr.values[:,:,2] ≈ [1. 0; 0 -im]
@time Bhr = ts2hr(Bt);
@test Bhr.values[:,:,1] ≈ [0.; 1] && Bhr.values[:,:,2] ≈ [1.0+im; -im]
@time Chr = ts2hr(Ct);
@test Chr.values[:,:,1] ≈ [0. 1] && Chr.values[:,:,2] ≈ [im 0] && Chr.values[:,:,3] ≈ [1 0]

#@time Affm = ts2ffm(At);

nperiod = 24
time1 = (0:N-1)*2*pi*nperiod/N;
At1 = PeriodicTimeSeriesMatrix(tA.(time1),2*pi*nperiod);
Ahr1 = ts2hr(At1);
@test convert(PeriodicFunctionMatrix,Ahr1).f(1) ≈ convert(PeriodicFunctionMatrix,Ahr).f(1)

@test iszero(hr2psm(Ahr,1:1) + hr2psm(Ahr,0:0) - hr2psm(Ahr))
@test iszero(hr2psm(Ahr1,1:1) + hr2psm(Ahr1,0:0) - hr2psm(Ahr1))

# harmonic vs. symbolic
@test norm(substitute.(convert(PeriodicSymbolicMatrix,Ahr).F - A, (Dict(t => rand()),))) < 1e-15
@test norm(substitute.(convert(PeriodicSymbolicMatrix,Bhr).F - B, (Dict(t => rand()),))) < 1e-15
@test norm(substitute.(convert(PeriodicSymbolicMatrix,Chr).F - C, (Dict(t => rand()),))) < 1e-15

# harmonic vs. time series
@test all(norm.(tvmeval(Ahr,tg).-At.values) .< 1.e-7)
@test all(norm.(tvmeval(Bhr,tg).-Bt.values) .< 1.e-7)
@test all(norm.(tvmeval(Chr,tg).-Ct.values) .< 1.e-7)

# check time values on the grid
for method in ("constant", "linear", "quadratic", "cubic")
      @test all(norm.(tvmeval(At,tg; method).-At.values) .< 1.e-7)
      @test all(norm.(tvmeval(Bt,tg; method).-Bt.values) .< 1.e-7)
      @test all(norm.(tvmeval(Ct,tg; method).-Ct.values) .< 1.e-7)
end
# check interpolated values: time series vs. harmonic
tt = rand(10)*2pi;
for method in ("linear", "quadratic", "cubic")
    @test all(norm.(tvmeval(At,tt; method).-tvmeval(Ahr,tt; exact = true)) .< 1.e-3)
    @test all(norm.(tvmeval(Bt,tt; method).-tvmeval(Bhr,tt; exact = false)) .< 1.e-3)
    @test all(norm.(tvmeval(Ct,tt; method).-tvmeval(Chr,tt; exact = true)) .< 1.e-3)
end

# check conversion to function form
Amat = convert(PeriodicFunctionMatrix,Ahr);
@test all(norm.(tvmeval(At,tt; method = "linear").-tvmeval(Ahr,tt)) .< 1.e-3)
@test iszero(convert(PeriodicSymbolicMatrix,Amat).F-Ap.F)
@test isperiodic(Amat)

Amat =  convert(PeriodicFunctionMatrix,At);
@test all(norm.(tvmeval(At,tt; method = "linear").-Amat.f.(tt)) .< 1.e-3)
#@test iszero(convert(PeriodicSymbolicMatrix,Amat).F-Ap.F)
@test isperiodic(Amat)

Amat = PeriodicFunctionMatrix(tA,2pi);
@test all(norm.(tvmeval(At,tt; method = "linear").-tvmeval(Amat,tt)) .< 1.e-3)
@test iszero(convert(PeriodicSymbolicMatrix,Amat).F-Ap.F)
@test isperiodic(Amat)

Amat =  convert(PeriodicFunctionMatrix,Ap);
@test all(norm.(tvmeval(At,tt; method = "linear").-tvmeval(Ap,tt)) .< 1.e-3)
@test iszero(convert(PeriodicSymbolicMatrix,Amat).F-Ap.F)
@test isperiodic(Amat) && isperiodic(Ap) && size(Amat) == size(Ap)

for method in ("constant", "linear", "quadratic", "cubic")
      Amat = ts2pfm(At; method);
      @test all(norm.(At.values.-Amat.f.(tg)) .< 1.e-10)
      Bmat = ts2pfm(Bt; method);
      @test all(norm.(Bt.values.-Bmat.f.(tg)) .< 1.e-10)
      Cmat = ts2pfm(Ct; method);
      @test all(norm.(Ct.values.-Cmat.f.(tg)) .< 1.e-10)
end

# example of Colaneri
at(t) = [0 1; -10*cos(t) -24-10*sin(t)];
Afun=PeriodicFunctionMatrix(at,2pi);
ev = pseig(Afun; solver = "non-stiff", reltol = 1.e-10, abstol = 1.e-10)
# @time Φ = tvstm(Afun.f, Afun.period; reltol = 1.e-10)
cvals = log.(complex(ev))/2/pi
println("cvals = $cvals  -- No one digit accuracy!!!")
@test maximum(abs.(ev)) ≈ 1

# using ApproxFun
s = Fourier(0..2π)
At = FourierFunctionMatrix(Fun(t -> [0 1; -10*cos(t) -24-10*sin(t)],s), 2pi)
Atfun = convert(PeriodicFunctionMatrix,At)
Ahrfun = convert(PeriodicFunctionMatrix,pfm2hr(Afun))

@time cvals = psceig(Afun, 500; solver = "non-stiff", reltol = 1.e-10, abstol = 1.e-10)
@time cvals1 = psceig(Atfun, 500; solver = "non-stiff", reltol = 1.e-10, abstol = 1.e-10)
@time cvals2 = psceig(Ahrfun, 500; solver = "non-stiff", reltol = 1.e-10, abstol = 1.e-10)
@test cvals ≈ cvals1 ≈ cvals2
# Tt = Fun(t -> [12+5*sin(t) 1/2; 1 0],s)
# Tinvt=inv(Tt)
# Atilde=Tt*At.M*Tinvt+Tt'*Tinvt
# Aref = Fun(t -> [0 0; 2 -24-10*sin(t)],s)
# @test norm(Aref-Atilde) < 1.e-10

# example Floquet analysis from ApproxFun.jl
a=0.15
at1(t) = -[0 -1 0 0; (2+a*cos(2t)) 0 -1 0; 0 0 0 -1; -1 0 (2+a*cos(2t)) 0]
Afun1=PeriodicFunctionMatrix(at1,pi);
ev1 = pseig(Afun1; solver = "non-stiff", reltol = 1.e-10, abstol = 1.e-10)
cvals1 = log.(complex(ev1))/pi

a=0.15
at2(t) = -[0 -1 0 0; (2+a*cos(2t)) 0 -1 0; 0 0 0 -1; -1 0 (2+a*cos(2t)) 0]
Afun2=PeriodicFunctionMatrix(at2,2*pi;nperiod=2);
ev2 = pseig(Afun2; solver = "non-stiff", reltol = 1.e-10, abstol = 1.e-10)
cvals2 = log.(complex(ev2))/(2pi)
@test ev1.^2 ≈ ev2 && real(cvals1) ≈ real(cvals2)

# full accuracy characteristic exponents
# solver = "symplectic"
# for solver in ("non-stiff", "stiff", "linear", "symplectic", "noidea")
#     @time M = monodromy(Afun, 500; solver, reltol = 1.e-10, abstol = 1.e-10);
#     cvals = log.(complex(pseig(M)))/2/pi
#     #println("solver = $solver cvals = $cvals")
#     @test isapprox(cvals, [0; -24], atol = 1.e-7)
# end

solver = "non-stiff"
for solver in ("non-stiff", "stiff", "linear", "symplectic", "noidea")
    @time cvals = psceig(Afun, 500; solver, reltol = 1.e-10, abstol = 1.e-10)
    @test isapprox(cvals, [0; -24], atol = 1.e-7)
end

# Vinograd example: unstable periodic system with all A(t) stable
#at(t) = [3.5 6;-6 -5.5]+[-4.5 6; 6 4.5]*cos(12*t)+[6 4.5;4.5 -6]*sin(12*t); T = pi/6; # does not work
function at(t::Real)
    [-1-9*(cos(6*t))^2+12*sin(6*t)*cos(6*t) 12*(cos(6*t))^2+9*sin(6*t)*cos(6*t);
    -12*(sin(6*t))^2+9*sin(6*t)*cos(6*t) -1-9*(sin(6*t))^2-12*sin(6*t)*cos(6*t)]
end
@test eigvals(at(rand())) ≈ [-10,-1]
T = pi/3;
Afun=PeriodicFunctionMatrix(at,T);
for solver in ("non-stiff", "stiff", "linear", "symplectic", "noidea")
      @time cvals = psceig(Afun, 500; solver, reltol = 1.e-10)
      @test cvals ≈ [2; -13]
end

# #using BenchmarkTools
# solver = "non-stiff"
# at1(t) = [0 1; -10*cos(t) -24-10*sin(t)];
# Afun1=PeriodicFunctionMatrix(at1,2pi);
# solvers = ("non-stiff", "stiff", "linear", "symplectic", "noidea")
# global k = 1
# for i = 1:5
#       println("solver = $(solvers[k])")
#       @btime  monodromy(Afun1, 500; solver = solvers[k], reltol = 1.e-10, abstol = 1.e-10);
#       k +=1
#       # @btime psceig(Afun, 500; solver, reltol = 1.e-10)
# end

# at1(t) = [0 1; -10*cos(t) -24-10*sin(t)];
# Afun1=PeriodicFunctionMatrix(at1,2pi);
# println("solver = non-stiff")
# @btime  monodromy(Afun1, 500; solver = "non-stiff", reltol = 1.e-10, abstol = 1.e-10);
# println("solver = stiff")
# @btime  monodromy(Afun1, 500; solver = "stiff", reltol = 1.e-10, abstol = 1.e-10);
# println("solver = linear")
# @btime  monodromy(Afun1, 500; solver = "linear", reltol = 1.e-10, abstol = 1.e-10);
# println("solver = symplectic")
# @btime  monodromy(Afun1, 500; solver = "symplectic", reltol = 1.e-10, abstol = 1.e-10);
# println("solver = noidea")
# @btime  monodromy(Afun1, 500; solver = "noidea", reltol = 1.e-10, abstol = 1.e-10);

# at(t) = [BigFloat(0) 1; -10*cos(t) -24-10*sin(t)];
# Afun=PeriodicFunctionMatrix(at,2pi);

# f(u,p,t) = at(t)*u;
# u0 = Matrix{BigFloat}(I,2,2);
# tspan = (BigFloat(0.),2BigFloat(pi))
# prob = ODEProblem(f, u0, tspan)
# sol = solve(prob, Vern9(), rtol = 1.e-50, atol = 1.e-50, save_evrystep = false)
# cvals = log.(complex(eigvals(sol[end])))/2/pi

# at(t) = [BigFloat(0) 1; -10*cos(t) -24-10*sin(t)];
# f!(du,u,p,t) = mul!(du,at(t),u)
# u0 = Matrix{BigFloat}(I,2,2);
# tspan = (BigFloat(0.),2BigFloat(pi))
# prob1 = ODEProblem(f!, u0, tspan)
# sol1 = solve(prob1, Vern9(), rtol = 1.e-50, atol = 1.e-50, save_evrystep = false)
# cvals1 = log.(complex(eigvals(sol1[end])))/2/pi

# at(t) = [BigFloat(0) 1; -10*cos(t) -24-10*sin(t)];
# Afun=PeriodicFunctionMatrix(at,2BigFloat(pi));
# @time Φ = tvstm(Afun, Afun.period; reltol = 1.e-50, abstol = 1.e-50)  # running forever???
# cvals = log.(complex(eigvals(Φ)))/2/pi
# println("cvals = $cvals  -- No one digit accuracy")
# @test maximum(abs.(eigvals(Φ))) > 1

# @variables x
# Cx(x::Real) = [ sin(x)+cos(2*x) 1]

# Anum = PeriodicSymbolicMatrix(Cx(x),2pi)

# # using the lifted periodic system
# ap = [A3,A2,A1]; ep = [E3,E2,E1]; bp = [rand(3,0), rand(3,0), rand(3,0)];
# cp = [rand(0,3), rand(0,3), rand(0,3)]; dp = [rand(0,0), rand(0,0), rand(0,0)];
# psys = PeriodicDiscreteDescriptorStateSpace{Float64}(ap,ep,bp,cp,dp,-1)
# poles = gpole(dss(psys))

# M, N = psreduc_fast([A3,A2,A1],[E3,E2,E1],atol = 1.e-7);
# poles = pzeros(M,N)[1]

# A1 = rand(3,4); A2 = rand(4,3); A = [A1,A2];
# E1 = rand(3,3); E2 = rand(4,4); E = [E1,E2];

# M, N = psreduc_fast(A,E,atol = 1.e-7)
# #E1 = eye(Float64,3); E2 = rand(4,4); E = [E1,E2];
# B = [rand(3,2), rand(4,1)];
# C = [rand(2,4), rand(3,3)];
# D = [rand(2,2), rand(3,1)];
# psys = PeriodicDiscreteDescriptorStateSpace{Float64}(A,E,B,C,D,-1)
# sys = dss(psys)

# ev = eigvals(inv(E2)*A2*inv(E1)*A1);
# poles = gpole(sys);

# @test sort(real(poles)) ≈ sort(real(ev)) &&
#       sort(imag(poles)) ≈ sort(imag(ev))

# sys = dss(psys, kstart = 2, compacted = true, atol = 1.e-7)

# ev = eigvals(inv(E1)*A1*inv(E2)*A2);
# poles = gpole(sys);

# @test sort(real(poles)) ≈ sort(real(ev)) &&
#       sort(imag(poles)) ≈ sort(imag(ev))

# psys = PeriodicDiscreteDescriptorStateSpace{Float64}(A,I,B,C,D,-1)
# sys = dss(psys)

# ev = eigvals(A2*A1);
# poles = gpole(sys);

# @test sort(real(poles)) ≈ sort(real(ev)) &&
#       sort(imag(poles)) ≈ sort(imag(ev))

# sys = dss(psys, kstart = 2, compacted = true, atol = 1.e-7)

# ev = eigvals(A1*A2);
# poles = gpole(sys);

# @test sort(real(poles)) ≈ sort(real(ev)) &&
#       sort(imag(poles)) ≈ sort(imag(ev))

end
end


module Select

using GaPLAC
using GaPLAC.CSV
using GaPLAC.DataFrames
using GaPLAC.StatsBase
using GaPLAC.AbstractGPs

function run(args)
    @info "running 'select'"
    @info args

    !isempty(args["chains"]) && !isempty(args["formulae"]) && throw(ArgumentError("'select' can only take one of '--formulae' or '--chains', not both"))

    if !isempty(args["chains"])
        lp1df = CSV.read(args["chains"][1], DataFrame)
        lp1 = log2(harmmean([BigFloat(2) ^ x for x in lp1df[!, :lp]]))
        lp2df = CSV.read(args["chains"][2], DataFrame)
        lp2 = log2(harmmean([BigFloat(2) ^ x for x in lp2df[!, :lp]]))
        bayes = log2(BigFloat(2)^lp1 / BigFloat(2)^lp2)
    elseif !isempty(args["formulae"])
        spec1 = GaPLAC.gp_formula(args["formulae"][1])
        spec2 = GaPLAC.gp_formula(args["formulae"][2])

        @debug "GP formulae" GaPLAC.formula(spec1) GaPLAC.formula(spec2)

        eq1, vars1 = GaPLAC._apply_vars(GaPLAC.formula(spec1))
        kernels1 = GaPLAC._walk_kernel(eq1)
        length(vars1) == length(kernels1) || error("Something went wrong with equation parsing, number of variables should == number of kernels")

        eq2, vars2 = GaPLAC._apply_vars(GaPLAC.formula(spec2))
        kernels2 = GaPLAC._walk_kernel(eq2)
        length(vars2) == length(kernels2) || error("Something went wrong with equation parsing, number of variables should == number of kernels")

        gp1 = GP(eq1)
        gp2 = GP(eq2)

        @debug "GPs:" gp1 gp2

        df = CSV.read(args["data"], DataFrame)
        df = disallowmissing(df[completecases(df),:])

        y1 = df[!, GaPLAC.response(spec1)]
        x1 = Matrix(df[!, vars1])
        pr1 = AbstractGPs.FiniteGP(gp1, x1, 0.1, obsdim=1)

        y2 = df[!, GaPLAC.response(spec2)]
        x2 = Matrix(df[!, vars2])
        pr2 = AbstractGPs.FiniteGP(gp2, x2, 0.1, obsdim=1)

        lp1 = logpdf(pr1, y1)
        lp2 = logpdf(pr2, y2)
        pst1 = posterior(pr1, y1)
        pst2 = posterior(pr2, y2)

        bayes = log2(BigFloat(2)^lp1 / BigFloat(2)^lp2)
    else
        throw(ArgumentError("'select' command requires either '--chains' or '--formulae' arguments"))
    end

    @info """
        **Log2 Bayes**: $(round(Float64(bayes), digits=3))

        - **Log(pdf)** - model 1: $(round(lp1, digits=4))
        - **Log(pdf)** - model 2: $(round(lp2, digits=4))

        _Note_ - Positive values indicate more evidence for model 1
        """

end

end


using JuMP

#essa função recebe um problema relaxado e os indices das variaveis que deveriam ser binarias
#apartir disso retorna uma solução viável para o problema, no caso de min isso nos da um upperbound
function fpump(mod::JuMP.Model, binI::Vector{Int64})
  m=deepcopy(mod)
  obf = JuMP.prepAffObjective(m)
  tam=m.numCols
  n=length(binI)

  v = Variable.(m, 1:tam)

  # conferir se o b&b da solve no modelo antes, se n der precisa fazer aqui
  solve(m)
  xvia=m.colVal

  xint=roundbin(xvia,binI,tam)

  if xint == false
    for i in binI
      @constraint(m,0 <= v[i])
      @constraint(m, v[i] <= 1)
    end
    xint=roundbin(xrel,binI,tam)
  end

  d=dist(xvia,xint,binI)

  cont=0
  while d>1e-7 && cont < 1e3
    cont = cont + 1

    @objective(m, Min, sum{ ifelse(xint[i] == 0  , v[i] , 0) + ifelse(xint[i] == 1  , xint[i] - v[i], 0), i in binI})

    solve(m)
    xvia = m.colVal
    d1 = m.objVal

    if d1 ≈ d
      #se o xint for igual ao anterior a distância permanecerá a mesma então criamos uma perturbação

      xint=mudaround(xvia,xint,binI,tam,n)
      d = dist(xvia,xint,binI)

    else
      xint=roundbin(xvia,binI,tam)
      d=d1
    end
  end

  objval=obf'*xvia

  if d <= 1e-7
    return xvia, objval
  end
  return false, false

end

function dist(x,xint,binI)
  dist=0
  for i  in binI
    if xint[i] == 0
      soma = x[i] - 0
    elseif xint[i] == 1
      soma = 1 - x[i]
    end
    dist=dist+soma
  end
  return dist
end

function roundbin(x::Vector, binind::Vector{Int64}, tam::Int64)
  xint=zeros(tam)
  for i  in 1:tam
    if i in binind
      xint[i] = round(x[i])
      if xint[i] != 1 && xint[i] != 0
        return false
      end
    else
      xint[i] = x[i]
    end
  end
  return xint
end

function mudaround(x::Vector, xint::Vector, binind::Vector{Int64}, tam::Int64, n::Int64)

  ind=indChange(x,xint,binind,tam,n)

  for i  in ind
    if xint[i] == 1
      xint[i] = 0
    else
      xint[i] = 1
    end
  end
  return xint
end

#essa função nos da os indices aleatorios das variaveis com maiores scores
function indChange(x::Vector, xint::Vector, binind::Vector{Int64}, tam::Int64, n::Int64)
  vet=zeros(n)

  #preenche um vetor de scores (dif entre xvia e xint das variaveis binarias)
  j=1;
  for i in binind
    vet[j] = abs(x[i] - xint[i])
    j=j+1
  end

  #p é um vetor que define a permutação dos indices de vet que o coloca em ordem decrescente
  p=sortperm(vet,rev=true)
  #como o util para mim são os índices das variaveis originais organizamos binind dessa forma
  novob=binind[p]

  #quero mudar apenas uma certa quantidade de indices, quanto menos melhor
  quant=Int.(ceil(n/400))

  if n>400
    div=0.01
  else
    div=0.1
  end

  #quero que os ind sejam das variaveis com maiores scores (uma certa proporção dos primeiros)

  percent=Int.(ceil(div*n))

  j=randperm(percent)[1:quant]
  ind=Int.(zeros(quant))
  for k in 1:quant
    ind[k]=novob[j[k]]
  end
  return ind
end


module LSDescent

export bfgs, bfgs_Stop, bfgs_StopLS
export L_bfgs, L_bfgs_Stop, L_bfgs_StopLS
export M_bfgs, M_bfgs_Stop, M_bfgs_StopLS
export Newton_StopLS, Newton_Stop, Newton_Spectral
export CG_generic, formula_FR, formula_PR, formula_HS, formula_HZ
export CG_FR, CG_PR, CG_HS, CG_HZ

using LinearAlgebra

using NLPModels
using OneDmin

using LinearOperators
using Logging
using SolverCore   # Pour avoir les utilitaires d'affichage log_header et log_row

using Stopping

include("Newton/NewtonStopLS.jl")
include("Newton/NewtonStop.jl")
include("Newton/NewtonSolver.jl")

export Newton_G_StopLS, Newton_G_Stop
include("Newton/Newton_G_StopLS.jl")
include("Newton/Newton_G_Stop.jl")

include("Newton/HessianDense.jl")
include("Newton/HessianOp.jl")
include("Newton/HessianSparse.jl")

include("Newton/NewtonSpectralAbs.jl")

include("Newton/ldlt_symm.jl")
export ldlt_symm
include("Newton/NewtonLDLtAbs.jl")
include("Newton/NewtonCG.jl")

include("BFGS/bfgsSolver.jl")
include("BFGS/bfgsStop.jl")
include("BFGS/bfgsStopLS.jl")
include("BFGS/FormuleN2.jl")

include("BFGS/Type.jl")
include("BFGS/FormuleN2Op.jl")

include("ConjugateGradient/CG_generic.jl")
include("ConjugateGradient/formulae.jl")

include("ConjugateGradient/CG_FR.jl")
include("ConjugateGradient/CG_PR.jl")
include("ConjugateGradient/CG_HS.jl")
include("ConjugateGradient/CG_HZ.jl")

end # module


using Test, OrdinaryDiffEq, HyperbolicDiffEq, DiffEqCallbacks

function runtest()
    xmin = -1.0
    xmax = +1.0
    p = 1
    N = 2^6
    tspan = (0., 0.02)
    ode_solver = OrdinaryDiffEq.SSPRK33()

    mesh = UniformPeriodicMesh1D(xmin, xmax, N)
    basis = LobattoLegendre(p)
    balance_law = QuarticNonconvex()
    fvol = EnergyConservativeFlux()
    fnum = FluxPlusDissipation(EnergyConservativeFlux(), LocalLaxFriedrichsDissipation(max_abs_speed))
    semidisc = UniformPeriodicFluxDiffDisc1D(balance_law, mesh, basis, fvol, fnum)
    ode = semidiscretise(semidisc, sinpi, tspan)

    save_func = (u,t,integrator) -> integrate(u->IntegralQuantitiesBurgers(u,balance_law), u, integrator.f.f)
    saved_values = SavedValues(Float64, IntegralQuantitiesBurgers{Float64})
    saving = SavingCallback(save_func, saved_values, saveat=range(tspan..., length=10^2))
    stepsize = StepsizeLimiter((u,p,t)->max_dt(t,u,semidisc), safety_factor=0.5)

    sol = solve(ode, ode_solver, dt=0.5*max_dt(tspan[1], ode.u0, semidisc), save_everystep=false, callback=CallbackSet(stepsize,saving))

    mass = first.(saving.affect!.saved_values.saveval)
    entropy = last.(saving.affect!.saved_values.saveval)
    @test mass[1] ≈ mass[end] atol=10*eps()
    @test entropy[1] >= entropy[end]

    nothing
end

runtest()


## Load Packages
using Omega
using StatsBase
using Random
using Plots
using Distributions

Random.seed!(123)

## Define Functions
function gillespie_(rng, transitions, N, t0=0.0)

    """
    Implementation of the Gillespie algorithm for Lotka-Volterra model

    Args:
    transitions(dict): A dictionary with keys "prey" and "pred", and each element
        is a list containing the row of the transition matrix.
    N(int): Number of desired simulation iterations.
    t0(float): initial time point.
    """

    initial = Dict("prey" => prey_init(rng), "pred" => pred_init(rng))
    trajectory = Dict("prey" => [initial["prey"]], "pred" => [initial["pred"]])
    times = [t0]
    ecology = initial

    theta = Dict("spawn_prey" => spawn_prey(rng),
                 "prey2pred" => prey2pred(rng),
                 "pred_dies" => pred_dies(rng))
    t = times[1]

    for i = 1:N
        hazards = get_hazards(ecology, theta)
        # Can we replace this with a random variable?
        # It already samples one of three hazards according to some weights
        transition = transitions[sample(collect(keys(hazards)), Weights(collect(values(hazards))))]
        t = t + sum(values(hazards))

        ecology["prey"] += transition[1]
        ecology["pred"] += transition[2]
        # Enforce only positive integers
        ecology["prey"] = max(1, ecology["prey"])
        ecology["pred"] = max(1, ecology["pred"])

        append!(trajectory["prey"], ecology["prey"])
        append!(trajectory["pred"], ecology["pred"])
        append!(times, t)
    end

    return Dict(
        "prey" => trajectory["prey"],
        "pred" => trajectory["pred"],
        "times" => times
        )
end

function get_hazards(ecology, theta)
    """
    Compute the hazard function given the current states.  Note that in this
        model the function depends only on state, not time, though this is not
        the case in general. "spawn_prey" represents the event of a prey being born,
        "prey2pred" represents a predator consuming a new prey and consequently spawning
        a new predator, "pred_dies" represents the death of a predator.

    args:
        ecology(dict): a dictionary containing species counts.
        theta(dict): A dictionary where keys represent events (reactions) and
            values are the location hyperparameters for rate parameter constants
            corresponding to those events.
    """

    return Dict(
        "spawn_prey" => theta["spawn_prey"] * ecology["prey"],
        "prey2pred" => theta["prey2pred"] * ecology["prey"] * ecology["pred"],
        "pred_dies" => theta["pred_dies"] * ecology["pred"]
        )
end

# Define gillespie parameters
Pre = [[1, 0], [1, 1], [0, 1]]
Post = [[2, 0], [0, 2], [0, 0]]
transition_mat = Post - Pre
transitions = Dict("spawn_prey" => transition_mat[1,],
                    "prey2pred" => transition_mat[2,],
                    "pred_dies" => transition_mat[3,])
t0=0.0
N = 10000

# Initiate starting values randomely
prey_init = normal(100, 10)
pred_init = normal(50, 10)

# Should these be random variables? maybe not? they are in the paper
spawn_prey = normal(1, .1)
prey2pred = normal(.01, .001)
pred_dies = normal(.5, .05)

# Make gillespie into random variable
gillespie = ciid(gillespie_, transitions, 10000, 0.0)

# Sample from random variables
samples = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie),
                5, alg = RejectionSample)

plot_vals = hcat(values(samples[2][6]["prey"]), values(samples[2][6]["pred"]))
plot(1:10001, plot_vals[1:10001,1:2],
        title = "Basic Simulation",
        label = ["Prey" "Pred"],
        lw = 1.25)

## Sample with intervention
samples_int = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie),
                prey_init > 120, 5, alg = RejectionSample)

plot_vals = hcat(values(samples[2][6]["prey"]), values(samples_int[2][6]["prey"]))
plot(1:10001, plot_vals[1:10001,1:2],
        title = "Intervention Comparison",
        label = ["Prey No Intervention" "Prey > 120"],
        lw = 1.25)

plot_vals = hcat(values(samples_int[2][6]["prey"]), values(samples_int[2][6]["pred"]))
plot(1:10001, plot_vals[1:10001,1:2],
        title = "Intervention Simulation",
        label = ["Prey" "Pred"],
        lw = 1.25)

# Histogram
function intervention_hist(sim1, sim2, b)

    sim1_final_vals = zeros(0)
    sim2_final_vals = zeros(0)
    for i in 1:length(sim1)
        append!(sim1_final_vals, sim1[1][6]["prey"][i])
        append!(sim2_final_vals, sim2[1][6]["prey"][i])
    end

    final_val_diff = sim1_final_vals - sim2_final_vals
    histogram(final_val_diff, bins = b)
end

samples_100 = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie),
                100, alg = RejectionSample)
samples_int_100 = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie),
                prey_init > 120, 100, alg = RejectionSample)

intervention_hist(samples,samples_int, 20)

# Sample Rate
samples_rate_int = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie),
                    spawn_prey > 1.2, 5, alg = RejectionSample)

plot_vals = hcat(values(samples[2][6]["prey"]), values(samples_rate_int[2][6]["prey"]))
plot(1:5001, plot_vals[1:5001,1:2],
        title = "Intervention Comparison",
        label = ["Prey No Intervention" "Prey - spawn_prey_rate > 1.2"],
        lw = 1.25)

## Idea for intervention at t_now
function sample_weights(rng, hazards)
    result = sample(collect(keys(hazards)), Weights(collect(values(hazards))))
    return result
end

function initial_simulation(rng, t, transitions)

    ecology = Dict("prey" => prey_init(rng), "pred" => pred_init(rng))
    theta = Dict("spawn_prey" => spawn_prey(rng),
                 "prey2pred" => prey2pred(rng),
                 "pred_dies" => pred_dies(rng))

    hazards = get_hazards(ecology, theta)

    # Does this make sense?
    hazard_sample = ciid(sample_weights, hazards)
    transition = transitions[hazard_sample(rng)]
    t_new = t + sum(values(hazards))

    ecology_new_prey = ecology["prey"] + transition[1]
    ecology_new_pred = ecology["pred"] + transition[2]
    # Enforce only positive integers
    ecology_new_prey = max(1, ecology_new_prey)
    ecology_new_pred = max(1, ecology_new_pred)

    return Dict("prey" => ecology_new_prey,
                "pred" => ecology_new_pred,
                "t" => t_new)
end

temp = ciid(initial_simulation, t0, transitions)
samples = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, temp),
                5, alg = RejectionSample)

function get_hazards_fix(ecology, theta)
    """
    A function to fix annoying julia things
    """
    return Dict(
        "spawn_prey" => theta["spawn_prey"][1] * ecology["prey"],
        "prey2pred" => theta["prey2pred"][1] * ecology["prey"] * ecology["pred"],
        "pred_dies" => theta["pred_dies"][1] * ecology["pred"]
        )
end

function one_simulation(rng, n, theta, transitions)
    ecology = functions_list[n](rng)
    ecology_temp = copy(ecology)
    t = ecology["t"]
    delete!(ecology_temp, "t")
    hazards = get_hazards_fix(ecology_temp, theta)
    # Does this make sense?
    #hazard_sample = ciid(sample_weights, hazards)
    transition = transitions[sample(collect(keys(hazards)), Weights(collect(values(hazards))))]
    t_new = t + sum(values(hazards))

    ecology_new_prey = ecology_temp["prey"] + transition[1]
    ecology_new_pred = ecology_temp["pred"] + transition[2]
    # Enforce only positive integers
    ecology_new_prey = max(1, ecology_new_prey)
    ecology_new_pred = max(1, ecology_new_pred)

    return Dict("prey" => ecology_new_prey,
                "pred" => ecology_new_pred,
                "t" => t_new)
end

temp = ciid(initial_simulation, t0, transitions)
functions_list = Any[]
push!(functions_list, temp)
N = 5
for f in 2:N
    last = f - 1
    temp = ciid(one_simulation, last, theta, transitions)
    push!(functions_list, temp)
end

push!(functions_list,prey_init)
push!(functions_list, pred_init)
push!(functions_list, spawn_prey)
push!(functions_list, prey2pred)
push!(functions_list, pred_dies)

samples = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies,
functions_list[1], functions_list[2]),
                5, alg = RejectionSample)

samples = rand(Tuple(x for x in functions_list),
                5, alg = RejectionSample)


using DelimitedFiles

function sum_sliding_window(depths, i)
  depths[i] + depths[i+1] + depths[i+2]
end

no_of_depth_increases = 0

depths = readdlm("input.txt", '\n', Int)
  for i = 2:length(depths)-2 # to avoid index error
    if sum_sliding_window(depths, i) > sum_sliding_window(depths, i-1)
      global no_of_depth_increases += 1
    end
  end
return no_of_depth_increases

print("Total Number of depth increases is: ", no_of_depth_increases)


import DelayEmbeddings: genembed, Dataset
import Statistics: mean

export SymbolicWeightedPermutation

"""
    SymbolicWeightedPermutation(; τ = 1, m = 3, lt = Entropies.isless_rand) <: PermutationProbabilityEstimator

See docstring for [`SymbolicPermutation`](@ref).
"""
struct SymbolicWeightedPermutation{F} <: PermutationProbabilityEstimator
    τ::Int
    m::Int
    lt::F
end
function SymbolicWeightedPermutation(; τ::Int = 1, m::Int = 3, lt::F = isless_rand) where {F <: Function}
    m >= 2 || error("Need m ≥ 2, otherwise no dynamical information is encoded in the symbols.")
    SymbolicWeightedPermutation{F}(τ, m, lt)
end

function weights_from_variance(x, m::Int)
    sum((x .- mean(x)) .^ 2)/m
end

function probabilities(x::AbstractDataset{m, T}, est::SymbolicWeightedPermutation) where {m, T}
    m >= 2 || error("Need m ≥ 2, otherwise no dynamical information is encoded in the symbols.")
    πs = symbolize(x, SymbolicPermutation(m = m, lt = est.lt))  # motif length controlled by dimension of input data
    wts = weights_from_variance.(x.data, m)

    Probabilities(symprobs(πs, wts, normalize = true))
end

function probabilities(x::AbstractVector{T}, est::SymbolicWeightedPermutation) where {T<:Real}
    τs = tuple([est.τ*i for i = 0:est.m-1]...)
    emb = genembed(x, τs)
    πs = symbolize(emb, SymbolicPermutation(m = est.m, lt = est.lt)) # motif length controlled by estimator m
    wts = weights_from_variance.(emb.data, est.m)

    Probabilities(symprobs(πs, wts, normalize = true))
end


module PointProcessGeneration

using Distributions
using Distances

# Write your package code here.
include("simple.jl")
include("metric.jl")
include("support.jl")
include("colored.jl")

# simple
export poisson, normal, circle_poisson, circle_uniform
# with support
export clusters, voronoi, circles
# colored
export normal_clusters, circle_clusters, anarchy, colored_clusters, target_circles

end


module Katz
using LightGraphs
using GraphMatrices
using IterativeSolvers
export katz, stdbasis, ikatz, ρ

function katz(A, a, v)
    F = lufact(I-a*A)
    return A*(F\v)
end

function ikatz(A::KatzAdjacency, b::AbstractVector, tol::Real=1e-8, maxiter::Integer=100)
    x, convhist = cg(A, b, tol=tol, maxiter=maxiter)
    return x, convhist
end

function ρ(convhist::ConvergenceHistory)
    rs = @show convhist.residuals
    return mean([rs[i+1]/rs[i] for i in 1:length(rs)-1])
end

function stdbasis(n,i)
    x = ones(n)
    x[i] + 1
    return x
end

end # module


using Test, OrdinaryDiffEq

prob = ODEProblem(function (du,u,p,t)
                    du[1] = 10.0(u[2]-u[1])
                    du[2] = u[1]*(28.0-u[3]) - u[2]
                    du[3] = u[1]*u[2] - (8/3)*u[3]
                  end,
                  [1.0;0.0;0.0], (0.0, 100.0))

function buildW(...)::NTuple{3}
    nothing #= left preconditioner =#, nothing #= right preconditioner =#
end

nlsolve = NLNewton(buildW, )
sol = solve(prob, TRBDF2(nlsolve=nlsolve))


# DL model for the Prima Indian data

# load packages
using Pkg
Pkg.status()
Pkg.build("Knet")
using Knet
using PyPlot
using LinearAlgebra
using Random
using Printf
import StatsBase.predict

# load help function
include(pwd()*"/utilities.jl")

Random.seed!(1234) # set random numbers

x_train,y_train,x_val,y_val,x_test,y_test = load_data([1,2],0.2)

# define network
n_input = 8
n_hidden_1 = 50
n_hidden_2 = 20
n_out = 2

# predict function

function predict(w,x,p)
    output = zeros(2,size(x,2))
    for i=1:2:length(w)-2
        if i > 1
            x = dropout(x,p)
        end
        x = relu.(w[i]*x .+ w[i+1])
    end
    return softmax(w[end-1]*x .+ w[end];dims=1)
end

loss(w,x,ygold,p, lambda) = Knet.nll(predict(w,x,p), ygold; average=true) + lambda[1]*sum(w[1].^2)  + lambda[2]*sum(w[3].^2)  + lambda[3]*sum(w[5].^2)

lossgradient = grad(loss)

function train(epoch::Int,dropout_percentage::Real,lambda::Vector, w, optim)

    loss_train = zeros(epoch)
    loss_val = zeros(epoch)

    for i in 1:epoch

        idx_train = rand(1:size(x_train,2),size(x_train,2))
        idx_val = rand(1:size(x_val,2),size(x_val,2))

        grads = lossgradient(w,x_train[:,idx_train[:]],y_train[idx_train],dropout_percentage,lambda)
        update!(w, grads, optim)

        loss_train[i] = loss(w,x_train[:,idx_train[:]], y_train[idx_train], 0,lambda)
        loss_val[i] = loss(w,x_val[:,idx_val], y_val[idx_val], 0,lambda)

    end

    return loss_train, loss_val

end


function POMDPs.observation(pomdp::SingleCovidPOMDP, s::SingleCovidState, a::CovidAction, sp::SingleCovidState)
    iszero(a.testing_prop) && return CovidPOMDPs.UninformedDist() # not getting any info; any transition equally likely
    tot_mean = 0.0
    tot_variance = 0.0
    p = s.params
    for (i,inf) in enumerate(s.I)
        num_already_tested = sum(@view s.T[:,i])
        num_tested = floor(Int,(inf-num_already_tested)*a.testing_prop)
        dist = Binomial(num_tested,p.pos_test_probs[i])
        tot_mean += Statistics.mean(dist)
        tot_variance += Statistics.var(dist)
    end
    return Normal(tot_mean, sqrt(tot_variance))
end


using HungarianAlg: hungarian, display

msize = 4
matrix = rand(0:200, (msize, msize))
println(string(join([string(matrix[a, :]) for a in 1:msize], "\n")))
println()

result = hungarian(matrix)
println(display(result))
println()

println(result.match)
println(result.revenues)
println(result.row_weights)
println(result.col_weights)
println(result.revenue_sum)


#=
Work out the first ten digits of the sum of the following one-hundred 50-digit numbers.
=#

x = [37107287533902102798797998220837590246510135740250,
46376937677490009712648124896970078050417018260538,
74324986199524741059474233309513058123726617309629,
91942213363574161572522430563301811072406154908250,
23067588207539346171171980310421047513778063246676,
89261670696623633820136378418383684178734361726757,
28112879812849979408065481931592621691275889832738,
44274228917432520321923589422876796487670272189318,
47451445736001306439091167216856844588711603153276,
70386486105843025439939619828917593665686757934951,
62176457141856560629502157223196586755079324193331,
64906352462741904929101432445813822663347944758178,
92575867718337217661963751590579239728245598838407,
58203565325359399008402633568948830189458628227828,
80181199384826282014278194139940567587151170094390,
35398664372827112653829987240784473053190104293586,
86515506006295864861532075273371959191420517255829,
71693888707715466499115593487603532921714970056938,
54370070576826684624621495650076471787294438377604,
53282654108756828443191190634694037855217779295145,
36123272525000296071075082563815656710885258350721,
45876576172410976447339110607218265236877223636045,
17423706905851860660448207621209813287860733969412,
81142660418086830619328460811191061556940512689692,
51934325451728388641918047049293215058642563049483,
62467221648435076201727918039944693004732956340691,
15732444386908125794514089057706229429197107928209,
55037687525678773091862540744969844508330393682126,
18336384825330154686196124348767681297534375946515,
80386287592878490201521685554828717201219257766954,
78182833757993103614740356856449095527097864797581,
16726320100436897842553539920931837441497806860984,
48403098129077791799088218795327364475675590848030,
87086987551392711854517078544161852424320693150332,
59959406895756536782107074926966537676326235447210,
69793950679652694742597709739166693763042633987085,
41052684708299085211399427365734116182760315001271,
65378607361501080857009149939512557028198746004375,
35829035317434717326932123578154982629742552737307,
94953759765105305946966067683156574377167401875275,
88902802571733229619176668713819931811048770190271,
25267680276078003013678680992525463401061632866526,
36270218540497705585629946580636237993140746255962,
24074486908231174977792365466257246923322810917141,
91430288197103288597806669760892938638285025333403,
34413065578016127815921815005561868836468420090470,
23053081172816430487623791969842487255036638784583,
11487696932154902810424020138335124462181441773470,
63783299490636259666498587618221225225512486764533,
67720186971698544312419572409913959008952310058822,
95548255300263520781532296796249481641953868218774,
76085327132285723110424803456124867697064507995236,
37774242535411291684276865538926205024910326572967,
23701913275725675285653248258265463092207058596522,
29798860272258331913126375147341994889534765745501,
18495701454879288984856827726077713721403798879715,
38298203783031473527721580348144513491373226651381,
34829543829199918180278916522431027392251122869539,
40957953066405232632538044100059654939159879593635,
29746152185502371307642255121183693803580388584903,
41698116222072977186158236678424689157993532961922,
62467957194401269043877107275048102390895523597457,
23189706772547915061505504953922979530901129967519,
86188088225875314529584099251203829009407770775672,
11306739708304724483816533873502340845647058077308,
82959174767140363198008187129011875491310547126581,
97623331044818386269515456334926366572897563400500,
42846280183517070527831839425882145521227251250327,
55121603546981200581762165212827652751691296897789,
32238195734329339946437501907836945765883352399886,
75506164965184775180738168837861091527357929701337,
62177842752192623401942399639168044983993173312731,
32924185707147349566916674687634660915035914677504,
99518671430235219628894890102423325116913619626622,
73267460800591547471830798392868535206946944540724,
76841822524674417161514036427982273348055556214818,
97142617910342598647204516893989422179826088076852,
87783646182799346313767754307809363333018982642090,
10848802521674670883215120185883543223812876952786,
71329612474782464538636993009049310363619763878039,
62184073572399794223406235393808339651327408011116,
66627891981488087797941876876144230030984490851411,
60661826293682836764744779239180335110989069790714,
85786944089552990653640447425576083659976645795096,
66024396409905389607120198219976047599490197230297,
64913982680032973156037120041377903785566085089252,
16730939319872750275468906903707539413042652315011,
94809377245048795150954100921645863754710598436791,
78639167021187492431995700641917969777599028300699,
15368713711936614952811305876380278410754449733078,
40789923115535562561142322423255033685442488917353,
44889911501440648020369068063960672322193204149535,
41503128880339536053299340368006977710650566631954,
81234880673210146739058568557934581403627822703280,
82616570773948327592232845941706525094512325230608,
22918802058777319719839450180888072429661980811197,
77158542502016545090413245809786882778948721859617,
72107838435069186155435662884062257473692284509516,
20849603980134001723930671666823555245252804609722,
53503534226472524250874054075591789781264330331690]

@time println(string(sum(x))[1:10])


## solving by the L&S method
begin
    using JLD2, FileIO
    using LinearAlgebra
    using Interpolations
    using Roots
    using Plots
    using Optim
    pyplot()
    #CPROB=[1.0/3.0 2.0/3.0;1.0/4.0 3.0/4.0];
    #CPROB=[0.35 0.65;0.22 0.78];
    CPROB=[1.0/3.0 2.0/3.0;1.0/4.0 3.0/4.0];

    nos=size(CPROB)[1];
    [sum(CPROB[i,:]) for i=1:nos];

    mM=CPROB'-I;
    mM[end,:]=[1 for i=1:nos];
    rhs=[floor(i/nos) for i=1:nos];
    BSP=mM\rhs; #the list of steady-state probabilities

    hours = [0.0,1.0];
    wage=1.0;
    inc=wage*hours;
    ι=inc'*BSP; #payoff from the bond (= the aggregate output in the economy)

    #ρ=0.05
    ρ=-log(0.96)
    R=3.0;
end

function U(ξ::Float64,R::Float64)
    if ξ <= 0.
        -9e10
        else ((1/ξ^(R-1.0))/(1.0 - R))
    end
end

function II(aa::Int64,k::Int64,a::Int64,tbl2::Array{Int64,2})
    if aa==tbl2[k,a]
        return 1
    else
        return 0
    end
end

## the step on the grid created above

function LS_method(intr::Float64,accu_pow1::Int64,accu_pow2::Int64,nos::Int64,hours::Array{Float64,1},wage::Float64,β::Float64,grda::Array{Float64,1},Lg::Int64,CPROB::Array{Float64,2},R::Float64)
    local TBL,TBL1,TBL2,tbl,tbl1,tbl2,Λ0,Λ1
    currentU=[U((1+intr)*grda[i]+wage*hours[k]-grda[j],R) for k=1:nos, i=1:Lg, j=1:Lg];
    nextV=[U((1+intr)*grda[j]+wage*hours[k],R) for k=1:nos, j=1:Lg];

    TBL=[findmax(currentU[k,i,:]+β*(CPROB[k,:]'*nextV)') for k=1:nos, i=1:Lg];
    TBL1=first.(TBL);TBL2=last.(TBL);

    tbl=[findmax(currentU[k,i,:]+β*(CPROB[k,:]'*TBL1)') for k=1:nos, i=1:Lg];
    tbl1=first.(tbl);tbl2=last.(tbl);

    while maximum(abs.(TBL2-tbl2))>0||maximum(abs.(TBL1-tbl1))>1/(10^accu_pow1)
        TBL1=tbl1;
        TBL2=tbl2;
        tbl=[findmax(currentU[k,i,:]+β*(CPROB[k,:]'*TBL1)') for k=1:nos, i=1:Lg];
        tbl1=first.(tbl);tbl2=last.(tbl);
    end

    Λ0=[Float64(1/(nos*Lg)) for k=1:nos, i=1:Lg];

    Λ1=[sum([Λ0[k,i]*CPROB[k,kk]*II(ii,k,i,tbl2) for k=1:nos, i=1:Lg]) for kk=1:nos, ii=1:Lg];

    Λ0=Λ1;
    Λ1=[sum([Λ0[k,i]*CPROB[k,kk]*II(ii,k,i,tbl2) for k=1:nos, i=1:Lg]) for kk=1:nos, ii=1:Lg];
    while maximum(abs.(Λ1-Λ0))>1/10^accu_pow2
        Λ0=Λ1;
        Λ1=[sum([Λ0[k,i]*CPROB[k,kk]*II(ii,k,i,tbl2) for k=1:nos, i=1:Lg]) for kk=1:nos, ii=1:Lg];
    end
    return Λ1,tbl2
end


module SamplingProcedures
using DataFrames, Distributions, RDatasets
iris = dataset("datasets","iris")
## Creation of SampleObject. Special attention to SampleStrat. This has to connect with something right? Sampling -> Doc.
## Notes to self: Follow Kish notation.
struct SampleObject
	N::Int
	n::Int
	SampleStrat::String
	Sample::AbstractDataFrame
end

function svySRS(DataFrame::AbstractDataFrame, SampleSize::Int)
	## Procedure: Random sort SRSWOR, based in Algorith 4.5 from Yves Tillé book.
	## For each of the DataFrame rows, we generate a random value from a [0,1] Uniform distribution.
	## After this, we sort the DataFrame and select the desired sample size (based in SampleSize value).
	μ = Uniform()
	Rand = rand(μ,nrow(DataFrame))
	DataFrame₂= hcat(DataFrame,Rand)
	sort!(DataFrame₂,:x1,rev=true)
	DataFrame₂ = DataFrame₂[1:SampleSize,:]
	## How to put logs here?
	return SampleObject(nrow(DataFrame),nrow(DataFrame₂),"SRSWOR",DataFrame₂)
end

test=svySRS(iris,30)
typeof(test)
test.N

function svyReservoirSRS(DataFrame::AbstractDataFrame,SampleSize::Int)
	## Procedure: Reservoir method for SRSWOR
	## Select the first n units, then for the remaining values
	## with probability n/k (Bernoulli) select this unit and replace it in the selected sample. This replacement
	## is at random.
	DataFrame₂ = first(DataFrame,SampleSize)
	β = Bernoulli(SampleSize/nrow(DataFrame))
	μ = DiscreteUniform(1,SampleSize)
	DataFrameₖ = last(DataFrame,nrow(DataFrame)-SampleSize)
	for row in eachrow(DataFrameₖ)
		Randᵦ = rand(β,1)
		if Randᵦ==[1]
			Randᵤ= rand(μ,1)
			DataFrame₂[Randᵤ,:] = DataFrame(row) ## current problem: can't replace row! broadcasting issues.
		end
	end
	return SampleObject(nrow(DataFrame),nrow(DataFrame₂),"SRSWOR",DataFrame₂)
end

svyReservoirSRS(iris,20)

function svyStrataSRS(DataFrame::AbstractDataFrame,Strata::Real, SampleSize::Int)

end

end


#=
decompose a matrix C into a sum of a sparse binary matrix A and a matrix B with small
spectral norm, constraining the number of nonzeros in A

given matrix C, solve:
min specnorm(C - A) :
Aᵢⱼ ∈ {0,1}
sum(A) = k
=#

struct MatrixDecomposition <: ExampleInstance
    m::Int
    n::Int
    use_nat::Bool # use nuclear norm cone, else PSD cone formulation
end

function build(inst::MatrixDecomposition)
    (m, n) = (inst.m, inst.n)
    B_rank = 1 + div(m, 2)
    @assert 1 <= B_rank <= m <= n

    # generate data
    A_sparsity = min(1.0, 2 / n)
    A0 = sprand(Bool, m, n, A_sparsity)
    k = nnz(A0)
    B0 = randn(m, B_rank) * randn(B_rank, n)
    C = A0 + B0

    # build model
    model = JuMP.Model()
    JuMP.@variable(model, A[1:m, 1:n], Bin)
    JuMP.@constraint(model, sum(A) == k)

    JuMP.@variable(model, u)
    JuMP.@objective(model, Min, u)

    add_nonsymm_specnuc(inst.use_nat, false, u, C - A, model)

    # save for use in tests
    model.ext[:u] = u
    model.ext[:A] = A
    model.ext[:B] = C - A

    return model
end

function test_extra(inst::MatrixDecomposition, model::JuMP.Model)
    stat = JuMP.termination_status(model)
    @test stat in OPT_OR_LIMIT
    (stat == MOI.OPTIMAL) || return

    tol = eps()^0.2
    A = JuMP.value.(model.ext[:A])
    @test A ≈ round.(Int, A) atol = tol rtol = tol
    @test all(-tol .<= A .<= 1 + tol)
    # check nuclear norm value
    tol = eps()^0.2
    u = JuMP.value.(model.ext[:u])
    B = JuMP.value.(model.ext[:B])
    @test u ≈ maximum(svdvals(B)) atol = tol rtol = tol
    @test JuMP.objective_value(model) ≈ u
    return
end


@testset "RecursiveVec - singular values full" begin
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        @testset for orth in (cgs2, mgs2, cgsr, mgsr)
            A = rand(T, (n,n))
            v = rand(T, (n,))
            v2 = RecursiveVec((v, zero(v)))
            alg = Lanczos(orth = orth, krylovdim = 2*n, maxiter = 1, tol = precision(T))
            D, V, info = eigsolve(v2, n, :LR, alg) do x
                        x1, x2 = x
                        y1 = A*x2
                        y2 = A'*x1
                        return RecursiveVec(y1, y2)
            end
            @test info.converged >= n
            S = D[1:n]
            @test S ≈ svdvals(A)
            UU = hcat((sqrt(2*one(T))*v[1] for v in V[1:n])...)
            VV = hcat((sqrt(2*one(T))*v[2] for v in V[1:n])...)
            @test UU*Diagonal(S)*VV' ≈ A
        end
    end
end

@testset "RecursiveVec - singular values iteratively" begin
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        @testset for orth in (cgs2, mgs2, cgsr, mgsr)
            A = rand(T, (N,2*N))
            v = rand(T, (N,))
            w = rand(T, (2*N,))
            v2 = RecursiveVec((v, w))
            alg = Lanczos(orth = orth, krylovdim = n, maxiter = 300, tol = precision(T))
            n1 = div(n,2)
            D, V, info = eigsolve(v2, n1, :LR, alg) do x
                        x1, x2 = x
                        y1 = A*x2
                        y2 = A'*x1
                        return RecursiveVec(y1, y2)
            end
            @test info.converged >= n1
            S = D[1:n1]
            @test S ≈ svdvals(A)[1:n1]
            UU = hcat((sqrt(2*one(T))*v[1] for v in V[1:n1])...)
            VV = hcat((sqrt(2*one(T))*v[2] for v in V[1:n1])...)
            @test Diagonal(S) ≈ UU'*A*VV
        end
    end
end


using Pkg.TOML, StaticArrays, LinearAlgebra

@testset "Magnetization.jl" begin
    params = TOML.parsefile("Magnetization.toml")

    separableMagnetization = MPISimulations.initialize(MPISimulations.SeparableMagnetization,"SeparableMagnetization",params,Float64)
    r = zeros(SVector{3,Float64})
    t = randn()
    H = SVector{3,Float64}(1,0,0)
    x = 2e3*norm(H)

    @test !MPISimulations.isZero(separableMagnetization,r)
    @test MPISimulations.magnetizationField(separableMagnetization,H,r,t) ≈ 42e-2*7e-18*(coth(x) - 1/x)*normalize(H)
	r = SVector{3,Float64}(1e-3,0,0)
    @test !MPISimulations.isZero(separableMagnetization,r)
    @test MPISimulations.magnetizationField(separableMagnetization,H,r,t) ≈ 42e-2*7e-18*(coth(x) - 1/x)*normalize(H)
	r = SVector{3,Float64}(1.01e-3,0,0)
    @test MPISimulations.isZero(separableMagnetization,r)
    @test MPISimulations.magnetizationField(separableMagnetization,H,r,t) ≈ zeros(SVector{3,Float64})
end


# const Scalar = Union{Number, StaticArrays.Scalar, SVector{1}}

jacobian(f, w::Number, n) = ForwardDiff.derivative(v -> f(v, n), w)
jacobian(f, w::AbstractVector, n) = ForwardDiff.jacobian(v -> f(v, n), w)

eigvals(f, w, n) = real.(LinearAlgebra.eigen(jacobian(f, w, n)).values)
function eigen(f, w, n)
    λc, Rc = LinearAlgebra.eigen(jacobian(f, w, n))
    return real.(λc), real.(Rc)
end

##############################

struct LinearAdvectionFlux{V}
    velocity::V
end
LinearAdvectionFlux(v::AbstractVector) = LinearAdvectionFlux{SVector{length(v), eltype(v)}}(SVector(v...))

(f::LinearAdvectionFlux)(w, n) = w * (f.velocity' * n)
jacobian(f::LinearAdvectionFlux, w, n) = f.velocity' * n
eigvals(f::LinearAdvectionFlux, w, n) = f.velocity' * n

##############################

struct AdvectionFlux{D, VF}
    velocity_at_face::VF
end

AdvectionFlux{D}(vf) where D = AdvectionFlux{D, typeof(vf)}(vf)

function numerical_flux(flux::AdvectionFlux, mesh, w, scheme, i_face, dt)
    numerical_flux(LinearAdvectionFlux(flux.velocity_at_face(mesh, i_face)), mesh, w, scheme, i_face, dt)
end

function RotationFlux(center)
    function rotation_velocity(mesh, i_face)
        x = face_center(mesh, i_face) .- center
        return SVector{2, Float64}(-x[2], x[1])
    end
    return AdvectionFlux{2}(rotation_velocity)
end

##############################

struct ShallowWater{T}
    g::T
end
ShallowWater(; g=9.81) = ShallowWater{typeof(g)}(g)

function (f::ShallowWater)(v::SVector{2}, n::Union{Number, Scalar, SVector{1}})
    h, hu = v[1], v[2]
    SVector(hu * n, hu^2/h + h^2 * f.g/2) * n
end

function (f::ShallowWater)(v::SVector{3}, n::SVector{2})
    h, ux, uy = v[1], v[2]/v[1], v[3]/v[1]
    un = ux*n[1] + uy*n[2]
    return SVector(h*un, h*ux*un + h^2*n[1]*f.g/2, h*uy*un + h^2*n[2]*f.g/2)
end

function jacobian(f::ShallowWater, v::AbstractVector, n::SVector{2})
    h, ux, uy = v[1], v[2]/v[1], v[3]/v[1]
    un = ux * n[1] + uy * n[2]
    g = f.g
    @SMatrix [0                 n[1]        n[2];
              -ux*un+g*h*n[1]   un+ux*n[1]  ux*n[2];
              -uy*un+g*h*n[2]   uy*n[1]     un+uy*n[2]]
end

function eigvals(f::ShallowWater, v, n::SVector{2})
    h, ux, uy = v[1], v[2]/v[1], v[3]/v[1]
    un = ux * n[1] + uy * n[2]
    g = f.g
    c = sqrt(g*h)
    return SVector(un - c, un, un + c)
end

function eigen(f::ShallowWater, v, n::SVector{2})
    h, ux, uy = v[1], v[2]/v[1], v[3]/v[1]
    un = ux * n[1] + uy * n[2]
    ut = -ux * n[2] + uy * n[1]
    g = f.g
    c = sqrt(g*h)
    vals = SVector(un - c, un, un + c)
    vects = 0.5 .* @SMatrix [1.0                    0.0   1.0;
                             n[1]*(un-c) + n[2]*ut  n[2]  n[1]*(c+un) + n[2]*ut;
                             n[1]*ut + n[2]*(un-c)  n[1]  n[1]*ut + n[2]*(c+un)]
    return vals, vects
end


module FastLegendre

# From
# O(1) Computation of Legendre polynomials and Gauss-Legendre nodes and weights for parallel computing
# Bogaert, Ignace and Michiels, Bart and Fostier, Jan
# SIAM JOURNAL ON SCIENTIFIC COMPUTING
# 34, 2012
# http://dx.doi.org/10.1137/110855442
#
# as mentioned https://github.com/JuliaMath/SpecialFunctions.jl/issues/124
#
# Only Float64
#
# Gives a speed up over Clenshaw, which is expected, as this is O(1), not O(n):
#
# n= 1000 Clenshaw might be faster
# n= 10_000 10x faster
# n= 100_000 10x faster
# n= 1_000_000 500x faster
# n= 10_000_000 10,000x faster
# (measurement of basis(P,n)(x) compared to Basis(P,n)(x)

using FastGaussQuadrature
using SpecialFunctions

# 4 regions of evaluation
# specialized for Float64 values
function fastlegendre(n, x)
    n == 0 && return one(x)
    n == 1 && return x
    n < 100 && return direct_32(n, x)
    x < 0 && return (iseven(n) ? 1 : -1) * fastlegendre(n, -x)
    θ = acos(x)
    (n + 1) * sin(θ) > 25 && return asy_32(n, x)
    asy_33(n, x)
end

## ---- n < 100 case

# precompute first 99
const xᵢs = [FastGaussQuadrature.gausslegendre(n)[1][1:(n ÷ 2)] for n in 1:99]
const xᵢs⁻¹ = [inv.(xᵢ) for xᵢ in xᵢs]
const xᵢs2⁻¹ = [inv.(1 .- xᵢ .^ 2) for xᵢ in xᵢs]

# precompute
const Cn = zeros(Float64, 99)
Cn[1] = 1.0
for i in 1:49
    n = 2i
    Cn[n] = sqrt(pi) / gamma(i + 1) / gamma((1 - n) / 2)
    Cn[n + 1] = (n + 1) * Cn[n]
end

# use cosine or sine
function direct_32(n, x)
    if abs(x) > 0.707
        if iseven(n)
            direct_sine(Val(true), n, x)
        else
            direct_sine(Val(false), n, x)
        end
    else
        if iseven(n)
            direct_cosine(Val(true), n, x)
        else
            direct_cosine(Val(false), n, x)
        end
    end
end

#x = cos(theta) 1 - x^2 = sintheta0
function direct_sine(::Val{true}, n, x)
    tot = one(Float64)
    xs = xᵢs2⁻¹[n]
    for aᵢ in xs
        tot *= 1 - (1 - x^2) * aᵢ
    end
    tot
end

direct_sine(::Val{false}, n, x) = x * direct_sine(Val(true), n, x)

function direct_cosine(p::Val{true}, n, x)
    tot = Cn[n]
    xs = xᵢs⁻¹[n]
    for aᵢ in xs
        tot *= (1 - (x * aᵢ)^2)
    end
    tot
end
direct_cosine(p::Val{false}, n, x) = x * direct_cosine(Val(true), n, x)

## ----
const Cl0s = [gamma(n + 1) / gamma(n + 3 / 2) for n in 1:9]
τ(x) = evalpoly(
    1 / x,
    (
        1 / 1,
        0,
        -1 / 64,
        0,
        21 / 8192,
        0,
        -671 / 524288,
        0,
        180323 / 80323134,
        0,
        -20898423 / 8589934592,
        -7426362705 / 1099511627776,
    ),
)

function Cl0(l)
    1 <= l < 10 && return Cl0s[l]
    1 / sqrt(l + 3 / 4) * τ(l + 3 / 4)
end

function asy_32(n, x)
    M = 17
    θ = acos(x)
    val = sqrt((2 / pi) * (1 / sin(θ)))
    tot = 0.0
    sθ = sin(θ)
    λ = Cl0(n)
    for m in 0:(M - 1)
        α = (n + m + 1 / 2) * θ - (m + 1 / 2) * pi / 2
        tot += λ * cos(α)
        λ *= (m + 1 / 2)^2 / 2 / (m + 1) / (n + m + 3 / 2) / sθ
    end
    val * tot
end

## ---
h(n, y) = y^n * besselj(n, y)

const cns = [
    [1 / 8, -1 / 12],
    [11 / 384, -7 / 160, 1 / 160],
    [173 / 15360, -101 / 3584, 671 / 80640, -61 / 120960],
    [22931 / 3440640, -90497 / 3870720, 217 / 20480, -1261 / 967680, 1261 / 29030400],
    [
        1319183 / 247726080,
        -10918993 / 454164480,
        +1676287 / 113541120,
        -7034857 / 2554675200,
        +1501 / 8110080,
        -79 / 20275200,
    ],
    [
        233526463 / 43599790080,
        -1396004969 / 47233105920,
        +2323237523 / 101213798400,
        -72836747 / 12651724800,
        +3135577 / 5367398400,
        -1532789 / 61993451520,
        +66643 / 185980354560,
    ],
]

function fn(i, y)
    n = 2i
    cs = cns[i]
    tot = 0.0
    Δ = i - 1
    for j in i:(2i)
        tot += cs[j - Δ] * h(j, y)
    end
    tot
end

function asy_33(n, x)
    v = n + 1 / 2
    θ = acos(x)
    y = v * θ

    # Eq (3.19)
    val = besselj(0, y)

    λ = 1 / v^2
    for i in 1:6
        val += fn(i, y) * λ
        λ /= v^2
    end
    val
end

end


module FastRounding

export add_round, sub_round, mul_round, square_round,
       inv_round, div_round, sqrt_round

using ErrorfreeArithmetic

const SysFloat = Union{Float64, Float32}  # fast iff fma is available in hardware

setrounding(Float64, RoundNearest)
setrounding(Float32, RoundNearest)

@inline function add_round(a::T, b::T, rounding::R)::T where {T<:SysFloat, R<:RoundingMode}
    hi, lo = two_sum(a, b)
    return round_errorfree(hi, lo, rounding)
end
add_round(a::T, b::T) where {T<:SysFloat} = a + b

@inline function sub_round(a::T, b::T, rounding::R)::T where {T<:SysFloat, R<:RoundingMode}
    hi, lo = two_diff(a, b)
    return round_errorfree(hi, lo, rounding)
end
sub_round(a::T, b::T) where {T<:SysFloat} = a - b

@inline function mul_round(a::T, b::T, rounding::R)::T where {T<:SysFloat, R<:RoundingMode}
    hi, lo = two_prod(a, b)
    return round_errorfree(hi, lo, rounding)
end
mul_round(a::T, b::T) where {T<:SysFloat} = a * b

@inline function inv_round(a::T, rounding::R)::T where {T<:SysFloat, R<:RoundingMode}
    hi, lo = two_inv(a)
    return round_errorfree(hi, lo, rounding)
end
inv_round(a::T) where {T<:SysFloat} = inv(a)

@inline function div_round(a::T, b::T, rounding::R)::T where {T<:SysFloat, R<:RoundingMode}
    hi, lo = two_div(a, b)
    return round_errorfree(hi, lo, rounding)
end
div_round(a::T, b::T) where {T<:SysFloat} = a / b

@inline function square_round(a::T, rounding::R)::T where {T<:SysFloat, R<:RoundingMode}
    hi, lo = two_square(a)
    return round_errorfree(hi, lo, rounding)
end
square_round(a::T) where {T<:SysFloat} = a^2

@inline function sqrt_round(a::T, rounding::R)::T where {T<:SysFloat, R<:RoundingMode}
    hi, lo = two_sqrt(a)
    return round_errorfree(hi, lo, rounding)
end
sqrt_round(a::T) where {T<:SysFloat} = sqrt(a)

#=
    To perform arithmetic with directed rounding more rapidly
      we use error-free transformations to control rounding.
=#

@inline function round_errorfree(hi::T, lo::T, ::RoundingMode{:Nearest})::T where {T<:SysFloat}
     !isinf(hi) && return hi
     return signbit(hi) ? T(-Inf) : T(Inf)
end

@inline function round_errorfree(hi::T, lo::T, ::RoundingMode{:ToZero})::T where {T<:SysFloat}
     !isinf(hi) && return signbit(hi) != signbit(lo) ? next_nearerto_zero(hi) : hi
     return signbit(hi) ? nextfloat(T(-Inf)) : prevfloat(T(Inf))
end

@inline function round_errorfree(hi::T, lo::T, ::RoundingMode{:FromZero})::T where {T<:SysFloat}
    !isinf(hi) && return signbit(hi) == signbit(lo) ? next_awayfrom_zero(hi) : hi
    return signbit(hi) ? T(-Inf) : T(Inf)
end

@inline function round_errorfree(hi::T, lo::T, ::RoundingMode{:Up})::T where {T<:SysFloat}
    !isinf(hi) && return (lo<=zero(T) || isnan(lo))  ? hi : nextfloat(hi)
    return signbit(hi) ? nextfloat(T(-Inf)) : T(Inf)
end

@inline function round_errorfree(hi::T, lo::T, ::RoundingMode{:Down})::T where {T<:SysFloat}
    !isinf(hi) && return (lo>=zero(T) || isnan(lo))  ? hi : prevfloat(hi)
    return signbit(hi) ? T(-Inf) : prevfloat(T(Inf))
end

@inline next_nearerto_zero(x::T) where {T<:SysFloat} = !signbit(x) ? prevfloat(x) : nextfloat(x)
@inline next_awayfrom_zero(x::T) where {T<:SysFloat} = !signbit(x) ? nextfloat(x) : prevfloat(x)

end # module


module Saw

import GreyDecision.GreyNumbers.GreyNumber
import GreyDecision.Utility: unitize, rowsums, weightize, reverseminmax

function saw(decisionMat::Array{GreyNumber{T}, 2}, weights::Array{Float64,1}, fns::Array{Function,1};  v::Float64=0.5) where {T <: Real}
    n, p = size(decisionMat)

    normalizedDecisionMat = similar(decisionMat)

    w = unitize(weights)

    colminmax = zeros(GreyNumber{Float64}, p)

    for i in 1:p
        colminmax[i] = decisionMat[:, i] |> fns[i]
        if fns[i] == maximum
            normalizedDecisionMat[:, i] = map(x -> x / colminmax[i], decisionMat[:, i])
        elseif fns[i] == minimum
            normalizedDecisionMat[:, i] = map(x -> colminmax[i] / x, decisionMat[:, i])
        else
            @error fns[i]
            error("Function not found")
        end
    end

    scores = weightize(normalizedDecisionMat, w) |> rowsums

    rankings = scores |> sortperm |> reverse

    bestIndex = rankings |> first

    return Dict(
        :scores => scores,
        :rankings => rankings,
        :bestindex => bestIndex
    )
end

end # End of module saw


using JFVM, GR
include("rel_perms.jl")

"""
# Imbibition cell, IMPES
# Buckley Leverett equation
# dependent variables: pressure and water saturation
# Prepared for educational purposes by ** AAE **
# load peteng relperm files before running this function
## define the geometry
"""
function imb_impes()
    Nx = 20 # number of cells in x direction
    Ny = 50 # number of cells in y direction
    W = 0.02 # [m] length of the domain in x direction
    H = 0.07 # [m] length of the domain in y direction
    # m = createMesh1D(Nx, W)
    m = createMesh2D(Nx, Ny, W, H) # creates a 2D mesh
    ## define the physical parametrs
    # all the Corey-type relperm parametrs are defined for the oil-wet and water-wet
    # cases
    k0 = 0.01e-12 # [m^2] average reservoir permeability
    phi0 = 0.45 # average porosity
    mu_oil = 2e-3 # [Pa.s] oil viscosity
    mu_water = 1e-3 # [Pa.s] water viscosity
    krw0_ww = 0.3
    krw0_ow = 1.0
    kro0_ww = 0.6
    kro0_ow = 0.76
    nw_ww = 2.4
    nw_ow= 2.4
    no_ww = 2.0
    no_ow= 2.0
    sor_ww=0.1
    sor_ow=0.12
    swc_ww=0.09
    swc_ow=0.09
    teta_ow=deg2rad(130)
    teta_ww=deg2rad(20)
    gama_ow=0.03 # N/m
    labda=2.4 # for Ekofisk chalk
    r_ave=sqrt(k0/phi0) #  meter average pore diameter
    pce0=2*gama_ow*cos(teta_ww)/r_ave # Pa capillary entry pressure

    SF=createFaceVariable(m, 0.0) # 1 is water wet, 0 is oil wet
    krw0=krw0_ww*SF+krw0_ow*(1-SF)
    kro0=kro0_ww*SF+kro0_ow*(1-SF)
    sor=sor_ww*SF+sor_ow*(1-SF)
    swc=swc_ww*SF+swc_ow*(1-SF)
    no= no_ww*SF+no_ow*(1-SF)
    nw= nw_ww*SF+nw_ow*(1-SF)
    teta=teta_ww*SF+teta_ow*(1-SF)
    pce=createFaceVariable(m, pce0)

    p0 = 100e5 # [bar] pressure
    pin = 150e5 # [bar] injection pressure at the left boundary
    u_in= 1.0/(24*3600) # [m/s] equal to 1 m/day
    sw0 = swc_ww+0.01
    # sw0(10:end-10, 10:end-10)=swc+0.2
    # sw0 = swc+0.1 # initial water saturation
    sw_in = 1.0

    eps1=1e-6
    clx=0.05
    cly=0.05
    V_dp=0.01 # Dykstra-Parsons coef.
    perm_val= k0 #field2d(Nx,Ny,k0,V_dp,clx,cly)

    k=createCellVariable(m, perm_val)
    phi=createCellVariable(m, phi0)

    lw = geometricMean(k)/mu_water
    lo = geometricMean(k)/mu_oil

    ## Define the boundaries: all fixed Sw=1, fixed pressure everywhere(?)
    BCp = createBC(m) # Neumann BC for pressure
    BCs = createBC(m) # Neumann BC for saturation
    # left boundary pressure gradient
    # BCp.left.a[:]=(krw(sw_in)*lw.xvalue(1,:)+kro(sw_in)*lo.xvalue(1,:)) BCp.left.b[:]=0 BCp.left.c[:]=-u_in
    # change the right boandary to constant pressure (Dirichlet)
    # BCp.left.a[:]=0 BCp.left.b[:]=1 BCp.left.c[:]=p0
    BCp.right.a[:].=0.0
    BCp.right.b[:].=1.0
    BCp.right.c[:].=p0
    BCp.top.a[:].=0.0
    BCp.top.b[:].=1.0
    BCp.top.c[:].=p0
    BCp.bottom.a[:].=0.0
    BCp.bottom.b[:].=1.0
    BCp.bottom.c[:].=p0
    # change the left boundary to constant saturation (Dirichlet)
    # BCs.left.a[:]=0 BCs.left.b[:]=1 BCs.left.c[:]=1.0-sor
    BCs.right.a[:].=0.0
    BCs.right.b[:].=1.0
    BCs.right.c[:].=1.0
    BCs.top.a[:].=0.0
    BCs.top.b[:].=1.0
    BCs.top.c[:].=1.0
    BCs.bottom.a[:].=0.0
    BCs.bottom.b[:].=1.0
    BCs.bottom.c[:].=1.0
    ## define the time step and solver properties
    # dt = 1000 # [s] time step
    # dt=(W/Nx)/u_in/20 # [s]
    dt=1
    t_end = 10*3600*24 # [s] final time
    eps_p = 1e-7 # pressure accuracy
    eps_sw = 1e-7 # saturation accuracy
    ## define the variables
    sw_old = createCellVariable(m, sw0, BCs)
    p_old = createCellVariable(m, p0, BCp)
    sw = copyCell(sw_old)
    oil_init=domainInt(1-sw_old)
    p = copyCell(p_old)
    uw = -gradientTerm(p_old) # an estimation of the water velocity
    ## start the main loop
    # generate intial pressure profile (necessary to initialize the fully
    # implicit solver)
    rec_fact=zeros(1)
    t_day=zeros(1)
    t = 0.0
    dt0=dt
    dsw_alwd= 0.001
    dp_alwd= 100.0 # Pa
    while (t<t_end)
    # for i=1:5
        error_p = 1e5
        error_sw = 1e5
        # Implicit loop
    #     while ((error_p>eps_p) || (error_sw>eps_sw))
        while(true)
            # calculate parameters
            pgrad = gradientTerm(p)
    #         pcgrad=gradientTerm(pc(sw))
            sw_face = upwindMean(sw, -pgrad) # average value of water saturation
            sw_grad=gradientTerm(sw)
            sw_ave=arithmeticMean(sw)
            pcgrad=faceEval(dpc_imb, sw_ave, pce, swc, sor, teta)
            # solve for pressure at known Sw
            labdao = lo.*faceEval(kro, sw_face, kro0, sor, swc, no)
            labdaw = lw.*faceEval(krw, sw_face, krw0, sor, swc, nw)
            labda = labdao+labdaw
            # compute [Jacobian] matrices
            Mdiffp1 = diffusionTerm(-labda)
            RHSpc1=divergenceTerm(labdao.*pcgrad)
            Mbcp, RHSbcp = boundaryConditionTerm(BCp)
            RHS1 = RHSpc1+RHSbcp # with capillary
            p_new=solvePDE(m, Mdiffp1+Mbcp, RHS1)

            # solve for Sw
            pgrad = gradientTerm(p_new)
            uw=-labdaw.*pgrad
            Mbcsw, RHSbcsw = boundaryConditionTerm(BCs)
            RHS_sw=-divergenceTerm(uw)
            sw_new=solveExplicitPDE(sw_old, dt, RHS_sw, BCs, phi)

            error_p = maximum(abs((internalCells(p_new)-internalCells(p))./internalCells(p_new)))
            error_sw = maximum(abs(internalCells(sw_new)-internalCells(sw)))
            dt_new=dt*min(dp_alwd/error_p, dsw_alwd/error_sw)
            # print("sw_error = $error_sw \n")
            # print("new time step = $dt \n")
            # print(internalCells(sw_new))
            # print(internalCells(sw_old))
            # sleep(1.0)
            # assign new values of p and sw
            if error_sw>dsw_alwd
                dt=dt*(dsw_alwd/error_sw)
                # print("new time step = $dt \n")
            else
                t=t+dt
                p = copyCell(p_new)
                sw = copyCell(sw_new)
                p_old = copyCell(p)
                sw_old = copyCell(sw)
                dt=min(dt*(dsw_alwd/error_sw), 10*dt)
                break
            end
        end

        rec_fact=push!(rec_fact, (oil_init-domainInt(1-sw))/oil_init)
        t_day=push!(t_day, t)
        # print(t)
        # GR.imshow(sw.value[2:end-1,2:end-1])
        # visualizeCells(1-sw)
        # GR.plot(t_day/3600/24, rec_fact)
        # xlabel('time [day]')
        # ylabel('recovery factor')
        # title([num2str(t/3600/24) ' day']) drawnow
    end
    return t_day, rec_fact, sw
end # imb_impes

t, r, s = imb_impes()
GR.plot(t, r)


using Test
using NeXLCore

@testset "MatUncertain" begin
    mat = "Au60Ag40"
    inplbls = [
        MassFractionLabel(mat, n"Ag"),
        AtomicWeightLabel(mat, n"Ag"),
        MassFractionLabel(mat, n"Au"),
        AtomicWeightLabel(mat, n"Au"),
    ]
    inpvals = Float64[0.4020, 107.87, 0.5950, 196.97]
    inpcovs = Float64[(0.0090)^2 0 0 0; 0 0 0 0; 0 0 (0.0120)^2 0; 0 0 0 0]
    inputs = uvs(inplbls, inpvals, inpcovs)

    @testset "AuAg - MFtoAF" begin

        afs = propagate(MFtoAF(mat), inputs)

        @test isapprox(value(afs[AtomicFractionLabel(mat, n"Ag")]), 0.5523, atol = 0.0001)
        @test isapprox(value(afs[AtomicFractionLabel(mat, n"Au")]), 0.4477, atol = 0.0001)

        @test isapprox(σ(afs[AtomicFractionLabel(mat, n"Ag")]), 0.0075, atol = 0.0001)
        @test isapprox(σ(afs[AtomicFractionLabel(mat, n"Au")]), 0.0075, atol = 0.0001)
        @test isapprox(
            NeXLUncertainties.correlation(
                afs,
                AtomicFractionLabel(mat, n"Ag"),
                AtomicFractionLabel(mat, n"Au"),

            ),
            -1.0,
            atol = 0.0001,
        )
    end

    @testset "AuAg - MFtoNMF" begin

        afs = propagate(MFtoNMF(mat), inputs)

        @test isapprox(value(afs[NormMassFractionLabel(mat, n"Ag")]), 0.4032, atol = 0.0001)
        @test isapprox(value(afs[NormMassFractionLabel(mat, n"Au")]), 0.5968, atol = 0.0001)

        @test isapprox(σ(afs[NormMassFractionLabel(mat, n"Ag")]), 0.0073, atol = 0.0001)
        @test isapprox(σ(afs[NormMassFractionLabel(mat, n"Au")]), 0.0073, atol = 0.0001)
        @test isapprox(
            NeXLUncertainties.correlation(
                afs,
                NormMassFractionLabel(mat, n"Ag"),
                NormMassFractionLabel(mat, n"Au"),
            ),
            -1.0,
            atol = 0.0001,
        )
    end

    @testset "AuAg - MatStats" begin

        afs = propagate(MatStats(mat), inputs)

        @test isapprox(value(afs[MeanZ(mat)]), 65.8990, atol = 0.0001)
        @test isapprox(value(afs[MeanAz(mat)]), 160.56, atol = 0.01)

        @test isapprox(σ(afs[MeanZ(mat)]), 1.0381, atol = 0.0001)
        @test isapprox(σ(afs[MeanAz(mat)]), 2.5552, atol = 0.0001)
        @test isapprox(
            NeXLUncertainties.correlation(afs, MeanZ(mat), MeanAz(mat)),
            1.0,
            atol = 0.001,
        )
    end

    @testset "AuAg - Comp" begin
        afs = mf2comp(mat, inputs)

        @test isapprox(value(afs[AtomicFractionLabel(mat, n"Ag")]), 0.5523, atol = 0.0001)
        @test isapprox(value(afs[AtomicFractionLabel(mat, n"Au")]), 0.4477, atol = 0.0001)

        @test isapprox(σ(afs[AtomicFractionLabel(mat, n"Ag")]), 0.0075, atol = 0.0001)
        @test isapprox(σ(afs[AtomicFractionLabel(mat, n"Au")]), 0.0075, atol = 0.0001)
        @test isapprox(
            NeXLUncertainties.correlation(
                afs,
                AtomicFractionLabel(mat, n"Ag"),
                AtomicFractionLabel(mat, n"Au")
            ),
            -1.0,
            atol = 0.0001,
        )

        @test isapprox(value(afs[NormMassFractionLabel(mat, n"Ag")]), 0.4032, atol = 0.0001)
        @test isapprox(value(afs[NormMassFractionLabel(mat, n"Au")]), 0.5968, atol = 0.0001)

        @test isapprox(σ(afs[NormMassFractionLabel(mat, n"Ag")]), 0.0073, atol = 0.0001)
        @test isapprox(σ(afs[NormMassFractionLabel(mat, n"Au")]), 0.0073, atol = 0.0001)
        @test isapprox(
            NeXLUncertainties.correlation(
                afs,
                NormMassFractionLabel(mat, n"Ag"),
                NormMassFractionLabel(mat, n"Au")
            ),
            -1.0,
            atol = 0.0001,
        )

        @test isapprox(value(afs[MeanZ(mat)]), 65.8990, atol = 0.0001)
        @test isapprox(value(afs[MeanAz(mat)]), 160.56, atol = 0.01)

        @test isapprox(σ(afs[MeanZ(mat)]), 1.0381, atol = 0.0001)
        @test isapprox(σ(afs[MeanAz(mat)]), 2.5552, atol = 0.0001)
        @test isapprox(
            NeXLUncertainties.correlation(afs, MeanZ(mat), MeanAz(mat)),
            1.0,
            atol = 0.001,
        )
    end

    @testset "AuAg - MacU" begin

        mat, cxr = material("Au60Ag40", n"Ag" => 0.4020, n"Au" => 0.5950), n"O K-L3"
        inplbls = [
            MassFractionLabel(mat.name, n"Ag"),
            AtomicWeightLabel(mat.name, n"Ag"),
            MassFractionLabel(mat.name, n"Au"),
            AtomicWeightLabel(mat.name, n"Au"),
            μoρElementLabel(n"Ag", cxr),
            μoρElementLabel(n"Au", cxr),
        ]
        macAg, macAu = macU(n"Ag", cxr), macU(n"Au", cxr)
        @test isapprox(value(macAg), mac(n"Ag", cxr), rtol = 1.0e-6)
        @test isapprox(value(macAu), mac(n"Au", cxr), rtol = 1.0e-6)

        inpvals = Float64[0.4020, 107.87, 0.5950, 196.97, value(macAg), value(macAu)]
        inpcovs = Float64[
            (0.0090)^2 0 0 0 0 0
            0 0 0 0 0 0
            0 0 (0.0120)^2 0 0 0
            0 0 0 0 0 0
            0 0 0 0 variance(macAg) 0
            0 0 0 0 0 variance(macAu)
        ]

        inputs = uvs(inplbls, inpvals, inpcovs)
        model = μoρMaterial(mat.name, cxr)
        res = model(inputs)
        mc_res = mcpropagate(
            model,
            inputs,
            1000,
            parallel = false,
            rng = MersenneTwister(0xBADF00D),
        )
        #println(res)
        #println(mc_res)
        lbl = μoρLabel(mat.name, cxr)
        @test isapprox(value(res[lbl]), mac(mat, cxr), rtol = 1.0e-6)
        @test isapprox(value(res[lbl]), value(mc_res[lbl]), rtol = 0.01)
        @test isapprox(σ(res[lbl]), σ(mc_res[lbl]), rtol = 0.01)
    end
end


"""
    Saturation function for quadratic saturation models for machines
        Se(x) = B * (x - A)^2 / x
"""
function saturation_function(
    machine::Union{PSY.RoundRotorQuadratic, PSY.SalientPoleQuadratic},
    x::Number,
)
    Sat_A, Sat_B = PSY.get_saturation_coeffs(machine)
    return Sat_B * (x - Sat_A)^2 / x
end

"""
    Saturation function for exponential saturation models for machines
        Se(x) = B * x^A
"""
function saturation_function(
    machine::Union{PSY.RoundRotorExponential, PSY.SalientPoleExponential},
    x::Number,
)
    Sat_A, Sat_B = PSY.get_saturation_coeffs(machine)
    return Sat_B * x^Sat_A
end

function rectifier_function(I::Float64)
    if I <= 0.0
        return 1.0
    elseif I <= 0.433
        return 1.0 - 0.577 * I
    elseif I < 0.75
        return sqrt(0.75 - I^2)
    elseif I <= 1.0
        return 1.732 * (1.0 - I)
    else
        return 0.0
    end
end

function saturation_function(avr::PSY.ESAC1A, x::Number)
    Sat_A, Sat_B = PSY.get_saturation_coeffs(avr)
    return Sat_B * (x - Sat_A)^2 / x
end

function rectifier_function(I::T) where {T <: Number}
    if I <= 0.0
        return one(T)
    elseif I <= 0.433
        return 1.0 - 0.577 * I
    elseif I < 0.75
        return sqrt(0.75 - I^2)
    elseif I <= 1.0
        return 1.732 * (1.0 - I)
    else
        return zero(T)
    end
end


module DTW

function dtw(x,y; distance::Function=Distance.default)
	D = zeros(Float64, length(x)+1, length(y)+1)
	D[1,2:end] = Inf
	D[2:end,1] = Inf
	for i=1:length(x), j=1:length(y)
		D[i+1, j+1] = distance(x[i], y[j])
	end
	for i=1:length(x), j=1:length(y)
		D[i+1, j+1] += min(D[i,j], D[i,j+1], D[i+1,j])
	end

	D = D[2:end, 2:end]
	dist = D[end,end]/sum(size(D))
	dist, D, trackback(D)
end

function trackback(D)
	i,j = size(D)
	p,q = [i],[j]

	while i > 1 && j > 1
		tb = indmin([D[i-1,j-1], D[i-1,j], D[i,j-1]])
		@show(i,j,tb)
		tb in [1,2] && (i-=1)
		tb in [1,3] && (j-=1)
		unshift!(p,i)
		unshift!(q,j)
	end
	unshift!(p,1)
	unshift!(q,1)
	p,q
end

module Distance # contains distance measurement functions
default(x,y) = norm(x-y,1)
end

export dtw

end # module


using PkgSRA, Test, Random
Random.seed!(13)

include("../utils/sindy_turing_utils.jl")

# Generate test data: Van der Pol (2-d)
include("../examples/example_vanDerPol.jl")
dat = Array(solve_vdp_system())
numerical_grad = numerical_derivative(dat, ts)

# Initialize main SRA object
this_model = sra_stateful_object(ts, tspan, dat, u0, numerical_grad)
prams = this_model.parameters
prams.sindyc_ensemble_parameters[:selection_criterion] =
    sindy_cross_validate;
# Several changes are required because there are only 2 variables here
prams.variable_names = ["x", "y"];
prams.sindy_terms_list = Iterators.product(1:3, 1:3)
prams.sindy_library["cross_terms"] = [2, 3] # Also include cubic terms

# Fit; iteration not necessary here
fit_first_model(this_model, 1);
calculate_subsampled_ind(this_model);
all_models = fit_model(this_model);

@testset "Exact Nonzero terms" begin
    names = get_nonzero_term_names(all_models[1]);
    @test names == ["dxdt_x", "dydt_x"]
    names = get_nonzero_term_names(all_models[2]);
    @test names == ["dxdt_x", "dydt_x", "dxdt_y"]
    names = get_nonzero_term_names(all_models[3]);
    @test names == ["dxdt_x", "dydt_x", "dxdt_y", "dxdt_xxx"]
    # TODO: finish
end;

@testset "Correct Regression" begin
    names = get_nonzero_term_names(this_model.sindy_model);
    @test names == ["dxdt_x",
                    "dydt_x",
                    "dxdt_y",
                    "dydt_y", # Spurious term
                    "dydt_", # Spurious term
                    "dxdt_xxx"]
end;


# https://arxiv.org/abs/2012.06678

using Flux
using Flux: @functor
using Parameters: @with_kw
using FastAI
import FastAI.Models.tabular_embedding_backbone
include("../layers/transformer_block.jl")
include("../layers/categorical_embeddings.jl")

# Struct to define hyperparameters
@with_kw mutable struct TabTransfortmerArgs
    embedding_dims::Vector{Tuple{Int64, Int64}}
    mha_input_dims::Int64
    cont_input_dim::Int64
    output_dim::Int64 = 2
    lr::Float64 = 3e-4		# learning rate
    epochs::Int64 = 10        # number of epochs
    use_cuda::Bool = true   # use gpu (if cuda available)
    dropout::Float64 = 0.10 # dropout from dense layers
    hidden_dims::Vector{Int64} = [128, 64] # Size of hidden layers
    embedding_dropout::Float64 = 0.10 # dropout for categorical embeddin
    batchnorm::Bool = true # batchnorm on dense layers
    linear_first::Bool = true # linear layer before or after batch norm
    activation = Flux.relu
    output_activation = sigmoid
    mha_heads::Int64 = 8
    mha_head_dims::Int64 = 32
    transformer_dropout::Float64 = 0.1
    transformer_dense_hidden_dim::Int64 = 64
    transformer_blocks::Int64 = 2
end

"""
Create a series of dense layers following the
embedding and continious data vcat

https://github.com/FluxML/FastAI.jl/blob/master/src/models/tabularmodel.jl
"""
function dense_layers(args)
    layers = []
    first_layer = linbndrop(args.mha_input_dims + args.cont_input_dim, first(args.hidden_dims);
                    use_bn=args.batchnorm, p=args.dropout, lin_first=args.linear_first,
                    act=args.activation)
    push!(layers, first_layer)

    for (isize, osize) in zip(args.hidden_dims[1:(end-1)], args.hidden_dims[2:end])
        layer = linbndrop(isize, osize; use_bn=args.batchnorm, p=args.dropout,
                        lin_first=args.linear_first, act=args.activation)
        push!(layers, layer)
    end

    return Chain(layers...)

end

struct TabTransformer
    cat_embeddings
    cat_backbone
    cont_backbone
    dense
    output
end
@functor TabTransformer

TabTransformer(args::TabTransfortmerArgs) = TabTransformer(
    tabular_embedding_backbone(args.embedding_dims, args.embedding_dropout),
    Chain([TransformerBlock(args) for i in 1:args.transformer_blocks]...,
            Flux.flatten),
    BatchNorm(args.cont_input_dim),
    dense_layers(args),
    Dense(args.hidden_dims[end], args.output_dim, args.output_activation)
)

function (t::TabTransformer)(X_cat, X_cont)
    h_cat = t.cat_embeddings(X_cat...)
    h_cat = t.cat_backbone(reshape(h_cat, (size(h_cat, 1), 1, size(h_cat, 2))))
    h_cont = t.cont_backbone(X_cont)
    h2 = t.dense(vcat(h_cat, h_cont))
    t.output(h2)
end


function findperiod(cf::PartialQuotients)
  return 0
end

"""
sqrt(n) = [a0; a1 a2 ... a2 a1 2*a0 a1 a2...] with period "a1 a2 ... a2 a1 2*a0".
"""
function findperiod(cf::PartialQuotients{QuadraticIrrational})
  frac = cf.partialquotients
  len = length(frac)
  (len < 2) && return 0
  (frac[end] != 2 * frac[1]) && return 0
  stop = false
  period_len = (len % 2 == 0) ? 1 : 2
  i = 1
  while !stop
    if 1 + i < len - i
      if frac[1 + i] != frac[len - i]
        return 0
      end
      period_len += 2
    else
      stop = true
    end
    i += 1
  end
  return period_len
end


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

robust geometric programming problem
given a convex set C in R_+^k (described by conic constraints) and B in R^{k, n},
calculate:
f(C, B) = sup_{c in C} (inf_{x in R^n, z in R_+^k} c'*z : B_i*x <= log(z_i), i = 1..k)
note the inner problem is an unconstrained GP, and C specifies coefficient uncertainty
for more details, see section 4.4 of:
"Relative entropy optimization and its applications"
(2017) by Chandrasekaran & Shah
the authors show that:
    f(C, B) = -inf_{c in C, v in R_+^k} (d(v, e*c) : B_j'*v = 0, j = 1..n)
where e = exp(1) and d(a, b) = ∑ᵢ a_i*log(a_i/b_i) is the relative entropy of a, b
=#

struct RobustGeomProgJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    n::Int
    k::Int
end

function build(inst::RobustGeomProgJuMP{T}) where {T <: Float64}
    (n, k) = (inst.n, inst.k)
    @assert n < k # want some degrees of freedom for v
    B = randn(T, k, n) # GP powers matrix
    model = JuMP.Model()
    JuMP.@variable(model, t)
    JuMP.@objective(model, Min, t)

    JuMP.@variable(model, c[1:k])
    JuMP.@variable(model, v[1:k])
    JuMP.@constraint(model, vcat(t, ℯ * c, v) in MOI.RelativeEntropyCone(1 + 2k))
    JuMP.@constraint(model, B' * v .== 0)

    # use bounded convex set C of R_+^k excluding origin
    # note that the entropy constraint already forces c >= 0
    JuMP.@constraint(model, vcat(sqrt(k) / 2, 1 .- c) in MOI.NormOneCone(1 + k))

    return model
end


"""
    calculate_prf_response - dot product of stimImage with
    proposed pRF over many time frames

    ported from matlab code with a julian twist...

    ds 2020-10-13

"""
function calculate_prf_response(stimImUnwrapped, pRFhat, hrf=[])

# idea: consider image as a 1d array!
# "vectorise" stimImage and also the pRF
    # but avoid allocation by not creating variable...?!
    # pRFhatUnwrapped = vec(pRFhat);

    # calculate response with one matrix product
    r = stimImUnwrapped' * vec(pRFhat);

    # if the user asks for the response w/ hrf included, calculate it
    # this will save time, if not!
    if isempty(hrf)
        return r
    else
        # r_hrf = conv(r, hrf)
        r_hrf = ImageFiltering.imfilter(r, reflect(hrf)) # CONV, not CORR!
        r_hrf .-= mean(r_hrf)
        return r_hrf[1:size(stimImUnwrapped, 2)]
    end
end

"""
mgauss - 2d gaussian (x,y, with params p x0,y0, sdxy)
   inputs: p (parameter vector, [x0,y0, sdxy] )
           x,y (meshgrid format)

          NB! params first, to make optimisation easier

ds 2019-03-20 (break out of mrTools/mgl version)
ds 2020-10-13 julia version
"""
function mgauss(p, x, y)

# unpack parameters for readability
    x0 = p[1];
    y0 = p[2];
    sdxy = p[3];

# compute 2d gaussian
    m = exp.(-(((x .- x0).^2) ./ (2 * (sdxy.^2)) .+
           ((y .- y0).^2) ./ (2 * (sdxy.^2))));

# clamp small values to 0
    m[m .< 0.01] .= 0;

# return m
    m
end

"""
fmribHRF - hrf function used in by FMRIB / Oxford crew

Simple port of the Matlab code used elsewhere

    inputs: t, p [vector of params], scaled [default=1]
    outputs: result
    purpose: calculate a haemodynamic response function according
    to the equations given in the Jezzard et al book (page 252).

               H(t) = ( t / d1 ).^a1 .* exp( -(t-d1)./b1 )
                      .- c*( t / d2 ).^a2 .* exp( -(t-d2)./b2 )

   where di = ai*bi (default params a1, a2, b1, b2, c [6 12 0.9 0.9 0.35] )
"""
function fmribHRF(t=0.0:1.5:30.0, p=[6.0; 12.0; 0.9; 0.9; 0.35], scaled=true)

# if no parameters are passed use defaults from worsley chapter
# and a TR of 1.5s - which often used by Schluppeck et al.
#
# FMRIB book (p252)
# default params: a1, a2, b1, b2, c

# calculate the HRF
    t = collect(t);
    result = ( t ./ ( p[1] * p[3] )  ).^p[1] .* exp.(-( t .- p[1] * p[3] ) ./ p[3]) .-
    p[5] * ( t ./ ( p[2] * p[4] )  ).^p[2] .* exp.(-( t .- p[2] * p[4] ) ./ p[4]);

    if scaled
  # normalize to SUM=1
        result = result ./ sum(result);
    end
end