using ML_Hopf
using LinearAlgebra
using Statistics
using DiffEqFlux
using PGFPlotsX
using LaTeXStrings
using DifferentialEquations

tol = 1e-8; #Zero tolerance
N = 100; # collocation points for continuation
sp = 800; # Number of continuation points
sU = 18.0; # Starting point of continuation (wind velocity)
ds = 0.02; # Arclength (note that this remains constant in Numerical_Cont.jl)
st = 1e-3
ind = [1, 3] #index of projecting plane (Not important if we do not use CBC)
BLAS.set_num_threads(1) # This prevents stack overfolowerror
slco = ML_Hopf.continuation_flutter(N, sp, sU, ds)
U = slco.V
P = slco.P
amp = ML_Hopf.amp_LCO(U, 1, N)
θ_l = 100

## Generating training data (Data of trajectory)
vel_l = 10
Vel = range(15.0; stop=18.0, length=vel_l)
nh = 10
s_ind = [600, 700]

Vel2 = [P[s_ind[i]] for i in 1:length(s_ind)]
U₀ = 18.27 #Estimated flutter speed
s = 3.65
θ_l = 100;

dat = ML_Hopf.generate_data(vel_l, Vel, nh, st, s_ind, U, N, θ_l)
AA = dat.data
Al = dat.data2
tl2 = 1.0
tt = Int(tl2 / st + 1)
t_series = [dat.ts[i][:, 1:tt] for i in 1:length(Vel)]
θ_series = [dat.theta_s[i][1:tt] for i in 1:length(Vel)]
u₀ = dat.d0[1]
v₀ = dat.d0[2]
u0 = [u₀, v₀]
siθ = dat.siθ;
coθ = dat.coθ;
uu0 = dat.d0
## Train Linear-transformation

rot = -π * 0.1
R = [cos(rot) -sin(rot); sin(rot) cos(rot)]
θ = vec(1e-2 * R * [2.0 0.0; 0.0 3])
θ = vcat(θ, zeros(2))
scale_f_l = 1e2 # optimization works for scale_f_l>=50 for small scale_f_l optimization does not work.
function loss_lt_(θ_t)
    return ML_Hopf.loss_lt(θ_t, U₀, s, Vel, Vel2, uu0, AA, scale_f_l, nh, θ_l, coθ, siθ)
end
res1 = DiffEqFlux.sciml_train(loss_lt_, θ, ADAM(0.01); maxiters=200)
θ_ = res1.minimizer
Ap = ML_Hopf.lt_pp(θ_, U₀, s, Vel, Vel2, scale_f_l, θ_l, coθ, siθ)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[1][1, :], Ap[1][2, :])),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(t_series[1][1, :], t_series[1][2, :])),
    LegendEntry("Ground truth"),
)

## Train nonlinear transformation
hidden = 21
ann = FastChain(
    FastDense(3, hidden, tanh), FastDense(hidden, hidden, tanh), FastDense(hidden, 2)
)
θn = initial_params(ann)
scale_f = 1e3
pp = [18.27, 3.65]
θn = vcat(pp, θn)
function loss_nt_(θ_t)
    return ML_Hopf.loss_nt(
        θ_t, θ_, θ_l, scale_f_l, Vel, Vel2, scale_f, AA, ann, nh, uu0, coθ, siθ
    )
end
res_n = DiffEqFlux.sciml_train(loss_nt_, θn, ADAM(0.01); maxiters=300)
θ_2 = res_n.minimizer

Ap = ML_Hopf.lt_pp_n(θ_2, θ_, θ_l, scale_f_l, Vel, Vel2, scale_f, ann, coθ, siθ)
ind = 1
a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -3e-2,
        xmax = 3e-2,
        ymax = 6e-2,
        ymin = -9e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[ind][1, :], Ap[ind][2, :])),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(t_series[ind][1, :], t_series[ind][2, :])),
    LegendEntry("Ground truth"),
)

#Checking the phase portrait of the model (Unstable LCO)
Ap = ML_Hopf.lt_pp_n_u(θ_2, θ_, θ_l, scale_f_l, Vel, Vel2, scale_f, ann, coθ, siθ)
ind = 1 # Near the fold
vv = Vel2[ind]
uu = ML_Hopf.get_sol(U[s_ind[ind]], N, 1, 3)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -3e-2,
        xmax = 3e-2,
        ymax = 6e-2,
        ymin = -9e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[1][1, :], Ap[1][2, :])),
    LegendEntry("Learnt model"),
    Plot(
        {color = "blue", no_marks},
        Coordinates(vcat(uu.u[1, :], uu.u[1, 1]), vcat(uu.u[3, :], uu.u[3, 1])),
    ),
    LegendEntry("Ground truth"),
)

##
ind = 2
vv = Vel2[ind]# Near the equilibrium
uu = ML_Hopf.get_sol(U[s_ind[ind]], 50, 1, 3)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -3e-2,
        xmax = 3e-2,
        ymax = 6e-2,
        ymin = -9e-2,
        xtick = -1e-2:4e-3:1e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[2][1, :], Ap[2][2, :])),
    LegendEntry("Learnt model"),
    Plot(
        {color = "blue", no_marks},
        Coordinates(vcat(uu.u[1, :], uu.u[1, 1]), vcat(uu.u[3, :], uu.u[3, 1])),
    ),
    LegendEntry("Ground truth"),
)

##
u_ind = [400, 750]
Vel3 = [P[u_ind[i]] for i in 1:length(u_ind)]
Ap = ML_Hopf.lt_pp_n_u(θ_2, θ_, θ_l, scale_f_l, Vel, Vel3, scale_f, ann, coθ, siθ)

ind = 1
vv = Vel3[ind]
uu = ML_Hopf.get_sol(U[u_ind[ind]], N, 1, 3)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -3e-2,
        xmax = 3e-2,
        ymax = 6e-2,
        ymin = -9e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[1][1, :], Ap[1][2, :])),
    LegendEntry("Learnt model"),
    Plot(
        {color = "blue", no_marks},
        Coordinates(vcat(uu.u[1, :], uu.u[1, 1]), vcat(uu.u[3, :], uu.u[3, 1])),
    ),
    LegendEntry("Ground truth"),
)

##
ind = 2 # Near the equilibrium
vv = Vel3[ind]
uu = ML_Hopf.get_sol(U[u_ind[ind]], 50, 1, 3)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north east",
        height = "9cm",
        width = "9cm",
        xmin = -3e-2,
        xmax = 3e-2,
        ymax = 6e-2,
        ymin = -9e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[2][1, :], Ap[2][2, :])),
    LegendEntry("Learnt model"),
    Plot(
        {color = "blue", no_marks},
        Coordinates(vcat(uu.u[1, :], uu.u[1, 1]), vcat(uu.u[3, :], uu.u[3, 1])),
    ),
    LegendEntry("Ground truth"),
)

θ_t = [θ_; θ_2]
vel_l_ = 800
ind = 1
bd = ML_Hopf.lt_b_dia(θ_t, ind, vel_l_, ann, θ_l, coθ, siθ, scale_f_l, scale_f)
h = [maximum(t_series[i][1, :]) - minimum(t_series[i][1, :]) for i in 1:length(Vel)]
d_amp = [amp[s_ind[i]] for i in 1:length(s_ind)]
d_P = [P[s_ind[i]] for i in 1:length(s_ind)]
P = vec(P)
amp = vec(amp)
vv = vcat(bd.v, bd.v)
aa = vcat(bd.s, bd.u)

#Plot bifurcation diagram
a = @pgf Axis(
    {
        xlabel = "Air speed (m/sec)",
        ylabel = "Heave amplitude (m)",
        legend_pos = "north west",
        height = "11cm",
        width = "15cm",
        ymin = 0,
        ymax = 9e-2,
        mark_options = {scale = 1.5},
    },
    Plot({color = "blue", only_marks}, Coordinates(Vel, h)),
    LegendEntry("Training data  (stable LCO)"),
    Plot({color = "red", only_marks, mark = "triangle*"}, Coordinates(d_P, d_amp)),
    LegendEntry("Training data (unstable LCO)"),
    Plot({color = "blue", no_marks}, Coordinates(vec(P), amp)),
    LegendEntry("Ground truth"),
    Plot({color = "red", no_marks}, Coordinates(bd.v, bd.s)),
    LegendEntry("Learnt model"),
    Plot({color = "red", no_marks}, Coordinates(bd.v, bd.u)),
)

## Train speed of the phase

function Inv_T_u(th0, vel) # This function gives initial conditions of the model
    p1, p2, p3, p4, p5, p6 = θ_[1:6]
    T = [p1 p3; p2 p4]
    np1, np2 = θ_2[1:2]
    pn = θ_2[3:end]
    s_amp = sqrt(np2 / 2 + sqrt(np2^2 + 4 * (vel - np1)) / 2)
    theta = range(-π; stop=π, length=100000)

    uu = [[s_amp * cos(theta[i]), s_amp * sin(theta[i])] for i in 1:length(theta)]
    u = [
        [p5, p6] * norm(uu[i])^2 / scale_f_l +
        T * uu[i] +
        ann([uu[i]; vel - np1], pn) / scale_f for i in 1:length(theta)
    ]
    t0 = [abs(atan(u[i][2], u[i][1]) - th0) for i in 1:length(theta)]
    er = minimum(t0)
    return theta[argmin(t0)]
end

function dudt_ph(u, p, t)
    θ = u[1]
    r = u[2]
    c = u[3]
    np1, np2 = θ_2[1:2]
    a2 = np2
    δ₀ = np1
    ν = (c - δ₀)
    ω₀ = p[1]
    uu = [r * cos(θ), r * sin(θ), ν]
    du₁ = ω₀ + ann3(uu, p[2:end])[1] / om_scale
    du₂ = 0
    du₃ = 0
    return [du₁, du₂, du₃]
end

function predict_time_T(p) #,uu_t0
    np1, np2 = θ_2[1:2]
    pn = θ_2[3:end]
    p1, p2, p3, p4, p5, p6 = θ_[1:6]
    T = [p1 p3; p2 p4]
    A1 = [
        Array(
            concrete_solve(
                ODEProblem(dudt_ph, u_t0[i], (0, tl2), p),
                Tsit5(),
                u_t0[i],
                p;
                saveat=st,
                abstol=1e-8,
                reltol=1e-8,
                sensealg=InterpolatingAdjoint(; autojacvec=ReverseDiffVJP()),
            ),
        ) for i in 1:length(Vel)
    ]
    uu = [
        transpose(
            hcat(
                A1[i][2, :] .* cos.(A1[i][1, :]),
                A1[i][2, :] .* sin.(A1[i][1, :]),
                A1[i][3, :],
            ),
        ) for i in 1:vl
    ]
    delU = zeros(2, spl)
    delU2 = -np1 * ones(1, spl)
    delU = vcat(delU, delU2)
    uu = [uu[i] + delU for i in 1:vl]
    dis = transpose([p5 * ones(spl) p6 * ones(spl)]) / scale_f_l

    vlT = [
        dis * norm(uu[i][1:2])^2 +
        T * uu[i][1:2, :] +
        ML_Hopf.Array_chain(uu[i], ann, pn) / scale_f for i in 1:vl
    ]
    Pr = zeros(0, spl)
    for i in 1:vl
        theta = vlT[i][[1, 2], :]
        Pr = vcat(Pr, theta)
    end
    return Pr
end

function loss_time_T(p)
    pred = predict_time_T(p)
    return sum(abs2, A3 .- pred)
end
##
tl2 = 1.0
spl = Int(tl2 / st + 1)
vl = length(Vel)
hidden = 31
ann3 = FastChain(FastDense(3, hidden, tanh), FastDense(hidden, 1, tanh))
np = initial_params(ann3)
omega = 15.3
p = vcat(omega, np)

np1, np2 = θ_2[1:2]
pn = θ_2[3:end]
s_amp = [sqrt(np2 / 2 + sqrt(np2^2 + 4 * (Vel[i] - np1)) / 2) for i in 1:length(Vel)]
theta0 = [θ_series[i][1] for i in 1:length(Vel)]
θ₀ = [Inv_T_u(theta0[i], Vel[i]) for i in 1:length(Vel)]
u_t0 = [[θ₀[i], s_amp[i], Vel[i]] for i in 1:length(Vel)]

spl = length(t_series[1][1, :])
t_s = zeros(vl * 2, spl)
for i in 1:vl
    t_s[[2 * (i - 1) + 1, 2 * (i - 1) + 2], :] = t_series[i]
end
A3 = t_s

om_scale = 0.3
loss_time_T(p)
res_t = DiffEqFlux.sciml_train(loss_time_T, p, ADAM(0.01); maxiters=100)
res_t = DiffEqFlux.sciml_train(
    loss_time_T, res_t.minimizer, BFGS(; initial_stepnorm=1e-3); maxiters=10000
)
p = res_t.minimizer

tv = range(0, 1; length=1001)
a = @pgf Axis(
    {
        xlabel = "Time (sec)",
        ylabel = L"$h$ (m)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = 0,
        xmax = 1,
        ymax = 6e-2,
        ymin = -4e-2,
        mark_options = {scale = 0.1},
    },
    Plot({color = "red", no_marks}, Coordinates(tv, predict_time_T(p)[2 * (1 - 1) + 1, :])),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(tv, A3[2 * (1 - 1) + 1, :])),
    LegendEntry("Ground truth"),
)

pgfsave("./Figures/num_flutter/time_u15.pdf", a)

a = @pgf Axis(
    {
        xlabel = "Time (sec)",
        ylabel = L"$h$ (m)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = 0,
        xmax = 1,
        ymax = 6e-2,
        ymin = -4e-2,
    },
    Plot(
        {color = "red", no_marks}, Coordinates(tv, predict_time_T(p)[2 * (10 - 1) + 1, :])
    ),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(tv, A3[2 * (10 - 1) + 1, :])),
    LegendEntry("Ground truth"),
)
pgfsave("./Figures/num_flutter/time_u18.pdf", a)


using OrdinaryDiffEq, Flux, DiffEqFlux, Test

@testset "NeuralODE struct" begin
    x = Float32[2.; 0.]
    tspan = (0.0f0,25.0f0)
    dudt = Chain(Dense(2,50,tanh),Dense(50,2))
    slver = Tsit5()
    slver_kwargs = Dict(:save_everystep=>false, :save_start=>false)

    funcsolv = neural_ode(dudt,x,tspan,slver;slver_kwargs...)

    node = NeuralODE(dudt,tspan,slver,slver_kwargs)

    @test node(x) ≈ funcsolv
end


using MetabolicFluxAnalysis

# return values: [ lambda_min , lambda_max , x_intersection_min , x_intersection_max ]
# xrayshoot considers the system of linear inequalities A*x<b, and assumes
# that all columns of x0 lie inside the solution set.
# Then this function returns the most negative / most postitive lambdas
# such that x0_i + lambda_i * v_i <= A*b holds.
#
#
#TEST:
#A1=[eye(2);-eye(2)]; b1=[1;1;1;1];
# lambda_min , lambda_max , x_intersection_min , x_intersection_max = XRayShoot(A1,b1,zeros(2,4),[[1;1],[1;0],[-1;1],[0.1;0.2]]);
# lambda_min , lambda_max , x_intersection_min , x_intersection_max = XRayShoot(A1,b1,zeros(2,1),[[1;1],[1;0],[-1;1],[0.1;0.2]]);
function xrayshoot( A::Array{Float64,2} , b::Array{Float64,1} , x0::Array{Float64,2} , v::Array{Float64,2} )
    # if only one point, then replicate it once for every supplied
    # direction
    if ( (size(x0,2)==1) && (size(v,2)>1) )
        x0 = repmat(x0,1,size(v,2))
    end

    ## 1. compute intersections with constraints:
    AZ = A*v
    S::Array{Float64,2} = broadcast(-,b,A*x0)./AZ
    smax = copy(S)
    smax[AZ.<=0] = Inf
    smax = minimum(smax,1)
    lambda_max = smax;

    smin = S
    smin[AZ.>=0] = -Inf
    #print(smin)
    smin = maximum(smin,1)
    lambda_min = smin

    ## also return the intersecting points:
    #if(nargout >= 3),
    x_intersection_min = x0 + broadcast( * , lambda_min , v )
    x_intersection_max = x0 + broadcast( * , lambda_max , v )

    return lambda_min , lambda_max , x_intersection_min , x_intersection_max
