function test_fltype()
  Tf = Float16
  Tc = Complex{Tf}
  Tf == fltype(Tf) && Tf == fltype(Tc) && Tf == fltype((Tc, Tf)) && Tf == fltype((Tf, Tc))
end

function test_cxtype()
  Tf = Float16
  Tc = Complex{Tf}
  Tc == cxtype(Tf) && Tc == cxtype(Tc)
end

function test_innereltype(T=Complex{Float32})
  a = [zeros(T, 3), zeros(T, 2, 5)]
  T == innereltype(a)
end

function test_superzeros()
  T = Float64
  dims1 = (2, 3)
  dims2a = (7,)
  dims2b = dims1
  dims2 = (dims2a, dims2b)

  a1 = rand(T, dims1)
  a2 = [(1+im)*rand(T, dims2a), (2+im)*rand(T, dims2b)]

  a1a = zeros(T, dims1)
  a2a = [zeros(Complex{T}, dims2a), zeros(Complex{T}, dims2b)]

  a1z = superzeros(a1)
  a2z = superzeros(a2)
  a1d = superzeros(dims1)
  a2d = superzeros(Complex{T}, dims2)

  @superzeros innereltype(a1) a1 a1ma dum
  @superzeros innereltype(a2) a2 a2ma dum
  @superzeros innereltype(a1) dims1 a1md dum
  @superzeros innereltype(a2) dims2 a2md dum

  ( a1a == a1z && a1a == a1d && a1a == a1ma && a1a == a1md &&
    a2a == a2z && a2a == a2d && a2a == a2ma && a2a == a2md )
end

function test_supertuplezeros(; T1=Float64, T2=Complex{Float64}, dims1=(1,), dims2=(3, 3))
  T = (T1, T2)
  dims = (dims1, dims2)
  a = superzeros(T, dims)
  ( eltype(a[1]) == T1 && eltype(a[2]) == T2 &&
    size(a[1]) == dims1 && size(a[2]) == dims2 )
end

# This test could use some further work.
function test_radialspectrum(dev::Device, n, ahkl, ahρ; debug=false, atol=0.1, rfft=false)
  grid = TwoDGrid(dev, n, 2π)
  if rfft==true
    ah = @. ahkl(grid.kr, grid.l)
  else
    ah = @. ahkl(grid.k, grid.l)
  end
  CUDA.@allowscalar ah[1, 1] = 0.0

  ρ, ahρ_estimate = FourierFlows.radialspectrum(ah, grid; refinement=16)

  if debug
    println(sum(ahρ.(ρ) - ahρ_estimate))
    return ah, ahρ_estimate, ahρ.(ρ)
  else
    normalizeddiff = sum(abs.(ahρ.(ρ) - ahρ_estimate)) / length(ρ)
    return isapprox(normalizeddiff, 0.0, atol=atol)
  end
end

function integralsquare(func, grid::OneDGrid)
    sum(abs2.(func))*grid.dx
end

function integralsquare(func, grid::TwoDGrid)
    sum(abs2.(func))*grid.dx*grid.dy
end

function test_parsevalsum(func, grid; realvalued=true)
  # Compute integral in physical space
  integral = integralsquare(func, grid)
  # Compute integral in wavenumber space using Parseval's theorem
  if realvalued==true;      funch = rfft(func)
  elseif realvalued==false; funch = fft(func)
  end
  parsevalsum = FourierFlows.parsevalsum(abs2.(funch), grid)
  isapprox(integral, parsevalsum; rtol=rtol_utils)
end

function test_parsevalsum2(func, grid; realvalued=true)
  # Compute integral in physical space
  integral = integralsquare(func, grid)
  # Compute integral in wavenumber space using Parseval's theorem
  if realvalued==true;      funch = rfft(func)
  elseif realvalued==false; funch = fft(func)
  end
  parsevalsum2 = FourierFlows.parsevalsum2(funch, grid)
  isapprox(integral, parsevalsum2; rtol=rtol_utils)
end

"""
Compute the jacobian J(a, b) and compare the result with `analytic`. Use `atol` for the comparison
to ensure validity when `analytic=0`.
"""
test_jacobian(a, b, analytic, grid) = isapprox(FourierFlows.jacobian(a, b, grid), analytic;
                                               atol=grid.nx*grid.ny*10*eps())

"""
Test the zeros macro.
"""
function test_zeros(T=Float64, dims=(13, 45))
  a1, b1 = zeros(T, dims), zeros(T, dims)
  @zeros T dims a2 b2
  a1 == a2 && b1 == b2
end

"""
Test supersize().
"""
function test_supersize()
  a = rand(16, 16)
  dimsa = size(a)

  b = [rand(1), rand(3, 34)]
  dimsb = ((1,), (3, 34))

  dimsa == supersize(a) && dimsb == supersize(b)
end

"""
Test arraytype().
"""
function test_arraytype(dev::Device)
  dev==CPU() ? ArrayType(dev)==Array : ArrayType(dev)==CuArray
end

function test_arraytypeTdim(dev::Device, T=Float64, dim=1)
  dev==CPU() ? ArrayType(dev, T, dim)<:Array{T, dim} : ArrayType(dev, T, dim)<:CuArray{T, dim}
end

function test_ongrid(dev::Device)
  nx, ny, nz = 6, 8, 10
  Lx, Ly, Lz = 2π, 2.0, 3.0

  g₁ = OneDGrid(dev, nx, Lx)
  X₁ = ArrayType(dev)(g₁.x)
  f₁(x) = x^2

  g₂ = TwoDGrid(dev, nx, Lx, ny, Ly)
  X₂, Y₂ = gridpoints(g₂)
  f₂(x, y) = x^2 - y^3

  g₃ = ThreeDGrid(dev, nx, Lx, ny, Ly, nz, Lz)
  X₃, Y₃, Z₃ = gridpoints(g₃)
  f₃(x, y, z) = x^2 - y^3 + sin(z)

  return (FourierFlows.on_grid(f₁, g₁) ≈ f₁.(X₁) && FourierFlows.on_grid(f₂, g₂) ≈ f₂.(X₂, Y₂) && FourierFlows.on_grid(f₃, g₃) ≈ f₃.(X₃, Y₃, Z₃))
end


@testset "argmin_x" begin

#test argmin_x.jl for CSC sparse matrix format
  A = sprandn(100,100,0.01)+sparse(I, 100, 100);
  A = A'*A;
  while rank(Array(A))<100
    A = A + sparse(I, 100, 100)
  end
  xt = randn(100)
  b  = A*xt;

  maxit = 100;
  pp    = 2
  p     = 2
  TF    = Float64
  x_solve_tol_ref = 1e-5

  log_PARSDMM = log_type_PARSDMM(zeros(maxit,pp),zeros(maxit,p),zeros(maxit,p),zeros(maxit),zeros(maxit),zeros(maxit),
  zeros(maxit),zeros(maxit,p),zeros(maxit,p),zeros(maxit),zeros(maxit),TimerOutput());

  #test for zero initial guess
    (x,iter,relres)=argmin_x(A,b,zeros(100),x_solve_tol_ref,5,log_PARSDMM)
    @test norm(A*x-b)/norm(b) <= x_solve_tol_ref
    @test isapprox(norm(A*x-b)/norm(b),relres,rtol=x_solve_tol_ref)

  #test for good initial guess
    (x2,iter2,relres2)=argmin_x(A,b,xt+randn(100).*1e-4,1e-5,5,log_PARSDMM)
    @test iter2<iter
    @test norm(A*x2-b)/norm(b) <= 1e-5
    @test isapprox(norm(A*x2-b)/norm(b),relres2,rtol=1e-5)

  #test if 10*eps(Float32/Float64) is achieved
    (x,iter,relres)=argmin_x(A,b,zeros(100),10*eps(),15,log_PARSDMM)
    @test norm(A*x-b)/norm(b) <= 20*eps()

#test argmin_x.jl for CDS sparse matrix format
  A=sprandn(100,100,0.01)+sparse(I, 100, 100);
  A=A'*A;
  while rank(Array(A))<100
    A=A+sparse(I, 100, 100)
  end
  xt=randn(100)
  b=A*xt;

  (R,offset) = mat2CDS(A)

  y=zeros(Float64,length(b))
  (x,iter,relres)=argmin_x(A,b,zeros(100),1e-5,5,log_PARSDMM,offset,y)

  @test norm(A*x-b)/norm(b) <= 1e-5
  @test norm(A*x-b)/norm(b) <= 2.0*relres

#test for more accurate solutions with CDS
  y=zeros(Float64,length(b))
  (x,iter,relres)=argmin_x(A,b,zeros(100),1e-10,5,log_PARSDMM,offset,y)

  @test norm(A*x-b)/norm(b) <= 1e-10
  @test norm(A*x-b)/norm(b) <= 2.0*relres

end


export ProductFunction

struct ProductFunction{T, N} <: AbstractMathFun{T, N}
    terms::MFCollection

    ProductFunction{T, N}() where {T, N} = new{T, N}(MFCollection(undef, 0))
    ProductFunction{T, N}(terms::MFCollection) where {T, N} = new{T, N}(terms)

    Base.insert!(p::ProductFunction, i::Integer, term::AbstractMathFun) = insert!(p.terms, i, term)
    Base.push!(p::ProductFunction, term::AbstractMathFun) = push!(p.terms, term)
    Base.deleteat!(p::ProductFunction, range) = deleteat!(p.terms, range)
    Base.length(p::ProductFunction) = length(p.terms)
    Base.getindex(p::ProductFunction, i::Integer) = getindex(p.terms, i)
    Base.setindex!(p::ProductFunction, val::AbstractMathFun, i::Integer) = setindex!(p.terms, val, i)
    Base.iterate(p::ProductFunction) = Base.iterate(p.terms)
    Base.iterate(p::ProductFunction, state) = Base.iterate(p.terms, state)

    function Functions.string(p::ProductFunction, emitter::FunctionEmitter)
        str = "∏("
        first_print = true
        for term in p.terms
            term_str = string(term, emitter)
            length(term_str) != 0 && (str *= (first_print ? term_str : ", " * term_str); first_print = false)
        end
        str * ")"
    end

    function (p::ProductFunction)(call::MFCall)
        product = 1.0
        for term in p.terms
            product *= term(call::MFCall)
        end
        call(product)
    end

    function Functions.emit(f::ProductFunction{T, N}, emitter::FunctionEmitter{T}) where {T, N}
        prod_expr = Expr(:call, :∏)
        for term in f.terms
            push!(prod_expr.args, emit(term, emitter))
        end
        push!(prod_expr.args, Expr(:kw, :T, Symbol(string(T))))
        prod_expr
    end

    function Functions.simplify(f::ProductFunction{T, N}, emit_constants::Bool) where {T, N}
        constprod = 1
        termcpy = MFCollection(undef, 0)

        for item in f.terms
            item = simplify(item, emit_constants)
            if is_constant(item, emit_constants)
                constprod *= item.value
            else
                if item isa ProductFunction
                    append!(termcpy, item.terms)
                else
                    push!(termcpy, item)
                end
            end
        end

        (constprod != 1 && constprod != 0) && push!(termcpy,  ConstFunction{T}(constprod))
        (length(termcpy) == 0) && return ConstFunction{T}(constprod)
        length(termcpy) == 1 && return termcpy[1]
        return ProductFunction{T, N}(termcpy)
    end
end


# solve the KS function and store data for plotting

include("kuramoto-sivashinsky.jl")
using kuramoto_sivashinsky

# set the parameters
tp = Float64            # floating point type to use
num_nodes = 127         # number of (interior) nodes to use
num_steps = 5000        # number of time steps to use
Lx = convert(tp, 128)   # domain size
order = 16              # order of accuracy
Time = convert(tp, 500) # time period for simulation

# set initial condition and set-up the KSData structure
dx = Lx/num_nodes
x = Array{tp}(0.0:dx:Lx-dx)
u = exp.(-((x-64).^2)./512)
ks = kuramoto_sivashinsky.buildKSDataPeriodic(order, num_nodes, tp)

# storage for solution over all time steps; for plotting
dt = Time/num_steps
num_steps = round(Int, Time/dt, RoundUp)
sol = zeros(num_nodes+1, num_steps+1)
sol[1:end-1,1] = u
sol[end,1] = u[1]

# solve using midpoint; don't do this in practice!!!
t = convert(tp, 0.0)
Jac = zeros(ks.linear_op)
A = zeros(ks.linear_op)
u_mid = zeros(u)
u_old = zeros(u)
r = zeros(u)
newt_tol = 1e-10 # We will need to ensure this is sufficiently small
max_newt_iter = 50
for n = 1:num_steps
    println("Iteration ",n,": time ",t)

    # Newton iterations
    u_old = u
    for k = 1:max_newt_iter
        u_mid = 0.5*(u + u_old)
        kuramoto_sivashinsky.getKSFuncPeriodic!(ks, t, u_mid, r)
        r = u - u_old - dt*r
        println("\tNewton iter ",k,": res norm = ",norm(r))
        if norm(r) < newt_tol break end
        kuramoto_sivashinsky.getKSJacPeriodic!(ks, t, u_mid, Jac)
        Jac *= 0.5*dt
        Jac -= speye(num_nodes)
        du = Jac\r
        u += du
    end

    t += dt
    sol[1:end-1,n+1] = u
    sol[end,n+1] = u[1]
end

# functionals; here I use only trapezoidal for the temporal integration
u_avg = zero(tp)
u2_avg = zero(tp)
for n = 1:num_steps+1
    wt = dx*dt
    if n == 1 || n == num_steps+1 wt *= 0.5 end
    for val in sol[1:end-1,n]
        u_avg += val*wt
        u2_avg += val*val*wt
    end
end
u_avg /= (Time*Lx)
u2_avg /= (Time*Lx)
println("u averaged = ", u_avg)
println("u^2 averaged = ", u2_avg)

# output to file for plotting
file = open("ks_solution.dat", "w")
for n = 1:num_steps+1
    for val in sol[:,n]
        print(file, val, " ")
    end
    println(file)
end
close(file)


"""
    get_powers(x, y, powers)

Obtain a DataFrame of the powers and cross products of x and y

# Arguments

- x = 1d array of floats
- y = 1d array of floats of same length as `x`
- powers = 1d array of integers indicating the powers to include. Defaults to
  [1, 2, 3, 4] to get everything up to the quadric term

# Value

DataFrame of the powers and cross products of `x` and `y`, as specified by the
powers argument.

"""
function get_powers(x::Array{Float64,1}, y::Array{Float64,1},
                    powers::Array{Int64,1}=[1, 2, 3, 4])

    # Check that x and y are the same length
    @assert length(x) == length(y)

    # Initialze DataFrame to return
    out = DataFrame()

    # Get powers of x and y
    for i in 1:length(powers)
        out[!, Symbol(string("X", powers[i]))] = x.^powers[i]
            out[!, Symbol(string("Y", powers[i]))] = y.^powers[i]
    end

    # Get cross products
    for i in 1:(length(powers)-1)
        for j in 1:(length(powers)-i)
            out[!, Symbol(string("X", powers[i], "Y", powers[j]))] =
                out[:, Symbol(string("X", powers[i]))] .*
                out[:, Symbol(string("Y", powers[j]))]
        end
    end

    return out
end


using Scat
using BenchmarkTools
using Plots