end

# hr step for each col of x will be performed in replace
# random numbers are generated by the Julia GLOBAL_RNG
function hrstep!(A::Array{Float64,2} , b::Array{Float64,1} , x::Array{Float64,2} , nsteps::Int64)
    (d,n) = size(x)
    for zi=1:nsteps
        rv = randn(d,n)
        #rv = broadcast( / , rv , sum(rv.^2,1) )
        ru = rand(1, n)
        (lambda_min , lambda_max , x_intersection_min , x_intersection_max) = xrayshoot( A , b , x , rv )
        x[:,:] = x_intersection_min + broadcast( * , x_intersection_max-x_intersection_min , ru )
    end
    #return xi
end

function xrayshoot_centerball(r::Float64,x0::Array{Float64,2},v::Array{Float64,2})
    n=size(x0,2)
    #compute  for (x'=x+t*v)
    # (x^2)=r^2
    # t^2*v^2+t*(2xv)+x^2-r^2=0
    ea = sum( v.*v , 1); eb= 2*sum(v.*x0,1); ec= sum(x0.*x0,1)-r^2;
    eD = eb.^2 - 4*ea.*ec
    if( any( eD.<0 )); print("eD<0 -> kaputt\n"); end

    ex = [ (-eb+sqrt.(eD))./(2*ea) ; (-eb-sqrt.(eD))./(2*ea) ]
    ex_g0 = ex[1,:].>0
    lambda_max = NaN * ones(1,n)
    lambda_min = NaN * ones(1,n)
    lambda_max[1,ex_g0]  = ex[1,ex_g0]
    lambda_max[1,.~ex_g0] = ex[2,.~ex_g0]
    lambda_min[1,ex_g0]  = ex[2,ex_g0]
    lambda_min[1,.~ex_g0] = ex[1,.~ex_g0]

    x_intersection_min  = x0 + broadcast(*,v,lambda_min)
    x_intersection_max  = x0 + broadcast(*,v,lambda_max)
    return (lambda_min,lambda_max,x_intersection_min,x_intersection_max)
end

"""
  hrstep!(A::Array{Float64,2} , b::Array{Float64,1} , r::Float64 , x::Array{Float64,2} , nsteps::Int64)

hr step with uniform direction choice for each col of x will be performed inside
the intersection of the polytope G*x<h and the ball with center origin and
radius r.
NOTE: random numbers are generated by the Julia GLOBAL_RNG
"""
function hrstep!(A::Array{Float64,2} , b::Array{Float64,1} , r::Float64 , x::Array{Float64,2} , nsteps::Int64)
    #TODO implement..

    (d,n) = size(x)
    for zi=1:nsteps
        rv = randn(d,n)
        #rv = broadcast( / , rv , sum(rv.^2,1) )
        ru = rand(1, n)

        (lambda_min , lambda_max , x_intersection_min , x_intersection_max) = xrayshoot( A , b , x , rv )
        # and
        (lambda_min_cb , lambda_max_cb , x_intersection_min_cb , x_intersection_max_cb) = xrayshoot_centerball( r , x , rv )

        ix_min_cb = lambda_min_cb .> lambda_min
        x_intersection_min[ : , ix_min_cb[:] ] = x_intersection_min_cb[:,ix_min_cb[:]]
        ix_max_cb = lambda_max_cb .< lambda_max
        x_intersection_max[ : , ix_max_cb[:] ] = x_intersection_max_cb[:,ix_max_cb[:]]

        x[:,:] = x_intersection_min + broadcast( * , x_intersection_max-x_intersection_min , ru )
    end
    #return xi
end

"""
  hrstep_cchr!(A::Array{Float64,2} , b::Array{Float64,1} , r::Float64 , x::Array{Float64,2} , nsteps::Int64)

hr step with cchr direction choice for each col of x will be performed inside the
intersection of the polytope G*x<h and the ball with center origin and radius r
NOTE: random numbers are generated by the Julia GLOBAL_RNG
"""
function hrstep_cchr!(A::Array{Float64,2} , b::Array{Float64,1} , r::Float64 , x::Array{Float64,2} , nsteps::Int64)
    #TODO implement..

    d  = size(A,2)
    m  = size(A,1)

    ns = size(x,2)

    rand_dirs  = rand( 1:d , nsteps , ns)
    rand_steps = rand( nsteps , ns)

    for zi in 1:size(x,2)
        xi = x[:,zi]
        mx_b_minus_ax = b - A * xi

        for zs in 1:nsteps
            rdi::Int = rand_dirs[zs,zi]

            #sb::Float64 =  Inf
            #sa::Float64 = -Inf

            # test lin. constraints
            #ti = mx_b_minus_ax ./ A[:,rdi]
            #ti_idx_pos = ti .> 0.
            #s_a::Float64 = maximum( [ -Inf ; ti[.~ti_idx_pos] ] )
            #s_b::Float64 = minimum( [  Inf ; ti[  ti_idx_pos] ] )
            #s_a::Float64 = maximum( ti[.~ti_idx_pos] )
            #s_b::Float64 = minimum( ti[  ti_idx_pos] )
            s_a::Float64 = -Inf
            s_b::Float64 =  Inf
            for zm in 1:m
                #if( ti[zm] > 0.0 )
                ti_m = mx_b_minus_ax[zm] / A[zm,rdi]
                if( ti_m > 0.0 )
                    s_b = min(s_b,ti_m)
                else
                    s_a = max(s_a,ti_m)
                end
            end

            # test radius:
            quad_a::Float64 = 1.0
            quad_b::Float64 = 2*xi[rdi]
            quad_c::Float64 = -(r*r) + sum( xi.*xi )
            quad_D::Float64 = quad_b*quad_b - 4*quad_a*quad_c

            quad_q = -(quad_b + (quad_b>0?1.0:-1.0) * sqrt(quad_D))/ 2.0;
            b_xa::Float64 = quad_q/quad_a;
            b_xb::Float64 = quad_c/quad_q;

            if(b_xa<b_xb)
                s_a = (b_xa>s_a)?b_xa:s_a;
                s_b = (b_xb<s_b)?b_xb:s_b;

            else
                s_a = (b_xb>s_a)?b_xb:s_a;
                s_b = (b_xa<s_b)?b_xa:s_b;
            end

            step_bounds_gap::Float64 = 0.000001
            t_i = s_a +  (s_b-s_a) *  step_bounds_gap + (1.0-2*step_bounds_gap) * rand_steps[zs,zi]
            #t_i = s_a +  (s_b-s_a) * rand_steps[zs,zi]

            xi[rdi] = xi[rdi] + t_i
            mx_b_minus_ax -= A[:,rdi] * t_i
        end
        x[:,zi] = xi
    end
end

function sample_unif_hr( p::Polyhedron , x::Array{Float64,1} , nsamples::Int64 , nsteps::Int64 )
    x2 = repmat(x,1,nsamples)
    hrstep!(p.G , p.h[:] , x2 , nsteps)
    return x2
end
function sample_unif_hr( p::Polyhedron , x::Array{Float64,2} , nsteps::Int64 )
    x2 = x #repmat(x,1,nsamples)
    hrstep!(p.G , p.h[:] , x2 , nsteps)
    return x2
end


using Test
using SWIP
using Distributions

function testintervention()
  # Normally distributed random variable with id 1
  x = 1 ~ ω -> Normal(ω, 0, 1)

  # Normally distributed random variable with id 2 and x as mean
  y = 2 ~ (ω -> Normal(ω, x(ω), 1)) <| (x,)
  x_ = 0.1
  y_ = 0.3

  # An omega object -- like a trace
  ω = Ω(Dict((1,) => x_, (2,) => y_))

  # model -- tuple-valued random variable ω -> (x(ω), y(ω)), becase we want joint pdf
  m = rt(x, y)

  # Log density of model wrt ω
  l = logpdf(m, ω)

  # Check it is what it should be
  @test l == logpdf(Normal(0, 1), x_) + logpdf(Normal(x_, 1), y_)

  # Intervened model
  v_ = 100.0

  # y had x been v_
  yⁱ = y | had(x => v_)

  # new model with Intervened variables
  mⁱ = rt(x, yⁱ)

  # log pdf of interved model on same ω
  lⁱ = logpdf(mⁱ, ω)
  logpdf(x, ω)

  @test lⁱ == logpdf(Normal(0, 1), x_) + logpdf(Normal(v_, 1), y_)
  @test lⁱ < l
end


using Printf
using FiniteDiff
using ForwardDiff
using KernelAbstractions
using LinearAlgebra
using Random
using SparseArrays
using Test
using TimerOutputs
using ExaPF
import ExaPF: PowerSystem, AutoDiff

const PS = PowerSystem