function time_subsamplingSize(Qs::Vector{Int64}, N::Int64; nTrials=1)
    # define the scattering transform
    # first layer parameters and filter bank
    J=8
    σ0=0.1

    times_subsamp = zeros(N)
    times_nosubsamp = zeros(N)
    for i=1:nTrials
        for n=1:N
            println(n)
            x=ones(2^n)
            Transform = ScatteringTransform1d(2, Qs, [J,J])
            FilterBanks = get_FilterBanks(n, Qs, J, σ0)
            times_subsamp[n] += @elapsed Scat.scatter1d(x, Transform, FilterBanks, subsample=true)
            times_nosubsamp[n] += @elapsed Scat.scatter1d(x, Transform, FilterBanks, subsample=false)
        end
    end
    times_subsamp./=nTrials
    times_nosubsamp./=nTrials

    return times_subsamp, times_nosubsamp
end
times_ss, times_nss = time_subsamplingSize([12, 1], 18, nTrials=5)
scatter(collect(1:18).+log2(64), times_nss, yscale=:log10, label="No subsampling", legend=:topleft)
scatter!(collect(1:18).+log2(64), times_ss, yscale=:log10, label="With subsampling")
xlabel!("log2(Input size [bits])")
ylabel!("Time [s]")
savefig("time_subsampling_size.png")
mean(times_nss./times_ss)

function time_scatteringTransform_Q(Q::Array{Int64,2}, N::Int64)
    # define the scattering transform
    # first layer parameters and filter bank
    J=8
    σ0=0.1
    x=ones(2^N)

    times = zeros(size(Q,1))
    nFilters = zeros(size(Q,1))
    for i=1:size(Q,1)
        FilterBanks = get_FilterBanks(N, Q[i,:], J, σ0)
        Transform = ScatteringTransform1d(2, Q[i,:], [J,J])
        println(length(FilterBanks[1].Λ) * length(FilterBanks[2].Λ))
        nFilters[i] = length(FilterBanks[1].Λ)*length(FilterBanks[2].Λ)
        times[i] = @elapsed Scat.scatter1d(x, Transform, FilterBanks)
    end

    return nFilters, times
end

# vary the number of filters in the first layer
nQ=100
Q = ones(nQ, 2)
#Q[:, 1] .= Q[:, 1].*collect(1:100)
Q[:,1].= rand(1:12,nQ)
Q[:,2].=rand(1:4,nQ)
Q = convert.(Int, Q)
N=12
nFilters, times = time_scatteringTransform_Q(Q, N)
scatter(Q[:,1].*Q[:,2], times, legend=:none)
xlabel!("Q1 ⋅ Q2")
ylabel!("Time [s]")

scatter(nFilters, times, legend=:none)
xlabel!("Total Number of Filters")
ylabel!("Time [s]")
savefig("nFilters_times.png")

function time_subsamplingQ(Q::Array{Int64,2}, N::Int64)
    # define the scattering transform
    # first layer parameters and filter bank
    J=8
    σ0=0.1
    x=ones(2^N)

    times_ss = zeros(size(Q,1))
    times_nss = zeros(size(Q,1))
    nFilters = zeros(size(Q,1))
    for i=1:size(Q,1)
        FilterBanks = get_FilterBanks(N, Q[i,:], J, σ0)
        Transform = ScatteringTransform1d(2, Q[i,:], [J,J])
        println(length(FilterBanks[1].Λ) * length(FilterBanks[2].Λ))
        nFilters[i] = length(FilterBanks[1].Λ)*length(FilterBanks[2].Λ)
        times_ss[i] = @elapsed Scat.scatter1d(x, Transform, FilterBanks)
        times_nss[i] = @elapsed Scat.scatter1d(x, Transform, FilterBanks, subsample=false)
    end

    return nFilters, times_ss, times_nss
end
# vary the number of filters in the first layer
nQ=100
Q = ones(nQ, 2)
#Q[:, 1] .= Q[:, 1].*collect(1:100)
Q[:,1].= rand(1:12,nQ)
Q[:,2].=rand(1:4,nQ)
Q = convert.(Int, Q)
N=12
nFilters, times_ss, times_nss = time_subsamplingQ(Q, N)
scatter(nFilters, times_nss, label="No subsampling")
scatter!(nFilters, times_ss, legend=:topleft, label="With subsampling")
xlabel!("Q1 ⋅ Q2")
ylabel!("Time [s]")
savefig("time_subsampling_nFilters.png")

function time_scatteringTransform_J(J::Vector{Int64}, N::Int64)
    # define the scattering transform
    # first layer parameters and filter bank
    σ0=0.1
    Q=[12, 2]

    times = zeros(length(J))
    for i=1:length(J)
        println(J[i])
        x=ones(2^N)
        FilterBanks = get_FilterBanks(N, Q, J[i], σ0)
        Transform = ScatteringTransform1d(2, Q, [J[i], J[i]])
        times[i] = @elapsed Scat.scatter1d(x, Transform, FilterBanks)
    end

    return times
end
N=12
J = collect(4:20)
times = time_scatteringTransform_J(J, N)
scatter(J, sqrt.(times), legend=:none)
xlabel!("log2(maximum scale)")
ylabel!("√Time")
savefig("maxScale_time.png")

function time_subsamplingJ(J::Vector{Int64}, N::Int64; nTrials=1)
    # define the scattering transform
    # first layer parameters and filter bank
    σ0=0.1
    Q=[12, 2]

    times = zeros(length(J))
    for n=1:nTrials
        for i=1:length(J)
            println(J[i])
            x=ones(2^N)
            FilterBanks = get_FilterBanks(N, Q, J[i], σ0)
            Transform = ScatteringTransform1d(2, Q, [J[i], J[i]])
            times[i] += @elapsed Scat.scatter1d(x, Transform, FilterBanks)
        end
    end

    return times./nTrials
end
N=12
J = collect(4:16)
times = time_subsamplingJ(J, N, nTrials=15)
scatter(J, sqrt.(times), legend=:none)
xlabel!("log2(maximum scale)")
ylabel!("√Time")
savefig("maxScale_time.png")

function time_scatteringTransform_BreadthDepth(N)
    # define the scattering transform
    J=[8,8]
    σ0=0.1
    Q=[12,6]

    times_depth = zeros(N)
    times_breadth = zeros(N)
    for n=1:N
        println(n)
        x=ones(2^n)
        FilterBanks = get_FilterBanks(n, Q, J, σ0)
        FilterBankBlocks = get_FilterBanks(n, Q, J, σ0, typ=FilterBank1dBlock)
        Transform = ScatteringTransform1d(2, Q, J)
        times_depth[n] = @elapsed Scat.scatter1d(x, Transform, FilterBanks)
        times_breadth[n] = @elapsed Scat.scatter1d(x, Transform, FilterBankBlocks)
    end

    return times_depth, times_breadth
end
td, tb = time_scatteringTransform_BreadthDepth(14)
scatter(collect(1:14).+log2(64), td, yscale=:log10, label="Depth-first", legend=:topleft)
scatter!(collect(1:14).+log2(64), tb, yscale=:log10, label="Breadth-first")
xlabel!("log2(Input size) [bits]")
ylabel!("Time [s]")
savefig("depth_vs_breadth_time.png")


# Test: ODESystem built with ModelingToolkit + Ensemble Simulation with DifferentialEquations + Optimisation with DiffEqFlux
using DiffEqFlux, DifferentialEquations, Plots, Statistics, ModelingToolkit
ann = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
dim_p = DiffEqFlux.paramlength(ann)

@variables t x1(t) x2(t)
@parameters p_NN[1:dim_p]
D = Differential(t)

deqs = [D(x1) ~ x2
        D(x2) ~ ann([t],p_NN)[1]^3]
sys = ODESystem(deqs, t, [x1, x2], p_NN; name = :sys)
# sys = ODESystem(deqs, iv, dvs, ps; )
# sys_simplified = structural_simplify(sys) # ---> not working

u0 = [x1 => -4f0,
      x2 => 0f0]
p_NN_0 = initial_params(ann)
p = [p_NN[i] => p_NN_0[i] for i in 1:dim_p]
tspan = (0.0f0, 8.0f0)
prob = ODEProblem(sys, u0, tspan, p; jac = true)
@time sol = solve(prob, Tsit5(), saveat = 0.1f0) # ---> not working

##--------
using DiffEqFlux, DifferentialEquations, Plots, Statistics, ModelingToolkit
dim_h = 32
tspan = (0.0f0,8.0f0)
ann = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
θ = initial_params(ann)
function dxdt_(dx,x,p,t)
    x1, x2 = x
    dx[1] = x[2]
    dx[2] = ann([t],p)[1]^3
end
x0 = [-4f0, 0f0]
prob = ODEProblem(dxdt_,x0,tspan,θ)
sys = structural_simplify(modelingtoolkitize(prob))
prob_mtk = ODEProblem(sys, [], tspan, jac = true)

@time sol = solve(prob_mtk, Tsit5(), saveat = 0.1f0) # ---> working

#------------
ts = Float32.(collect(0.0:0.01:tspan[2]))
function predict_adjoint(θ)
  Array(solve(prob_mtk,Tsit5(),p=θ,saveat=ts,sensealg=InterpolatingAdjoint(autojacvec=ZygoteVJP())))
end
function loss_adjoint(θ)
  x = predict_adjoint(θ)
  mean(abs2,4.0 .- x[1,:]) + 2mean(abs2,x[2,:]) + mean(abs2,[first(ann([t],θ)) for t in ts])/10
end
l = loss_adjoint(θ)
cb = function (θ,l)
  println(l)
  # p = plot(solve(remake(prob,p=θ),Tsit5(),saveat=0.01),ylim=(-6,6),lw=3)
  # plot!(p,ts,[first(ann([t],θ)) for t in ts],label="u(t)",lw=3)
  # display(p)
  return false
end
# Display the ODE with the current parameter values.
cb(θ,l)
loss1 = loss_adjoint(θ)
res1 = DiffEqFlux.sciml_train(loss_adjoint, θ, ADAM(0.005), cb = cb,maxiters=100)
res2 = DiffEqFlux.sciml_train(loss_adjoint, res1.u,
                              BFGS(initial_stepnorm=0.01), cb = cb,maxiters=100,
                              allow_f_increases = false)

## ------------------------------------------------------------------------------------
using ModelingToolkit

function decay(;name)
  @parameters t a
  @variables x(t) f(t)
  D = Differential(t)
  ODESystem([
      D(x) ~ -a*x + f
    ];
    name=name)
end

@named decay1 = decay()
@named decay2 = decay()

@parameters t
D = Differential(t)
connected = compose(ODESystem([
                        decay2.f ~ decay1.x
                        D(decay1.f) ~ 0
                      ], t; name=:connected), decay1, decay2)

equations(connected)

#4-element Vector{Equation}:
# Differential(t)(decay1₊f(t)) ~ 0
# decay2₊f(t) ~ decay1₊x(t)
# Differential(t)(decay1₊x(t)) ~ decay1₊f(t) - (decay1₊a*(decay1₊x(t)))
# Differential(t)(decay2₊x(t)) ~ decay2₊f(t) - (decay2₊a*(decay2₊x(t)))

simplified_sys = structural_simplify(connected)

equations(simplified_sys)

#3-element Vector{Equation}:
# Differential(t)(decay1₊f(t)) ~ 0
# Differential(t)(decay1₊x(t)) ~ decay1₊f(t) - (decay1₊a*(decay1₊x(t)))
# Differential(t)(decay2₊x(t)) ~ decay1₊x(t) - (decay2₊a*(decay2₊x(t)))

x0 = [
  decay1.x => 1.0
  decay1.f => 0.0
  decay2.x => 1.0
]
p = [
  decay1.a => 0.1
  decay2.a => 0.2
]

using DifferentialEquations
prob = ODEProblem(simplified_sys, x0, (0.0, 100.0), p)
sol = solve(prob, Tsit5())
sol[decay2.f]

## ------------------------------------------------------------------------------------
using DifferentialEquations, DiffEqFlux
pa = [1.0]
u0 = [3.0]
θ = [u0;pa]

function model1(θ,ensemble)
  prob = ODEProblem((u, p, t) -> 1.01u .* p, [θ[1]], (0.0, 1.0), [θ[2]])

  function prob_func(prob, i, repeat)
    remake(prob, u0 = 0.5 .+ i/100 .* prob.u0)
  end

  ensemble_prob = EnsembleProblem(prob, prob_func = prob_func)
  sim = solve(ensemble_prob, Tsit5(), ensemble, saveat = 0.1, trajectories = 100)
end

# loss function
loss_serial(θ)   = sum(abs2,1.0.-Array(model1(θ,EnsembleSerial())))
loss_threaded(θ) = sum(abs2,1.0.-Array(model1(θ,EnsembleThreads())))

cb = function (θ,l) # callback function to observe training
  @show l
  false
end

opt = ADAM(0.1)
l1 = loss_serial(θ)
res_serial = DiffEqFlux.sciml_train(loss_serial, θ, opt; cb = cb, maxiters=100)
res_threads = DiffEqFlux.sciml_train(loss_threaded, θ, opt; cb = cb, maxiters=100)

## ------------------------------------------------------------------------------------

using DifferentialEquations
prob = ODEProblem((u,p,t)->1.01*[u[2];u[1]],[0.5, 1.0],(0.0,1.0))
function prob_func(prob,i,repeat)
  remake(prob,u0=rand()*prob.u0)
end
ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
sim = solve(ensemble_prob,Tsit5(),EnsembleThreads(),trajectories=100)

## ------------------------------------------------------------------------------------

using DifferentialEquations, Plots, GalacticOptim

function lotka_volterra!(du,u,p,t)
    rab, wol = u
    α,β,γ,δ=p
    du[1] = drab = α*rab - β*rab*wol
    du[2] = dwol = γ*rab*wol - δ*wol
    nothing
end

u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob = ODEProblem(lotka_volterra!,u0,tspan,p)
sol = solve(prob,saveat=0.1)
plot(sol)