@testset "Compute reduced gradient on CPU" begin
    @testset "Case $case" for case in ["case9.m", "case30.m"]
        datafile = joinpath(dirname(@__FILE__), "..", "data", case)
        tolerance = 1e-8
        pf = PS.PowerNetwork(datafile, 1)

        polar = PolarForm(pf, CPU())
        cache = ExaPF.get(polar, ExaPF.PhysicalState())

        xk = ExaPF.initial(polar, State())
        u = ExaPF.initial(polar, Control())
        p = ExaPF.initial(polar, Parameters())

        jx, ju, ∂obj = ExaPF.init_autodiff_factory(polar, cache)

        # solve power flow
        conv = powerflow(polar, jx, cache, tol=1e-12)
        ExaPF.get!(polar, State(), xk, cache)
        # No need to recompute ∇gₓ
        ∇gₓ = jx.J
        ∇gᵤ = ExaPF.jacobian(polar, ju, cache)
        # test jacobian wrt x
        ∇gᵥ = ExaPF.jacobian(polar, jx, cache)
        @test isapprox(∇gₓ, ∇gᵥ)

        # Test with Matpower's Jacobian
        V = cache.vmag .* exp.(im * cache.vang)
        Ybus = pf.Ybus
        J = ExaPF.residual_jacobian(V, Ybus, pf.pv, pf.pq)
        @test isapprox(∇gₓ, J)

        # Test gradients
        @testset "Reduced gradient" begin
            # Refresh cache with new values of vmag and vang
            ExaPF.refresh!(polar, PS.Generator(), PS.ActivePower(), cache)
            # We need uk here for the closure
            uk = copy(u)
            ExaPF.∂cost(polar, ∂obj, cache)
            ∇fₓ = ∂obj.∇fₓ
            ∇fᵤ = ∂obj.∇fᵤ

            ## ADJOINT
            # lamba calculation
            λk  = -(∇gₓ') \ ∇fₓ
            grad_adjoint = ∇fᵤ + ∇gᵤ' * λk
            # ## DIRECT
            S = - inv(Array(∇gₓ)) * ∇gᵤ
            grad_direct = ∇fᵤ + S' * ∇fₓ
            @test isapprox(grad_adjoint, grad_direct)

            # Compare with finite difference
            function reduced_cost(u_)
                # Ensure we remain in the manifold
                ExaPF.transfer!(polar, cache, xk, u_, p)
                convergence = powerflow(polar, jx, cache, tol=1e-14)
                ExaPF.refresh!(polar, PS.Generator(), PS.ActivePower(), cache)
                return ExaPF.cost_production(polar, cache.pg)
            end

            grad_fd = FiniteDiff.finite_difference_gradient(reduced_cost, uk)
            @test isapprox(grad_fd, grad_adjoint, rtol=1e-4)
        end
        @testset "Reduced Jacobian" begin
            for cons in [ExaPF.state_constraint, ExaPF.power_constraints]
                m = ExaPF.size_constraint(polar, cons)
                for icons in 1:m
                    ExaPF.jacobian(polar, cons, icons, ∂obj, cache)
                end
            end
        end
    end
end


"""
Contains structures, functions, and problem definitions
for core astrodynamics problems, including the
Restricted Two-body Problem, the Circular Restricted
Three-body Problem, and the N-body problem.
"""
module AstrodynamicsCore

# Common stuctures and functions
export AbstractUnitfulStructure, AbstractState, AbstractFrame, AbstractSystem, AbstractOrbit, CartesianState
export MassParameter, lengthunit, timeunit, velocityunit, massparameterunit, coordinateframe
export position_vector, velocity_vector, scalar_position, scalar_velocity
export AbstractFrame, Inertial, Synodic, Perifocal, ECI, HCI, convert, epoch
export NormalizedLengthUnit, NormalizedTimeUnit, Trajectory, Manifold
export convert, show, eltype, isapprox, isequal, string

# Core data structures and functions for R2BP calculations
export AbstractConic, Circular, Elliptical, Parabolic, Hyperbolic, Invalid
export KeplerianState, RestrictedTwoBodySystem, RestrictedTwoBodyOrbit
export R2BPState, R2BPSystem, R2BPOrbit, Orbit, CartesianOrbit, KeplerianOrbit
export mass_parameter, C3, SOI, SOA, hohmann
export eccentricity, semimajor_axis, inclination, RAAN, argument_of_periapsis, true_anomoly
export conic, keplerian, cartesian, perifocal, semi_parameter, periapsis_radius
export specific_angular_momentum_vector, specific_angular_momentum, eccentricity_vector
export specific_energy, specific_potential_energy, mean_motion, mean_motion_vector
export eccentric_anomoly, time_since_periapsis, period
export kepler, lambert, lambert_universal, lambert_lancaster_blanchard
export Sun, Mercury, Venus, Earth, Moon, Luna, Mars, Jupiter, Saturn, Uranus, Neptune, Pluto

# Core data structures and functions for CR3BP calculations
export CR3BPFrames, CR3BPSystem, CR3BPOrbit, CR3BPState
export SynodicCR3BPOrbit, NormalizedSynodicCR3BPOrbit, NormalizedSynodicSTMCR3BPOrbit
export SynodicCartesianState, InertialCartesianState, SynodicCartesianSTMState
export NormalizedCartesianState, MinimalCircularRestrictedThreeBodySystem
export CircularRestrictedThreeBodySystem, CircularRestrictedThreeBodyOrbit
export normalized_length_unit, normalized_time_unit, normalized_mass_parameter
export mass_parameters, primary_mass_parameter, secondary_mass_parameter
export primary_synodic_position, secondary_synodic_position, distance_to_primary, distance_to_secondary
export time_scale_factor, nondimensionalize, redimensionalize, nondimensional_radius
export nondimensionalize_length, nondimensionalize_time, nondimensionalize_velocity
export redimensionalize_length, redimensionalize_time, redimensionalize_velocity
export normalize, lagrange, inertial, synodic, accel, accel!, analyticalhalo
export potential_energy, jacobi_constant, zerovelocity_curves
export closest_approach, optimal_approach, body_position, position_wrt_body
export transform, transform_to_primary, transform_to_secondary
export SunVenus, SunEarth, EarthMoon, SunMars, SunJupiter, SunSaturn

# Module Dependencies
using Contour
using StaticArrays
using LinearAlgebra
using Roots: find_zero
using Unitful, UnitfulAstro, UnitfulAngles

# Source Code
include("Common/CommonTypes.jl")

include("R2BP/R2BPStates.jl")
include("R2BP/R2BPCalculations.jl")
include("R2BP/Kepler.jl")
include("R2BP/Lambert.jl")
include("R2BP/R2BPSystems.jl")

include("CR3BP/CR3BPStates.jl")
include("CR3BP/CR3BPCalculations.jl")
include("CR3BP/CR3BPSystems.jl")
include("CR3BP/CR3BPTransfers.jl")

end


using LinearAlgebra
@testset "rational" begin
    A = [-3 1; 5 6.6]

    # Simple calls with coefficients
    den_coeff = [3; -1; 2.0; 0.1; 1.2; 0.9; -0.3; 1.6]
    num_coeff = [1.4, 2.1, -4, -0.1, 0.2, -1.2, 2.1, 0.4, 1.4]
    den =
        den_coeff[1] * I +
        den_coeff[2] * A +
        den_coeff[3] * A^2 +
        den_coeff[4] * A^3 +
        den_coeff[5] * A^4 +
        den_coeff[6] * A^5 +
        den_coeff[7] * A^6 +
        den_coeff[8] * A^7
    num =
        num_coeff[1] * I +
        num_coeff[2] * A +
        num_coeff[3] * A^2 +
        num_coeff[4] * A^3 +
        num_coeff[5] * A^4 +
        num_coeff[6] * A^5 +
        num_coeff[7] * A^6 +
        num_coeff[8] * A^7 +
        num_coeff[9] * A^8
    P = num \ den

    (graph, cref) = graph_rational(den_coeff, num_coeff, graph_monomial)
    @test eval_graph(graph, A) ≈ P

    (graph, cref) = graph_rational(den_coeff, num_coeff, graph_ps)
    @test eval_graph(graph, A) ≈ P

    # Generate more complicated graphs before
    x1a = [1.2; 4.4]
    x1b = [0.3; -0.3]
    x2a = [2.1; 0.1; -0.1]
    x2b = [0.3; 0.2; 0.1]
    x = [(x1a, x1b), (x2a, x2b)]
    z = [1.0; 0.1; -0.1; 0.01]
    den_coeffs = [(x1a, x1b), (x2a, x2b)]
    β1 = (x1a[1] * I + x1a[2] * A) * (x1b[1] * I + x1b[2] * A)
    β2 =
        (x2a[1] * I + x2a[2] * A + x2a[3] * β1) *
        (x2b[1] * I + x2b[2] * A + x2b[3] * β1)
    den = z[1] * I + z[2] * A + z[3] * β1 + z[4] * β2
    (den_graph, den_cref) = graph_degopt(x, z)

    x1a = [2.1; 1.2]
    x1b = [1.3; 0.3]
    x2a = [0.1; -1; 1]
    x2b = [0; 2; -0.3]
    x = [(x1a, x1b), (x2a, x2b)]
    z = [1.0; -1; 0.1; -0.01]
    num_coeffs = [(x1a, x1b), (x2a, x2b)]
    β1 = (x1a[1] * I + x1a[2] * A) * (x1b[1] * I + x1b[2] * A)
    β2 =
        (x2a[1] * I + x2a[2] * A + x2a[3] * β1) *
        (x2b[1] * I + x2b[2] * A + x2b[3] * β1)
    num = z[1] * I + z[2] * A + z[3] * β1 + z[4] * β2
    (num_graph, num_cref) = graph_degopt(x, z)

    P = num \ den
    (graph, cref) = graph_rational(
        den_graph,
        num_graph,
        den_cref = den_cref,
        num_cref = num_cref,
    )
    @test eval_graph(graph, A) ≈ P
end


import Flux
import LinearAlgebra

include("TensorOps.jl")

struct MultiHeadAttention{T}
    WQueries::AbstractArray{T, 3}
    WKeys::AbstractArray{T, 3}
    WValues::AbstractArray{T, 3}
    HeadDense::Flux.TrackedArray
    L::Integer
end

function MultiHeadAttention(dims::Integer, layers::Integer)
    OutDims = Int32(floor(dims / layers))
    MultiHeadAttention(
        cat([randn(Float32, dims, OutDims) for _ in 1:layers]..., dims=3) |> Flux.gpu |> Flux.param,
        cat([randn(Float32, dims, OutDims) for _ in 1:layers]..., dims=3) |> Flux.gpu |> Flux.param,
        cat([randn(Float32, dims, OutDims) for _ in 1:layers]..., dims=3) |> Flux.gpu |> Flux.param,
        randn(Float32, dims, dims) |> Flux.gpu |> Flux.param,
        layers
    )
end
Flux.@treelike MultiHeadAttention

function ScalarDotAttention(Query::AbstractArray{T, 3}, Key::AbstractArray{T, 3}, Value::AbstractArray{T, 3}; futuremask = false) where T
    Numerator = BatchMatMul(Query, permutedims(Key, [2, 1, 3]))
    Denominator = Float32(sqrt(size(Key, 1)))

    if futuremask
        Mask = ones(Float32, size(Numerator)[1:2])

        # lower triangular matrix with very large negative numbers
        Mask = LinearAlgebra.tril!(trues(size(Mask)), -1) * Float32(-1e9)
        Mask = Flux.param(Mask) |> Flux.gpu
        Numerator = Numerator .* Mask
    end

    x = TensorSoftmax(Numerator / Denominator, dims=2)
    return BatchMatMul(x, Value)
end

function _MHALayer(MHA::MultiHeadAttention, MHAData::AbstractArray{T, 1} where T, L::Integer; futuremask = false)
    (Query, Key, Value) = MHAData
    Nbatches = size(Query, 3)

    HQueries = TensorDot(Query, MHA.WQueries[:, :, L])
    HKeys = TensorDot(Key, MHA.WKeys[:, :, L])
    HValues = TensorDot(Value, MHA.WValues[:, :, L])

    return ScalarDotAttention(HQueries, HKeys, HValues, futuremask=futuremask)
end

function (m::MultiHeadAttention)(Query::AbstractArray{T, 3}, Key::AbstractArray{T, 3}, Value::AbstractArray{T, 3}; layers = 1, futuremask = false) where T
    TotalHeads = cat([_MHALayer(m, [Query, Key, Value], layer, futuremask=futuremask) for layer in 1:m.L]..., dims=2)
    return TensorDot(TotalHeads, m.HeadDense)
end


module SfRelativity

using PhysicalConstants.CODATA2018: c_0, g_n, G
using Unitful
using UnitfulAstro

using ..SfUnits

import ..SfPhysics: kinetic_energy

export lorentz_factor, lorentz_velocity, relativistic_kinetic_energy, relativistic_brachistochrone_transit_time,
	proper_relativistic_brachistochrone_transit_time, relativistic_velocity, relativistic_delta_v

"""
	lorentz_factor(v::Unitful.Velocity)

Compute the Lorentz factor associated with the given velocity.

Note that for small values of `v` floating point precision issues can produce substantial
inaccuracies unless velocities are specified as `BigInt` or `BigFloat`.
"""
lorentz_factor(v::Unitful.Velocity) = 1/sqrt(1-(v/c_0)^2) |> u"m/m"
lorentz_factor(t::Unitful.Time, acc::Unitful.Acceleration) = sqrt(1 + (acc*t/c_0)^2) |> u"m/m"
lorentz_factor(d::Unitful.Length, acc::Unitful.Acceleration) = 1 + (acc * d)/c_0^2 |> u"m/m"

lorentz_velocity(γ) = sqrt(1 - (1/γ)^2) * c_0 |> u"c"

"""
	relativistic_kinetic_energy(m::Unitful.Mass, v::Unitful.Velocity)

Compute the kinetic energy of a body with mass `m` travelling at velocity `v`.

Relativistic corrections are included. Note that for small values of `v` floating point precision issues
can produce substantial inaccuracies unless velocities are specified as `BigInt` or `BigFloat`.
"""
relativistic_kinetic_energy(m::SfUnits.Mass, v::Unitful.Velocity) = (m * big(c_0)^2)* (lorentz_factor(v) - 1) |> u"J"

"""
	relativistic_velocity(m::SfUnits.Mass, ke::Unitful.Energy)

Velocity of a particle of mass `m` with kinetic energy `ke`, accounting for relativity.
"""
function relativistic_velocity(m::SfUnits.Mass, ke::Unitful.Energy)
	x = 1 + ke / (m * c_0^2)
	return sqrt(c_0^2 - c_0^2 / x^2)
end

function relativistic_brachistochrone_transit_time(dist::Unitful.Length, acc::Unitful.Acceleration)
	hd = dist/2

	2 * sqrt((hd/c_0)^2 + (2 * hd / acc))
end

function proper_relativistic_brachistochrone_transit_time(d::Unitful.Length, acc::Unitful.Acceleration)
	hd = d/2

	2 * (c_0 / acc) * acosh(((acc * hd) / c_0^2) + 1)
end

"""
    relativistic_velocity(t::Unitful.Time, acc::Unitful.Acceleration)

Velocity attained by a body undergoing constant acceleration `acc` for time `t`, accounting for relativistic effects.
"""
relativistic_velocity(t::Unitful.Time, acc::Unitful.Acceleration) = (acc*t)/sqrt(1+((acc*t)/c_0)^2|>u"s/s") |>u"m/s"

"""
    relativistic_delta_v(ve::Unitful.Velocity, mr)

Delta-V of a rocket with exhaust velocity `ve` and mass ration `mr` accounting for relativistic effects.
"""
relativistic_delta_v(ve::Unitful.Velocity, mr) = c_0 * tanh((ve/c_0) * log(mr)) |>u"m/s"

end


using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D
using Basis2DQuad
using UniformQuadMesh

using SetupDG

push!(LOAD_PATH, "./examples/EntropyStableEuler.jl/src")
include("../EntropyStableEuler.jl/src/logmean.jl")
#using EntropyStableEuler
#using EntropyStableEuler.Fluxes1D

function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,rhou,E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end

function pfun_nd(rho, rhou, rhov, E)
    rhoUnorm2 = (rhou^2+rhov^2)/rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function pfun_nd(rho, rhou, E)
    rhoUnorm2 = (rhou^2)/rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function primitive_to_conservative_hardcode(rho,u,v,p)
    Unorm = u^2+v^2
    E = @. p/(γ-1) + .5*rho*Unorm
    return (rho,rho*u,rho*v,E)
end

function euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
    rhoR,uR,vR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))
    vavg   = (@. .5*(vL+vR))

    unorm = (@. uL*uR + vL*vR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. FxS1*vavg)
    FxS4 = (@. f4aux*uavg)

    FyS1 = (@. rholog*vavg)
    FyS2 = (@. FxS3)
    FyS3 = (@. FyS1*vavg + pa)
    FyS4 = (@. f4aux*vavg)
    return (FxS1,FxS2,FxS3,FxS4),(FyS1,FyS2,FyS3,FyS4)
end

function euler_fluxes(rhoL,uL,vL,betaL,rhoR,uR,vR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
                           rhoR,uR,vR,betaR,rhologR,betalogR)
end

function is_zero(a)
    return abs(a) < TOL
end

function Riemann_2D(x,y)
    if x >= 0 && y >= 0
        return .5313, 0.0, 0.0, 0.4
    elseif x < 0 && y >= 0
        return 1.0, .7276, 0.0, 1.0
    elseif x < 0 && y < 0
        return .8, 0.0, 0.0, 1.0
    else
        return 1.0, 0.0, .7276, 1.0
    end

    # if x >= 0 && y >= 0
    #     return .5313, 0.0, 0.0, 0.004#0.4
    # elseif x < 0 && y >= 0
    #     return 1.0, .7276, 0.0, 1.0
    # elseif x < 0 && y < 0
    #     return .8, 0.0, 0.0, 1.0
    # else
    #     return 1.0, 0.0, .7276, 1.0
    # end

    # if x >= 0 && y >= 0
    #     return 1.5, 0.0, 0.0, 1.5
    # elseif x < 0 && y >= 0
    #     return 0.5323, 1.206, 0.0, 0.3
    # elseif x < 0 && y < 0
    #     return .138, 1.206, 1.206, 0.029
    # else
    #     return 0.5323, 0.0, 1.206, 0.3
    # end
end

function vortex(x,y,t,γ=1.4)
    x0 = 0
    y0 = 0
    beta = 5
    r2 = @. (x-x0-t)^2 + (y-y0)^2

    u = @. 1 - beta*exp(1-r2)*(y-y0)/(2*pi)
    v = @. beta*exp(1-r2)*(x-x0-t)/(2*pi)
    rho = @. 1 - (1/(8*γ*pi^2))*(γ-1)/2*(beta*exp(1-r2))^2
    rho = @. rho^(1/(γ-1))
    p = @. rho^γ

    return (rho, u, v, p)
end

const γ = 1.4
const TOL = 1e-14
const Nc = 4 # number of components

"Approximation parameters"
# N = 3 # The order of approximation
# K1D = 20
N = 2
K1D = 32
# N = 3
# K1D = 64
T = 0.25
#T = 0.25

"Mesh related variables"
Kx = K1D
Ky = K1D
VX, VY, EToV = uniform_quad_mesh(Kx, Ky)

"Initialize reference element"
r,_ = gauss_lobatto_quad(0,0,N)   # Reference nodes
VDM = vandermonde_1D(N,r)         # modal to nodal
Dr = grad_vandermonde_1D(N,r)/VDM # nodal differentiation
V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1]) # nodal linear interpolation

Nq = N
rq,wq = gauss_lobatto_quad(0,0,Nq)
M = diagm(wq)
Mlump = zeros(size(M))
Mlump_inv = zeros(size(M))
for i = 1:Nq+1
    Mlump[i,i] = sum(M[i,:])
    Mlump_inv[i,i] = 1.0/Mlump[i,i]
end

# 1D operators
M1D = Mlump
M1D_inv = Mlump_inv
Q1D = M*Dr
B1D = zeros(N+1,N+1)
B1D[1,1] = -1
B1D[end,end] = 1
L = Array(spdiagm(0=>-2*ones(N+1), 1=>ones(N), -1=>ones(N)))
L[1,1] = -1
L[end,end] = -1
psi = pinv(L)*-1/2*B1D*ones(N+1)
S0 = zeros(N+1,N+1)
for i = 1:N+1
    for j = 1:N+1
        if L[i,j] != 0
            S0[i,j] = psi[j] - psi[i]
        end
    end
end
Q01D = S0+1/2*B1D