dataset = Array(sol)
scatter!(sol.t,dataset')

tmp_prob = remake(prob, p=[1.2,0.8,2.5,0.8])
tmp_sol = solve(tmp_prob)
plot(tmp_sol)
scatter!(sol.t,dataset')

function loss(p)
  tmp_prob = remake(prob, p=p)
  tmp_sol = solve(tmp_prob,Tsit5(),saveat=0.1)
  if tmp_sol.retcode == :Success
    return sum(abs2,Array(tmp_sol) - dataset)
  else
    return 1E10
  end
end

# function loss(p)
#   tmp_prob = remake(prob, p=p)
#   tmp_sol = solve(tmp_prob,Tsit5(),saveat=0.1)
#   if size(tmp_sol) == size(dataset)
#     return sum(abs2,Array(tmp_sol) .- dataset)
#   else
#     return Inf
#   end
# end

using DiffEqFlux

pinit = [1.2,0.8,2.5,0.8]
res = DiffEqFlux.sciml_train(loss,pinit,ADAM(), maxiters = 100)

# res = DiffEqFlux.sciml_train(loss,pinit,BFGS(), maxiters = 100) ### errors!

#try Newton method of optimization
res = DiffEqFlux.sciml_train(loss,pinit,BFGS())

## ------------------------------------------------------------------------------------

# compatibility of modelingtoolkitize and ensemble parellel simulation
using DifferentialEquations, ModelingToolkit
function rober(du,u,p,t)
  y₁,y₂,y₃ = u
  k₁,k₂,k₃ = p
  du[1] = -k₁*y₁+k₃*y₂*y₃
  du[2] =  k₁*y₁-k₂*y₂^2-k₃*y₂*y₃
  du[3] =  k₂*y₂^2
  nothing
end

function prob_func(prob,i,repeat)
    p_i = collect(prob.p)
    p_i[1] = 1E-2 * i
    remake(prob,p=p_i)
    # remake(prob,u0=rand()*prob.u0)
end

prob = ODEProblem(rober,[1.0,0.0,0.0],(0.0,1e1),(0.04,3e7,1e4))
sys = modelingtoolkitize(prob)
prob_jac = ODEProblem(sys,[],(0.0,1e1),jac=true)

ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
ensemble_prob_jac = EnsembleProblem(prob_jac,prob_func=prob_func)

@time sim = solve(ensemble_prob,Tsit5(),EnsembleSerial(),trajectories=10)
@time sim_jac = solve(ensemble_prob_jac,Tsit5(),EnsembleSerial(),trajectories=10)

# conclusion 1) sim_jac is faster than sim if ensemble algorithm is EnsembleSerial().
# ---> sim_jac is even slower than sim when EnsembleThreads() is used.
# conclusion 2) sim_jac[i].prob.p is determined as intended by prob_func even though the modelingtoolkitized system is used.

## ------------------------------------------------------------------------------------

using Flux

d = Dense(5,2)
p1 = d.weight, d.bias
p2 = Flux.params(d)
p1[1] == p2[1]
p2[1] == p2[2]

m = Chain(Dense(10,5), Dense(5,2))
x = rand(10)
# m(x) == m[2](m[1](x))

# loss(ŷ, y)

p, re = Flux.destructure(m)

m(x) - re(p)(x)

## ------------------------------------------------------------------------------------
using Flux

W = rand(2, 5)
b = rand(2)

predict(x) = (W * x) .+ b
loss(x, y) = sum((predict(x) .- y).^2)

x, y = rand(5), rand(2) # Dummy data
l = loss(x, y) # ~ 3

θ = params(W, b)
grads = gradient(() -> loss(x, y), θ)

using Flux.Optimise: update!

η = 0.1 # Learning Rate
for p in (W, b)
  update!(p, η * grads[p])
end

opt = Descent(0.1) # Gradient descent with learning rate 0.1

for p in (W, b)
  update!(opt, p, grads[p])
end

## ------------------------------------------------------------------------------------
# Comparison between the automatic differentiation results for ForwardDiffSensitivity and InterpolatingAdjoint
using DiffEqSensitivity, OrdinaryDiffEq, Zygote, LinearAlgebra, QuadGK

function fiip(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
  du[3] = u[2]^2
end

p       = [1.5,1.0,3.0,1.0]
u0      = [1.0;1.0;0.0]
prob    = ODEProblem(fiip,u0,(0.0,10.0),p)
# y_sense = solve(prob,Tsit5(),u0=u0,p=p,saveat=0.1;sensealg = InterpolatingAdjoint(autojacvec=ZygoteVJP()))

# loss(u0,p) = sum(solve(prob,Tsit5(),u0=u0,p=p,saveat=0.1,abstol=1e-14,reltol=1e-14))
# loss_adjoint(u0,p) = sum(solve(prob,Tsit5(),u0=u0,p=p,saveat=0.1,abstol=1e-14,reltol=1e-14;sensealg = InterpolatingAdjoint()))
function loss(u0,p)
  y = Array(solve(prob,Tsit5(),u0=u0,p=p,saveat=0.1,abstol=1e-14,reltol=1e-14))
  return y[3,end]
end
function loss_adjoint(u0,p)
  y = Array(solve(prob,Tsit5(),u0=u0,p=p,saveat=0.1,abstol=1e-14,reltol=1e-14;sensealg = InterpolatingAdjoint()))
  return y[3,end]
end
function loss_discrete(u0,p)
  y = Array(solve(prob,Tsit5(),u0=u0,p=p,saveat=0.1,abstol=1e-14,reltol=1e-14;sensealg = InterpolatingAdjoint()))
  # return quadgk((t) -> (y(t)[2])^2, 0.0, 10.0)[1] # atol=1e-14,rtol=1e-10)  # this is incorrect and not working
  return 0.1*sum(y[2,:].^2)
end
function loss_adjoint_Bolza(u0,p)
  y = Array(solve(prob,Tsit5(),u0=u0,p=p,abstol=1e-14,reltol=1e-14;sensealg = InterpolatingAdjoint()))
  return y[3,end] + 100.0*y[1:2,end]'*y[1:2,end]
end

du01,dp1 = Zygote.gradient(loss,u0,p)
du02,dp2 = Zygote.gradient(loss_adjoint,u0,p)
du03,dp3 = Zygote.gradient(loss_discrete,u0,p)
du04,dp4 = Zygote.gradient(loss_adjoint_Bolza,u0,p)

# Comparison between the results of automatic differentiation and continuous adjoint sensitivity analysis method
y = solve(prob,Tsit5(),u0=u0,p=p,abstol=1e-14,reltol=1e-14;sensealg = InterpolatingAdjoint())
bwd_sol = solve(ODEAdjointProblem(y, InterpolatingAdjoint(), (out, x, p, t, i) -> (out[:] = [200.0*y[end][1:2]; 1]'), [10.0]), Tsit5(),
                # dense = false, save_everystep = false, save_start = false, # saveat = 0.1,
                abstol=1e-14,reltol=1e-14)
du05, dp5 = bwd_sol[end][1:length(u0)], -bwd_sol[end][1+length(u0):end]

isapprox(du04, du05)
isapprox(dp4, dp5)

## ------------------------------------------------------------------------------------
# Comparison between forward and adjoint sensitivity calculation results
using DiffEqSensitivity, OrdinaryDiffEq, Zygote

function f(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + u[1]*u[2]
end

p = [1.5,1.0,3.0]
u0 = [1.0;1.0]
prob = ODEProblem(f,u0,(0.0,10.0),p)
sol = solve(prob,Vern9(),abstol=1e-10,reltol=1e-10)

g(u,p,t) = (sum(u).^2) ./ 2 + p'*p / 2

function dgdu(out,u,p,t)
  out[1]= u[1] + u[2]
  out[2]= u[1] + u[2]
end

function dgdp_s(out,u,p,t)
  out[1] = p[1]
  out[2] = p[2]
  out[3] = p[3]
end

function dgdp_v1(out,u,p,t)
  out .= p
end

function dgdp_v2(out,u,p,t)
  out = p
end

function dgdp_v3(out,u,p,t)
  out = p'
end

res1 = adjoint_sensitivities(sol,Vern9(),g,nothing,abstol=1e-8, reltol=1e-8,iabstol=1e-8,ireltol=1e-8)                # this is incorrect
res2 = adjoint_sensitivities(sol,Vern9(),g,nothing,dgdu,abstol=1e-8, reltol=1e-8,iabstol=1e-8,ireltol=1e-8)           # this is incorrect
res3 = adjoint_sensitivities(sol,Vern9(),g,nothing,(dgdu,dgdp_s),abstol=1e-8, reltol=1e-8,iabstol=1e-8,ireltol=1e-8)  # this is correct
res4 = adjoint_sensitivities(sol,Vern9(),g,nothing,(dgdu,dgdp_v1),abstol=1e-8, reltol=1e-8,iabstol=1e-8,ireltol=1e-8) # this is correct
res5 = adjoint_sensitivities(sol,Vern9(),g,nothing,(dgdu,dgdp_v2),abstol=1e-8, reltol=1e-8,iabstol=1e-8,ireltol=1e-8) # this is incorrect
res6 = adjoint_sensitivities(sol,Vern9(),g,nothing,(dgdu,dgdp_v3),abstol=1e-8, reltol=1e-8,iabstol=1e-8,ireltol=1e-8) # this is incorrect

using ForwardDiff, Calculus
using QuadGK
function G(p)
  tmp_prob = remake(prob,p=p)
  sol = solve(tmp_prob,Vern9(),abstol=1e-14,reltol=1e-14)
  res,err = quadgk((t)-> (sum(sol(t)).^2)./2 + p'*p/2,0.0,10.0,atol=1e-14,rtol=1e-10)
  res
end
res7 = ForwardDiff.gradient(G,p)
res8 = Calculus.gradient(G,p)
# res9 = Zygote.gradient(G,p) # not working

## ------------------------------------------------------------------------------------
## ------------------------------------------------------------------------------------
## ------------------------------------------------------------------------------------
# Samuel Aintworth's code

import DiffEqBase
import DiffEqSensitivity:
    solve,
    ODEProblem,
    ODEAdjointProblem,
    InterpolatingAdjoint
import Zygote
import Statistics: mean

function extract_loss_and_xT(fwd_sol)
    fwd_sol[end][1], fwd_sol[end][2:end]
end

"""Returns a differentiable loss function that rolls out a policy in an
environment and calculates its cost."""
function ppg_goodies(dynamics, cost, policy, T)

    function aug_dynamics(z, policy_params, t)
        x = @view z[2:end]
        u = policy(x, t, policy_params)
        [cost(x, u); dynamics(x, u)]
    end

    function loss_pullback(x0, policy_params, solvealg, solve_kwargs)
        z0 = vcat(0.0, x0)
        fwd_sol = solve(
            ODEProblem(aug_dynamics, z0, (0, T), policy_params),
            solvealg,
            u0 = z0,
            p = policy_params;
            solve_kwargs...,
        )

        function _adjoint_solve(g_zT, sensealg; kwargs...)
            # See https://diffeq.sciml.ai/stable/analysis/sensitivity/#Syntax-1
            # and https://github.com/SciML/DiffEqSensitivity.jl/blob/master/src/local_sensitivity/sensitivity_interface.jl#L9.
            solve(
                ODEAdjointProblem(
                    fwd_sol,
                    sensealg,
                    (out, x, p, t, i) -> (out[:] = g_zT),
                    [T],
                ),
                solvealg;
                kwargs...,
            )
        end

        # This is the pullback using the augmented system and a discrete
        # gradient input at time T. Alternatively one could use the continuous
        # adjoints on the non-augmented system although this seems to be slower
        # and a less stable feature.

        function pullback(g_zT, sensealg::InterpolatingAdjoint)
            bwd_sol = _adjoint_solve(
                g_zT,
                sensealg,
                dense = false,
                save_everystep = false,
                save_start = false,
                # reltol = 1e-3,
                # abstol = 1e-3,
            )

            # The first z_dim elements of bwd_sol.u are the gradient wrt z0,
            # next however many are the gradient wrt policy_params.
            p = fwd_sol.prob.p
            l = p === nothing || p === DiffEqBase.NullParameters() ? 0 :
                length(fwd_sol.prob.p)
            g_x0 = bwd_sol[end][1:length(fwd_sol.prob.u0)]

            # We do exactly as many f calls as there are function calls in the
            # forward pass, and in the backward pass we don't need to call f,
            # but instead we call ∇f.
            (
                g = -bwd_sol[end][(1:l).+length(fwd_sol.prob.u0)],
                nf = 0,
                n∇ₓf = bwd_sol.destats.nf,
                n∇ᵤf = bwd_sol.destats.nf,
            )
        end

        fwd_sol, pullback
    end

    function ez_loss_and_grad(
        x0,
        policy_params,
        solvealg,
        sensealg;
        fwd_solve_kwargs = Dict(),
    )
        # @info "fwd"
        fwd_sol, vjp = loss_pullback(x0, policy_params, solvealg, fwd_solve_kwargs)
        # @info "bwd"
        bwd = vjp(vcat(1, zero(x0)), sensealg)
        loss, _ = extract_loss_and_xT(fwd_sol)
        # @info "fin"
        loss, bwd.g, (nf = fwd_sol.destats.nf + bwd.nf, n∇ₓf = bwd.n∇ₓf, n∇ᵤf = bwd.n∇ᵤf)
    end

    function _aggregate_batch_results(res)
        (
            mean(loss for (loss, _, _) in res),
            mean(g for (_, g, _) in res),
            (
                nf = sum(info.nf for (_, _, info) in res),
                n∇ₓf = sum(info.n∇ₓf for (_, _, info) in res),
                n∇ᵤf = sum(info.n∇ᵤf for (_, _, info) in res),
            ),
        )
    end

    function ez_loss_and_grad_many(
        x0_batch,
        policy_params,
        solvealg,
        sensealg;
        fwd_solve_kwargs = Dict(),
    )
        # Using tmap here gives a segfault. See https://github.com/tro3/ThreadPools.jl/issues/18.
        _aggregate_batch_results(
            map(x0_batch) do x0
                ez_loss_and_grad(
                    x0,
                    policy_params,
                    solvealg,
                    sensealg,
                    fwd_solve_kwargs = fwd_solve_kwargs,
                )
            end,
        )
    end

    (
        aug_dynamics = aug_dynamics,
        loss_pullback = loss_pullback,
        ez_loss_and_grad = ez_loss_and_grad,
        ez_loss_and_grad_many = ez_loss_and_grad_many,
        ez_euler_bptt = ez_euler_bptt,
        ez_euler_loss_and_grad_many = ez_euler_loss_and_grad_many,
    )
end

function policy_dynamics!(dx, x, policy_params, t)
    u = policy(x, policy_params)
    dx .= dynamics(x, u)
end

function cost_functional(x, policy_params, t)
    cost(x, policy(x, policy_params))
end

# See https://github.com/SciML/DiffEqSensitivity.jl/issues/302 for context.
dcost_tuple = (
    (out, u, p, t) -> begin
        ū, _, _ = Zygote.gradient(cost_functional, u, p, t)
        out .= ū
    end,
    (out, u, p, t) -> begin
        _, p̄, _ = Zygote.gradient(cost_functional, u, p, t)
        out .= p̄
    end,
)

function gold_standard_gradient(x0, policy_params)
    # Actual/gold standard evaluation. Using high-fidelity Vern9 method with
    # small tolerances. We want to use Float64s for maximum accuracy. Also 1e-14
    # is recommended as the minimum allowable tolerance here: https://docs.sciml.ai/stable/basics/faq/#How-to-get-to-zero-error-1.
    x0_f64 = convert(Array{Float64}, x0)
    policy_params_f64 = convert(Array{Float64}, policy_params)
    fwd_sol = solve(
        ODEProblem(policy_dynamics!, x0, (0, T), policy_params),
        Vern9(),
        u0 = x0_f64,
        p = policy_params_f64,
        abstol = 1e-14,
        reltol = 1e-14,
    )
    # Note that specifying dense = false is essential for getting acceptable
    # performance. save_everystep = false is another small win.
    bwd_sol = solve(
        ODEAdjointProblem(
            fwd_sol,
            InterpolatingAdjoint(),
            cost_functional,
            nothing,
            dcost_tuple,
        ),
        Vern9(),
        dense = false,
        save_everystep = false,
        abstol = 1e-14,
        reltol = 1e-14,
    )
    @assert typeof(fwd_sol.u) == Array{Array{Float64,1},1}
    @assert typeof(bwd_sol.u) == Array{Array{Float64,1},1}

    # Note that the backwards solution includes the gradient on x0, as well as
    # policy_params. The full ODESolution things can't be serialized easily
    # since they `policy_dynamics!` and shit...
    (xT = fwd_sol.u[end], g = bwd_sol.u[end])
end

function eval_interp(x0, policy_params, abstol, reltol)
    fwd_sol = solve(
        ODEProblem(policy_dynamics!, x0, (0, T), policy_params),
        Tsit5(),
        u0 = x0,
        p = policy_params,
        abstol = abstol,
        reltol = reltol,
    )
    bwd_sol = solve(
        ODEAdjointProblem(
            fwd_sol,
            InterpolatingAdjoint(),
            cost_functional,
            nothing,
            dcost_tuple,
        ),
        Tsit5(),
        dense = false,
        save_everystep = false,
        abstol = abstol,
        reltol = reltol,
    )
    @assert typeof(fwd_sol.u) == Array{Array{floatT,1},1}
    @assert typeof(bwd_sol.u) == Array{Array{floatT,1},1}

    # Note that g includes the x0 gradient and the gradient on parameters.
    # We do exactly as many f calls as there are function calls in the forward
    # pass, and in the backward pass we don't need to call f, but instead we
    # call ∇f.
    (
        xT = fwd_sol.u[end],
        g = bwd_sol.u[end],
        nf = fwd_sol.destats.nf,
        n∇f = bwd_sol.destats.nf,
    )
end

## ------------------------------------------------------------------------------------
# using DiffEqFlux, DifferentialEquations, Plots
# using ComponentArrays
# using FSimBase  # for apply_inputs

# u0 = Float32[2.; 0.]
# u0_aug = vcat(u0, zeros(1))
# datasize = 30
# tspan = (0.0f0,1.5f0)

# L(u) = u' * u  # running cost

# function trueODEfunc(dx,x,p,t; a)
#     true_A = [-0.1 2.0; -2.0 -0.1]
#     dx[1:2, :] .= ((x[1:2, :] .^ 3)'true_A)' + a
#     dx[3, :] = L(x[1:2, :])
# end
# t = range(tspan[1],tspan[2],length=datasize)
# # prob = ODEProblem(trueODEfunc,u0,tspan)
# # prob = ODEProblem(trueODEfunc,u0_aug,tspan)
# # ode_data = Array(solve(prob,Tsit5(),saveat=t))

# # dudt2 = Chain(x -> x.^3,
# #              Dense(2,50,tanh),
# #              Dense(50,2))
# # p,re = Flux.destructure(dudt2) # use this p as the initial condition!
# # dudt(u,p,t) = re(p)(u) # need to restructure for backprop!
# # prob = ODEProblem(dudt,u0_aug,tspan)
# controller2 = Chain(u -> u.^3,
#                    Dense(2,50,tanh),
#                    Dense(50,2))
# p,re = Flux.destructure(controller2)
# controller(u, p, t) = re(p)(u)
# prob = ODEProblem(apply_inputs(trueODEfunc; a=(x, p, t) -> controller(x[1:2], p, t)),u0_aug,tspan)

# function predict_n_ode()
#   Array(solve(prob,Tsit5(),u0=u0_aug,p=p,saveat=t))  # it seems that the output should be an Array
# end

# function loss_n_ode()
#     pred = predict_n_ode()
#     # loss = sum(abs2,ode_data .- pred)
#     loss = pred[end][end]  # to make it scalar for gradient calculation
#     loss
# end

# loss_n_ode() # n_ode.p stores the initial parameters of the neural ODE

# cb = function (;doplot=false) # callback function to observe training
#     pred = predict_n_ode()
#     @show pred[end][end]
#     # display(sum(abs2,ode_data .- pred))
#     # # plot current prediction against data
#     # pl = scatter(t,ode_data[1,:],label="data")
#     # scatter!(pl,t,pred[1,:],label="prediction")
#     # display(plot(pl))
#     return false
# end

# # Display the ODE with the initial parameter values.
# cb()

# data = Iterators.repeated((), 1000)
# # Flux.train!(loss_n_ode, Flux.params(u0,p), data, ADAM(0.05), cb = cb)
# Flux.train!(loss_n_ode, Flux.params(p), data, ADAM(0.05), cb = cb)

#
# #-------------------------------------------------------------------------------
# using DiffEqFlux, DifferentialEquations, Plots, Statistics
# tspan = (0.0f0,8.0f0)
# ann = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
# θ = initial_params(ann)
# function dxdt_(dx,x,p,t)
#     x1, x2 = x
#     dx[1] = x[2]
#     dx[2] = ann([t],p)[1]^3
# end
# x0 = [-4f0,0f0]
# ts = Float32.(collect(0.0:0.01:tspan[2]))
# prob = ODEProblem(dxdt_,x0,tspan,θ)
# solve(prob,Vern9(),abstol=1e-10,reltol=1e-10)
# ##

# function predict_adjoint(θ)
#   Array(solve(prob,Vern9(),p=θ,saveat=ts,sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true))))
# end
# function loss_adjoint(θ)
#   x = predict_adjoint(θ)
#   mean(abs2,4.0 .- x[1,:]) + 2mean(abs2,x[2,:]) + mean(abs2,[first(ann([t],θ)) for t in ts])/10
# end
# l = loss_adjoint(θ)
# cb = function (θ,l)
#   println(l)
#   p = plot(solve(remake(prob,p=θ),Tsit5(),saveat=0.01),ylim=(-6,6),lw=3)
#   plot!(p,ts,[first(ann([t],θ)) for t in ts],label="u(t)",lw=3)
#   display(p)
#   return false
# end
# # Display the ODE with the current parameter values.
# cb(θ,l)
# loss1 = loss_adjoint(θ)
# res1 = DiffEqFlux.sciml_train(loss_adjoint, θ, ADAM(0.005), cb = cb,maxiters=100)
# res2 = DiffEqFlux.sciml_train(loss_adjoint, res1.u,
#                               BFGS(initial_stepnorm=0.01), cb = cb,maxiters=100,
#                               allow_f_increases = false)

##-------------------------------------------------------------------------------
using DiffEqFlux, DiffEqSensitivity, DifferentialEquations, Plots, Statistics
tspan = (0.0f0,8.0f0)
ann = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
θ = initial_params(ann)
function dxdt_(x,p,t)
    x1, x2 = x
    return dx = [x2; ann([t],p)[1]^3]
end
x0 = [-4f0,0f0]
ts = Float32.(collect(0.0:0.01:tspan[2]))
prob = ODEProblem(dxdt_,x0,tspan,θ)
# solve(prob,Vern9(),abstol=1e-10,reltol=1e-10)
# function predict_adjoint(θ)
#   Array(solve(prob,Vern9(),p=θ,saveat=0.01f0,sensealg=InterpolatingAdjoint()))
# end
function loss_adjoint_test(θ)
#   x = predict_adjoint(θ)
  x = Array(solve(prob,Tsit5(),p=θ,saveat=0.01f0,sensealg=InterpolatingAdjoint( autojacvec = ZygoteVJP() )))
  mean(abs2,4.0f0 .- x[1,:]) + 2mean(abs2,x[2,:]) + mean(abs2,[first(ann([t],θ)) for t in ts])/10.0f0
end
l = loss_adjoint_test(θ)
cb = function (θ,l)
  println(l)
  p = plot(solve(remake(prob,p=θ),Tsit5(),saveat=0.01),ylim=(-6,6),lw=3)
  plot!(p,ts,[first(ann([t],θ)) for t in ts],label="u(t)",lw=3)
  display(p)
  return false
end
# Display the ODE with the current parameter values.

cb(θ,l)
loss1 = loss_adjoint_test(θ)
res1 = DiffEqFlux.sciml_train(loss_adjoint_test, θ, ADAM(0.005); cb = cb, maxiters=100)
res2 = DiffEqFlux.sciml_train(loss_adjoint_test, res1.u,
                              BFGS(initial_stepnorm=0.01), cb = cb,maxiters=100,
                              allow_f_increases = false)

## -------------------------------------------------------------------------
using OrdinaryDiffEq, DiffEqFlux
pa = [1.0]
u0 = [3.0]
θ = [u0;pa]

function model1(θ,ensemble)
  prob = ODEProblem((u, p, t) -> 1.01u .* p, [θ[1]], (0.0, 1.0), [θ[2]])

  function prob_func(prob, i, repeat)
    remake(prob, u0 = 0.5 .+ i/100 .* prob.u0)
  end

  ensemble_prob = EnsembleProblem(prob, prob_func = prob_func)
  sim = solve(ensemble_prob, Tsit5(), ensemble, saveat = 0.1, trajectories = 100, sensealg = InterpolatingAdjoint(autojacvec = ZygoteVJP()))
end

# loss function
loss_serial(θ)   = sum(abs2,1.0.-Array(model1(θ,EnsembleSerial())))
loss_threaded(θ) = sum(abs2,1.0.-Array(model1(θ,EnsembleThreads())))

prob_base = ODEProblem((u, p, t) -> 1.01u .* p, [θ[1]], (0.0, 1.0), [θ[2]])
function prob_gen(prob, i, repeat)
  remake(prob, u0 = 0.5 .+ i/100 .* prob.u0)
end
function loss_test(θ)
    prob = remake(prob_base, u0 = [θ[1]], p = [θ[2]])
    fwd_ensemble_sol = Array( solve(EnsembleProblem(prob, prob_func = prob_gen), Tsit5(), EnsembleThreads(), saveat = 0.1, trajectories = 100, sensealg = InterpolatingAdjoint(autojacvec = ZygoteVJP())))
    loss_val = sum(abs2, 1.0 .- fwd_ensemble_sol)
    return loss_val
end

cb = function (θ,l) # callback function to observe training
  @show l
  false
end

opt = ADAM(0.1)
l1 = loss_serial(θ)
# res_serial = DiffEqFlux.sciml_train(loss_serial, θ, opt; cb = cb, maxiters=100)
# res_threads = DiffEqFlux.sciml_train(loss_threaded, θ, opt; cb = cb, maxiters=100)
res_test = DiffEqFlux.sciml_train(loss_test, θ, opt; cb = cb, maxiters = 100)

## ----------------
using DiffEqGPU, OrdinaryDiffEq
function lorenz(du,u,p,t)
    du[1] = p[1]*(u[2]-u[1])
    du[2] = u[1]*(p[2]-u[3]) - u[2]
    du[3] = u[1]*u[2] - p[3]*u[3]
end

u0 = Float32[1.0;0.0;0.0]
tspan = (0.0f0,100.0f0)
p = [10.0f0,28.0f0,8/3f0]
prob = ODEProblem(lorenz,u0,tspan,p)
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy=false)
@time sol = solve(monteprob,Tsit5(),EnsembleGPUArray(),trajectories=10000,saveat=1.0f0)

## -------------------
using OrdinaryDiffEq, Flux, Optim, DiffEqFlux, DiffEqSensitivity

model_gpu = Chain(Dense(2, 50, tanh), Dense(50, 2)) |> gpu
p, re = Flux.destructure(model_gpu)
dudt!(u, p, t) = re(p)(u)

# Simulation interval and intermediary points
tspan = (0.0, 10.0)
tsteps = 0.0:0.1:10.0

u0 = Float32[2.0; 0.0] |> gpu
prob_gpu = ODEProblem(dudt!, u0, tspan, p)

# Runs on a GPU
sol_gpu = solve(prob_gpu, Tsit5(), saveat = tsteps)


function makePlots()

println("Generating plots")

plotsDir = (pwd()[end-3:end] == "docs") ? "build/plots" : "docs/build/plots"
mkdir(plotsDir)
Plots.gr()

# LQG design
h = 0.1
A = [1 h; 0 1]
B = [0;1]
C = [1 0]
sys = ss(A,B,C,0, h)
Q = Matrix{Float64}(I,2,2)
R = Matrix{Float64}(I,1,1)
L = dlqr(A,B,Q,R) # lqr(sys,Q,R) can also be used

u(x,t)  = -L*x .+ 1.5(t>=2.5)# Form control law (u is a function of t and x), a constant input disturbance is affecting the system from t≧2.5
t=0:h:5
x0 = [1,0]
y, t, x, uout = lsim(sys,u,t,x0=x0)
plot(t,x, lab=["Position" "Velocity"], xlabel="Time [s]")
Plots.savefig(plotsDir*"/lqrplot.svg")

# PID design functions
P = tf(1,[1,1])^4
gangoffourplot(P,tf(1))
Plots.savefig(plotsDir*"/pidgofplot.svg")

ωp = 0.8
kp,ki,C = loopshapingPI(P,ωp,phasemargin=60, doplot=false)
gangoffourplot(P, [tf(1), C])
Plots.savefig(plotsDir*"/pidgofplot2.svg")
nyquistplot([P, P*C])
Plots.savefig(plotsDir*"/pidnyquistplot.svg")

ωp = 2
kp,ki,C60 = loopshapingPI(P,ωp,rl=1,phasemargin=60, doplot=true)
gangoffourplot(P, [tf(1), C60])
Plots.savefig(plotsDir*"/pidgofplot3.svg")
nyquistplot([P, P*C60])
Plots.savefig(plotsDir*"/pidnyquistplot2.svg")

# Advanced pole placement
ζ = 0.2
ω = 1
B = [1]
A   = [1, 2ζ*ω, ω^2]
P  = tf(B,A)
# Control design
ζ0 = 0.7
ω0 = 2
Am = [1, 2ζ0*ω0, ω0^2]
Ao = conv(2Am, [1/2, 1]) # Observer polynomial
AR = [1,0] # Force the controller to contain an integrator

B⁺  = [1] # The process numerator polynomial can be facored as B = B⁺B⁻ where B⁻ contains the zeros we do not want to cancel (non-minimum phase and poorly damped zeros)
B⁻  = [1]
Bm  = conv(B⁺, B⁻) # In this case, keep the entire numerator polynomial of the process

R,S,T = rstc(B⁺,B⁻,A,Bm,Am,Ao,AR) # Calculate the 2-DOF controller polynomials

Gcl = tf(conv(B,T),zpconv(A,R,B,S)) # Form the closed loop polynomial from reference to output

stepplot([P,Gcl]) # Visualize the open and closed loop responses.
Plots.savefig(plotsDir*"/ppstepplot.svg")
gangoffourplot(P, tf(-S,R)) # Plot the gang of four to check that all tranfer functions are OK
Plots.savefig(plotsDir*"/ppgofplot.svg")

P1(s) = exp(-sqrt(s))
f1 = stabregionPID(P1,exp10.(range(-5, stop=1, length=1000))); Plots.savefig(plotsDir*"/stab1.svg")
P2 = s -> 100*(s+6).^2 ./(s.*(s+1).^2 .*(s+50).^2)
f2 = stabregionPID(P2,exp10.(range(-5, stop=2, length=1000))); Plots.savefig(plotsDir*"/stab2.svg")
P3 = tf(1,[1,1])^4
f3 = stabregionPID(P3,exp10.(range(-5, stop=0, length=1000))); Plots.savefig(plotsDir*"/stab3.svg")

# PID plots

P = tf([1.],[1., 1])
ζ = 0.5 # Desired damping
ws = exp10.(range(-1, stop=2, length=8)) # A vector of closed-loop bandwidths
kp = 2*ζ*ws.-1 # Simple pole placement with PI given the closed-loop bandwidth, the poles are placed in a butterworth pattern
ki = ws.^2
pidplots(P,:nyquist,;kps=kp,kis=ki, ω= exp10.(range(-2, stop=2, length=500)))
Plots.savefig(plotsDir*"/pidplotsnyquist1.svg")
pidplots(P,:gof,;kps=kp,kis=ki, ω= exp10.(range(-2, stop=2, length=500)))
Plots.savefig(plotsDir*"/pidplotgof1.svg")

kp = range(-1, stop=1, length=8) # Now try a different strategy, where we have specified a gain crossover frequency of 0.1 rad/s
ki = sqrt.(1 .-kp.^2)./10
pidplots(P,:nyquist,;kps=kp,kis=ki)
Plots.savefig(plotsDir*"/pidplotsnyquist2.svg")
pidplots(P,:gof,;kps=kp,kis=ki)
Plots.savefig(plotsDir*"/pidplotsgof2.svg")

#AR = [1,0]
#B⁺  = [1]
#B⁻  = [1]
#Bm  = conv(B⁺, B⁻)
#R,S,T = rstc(B⁺,B⁻,A,bm,Am,Ao,AR)
#Gcl = tf(conv(B,T),zpconv(A,R,B,S))
#f1 = stepplot([P,Gcl])
#f2 = gangoffourplot(P, tf(-S,R))
#Plots.savefig(f1, "$plotsDir/rstcstepplot.svg")
#Plots.savefig(f2, "$plotsDir/rstcgofplot.svg")
end


module AlgebraicDependencies

export @deps, dependencies

using LinearAlgebra
using SymEngine
using SymPy
using Nemo
using Hecke
using Singular
using Polynomials.PolyCompat

import Base: length, iterate, collect

include("ContinuedFractions.jl")

"""
    @deps r1 r2 r3 ...

Compute a basis for the algebraic dependencies among the ``r_i^n``.

# Examples
```julia-repl
julia> @deps 2 1/2
1-element Array{Expr,1}:
 :(v1 * v2 - 1)

julia> @deps (1-sqrt(5))/2 (1+sqrt(5))/2
1-element Array{Expr,1}:
 :(v1 ^ 2 * v2 ^ 2 - 1)
```
"""
macro deps(args...)
    dependencies(collect(args))
end

function dependencies(::Type{sideal}, roots::Vector{Basic}; variables=Basic[])
    if length(roots) < 2
        return nothing
    end
    if isempty(variables)
        variables = [Basic("v$i") for i in 1:length(roots)]
    elseif length(variables) != length(roots)
        throw("Number of variables does not match number exponentials. Got $(length(variables)), need $(length(roots))")
    end
    lattice = findrelations(roots)
    @debug "Lattice" lattice
    if isempty(lattice)
        return nothing
    end
    return ideal(lattice, variables)
end

dependencies(::Type{sideal}, roots; variables=[]) =
    dependencies(sideal, map(Basic ∘ string, roots); variables = map(Basic ∘ string, variables))

function dependencies(::Type{Basic}, roots::Vector{Basic}; variables=Basic[])
    ideal = dependencies(sideal, roots; variables = variables)
    if ideal === nothing
        return Basic[]
    end
    [Basic(string(ideal[i])) for i in 1:Singular.ngens(ideal)]
end

function dependencies(roots; variables=[])
    ideal = dependencies(sideal, map(Basic ∘ string, roots); variables = map(Basic ∘ string, variables))
    if ideal === nothing
        return Expr[]
    end
    [Meta.parse(string(ideal[i])) for i in 1:Singular.ngens(ideal)]
end

is_rational(x::Basic) = SymEngine.BasicType(x) isa SymEngine.BasicType{Val{:Integer}} || SymEngine.BasicType(x) isa SymEngine.BasicType{Val{:Rational}}

function minpoly(r::Basic, z::Symbol)
    s = sympify(z)
    p = sympy.minpoly(sympify(string(r)), s)
    cs = sympy.Poly(p, s).coeffs()
    bs = map(Basic ∘ string, cs)
    Basic(string(p)), PolyCompat.Poly(bs)
end

function common_number_field(roots::Vector{Basic})
    v = :a
    minpolys = Basic[]
    ps = PolyCompat.Poly[]
    for r in roots
        mp, p = minpoly(r, v)
        if is_rational(r)
            push!(minpolys, r)
        else
            push!(minpolys, mp)
        end
        push!(ps, p)
    end
    indices = findall(is_rational, minpolys)
    ratroots = [minpolys[i] for i in indices]
    deleteat!(minpolys, indices)
    mps = [convert(Expr, b) for b in minpolys]
    if isempty(mps)
        return 1, roots, ps
    end
    deleteat!(roots, indices)

    uniquemps = Base.unique(mps)
    mpsindices = [findfirst(x->x==m, uniquemps) for m in mps]
    R, _ = Nemo.PolynomialRing(Nemo.QQ, string(v))
    K, g = number_field(map(R, uniquemps))
    @debug "Number field" K, mpsindices
    S, mS = simple_extension(K)
    @debug "Simple extension" S
    gs = [SymEngine.lambdify(Basic(replace(string(mS\(x)), "//"=>"/"))) for x in g]
    rs = [gs[j](roots[i]) for (i, j) in enumerate(mpsindices)]
    for (i, ind) in enumerate(indices)
        splice!(rs, ind:ind-1, ratroots[i])
    end
    Hecke.degree(S), rs, ps
end

function masser_bound(degree::Int, roots::Vector{Basic}, mpolys::Vector{PolyCompat.Poly})
    k = length(roots)
    # assume all roots belong to the same number field
    d = degree
    # h = maximum height of the a[i]. The height of an algebraic number is the
    # sum of the degree and the binary length of all coefficients in the
    # defining equation over Q
    h = 0
    for p in mpolys
        cs = filter(x->!iszero(x), PolyCompat.coeffs(p))
        h0 = ceil(PolyCompat.degree(p) + sum([log(abs(c)) for c in cs]))
        if h0 > h
            h = h0
        end
    end
    @debug "Parameters for Masser bound" d h k
    return ceil(d^2 * (4*h*k*d* (log(2+d)/log(log(2+d)))^3)^(k-1) + 1)
end

gensym_unhashed(s::Symbol) = Symbol(replace(string(gensym(s)), "#"=>""))

function ideal(m::Matrix{BigInt}, x::Array{Basic,1})
    if isempty(m)
        return []
    end

    y = [Basic(gensym_unhashed(:vv)) for i in 1:length(x)]
    base = Basic[]
    for i in 1:nrows(m)
        r = 1
        for j in 1:ncols(m)
            exp = m[i,j]
            v = exp > 0 ? x[j] : y[j]
            r *= v^abs(exp)
        end
        base = [base; r - 1]
    end

    inv = [xi*yi - 1 for (xi,yi) in zip(x,y)]
    base = [base; inv]

    R, g = Singular.PolynomialRing(Nemo.QQ, map(string, [y; x]))
    sbasis = [R(convert(Expr, b)) for b in base]
    I = Singular.Ideal(R, sbasis)
    return Singular.eliminate(I, g[1:length(y)]...)
end

function findrelations(roots::Vector{Basic})
    # first treat zeros in the root list
    zpos = findall(iszero, roots)
    if length(zpos) == length(roots)
        return []
    end
    if !isempty(zpos)
        b = findrelations(filter(!iszero, roots))
        for i in zpos
            b = hcat(b[:,1:i-1], zeros(nrows(b)), b[:,i:end])
        end
        return b
    end

    degree, roots, mpolys = common_number_field(roots)
    @debug "Common number field" degree roots mpolys
    rs = [convert(Complex{BigFloat}, x) for x in roots]
    @debug "Complex roots" rs

    relog = [real(log(x)) for x in rs]
    imlog = [imag(log(x)) for x in rs]
    imlog = [imlog; 2*BigFloat(pi)]

    # Replace implicit zeros by explicit ones
    for i in 1:length(rs)
        z = rs[i]

        # abs(z) == 1
        if abs(abs(z) - 1) < 0.1 && abs(z) == 1
            relog[i] = 0
        end

        # z is real and z >= 0
        if isreal(z) && isreal(sqrt(z))
            imlog[i] = 0
        else
            # TODO: try harder: If[ Element[RootReduce[z], Reals] && Element[Sqrt[RootReduce[z]], Reals], imLog[[i]] = 0 ];
        end
    end
    @debug "reLog and imLog" relog imlog

    # Compute a bound for the coefficients of the generators
    bound = Int(masser_bound(degree, roots, mpolys))
    @debug "Masser bound" bound

    m = eye(Int, length(roots)) # identity matrix

    # Successively refine the approximation until only valid generators are returned
    level = Int(ceil(log2(bound) + 1))
    while prod(Bool[check_relation(roots, m[i,:]) for i in 1:size(m)[1]]) == 0
        m1 = getbasis(relog, level, bound)
        m2 = getbasis(imlog, level, bound)
        @debug "" m1 m2[:,1:end-1]
        m = z_module_intersect(m1, m2[:,1:end-1])
        level = level + 1
    end

    return m
end

hnf_with_transform(m::Matrix{Int}) = Matrix{Int}.(Nemo.hnf_with_transform(matrix(FlintZZ, m)))

function lattice_divide(l::Matrix{Int}, d::Int)
    n = ncols(l)
    return z_module_intersect(l, d*eye(Int, n)) / d
end

function check_relation(a::Vector{Basic}, e::Vector{<:Integer})
    return SymEngine.expand(prod([ax^ex for (ax, ex) in zip(a,e)])) == 1
end

function convergent(x, n)
    cf = ContinuedFraction(x)
    co = convergents(cf)
    return iterate(co, n)[1]
end

nrows(a::Matrix{<:Any}) = size(a)[1]
ncols(a::Matrix{<:Any}) = size(a)[2]

function getbasis(l::Vector{BigFloat}, level::Int, bound::Int)
    n = length(l)

    # First treat zero elements in l as special case
    zpos = findall(iszero, l)
    if length(zpos) == length(l)
        return eye(Int, n)
    end

    if length(zpos) > 0
        ll = deleteat!(copy(l), zpos)
        b = getbasis(ll, level, bound) # basis for nonzero numbers
        zvec = zeros(nrows(b), 1)
        # Insert new dimensions
        for pos in zpos
            b = hcat(b[:,1:pos-1], zvec, b[:,pos:end])
        end
        # Add unit vectors at the zero positions
        b = vcat(b, eye(Int, n)[zpos,:])
        return Matrix{Int}(b)
    end

    # Now we assume that l is a list of nonzero real numbers
    c1 = [convergent(x, level) for x in l]
    c2 = [convergent(x, level+1) for x in l]

    e = [1//denominator(x1*x2) for (x1, x2) in zip(c1, c2)]
    # Refine the approximation such that all errors are smaller than the smallest
    # element of l in absolute value
    lev = level + 1
    while length(filter(x -> maximum(e) >= abs(x), c1)) > 0 && lev < level + 5
        ex = findall(x->x==maximum(e), e) # indices with greates error

        lev++
        for i in 1:length(ex)
            j = ex[i]
            c1[j] = c2[j]
            c2[j] = convergent(l[j], lev)
            e[j] = c1[j] == l[j] ? 0 : 1/denominator(c1[j]*c2[j])
        end
    end

    # now: max e[i] < min |c1[i]|
    # this bound guarantees that generators whose norm is at most bound will
    # appear in the LLL-reduced basis
    minc1 = minimum([abs(c) for c in c1])
    maxe = maximum([abs(c) for c in e])
    d = BigInt(ceil(2^((length(c1) - 1)/2)*bound/(minc1 - maxe)))
    identity = eye(Int, n)
    row = c1 * d
    b = hcat(identity, row)
    b = lll(b)
    # Vectors whose right hand side is greater than the errors permit can be
    #   discarded; they cannot correspond to integer relations.
    # b = vcat([b[i,:] for i in 1:nrows(b) ])
    # TODO: find better way to filter rows
    res = Matrix{Rational{Int}}(undef, 0,n+1)
    for i in 1:nrows(b)
        if (abs(b[i,:][end]) <= d*sum(abs(bb*ee) for (bb, ee) in zip(b[i,1:n],e)))
            res = vcat(res, transpose(b[i,:]))
        end
    end

    # all remaining vectors are returned as candidates
    # TODO: result should be integer matrix?
    return Matrix{Int}(res[:,1:end-1])
end

function z_nullspace(matrix::Matrix{Int})
    h, t = hnf_with_transform(matrix)
    t = t * -1
    # println("HNF: $(h) | $(t)")

    # kernel is generated by the rows of t that correspond to zero rows in h
    zvec = zeros(size(h, 2))

    # TODO: find better way to filter zero vectors
    res = Matrix{Int}(undef, 0, ncols(t))
    for i in 1:nrows(t)
        if iszero(h[i,:])
            res = vcat(res, transpose(t[i,:]))
        end
    end
    return res
    # return [t[i,:] for i in 1:size(h, 1) if h[i,:] == zvec]
end

function z_module_intersect(base1::Matrix{Int}, base2::Matrix{Int})
    if isempty(base1) || isempty(base2)
        return []
    end

    sol = z_nullspace(vcat(base1, base2))

    if isempty(sol)
        return []
    end

    m1 = transpose(base1)
    m2 = transpose(sol[:, 1:nrows(base1)])
    return lll(collect(transpose(m1 * m2)))
end

function clear_denom(a::Matrix{Rational{BigInt}})
    d = lcm(denominator.(a))
    return a*d, d
end

import Nemo.lll

function lll(a::Matrix{Rational{BigInt}})
    m, d = clear_denom(a)
    m = numerator.(m)
    m = Matrix{BigInt}(Nemo.lll(matrix(FlintZZ, m)))
    return m // d
end

lll(m::Matrix{Int}) = Matrix{BigInt}(Nemo.lll(matrix(FlintZZ, m)))

eye(::Type{T}, n) where {T} = Matrix{T}(I, n, n)

function (R::FmpqPolyRing)(p::Expr)
    v = gen(R)
    vs = [:($(Symbol(string(v))) = $v)]
    q = quote
        let $(vs...)
            $(p)
        end
    end
    eval(q)
end

function (R::Singular.PolyRing)(p::Expr)
    vs = gens(R)
    vs = [:($(Symbol(string(v))) = $v) for v in vs]
    q = quote
        let $(vs...)
            $(p)
        end
    end
    eval(q)
end

function (K::NfAbsNS)(a::fmpq_poly)
    q, w = divrem(a, K.pol)
    z = NfAbsNSElem(w)
    z.parent = K
    return z
end

end # module


using DifferentialEquations, GeometricIntegratorsDiffEq, ParameterizedFunctions, Plots, LinearAlgebra, DelimitedFiles, LaTeXStrings

## 3.1 Determining A Benchmark Value for the Infinite-Time average

include("..\\NewtonCotes\\NewtonCotes11pt.jl")

# Define the Chen system and parameters
chen = @ode_def begin
 dx = a * (y - x)
 dy = (c - a) * x - x * z + c * y
 dz = x * y - b * z
end a b c

p = [35.0, 3.0, 28.0]

probChenTrue = ODEProblem(chen, [1, 1, 1], (0, 1000), p)
solChenTrue = solve(probChenTrue, alg=GIGLRK(4), dt=0.0005)

# Must transfer z component of solution to single vector for integration
z = Array{Float64}(undef, length(solChen.u))
for i = 1:length(solChen.u)
    if solChenTrue.t[i] >= 5
        z[i] = solChenTrue.u[i][3]
    end
end

# Integrate and then average over t ∈ [5, 1000]
Z_true_Chen = NewtonCotes.NewtonCotes11pt(z, 0.0005)/(1000-5)

## 3.2 Initial Condition Determination

prob = ODEProblem(chen, [1, 1, 1], (0, 5), p)
sol = solve(prob, alg=GIGLRK(4), dt=0.001)
plot(sol, vars = (1, 2, 3), label=L"u_0 = [1, 1, 1]", linecolor=:blue)

# Observe the long, sweeping curve at the start, which travels up to around z = 60
# with a large initial radius before spiraling inwards and downwards.

# Simulate again with two sets of initial conditions over a long period of time.
prob2 = ODEProblem(chen, [1, 1, 1], (0, 500), p)
sol2 = solve(prob2, alg=GIGLRK(4), dt=0.001)
prob3 = ODEProblem(chen, [-1, 0, 0], (0, 500), p)
sol3 = solve(prob3, alg=GIGLRK(4), dt=0.001)
plot(sol2, vars = (2, 3), label=L"u_0 = [1, 1, 1]", linecolor=:blue)
plot!(sol3, vars = (2, 3), label=L"u_0 = [-1, 0, 0]", linecolor=:orange)

# Examining the above plots demonstrates that, similar to the Lorenz system, there is a
# spin-up period after which no similar behavior is noticed over a very long time.
# In this case the solution trajectory at t = 5 from the 8th order GLRK method with
# timestep h = 0.001 is taken to be outside the spinup period and on the attractor,
# so the coordinates at that point will comprise the new initial conditions.

u0Chen = sol.u[end] # = [10.747927627254173, 12.40061941590775,21.207306876020063]

## 3.3 Step Size Determination

# The same step sizes as determined for the Lorenz system are used.
# No extra computation needed here.
# However note that attempts to simulate with some larger step sizes revealed that
# for the Chen system, GLRK and RK4 are only stable for step sizes up to 0.04,
# and Heun's method is only stable up to 0.02.

## 3.4 Time-March Testing
# NOTE: At time of initial investigation, the programmer was more experienced with MATLAB
# than Julia. Thus a good portion of the statistical analysis code was written in MATLAB, and
# the workflow for gathering statistics looks like this:
#    1. Run simulations in Julia
#    2. Export simulation results to .txt file
#    3. Import results from .txt file into MATLAB
#    4. Perform statistical analysis and plotting within MATLAB
#
# The code may, in the future, be updated to use Julia entirely. However, the code below shows
# steps 1 and 2 of the above workflow; while step 3 is performed by importer.m, and step 4
# is performed by runner.m and stats_plotter.m
#

probChen = ODEProblem(chen, u0Chen, (0, 300), p)

h = [0.04, 0.02, 0.01, 0.005, 0.0025]
for i = h
    # GLRK stable for all step sizes
    solGLRK = solve(probChen, alg=GIGLRK(4), dt=i)
    writedlm("chen\\Simulated Solutions\\GLRK(4)_u_"*string(i)*".txt", solGLRK.u)

    # Explicit RK4 stable for all step sizes
    solRK4 = solve(probChen, alg=GIERK4(), dt=i)
    writedlm("chen\\Simulated Solutions\\RK4_u_"*string(i)*".txt", solRK4.u)

    # Heun's method stable for h < 0.04
    if i < 0.04
        solHeun = solve(probChen, alg=GIHeun(), dt=i)
        writedlm("chen\\Simulated Solutions\\Heun_u_"*string(i)*".txt", solHeun.u)
    end
end


n = rand(1:10)
A = matrixdepot("hilb", n)
B = matrixdepot("lotkin", n)

@test A[2:n, :] == B[2:n, :]
@test B[1,:][:] == ones(n)
println("'lotkin' passed test...")


export compose_rotation

################################################################################
#                              Compose Rotations
################################################################################

"""
    compose_rotation(R1, [, R2, R3, R4, R5, ...])

Compute a composed rotation using the rotations `R1`, `R2`, `R3`, `R4`, ..., in
the following order:

     First rotation
     |
     |
    R1 => R2 => R3 => R4 => ...
           |
           |
           Second rotation

The rotations can be described by:

* A direction cosina matrix (`DCM`);
* An Euler angle and axis (`EulerAngleAxis`);
* A set of Euler anlges (`EulerAngles`); or
* A quaternion (`Quaternion`).

Notice, however, that all rotations **must be** of the same type (DCM or
quaternion).

The output will have the same type as the inputs.

# Example

```julia-repl
julia> D1 = angle_to_dcm(+pi/3,+pi/4,+pi/5,:ZYX);

julia> D2 = angle_to_dcm(-pi/5,-pi/4,-pi/3,:XYZ);

julia> compose_rotation(D1,D2)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
 1.0          0.0          5.55112e-17
 0.0          1.0          5.55112e-17
 5.55112e-17  5.55112e-17  1.0

julia> ea1 = EulerAngleAxis(30*pi/180, [0;1;0]);

julia> ea2 = EulerAngleAxis(45*pi/180, [0;1;0]);

julia> compose_rotation(ea1,ea2)
EulerAngleAxis{Float64}:
  Euler angle:   1.3090 rad ( 75.0000 deg)
   Euler axis: [  0.0000,   1.0000,   0.0000]

julia> Θ1 = EulerAngles(1,2,3,:ZYX);

julia> Θ2 = EulerAngles(-3,-2,-1,:XYZ);

julia> compose_rotation(Θ1, Θ2)
EulerAngles{Float64}:
  R(X):  -0.0000 rad (  -0.0000 deg)
  R(Y):   0.0000 rad (   0.0000 deg)
  R(Z):  -0.0000 rad (  -0.0000 deg)

julia> q1 = angle_to_quat(+pi/3,+pi/4,+pi/5,:ZYX);

julia> q2 = angle_to_quat(-pi/5,-pi/4,-pi/3,:XYZ);

julia> compose_rotation(q1,q2)
Quaternion{Float64}:
  + 1.0 + 0.0.i + 2.0816681711721685e-17.j + 5.551115123125783e-17.k
```
"""
@inline compose_rotation(D::DCM) = D
@inline compose_rotation(D::DCM, Ds::DCM...) = compose_rotation(Ds...)*D

@inline compose_rotation(ea::EulerAngleAxis) = ea
@inline compose_rotation(ea::EulerAngleAxis, eas::EulerAngleAxis...) =
    compose_rotation(eas...)*ea

@inline compose_rotation(Θ::EulerAngles) = Θ
@inline compose_rotation(Θ::EulerAngles, Θs::EulerAngles...) =
    compose_rotation(Θs...)*Θ

@inline compose_rotation(q::Quaternion) = q
@inline compose_rotation(q::Quaternion, qs::Quaternion...) =
    q*compose_rotation(qs...)

# This algorithm was proposed by @Per in
#
#   https://discourse.julialang.org/t/improve-the-performance-of-multiplication-of-an-arbitrary-number-of-matrices/10835/24


# total_cpi_rebase.jl - Función de inflación total con cambio de base sintético.

# El parámetro de maxchanges permite realizar hasta un máximo de cambios de base
# sintéticos en cada VarCPIBase. El valor cero indica que se utilizarán todos
# los cambios de base sintéticos posibles dentro de cada base
Base.@kwdef struct InflationTotalRebaseCPI <: InflationFunction
    period::Int = 60 # períodos para realizar el cambio de base sintético
    maxchanges::Int = 0
end

# Constructor de convenienica para número máximo de cambios de base en cada
# VarCPIBase
InflationTotalRebaseCPI(period::Int) = InflationTotalRebaseCPI(period, 0)

# Nombre de la medida
measure_name(inflfn::InflationTotalRebaseCPI) =
    "Variación interanual IPC con cambios de base sintéticos ($(inflfn.period), $(inflfn.maxchanges))"

# Parámetros
params(totalrebasefn::InflationTotalRebaseCPI) = (totalrebasefn.period, )

# Computar variación intermensual resumen de medida de inflación aplicando
# metodología de cambio de base sintético
function (totalrebasefn::InflationTotalRebaseCPI)(base::VarCPIBase)

    # Número de períodos para realizar cambio de base
    period = totalrebasefn.period

    # Obtener períodos de la base
    T = periods(base)

    # Obtener vector de índices
    startidxs, endidxs = _getranges(T, period, totalrebasefn.maxchanges)

    # Mapear cada rango de índices en el resumen intermensual obtenido con
    # fórmula del IPC
    vinterm = mapreduce(vcat, startidxs, endidxs) do startidx, endidx
        cpi = capitalize(view(base.v, startidx:endidx, :)) * base.w / 100
        varinterm(cpi)
    end

    # Resumen intermensual es la concatenación vertical del resumen intermensual
    # en las distintas bases sintéticas
    vinterm

end

# Función de ayuda para obtener vector de rangos
function _getranges(T, period, max_changes=0)

    # Si hay menos observaciones que el período, devolver el único rango
    T <= period && return 1, T

    # max_changes == 0 => todos los cambios de base posibles
    if max_changes == 0
        blocks = cld(T, period)
    else
        blocks = min(cld(T, period), max_changes+1)
    end

    startidxs = 1:period:T # índices iniciales
    finalidxs = Vector{Int}(undef, blocks) # vector de índices finales
    c = 0 # contador de cambios de base

    # Completar los índices finales
    for j in 1:blocks
        finalidxs[j] = startidxs[j] + period - 1
        c += 1
        if finalidxs[j] >= T || (max_changes != 0 && c > max_changes)
            finalidxs[j] = T
            break
        end
    end

    startidxs[1:blocks], finalidxs
end


using FinancialToolbox, DualNumbers, HyperDualNumbers, FinancialFFT, FinancialMonteCarlo

A = 600.0;
N = 18;

method_cm = CarrMadanMethod(A, N);
S0 = 100.0
K = 100.0;
r = hyper(0.02, 1.0, 1.0, 0.0);

T = 1.0;
d = 0.01
sigma = 0.2
lam = 15.0;
mu1 = 0.03;
sigma1 = 0.02;
zero_rate = ZeroRate(r);

Model = VarianceGammaProcess(sigma, 0.03, 0.01, Underlying(S0, d))

EUData = EuropeanOption(T, K)
Nsim = 100000;
Nstep = 30;
function test_hyper(x, y, toll)
    @test abs(x.value - y.value) < toll
    @test abs(x.epsilon1 - y.epsilon1) < toll
    @test abs(x.epsilon2 - y.epsilon2) < toll
    @test abs(x.epsilon12 - y.epsilon12) < toll
    nothing
end

result_cm = pricer(Model, zero_rate, method_cm, EUData);
method_lewis = LewisMethod(700.0, 200000);
analytic_result = pricer(Model, zero_rate, method_lewis, EUData);
method_cm_lewis = CarrMadanLewisMethod(A, N);
result_cm_lewis = pricer(Model, zero_rate, method_cm_lewis, EUData);
toll = 1e-3;
#test_hyper(analytic_result, result_cm, toll)
@testset "hyper test" begin
    test_hyper(analytic_result, result_cm_lewis, toll)
end


function calculaAnguloRotacion(x,y)

x=x[:,1];
y=y[:,1];

numLados=size(x,1);
largoLados=zeros(numLados,1);
theta=zeros(numLados,1);
cr=zeros(numLados,2);
for i=1:numLados
    if i<numLados
        x1=x[i];
        y1=y[i];
        x2=x[i+1];
        y2=y[i+1];
    else
        x1=x[end];
        y1=y[end];
        x2=x[1];
        y2=y[1];
    end
    d=sqrt((x2-x1)^2+(y2-y1)^2);
    largoLados[i,1]=d;
    theta1=atan((y1-y2)/(x1-x2));
    theta2=abs(theta1)-pi/2;
    if abs(theta1)<=abs(theta2)
        theta[i]=theta1;
    else
        theta[i]=theta2;
    end
    cr[i,:]=[x1 y1];
end
largoLados=largoLados[:,1];
isort = sortperm(largoLados, rev=true);
theta=theta[isort[1:numLados]];
cr=cr[isort[1:numLados],:];

theta, cr

end


module DA

export etkf, ensrf, gaspari_cohn

using Statistics
using LinearAlgebra
using Distributions

function gaspari_cohn(r)
    if 0 <= r < 1
        G = 1 - 5/3*r^2 + 5/8*r^3 + 1/2*r^4 - 1/4*r^5
    elseif 1 <= r < 2
        G = 4 - 5*r + 5/3*r^2 + 5/8*r^3 - 1/2*r^4 + 1/12*r^5 - 2/(3*r)
    elseif r >= 2
        G = 0
    end
    return G
end

function gaspari_cohn_localization(c, D; cyclic=false)
    localization = zeros(D, D)
    for i=1:D
        for j=1:i
            if cyclic
                r = min(mod(i - j, 0:D), mod(j - i, 0:D))/c
            else
                r = abs(i - j)
            end
            localization[i, j] = DA.gaspari_cohn(r)
        end
    end
    return Symmetric(localization, :L)
end

"""
Ensemble transform Kalman filter (ETKF)
"""
function etkf(; E::AbstractMatrix{float_type}, R::Symmetric{float_type},
                R_inv::Symmetric{float_type},
                inflation::float_type=1.0, H::AbstractMatrix,
                y::AbstractVector{float_type}) where {float_type<:AbstractFloat}
    D, m = size(E)

    x_m = mean(E, dims=2)
    X = (E .- x_m)/sqrt(m - 1)

    X = sqrt(inflation)*X

    y_m = H*x_m
    Y = (H*E .- y_m)/sqrt(m - 1)
    Ω = inv(Symmetric(I + Y'*R_inv*Y))
    w = Ω*Y'*R_inv*(y - y_m)

    E = x_m .+ X*(w .+ sqrt(m - 1)*sqrt(Ω))

    return E
end

function ensrf(; E::AbstractMatrix{float_type}, R::Symmetric{float_type},
                 R_inv::Symmetric{float_type},
                 inflation::float_type=1.0, H::AbstractMatrix,
                 y::AbstractVector{float_type},
                 localization=nothing) where {float_type<:AbstractFloat}
    D, m = size(E)

    x_m = mean(E, dims=2)
    A = E .- x_m

    if localization === nothing
        P = inflation*A*A'/(m - 1)
    else
        P = inflation*localization.*(A*A')/(m - 1)
    end

    K = P*H'*inv(H*P*H' + R)
    x_m .+= K*(y - H*x_m)

    E = x_m .+ real((I + P*H'*R_inv*H)^(-1/2))*A

    return E
end

end


using Interpolations

"""
Does a cubic spline interpolation and then takes a derivative across
the row for each column
"""
function numerical_derivative(dat, ts)
    # Take the approximate derivative in the row direction
    interp = []
    dat_grad = zeros(size(dat))
    if ts isa Array
        # Breaks the interpolation function if 'ts' is an array
        dt = ts[2]-ts[1]
        ts = ts[1]:dt:ts[end]
    end

    for i in 1:size(dat,1)
        push!(interp,CubicSplineInterpolation(ts, dat[i,:]))
        for (i2, t) in enumerate(ts)
            dat_grad[i,i2] = Array(Interpolations.gradient(interp[i], t))[1]
        end
    end
    return dat_grad
end

export numerical_derivative


@enum Trit False Maybe True
const trits = (False, Maybe, True)

Base.:!(a::Trit) = a == False ? True : a == Maybe ? Maybe : False
∧(a::Trit, b::Trit) = a == b == True ? True : (a, b) ∋ False ? False : Maybe
∨(a::Trit, b::Trit) = a == b == False ? False : (a, b) ∋ True ? True : Maybe
⊃(a::Trit, b::Trit) = a == False || b == True ? True : (a, b) ∋ Maybe ? Maybe : False
≡(a::Trit, b::Trit) = (a, b) ∋ Maybe ? Maybe : a == b ? True : False

println("Not (!):")
println(join(@sprintf("%10s%s is %5s", "!", t, !t) for t in trits))
println("And (∧):")
for a in trits
    println(join(@sprintf("%10s ∧ %5s is %5s", a, b, a ∧ b) for b in trits))
end
println("Or (∨):")
for a in trits
    println(join(@sprintf("%10s ∨ %5s is %5s", a, b, a ∨ b) for b in trits))
end
println("If Then (⊃):")
for a in trits
    println(join(@sprintf("%10s ⊃ %5s is %5s", a, b, a ⊃ b) for b in trits))
end
println("Equivalent (≡):")
for a in trits
    println(join(@sprintf("%10s ≡ %5s is %5s", a, b, a ≡ b) for b in trits))
end


"""
    change_basis(basis,changeofbasis::AbstractMatrix)

# Examples

Compute the Lagrangian basis associated with a set of nodes

```jldoctests
using LinearAlgebra
using InplaceArrays.Fields
using InplaceArrays.Polynomials

D = 2
order = 1
f = MonomialBasis{D}(Float64,order)

nodes = Point{2,Int}[(0,0),(1,0),(0,1),(1,1)]
change = inv(evaluate(f,nodes))

g = change_basis(f,change)
println(evaluate(g,nodes))

# output
[1.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0; 0.0 0.0 1.0 0.0; 0.0 0.0 0.0 1.0]

```
"""
function change_basis(basis,changeofbasis::AbstractMatrix)
  BasisFromChangeOfBasis(basis,changeofbasis)
end

struct BasisFromChangeOfBasis{B,M} <: Field
  basis::B
  change::M
  function BasisFromChangeOfBasis(basis,change::AbstractMatrix)
    B = typeof(basis)
    M = typeof(change)
    new{B,M}(basis,change)
  end
end

function field_cache(b::BasisFromChangeOfBasis,x)
  cb = field_cache(b.basis,x)
  bx = evaluate_field!(cb,b.basis,x)
  c = CachedArray(bx*b.change)
  (c,cb)
end

function evaluate_field!(cache,b::BasisFromChangeOfBasis,x)
  c, cb = cache
  bx = evaluate_field!(cb,b.basis,x)
  setsize!(c,size(bx))
  mul!(c.array,bx,b.change)
  c.array
end

function gradient_cache(b::BasisFromChangeOfBasis,x)
  cb = gradient_cache(b.basis,x)
  bx = evaluate_gradient!(cb,b.basis,x)
  c = CachedArray(bx*b.change)
  (c,cb)
end

function evaluate_gradient!(cache,b::BasisFromChangeOfBasis,x)
  c, cb = cache
  bx = evaluate_gradient!(cb,b.basis,x)
  setsize!(c,size(bx))
  mul!(c.array,bx,b.change)
  c.array
end

function hessian_cache(b::BasisFromChangeOfBasis,x)
  cb = hessian_cache(b.basis,x)
  bx = evaluate_hessian!(cb,b.basis,x)
  c = CachedArray(bx*b.change)
  (c,cb)
end

function evaluate_hessian!(cache,b::BasisFromChangeOfBasis,x)
  c, cb = cache
  bx = evaluate_hessian!(cb,b.basis,x)
  setsize!(c,size(bx))
  mul!(c.array,bx,b.change)
  c.array
end


using Base.Test
using Tensors

# Single mode QR

t = rand([Mode(k,2) for k in 1:4])
q,r = qr(t, 3)
@test t ≈ q*r

# Multiple modes QR

t = rand([Mode(k,2) for k in 1:4])
q,r = qr(t, Any[3,1], 5)
@test t ≈ q*r

# Single mode SVD

t = rand([Mode(k,2) for k in 1:4])
u,s,v = svd(t, 3)
@test t ≈ u*scale(s,v)

# Multiple modes SVD

t = rand([Mode(k,2) for k in 1:4])
u,s,v = svd(t, Any[4,2], 5)
@test t ≈ scale(u,s)*v


import FastTransforms: chebyshevpoints#, transform, itransform
import OffsetArrays: OffsetArray

@inline function _chebyshev(::Val{1}, j::Int, x::T) where {T}
    if j < 0
        return zero(T)
    elseif j == 0
        return one(T)
    elseif j == 1
        return x
    else
        return _chebyshev(Val(1), j-1, x) * 2x - _chebyshev(Val(1), j-2, x)
    end
end

@inline function _chebyshev(::Val{2}, j::Int, x::T) where {T}
    if j < 0
        return zero(T)
    elseif j == 0
        return one(T)
    elseif j == 1
        return 2x
    else
        return _chebyshev(Val(2), j-1, x) * 2x - _chebyshev(Val(2), j-2, x)
    end
end

"""
Chebyshev basis on the interval [-1..+1].
"""
struct Chebyshev{kind, T, BT, XT <: AbstractVector{T}} <: Basis{T}
    b::BT
    x::XT

    function Chebyshev{kind, T}(n::Int) where {kind, T}
        p = n-1
        x = chebyshevpoints(T, n, Val(kind))
        b = OffsetArray([y -> _chebyshev(Val(kind), i, y) for i in 0:p], 0:p)
        new{kind, T, typeof(b), typeof(x)}(b, x)
    end

    Chebyshev{kind}(::Type{T}, n::Int) where {kind, T} = Chebyshev{kind, T}(n)
    Chebyshev{kind}(n::Int) where {kind} = Chebyshev{kind, Float64}(n)
end

Chebyshev(::Type{T}, n::Int, ::Val{kind}) where {kind, T} = Chebyshev{kind, T}(n)
Chebyshev(n::Int, ::Val{kind}) where {kind} = Chebyshev(Float64, n, Val(kind))

const ChebyshevT = Chebyshev{1}
const ChebyshevU = Chebyshev{2}

(C::Chebyshev)(x::Number, j::Int) = C.b[j](x)

basis(C::Chebyshev) = C.b
nodes(C::Chebyshev) = C.x
nbasis(C::Chebyshev) = length(basis(C))
nnodes(C::Chebyshev) = length(nodes(C))
order(C::Chebyshev)  = nnodes(C)
degree(C::Chebyshev) = nnodes(C) - 1

Base.eltype(::Chebyshev{kind,T}) where {kind,T} = T
Base.eachindex(C::Chebyshev) = eachindex(C.b)
Base.axes(C::Chebyshev) = (Inclusion(0..1), eachindex(C))
ContinuumArrays.grid(C::Chebyshev) = nodes(C)

Base.hash(C::Chebyshev{kind}, h::UInt) where {kind} = hash(C.x, hash(kind, h))
Base.:(==)(C1::Chebyshev{kind1}, C2::Chebyshev{kind2}) where {kind1,kind2} = (C1.x == C2.x && kind1 == kind2)
Base.isequal(C1::Chebyshev{kind1,T1}, C2::Chebyshev{kind2,T2}) where {kind1,T1,kind2,T2} = (T1 == T2 && C1 == C2)
Base.isapprox(C1::Chebyshev, C2::Chebyshev; kwargs...) = isapprox(C1.x, C2.x; kwargs...)

Base.getindex(C::Chebyshev, x::Number, j::Int) = C(x,j)
Base.getindex(C::Chebyshev, x::Number,  ::Colon) = [b(x) for b in C.b]
Base.getindex(C::Chebyshev, X::AbstractVector, j::Int) = C.(X,j)
Base.getindex(C::Chebyshev, X::AbstractVector,  ::Colon) = [b(x) for x in X, b in C.b]

## Derivative

@simplify *(D::Derivative, C::Chebyshev) = Mul(D,C)

const ChebyshevDerivative  = QMul2{<:Derivative,<:Chebyshev}
const ChebyshevTDerivative = QMul2{<:Derivative,<:ChebyshevT}
const ChebyshevUDerivative = QMul2{<:Derivative,<:ChebyshevU}

function _eval(D::ChebyshevTDerivative, x::DT, i::Int) where {DT}
    local C = D.B
    local x̃ = promote_type(eltype(C), DT)(x)
    @assert i ≥ 0 && i < nbasis(C)
    _chebyshev(Val(2), i-1, x̃) * i
end

function _eval(D::ChebyshevUDerivative, x::DT, i::Int) where {DT}
    local C = D.B
    local x̃ = promote_type(eltype(C), DT)(x)
    @assert i ≥ 0 && i < nbasis(C)
    ( _chebyshev(Val(1), i+1, x̃) * (i+1) - _chebyshev(Val(2), i, x̃) * x̃ ) / (x̃^2 - 1)
end

Base.getindex(D::ChebyshevDerivative, x::Number, j::Integer) = _eval(D, x, j)
Base.getindex(D::ChebyshevDerivative, x::Number,  ::Colon) = [_eval(D, x, j) for j in eachindex(D.B)]
Base.getindex(D::ChebyshevDerivative, X::AbstractVector, j::Integer) = [_eval(D, x, j) for x in X]
Base.getindex(D::ChebyshevDerivative, X::AbstractVector,  ::Colon) = [_eval(D, x, j) for x in X, j in eachindex(D.B)]

Base.adjoint(C::Chebyshev) = Derivative(axes(C,1)) * C


if !isdefined(Main, :CustomStencil)
    include("../module/CustomStencil.jl")
    using Main.CustomStencil
end
include("misc.jl")
include("cpu_stencils.jl")

nx = 9
ny = 9
nz = 8
dx = 1<<(nx)
dy = 1<<(ny)
dz = 1<<(nz)

original_data = CreateData(dx,dy,dz)
## Stencil Definition
t_cases = [32 16 8 4]
i = 16
radius = i
t_steps = Int(max(t_cases...)/i)
coefs = [Float64(round(1/i, digits=2)) for i = (radius+1):-1:1]
star_stencil = @def_stencil_expression c[0]D[x,y,z] + @sum(i, 1, 16, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))
st_def = CreateStencilDefinition(star_stencil, coefs)
st_inst = NewStencilInstance(st_def)

## Input Data size Definition
data = PadData(radius, original_data)
@time dout = ApplyStencil(st_inst, data, t_steps)
gpu_out = Array(dout)

# t_cases = [32 16 8 4]
# i = t_cases[2]
# radius = i
# t_steps = Int(max(t_cases...)/i)
# coefs = [Float64(round(1/i, digits=2)) for i = (radius+1):-1:1]
# star_stencil = @def_stencil_expression c[0]D[x,y,z] + @sum(i, 1, 16, c[i]*(
#             D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
#             D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))
# st_def = CreateStencilDefinition(star_stencil, coefs)
# st_inst = NewStencilInstance(st_def)
#
# ## Input Data size Definition
# data = PadData(radius, original_data)
# @time dout = ApplyStencil(st_inst, data, t_steps)
# gpu_out = Array(dout)
#
# t_cases = [32 16 8 4]
i = t_cases[3]
radius = i
t_steps = Int(max(t_cases...)/i)
coefs = [Float64(round(1/i, digits=2)) for i = (radius+1):-1:1]
star_stencil = @def_stencil_expression c[0]D[x,y,z] + @sum(i, 1, 8, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))
st_def = CreateStencilDefinition(star_stencil, coefs)
st_inst = NewStencilInstance(st_def)

## Input Data size Definition
data = PadData(radius, original_data)
@time dout = ApplyStencil(st_inst, data, t_steps)
gpu_out = Array(dout)

t_cases = [32 16 8 4]
i = t_cases[4]
radius = i
t_steps = Int(max(t_cases...)/i)
coefs = [Float64(round(1/i, digits=2)) for i = (radius+1):-1:1]
star_stencil = @def_stencil_expression c[0]D[x,y,z] + @sum(i, 1, 4, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))
st_def = CreateStencilDefinition(star_stencil, coefs)
st_inst = NewStencilInstance(st_def)

## Input Data size Definition
data = PadData(radius, original_data)
@time dout = ApplyStencil(st_inst, data, t_steps)
gpu_out = Array(dout)


z = katz_centrality(g5, 0.4)
@test round(z, 2) == [0.32, 0.44, 0.62, 0.56]


#===========================================
Tests on the RTN generation routine function
===========================================#

#=
Here we calculate the autocorrelation functions of the two methods of generating RTN histories. The functions need a huge time in input (40000. is ok) in order to have ergodicity, and they return the autocorrelation array [<RTN(0)RTN(0)>-<RTN(0)><RTN(0)>,<RTN(1)RTN(0)>-<RTN(1)><RTN(0)>,...] up to the lag corresponding to an evolution time of tc=4./mu.
=#
using Test
using Statistics
using StatsBase

function autocorrelation(noisefunction, time, mu, dt)
    t = range(0., stop=time, step=dt)

    signal = noisefunction(t, mu)[:]
    tc = 4. / mu
    dt = Float64(t.step)

    lags = [x for x = 0:(Int(fld(tc, dt))-1)]
    autoC = autocov(signal, lags)

    tvec = range(0., stop=tc, length=Int(fld(tc,dt)))
    return (tvec, autoC)
end

function spectrum(t, signal)
    n = length(signal)
    p = fft(signal)
    sampFreq = 1/(t[2] - t[1])
    nUniquePts = ceil(Int, (n+1)/2)
    p = p[1:nUniquePts]
    p = abs2.(p) / n^2
    # the fourier transform of the tone returned by the fft function contains both magnitude and phase information and is given in a complex representation (i.e. returns complex numbers). By taking the absolute value of the fourier transform we get the information about the magnitude of the frequency components.

    # odd nfft excludes Nyquist point
    if n % 2 > 0
        p[2 : length(p)] = p[2:length(p)]*2 # we've got odd number of   points fft
    else
        p[2: (length(p) -1)] = p[2: (length(p) -1)]*2 # we've got even number of points fft
    end

    freqArray = (0:(nUniquePts-1)) * (sampFreq / n)
    return (freqArray, p)
end

function test_autocorrelation(noisefunction, mu; atol=1e-1)
    (t, ac) = autocorrelation(noisefunction, 20000. /mu, mu, .05 / mu)
    @test ac ≈ exp.(-2 * mu * t) atol = atol
end

# Test that the mean is zero
@testset "RTN mean" begin
    @testset "γ = $γ" for γ in [0.01, 0.1 , 1., 10.]
        @test mean(generateRTN(0. : .05 / γ : 1000 / γ, γ, 100)) ≈ 0 atol=5e-3
    end
end

# Test that the autocorrelation function is the decaying exponential
@testset "RTN autocorrelation" begin
    @testset "γ = $γ" for γ in [0.01, 0.1 , 1., 10.]
        @testset "RTNgenTime" for i in 1:100
            test_autocorrelation(generateRTN, γ, atol=1e-1)
        end
    end
end


function mimo_data(seed, m, n)
    Random.seed!(seed)
    # Channel
    H = randn((m, n))
    # Gaussian noise
    v = randn((m, 1))
    # True signal
    s = rand([-1, 1], n)
    # Received signal
    sigma = .0001
    y = H * s + sigma * v
    L = [hcat(H' * H, -H' * y); hcat(-y' * H, y' * y)]
    return s, H, y, L
end

function mimo_eval(s, H, y, L, XX)
    x_hat = sign.(XX[1:end-1, end])
    rank = length([eig for eig in eigen(XX).values if eig > 1e-7])
    @show decode_error = sum(abs.(x_hat - s))
    @show rank
    @show norm(y - H * x_hat)
    @show norm(y - H * s)
    @show tr(L * XX)
    return nothing
end


using ConvDiffMIPDECO
using jInv.Mesh
using jInv.ForwardShare
using jInv.InverseSolve
using jInv.LinearSolvers
using Statistics

function runInv()
domain = [0 3. -0.5 1.5]
n      = [60 40]
M      = getRegularMesh(domain,n)
display(M)

ftrue = getFEMsource2D(M)
#pFor  = getConvDiffFEMParam(M,sig=.01,Ainv = getMUMPSsolver())
 pFor  = getConvDiffFEMParam(M,sig=.01,Ainv = getJuliaSolver())
dobs, = getData(vec(ftrue),pFor)
dobs += 0.4*randn(size(dobs))*mean(abs.(vec(dobs))) # add noise

sigback     = 0.0
Wt          = ones(prod(n.+1))
pMis        = getMisfitParam(pFor,Wt,dobs,SSDFun);

# configure regularization
alpha	   	= 1e-2;
reg         = (m,mr,M,I=1.0) -> wTVReg(m,mr,M,eps=1e-6)
mref        = zeros(M.nc);
# configure optimization
HesPrec     = getSSORRegularizationPreconditioner(1.0,1e-15,50);
cgit       	= 5;
pcgTol     	= 1e-1;
maxIter    	= 5;
minUpdate 	= 1e-3;

# use identity model
modFun      = identityMod
boundsLow  	= 0*ones(prod(n));
boundsHigh 	= 1*ones(prod(n));

# # use level set approach
modFun      = boundMod
boundsLow  	= -Inf*ones(prod(n));
boundsHigh 	= Inf *ones(prod(n));

maxStep		= 0.1*maximum(boundsHigh);

pInv = getInverseParam(M,modFun,reg,alpha,mref,
         boundsLow,boundsHigh,maxStep=maxStep,pcgMaxIter=cgit,pcgTol=pcgTol,
         minUpdate=minUpdate, maxIter = maxIter,HesPrec=HesPrec);

mc,Dc,flag,His = projGNCG(0*ftrue[:],pInv,pMis);
fc = pInv.modelfun(mc)[1];

println("time grad misfit: $(sum(His.timeGradMisfit))")
println("time GN solve:    $(sum(His.timeLinSol))")
println("time misfit       $(sum(His.timeMisfit))")
println("time reg          $(sum(His.timeReg))")
end

runInv()

# using ProfileView
# Profile.clear()
#
# @profile runInv()
#
# ProfileView.view()


using EmpiricalBayes
using Test

using Distributions

marginal_grid = collect(range(-6, stop=6, length=1001));
true_dist = Normal(0,1)
d_true = NormalConvolutionProblem(true_dist, marginal_grid);

m = 7_000
Xs = rand(d_true, m)

estimate(Xs, BrownGreenshtein, PosteriorTarget(PosteriorMeanNumerator(2.0)))
posterior_stats(d_true, PosteriorTarget(PosteriorMeanNumerator(2.0)))

estimate(Xs, BrownGreenshtein, PosteriorTarget(PosteriorMeanNumerator(4.0)))
posterior_stats(d_true, PosteriorTarget(PosteriorMeanNumerator(4.0)))

estimate(Xs, BrownGreenshtein, PosteriorTarget(PosteriorMeanNumerator(3.0)))
posterior_stats(d_true, PosteriorTarget(PosteriorMeanNumerator(3.0)))


@testset "WreathProducts" begin
   S_3 = PermutationGroup(3)
   S_2 = PermutationGroup(2)
   b = perm"(1,2,3)"
   a = perm"(1,2)"

   @testset "Constructors" begin
      @test Groups.WreathProduct(S_2, S_3) isa AbstractAlgebra.Group
      B3 = Groups.WreathProduct(S_2, S_3)
      @test B3 isa Groups.WreathProduct
      @test B3 isa WreathProduct{3, Generic.PermGroup{Int}, Generic.PermGroup{Int}}

      aa = Groups.DirectPowerGroupElem((a^0 ,a, a^2))

      @test Groups.WreathProductElem(aa, b) isa AbstractAlgebra.GroupElem
      x = Groups.WreathProductElem(aa, b)
      @test x isa Groups.WreathProductElem
      @test x isa
         Groups.WreathProductElem{3, perm{Int}, perm{Int}}

      @test B3.N == Groups.DirectPowerGroup(S_2, 3)
      @test B3.P == S_3

      @test B3(aa, b) == Groups.WreathProductElem(aa, b)
      w = B3(aa, b)
      @test B3(w) == w
      @test B3(b) == Groups.WreathProductElem(B3.N(), b)
      @test B3(aa) == Groups.WreathProductElem(aa, S_3())

      @test B3((a^0 ,a, a^2), b) isa WreathProductElem

      @test B3((a^0 ,a, a^2), b) == B3(aa, b)
   end

   @testset "Types" begin
      B3 = Groups.WreathProduct(S_2, S_3)

      @test elem_type(B3) == Groups.WreathProductElem{3, perm{Int}, perm{Int}}

      @test parent_type(typeof(B3())) == Groups.WreathProduct{3, parent_type(typeof(B3.N.group())), Generic.PermGroup{Int}}

      @test parent(B3()) == Groups.WreathProduct(S_2,S_3)
      @test parent(B3()) == B3
   end

   @testset "Basic operations on WreathProductElem" begin
      aa = Groups.DirectPowerGroupElem((a^0 ,a, a^2))
      B3 = Groups.WreathProduct(S_2, S_3)
      g = B3(aa, b)

      @test g.p == b
      @test g.n == DirectPowerGroupElem(aa.elts)

      h = deepcopy(g)
      @test h == g
      @test !(g === h)

      g = B3(Groups.DirectPowerGroupElem((a ,a, a^2)), g.p)

      @test g.n[1] == parent(g.n[1])(a)
      @test g != h

      @test hash(g) != hash(h)
   end

   @testset "Group arithmetic" begin
      B4 = Groups.WreathProduct(AdditiveGroup(GF(3)), PermutationGroup(4))

      x = B4((0,1,2,0), perm"(1,2,3)(4)")
      @test inv(x) == B4((1,0,2,0), perm"(1,3,2)(4)")

      y = B4((1,0,1,2), perm"(1,4)(2,3)")
      @test inv(y) == B4((1,2,0,2), perm"(1,4)(2,3)")

      @test x*y == B4((0,2,0,2), perm"(1,3,4)(2)")

      @test y*x == B4((1,2,2,2), perm"(1,4,2)(3)")

      @test inv(x)*y == B4((2,1,2,2), perm"(1,2,4)(3)")

      @test y*inv(x) == B4((1,2,1,0), perm"(1,4,3)(2)")

      @test (x*y)^6 == ((x*y)^2)^3

   end

   @testset "Iteration" begin
      B3_a = Groups.WreathProduct(AdditiveGroup(GF(3)), S_3)
      @test order(B3_a) == 3^3*6
      @test collect(B3_a) isa Vector{
      WreathProductElem{3, AddGrpElem{AbstractAlgebra.gfelem{Int}}, perm{Int}}}

      B3_m = Groups.WreathProduct(MultiplicativeGroup(GF(3)), S_3)
      @test order(B3_m) == 2^3*6
      @test collect(B3_m) isa Vector{
      WreathProductElem{3, MltGrpElem{AbstractAlgebra.gfelem{Int}}, perm{Int}}}

      @test length(Set([B3_a, B3_m, B3_a])) == 2

      Wr = WreathProduct(PermutationGroup(2),PermutationGroup(4))

      elts = collect(Wr)
      @test elts isa Vector{Groups.WreathProductElem{4, perm{Int}, perm{Int}}}
      @test order(Wr) == 2^4*factorial(4)

      @test length(elts) == order(Wr)
      @test all([g*inv(g) == Wr() for g in elts])
      @test all(inv(g*h) == inv(h)*inv(g) for g in elts for h in elts)
   end

   @testset "Misc" begin
      B3_a = Groups.WreathProduct(AdditiveGroup(GF(3)), S_3)
      @test string(B3_a) == "Wreath Product of The additive group of Finite field F_3 by Permutation group over 3 elements"

      @test string(B3_a(perm"(1,3)")) == "([0,0,0]≀(1,3))"

      B3_m = Groups.WreathProduct(MultiplicativeGroup(GF(3)), S_3)
      @test string(B3_m) == "Wreath Product of The multiplicative group of Finite field F_3 by Permutation group over 3 elements"

      @test string(B3_m(perm"(1,3)")) == "([1,1,1]≀(1,3))"

   end

end


function solve_sai(U0, t_grid, dt, lambda::Number)

    grid_len = length(t_grid)
    s = t_grid[1]

    t_outer = t_grid .- t_grid'

    K0 = make_volterra_kernel(U0)
    pert = make_volterra_kernel(t_outer)

    #c1 = airyai(s) * exp(-lambda * sqrt(s))
    #c2 = -sqrt(s) * c1

    c1 = airyai(s + lambda)
    c2 = airyaiprime(s + lambda)

    v1 = ones(grid_len)
    v2 = U0[:,1]
    v3 = t_outer[:,1]

    v = (c2 * v1) + (c1 * v2) + (c1 * lambda * v3)

    A = I - dt/2 * K0
    B = -dt/2 * pert

    sai_path = (A + lambda * B) \ v

    trapz_matrix = make_volterra_kernel(ones(size(U0))) * dt/2

    return (
        SAi = trapz_matrix * sai_path,
        DSAi = sai_path
    )
end

function solve_phi(U0, t_grid, dt, lambda_grid; verbose=false)

    grid_len = length(t_grid)

    s = t_grid[1]

    t_outer = t_grid .- t_grid'

    K0 = make_volterra_kernel(U0)
    pert = make_volterra_kernel(t_outer)

    u_deriv = [zeros(grid_len - 1); 1]
    u_full = dt/2 * ([0; ones(grid_len - 1)] + [ones(grid_len - 1); 0]);

    u = [u_deriv u_full]

    v1 = ones(grid_len)
    v2 = U0[:,1]
    v3 = t_outer[:,1]

    A = I - dt/2 * K0
    B = -dt/2 * pert

    BA_inv = B / A
    uA_inv = A' \ u

    N = BA_inv

    Nv_curr = [v1 v2 v3]

    # the 2*3 matrix u^t A^-1 v for u1-u2, v1-v3 for each power
    coefs_v = Array{Complex{Float64}, 3}(undef, (grid_len, 2, 3))

    for pow in 1:grid_len

        coefs_v[pow,:,:] = uA_inv' * Nv_curr * (-1)^(pow + 1)
        Nv_curr = N * Nv_curr

     end

    log_coefs_v = log.(coefs_v)

    phi_deriv_grid = undef_copy(lambda_grid) * 1.0
    phi_full_grid = undef_copy(lambda_grid) * 1.0

    for (i, lambda) in pairs(lambda_grid)

        c1 = airyai(s + lambda)
        c2 = airyaiprime(s + lambda)

        log_lambda = log(Complex(lambda))

        # Make the array [0, log(lambda), 2 log(lambda), ...]
        log_lambda_pows = accumulate(+, [0; log_lambda * ones(grid_len - 1)])

        # Make the array [C_0(lambda) lambda^0, C_1(lambda) lambda^1, ...]
        sums_deriv = sum(exp.(log_lambda_pows .+ log_coefs_v[:,1,:]), dims = 1)
        sums_full = sum(exp.(log_lambda_pows .+ log_coefs_v[:,2,:]), dims = 1)

        if any(isnan.(sums_deriv)) | any(isnan.(sums_full))
            ss = log_lambda_pows .+ log_coefs_v + maxval_log_cum
            display(ss)
            display(cumsum(ss, dims = 1))
            display(sums)
            throw("NAN sums")
        end

        # If lamba grid is purely real, make phi values real
        if eltype(lambda_grid) <: Real
            #@assert maximum(abs.(imag.(sums_deriv))/abs.(sums_deriv)) < 1e-6
            #@assert maximum(abs.(imag.(sums_full))/abs.(sums_deriv)) < 1e-6

            sums_deriv = real.(sums_deriv)
            sums_full = real.(sums_full)
        end

        phi_deriv_grid[i] = (c2 * sums_deriv[1,1] + c1 * sums_deriv[1,2] + c1 * lambda * sums_deriv[1,3]) + c1
        phi_full_grid[i] = (c2 * sums_full[1,1] + c1 * sums_full[1,2] + c1 * lambda * sums_full[1,3]) + c1

        if isnan(phi_deriv_grid[i]) | isnan(phi_full_grid[i])
            println(i)
            throw("NAN phi")
        end
    end

    if verbose
        return (
            SAi = phi_full_grid,
            DSAi = phi_deriv_grid,
            Vcoefs = coefs_v,
            LogVcoefs = log_coefs_v,
            MaxLogs = maxval_logs)
    else
        return (
            SAi = phi_full_grid,
            DSAi = phi_deriv_grid)
    end
end;


# This file includes the plugin for calculation of fast fourier transform of data
using FFTW

"""
    Fft(dims::Int)

Constructs an `Fft` plugin. The  [`process(plg::Fft, x)`](@ref) function performes an `fft` operatinon along `dims` of `x`. See also: [`fft`](https://juliamath.github.io/AbstractFFTs.jl/stable/api/#AbstractFFTs.fft)
"""
struct Fft <: AbstractPlugin
    dims::Int
end
Fft(;dims::Int=1) = Fft(dims)

show(io::IO, plg::Fft) = print(io, "Fft(dims:$(plg.dims))")

"""
    process(plg::Fft, x)

Performes an `fft` transformation for the input data `x`.

# Example
```julia
julia> x = collect(reshape(1:16, 4,4))
4×4 Array{Int64,2}:
 1  5   9  13
 2  6  10  14
 3  7  11  15
 4  8  12  16

julia> plg = Plugins.Fft(dims=1)
Fft(dims:1)

julia> process(plg, x)
4×4 Array{Complex{Float64},2}:
 10.0+0.0im  26.0+0.0im  42.0+0.0im  58.0+0.0im
 -2.0+2.0im  -2.0+2.0im  -2.0+2.0im  -2.0+2.0im
 -2.0+0.0im  -2.0+0.0im  -2.0+0.0im  -2.0+0.0im
 -2.0-2.0im  -2.0-2.0im  -2.0-2.0im  -2.0-2.0im
```
"""
process(plg::Fft, x) = fft(x, plg.dims)


@defcomp grosseconomy begin
	GA		= Variable(index=[time])	#Growth rate of productivity
	AL		= Variable(index=[time])	#Level of total factor productivity
    K       = Variable(index=[time])    #Capital stock (trillions 2010 US dollars)
    YGROSS  = Variable(index=[time])    #Gross world product GROSS of abatement and damages (trillions 2010 USD per year)

	a0		= Parameter()				#Initial level of total factor productivity
	ga0		= Parameter()				#Initial growth rate for TFP per 5 years
	dela	= Parameter()				#Decline rate of TFP per 5 years
    I       = Parameter(index=[time])   #Investment (trillions 2010 USD per year)
    l       = Parameter(index=[time])   #Level of population and labor
    dk      = Parameter()               #Depreciation rate on capital (per year)
    gama    = Parameter()               #Capital elasticity in production function
    k0      = Parameter()               #Initial capital value (trill 2010 USD)

    function run_timestep(p, v, d, t)
		#Define function for GA
		for tm = 1:100
			v.GA[tm] = p.ga0 * exp(-p.dela * 5 * (tm-1))
		end

		#Define function for AL
		if is_first(t)
			v.AL[t] = p.a0
		else
			v.AL[t] = v.AL[t-1]/(1 - v.GA[t-1])
		end

        #Define function for K
        if is_first(t)
            v.K[t] = p.k0
        else
            v.K[t] = (1 - p.dk)^5 * v.K[t-1] + 5 * p.I[t-1]
        end

        #Define function for YGROSS
        v.YGROSS[t] = (v.AL[t] * (p.l[t]/1000)^(1-p.gama)) * (v.K[t]^p.gama)
    end
end