# Drop zeros
Q1D = Matrix(droptol!(sparse(Q1D),TOL))
Q01D = Matrix(droptol!(sparse(Q01D),TOL))
B1D = Matrix(droptol!(sparse(B1D),TOL))
S01D = Matrix(droptol!(sparse(S0),TOL))
S1D = Matrix(droptol!(sparse((Q1D-Q1D')/2),TOL))

# Tensor product operators
M = droptol!(sparse(kron(M1D,M1D)),TOL)
Minv = spdiagm(0 => 1 ./ diag(M))
Qr = droptol!(sparse(kron(M1D,Q1D)),TOL)
Qs = droptol!(sparse(kron(Q1D,M1D)),TOL)
Qr0 = droptol!(sparse(kron(M1D,Q01D)),TOL)
Qs0 = droptol!(sparse(kron(Q01D,M1D)),TOL)
Sr = droptol!(sparse((Qr-Qr')/2),TOL)
Ss = droptol!(sparse((Qs-Qs')/2),TOL)
S0r = droptol!(sparse(kron(M1D,S01D)),TOL)
S0s = droptol!(sparse(kron(S01D,M1D)),TOL)
Br = droptol!(sparse(Qr+Qr'),TOL)
Bs = droptol!(sparse(Qs+Qs'),TOL)

# Make domain periodic, TODO: hardcoded
Jf = 1/K1D # Assume uniform quad mesh
K = K1D^2
Nfaces = 4
Nfp = Nfaces*N

# physical normal, assume standard uniform quad mesh
#=
nxJ = zeros(Nfp,K)
nyJ = zeros(Nfp,K)
for k = 1:K
    nxJ[1,k] = -Jf
    nxJ[2:N,k] .= zeros(N-1)
    nxJ[N+1,k] = Jf
    for i = 0:N-2
        nxJ[N+2+i*2,k] = -Jf
        nxJ[N+2+i*2+1,k] = Jf
    end
    nxJ[end-N,k] = -Jf
    nxJ[end-(N-1):end-1,k] .= zeros(N-1)
    nxJ[end,k] = Jf

    nyJ[1:N+1,k] .= -Jf*ones(N+1)
    nyJ[end-N:end,k] .= Jf*ones(N+1)
end
=#
nxJ = zeros((N+1)*(N+1),K)
nyJ = zeros((N+1)*(N+1),K)
for k = 1:K
    for i = 0:N
        nxJ[i*(N+1)+1,k] = -Jf
        nxJ[i*(N+1)+N+1,k] = Jf
    end
    nyJ[1:N+1,k] .= -Jf*ones(N+1)
    nyJ[end-N:end,k] .= Jf*ones(N+1)
end

# indices of nodal points on faces
face_idx = zeros(Nfp)
face_idx[1:N+1] .= 1:N+1
face_idx[N+2:2*N+2] .= N*(N+1)+1:(N+1)*(N+1)
for i = 1:N-1
    face_idx[2*N+3+2*(i-1)] = i*(N+1)+1
    face_idx[2*N+3+2*(i-1)+1] = i*(N+1)+N+1
end
sort!(face_idx)
face_idx = Int64.(face_idx)

# face connectivity via global node number
mapM = zeros(Int64,Nfp,K)
mapP = zeros(Int64,Nfp,K)
mapP_x = zeros(Int64,Nfp,K)
mapP_y = zeros(Int64,Nfp,K)
for kj = 1:K1D
    for ki = 1:K1D
        k = (kj-1)*K1D+ki
        mapM[:,k] .= (k-1)*(N+1)*(N+1).+face_idx
        mapP[1:N+1,k] = mod(k-1-K1D,K)*(N+1)*(N+1).+face_idx[end-N:end]
        mapP[end-N:end,k] = mod(k-1+K1D,K)*(N+1)*(N+1).+face_idx[1:N+1]
        mapP[N+2:2:end-N-2,k] = ((kj-1)*K1D+mod(ki-2,K1D))*(N+1)*(N+1).+face_idx[N+3:2:end-N-1] # left interface
        mapP[N+3:2:end-N-1,k] = ((kj-1)*K1D+mod(ki,K1D))*(N+1)*(N+1).+face_idx[N+2:2:end-N-2] # right interface
        mapP_x[[1;N+2:2:end-N],k] = ((kj-1)*K1D+mod(ki-2,K1D))*(N+1)*(N+1).+face_idx[[N+1:2:end-N-1;end]]
        mapP_x[[N+1:2:end-N-1;end],k] = ((kj-1)*K1D+mod(ki,K1D))*(N+1)*(N+1).+face_idx[[1;N+2:2:end-N]]
        mapP_y[1:N+1,k] = mod(k-1-K1D,K)*(N+1)*(N+1).+face_idx[end-N:end]
        mapP_y[end-N:end,k] = mod(k-1+K1D,K)*(N+1)*(N+1).+face_idx[1:N+1]
    end
end

rd = init_reference_quad(N,gauss_lobatto_quad(0,0,N))
md = init_mesh((VX,VY),EToV,rd)
@unpack x,y,rxJ,ryJ,sxJ,syJ = md

function translate_bump(x,y)
    rho = exp(1/10/(x^2+y^2+0.1))
    u = 1.0
    v = 0.0
    p = 1.0
    return rho,u,v,p
end
#U = @. translate_bump(x,y)
U = @. Riemann_2D(x,y)
#U = @. vortex(x,y,0.0)
rho,u,v,p = [(x->x[i]).(U) for i = 1:Nc]
U = @. primitive_to_conservative_hardcode(rho,u,v,p)
rho,rhou,rhov,E = [(x->x[i]).(U) for i = 1:Nc]
U = (rho,rhou,rhov,E)

function limiting_param(U_low, P_ij)
    l = 1.0
    # Limit density
    if U_low[1] + P_ij[1] < -TOL
        l = min(abs(U_low[1])/(abs(P_ij[1])+1e-14), 1.0)
    end

    # limiting internal energy (via quadratic function)
    # a = P_ij[1]*P_ij[3]-1.0/2.0*P_ij[2]^2
    # b = U_low[3]*P_ij[1]+U_low[1]*P_ij[3]-U_low[2]*P_ij[2]
    # c = U_low[3]*U_low[1]-1.0/2.0*U_low[2]^2
    a = P_ij[1]*P_ij[4]-1.0/2.0*(P_ij[2]^2+P_ij[3]^2)
    b = U_low[4]*P_ij[1]+U_low[1]*P_ij[4]-U_low[2]*P_ij[2]-U_low[3]*P_ij[3]
    c = U_low[4]*U_low[1]-1.0/2.0*(U_low[2]^2+U_low[3]^2)

    l_eps_ij = 1.0
    if b^2-4*a*c >= 0
        r1 = (-b+sqrt(b^2-4*a*c))/(2*a)
        r2 = (-b-sqrt(b^2-4*a*c))/(2*a)
        if r1 > TOL && r2 > TOL
            l_eps_ij = min(r1,r2)
        elseif r1 > TOL && r2 < -TOL
            l_eps_ij = r1
        elseif r2 > TOL && r1 < -TOL
            l_eps_ij = r2
        end
    end

    l = min(l,l_eps_ij)
    return l
end

function rhs_IDP(U,K1D,N,Nfaces,nxJ,nyJ,Minv,Sr,Ss,S0r,S0s,Br,Bs,face_idx,mapM,mapP_x,mapP_y,step)
    p = pfun_nd.(U[1],U[2],U[3],U[4])
    flux_x = zero.(U)
    flux_y = zero.(U)
    @. flux_x[1] = U[2]
    @. flux_x[2] = U[2]^2/U[1]+p
    @. flux_x[3] = U[2]*U[3]/U[1]
    @. flux_x[4] = U[4]*U[2]/U[1]+p*U[2]/U[1]
    @. flux_y[1] = U[3]
    @. flux_y[2] = U[2]*U[3]/U[1]
    @. flux_y[3] = U[3]^2/U[1]+p
    @. flux_y[4] = U[4]*U[3]/U[1]+p*U[3]/U[1]

    Ub = zero.(U)
    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[3]/U[1]
    @. Ub[4] = U[1]/(2*p)

    K = K1D^2
    Nfp = Nfaces*N
    Np = (N+1)*(N+1)
    J = (1/K1D)^2
    Jf = 1/K1D
    rxJ = 1/K1D
    syJ = 1/K1D
    sJ = 1/K1D

    F_low  = [zeros(Np,Np) for i = 1:Nc]
    F_high = [zeros(Np,Np) for i = 1:Nc]
    F_P    = [zeros(Nfp) for i = 1:Nc]
    rhsU   = [zeros(Np,K) for i = 1:Nc]
    U_low  = [zeros(Np) for i = 1:Nc]
    L = zeros(Np,Np)

    d_ii_arr = zeros(Np,K)
    # TODO: precompute
    for k = 1:K
        for i = 1:Np
            for j = 1:Np
                c_ij_norm = sqrt(S0r[i,j]^2+S0s[i,j]^2)
                if abs(c_ij_norm) >= TOL
                    n_ij = [S0r[i,j]; S0s[i,j]]./c_ij_norm
                    wavespd_i = wavespeed_1D(U[1][i,k],(n_ij[1]*U[2][i,k]+n_ij[2]*U[3][i,k]),U[4][i,k])
                    wavespd_j = wavespeed_1D(U[1][j,k],(n_ij[1]*U[2][j,k]+n_ij[2]*U[3][j,k]),U[4][j,k])
                    wavespd = max(wavespd_i,wavespd_j)
                    d_ij = wavespd*c_ij_norm
                    d_ii_arr[i,k] -= d_ij
                end
            end
        end

        S0r1 = sum(S0r,dims=2)
        S0s1 = sum(S0s,dims=2)
        for i = 1:Nfp
            fidM = mapM[i,k]
            S0rP = -S0r1[face_idx[i]]
            S0sP = -S0s1[face_idx[i]]

            # flux in x direction
            fidP = mapP_x[i,k]
            if fidP != 0
                wavespd_M = wavespeed_1D(U[1][fidM],U[2][fidM],U[4][fidM])
                wavespd_P = wavespeed_1D(U[1][fidP],U[2][fidP],U[4][fidP])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0rP)
                d_ii_arr[face_idx[i],k] -= d_ij
            end

            # flux in y direction
            fidP = mapP_y[i,k]
            if fidP != 0
                wavespd_M = wavespeed_1D(U[1][fidM],U[3][fidM],U[4][fidM])
                wavespd_P = wavespeed_1D(U[1][fidP],U[3][fidP],U[4][fidP])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0sP)
                d_ii_arr[face_idx[i],k] -= d_ij
            end
        end
    end

    dt = minimum(-J/2*M*(1 ./d_ii_arr))
    dt = dt

    for k = 1:K
        for c = 1:Nc
            for i = 1:Nfp
                F_P[c][i] = 0.0
            end
            for i = 1:Np
                U_low[c][i] = 0.0
            end
        end

        # Calculate interior low and high order flux
        for i = 1:Np
            for j = 1:Np
                c_ij_norm = sqrt(rxJ^2*S0r[i,j]^2+syJ^2*S0s[i,j]^2)
                if abs(c_ij_norm) >= TOL
                    n_ij = [rxJ*S0r[i,j]; syJ*S0s[i,j]]./c_ij_norm
                    wavespd_i = wavespeed_1D(U[1][i,k],(n_ij[1]*U[2][i,k]+n_ij[2]*U[3][i,k]),U[4][i,k])
                    wavespd_j = wavespeed_1D(U[1][j,k],(n_ij[1]*U[2][j,k]+n_ij[2]*U[3][j,k]),U[4][j,k])
                    wavespd = max(wavespd_i,wavespd_j)
                    d_ij = wavespd*c_ij_norm

                    for c = 1:Nc
                        F_low[c][i,j]  = (rxJ*S0r[i,j]*(flux_x[c][i,k]+flux_x[c][j,k])
                                        + syJ*S0s[i,j]*(flux_y[c][i,k]+flux_y[c][j,k])
                                        - d_ij*(U[c][j,k]-U[c][i,k]))
                    end
                end
            end
        end

        # High order flux
        for i = 1:Np-1
            for j = i+1:Np
                # TODO: can preallocate nonzero entries
                if Sr[i,j] != 0 || Ss[i,j] != 0
                    F1,F2 = euler_fluxes(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[4][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],Ub[4][j,k])
                    for c = 1:Nc
                        val = 2*rxJ*Sr[i,j]*F1[c]+2*syJ*Ss[i,j]*F2[c]
                        F_high[c][i,j] = val
                        F_high[c][j,i] = -val
                    end
                end
            end
        end

        # Calculate interface flux
        S0r1 = sum(S0r,dims=2)
        S0s1 = sum(S0s,dims=2)
        for i = 1:Nfp
            fidM = mapM[i,k]
            S0rP = -S0r1[face_idx[i]]
            S0sP = -S0s1[face_idx[i]]

            # flux in x direction
            fidP = mapP_x[i,k]
            if fidP != 0
                wavespd_M = wavespeed_1D(U[1][fidM],U[2][fidM],U[4][fidM])
                wavespd_P = wavespeed_1D(U[1][fidP],U[2][fidP],U[4][fidP])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0rP)
                for c = 1:Nc
                    F_P[c][i] += (Jf*S0rP*(flux_x[c][fidM]+flux_x[c][fidP])
                                - Jf*d_ij*(U[c][fidP]-U[c][fidM]))
                end
            end

            # flux in y direction
            fidP = mapP_y[i,k]
            if fidP != 0
                wavespd_M = wavespeed_1D(U[1][fidM],U[3][fidM],U[4][fidM])
                wavespd_P = wavespeed_1D(U[1][fidP],U[3][fidP],U[4][fidP])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0sP)
                for c = 1:Nc
                    F_P[c][i] += (Jf*S0sP*(flux_y[c][fidM]+flux_y[c][fidP])
                                - Jf*d_ij*(U[c][fidP]-U[c][fidM]))
                end
            end
        end

        # Calculate low order solution
        for i = 1:Np
            for c = 1:Nc
                U_low[c][i] = sum(F_low[c][i,:])
            end
        end
        # TODO: redundant
        for i = 1:Nfp
            for c = 1:Nc
                U_low[c][face_idx[i]] += F_P[c][i]
            end
        end
        for c = 1:Nc
            for i = 1:Np
                U_low[c][i] = U[c][i,k]-dt/J/M[i,i]*U_low[c][i]
            end
        end

        # Calculate limiting parameters
        P_ij = [0.0;0.0;0.0;0.0]
        for i = 1:Np
            m_i = J*M[i,i]
            lambda_j = 1/(Np-1)
            U_low_i = [U_low[c][i] for c = 1:Nc]
            for j = 1:Np
                if i != j
                    for c = 1:Nc
                        P_ij[c] = dt/(m_i*lambda_j)*(F_low[c][i,j]-F_high[c][i,j])
                    end
                    L[i,j] = limiting_param(U_low_i, P_ij)
                end
            end
        end

        # Symmetrize limiting parameters
        for i = 1:Np
            for j = 1:Np
                if i != j
                    l = min(L[i,j],L[j,i])
                    L[i,j] = l
                    L[j,i] = l
                end
            end
        end

        # L = ones(size(L))
        # if sum(L)-Np*Np+Np < 0
        #     @show sum(L)-Np*Np+Np
        # end

        for c = 1:Nc
            rhsU[c][:,k] = sum((L.-1).*F_low[c] - L.*F_high[c],dims=2)
            #rhsU[c][:,k] = sum((-1).*F_high[c],dims=2)
            #rhsU[c][:,k] = sum((-1).*F_low[c],dims=2)
            for i = 1:Nfp
                rhsU[c][face_idx[i],k] -= F_P[c][i]
            end
        end
    end

    for c = 1:Nc
        rhsU[c] = 1/J*Minv*rhsU[c]
    end

    return rhsU,dt
end

function rhs_IDP(U,K1D,N,Nfaces,nxJ,nyJ,Minv,Sr,Ss,S0r,S0s,Br,Bs,face_idx,mapM,mapP_x,mapP_y,step,dt)
    p = pfun_nd.(U[1],U[2],U[3],U[4])
    flux_x = zero.(U)
    flux_y = zero.(U)
    @. flux_x[1] = U[2]
    @. flux_x[2] = U[2]^2/U[1]+p
    @. flux_x[3] = U[2]*U[3]/U[1]
    @. flux_x[4] = U[4]*U[2]/U[1]+p*U[2]/U[1]
    @. flux_y[1] = U[3]
    @. flux_y[2] = U[2]*U[3]/U[1]
    @. flux_y[3] = U[3]^2/U[1]+p
    @. flux_y[4] = U[4]*U[3]/U[1]+p*U[3]/U[1]

    Ub = zero.(U)
    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[3]/U[1]
    @. Ub[4] = U[1]/(2*p)

    K = K1D^2
    Nfp = Nfaces*N
    Np = (N+1)*(N+1)
    J = (1/K1D)^2
    Jf = 1/K1D
    rxJ = 1/K1D
    syJ = 1/K1D
    sJ = 1/K1D

    F_low  = [zeros(Np,Np) for i = 1:Nc]
    F_high = [zeros(Np,Np) for i = 1:Nc]
    F_P    = [zeros(Nfp) for i = 1:Nc]
    rhsU   = [zeros(Np,K) for i = 1:Nc]
    U_low  = [zeros(Np) for i = 1:Nc]
    L = zeros(Np,Np)

    for k = 1:K
        for c = 1:Nc
            for i = 1:Nfp
                F_P[c][i] = 0.0
            end
            for i = 1:Np
                U_low[c][i] = 0.0
            end
        end

        # Calculate interior low and high order flux
        for i = 1:Np
            for j = 1:Np
                c_ij_norm = sqrt(rxJ^2*S0r[i,j]^2+syJ^2*S0s[i,j]^2)
                if abs(c_ij_norm) >= TOL
                    n_ij = [rxJ*S0r[i,j]; syJ*S0s[i,j]]./c_ij_norm
                    wavespd_i = wavespeed_1D(U[1][i,k],(n_ij[1]*U[2][i,k]+n_ij[2]*U[3][i,k]),U[4][i,k])
                    wavespd_j = wavespeed_1D(U[1][j,k],(n_ij[1]*U[2][j,k]+n_ij[2]*U[3][j,k]),U[4][j,k])
                    wavespd = max(wavespd_i,wavespd_j)
                    d_ij = wavespd*c_ij_norm

                    for c = 1:Nc
                        F_low[c][i,j]  = (rxJ*S0r[i,j]*(flux_x[c][i,k]+flux_x[c][j,k])
                                        + syJ*S0s[i,j]*(flux_y[c][i,k]+flux_y[c][j,k])
                                        - d_ij*(U[c][j,k]-U[c][i,k]))
                    end
                end
            end
        end

        # High order flux
        for i = 1:Np-1
            for j = i+1:Np
                # TODO: can preallocate nonzero entries
                if Sr[i,j] != 0 || Ss[i,j] != 0
                    F1,F2 = euler_fluxes(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[4][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],Ub[4][j,k])
                    for c = 1:Nc
                        val = 2*rxJ*Sr[i,j]*F1[c]+2*syJ*Ss[i,j]*F2[c]
                        F_high[c][i,j] = val
                        F_high[c][j,i] = -val
                    end
                end
            end
        end

        # Calculate interface flux
        S0r1 = sum(S0r,dims=2)
        S0s1 = sum(S0s,dims=2)
        for i = 1:Nfp
            fidM = mapM[i,k]
            S0rP = -S0r1[face_idx[i]]
            S0sP = -S0s1[face_idx[i]]

            # flux in x direction
            fidP = mapP_x[i,k]
            if fidP != 0
                wavespd_M = wavespeed_1D(U[1][fidM],U[2][fidM],U[4][fidM])
                wavespd_P = wavespeed_1D(U[1][fidP],U[2][fidP],U[4][fidP])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0rP)
                for c = 1:Nc
                    F_P[c][i] += (Jf*S0rP*(flux_x[c][fidM]+flux_x[c][fidP])
                                - Jf*d_ij*(U[c][fidP]-U[c][fidM]))
                end
            end

            # flux in y direction
            fidP = mapP_y[i,k]
            if fidP != 0
                wavespd_M = wavespeed_1D(U[1][fidM],U[3][fidM],U[4][fidM])
                wavespd_P = wavespeed_1D(U[1][fidP],U[3][fidP],U[4][fidP])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0sP)
                for c = 1:Nc
                    F_P[c][i] += (Jf*S0sP*(flux_y[c][fidM]+flux_y[c][fidP])
                                - Jf*d_ij*(U[c][fidP]-U[c][fidM]))
                end
            end
        end

        # Calculate low order solution
        for i = 1:Np
            for c = 1:Nc
                U_low[c][i] = sum(F_low[c][i,:])
            end
        end
        # TODO: redundant
        for i = 1:Nfp
            for c = 1:Nc
                U_low[c][face_idx[i]] += F_P[c][i]
            end
        end
        for c = 1:Nc
            for i = 1:Np
                U_low[c][i] = U[c][i,k]-dt/J/M[i,i]*U_low[c][i]
            end
        end

        # Calculate limiting parameters
        P_ij = [0.0;0.0;0.0;0.0]
        for i = 1:Np
            m_i = J*M[i,i]
            lambda_j = 1/(Np-1)
            U_low_i = [U_low[c][i] for c = 1:Nc]
            for j = 1:Np
                if i != j
                    for c = 1:Nc
                        P_ij[c] = dt/(m_i*lambda_j)*(F_low[c][i,j]-F_high[c][i,j])
                    end
                    L[i,j] = limiting_param(U_low_i, P_ij)
                end
            end
        end

        # Symmetrize limiting parameters
        for i = 1:Np
            for j = 1:Np
                if i != j
                    l = min(L[i,j],L[j,i])
                    L[i,j] = l
                    L[j,i] = l
                end
            end
        end

        # L = zeros(size(L))
        # if sum(L)-Np*Np+Np < 0
        #     @show sum(L)-Np*Np+Np
        # end

        for c = 1:Nc
            rhsU[c][:,k] = sum((L.-1).*F_low[c] - L.*F_high[c],dims=2)
            #rhsU[c][:,k] = sum((-1).*F_high[c],dims=2)
            #rhsU[c][:,k] = sum((-1).*F_low[c],dims=2)
            for i = 1:Nfp
                rhsU[c][face_idx[i],k] -= F_P[c][i]
            end
        end
    end

    for c = 1:Nc
        rhsU[c] = 1/J*Minv*rhsU[c]
    end

    return rhsU,dt
end

function rhs_IDPlow(U,K1D,N,Nfaces,nxJ,nyJ,Minv,Sr,Ss,S0r,S0s,Br,Bs,face_idx,mapM,mapP_x,mapP_y)
    p = pfun_nd.(U[1],U[2],U[3],U[4])
    flux_x = zero.(U)
    flux_y = zero.(U)
    @. flux_x[1] = U[2]
    @. flux_x[2] = U[2]^2/U[1]+p
    @. flux_x[3] = U[2]*U[3]/U[1]
    @. flux_x[4] = U[4]*U[2]/U[1]+p*U[2]/U[1]
    @. flux_y[1] = U[3]
    @. flux_y[2] = U[2]*U[3]/U[1]
    @. flux_y[3] = U[3]^2/U[1]+p
    @. flux_y[4] = U[4]*U[3]/U[1]+p*U[3]/U[1]

    K = K1D^2
    Nfp = Nfaces*N
    Np = (N+1)*(N+1)
    J = (1/K1D)^2
    Jf = 1/K1D
    rxJ = 1/K1D
    syJ = 1/K1D
    sJ = 1/K1D

    # Initialize sparse matrices to store numerical fluxes across interfaces
    I_arr = Int64[]
    J_arr = Int64[]
    for k = 1:K
        for n = 1:Nfp
            i = mapM[n,k]
            jx = mapP_x[n,k]
            jy = mapP_y[n,k]
            if jx != 0
                push!(I_arr,i)
                push!(J_arr,jx)
            end
            if jy != 0
                push!(I_arr,i)
                push!(J_arr,jy)
            end
        end
    end
    F_P_low = [sparse(I_arr,J_arr,zeros(size(I_arr))) for i = 1:Nc]
    # Low order and high order algebraic fluxes
    F_low   = [zeros(Np,Np) for i = 1:Nc]
    rhsU    = [zeros(Np,K) for i = 1:Nc]

    for k = 1:K
        # TODO: redundant iterations
        for i = 1:Np
            for j = 1:Np
                c_ij_norm = sqrt(rxJ^2*S0r[i,j]^2+syJ^2*S0s[i,j]^2)
                if abs(c_ij_norm) >= TOL
                    n_ij = [rxJ*S0r[i,j]; syJ*S0s[i,j]]./c_ij_norm
                    wavespd_i = wavespeed_1D(U[1][i,k],(n_ij[1]*U[2][i,k]+n_ij[2]*U[3][i,k]),U[4][i,k])
                    wavespd_j = wavespeed_1D(U[1][j,k],(n_ij[1]*U[2][j,k]+n_ij[2]*U[3][j,k]),U[4][j,k])
                    wavespd = max(wavespd_i,wavespd_j)
                    d_ij = wavespd*c_ij_norm

                    for c = 1:Nc
                        F_low[c][i,j] = (rxJ*S0r[i,j]*(flux_x[c][i,k]+flux_x[c][j,k])
                                       + syJ*S0s[i,j]*(flux_y[c][i,k]+flux_y[c][j,k])
                                       - d_ij*(U[c][j,k]-U[c][i,k]))
                    end
                end
            end
        end

        S0r1 = sum(S0r,dims=2)
        S0s1 = sum(S0s,dims=2)
        # Interface low order flux
        for i in 1:Nfp
            fidM = mapM[i,k]

            S0rP = -S0r1[face_idx[i]]
            S0sP = -S0s1[face_idx[i]]

            # numerical flux in x direction
            fidP = mapP_x[i,k]
            if fidP != 0
                c_ij_norm = abs(S0rP)
                if abs(c_ij_norm) >= TOL
                    n_ij_x = S0rP/c_ij_norm
                    wavespd_M = wavespeed_1D(U[1][fidM],n_ij_x*U[2][fidM],U[4][fidM])
                    wavespd_P = wavespeed_1D(U[1][fidP],n_ij_x*U[2][fidP],U[4][fidP])
                    wavespd = max(wavespd_M,wavespd_P)
                    d_ij = wavespd*c_ij_norm
                    for c = 1:Nc
                        F_P_low[c][fidM,fidP] += (Jf*S0rP*(flux_x[c][fidM]+flux_x[c][fidP])
                                                - Jf*d_ij*(U[c][fidP]-U[c][fidM]))
                    end
                end
            end

            # numerical flux in y direction
            fidP = mapP_y[i,k]
            if fidP != 0
                c_ij_norm = abs(S0sP)
                if abs(c_ij_norm) >= TOL
                    n_ij_y = S0sP/c_ij_norm
                    wavespd_M = wavespeed_1D(U[1][fidM],n_ij_y*U[3][fidM],U[4][fidM])
                    wavespd_P = wavespeed_1D(U[1][fidP],n_ij_y*U[3][fidP],U[4][fidP])
                    wavespd = max(wavespd_M,wavespd_P)
                    d_ij = wavespd*c_ij_norm
                    for c = 1:Nc
                        F_P_low[c][fidM,fidP] += (Jf*S0sP*(flux_y[c][fidM]+flux_y[c][fidP])
                                                - Jf*d_ij*(U[c][fidP]-U[c][fidM]))
                    end
                end
            end
        end

        for i = 1:Np
            for c = 1:Nc
                rhsU[c][i,k] += sum(F_low[c][i,:])
                rhsU[c][i,k] += sum(F_P_low[c][i+(k-1)*Np,:])
            end
        end

    end

    for c = 1:Nc
        rhsU[c] = -1/J*Minv*rhsU[c]
    end

    return rhsU
end

function rhs_ESDG(U,K1D,N,Nfaces,nxJ,nyJ,Minv,Sr,Ss,S0r,S0s,Br,Bs,face_idx,mapM,mapP_x,mapP_y)
    p = pfun_nd.(U[1],U[2],U[3],U[4])

    Ub = zero.(U)
    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[3]/U[1]
    @. Ub[4] = U[1]/(2*p)

    K = K1D^2
    Nfp = Nfaces*N
    Np = (N+1)*(N+1)
    J = (1/K1D)^2
    Jf = 1/K1D

    # Low order and high order algebraic fluxes
    # Initialize sparse matrices to store numerical fluxes across interfaces
    I_arr = Int64[]
    J_arr = Int64[]
    for k = 1:K
        for n = 1:Nfp
            i = mapM[n,k]
            jx = mapP_x[n,k]
            jy = mapP_y[n,k]
            if jx != 0
                push!(I_arr,i)
                push!(J_arr,jx)
            end
            if jy != 0
                push!(I_arr,i)
                push!(J_arr,jy)
            end
        end
    end
    F_P_high = [sparse(I_arr,J_arr,zeros(size(I_arr))) for i = 1:Nc]
    F_high   = [zeros(Np,Np) for i = 1:Nc]
    rhsU     = [zeros(Np,K) for i = 1:Nc]

    # TODO: assume mesh is uniform, so geometric factors constant
    #       and mesh is not rotated, i.e. ryJ = sxJ = 0.0
    rxJ = 1/K1D
    syJ = 1/K1D
    sJ = 1/K1D

    for k = 1:K
        for i = 1:Np-1
            for j = i+1:Np
                # TODO: can preallocate nonzero entries
                if Sr[i,j] != 0 || Ss[i,j] != 0
                    F1,F2 = euler_fluxes(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[4][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],Ub[4][j,k])
                    for c = 1:Nc
                        val = 2*rxJ*Sr[i,j]*F1[c]+2*syJ*Ss[i,j]*F2[c]
                        F_high[c][i,j] = val
                        F_high[c][j,i] = -val
                    end
                end
            end
        end

        # Interface high order flux
        for i in 1:Nfp
            fidM = mapM[i,k]
            nxJM = nxJ[fidM]
            nyJM = nyJ[fidM]
            rhou_nM_x = (U[2][fidM]*nxJM)/sJ
            rhou_nM_y = (U[3][fidM]*nyJM)/sJ
            wavespdM_x = wavespeed_1D(U[1][fidM],rhou_nM_x,U[4][fidM])
            wavespdM_y = wavespeed_1D(U[1][fidM],rhou_nM_y,U[4][fidM])

            # numerical flux in x direction
            fidP = mapP_x[i,k]
            if fidP != 0
                nxJP = nxJ[fidP]
                nyJP = nyJ[fidP]
                #rhou_nP = (U[2][fidP]*nxJP+U[3][fidP]*nyJP)/sJ
                rhou_nP = (U[2][fidP]*nxJP)/sJ
                wavespdP = wavespeed_1D(U[1][fidP],rhou_nP,U[4][fidP])
                wavespd = max(wavespdM_x,wavespdP)
                F1,_ = euler_fluxes(Ub[1][fidM],Ub[2][fidM],Ub[3][fidM],Ub[4][fidM],Ub[1][fidP],Ub[2][fidP],Ub[3][fidP],Ub[4][fidP])
                for c = 1:Nc
                    F_P_high[c][fidM,fidP] += Jf*Br[face_idx[i],face_idx[i]]*F1[c]-sJ*wavespd/2*(U[c][fidP]-U[c][fidM])#nxJM*Br[face_idx[i],face_idx[i]]*F1[c]#-wavespd/2*(U[c][fidP]-U[c][fidM])
                end
            end

            # numerical flux in y direction
            fidP = mapP_y[i,k]
            if fidP != 0
                nxJP = nxJ[fidP]
                nyJP = nyJ[fidP]
                #rhou_nP = (U[2][fidP]*nxJP+U[3][fidP]*nyJP)/sJ
                rhou_nP = (U[3][fidP]*nyJP)/sJ
                wavespdP = wavespeed_1D(U[1][fidP],rhou_nP,U[4][fidP])
                wavespd = max(wavespdM_y,wavespdP)
                _,F2 = euler_fluxes(Ub[1][fidM],Ub[2][fidM],Ub[3][fidM],Ub[4][fidM],Ub[1][fidP],Ub[2][fidP],Ub[3][fidP],Ub[4][fidP])
                for c = 1:Nc
                    F_P_high[c][fidM,fidP] += Jf*Bs[face_idx[i],face_idx[i]]*F2[c]-sJ*wavespd/2*(U[c][fidP]-U[c][fidM])#nyJM*Bs[face_idx[i],face_idx[i]]*F2[c]#-wavespd/2*(U[c][fidP]-U[c][fidM])
                end
            end
        end

        for i = 1:Np
            for c = 1:Nc
                rhsU[c][i,k] += sum(F_high[c][i,:])
                rhsU[c][i,k] += sum(F_P_high[c][i+(k-1)*Np,:])
            end
        end
    end

    for c = 1:Nc
        rhsU[c] = -1/J*Minv*rhsU[c]
    end

    return rhsU

end

# Time stepping
"Time integration"
t = 0.0
U = collect(U)
resU = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]
resW = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]
resZ = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]

@unpack VDM = rd
rp,sp = equi_nodes_2D(20)
Vp = vandermonde_2D(N,rp,sp)/VDM
gr(aspect_ratio=:equal,legend=false,
   markerstrokewidth=0,markersize=2,axis=nothing)
const GIFINTERVAL = 100
#plot()

dt_hist = []
anim = Animation()
i = 1

while t < T
    #rhsU,dt = rhs_IDP(U,K1D,N,Nfaces,nxJ,nyJ,Minv,Sr,Ss,S0r,S0s,Br,Bs,face_idx,mapM,mapP_x,mapP_y,i)
    # @. U = U + dt*rhsU

    # SSPRK(3,3)
    dt = min(1e-4,T-t)
    rhsU,dt = rhs_IDP(U,K1D,N,Nfaces,nxJ,nyJ,Minv,Sr,Ss,S0r,S0s,Br,Bs,face_idx,mapM,mapP_x,mapP_y,i,dt)
    @. resW = U + dt*rhsU
    rhsU,_ = rhs_IDP(resW,K1D,N,Nfaces,nxJ,nyJ,Minv,Sr,Ss,S0r,S0s,Br,Bs,face_idx,mapM,mapP_x,mapP_y,i,dt)
    @. resZ = resW+dt*rhsU
    @. resW = 3/4*U+1/4*resZ
    rhsU,_ = rhs_IDP(resW,K1D,N,Nfaces,nxJ,nyJ,Minv,Sr,Ss,S0r,S0s,Br,Bs,face_idx,mapM,mapP_x,mapP_y,i,dt)
    @. resZ = resW+dt*rhsU
    @. U = 1/3*U+2/3*resZ

    push!(dt_hist,dt)
    global t = t + dt
    println("Current time $t with time step size $dt, and final time $T, at step $i")
    global i = i + 1
    # if mod(i,50) == 1
    #     scatter(Vp*x,Vp*y,Vp*U[1],zcolor=Vp*U[1],camera=(0,90))
    #     frame(anim)
    # end
end

xp = Vp*x
yp = Vp*y
rhop = Vp*U[1]
p = pfun_nd.(U[1],U[2],U[3],U[4])
pp = Vp*p
#plt = scatter(Vp*x,Vp*y,Vp*U[1],zcolor=Vp*U[1],camera=(0,90))
plt = scatter(xp,yp,rhop,zcolor=rhop,camera=(0,90))
display(plt)
#savefig(plt,"~/Desktop/tmp.png")
#gif(anim,"~/Desktop/tmp.gif",fps=15)


using Sundials

## Define the system residual function.
function resrob(tres, y, yp, r)
    r[1] = -0.04 * y[1] + 1.0e4 * y[2] * y[3]
    r[2] = -r[1] - 3.0e7 * y[2] * y[2] - yp[2]
    r[1] -= yp[1]
    r[3] = y[1] + y[2] + y[3] - 1.0
end

t = [0.0; 4 * exp10.(range(-1.0, stop = 5.0, length = 7))]
yout, ypout = Sundials.idasol(resrob, [1.0, 0, 0], [-0.04, 0.04, 0.0], t)


using JuMP, BARON, Base.Test

m = Model(solver=BaronSolver())
ub = [2, 2, 1]
@variable(m, 12 ≤ x[i=1:4] ≤ 60, start = 24)

@NLconstraints(m, begin
    x[3] ≤ x[4]
    x[2] ≤ x[1]
end)

@NLobjective(m, Min, 6.931 - x[1]*x[2]/(x[3]*x[4])^2 + 1)

solve(m)


using InvertibleNetworks, NNlib, LinearAlgebra, Test

# Data
nx = 32
ny = 32
nz = 32
n_in = 3
batchsize = 4
k = 3   # kernel size
s = 1   # stride
p = 1   # padding

###################################################################################################
# Hyperbolic layer

# Data
X_prev = randn(Float32, nx, ny, nz, n_in, batchsize)
X_curr = randn(Float32, nx, ny, nz, n_in, batchsize)

# Network architecture ((decrease channel no?, no. of hidden units), ...)
architecture = ((0, 3), (-1, 3,), (0, 3), (1, 4), (0, 8))

# Hyperbolic network
HN = H = NetworkHyperbolic3D(n_in, architecture; k=3, s=1, p=1, logdet=true, α=1f0)

# Forward/inverse pass
Y_curr_, Y_new_, lgdet = HN.forward(X_prev, X_curr)
X_prev_, X_curr_ = HN.inverse(Y_curr_, Y_new_)

# Test invertibility
@test isapprox(norm(X_prev - X_prev_)/norm(X_prev_), 0f0; atol=1f-3)
@test isapprox(norm(X_curr - X_curr_)/norm(X_prev_), 0f0; atol=1f-3)


using SparseArrays, Arpack            # 引入稀疏矩阵库和 Arpack 数学库的接口
Ksp=Symmetric(sprand(5000,5000,0.01)) # 生成一个对称的稀疏矩阵
                                      # 注意，您的内存可能不够大，5000 可以适当调小
Kde=Symmetric(Matrix(K))              # 对应的稠密矩阵
@time eigen(Kde)                      # 计算本征对，计时
  # 19.264927 秒，您的时间大概和这里不同
@time Arpack.eigs(Ksp, nev=100)       # 计算前 100 个本征对
  #  1.389663 秒


#using SurfaceGeometry

struct Zinchenko1997
    VertexRing
end

# function Zinchenko1997(points,faces,n;VertexRing=i -> VertexRing(i,faces))
#     Zinchenko1997(VertexRing)
# end

eps = 1e-6

function f(i::Int,ni::Array{Float64,1},points,v,ring)
    s = zeros(3)
    pi = points[:,i]
    vi = v[:,i]
    for j in ring
        xij = points[:,j] - pi
        vij = v[:,j] - vi
        s += dot(xij,vij)*xij
    end
    P = eye(3) - ni*ni'
    return P*s
end

function fvec(i,points,v,ring)
    s = zeros(3)
    pi = points[:,i]
    vi = v[:,i]
    for j in ring
        xij = points[:,j] - pi
        vij = v[:,j] - vi
        s += dot(xij,vij)*xij
    end
    return s
end

function F(points,faces,v,mc::Zinchenko1997)

    function f(v1,v2)
        xij = points[:,v1] - points[:,v2]
        vij = v[:,v1] - v[:,v2]
        return dot(xij,vij)^2
    end

    N = size(faces,2)
    s = 0
    for i in 1:N
        v1,v2,v3 = faces[:,i]
        s += v2>v1 ? f(v1,v2) : 0
        s += v3>v2 ? f(v2,v3) : 0
        s += v3>v1 ? f(v1,v3) : 0
    end
    return 4*s
end

# Returns ksi and eta for conjugate gradient iterations
function conjugrad(points,faces,va,vb,f)
    function A(i,j)
        xij = points[:,j] - points[:,i]
        vij = vb[:,j] - vb[:,i]
        return dot(xij,vij)
    end

    function B(i,j)
        xij = points[:,j] - points[:,i]
        fij = f[:,j] - f[:,i]
        return dot(xij,fij)
    end

    function C(i,j)
        xij = points[:,j] - points[:,i]
        Dvij = vb[:,j] - va[:,j] - (vb[:,i] - va[:,i])
        return dot(xij,Dvij)
    end

    S = zeros(2,2)
    P = zeros(2)

    function add!(i,j)
        a = A(i,j)
        b = B(i,j)
        c = C(i,j)

        S[1,1] += b^2
        S[1,2] += b*c
        S[2,1] += b*c
        S[2,2] += c^2

        P[1] += -a*b
        P[2] += -a*c
    end

    for i in 1:size(points,2)
        v1,v2,v3 = faces[:,i]

        if v2>v1
            add!(v1,v2)
        end

        if v3>v2
            add!(v3,v2)
        end

        if v3>v1
            add!(v1,v3)
        end
    end

    ksi, eta = S \ P
    return ksi,eta
end

# A simple conjugate gradient
function conjugrad0(points,faces,v,f)
    function A(i,j)
        xij = points[:,j] - points[:,i]
        vij = v[:,j] - v[:,i]
        return dot(xij,vij)
    end

    function B(i,j)
        xij = points[:,j] - points[:,i]
        fij = f[:,j] - f[:,i]
        return dot(xij,fij)
    end

    P = 0
    S = 0

    for i in 1:size(points,2)
        v1,v2,v3 = faces[:,i]

        if v2>v1
            a = A(v1,v2)
            b = B(v1,v2)
            S += b^2
            P += - a*b
        end

        if v3>v2
            a = A(v2,v3)
            b = B(v2,v3)
            S += b^2
            P += - a*b
        end

        if v3>v1
            a = A(v1,v3)
            b = B(v1,v3)
            S += b^2
            P += - a*b
        end
    end

    return P/S
end

function stabilise!(points,faces,n,va,mc::Zinchenko1997)

    VertexRing = mc.VertexRing

    ff = Array{Float64}(undef,size(points)...)
    for j in 1:size(points,2)
        ring = VertexRing(j)
        ff[:,j] = f(j,n[:,j],points,va,ring)
    end
    ksi0 = conjugrad0(points,faces,va,ff)
    vb = va + ksi0*ff
    vtemp = Array{Float64}(undef,size(va)...)

    Fitemp = Inf

    for i in 1:200

        for j in 1:size(points,2)
            ring = VertexRing(j)
            ff[:,j] = f(j,n[:,j],points,vb,ring)
        end

        ksi,eta = conjugrad(points,faces,va,vb,ff)

        vtemp[:,:] = vb[:,:]
        for j in 1:size(points,2)
            vb[:,j] = vb[:,j] + ksi*ff[:,j] + eta*(vb[:,j]-va[:,j])
        end
        va[:,:] = vtemp[:,:]

        Fi = F(points,faces,vb,mc)
        #println("F is $Fi and ksi is $ksi")

        if abs(Fitemp-Fi)<eps*Fi
            break

            if i>100
                warn("Stabilisation had $i iterations")
            end
        end
        Fitemp = Fi
    end
end


@testset "vernal point" begin
    hours = 0:24
    lat,long = 0.0,0.0
    deg2second = 24*3600/360
    datetimes = DateTime(2021,3,20) .+ Hour.(hours) .- Second(round(long*deg2second))
    res2 = @pipe calc_sun_position_MOD.(datetime2julian.(datetimes)) |> DataFrame(_)
    # latitude over ecliptic of sun is zero
    @test all(res2.β .≈ 0.0)
    au2m = 149597870700.0
    # distance to sun about one astromical unit (150_000 km)
    @test all(.≈(res2.r,  au2m, rtol = 0.02))
    # declination crossing from negative to positive
    @test res2.δ[1] < 0
    @test res2.δ[end] > 0
    # azimuth crossing zero longitude
    @test res2.α[1] < 0
    @test res2.α[end] > 0
    # function tmpf()
    #     using Plots, StatsPlots
    #     @df res2 scatter(hours, cols([:α, :δ]), legend = :topleft, xlab="hours")
    #     @df res3 scatter(hours, cols(1:2), legend = :topleft, xlab="hours")
    # end
    #
    res3 = @pipe calc_sun_position_hor.(datetimes, lat, long) |> DataFrame(_)
    # altitude (-pi, +pi)
    @test minimum(res3.altitude) ≈ -π/2 atol = 0.1
    @test maximum(res3.altitude) ≈ +π/2 atol = 0.1
    #
    lat,long = 45.0,0.0
    datetimes = DateTime(2021,3,20) .+ Hour.(hours) .- Second(round(long*deg2second))
    res3 = @pipe calc_sun_position_hor.(datetimes, lat, long) |> DataFrame(_)
    # altitude (-pi, +pi)
    @test minimum(res3.altitude) ≈ -π/4 atol = 0.1
    @test maximum(res3.altitude) ≈ +π/4 atol = 0.1
    # azimuth increasing (aside edge cases)
    @test all(diff(res3.azimuth[2:(end-1)]) .> 0)
end

@testset "Dresden summer" begin
    hours = 0:24
    lat,long = 51.0, 13.6
    deg2second = 24*3600/360
    doy = 160
    datetimes = DateTime(2021) .+Day(doy-1) .+ Hour.(hours) .- Second(round(long*deg2second))
    # function tmpf()
    #     using Plots, StatsPlots
    #     @df res2 scatter(hours, cols([:α, :δ]), legend = :topleft, xlab="hours")
    #     @df res3 scatter(hours, cols(1:2), legend = :topleft, xlab="hours")
    # end
    #
    res3 = @pipe calc_sun_position_hor.(datetimes, lat, long) |> DataFrame(_)
    # maximum altitude at noon
    @test argmax(res3.altitude) == 1 + (length(hours)-1) ÷ 2
    # azimuth increasing (aside edge cases)
    @test all(diff(res3.azimuth[2:(end-1)]) .> 0)
end


using PDELib
using Triangulate

function fv_laplace_circle(;Plotter=nothing,nref=0)

    function circle!(builder, center, radius; n=20*2^nref)
        points=[point!(builder, center[1]+radius*sin(t),center[2]+radius*cos(t)) for t in range(0,2π,length=n)]
        for i=1:n-1
            facet!(builder,points[i],points[i+1])
        end
        facet!(builder,points[end],points[1])
    end

    builder=SimplexGridBuilder(Generator=Triangulate)
    cellregion!(builder,1)
    maxvolume!(builder,0.1)
    regionpoint!(builder,0,0)

    facetregion!(builder,1)
    circle!(builder,(0,0),1)

    facetregion!(builder,2)
    circle!(builder,(-0.2,-0.2),0.2)
    holepoint!(builder,-0.2,-0.2)

    grid=simplexgrid(builder,maxvolume=0.1*4.0^(-nref))

    function flux(f,u0,edge)
        u=unknowns(edge,u0)
        f[1]=u[1,1]-u[1,2]
    end
    physics=VoronoiFVM.Physics(num_species=1,flux=flux)

    sys=VoronoiFVM.System(grid,physics)
    ispec=1
    enable_species!(sys,ispec,[1])
    boundary_dirichlet!(sys,ispec,1,0.0)
    boundary_dirichlet!(sys,ispec,2,1.0)
    inival=unknowns(sys,inival=0)
    solution=unknowns(sys)
    VoronoiFVM.solve!(solution,inival,sys)

    visualizer=GridVisualizer(Plotter=Plotter,layout=(1,2),resolution=(800,400))
    gridplot!(visualizer[1,1],grid)
    scalarplot!(visualizer[1,2],grid,solution[1,:])
    reveal(visualizer)
    true
end


type DH
    dhpar::Matrix{Float64}
    offset::Vector{Float64}
    GR::Matrix{Float64}
end

DH(dhpar,offset) = DH(dhpar,offset,eye(length(offset)))

function DH2400()
    DH_alpha0 = -π/2
    DH_alpha1 =  0
    DH_alpha2 = -π/2
    DH_alpha3 = -π/2
    DH_alpha4 = +π/2
    DH_alpha5 =  0

    DH_a0= 100
    DH_a1= 705
    DH_a2= 135

    DH_d0 =615
    DH_d1 =0
    DH_d2 =0
    DH_d3 =755
    DH_d4 =0
    DH_d5 =85

    offset = [0, -90, 0, 180, 0, 0].*π/180

    dhpar   = [ DH_alpha0      DH_a0        0         DH_d0        0;
    DH_alpha1      DH_a1        0         DH_d1        0;
    DH_alpha2      DH_a2        0         DH_d2        0;
    DH_alpha3       0           0         DH_d3        0;
    DH_alpha4       0           0         DH_d4        0;
    DH_alpha5       0           0         DH_d5        0 ]

    dhpar[:,[2,4]] /= 1000
    return DH(dhpar,offset)
end

function DHYuMi(left=true)
    # 100 100 -100 100 -101 100 100
    # baseAnglesRight = [0.63 ; 0.95 ; 0.18 ];
    # rotRight = rotx(0.63)*roty(0.95)*rotz(0.18);
    # tRight = [39.3 -64.5 403.5]';
    # baseRight = [rotRight tRight; 0 0 0 1]
    #
    # baseAnglesLeft = [-0.629355 ; 0.950621 ; -0.184101 ];
    # rotLeft = rotx(-0.629355)*roty(0.950621)*rotz(-0.184101);
    # tLeft = [39.1 66.5 407.3]';
    # baseLeft = [rotLeft tLeft; 0 0 0 1]
    baseAnglesRight = [0.63 , 0.95 , 0.18 ]
    baseAnglesLeft = [-0.63 , 0.95 , -0.18]

    # gear ratio must be applied AFTER abb2logical
    gearRatio = zeros(7,7)
    gearRatio[1,1] = 1/100
    gearRatio[2,2] = 1/100

    gearRatio[7,7] = 1/100
    gearRatio[3,3] = 1/100
    gearRatio[4,4] = -1/100
    gearRatio[5,5] = 1/100
    gearRatio[6,6] = -1/101

#     tmp(1:2) = tmp(1:2)/100;
# tmp(3) = -1*tmp(3)/100;
# tmp(4) = tmp(4)/100;
# tmp(5) = -1*tmp(5)/101;
# tmp(6:7) = tmp(6:7)/100;

    DH_alpha0 = -pi/2;
    DH_alpha1 = pi/2;
    DH_alpha2 = -pi/2;
    DH_alpha3 = pi/2;
    DH_alpha4 = -pi/2;
    DH_alpha5 =  pi/2;
    DH_alpha6 = 0;

    upArmLen =      0.2465;
    lowArmLen =     0.265;
    handLen =       0.032;
    shoulderLen =   0.110;

    shoulderOffs = -0.030;
    elbowOffs =     0.0405;
    wristOffs =     0.0135;
    handOffs =     -0.027;
    n_joints = 7;
    offset = [0,0,0,pi/2,0,0,0]

    #          alpha                a           theta       d           Rev=0/Pris=1    theta-offset
    dhpar   = [ DH_alpha0      shoulderOffs     0       shoulderLen         0
                DH_alpha1     -shoulderOffs     0           0               0
                DH_alpha2      elbowOffs        0       upArmLen            0
                DH_alpha3     -elbowOffs        0           0               0
                DH_alpha4      wristOffs        0       lowArmLen           0
                DH_alpha5      handOffs         0           0               0
                DH_alpha6      0                0       handLen             0  ];

    return DH(dhpar,offset,gearRatio)

end

function DH7600()
    DH_alpha= [-pi/2 0 -pi/2 pi/2 -pi/2 0];
    DH_a0= 410;
    DH_a1= 1075;
    DH_a2= 165;

    DH_d0 =780;
    DH_d1 =0;
    DH_d2 =0;
    DH_d3 = 1056;
    DH_d4 =0;
    DH_d5 =250;

    offset = [0, -90, 0, 0, 0, 180]*pi/180;

    dhpar   = [ DH_alpha[1]      DH_a0        0         DH_d0        0
    DH_alpha[2]      DH_a1        0         DH_d1        0
    DH_alpha[3]      DH_a2        0         DH_d2        0
    DH_alpha[4]       0           0         DH_d3        0
    DH_alpha[5]       0           0         DH_d4        0
    DH_alpha[6]       0           0         DH_d5        0 ];

    dhpar[:,[2,4]] /= 1000

    return DH(dhpar,offset)
end

function DHtest()
    DH_alpha= [1.0,2,3,4,5,6]
    DH_a0= 0;
    DH_a1= 0;
    DH_a2= 0;

    DH_d0 =780;
    DH_d1 =10;
    DH_d2 =20;
    DH_d3 = 1056;
    DH_d4 =30;
    DH_d5 =250;

    offset = zeros(6);

    dhpar   = [ DH_alpha[1]      DH_a0        0         DH_d0        0
    DH_alpha[2]      DH_a1        0         DH_d1        0
    DH_alpha[3]      DH_a2        0         DH_d2        0
    DH_alpha[4]       0           0         DH_d3        0
    DH_alpha[5]       0           0         DH_d4        0
    DH_alpha[6]       0           0         DH_d5        0 ];

    return DH(dhpar/1000,offset)
end

function abb2logical!(q::Matrix)
    q[:,:] = q[:,[1, 2, 7, 3, 4, 5, 6]]
end

function logical2abb!(q::Matrix)
    q[:,:] = q[:,[1, 2, 4, 5, 6, 7, 3]]
end

function abb2logical!(q::Vector)
    q[:,:] = q[[1, 2, 7, 3, 4, 5, 6]]
end

function logical2abb!(q::Vector)
    q[:,:] = q[[1, 2, 4, 5, 6, 7, 3]]
end

abb2logical(q::Matrix) = q[:,[1, 2, 7, 3, 4, 5, 6]]
logical2abb(q::Matrix) = q[:,[1, 2, 4, 5, 6, 7, 3]]
abb2logical(q::Vector) = q[[1, 2, 7, 3, 4, 5, 6]]
logical2abb(q::Vector) = q[[1, 2, 4, 5, 6, 7, 3]]


# Base Function Test

struct Euclidean <: Metric end

MLK.base_aggregate(::Euclidean, s::T, x::T, y::T) where {T} = s + (x-y)^2
MLK.base_return(::Euclidean, s::T) where {T} = sqrt(s)

# Base Function References

const base_functions = (
    SquaredEuclidean,
    ScalarProduct,
    Euclidean
)

const base_functions_initiate = Dict(
    SquaredEuclidean => 0,
    ScalarProduct    => 0,
    Euclidean        => 0
)

const base_functions_aggregate = Dict(
    SquaredEuclidean => (s,x,y) -> s + (x-y)^2,
    ScalarProduct    => (s,x,y) -> s + x*y,
    Euclidean        => (s,x,y) -> s + (x-y)^2
)

const base_functions_return = Dict(
    SquaredEuclidean => s -> s,
    ScalarProduct    => s -> s,
    Euclidean        => s -> sqrt(s)
)

const base_functions_properties = Dict(
                       #|stnry |isotrop
    ScalarProduct    => (false, false),
    SquaredEuclidean => (true,  true),
    Euclidean        => (false, false) # just default values test
)

# Kernel Function References

const kernel_functions = (
    ExponentialKernel,
    SquaredExponentialKernel,
    GammaExponentialKernel,
    RationalQuadraticKernel,
    GammaRationalQuadraticKernel,
    MaternKernel,
    PolynomialKernel,
    ExponentiatedKernel,
    PowerKernel,
    LogKernel,
    SigmoidKernel
)

const kernel_functions_arguments = Dict(
    ExponentialKernel            => ((1.0,),        (2.0,)),
    SquaredExponentialKernel     => ((1.0,),        (2.0,)),
    GammaExponentialKernel       => ((1.0,1.0),     (2.0,0.5)),
    RationalQuadraticKernel      => ((1.0,1.0),     (2.0,2.0)),
    GammaRationalQuadraticKernel => ((1.0,1.0,1.0), (2.0,2.0,0.5)),
    MaternKernel                 => ((1.0,1.0),     (2.0,2.0)),
    PolynomialKernel             => ((1.0,1.0,3),   (2.0,2.0,2)),
    ExponentiatedKernel          => ((1.0,),        (2.0,)),
    PowerKernel                  => ((1.0,),        (0.5,)),
    LogKernel                    => ((1.0,1.0),     (2.0,0.5)),
    SigmoidKernel                => ((1.0,1.0),     (2.0,2.0))
)

const kernel_functions_kappa = Dict(
    ExponentialKernel            => (z,α)     -> exp(-α*sqrt(z)),
    SquaredExponentialKernel     => (z,α)     -> exp(-α*z),
    GammaExponentialKernel       => (z,α,γ)   -> exp(-α*z^γ),
    RationalQuadraticKernel      => (z,α,β)   -> (1 + α*z)^(-β),
    GammaRationalQuadraticKernel => (z,α,β,γ) -> (1 + α*z^γ)^(-β),
    MaternKernel                 => (z,ν,ρ)   -> begin
                                                   d = √(z)
                                                   T = typeof(z)
                                                   d = d < eps(T) ? eps(T) : d
                                                   tmp1 = √(2*ν)*d/ρ
                                                   tmp2 = 2^(1 - ν)
                                                   tmp2*(tmp1^ν)*besselk(ν, tmp1)/gamma(ν)
                                                end,
    PolynomialKernel             => (z,a,c,d) -> (a*z+c)^d,
    ExponentiatedKernel          => (z,a)     -> exp(a*z),
    PowerKernel                  => (z,γ)     -> z^γ,
    LogKernel                    => (z,α,γ)   -> log(α*z^γ+1),
    SigmoidKernel                => (z,a,c)   -> tanh(a*z+c)
)

const kernel_functions_base = Dict(
    PolynomialKernel         => ScalarProduct,
    ExponentiatedKernel      => ScalarProduct,
    SigmoidKernel            => ScalarProduct
)


# solve the ground state of Lieb-Liniger model
# also demonstrates the usage of preconditioners in Optim.jl

using TensorKit
using TensorOperations
using TensorKitAD
using KrylovKit
using Zygote
using KrylovKit
using LinearAlgebra
using Plots
using Optim
using ChainRules
using ChainRulesCore
using JLD2
using Printf

using Revise
using statmech_tm_solver

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

# parameters
c = 1
μ = 1
L = 50
d = 1
χ = 16

#########################################################
ψm = quickload("lieb_liniger_c$(c)_mu$(μ)_L$(L)_chi$(χ)") |> convert_to_cmps
ψm = normalize(ψm, L, sym=false)
E = energy_lieb_liniger(ψm, c, L, μ)

K = K_mat(ψm, ψm)
W, _ = eig(K)
wvec = W.data |> diag

env = finite_env(K_mat(ψm, ψm), L)
env = permute(env, (2, 3), (4, 1))
op_ρ = particle_density(ψm)
op_k = kinetic(ψm)
ρ_value = tr(env * op_ρ) |> real
k_value = tr(env * op_k) |> real
γ = c / ρ_value

#plot(angle.(wvec) .- sign.(angle.(wvec)) .* pi, -real.(wvec), seriestype=:scatter)

#sort(real.(wvec))

Ees = []
kmax = L
for k in -kmax:kmax
    p = k * 2 * pi / L
    h_lop = lieb_liniger_h_tangent_map(ψm, p, 0, c, L, μ)
    n_lop = tangent_map(ψm, L, p)

    h_lop_mat = zeros(ComplexF64, d*χ^2, d*χ^2)
    n_lop_mat = zeros(ComplexF64, d*χ^2, d*χ^2)
    Vr = TensorMap(zeros, ComplexF64, ℂ^(χ*d), ℂ^χ)
    indices = reshape(1:d*χ^2, (d*χ, χ))
    for ix in 1:d*χ
        for iy in 1:χ
            global Vr, h_lop_mat
            @printf "k, ix, iy = %2d %2d %2d \r" k ix iy
            Vr.data[ix, iy] = 1
            h_lop_mat[:, indices[ix, iy]] = vec(h_lop(Vr).data)
            n_lop_mat[:, indices[ix, iy]] = vec(n_lop(Vr).data)
            Vr.data[ix, iy] = 0
        end
    end
    @assert h_lop_mat ≈ h_lop_mat'
    @assert n_lop_mat ≈ n_lop_mat'

    excitation_data = Dict("hlop" => h_lop_mat, "nlop" => n_lop_mat)
    save("lieb_liniger_excitation_c$(c)_mu$(μ)_L$(L)_chi$(χ)_k$(k).jld2", excitation_data)

    Ee = eigvals(h_lop_mat * inv(Hermitian(n_lop_mat))) ./ L
    push!(Ees, Ee)
end

#ks = [0]
#Es = [0]
#ΔE = Ees[kmax+1][1] - E
#
#for (ix, Ee) in zip(ixs, Ees)
#    global ps, Es
#    ks = cat(ks, [ix for Ei in Ee]; dims=1)
#    Es = cat(Es, (Ee .- E) / ΔE; dims=1)
#end
#plot(ks, real.(Es), seriestype=:scatter)
#xlims!(-5, 5)
#ylims!(0, 20)
#yticks!(0:20)


import GreyDecision.Utility: makeminmax
import GreyDecision.GreyNumbers: GreyNumber
import GreyDecision.Marcos: marcos

@testset "Marcos" begin
    tol = 0.0001

    weights = [0.19019, 0.15915, 0.19819, 0.19019, 0.15115, 0.11111]

    fns = [maximum, maximum, maximum, maximum, maximum, maximum]
    Fns = convert(Array{Function,1}, fns)

    decmat = [
        8.675 8.433 8.000 7.800 8.025 8.043
        8.825 8.600 7.420 7.463 7.825 8.229
        8.325 7.600 8.040 7.700 7.925 7.600
        8.525 8.667 7.180 7.375 7.750 8.071
    ]

    nrow, ncol = size(decmat)
    # Converting white decision matrix to
    # grey decision matrix
    greydecmat = Array{GreyNumber{Float64}, 2}(undef, nrow, ncol)

    for i in 1:nrow
        for j in 1:ncol
            greydecmat[i, j] = GreyNumber(decmat[i, j])
        end
    end

    result = marcos(greydecmat, weights, Fns)

    @test result[:bestIndex] == 1
    @test isapprox(result[:scores][1].a, 0.684865943528, atol = tol)
    @test isapprox(result[:scores][2].a, 0.672767106696, atol = tol)
    @test isapprox(result[:scores][3].a, 0.662596906139, atol = tol)
    @test isapprox(result[:scores][4].a, 0.661103207660, atol = tol)

end


module Fit

# Simulate a dataset and test functions: loglikelihood!() and fit!()
using MixedModelsBLB, LinearAlgebra, Random, Distributions, Test

Random.seed!(123)
N = 100; p = 4; q = 2; reps = 5
βtrue  = [0.1; 6; -3; 1]
σ²true = 1.5
σtrue  = sqrt(σ²true)
Σtrue  = [1.0 0; 0 3.0] # Matrix(Diagonal([2.0; 1.2; 1.0]))
Ltrue  = Matrix(cholesky(Symmetric(Σtrue)).L)
obsvec = Vector{blblmmObs{Float64}}(undef, N)
X = Matrix{Float64}(undef, reps, p)
X[:, 1] = ones(reps)
Z = Matrix{Float64}(undef, reps, q)
Z[:, 1] = ones(reps)
storage_q = Vector{Float64}(undef, q)
re_storage = Vector{Float64}(undef, q)
y = Vector{Float64}(undef, reps)
fenames = vcat("Intercept", "x" .* string.([1:1:(p-1);]))
renames = ["Intercept", "z1"]

# Generate data
for i in 1:N
    randn!(y) # y = standard normal error
    # first column intercept, remaining entries iid std normal
    @views randn!(X[:, 2:p]) #Distributions.rand!(Normal(), X[:, 2:p])
    BLAS.gemv!('N', 1., X, βtrue, σtrue, y) # y = Xβ + σtrue * standard normal error
    randn!(storage_q)
    BLAS.gemv!('N', 1., Ltrue, storage_q, 0., re_storage)
    # first column intercept, remaining entries iid std normal
    @views randn!(Z[:, 2:q]) #Distributions.rand!(Normal(), Z[:, 2:q])
    BLAS.gemv!('N', 1., Z, re_storage, 1., y) # y = Xβ + Zα + error
    # y = X * βtrue .+ Z * (Ltrue * randn(q)) .+ σtrue * randn(ns[i])
    # form a blblmmObs instance
    obsvec[i] = blblmmObs(copy(y), copy(X), copy(Z))
end

# form a LmmModel instance
lmm = blblmmModel(obsvec, fenames, renames, N)
obsN = lmm.data[N]

# Data was simulated correctly
@test obsN.y[1] ≈ -13.248956453781018 atol = 1e-5

@testset "loglikelihood!()" begin
    @test loglikelihood!(obsN, βtrue, [σ²true], Ltrue, true, true) ≈ -10.753503805549846 atol = 1e-5

    init_ls!(lmm)
    @test loglikelihood!(lmm, true, true) ≈ -1090.2939385079767 atol = 1e-5
end

@testset "fit!()" begin
    fit!(lmm)
    @test loglikelihood!(lmm, true, false) ≈ -994.0524142638102 atol = 1e-5

    @test lmm.β[1] ≈ 0.20516618221633012 atol = 1e-5

    @test norm(lmm.∇β) < 1e-5
    @test norm(lmm.∇σ²) < 1e-5
    @test norm(lmm.∇L) < 1e-5
end

end


using PQEcolaPoint
using Dates

# short leg -
short_put_contract = PutContractModel()
short_put_contract.ticker = "XYZ"
short_put_contract.expiration_date = Date(2022, 04, 14)
short_put_contract.strike_price = 105.0
short_put_contract.premium = 0.0
short_put_contract.number_of_contracts = 1
short_put_contract.direction = 1

# long leg -
long_put_contract = PutContractModel()
long_put_contract.ticker = "XYZ"
long_put_contract.expiration_date = Date(2022, 04, 14)
long_put_contract.strike_price = 100.0
long_put_contract.premium = 0.0
long_put_contract.number_of_contracts = 1
long_put_contract.direction = 1

# contracts -
contracts = Array{PutContractModel,1}()
push!(contracts, short_put_contract)
push!(contracts, long_put_contract)

# compute the theta array -
Sₒ = 113.46
model = build(CRRLatticeModel; Sₒ = Sₒ, number_of_levels=80, σ=0.5212, T=(37.0 / 365), μ=0.0045)
C = premium(contracts, model)
θ_array = θ(contracts; Sₒ = Sₒ, number_of_levels=80, σ=0.5212, T=(37.0 / 365), μ=0.0047)
δ_array = δ(contracts; Sₒ = Sₒ, number_of_levels=80, σ=0.5212, T=(37.0 / 365), μ=0.0047)


function petk03()
#
# M. Petkovic testing polynomials, p134
#
    p1 = [1.0;-1-99*im/70];
    p2 = [1;2;3];
    p2 = conv(p2,p2);
    p = conv(p1,p2);
    p = conv(p,[1;1]);
    z = [-1.00000000000000 + 1.41421356237309im 2;
         -1.00000000000000 - 1.41421356237309im 2;
         -1-99*im/70 1;
         -1 1];
    p, PolyZeros(z)
end


###############################################################################
"""
Provides utilities to transform between physical (x,y,z) and pixel (u,v) coordinate systems.
"""
module ImageRectification
###############################################################################

###############################################################################
using StaticArrays              # provides SMatrix, SVector, for faster linear algebra
using ProgressMeter
###############################################################################

###############################################################################
"""
FindXYFromUVZ(u, v, z, DLT)

    Returns XY coordinates of a point.

    Calculated using the DLT formulation from eqn 3 in Holland et al. z is
    known a-priori (usually the water surface):
"""
function FindXYFromUVZ end
FindXYFromUVZ(u::Number, v::Number, z::Number, DLT::NamedTuple) = uvz2xy(u, v, z, DLT)
FindXYFromUVZ(uv::Array, z::Number, DLT::NamedTuple) = FindXYFromUVZ(view(uv,:,1), view(uv,:,2), z::Number, DLT)
FindXYFromUVZ(uv::Array, z::Array, DLT::NamedTuple) = FindXYFromUVZ(view(uv,:,1), view(uv,:,2), z::Array, DLT)
#
function FindXYFromUVZ(u, v, z::Number, DLT::NamedTuple)
    x = Array{Float64}(undef, size(u))
    y = similar(x)
    for i in eachindex(u); x[i],y[i] = uvz2xy(u[i], v[i], z, DLT); end
    return x,y
end
function FindXYFromUVZ(u, v, z::Array, DLT::NamedTuple)
    x = Array{Float64}(undef, size(u))
    y = similar(x)
    for i in eachindex(u); x[i],y[i] = uvz2xy(u[i], v[i], z[i], DLT); end
    return x,y
end
# Helper functions for FindXYFromUVZ():
Amatrix(u, v, z, DLT::NamedTuple) = SMatrix{2,2}( u*DLT.L₉-DLT.L₁,  v*DLT.L₉-DLT.L₅, u*DLT.L₁₀-DLT.L₂, v*DLT.L₁₀-DLT.L₆ )
bvector(u, v, z, DLT::NamedTuple) = SVector( DLT.L₄-u + (DLT.L₃-u*DLT.L₁₁)*z, DLT.L₈-v + (DLT.L₇-v*DLT.L₁₁)*z )
uvz2xy(u, v, z, DLT) = Amatrix(u, v, z, DLT) \ bvector(u, v, z, DLT)
###############################################################################

###############################################################################
""" Calculate the uv -> xyz coordinate conversion using each of the sets of DLTs: """
function CalculateModifiedxyCoordinates(u,v, WaterSurfaceElev, DLTs;
                                        wseNoiseFactor = 0.0)
    Perturb_wse = randn(size(DLTs,2)) .* wseNoiseFactor
    Modified_x, Modified_y = Rect.FindXYFromUVZ_DLT_z_sensitivity(u,v,
                                                                  # WaterSurfaceElev .+ randn(size(DLTs,2)) .* wseNoiseFactor,
                                                                  WaterSurfaceElev .+ Perturb_wse,
                                                                  DLTs)
    return Modified_x, Modified_y, Perturb_wse
end

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

###############################################################################
# This option is for sensitivity analysis of water surface elevation (and a single DLT set):
function FindXYFromUVZ_WSE_sensitivity(u::Array, v::Array, z::Array, DLT::NamedTuple)
    x = Array{Float64}(undef, size(u)..., length(z))
    y = Array{Float64}(undef, size(v)..., length(z))
    @showprogress 1 "Calculating $(length(u)) xy coordinates using $(size(DLT,2)) DLT sets:" for d in 1:length(z)
        for j in 1:size(u,2), i in 1:size(u,1)
            x[i,j,d],y[i,j,d] = FindXYFromUVZ(u[i,j], v[i,j], z[d], DLT)
        end
    end
    return x,y
end
# This option is for sensitivity analysis of multiple DLT sets:
function FindXYFromUVZ_DLT_sensitivity(u::Array, v::Array, z::Float64, DLTs::Array)
    x = Array{Float64}(undef, size(u)..., size(DLT,2))
    y = Array{Float64}(undef, size(v)..., size(DLT,2))
    @showprogress 1 "Calculating $(length(u)) xy coordinates using $(size(DLT,2)) DLT sets:" for d in 1:size(DLT,2)
        for j in 1:size(u,2), i in 1:size(u,1)
            x[i,j,d],y[i,j,d] = FindXYFromUVZ(u[i,j], v[i,j], z, DLT[d])
        end
    end
    return x,y
end
# This option is for sensitivity analysis of multiple DLT sets, along with
# multiple z values, for analysis of sensitivity to, e.g., water surface
# elevation:
function FindXYFromUVZ_DLT_z_sensitivity(u::Array, v::Array, z::Array, DLT::Array)
    (length(z) != size(DLT,2)) ? @error("lengths of z and DLT do not match!") : nothing
    x = Array{Float64}(undef, size(u)..., size(DLT,2))
    y = Array{Float64}(undef, size(v)..., size(DLT,2))
    @showprogress 1 "Calculating $(length(u)) xy coordinates using $(size(DLT,2)) DLT sets:" for d in 1:size(DLT,2) for j in 1:size(u,2), i in 1:size(u,1)
            x[i,j,d],y[i,j,d] = FindXYFromUVZ(u[i,j], v[i,j], z[d], DLT[d])
        end
    end
    return x,y
end
###############################################################################

###############################################################################
###############################################################################
###############################################################################
end                             # end of module
###############################################################################
###############################################################################
###############################################################################


include("../src/strip2.jl")
#using strip2
using DelimitedFiles
#using Plots
#gr()
# Strip2: 54deg x 54deg
# General constants
NN = 4024
pix_size = 0.5
beam_waist = 11.0
# White noise: evaluation of σ

dc = 0.35
observation = 24 # months
cmb = 1.1 # at 95 GHz
system = 8.5
atmosphere = 25 #102
bandwidth  = 31.5 # GHz
Tsys = cmb + system + atmosphere
Number_of_sensors = 1000.0
time_per_pixel = (observation * 30 * 24 * 60 * 60 * dc) / (NN^2)  # sec/pix

# Ma osservo 10 volte
time_per_pixel /=  10.0

sigma = ((Tsys / sqrt(31.5*1E9 * time_per_pixel )) / sqrt(Number_of_sensors)) * 1e6

# SZ Effect
N_SZ_Clusters  = 100
A_SZ_Clusters = 50.0
SZ_beta = 0.86
SZ_Theta_core = 1.0

# Point Sources - poisson distribution
N_Sources_POIS = 5000
A_Sources_POIS = 200.0

# Point Sources - exponential distrubution
N_Sources_EXP = 50
A_Sources_EXP = 1000.0

# Read the CAMB cmb power spectrum.
m = readdlm("power_spectrum/CAMB_totalCls.dat")

ell  = m[:, 1]
DlTT = m[:, 2]
DlEE = m[:, 3]
DlBB = m[:, 4]
DlTE = m[:, 5]

Tm, Qm, Um, Em, Bm = strip2.make_CMB_pol_maps(NN, pix_size, ell,
                                              DlTT, DlEE, DlTE, DlBB)
# Add the SZ contribute
sz_map = strip2.SZ_source(NN, pix_size, N_SZ_Clusters,
                   A_SZ_Clusters, SZ_beta,
                   SZ_Theta_core)

# Add the point sources contribute
point_pois = strip2.PS_poisson(NN, pix_size, N_Sources_POIS, A_Sources_POIS)
point_expo = strip2.PS_exp(NN, pix_size, N_Sources_EXP, A_Sources_EXP)

# Overall map (cosmological sources and noise) convolved with the
# telescope's beam
# mapp = Tm .+ sz_map .+ point_pois .+ point_expo
mapp = sz_map
conv_map = strip2.convolve_beam(NN, pix_size, beam_waist, mapp)

gradient=ColorGradient([:blue, :white, :red])
p2 = strip2.Plot_Sky(NN, pix_size, conv_map, cmin = -400, cmax = 400, lab_title="T Map")


struct GaussProcess{A<:Array, D, F0, F1, F2, F3}
    γ::A
    β::D
    μ::F0
    ∇ₓμ::F1
    Σ::F2
    ∇ₓΣ::F3
end

function GaussProcess(γ::Array, β::Real, Σ, ∇ₓΣ)
    μ(x) = zeros(size(x, 2))
    ∇ₓμ(x) = zeros(size(x))
    return GaussProcess(γ, β, μ, ∇ₓμ, Σ, ∇ₓΣ)
end

function negloglik(GP::GaussProcess, X̄::Matrix, ȳ::Vector)
    K = GP.Σ(relu.(GP.γ), X̄, X̄)
    K += K'
    K += (GP.β^2)*I
    μ = GP.μ(X̄)
    likelihood = MvNormal(μ, K)
    nll = -1*logpdf(likelihood, ȳ)
    return nll
end

function predict(GP::GaussProcess, X̄::Matrix, ȳ::Vector, x::Matrix)
    K = GP.Σ(GP.γ, X̄, X̄) + (GP.β^2)*I
    k = GP.Σ(GP.γ, x, X̄)
    reduced_ȳ = ȳ - GP.μ(X̄)

    Cho_K = cholesky(K)
    μ = GP.μ(x) + k*(Cho_K\reduced_ȳ)
    Σp = GP.Σ(GP.γ, x, x) - k*(Cho_K\(k'))

    σ = sqrt.(diag(Σp))
    return μ, σ
end

function ∇ₓpredict(GP::GaussProcess, X̄::Matrix, ȳ::Vector, x::Matrix)
    F, N₁, N₂ = size(X̄, 1), size(X̄, 2), size(x, 2)
    K = GP.Σ(relu.(GP.γ), X̄, X̄)
    K += K'
    K += (GP.β^2)*I
    ∇ₓk = GP.∇ₓΣ(relu.(GP.γ), x, X̄)
    ∇ₓμ = convert.(eltype(GP.γ), GP.∇ₓμ(x))
    reduced_ȳ = ȳ - GP.μ(X̄)

    Cho_K = cholesky(K)
    α = Cho_K\reduced_ȳ
    for j in 1:N₁
        for i in 1:N₂
            for l in 1:F
                ∇ₓμ[l,i] += ∇ₓk[l,i,j]*α[j]
            end
        end
    end
    return ∇ₓμ
end


using MinFEM

function parabolic(;T::Float64, tsteps::Int, theta=1.0)
  mesh = unit_square(100)

  boundary = union(mesh.Boundaries[1001].Nodes,
                   mesh.Boundaries[1002].Nodes,
                   mesh.Boundaries[1003].Nodes,
                   mesh.Boundaries[1004].Nodes)

  L = asmLaplacian(mesh)
  M = asmMassMatrix(mesh)

  # time increment assuming that t0=0
  dt = (T-0)/tsteps

  source(x) = 1.0
  f = evaluateMeshFunction(mesh, source)

  initial_condition(x) = 0.0
  u = evaluateMeshFunction(mesh, initial_condition)

  for i=1:tsteps

    # in first timestep additional output of u0
    if i==1
      vtkfile = open_vtk_file(mesh, "parabolic_"*lpad(string(0), 3, '0')*".vtu")
      write_point_data(vtkfile, u, "u")
      save_vtk_file(vtkfile)
    end

    # we have to solve the following equation depending on theta
    # with the prescribed Dirichlet conditions
    # M*(u_new - u)/dt + L(theta*u_new + (1.0-theta)*u) = M*f
    pde = PDESystem(A=(M + theta*dt*L), b=(M - (1.0-theta)*dt*L)*u + dt*M*f, bc=zeros(mesh.nnodes), DI=boundary)
    solve(pde)

    u = copy(pde.state)

    vtkfile = open_vtk_file(mesh, "parabolic_"*lpad(string(i), 3, '0')*".vtu")
    write_point_data(vtkfile, u, "u")
    save_vtk_file(vtkfile)
  end

end
parabolic(T=0.001, tsteps=10, theta=1.0)