text:
"""
    read_parquet(path; kwargs...)

Returns the table contained in the parquet file or dataset (partitioned parquet files in a folder)
in a Tables.jl compatible format.

Options:
- `rows`: The row range to iterate through, all rows by default. Applicable only when reading a single file.
- `filter`: Filter function to apply while loading only a subset of partitions from a dataset.
- `batchsize`: Maximum number of rows to read in each batch (default: row count of first row group). Applied only when reading a single file, and to each file when reading a dataset.
- `use_threads`: Whether to use threads while reading the file; applicable only for Julia v1.3 and later and switched on by default if julia processes is started with multiple threads.
- `column_generator`: Function to generate a partitioned column when not found in the partitioned table. Parameters provided to the function: table, column index, length of column to generate. Default implementation determines column values from the table path.

One can easily convert the returned object to any Tables.jl compatible table e.g. DataFrames.DataFrame via

```
using DataFrames
df = DataFrame(read_parquet(path))
```
"""
function read_parquet(path; kwargs...)
    if isdir(path)
        Parquet.Dataset(path; kwargs...)
    elseif isfile(path)
        Parquet.Table(path; kwargs...)
    else
        error("Invalid path to parquet file or dataset - $path")
    end
end

"""
    Parquet.Table(path; kwargs...)

Returns the table contained in the parquet file in a Tables.jl compatible format.

Options:
- `rows`: The row range to iterate through, all rows by default.
- `batchsize`: Maximum number of rows to read in each batch (default: row count of first row group).
- `use_threads`: Whether to use threads while reading the file; applicable only for Julia v1.3 and later and switched on by default if julia processes is started with multiple threads.

One can easily convert the returned object to any Tables.jl compatible table e.g. DataFrames.DataFrame via

```
using DataFrames
df = DataFrame(read_parquet(path))
```
"""
struct Table <: Tables.AbstractColumns
    path::String
    ncols::Int
    rows::Union{Nothing,UnitRange}
    batchsize::Union{Nothing,Signed}
    use_threads::Bool
    parfile::Parquet.File
    schema::Tables.Schema
    lookup::Dict{Symbol, Int} # map column name => index
    columns::Vector{AbstractVector}
    column_generator::Function

    Table(path, sch::Tables.Schema; kwargs...) = Table(path, Parquet.File(path), sch; kwargs...)
    function Table(path, parfile::Parquet.File=Parquet.File(path), sch::Tables.Schema=tables_schema(parfile);
            rows::Union{Nothing,UnitRange}=nothing,
            batchsize::Union{Nothing,Signed}=nothing,
            column_generator::Function=column_generator,
            use_threads::Bool=(nthreads() > 1))
        ncols = length(sch.names)
        lookup = Dict{Symbol, Int}(nm => i for (i, nm) in enumerate(sch.names))
        new(path, ncols, rows, batchsize, use_threads, parfile, sch, lookup, AbstractVector[], column_generator)
    end
end

function close(table::Table)
    empty!(getfield(table, :columns))
    close(getfield(table, :parfile))
end

function column_generator(table::Table, colidx::Int, len::Int)
    schema = getfield(table, :schema)
    coltype = schema.types[colidx]
    missingval = nonmissingtype(coltype) === coltype ? undef : missing
    Array{coltype}(missingval, len)
end

"""
Represents one partition of the parquet file.
Typically a row group, but could be any other unit as mentioned while opening the table.
"""
struct TablePartition <: Tables.AbstractColumns
    table::Table
    columns::Vector{AbstractVector}
end

"""
Iterator to iterate over partitions of a parquet file, returned by the `Tables.partitions(table)` method.
Each partition is typically a row group, but could be any other unit as mentioned while opening the table.
"""
struct TablePartitions
    table::Table
    ncols::Int
    schema::Tables.Schema
    cursor::BatchedColumnsCursor

    function TablePartitions(table::Table)
        new(table, getfield(table, :ncols), getfield(table, :schema), cursor(table))
    end
end
length(tp::TablePartitions) = length(tp.cursor)
function iterated_partition(partitions::TablePartitions, iterresult)
    (iterresult === nothing) && (return nothing)
    chunk, batchid = iterresult
    columns = AbstractVector[]
    for colidx in 1:partitions.ncols
        colname = partitions.schema.names[colidx]
        if hasproperty(chunk, colname)
            push!(columns, getproperty(chunk, colname))
        else
            generator = getfield(partitions.table, :column_generator)
            push!(columns, generator(partitions.table, colidx, length(first(chunk))))
        end
    end
    TablePartition(partitions.table, columns), batchid
end
Base.iterate(partitions::TablePartitions, batchid) = iterated_partition(partitions, iterate(partitions.cursor, batchid))
Base.iterate(partitions::TablePartitions) = iterated_partition(partitions, iterate(partitions.cursor))

function cursor(table::Table)
    kwargs = Dict{Symbol,Any}(:use_threads => getfield(table, :use_threads), :reusebuffer => false)
    (getfield(table, :rows) === nothing) || (kwargs[:rows] = getfield(table, :rows))
    (getfield(table, :batchsize) === nothing) || (kwargs[:batchsize] = getfield(table, :batchsize))
    BatchedColumnsCursor(getfield(table, :parfile); kwargs...)
end

loaded(table::Table) = !isempty(getfield(table, :columns))
load(table::Table) = load(table, cursor(table))
function load(table::Table, colcursor::BatchedColumnsCursor)
    chunks = [chunk for chunk in colcursor]
    ncols = getfield(table, :ncols)
    columns = getfield(table, :columns)
    schema = getfield(table, :schema)
    generator = getfield(table, :column_generator)

    empty!(columns)
    nchunks = length(chunks)
    if nchunks == 1
        chunk = chunks[1]
        for colidx in 1:ncols
            colname = schema.names[colidx]
            if hasproperty(chunk, colname)
                push!(columns, getproperty(chunk, colname))
            else
                push!(columns, generator(table, colidx, length(first(chunk))))
            end
        end
    elseif nchunks > 1
        for colidx in 1:ncols
            colname = schema.names[colidx]
            coltype = schema.types[colidx]
            vecs = Vector{coltype}[]
            for chunkidx in 1:nchunks
                chunk = chunks[chunkidx]
                if hasproperty(chunk, colname)
                    push!(vecs, getproperty(chunk, colname))
                else
                    push!(vecs, generator(table, colidx, length(first(chunk))))
                end
            end
            push!(columns, ChainedVector(vecs))
        end
    else
        schema = getfield(table, :schema)
        coltypes = schema.types
        for colidx in 1:ncols
            push!(columns, (coltypes[colidx])[])
        end
    end
    nothing
end

Tables.istable(::Table) = true
Tables.columnaccess(::Table) = true
Tables.schema(t::Table) = getfield(t, :schema)
Tables.columnnames(t::Table) = getfield(t, :schema).names
Tables.columns(t::Table) = Tables.CopiedColumns(t)
Tables.getcolumn(t::Table, nm::Symbol) = Tables.getcolumn(t, getfield(t, :lookup)[nm])
function Tables.getcolumn(t::Table, i::Int)
    loaded(t) || load(t)
    getfield(t, :columns)[i]
end
Tables.partitions(t::Table) = TablePartitions(t)

Tables.istable(::TablePartition) = true
Tables.columnaccess(::TablePartition) = true
Tables.schema(tp::TablePartition) = Tables.schema(getfield(tp, :table))
Tables.columnnames(tp::TablePartition) = Tables.columnnames(getfield(tp, :table))
Tables.columns(tp::TablePartition) = Tables.CopiedColumns(tp)
Tables.getcolumn(tp::TablePartition, nm::Symbol) = Tables.getcolumn(tp, getfield(getfield(tp, :table), :lookup)[nm])
Tables.getcolumn(tp::TablePartition, i::Int) = getfield(tp, :columns)[i]


text:
function plotCirc10BA(fg_::AbstractDFG,
                      fg::AbstractDFG;
                      pointsList_::Vector{Symbol}=Symbol[],
                      pointsList::Vector{Symbol}=Symbol[],
                      filepath::AbstractString="/tmp/test.pdf",
                      lineColor_::String="gray40",
                      lineColor::String="green",
                      ellipseList_=ls(fg_),
                      ellipseColor_::String="gray60",
                      ellipseColor::String="gray20",
                      drawEllipse::Bool=true,
                      contourList_::Vector{Symbol}=Symbol[],
                      contourColor_::String="red",
                      contourLineWidth_=1pt,
                      levels::Int=3,
                      show::Bool=true,
                      width=10cm, height=8cm)
  #
  plfl4 = drawPosesLandms(fg_, spscale=1.5, manualColor=lineColor_, point_size=4pt, drawhist=false, contour=false, levels=2, lbls=false)

  for vsym in ellipseList_
    PL = plotCovEllipseLayer(fg_, vsym, points=vsym in pointsList_ || vsym in pointsList, ellipseColor=ellipseColor_)
    if vsym != :x10
      for ll in PL
        push!(plfl4.layers, ll)
      end
    end
  end

  if exists(fg, :x9)
    PL = plotCovEllipseLayer(fg, :x9, points=:x9 in pointsList, ellipseColor=ellipseColor, drawEllipse=drawEllipse)
    for ll in PL
      push!(plfl4.layers, ll)
    end
  end

  plfl5 = drawPosesLandms(fg, spscale=1.5, point_size=3pt, manualColor=lineColor, drawhist=false, contour=false, line_width=2pt)

  for ll in plfl4.layers
    push!(plfl5.layers, ll)
  end

  for conl in contourList_
    X = getKDE(fg_, conl)
    plcon = plotKDEContour(marginal(X,[1;2]), levels=levels, c=[contourColor_], line_width=contourLineWidth_)
    for ll in plcon.layers
      push!(plfl5.layers, ll)
    end
  end

  plfl5.coord = Coord.Cartesian(xmin=-20,xmax=30,ymin=-5,ymax=40)

  plfl5 |> PDF("$filepath", width, height);
  if show
    @async run(`evince $filepath`)
  end
  plfl5
end

#

ext: jl


text:
# This file is a part of TypeDBClient.  License is MIT: https://github.com/Humans-of-Julia/TypeDBClient.jl/blob/main/LICENSE

#
# package typedb.client.api.logic;
#
# import graql.lang.pattern.Pattern;
#
# import javax.annotation.CheckReturnValue;
# import javax.annotation.Nullable;
# import java.util.stream.Stream;
#
# public interface LogicManager {
#
#     @Nullable
#     @CheckReturnValue
#     Rule getRule(String label);
#
#     @CheckReturnValue
#     Stream<? extends Rule> getRules();
#
#     Rule putRule(String label, Pattern when, Pattern then);
# }


text: "$greet, $whom."

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# can I make separate rate functions and simply use the general poissonSpikes function

"""
		poissonSpikes(T, freq, interval, seed=1)

Simulate a Poisson spiking model.

# Arguments
T=Length of simulation in seconds\\
\\
freq=A constant denoting the firing frequency or a 1D Array of frequencies over the entire simulated
interval. If an array is given the array length must be equal to T/interval.\\
\\
interval=Sampling interval for the simulation in seconds.\\
\\
seed=set seed (optional)

# Output
spikes=An array of spike times in seconds.

# Example
```julia-repl
julia> poissonSpikes(0.25, 50, 0.001)
Float64[0.004, 0.018, 0.057, 0.113, 0.168]
```
"""
function poissonSpikes(T,freq,interval,seed=1)
	rng = MersenneTwister(seed)
	randArr = rand!(rng, zeros(Int(T/interval)))
	if isa(freq, Array)
		spikes = findall(x -> x == true, randArr .< freq .* interval) .* interval
	else
		spikes = findall(x -> x < freq * interval, randArr)/(1/interval)
	end

	return spikes
end

function poissonSpikesRefractory(T,freq0,tau,interval,seed=1)
	rng = MersenneTwister(seed)
	df = DataFrame(time = interval:interval:T,
				   prob = rand!(rng, zeros(Int(T/interval))))
	df.freq = zeros(length(df.time))
	df.spikes = zeros(length(df.time))
	prevSpike = -10000
	for i in 1:size(df, 1)
		dt = df.time[i] - prevSpike
		freq = freq0 * (1 - exp(-1 * ( dt/tau)))
		df.freq[i] = freq
		if df.prob[i] < freq*interval
			df.spikes[i] = 1.0
			prevSpike = df.time[i]
		else
			df.spikes[i] = 0
		end
	end

	return df.time[df[:spikes] .== 1, :], df
end

function poissonSpikesVarRate(T, interval, func, seed=1)
	rng = MersenneTwister(seed)
	df = DataFrame(time = interval:interval:T,
				   prob = rand!(rng, zeros(Int(T/interval))))
	df.freq = zeros(length(df.time))
	df.spikes = zeros(length(df.time))
	for i in 1:size(df, 1)
		freq = func(i*interval)
		df.freq[i] = freq
		if df.prob[i] < freq*interval
			df.spikes[i] = 1.0
			prevSpike = df.time[i]
		else
			df.spikes[i] = 0
		end
	end

	return df.time[df[:spikes] .== 1, :], df
end

function poissonSpikesStim(T, s, freq0,tau,interval,seed=1)
	rng = MersenneTwister(seed)
	df = DataFrame(time = interval:interval:T,
					 prob = rand!(rng, zeros(Int(T/interval))))
	df.freq = zeros(length(df.time))
	df.spikes = zeros(length(df.time))
	prevSpike = 0
	for i in 1:size(df, 1)
		dt = df.time[i] - prevSpike
		df.freq[i] = abs.(s[i] .+ freq0) .- exp(-dt/tau)
		if df.prob[i] < df.freq[i] * interval
			df.spikes[i] = 1.0
			prevSpike = df.time[i]
		else
			df.spikes[i] = 0
		end
	end

	return df.time[df[:spikes] .== 1, :], df
end


text:
module DiffEqFlux

using DiffEqBase, Flux, DiffResults, DiffEqSensitivity, ForwardDiff

include("Flux/layers.jl")
include("Flux/neural_de.jl")
include("Flux/utils.jl")

export diffeq_fd, diffeq_rd, diffeq_adjoint
export neural_ode, neural_ode_rd
export neural_dmsde
end


text:
"""
    SecondChebyshevPlans

implementation of computation of second Poincare invariant by approximating surface with
Chebyshev polynomials
"""
module SecondChebyshevPlans

import ..PoincareInvariants: compute!, getpoints, getpointspec

using ..PoincareInvariants: SecondPoincareInvariant, @argcheck

using LinearAlgebra

include("PaduaTransforms.jl")
using .PaduaTransforms

## Differentiation ##

function getdiffmat(::Type{T}, degree::Integer) where T
    D = zeros(T, degree+1, degree+1)

    D[1, 2:2:end] .= 1:2:degree

    for i in 3:2:degree+1
        D[2:2:i-1, i] .= 2 * (i-1)
    end

    for i in 4:2:degree+1
        D[3:2:i-1, i] .= 2 * (i-1)
    end

    D
end

struct DiffPlan{T}
    D::Matrix{T}
end

DiffPlan{T}(degree::Integer) where T = DiffPlan{T}(getdiffmat(T, degree))

# differentiate in the Chebyshev basis
function differentiate!(∂x::AbstractMatrix, ∂y::AbstractMatrix, P::DiffPlan, coeffs::AbstractMatrix)
    ∂x[:, :] = coeffs  # differentiate each row
    rmul!(∂x, LowerTriangular(P.D'))

    ∂y[:, :] = coeffs  # differentiate each column
    lmul!(UpperTriangular(P.D), ∂y)

    ∂x, ∂y
end

function differentiate!(∂x, ∂y, P::DiffPlan, coeffs)
    eltype(∂x) <: AbstractMatrix && eltype(∂y) <: AbstractMatrix &&
        eltype(coeffs) <: AbstractMatrix || throw(ArgumentError(
            "coefficients must be AbstractMatrix or iterable thereof"))
    length(∂x) == length(∂y) == length(coeffs) || throw(ArgumentError(
        "number of coefficient matrices must match number of derivative matrices to write to"))

    for (∂xi, ∂yi, coeffsi) in zip(∂x, ∂y, coeffs)
        differentiate!(∂xi, ∂yi, P, coeffsi)
    end

    ∂x, ∂y
end

## Integration ##

getintweights(::Type{T}, n) where T = T[isodd(i) ? 0 : T(2) / T(1 - i^2) for i in 0:n]
getintweights(n) = getintweights(Float64, n)

integrate(coeffs, intweights) = dot(intweights, coeffs, intweights)

## getintegrand for ω Callable and out-of-place ##

struct CallIntPlan{T, D, IP, P}
    invpaduaplan::IP
    ∂xvals::Matrix{T}
    ∂yvals::Matrix{T}
    intvals::Vector{T}
    paduaplan::P
end

function CallIntPlan{T, D}(degree) where {T, D}
    invpaduaplan = InvPaduaTransformPlan{Float64}(degree)
    ∂xvals = Matrix{T}(undef, getpaduanum(degree), D)
    ∂yvals = Matrix{T}(undef, getpaduanum(degree), D)
    intvals = Vector{T}(undef, getpaduanum(degree))
    paduaplan = PaduaTransformPlan{Float64}(degree)

    CallIntPlan{T, D, typeof(invpaduaplan), typeof(paduaplan)}(
        invpaduaplan, ∂xvals, ∂yvals, intvals, paduaplan
    )
end

function getintegrand!(
    intcoeffs::AbstractMatrix, plan::CallIntPlan{T, D}, ω::ωT,
    points, t, p, ∂xcoeffs, ∂ycoeffs
) where {T, D, ωT}
    invpaduatransform!(plan.∂xvals, plan.invpaduaplan, ∂xcoeffs)
    invpaduatransform!(plan.∂yvals, plan.invpaduaplan, ∂ycoeffs)

    for i in axes(plan.intvals, 1)
        # This if statement should hopefully get optimised away by the compiler
        if ωT <: AbstractMatrix
            ωi = ω
        else
            pnti = view(points, i, :)
            ωi = ω(pnti, t, p)
        end

        ∂xi = view(plan.∂xvals, i, :)
        ∂yi = view(plan.∂yvals, i, :)
        plan.intvals[i] = dot(∂yi, ωi, ∂xi)
    end

    paduatransform!(intcoeffs, plan.paduaplan, plan.intvals)

    intcoeffs
end

## SecondChebyshevPlan and compute! ##

getintplan(::Type{T}, ::Any, ::Val{D}, degree) where {T, D} = CallIntPlan{T, D}(degree)
# getintplan(::Type{T}, ω::AbstractMatrix, D, degree) where T = ConstIntPlan{T}(ω, D, degree)

struct SecondChebyshevPlan{T, D, IP, PP<:PaduaTransformPlan}
    degree::Int
    paduaplan::PP
    phasecoeffs::NTuple{D, Matrix{T}}
    diffplan::DiffPlan{T}
    ∂x::NTuple{D, Matrix{T}}
    ∂y::NTuple{D, Matrix{T}}
    intplan::IP  # getting coefficients of integrand to integrate
    intcoeffs::Matrix{T}
    intweights::Vector{T}
end

function SecondChebyshevPlan{T, D}(ω, N::Integer) where {T, D}
    degree = nextdegree(N)

    paduaplan = PaduaTransformPlan{Float64}(degree)
    phasecoeffs = ntuple(_ -> zeros(T, degree+1, degree+1), D)
    diffplan = DiffPlan{T}(degree)
    ∂x = ntuple(_ -> Matrix{T}(undef, degree+1, degree+1), D)
    ∂y = ntuple(_ -> Matrix{T}(undef, degree+1, degree+1), D)
    intplan = getintplan(T, ω, Val(D), degree)
    intcoeffs = zeros(T, degree+1, degree+1)
    intweights = getintweights(T, degree)

    SecondChebyshevPlan{T, D, typeof(intplan), typeof(paduaplan)}(
        degree, paduaplan, phasecoeffs, diffplan, ∂x, ∂y, intplan, intcoeffs, intweights
    )
end

function compute!(
    pinv::SecondPoincareInvariant{<:Any, <:Any, <:Any, <:Any, P}, t, p
) where P <: SecondChebyshevPlan
    plan = pinv.plan
    points = pinv.points
    paduatransform!(plan.phasecoeffs, plan.paduaplan, points)
    differentiate!(plan.∂x, plan.∂y, plan.diffplan, plan.phasecoeffs)
    getintegrand!(plan.intcoeffs, plan.intplan, pinv.ω, points, t, p, plan.∂x, plan.∂y)
    integrate(plan.intcoeffs, plan.intweights)
end

# function _compute!(plan::SecondChebyshevPlan, ω::AbstractMatrix, D, points)
#     paduatransform!(plan.phasecoeffs, plan.paduaplan, points)
#     differentiate!(plan.∂x, plan.∂y, plan.diffplan, plan.phasecoeffs)
#     getintegrand!(plan.intcoeffs, plan.intplan, ω, plan.∂x, plan.∂y)
#     integrate(plan.intcoeffs, plan.intplan)
# end

## getpoints, getpointspec and getpointnum ##

getpointspec(N::Integer, ::Type{<:SecondChebyshevPlan}) = nextpaduanum(N)

function getpoints(f, ::Type{T}, N, ::Type{<:SecondChebyshevPlan}) where T
    degree = getdegree(getpointspec(N, SecondChebyshevPlan))
    return getpaduapoints(T, degree) do x, y
        f((x + T(1)) / T(2), (y + T(1)) / T(2))
    end
end

end  # module Chebyshev


text:
# Nonlinear Model of F16 Aircraft with NASA _table
using ForwardDiff, Ipopt
include("AeroTableInterpolation.jl")

# -------- Atmosphere Model ---------
function atmos(alt::Real,vt::Real)::Tuple{Real,Real,Real}
    rho0 = 2.377e-3;
    tfac =1 - .703e-5*(alt);

    temp = 519.0*tfac;
    if (alt >= 35000.0)
       temp=390;
    end

    rho=rho0*tfac^4.14;
    mach = (vt)/sqrt(1.4*1716.3*temp);
    qbar = .5*rho*vt^2;
    ps   = 1715.0*rho*temp;

    if ps == 0
      ps = 1715;
    end

  return (mach, qbar, ps);
end

function StateAndControlBounds()::Matrix
  d2r = pi/180;
  bounds = [5000       40000;  # h ft
            -90*d2r    90*d2r;     # phi rad
            -90*d2r    90*d2r;     # theta rad
            -90*d2r    90*d2r;     # psi rad
             300       900;    # Vt ft/s
             -20*d2r   45*d2r;     # Alpha rad
             -30*d2r   30*d2r;     # Beta rad
            -Inf       Inf;    # p rad/s
            -Inf       Inf;    # q rad/s
            -Inf       Inf;    # r rad/s
            1000       19000;  # Thrust lbs
             -25       25;     # dele degrees
             -21.5     21.5;   # dail degrees
             -30       30;     # drud degrees
               0       25;     # dlef degrees
            ];

  return bounds;
end

function PrintStateAndControlBounds()
  bounds = StateAndControlBounds();

  stateNames = ("h","ϕ","θ","ψ","V","α","β","p","q","r");
  stateUnits = ("ft","rad","rad","rad","ft/s","rad","rad","rad/s","rad/s","rad/s");

  controlNames = ("T","δe","δa","δr","δlef");
  controlUnits = ("lbs","deg","deg","deg","deg");

  println("===== State Bounds =====")
  println("-Inf ≦ Npos (ft) ≦ Inf");
  println("-Inf ≦ Epos (ft) ≦ Inf");
  for i in 1:10
    str = " ≦ "*stateNames[i]*" ("*stateUnits[i]*") ≦ ";
    println(bounds[i,1], str, bounds[i,2]);
  end

  println(" ");
  println("===== Control Bounds =====")
    for i in 1:5
      str = " ≦ "*controlNames[i]*" ("*controlUnits[i]*") ≦ ";
      println(bounds[i+10,1], str, bounds[i+10,2]);
    end
end

# -------- Nonlinear Dynamic of F16 model ---------
function Dynamics!(xdot::Vector, x::Vector, u::Vector)
    # Define Constants for the Aircraft Model
    g = 32.17;          # gravity, ft/s^2
    m = 636.94;         # mass, slugs
    B = 30.0;           # span, ft
    S = 300.0;          # planform area, ft^2
    cbar = 11.32;       # mean _ chord, ft
    xcgr = 0.35;        # reference center of gravity as a fraction of cbar
    xcg  = 0.30;        # center of gravity as a fraction of cbar.
    Heng = 0.0;         # turbine momentum along roll axis.

    # NasaData -- translated via eq. 2.4-6 on pg 80 of Stevens and Lewis
    Jy  = 55814.0;       # slug-ft^2
    Jxz = 982.0;         # slug-ft^2
    Jz  = 63100.0;       # slug-ft^2
    Jx  = 9496.0;        # slug-ft^2
    r2d = 180.0/pi;

    # Collect the states and control variables
    npos  = x[1];
    epos  = x[2];
    alt   = x[3];   # altitude
    phi   = x[4];   # Euler angles are in rad.
    theta = x[5];
    psi   = x[6];

    vvt   = x[7];     # total velocity
    alpha = x[8]*r2d; # angle of attack in degrees
    beta  = x[9]*r2d; # sideslip angle in degrees
    P     = x[10];    # Roll Rate --- rolling  moment is Lbar
    Q     = x[11];    # Pitch Rate--- pitching moment is M
    R     = x[12];    # Yaw Rate  --- yawing   moment is N

    sa    = sin(x[8]); # sin(alpha)
    ca    = cos(x[8]); # cos(alpha)
    sb    = sin(x[9]); # sin(beta)
    cb    = cos(x[9]); # cos(beta)
    tb    = tan(x[9]); # tan(beta)

    st    = sin(theta);
    ct    = cos(theta);
    tt    = tan(theta);
    sphi  = sin(phi);
    cphi  = cos(phi);
    spsi  = sin(psi);
    cpsi  = cos(psi);

    # Control variables
    T     = u[1];   # Thrust
    el    = u[2];   # Elevator setting in degrees
    ail   = u[3];   # Ailerons mex setting in degrees
    rud   = u[4];   # Rudder setting in degrees
    lef   = u[5];   # Leading edge flap setting in degrees

    # Normalization and scaling
    if vvt <= 0.1
      vt = 0.1;
    else
      vt = vvt;
    end

    dail  = ail/21.5; # Aileron angle normalized against max deflection
    drud  = rud/30.0;  # Rudder normalized against max angle
    dlef  = (1 - lef/25.0);  # Leading edge flap normalized against max angle

    # Atmospheric effects
    mach, qbar, ps = atmos(alt,vt);

    # Navigation equations
    U = vt*ca*cb;  # directional velocities.
    V = vt*sb;
    W = vt*sa*cb;

    nposd = U*(ct*cpsi) + V*(sphi*cpsi*st - cphi*spsi) + W*(cphi*st*cpsi + sphi*spsi); # npos dot
    eposd = U*(ct*spsi) + V*(sphi*spsi*st + cphi*cpsi) + W*(cphi*st*spsi - sphi*cpsi); # epos dot
    altd = U*st - V*(sphi*ct) - W*(cphi*ct); # alt dot

    # Kinematic equations
    phid = P + tt*(Q*sphi + R*cphi); # phi dot
    thd = Q*cphi - R*sphi; # theta dot
    psid = (Q*sphi + R*cphi)/ct; # psi dot

    # Aerodynamic forces and moments
    Cx = _Cx(alpha,beta,el);
    Cy = _Cy(alpha,beta);
    Cz = _Cz(alpha,beta,el);
    Cm = _Cm(alpha,beta,el);
    Cn = _Cn(alpha,beta,el);
    Cl = _Cl(alpha,beta,el);

    delta_Cx_lef, delta_Cy_lef, delta_Cz_lef,delta_Cm_lef, delta_Cn_lef, delta_Cl_lef  = Delta_lef(alpha,beta);

    # Aerodynamic Damping
    Cxq, Cyr, Cyp, Czq, Clr, Clp, Cmq, Cnr, Cnp  = _Damping(alpha);
    delta_Cxq_lef, delta_Cyr_lef, delta_Cyp_lef, delta_Czq_lef, delta_Clr_lef, delta_Clp_lef, delta_Cmq_lef, delta_Cnr_lef, delta_Cnp_lef  = _Damping_lef(alpha);

    # Rudder Influence
    delta_Cy_r30, delta_Cn_r30, delta_Cl_r30 = _RudderInfluence(alpha,beta);

    # Aileron Influence
    delta_Cy_a20, delta_Cy_a20_lef, delta_Cn_a20, delta_Cn_a20_lef, delta_Cl_a20, delta_Cl_a20_lef = _AileronInfluence(alpha,beta);

    # Other Coefficients
    delta_Cnbeta, delta_Clbeta, delta_Cm, eta_el = _OtherCoefficients(alpha,el);
    delta_Cm_ds = 0; # Ignore deep-stall effects.

    # ========= Dynamics -- Euler's first and second law ==============

    # Compute total Cx
    dXdQ = (cbar/(2*vt))*(Cxq + delta_Cxq_lef*dlef);
    Cx_tot = Cx + delta_Cx_lef*dlef + dXdQ*Q;

    # Compute total Cz
    dZdQ = (cbar/(2*vt))*(Czq + delta_Cz_lef*dlef);
    Cz_tot = Cz + delta_Cz_lef*dlef + dZdQ*Q;

    # Compute total Cm
    dMdQ = (cbar/(2*vt))*(Cmq + delta_Cmq_lef*dlef);
    Cm_tot = Cm*eta_el + Cz_tot*(xcgr-xcg) + delta_Cm_lef*dlef + dMdQ*Q + delta_Cm + delta_Cm_ds;

    # Compute total Cy
    dYdail = delta_Cy_a20 + delta_Cy_a20_lef*dlef;
    dYdR = (B/(2*vt))*(Cyr + delta_Cyr_lef*dlef);
    dYdP = (B/(2*vt))*(Cyp + delta_Cyp_lef*dlef);
    Cy_tot = Cy + delta_Cy_lef*dlef + dYdail*dail + delta_Cy_r30*drud + dYdR*R + dYdP*P;

    # Compute total Cn
    dNdail = delta_Cn_a20 + delta_Cn_a20_lef*dlef;
    dNdR = (B/(2*vt))*(Cnr + delta_Cnr_lef*dlef);
    dNdP = (B/(2*vt))*(Cnp + delta_Cnp_lef*dlef);
    Cn_tot = Cn + delta_Cn_lef*dlef - Cy_tot*(xcgr-xcg)*(cbar/B) + dNdail*dail + delta_Cn_r30*drud + dNdR*R + dNdP*P + delta_Cnbeta*beta;

    # Compute total Cl
    dLdail = delta_Cl_a20 + delta_Cl_a20_lef*dlef;
    dLdR = (B/(2*vt))*(Clr + delta_Clr_lef*dlef);
    dLdP = (B/(2*vt))*(Clp + delta_Clp_lef*dlef);
    Cl_tot = Cl + delta_Cl_lef*dlef + dLdail*dail + delta_Cl_r30*drud + dLdR*R + dLdP*P + delta_Clbeta*beta;

    # Compute Udot,Vdot, Wdot,(as on NASA report p36)
    Udot = R*V - Q*W - g*st + qbar*S*Cx_tot/m + T/m;
    Vdot = P*W - R*U + g*ct*sphi + qbar*S*Cy_tot/m;
    Wdot = Q*U - P*V + g*ct*cphi + qbar*S*Cz_tot/m;

    # vt_dot equation
    vd = (U*Udot + V*Vdot + W*Wdot)/vt;

    # alpha_dot equation
    αd = (U*Wdot - W*Udot)/(U*U + W*W);

    # beta_dot equation
    βd = (Vdot*vt - V*vd)/(vt*vt*cb);

    # Equations for Pdot, Qdot, and Rdot
    L_tot = Cl_tot*qbar*S*B;       # Get moments from coefficients
    M_tot = Cm_tot*qbar*S*cbar;
    N_tot = Cn_tot*qbar*S*B;
    denom = Jx*Jz - Jxz*Jxz;

    # Pdot
    pd =  (Jz*L_tot + Jxz*N_tot - (Jz*(Jz-Jy)+Jxz*Jxz)*Q*R + Jxz*(Jx-Jy+Jz)*P*Q + Jxz*Q*Heng)/denom;

    # Qdot
    qd = (M_tot + (Jz-Jx)*P*R - Jxz*(P*P-R*R) - R*Heng)/Jy;

    # Rdot
    rd = (Jx*N_tot + Jxz*L_tot + (Jx*(Jx-Jy)+Jxz*Jxz)*P*Q - Jxz*(Jx-Jy+Jz)*Q*R +  Jx*Q*Heng)/denom;

    xdot[1] = nposd;
    xdot[2] = eposd;
    xdot[3] = altd;
    xdot[4] = phid;
    xdot[5] = thd;
    xdot[6] = psid;
    xdot[7] = vd;
    xdot[8] = αd;
    xdot[9] = βd;
    xdot[10] = pd;
    xdot[11] = qd;
    xdot[12] = rd;
end

# Redefine the function for gradient calculation.
function Dynamics(x::Vector, u::Vector)::Vector
  # Define Constants for the Aircraft Model
  g = 32.17;          # gravity, ft/s^2
  m = 636.94;         # mass, slugs
  B = 30.0;           # span, ft
  S = 300.0;          # planform area, ft^2
  cbar = 11.32;       # mean _ chord, ft
  xcgr = 0.35;        # reference center of gravity as a fraction of cbar
  xcg  = 0.30;        # center of gravity as a fraction of cbar.
  Heng = 0.0;         # turbine momentum along roll axis.

  # NasaData -- translated via eq. 2.4-6 on pg 80 of Stevens and Lewis
  Jy  = 55814.0;       # slug-ft^2
  Jxz = 982.0;         # slug-ft^2
  Jz  = 63100.0;       # slug-ft^2
  Jx  = 9496.0;        # slug-ft^2
  r2d = 180.0/pi;

  # Collect the states and control variables
  npos  = x[1];
  epos  = x[2];
  alt   = x[3];    # altitude
  phi   = x[4];   # Euler angles are in rad.
  theta = x[5];
  psi   = x[6];

  vvt   = x[7];     # total velocity
  alpha = x[8]*r2d; # angle of attack in degrees for _ tables
  beta  = x[9]*r2d; # sideslip angle in degrees for _ tables
  P     = x[10];    # Roll Rate --- rolling  moment is Lbar
  Q     = x[11];    # Pitch Rate--- pitching moment is M
  R     = x[12];    # Yaw Rate  --- yawing   moment is N

  sa    = sin(x[8]); # sin(alpha)
  ca    = cos(x[8]); # cos(alpha)
  sb    = sin(x[9]); # sin(beta)
  cb    = cos(x[9]); # cos(beta)
  tb    = tan(x[9]); # tan(beta)

  st    = sin(theta);
  ct    = cos(theta);
  tt    = tan(theta);
  sphi  = sin(phi);
  cphi  = cos(phi);
  spsi  = sin(psi);
  cpsi  = cos(psi);

  # Control variables
  T     = u[1];   # Thrust
  el    = u[2];   # Elevator setting in degrees
  ail   = u[3];   # Ailerons mex setting in degrees
  rud   = u[4];   # Rudder setting in degrees
  lef   = u[5];   # Leading edge flap setting in degrees

  # Normalization and scaling
  if vvt <= 0.1
    vt = 0.1;
  else
    vt = vvt;
  end

  dail  = ail/21.5; # Aileron angle normalized against max deflection
  drud  = rud/30.0;  # Rudder normalized against max angle
  dlef  = (1 - lef/25.0);  # Leading edge flap normalized against max angle

  # Atmospheric effects
  mach, qbar, ps = atmos(alt,vt);

  # Navigation equations
  U = vt*ca*cb;  # directional velocities.
  V = vt*sb;
  W = vt*sa*cb;

  nposd = U*(ct*cpsi) + V*(sphi*cpsi*st - cphi*spsi) + W*(cphi*st*cpsi + sphi*spsi); # npos dot
  eposd = U*(ct*spsi) + V*(sphi*spsi*st + cphi*cpsi) + W*(cphi*st*spsi - sphi*cpsi); # epos dot
  altd = U*st - V*(sphi*ct) - W*(cphi*ct); # alt dot

  # Kinematic equations
  phid = P + tt*(Q*sphi + R*cphi); # phi dot
  thd = Q*cphi - R*sphi; # theta dot
  psid = (Q*sphi + R*cphi)/ct; # psi dot

  # Aerodynamic forces and moments
  Cx = _Cx(alpha,beta,el);
  Cy = _Cy(alpha,beta);
  Cz = _Cz(alpha,beta,el);
  Cm = _Cm(alpha,beta,el);
  Cn = _Cn(alpha,beta,el);
  Cl = _Cl(alpha,beta,el);

  delta_Cx_lef, delta_Cy_lef, delta_Cz_lef,delta_Cm_lef, delta_Cn_lef, delta_Cl_lef  = Delta_lef(alpha,beta);

  # Aerodynamic Damping
  Cxq, Cyr, Cyp, Czq, Clr, Clp, Cmq, Cnr, Cnp  = _Damping(alpha);
  delta_Cxq_lef, delta_Cyr_lef, delta_Cyp_lef, delta_Czq_lef, delta_Clr_lef, delta_Clp_lef, delta_Cmq_lef, delta_Cnr_lef, delta_Cnp_lef  = _Damping_lef(alpha);

  # Rudder Influence
  delta_Cy_r30, delta_Cn_r30, delta_Cl_r30 = _RudderInfluence(alpha,beta);

  # Aileron Influence
  delta_Cy_a20, delta_Cy_a20_lef, delta_Cn_a20, delta_Cn_a20_lef, delta_Cl_a20, delta_Cl_a20_lef = _AileronInfluence(alpha,beta);

  # Other Coefficients
  delta_Cnbeta, delta_Clbeta, delta_Cm, eta_el = _OtherCoefficients(alpha,el);
  delta_Cm_ds = 0; # Ignore deep-stall effects.

  # ========= Dynamics -- Euler's first and second law ==============

  # Compute total Cx
  dXdQ = (cbar/(2*vt))*(Cxq + delta_Cxq_lef*dlef);
  Cx_tot = Cx + delta_Cx_lef*dlef + dXdQ*Q;

  # Compute total Cz
  dZdQ = (cbar/(2*vt))*(Czq + delta_Cz_lef*dlef);
  Cz_tot = Cz + delta_Cz_lef*dlef + dZdQ*Q;

  # Compute total Cm
  dMdQ = (cbar/(2*vt))*(Cmq + delta_Cmq_lef*dlef);
  Cm_tot = Cm*eta_el + Cz_tot*(xcgr-xcg) + delta_Cm_lef*dlef + dMdQ*Q + delta_Cm + delta_Cm_ds;

  # Compute total Cy
  dYdail = delta_Cy_a20 + delta_Cy_a20_lef*dlef;
  dYdR = (B/(2*vt))*(Cyr + delta_Cyr_lef*dlef);
  dYdP = (B/(2*vt))*(Cyp + delta_Cyp_lef*dlef);
  Cy_tot = Cy + delta_Cy_lef*dlef + dYdail*dail + delta_Cy_r30*drud + dYdR*R + dYdP*P;

  # Compute total Cn
  dNdail = delta_Cn_a20 + delta_Cn_a20_lef*dlef;
  dNdR = (B/(2*vt))*(Cnr + delta_Cnr_lef*dlef);
  dNdP = (B/(2*vt))*(Cnp + delta_Cnp_lef*dlef);
  Cn_tot = Cn + delta_Cn_lef*dlef - Cy_tot*(xcgr-xcg)*(cbar/B) + dNdail*dail + delta_Cn_r30*drud + dNdR*R + dNdP*P + delta_Cnbeta*beta;

  # Compute total Cl
  dLdail = delta_Cl_a20 + delta_Cl_a20_lef*dlef;
  dLdR = (B/(2*vt))*(Clr + delta_Clr_lef*dlef);
  dLdP = (B/(2*vt))*(Clp + delta_Clp_lef*dlef);
  Cl_tot = Cl + delta_Cl_lef*dlef + dLdail*dail + delta_Cl_r30*drud + dLdR*R + dLdP*P + delta_Clbeta*beta;

  # Compute Udot,Vdot, Wdot,(as on NASA report p36)
  Udot = R*V - Q*W - g*st + qbar*S*Cx_tot/m + T/m;
  Vdot = P*W - R*U + g*ct*sphi + qbar*S*Cy_tot/m;
  Wdot = Q*U - P*V + g*ct*cphi + qbar*S*Cz_tot/m;

  # vt_dot equation
  vd = (U*Udot + V*Vdot + W*Wdot)/vt;

  # alpha_dot equation
  αd = (U*Wdot - W*Udot)/(U*U + W*W);

  # beta_dot equation
  βd = (Vdot*vt - V*vd)/(vt*vt*cb);

  # Equations for Pdot, Qdot, and Rdot
  L_tot = Cl_tot*qbar*S*B;       # Get moments from coefficients
  M_tot = Cm_tot*qbar*S*cbar;
  N_tot = Cn_tot*qbar*S*B;
  denom = Jx*Jz - Jxz*Jxz;

  # Pdot
  pd =  (Jz*L_tot + Jxz*N_tot - (Jz*(Jz-Jy)+Jxz*Jxz)*Q*R + Jxz*(Jx-Jy+Jz)*P*Q + Jxz*Q*Heng)/denom;

  # Qdot
  qd = (M_tot + (Jz-Jx)*P*R - Jxz*(P*P-R*R) - R*Heng)/Jy;

  # Rdot
  rd = (Jx*N_tot + Jxz*L_tot + (Jx*(Jx-Jy)+Jxz*Jxz)*P*Q - Jxz*(Jx-Jy+Jz)*Q*R +  Jx*Q*Heng)/denom;

  xdot = [nposd, eposd, altd, phid, thd, psid, vd, αd, βd, pd, qd, rd];

  return xdot;
end

function Linearize(x0::Vector,u0::Vector)::Tuple{Matrix,Matrix}
  f1 = x -> Dynamics(x,u0);
  f2 = u -> Dynamics(x0,u);
  A = ForwardDiff.jacobian(f1,x0);
  B = ForwardDiff.jacobian(f2,u0);
  return (A,B)
end

function Trim(h, V; γ=0, ψdot=0, ϕ = (0,0), θ = (0,0), ψ = (0,0), α = (0,0), β = (0,0), p = (0,0), q = (0,0), r = (0,0));
  # States in trim: h, ϕ, θ, ψ, Vt, α, β, p, q, r
  d2r = pi/180;

  # Initial guesses
  x0 = [h, ϕ[1], θ[1], ψ[1], V, α[1], β[1], p[1], q[1], r[1]];

  T0, dele0, dail0, drud0, dlef0 = 9000,0,0,0,0;
  u0 = [T0,dele0,dail0,drud0,dlef0];

  # Determine optimization variables
  ix = [1, ϕ[2], θ[2], ψ[2], 1, α[2], β[2], p[2], q[2], r[2]]; # 1 => Trim values are fixed in the optimization, 0 => Trim values are optimization variables.
  ii = findall(x->x==0, ix);

  # Extract states and control from optimization variables
  function getXU(z)
    z0 = x0
    z0[ii] .= 0;

    Mz = zeros(length(x0),length(ii));
    for (i,ind) in enumerate(ii)
        Mz[ind,i] = 1;
    end

    v = z0 + Mz*z[1:length(ii)];

    X = [0;0;v];
    U = z[(length(ii)+1):end];
    return (X,U)
  end

  # Nonlinear constraints from dynamics and trim types
  function _nonlinear_constraints(z)

    var = z[1:end-1];
    X,U = getXU(var);

    ϕ = X[4]; θ = X[5]; #ψ = X[6];
    α = X[8]; β = X[9];
    g = 9.806;

    a = cos(α)*cos(β);
    b = sin(ϕ)*sin(β) + cos(ϕ)*sin(α)*cos(β);
    G = ψdot*V/g;

    # Zero acceleration constraints
    C1 = Dynamics(X, U)[7:12]; # Only Vdot, αdot, βdot, pdot, qdot, and rdot

    # Rate of climb constraint
    v1 = a^2 - sin(γ)^2;
    C2 = tan(θ) - (a*b + sin(γ)*sqrt(v1 + b^2))/v1;

    # Coordinated turn constraint
    C3 = G*cos(β)*(sin(α)*tan(θ) + cos(α)*cos(ϕ));

    constr = [C1;C2;C3]; # Total 8 constraints

    # Weights
    W = [0.1 1 1 1 1 1];
    J = sum(constr.*constr) - z[end];
    return J
  end

# ---- IpOpt Specific Functions ----
  function eval_f(z) # Cost function
    return z[end];
  end

  function eval_g(z, g) # Nonlinear inequality constraint
    # Bad: g    = zeros(2)  # Allocates new array
    # OK:  g[:] = zeros(2)  # Modifies 'in place'
    g[:] .= _nonlinear_constraints(z);
  end

  function eval_grad_f(z, gradf) # Gradient of cost function.
    # Bad: grad_f    = zeros(4)  # Allocates new array
    # OK:  grad_f[:] = zeros(4)  # Modifies 'in place'
    f = x-> ForwardDiff.gradient(eval_f,x);
    gradf[:] = f(z);
  end

  function eval_jac_g(z, mode, rows, cols, values) # Jacobian of nonlinear inequality constraint.
    nvar = length(z);
    ng = 1;

    if mode == :Structure
        index = 1;
        for i=1:ng
            for j=1:nvar
                rows[index] = i;
                cols[index] = j;
                index += 1;
            end
        end
    else
        M = x->ForwardDiff.gradient(_nonlinear_constraints,x);
        values[:] = M(z);
    end
  end

  function eval_h(z, mode, rows, cols, obj_factor, lambda, values) # Hessian of cost.
    nvar = length(z);
    if mode == :Structure
      # Symmetric matrix, fill the lower left triangle only
      idx = 1
      for row = 1:nvar
        for col = 1:row
          rows[idx] = row
          cols[idx] = col
          idx += 1
        end
      end
    else
      H = obj_factor*zeros(nvar,nvar);
      f(z) = _nonlinear_constraints(z); # Only one constraint
      ∇2f = x->ForwardDiff.hessian(f,x);
      H = H + lambda[1]*∇2f(z);

      # Copy them into return variable -- only lower triangle.
      idx = 1;
      for row = 1:nvar
        for col = 1:row
          values[idx] = H[row,col];
          idx += 1;
        end
      end
    end
  end

  # Call IpOpt to solve the constraint nonlinear optmization problem
  bounds = StateAndControlBounds();
  bx = bounds[ii,:]; # State bounds
  bu = bounds[11:15,:]; # Control bounds

  n = length(ii) + 5 + 1;
  x_L = [bx[:,1];bu[:,1];0];
  x_U = [bx[:,2];bu[:,2];Inf];
  m = 1; # Only one nonlinear constraint ... true for all trim types
  g_L = [-Inf];
  g_U = [0.0];

  nele_jac  = m*n;
  nele_hess = Integer(n*(n+1)/2);
  prob = createProblem(n, x_L, x_U, m, g_L, g_U, nele_jac, nele_hess, eval_f, eval_g, eval_grad_f, eval_jac_g, eval_h);

  prob.x = [x0[ii];u0;0]; # Initial guess.
  status = solveProblem(prob);
  xbar,ubar = getXU(prob.x[1:end-1]);
  return (xbar, ubar, status, prob);
end


text:
# scalar version

"""
    ScalarFunctionizeBridge{T, S}

The `ScalarFunctionizeBridge` converts a constraint `SingleVariable`-in-`S`
into the constraint `ScalarAffineFunction{T}`-in-`S`.
"""
struct ScalarFunctionizeBridge{T, S} <: AbstractFunctionConversionBridge{MOI.ScalarAffineFunction{T}, S}
    constraint::CI{MOI.ScalarAffineFunction{T}, S}
end
function bridge_constraint(::Type{ScalarFunctionizeBridge{T, S}}, model,
                           f::MOI.SingleVariable, s::S) where {T, S}
    constraint = MOI.add_constraint(model, MOI.ScalarAffineFunction{T}(f), s)
    return ScalarFunctionizeBridge{T, S}(constraint)
end

# start allowing everything (scalar)
MOI.supports_constraint(::Type{ScalarFunctionizeBridge{T}},
                        ::Type{<:MOI.SingleVariable},
                        ::Type{<:MOI.AbstractScalarSet}) where {T} = true
MOIB.added_constrained_variable_types(::Type{<:ScalarFunctionizeBridge}) = Tuple{DataType}[]
function MOIB.added_constraint_types(::Type{ScalarFunctionizeBridge{T, S}}) where {T, S}
    return [(MOI.ScalarAffineFunction{T}, S)]
end
function concrete_bridge_type(::Type{<:ScalarFunctionizeBridge{T}},
                              ::Type{MOI.SingleVariable},
                              S::Type{<:MOI.AbstractScalarSet}) where T
    return ScalarFunctionizeBridge{T, S}
end

# Attributes, Bridge acting as a model
MOI.get(b::ScalarFunctionizeBridge{T, S}, ::MOI.NumberOfConstraints{MOI.ScalarAffineFunction{T}, S}) where {T, S} = 1
MOI.get(b::ScalarFunctionizeBridge{T, S}, ::MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{T}, S}) where {T, S} = [b.constraint]

# Indices
function MOI.delete(model::MOI.ModelLike, c::ScalarFunctionizeBridge)
    MOI.delete(model, c.constraint)
    return
end

# Constraints
function MOI.set(model::MOI.ModelLike, ::MOI.ConstraintFunction,
                 c::ScalarFunctionizeBridge{T}, f::MOI.SingleVariable) where {T}
    MOI.set(model, MOI.ConstraintFunction(), c.constraint, MOI.ScalarAffineFunction{T}(f))
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 b::ScalarFunctionizeBridge)
    f = MOIU.canonical(MOI.get(model, attr, b.constraint))
    return convert(MOI.SingleVariable, f)
end

# vector version

"""
    VectorFunctionizeBridge{T, S}

The `VectorFunctionizeBridge` converts a constraint `VectorOfVariables`-in-`S`
into the constraint `VectorAffineFunction{T}`-in-`S`.
"""
mutable struct VectorFunctionizeBridge{T, S} <: AbstractFunctionConversionBridge{MOI.VectorAffineFunction{T}, S}
    constraint::CI{MOI.VectorAffineFunction{T}, S}
end
function bridge_constraint(::Type{VectorFunctionizeBridge{T, S}}, model,
                           f::MOI.VectorOfVariables, s::S) where {T, S}
    constraint = MOI.add_constraint(model, MOI.VectorAffineFunction{T}(f), s)
    return VectorFunctionizeBridge{T, S}(constraint)
end

MOI.supports_constraint(::Type{VectorFunctionizeBridge{T}},
                        ::Type{MOI.VectorOfVariables},
                        ::Type{<:MOI.AbstractVectorSet}) where {T} = true
MOIB.added_constrained_variable_types(::Type{<:VectorFunctionizeBridge}) = Tuple{DataType}[]
function MOIB.added_constraint_types(::Type{VectorFunctionizeBridge{T, S}}) where {T, S}
    return [(MOI.VectorAffineFunction{T}, S)]
end
function concrete_bridge_type(::Type{<:VectorFunctionizeBridge{T}},
                              ::Type{<:MOI.AbstractVectorFunction},
                              S::Type{<:MOI.AbstractVectorSet}) where T
    return VectorFunctionizeBridge{T, S}
end

# Attributes, Bridge acting as a model
MOI.get(b::VectorFunctionizeBridge{T, S},
        ::MOI.NumberOfConstraints{MOI.VectorAffineFunction{T}, S}) where {T, S} = 1
MOI.get(b::VectorFunctionizeBridge{T, S},
        ::MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T}, S}) where {T, S} = [b.constraint]

# Indices
function MOI.delete(model::MOI.ModelLike, bridge::VectorFunctionizeBridge)
    MOI.delete(model, bridge.constraint)
    return
end
function MOI.delete(model::MOI.ModelLike, bridge::VectorFunctionizeBridge,
                    i::IndexInVector)
    func = MOI.get(model, MOI.ConstraintFunction(), bridge.constraint)
    idx = setdiff(1:MOI.output_dimension(func), i.value)
    new_func = MOIU.eachscalar(func)[idx]
    set = MOI.get(model, MOI.ConstraintSet(), bridge.constraint)
    new_set = MOI.update_dimension(set, MOI.dimension(set) - 1)
    MOI.delete(model, bridge.constraint)
    bridge.constraint = MOI.add_constraint(model, new_func, new_set)
end

# Constraints
function MOI.set(model::MOI.ModelLike, ::MOI.ConstraintFunction,
                 bridge::VectorFunctionizeBridge{T},
                 func::MOI.VectorOfVariables) where {T}
    MOI.set(model, MOI.ConstraintFunction(), bridge.constraint,
            MOI.VectorAffineFunction{T}(func))
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 b::VectorFunctionizeBridge)
    f = MOI.get(model, attr, b.constraint)
    return MOIU.convert_approx(MOI.VectorOfVariables, f)
end


text:
using DataStructures
using Test

#  Setup for pycall tests - done by travis, uncomment for manual run
# using Conda; ENV["PYTHON"] = Conda.PYTHONDIR
# using Pkg
# Pkg.build("PyCall")

# Conda.add_channel("conda-forge")
# Conda.add("cfgrib")

using PyCall

@testset "era5-levels-members DataSet parity" begin
    test_file = joinpath(dir_testfiles, "era5-levels-members.grib")
    res = CfGRIB.DataSet(test_file)

    @testset "Variables" begin
        @test res.variables["number"].dimensions == ("number",)
        @test res.variables["number"].data == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        @test res.variables["number"].attributes == Dict(
            "units" => "1",
            "long_name" => "ensemble member numerical id",
            "standard_name" => "realization"
        )

        @test res.variables["time"].dimensions == ("time",)
        @test res.variables["time"].data == [
            1483228800, 1483272000, 1483315200, 1483358400
        ]
        @test res.variables["time"].attributes == Dict(
            "units" => "seconds since 1970-01-01T00:00:00",
            "calendar" => "proleptic_gregorian",
            "long_name" => "initial time of forecast",
            "standard_name" => "forecast_reference_time"
        )

        @test res.variables["step"].dimensions == ()
        @test res.variables["step"].data == 0 #  In python this is a float
        @test res.variables["step"].attributes == Dict(
            "units" => "hours",
            "long_name" => "time since forecast_reference_time",
            "standard_name" => "forecast_period"
        )

        @test res.variables["isobaricInhPa"].dimensions == ("isobaricInhPa",)
        @test res.variables["isobaricInhPa"].data == [850, 500]
        @test res.variables["isobaricInhPa"].attributes == Dict(
            "units" => "hPa",
            "stored_direction" => "decreasing",
            "long_name" => "pressure",
            "standard_name" => "air_pressure",
            "positive" => "down"
        )

        @test res.variables["latitude"].dimensions == ("latitude",)
        @test res.variables["latitude"].data == [
            90.,   87.,  84.,  81.,  78.,  75.,  72.,  69.,  66.,  63.,  60.,
            57.,   54.,  51.,  48.,  45.,  42.,  39.,  36.,  33.,  30.,  27.,
            24.,   21.,  18.,  15.,  12.,   9.,   6.,   3.,   0.,  -3.,  -6.,
            -9.,  -12., -15., -18., -21., -24., -27., -30., -33., -36., -39.,
            -42., -45., -48., -51., -54., -57., -60., -63., -66., -69., -72.,
            -75., -78., -81., -84., -87., -90.
        ]
        @test res.variables["latitude"].attributes == Dict(
            "units" => "degrees_north",
            "stored_direction" => "decreasing",
            "long_name" => "latitude",
            "standard_name" => "latitude"
        )

        @test res.variables["longitude"].dimensions == ("longitude",)
        @test res.variables["longitude"].data == [
            0.,   3.,   6.,   9.,  12.,  15.,  18.,  21.,  24.,  27.,  30.,
            33.,  36.,  39.,  42.,  45.,  48.,  51.,  54.,  57.,  60.,  63.,
            66.,  69.,  72.,  75.,  78.,  81.,  84.,  87.,  90.,  93.,  96.,
            99., 102., 105., 108., 111., 114., 117., 120., 123., 126., 129.,
            132., 135., 138., 141., 144., 147., 150., 153., 156., 159., 162.,
            165., 168., 171., 174., 177., 180., 183., 186., 189., 192., 195.,
            198., 201., 204., 207., 210., 213., 216., 219., 222., 225., 228.,
            231., 234., 237., 240., 243., 246., 249., 252., 255., 258., 261.,
            264., 267., 270., 273., 276., 279., 282., 285., 288., 291., 294.,
            297., 300., 303., 306., 309., 312., 315., 318., 321., 324., 327.,
            330., 333., 336., 339., 342., 345., 348., 351., 354., 357.
        ]
        @test res.variables["longitude"].attributes == Dict(
            "units" => "degrees_east",
            "long_name" => "longitude",
            "standard_name" => "longitude"
        )

        @test res.variables["valid_time"].dimensions == ("time",)
        @test res.variables["valid_time"].data == [
            1.4832288e+09, 1.4832720e+09, 1.4833152e+09, 1.4833584e+09
        ]
        @test res.variables["valid_time"].attributes == Dict(
            "units" => "seconds since 1970-01-01T00:00:00",
            "calendar" => "proleptic_gregorian",
            "long_name" => "time",
            "standard_name" => "time"
    )
    end

    @testset "OnDiskArray" begin
        @test res.variables["z"].dimensions == (
            "number", "time", "isobaricInhPa", "longitude", "latitude"
        )
        @test typeof(res.variables["z"].data) == CfGRIB.OnDiskArray
        @test size(res.variables["z"].data) == (10, 4, 2, 120, 61)
        @test_broken res.variables["z"].missing_value == 9999
        #  Currently these are not `OrderedDict`, and in the wrong order
        #  TODO: Make these into 'OrderedDict', and order correctly
        @test_broken typeof(res.variables["z"].attributes) == OrderedDict
        @test res.variables["z"].attributes == OrderedDict(
            "GRIB_paramId"                            => 129,
            "GRIB_shortName"                          => "z",
            "GRIB_units"                              => "m**2 s**-2",
            "GRIB_name"                               => "Geopotential",
            "GRIB_cfName"                             => "geopotential",
            "GRIB_cfVarName"                          => "z",
            "GRIB_dataType"                           => "an",
            "GRIB_missingValue"                       => 9999,
            "GRIB_numberOfPoints"                     => 7320,
            "GRIB_totalNumber"                        => 10,
            "GRIB_typeOfLevel"                        => "isobaricInhPa",
            "GRIB_NV"                                 => 0,
            "GRIB_stepUnits"                          => 1,
            "GRIB_stepType"                           => "instant",
            "GRIB_gridType"                           => "regular_ll",
            "GRIB_gridDefinitionDescription"          => "Latitude/Longitude Grid",
            "GRIB_Nx"                                 => 120,
            "GRIB_iDirectionIncrementInDegrees"       => 3.0,
            "GRIB_iScansNegatively"                   => 0,
            "GRIB_longitudeOfFirstGridPointInDegrees" => 0.0,
            "GRIB_longitudeOfLastGridPointInDegrees"  => 357.0,
            "GRIB_Ny"                                 => 61,
            "GRIB_jDirectionIncrementInDegrees"       => 3.0,
            "GRIB_jPointsAreConsecutive"              => 0,
            "GRIB_jScansPositively"                   => 0,
            "GRIB_latitudeOfFirstGridPointInDegrees"  => 90.0,
            "GRIB_latitudeOfLastGridPointInDegrees"   => -90.0,
            "long_name"                               => "Geopotential",
            "units"                                   => "m**2 s**-2",
            "standard_name"                           => "geopotential",
            "coordinates"                             => "number time step isobaricInhPa latitude longitude valid_time"
        )
        #  Should be approx. equal to python:
        #    `res.variables["z"].data[:, :, 1, 1, 1]`
        @test res.variables["z"].data[:, :, 2, 2, 2] ≈ [
            51213.703 51188.277 50836.094 50185.605
            51216.605 51180.03  50833.027 50210.992
            51232.766 51160.65  50823.082 50199.85
            51203.06  51182.15  50816.246 50177.18
            51235.203 51181.234 50823.41  50182.023
            51207.727 51179.25  50813.094 50190.875
            51210.867 51179.35  50818.527 50183.945
            51193.94  51153.36  50794.383 50159.473
            51203.395 51183.875 50835.9   50207.61
            51197.383 51183.77  50819.78  50192.44
        ]
        @test size(res.variables["z"].data[1, 1, 1, :, :]) == (120, 61) # ' vs python
        #  Should be approx. equal to python:
        #    `res.variables["z"].data[0, 0, 0, :, :][:4, :4]`
        #  Used `adjoint` explicitly for clarity, the ' operator works here too
        @test adjoint(res.variables["z"].data[1, 1, 1, :, :])[1:4, 1:4] ≈ [
            14201.754 14201.754 14201.754 14201.754
            14323.254 14333.129 14342.879 14353.504
            14235.629 14252.129 14266.254 14275.879
            14104.254 14085.879 14069.504 14048.379
        ]
        #  Should be approx. equal to python:
        #    `res.variables["z"].data[0, 0, 0, :, :][-4:, -4:]`
        @test adjoint(res.variables["z"].data[1, 1, 1, :, :])[end-3:end, end-3:end] ≈ [
            12140.004 12226.504 12312.629 12406.004
            12417.629 12470.379 12520.129 12568.629
            12633.879 12658.754 12682.879 12705.879
            12660.379 12660.379 12660.379 12660.379
        ]

        @test res.variables["t"].dimensions == (
            "number", "time", "isobaricInhPa", "longitude", "latitude"
        )
        @test typeof(res.variables["t"].data) == CfGRIB.OnDiskArray
        @test size(res.variables["t"].data) == (10, 4, 2, 120, 61)
        @test_broken res.variables["t"].missing_value == 9999
        #  TODO: Make these into 'OrderedDict', and order correctly
        @test_broken typeof(res.variables["t"].attributes) == OrderedDict
        @test res.variables["t"].attributes == OrderedDict(
            "GRIB_paramId"                            => 130,
            "GRIB_shortName"                          => "t",
            "GRIB_units"                              => "K",
            "GRIB_name"                               => "Temperature",
            "GRIB_cfName"                             => "air_temperature",
            "GRIB_cfVarName"                          => "t",
            "GRIB_dataType"                           => "an",
            "GRIB_missingValue"                       => 9999,
            "GRIB_numberOfPoints"                     => 7320,
            "GRIB_totalNumber"                        => 10,
            "GRIB_typeOfLevel"                        => "isobaricInhPa",
            "GRIB_NV"                                 => 0,
            "GRIB_stepUnits"                          => 1,
            "GRIB_stepType"                           => "instant",
            "GRIB_gridType"                           => "regular_ll",
            "GRIB_gridDefinitionDescription"          => "Latitude/Longitude Grid",
            "GRIB_Nx"                                 => 120,
            "GRIB_iDirectionIncrementInDegrees"       => 3.0,
            "GRIB_iScansNegatively"                   => 0,
            "GRIB_longitudeOfFirstGridPointInDegrees" => 0.0,
            "GRIB_longitudeOfLastGridPointInDegrees"  => 357.0,
            "GRIB_Ny"                                 => 61,
            "GRIB_jDirectionIncrementInDegrees"       => 3.0,
            "GRIB_jPointsAreConsecutive"              => 0,
            "GRIB_jScansPositively"                   => 0,
            "GRIB_latitudeOfFirstGridPointInDegrees"  => 90.0,
            "GRIB_latitudeOfLastGridPointInDegrees"   => -90.0,
            "long_name"                               => "Temperature",
            "units"                                   => "K",
            "standard_name"                           => "air_temperature",
            "coordinates"                             => "number time step isobaricInhPa latitude longitude valid_time",
        )
        #  Should be approx. equal to python:
        #    `res.variables["t"].data[:, :, 0, 1, 2]`
        @test res.variables["t"].data[:, :, 1, 3, 2] ≈ [
            252.73932 253.1473  252.33952 253.19273
            252.73364 253.0484  252.51244 253.55228
            252.90796 253.02562 252.53241 253.17216
            252.62859 252.69759 252.48486 253.37857
            252.89291 253.22601 252.22084 254.01117
            252.91525 253.18896 252.48277 254.20941
            252.78569 253.13112 252.4354  253.75746
            252.38048 252.64601 252.12773 253.89052
            252.68658 252.9497  252.17213 254.07466
            252.59758 252.99641 252.4409  253.73424
        ]
        @test size(res.variables["t"].data[1, 1, 1, :, :]) == (120, 61) # ' vs python
        #  Should be approx. equal to python:
        #    `res.variables["t"].data[0, 0, 0, :, :][:4, :4]`
        @test adjoint(res.variables["t"].data[1, 1, 1, :, :])[1:4, 1:4] ≈ [
            252.66315 252.66315 252.66315 252.66315
            252.61041 252.68658 252.73932 252.78229
            251.11041 251.36627 251.6026  251.8487
            253.93658 254.41315 254.57526 254.62994
        ]
        #  Should be approx. equal to python:
        #    `res.variables["t"].data[0, 0, 0, :, :][-4:, -4:]`
        @test adjoint(res.variables["t"].data[1, 1, 1, :, :])[end-3:end, end-3:end] ≈ [
            260.59283 260.75104 260.6983  260.6983
            259.23737 259.14752 259.04987 258.8897
            258.56158 258.40143 258.27643 258.1983
            258.5401  258.5401  258.5401  258.5401
        ]
    end
end

cfgrib_dataset_py = pyimport("cfgrib.dataset")

test_files = [
    # "era5-levels-corrupted.grib",# - skip corrupted file tests
    "era5-levels-members.grib",
    "fields_with_missing_values.grib",
    "forecast_monthly_ukmo.grib",
    # "hpa_and_pa.grib",# - DatasetBuildError("multiple values for unique key, try re-open the file with one of
    "lambert_grid.grib",
    "multi_param_on_multi_dims.grib",
    # "reduced_gg.grib", - TODO: enable test with next release of Julia
    "regular_gg_ml.grib",
    "regular_gg_ml_g2.grib",
    "regular_gg_pl.grib",
    "regular_gg_sfc.grib",
    # "regular_gg_wrong_increment.grib", # Not clear how to handle data that does not fit into given array size
    "regular_ll_msl.grib",
    "regular_ll_sfc.grib",
    "regular_ll_wrong_increment.grib",
    "scanning_mode_64.grib",
    # "spherical_harmonics.grib", - ecCodes provides no latitudes/longitudes for gridType='sh'
    "t_analysis_and_fc_0.grib",
    # "t_on_different_level_types.grib", - DatasetBuildError("multiple values for unique key, try re-open the file with one of
    # "tp_on_different_grid_resolutions.grib", - DatasetBuildError("multiple values for unique key, try re-open the file with one of
    # "uv_on_different_levels.grib" - cfgrib.dataset.DatasetBuildError: key present and new value is different
]

attributes_key_blacklist = [
    "history"
]

python_type_mapping = Dict(
    "<class 'cfgrib.dataset.Variable'>" => CfGRIB.Variable,
    "<class 'cfgrib.dataset.OnDiskArray'>" => CfGRIB.OnDiskArray
)

@testset "pycall tests for $test_file" for test_file in test_files
    test_file_path = joinpath(dir_testfiles, test_file)
    res_py = cfgrib_dataset_py.open_file(test_file_path)
    res_jl = CfGRIB.DataSet(test_file_path)

    res_py_attributes = copy(res_py.attributes)
    [delete!(res_py_attributes, key) for key in attributes_key_blacklist]
    [delete!(res_jl.attributes, key) for key in attributes_key_blacklist]
    @test res_jl.attributes == res_py_attributes

    @test res_jl.dimensions == res_py.dimensions

    @test Set(keys(res_py.variables)) == Set(keys(res_jl.variables))

    @testset "variable $var" for var in keys(res_jl.variables)
        var_jl = res_jl.variables[var]
        var_py = res_py.variables[var]

        var_d_type_py = typeof(var_py.data)
        var_d_type_py_native = pystr(py"type($(var_py.data))")

        #  If the returned value is not a PyObject, it has been automatically
        #  converted to a native julia type, e.g. numpyarray/list -> Julia Array
        if var_d_type_py != PyObject
            if var_d_type_py <: Number
                @test var_jl.data ≈ var_py.data
            elseif var_d_type_py <: Array{T,1} where T <: Number
                #  One-dimensional arrays should be identical
                @test var_jl.data ≈ var_py.data
            elseif var_d_type_py <: Array
                #  Too much effort to try and figure out which dimension(s) have
                #  been transposed due to the row/col differences between julia
                #  and python, easier to just test if any case works
                if size(var_jl.data) == size(var_py.data)
                    if var_jl.data ≈ var_py.data
                        @test true
                    else
                        @test adjoint(var_jl.data) ≈ var_py.data
                    end
                elseif size(adjoint(var_jl.data)) == size(var_py.data)
                    @test adjoint(var_jl.data) ≈ var_py.data
                else
                    #  Values not equal for adjoint or direct comparison
                    @test false
                end
            else
                @warn "Julia Native var $var of type $var_d_type_py not tested"
            end
        #  If it is still a PyObject, it's most likely a CfGRIB 'OnDiskArray'
        elseif var_d_type_py == PyObject
            var_d_type_py_jl = python_type_mapping[var_d_type_py_native]
            if var_d_type_py_jl == CfGRIB.OnDiskArray
                @test var_jl.data.geo_ndim == var_py.data.geo_ndim

                jl_geo_dims = var_jl.dimensions[end-var_jl.data.geo_ndim+1:end]
                py_geo_dims = var_py.dimensions[end-var_jl.data.geo_ndim+1:end]
                transpose_flag = !(jl_geo_dims == py_geo_dims)

                multidim_idx = tuple(
                    ones(Int, length(size(var_jl.data)) - var_jl.data.geo_ndim)...,
                    repeat([Colon()], var_jl.data.geo_ndim)...
                )
                if transpose_flag
                    var_jl_dims = (
                        var_jl.dimensions[1:end-var_jl.data.geo_ndim]...,
                        var_jl.dimensions[end-var_jl.data.geo_ndim+1:end][end:-1:1]...
                    )
                    var_jl_data = adjoint(var_jl.data[multidim_idx...])
                elseif var_jl.data.geo_dims[end:-1:1] == var_py.data.geo_dims
                    var_jl_dims = var_jl.dimensions
                    var_jl_data = var_jl.data[multidim_idx...]
                end

                var_py_data = var_py.data.build_array()[multidim_idx...]
                var_py_data = replace(var_py_data, NaN=>missing)

                @test var_jl_dims == var_py.dimensions
                @test ismissing.(var_jl_data) ≈ ismissing.(var_py_data)
                @test collect(skipmissing(var_jl_data)) ≈ collect(skipmissing(var_py_data))
            else
                @warn "PyObject var $var of type $var_d_type_py not tested"
            end
        else
            @warn "var $var of type $var_d_type_py not tested"
        end

        @test var_jl.attributes == var_py.attributes
    end
end


text:
export rmap

rmap(f, x) = f(x)

function rmap(f, t::Tuple)
    map(x -> rmap(f,x), t)
end

function rmap(f, nt::NamedTuple{N,T}) where {N,T}
    NamedTuple{N}(map(x -> rmap(f,x), values(nt)))
end


text:
function boys_function(n::Int, x::Float64)::Float64
    x < 1e-6 && return 1/(2n+1)
    n == 0 && return 0.5*√(π/x)*erf(√x)
    return ((2n-1)*boys_function(n-1, x) - exp(-x))/(2x)
end

"S(gs, gs)"
function overlap_integral(g1::Gaussian_s, g2::Gaussian_s)::Float64
    p = (g1.α + g2.α)
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    return γ^(3/4) * exp(-μ * abs2(g1.R - g2.R))
end

"S(gs, gp)"
function overlap_integral(g1::Gaussian_s, g2::Gaussian_p)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    return √(g1.α) * γ^(5/4) * exp(-μ * abs2(Rab)) * (Rab * g2.n)
end

"S(gp, gs)"
function overlap_integral(g1::Gaussian_p, g2::Gaussian_s)::Float64
    overlap_integral(g2, g1)
end

"S(gp, gp)"
function overlap_integral(g1::Gaussian_p, g2::Gaussian_p)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    return γ^(5/4) * ((g1.n * g2.n) -√(g1.α*g2.α*γ) * (Rab * g1.n) * (Rab * g2.n)) * exp(-μ * abs2(Rab))
end

"T(gs, gs)"
function kinetic_integral(g1::Gaussian_s, g2::Gaussian_s)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rab2 = abs2(g1.R - g2.R)
    return γ^(3/4) * μ * (3-2μ*Rab2) * exp(-μ * Rab2)
end

"T(gp, gs)"
function kinetic_integral(g1::Gaussian_p, g2::Gaussian_s)::Float64
    p = g1.α + g2.α
    μ = g1.α * g2.α / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    Rab2 = abs2(Rab)
    return γ^(5/4) * exp(-μ * Rab2) * μ * (2μ*Rab2-5) * √g2.α * (Rab * g1.n)
end

"T(gs, gp)"
function kinetic_integral(g1::Gaussian_s, g2::Gaussian_p)::Float64
    kinetic_integral(g2, g1)
end

"T(gp, gp)"
function kinetic_integral(g1::Gaussian_p, g2::Gaussian_p)::Float64
    p = g1.α + g2.α
    μ = g1.α * g2.α / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    dotnab = g1.n * g2.n
    μRab2 = μ * abs2(Rab)
    M = (g1.n * Rab) * (g2.n * Rab)
    return γ^(5/4) * exp(-μRab2) * μ * (5dotnab - 14μ * M - 2μRab2 * dotnab + 4μ * μRab2 * M)
end

"Vn(gs, gs)"
function nuclear_acctaction_integral(Zc::Int, Rc::Vec3, g1::Gaussian_s, g2::Gaussian_s)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rμ = (g1.α / p) * g1.R + (g2.α / p) * g2.R
    return -2Zc * γ^(3/4) * √(p/π) * exp(-μ * abs2(g1.R - g2.R)) * boys_function(0, p*abs2(Rμ-Rc))
end

"Vn(gp,gs)"
function nuclear_acctaction_integral(Zc::Int, Rc::Vec3, g1::Gaussian_p, g2::Gaussian_s)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rμ = (g1.α / p) * g1.R + (g2.α / p) * g2.R
    pR2 = p*abs2(Rμ - Rc)
    return -4Zc * γ^(3/4) * √(g1.α * p / π) * exp(-μ * abs2(g1.R-g2.R)) *
        (-(g1.n * (Rμ-Rc)) * boys_function(1, pR2) + (g1.n * (Rμ-g1.R)) * boys_function(0, pR2))
end

"Vn(gs,gp)"
function nuclear_acctaction_integral(Zc::Int, Rc::Vec3, g1::Gaussian_s, g2::Gaussian_p)::Float64
    nuclear_acctaction_integral(Zc, Rc, g2, g1)
end

"Vn(gp, gp)"
function nuclear_acctaction_integral(Zc::Int, Rc::Vec3, g1::Gaussian_p, g2::Gaussian_p)::Float64
    p = g1.α + g2.α
    μ = g1.α * g2.α / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    Rμ = (g1.α / p) * g1.R + (g2.α / p) * g2.R
    Rμc = Rμ - Rc
    Dac = g1.n * Rμc
    Dbc = g2.n * Rμc
    Daa = g1.n * (Rμ - g1.R)
    Dbb = g2.n * (Rμ - g2.R)
    x = p * abs2(Rμc)
    F0x = boys_function(0, x)
    F1x = boys_function(1, x)
    F2x = boys_function(2, x)
    return -2Zc * γ^(5/4) * √(p/π) * exp(-μ*abs2(g1.R - g2.R)) * (
        (g1.n * g2.n) * (F0x - F1x) + 2p * Dac * Dbc * F2x -
        2p * (Dac * Dbb + Dbc * Daa) * F1x + 2p * Daa * Dbb * F0x
    )
end

function electron_repulsion_integral(g1::Gaussian_s, g2::Gaussian_s, g3::Gaussian_s, g4::Gaussian_s)
    p1 = g1.α + g4.α
    p2 = g2.α + g3.α
    μ1 = (g1.α * g4.α) / p1
    μ2 = (g2.α * g3.α) / p2
    γ1 = 4μ1 / p1
    γ2 = 4μ2 / p2
    Rμ1 = (g1.α / p1) * g1.R + (g4.α / p1) * g4.R
    Rμ2 = (g2.α / p2) * g2.R + (g3.α / p2) * g3.R
    λ = p1 * p2 / (p1 + p2)
    R = abs(Rμ1 - Rμ2)
    if R < 1e-6
        return (γ1*γ2)^(3/4) * exp(-μ1 * abs2(g1.R-g4.R)-μ2 * abs2(g2.R-g3.R)) * 2 * √(λ / π)
    else
        return (γ1*γ2)^(3/4) * exp(-μ1 * abs2(g1.R-g4.R)-μ2 * abs2(g2.R-g3.R)) * erf(√λ * R) / R
    end
end

function two_center_integral(s1::STONG{N}, s2::STONG{N}, integral::Function) where N
    total = 0.0
    for i = 1:N
        for j = 1:N
            total += s1.d[i] * s2.d[j] * integral(s1.g[i], s2.g[j])
        end
    end
    return total
end

function four_center_integral(s1::STONG{N}, s2::STONG{N}, s3::STONG{N}, s4::STONG{N}, integral::Function) where N
    total = 0.0
    for i = 1:N
        for j = 1:N
            for k = 1:N
                for l = 1:N
                    total += s1.d[i]*s2.d[j]*s3.d[k]*s4.d[l] * integral(s1.g[i], s2.g[j], s3.g[k], s4.g[l])
                end
            end
        end
    end
    return total
end

overlap_integral(s1::STONG, s2::STONG) = two_center_integral(s1, s2, overlap_integral)
kinetic_integral(s1::STONG, s2::STONG) = two_center_integral(s1, s2, kinetic_integral)
nuclear_acctaction_integral(Zc::Int, Rc::Vec3, s1::STONG, s2::STONG) = two_center_integral(
    s1, s2,
    (g1,g2)->nuclear_acctaction_integral(Zc, Rc, g1, g2)
)
electron_repulsion_integral(s1::STONG, s2::STONG, s3::STONG, s4::STONG) = four_center_integral(
    s1, s2, s3, s4,
    electron_repulsion_integral
)


text:
# Loading packages can often take a long time, but Julia is very fast once packages are loaded
# using Zygote  # for reverse-mode auto-differentiation
using Random  # Random numbers
using Distributions
using DataFrames
using DataFramesMeta
using CategoricalArrays
using LinearAlgebra  # Diagonal
using StatsModels  # formulas and modelmatrix
using GLM  # OLS (lm) and GLS (glm)
using Optim  # optimize (minimizing)
# using Symbolics  # symbolic functions
import Dates
using Latexify  # output symbolic expressions as latex code
# using ModelingToolkit  # more symbolic matrix functions
# using Reduce  # more symbolic matrix functions (specifically Reduce.Algebra.^)
# @force using Reduce.Algebra  # to extend native functions to Sybmol/Expr types
using Logging
import CSV

"""
NOTES:
Purpose: Estimate the parameters ρ, σₐ², σᵤ² of Larry Karp's emissions
    model using iterated Feasible GLS and MLE

How: Most of this script defines functions that represent equations from
    the model. Finally, in the "Iterate to convergence!" section at the
    end, the equations are used to iteratively estimate:
    (1) the GLS parameters b₀ᵢ and β (time trend) as inputs to (2)
    (2) the MLE parameters ρ, σₐ², σᵤ²

First, I test this on generated data that we know the parameter values
    for to see how close the estimates are to the "true" values. I test
    the estimates using different values of T, N, σₐ², σᵤ².

Lastly, I (still need to) estimate the paramters on the actual data that
    Larry provided.

Time to run: Most of the functions at the end that actually run the
    estimations also have time estimates in a comment after them. These
    were run on a linux computer with 32 GB of RAM and
    AMD Ryzen 3 2300x quad-core processor. Times may be more or less
    depending on your system. Julia does multithreading automatically,
    meaning it will take up all available CPU during most compuations
    until the script has complete.
"""

#######################################################################
#           Setup
#######################################################################
io = open("log.txt", "a")
logger = SimpleLogger(io)
# debuglogger = ConsoleLogger(stderr, Logging.Debug)
global_logger(logger)
s = "\n"^20 * "="^60 * "\nCOVARIANCE MATRIX ESTIMATION BEGIN\n" * "="^60
@info s Dates.format(Dates.now(), "yyyy-mm-dd HH:MM:SS")

#######################################################################
#           Equation Functions
#######################################################################

# Building the covariance matrix for the N=2, T=3 case
# Not needed anymore but nice to keep for tests
N = 2;
T = 3;
A(ρ, σₐ², σᵤ²) = 1 / (1 - ρ^2) * σₐ² + (1 + ρ^2 / (N * (1 - ρ^2))) * σᵤ²;
B(ρ, σₐ², σᵤ²) = 1 / (1 - ρ^2) * σₐ² + (ρ^2 / (N * (1 - ρ^2))) * σᵤ²;
C(ρ, σₐ², σᵤ²) = ρ / (1 - ρ^2) * σₐ² + (ρ / N + ρ^3 / (N * (1 - ρ^2))) * σᵤ²;
E(ρ, σₐ², σᵤ²) = ρ * C(ρ, σₐ², σᵤ²);

function Σ(ρ, σₐ², σᵤ²)
    [A(ρ, σₐ², σᵤ²) B(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) E(ρ, σₐ², σᵤ²) E(ρ, σₐ², σᵤ²)
        B(ρ, σₐ², σᵤ²) A(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) E(ρ, σₐ², σᵤ²) E(ρ, σₐ², σᵤ²)
        C(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) A(ρ, σₐ², σᵤ²) B(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²)
        C(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) B(ρ, σₐ², σᵤ²) A(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²)
        E(ρ, σₐ², σᵤ²) E(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) A(ρ, σₐ², σᵤ²) B(ρ, σₐ², σᵤ²)
        E(ρ, σₐ², σᵤ²) E(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) C(ρ, σₐ², σᵤ²) B(ρ, σₐ², σᵤ²) A(ρ, σₐ², σᵤ²)]
end

# Building the covariance matrix for general N, T
# Indicator functions
ι(i, j) = (i == j ? 1 : 0)
κ(s) = (s == 0 ? 0 : 1)
χ(ρ, i, j, s, N) = (1 - κ(s)) * ι(i, j) + κ(s) * ρ^s / N + ρ^(2 + s) / (N * (1 - ρ))

# Element of the covariance matrix: Σ = E[vᵢₜvⱼₜ₊ₛ]
Evᵢₜvⱼₜ₊ₛ(ρ, σₐ², σᵤ², i, j, s, N; b = 1) = 1 / b^2 * (σₐ² * ρ^s / (1 - ρ^2) + χ(ρ, i, j, s, N) * σᵤ²)

# i,t, j,t+s Elements of ∂Σ/∂y for y ∈ {ρ, σₐ², σᵤ²}
∂Σ∂ρ(ρ, σₐ², σᵤ², i, j, s, N) = ρ^(s-1) / N / (ρ^2 - 1)^2 * (
    (s*κ(s) + (2 + s - 2*s*κ(s))*ρ^2 + (s*κ(s) - s)*ρ^4)*σᵤ² +
    (N*(2-s)*ρ^2 + N*s)*σₐ²
)

∂Σ∂σₐ²(ρ, σₐ², σᵤ², i, j, s, N) = ρ^s / (1 - ρ^2)

∂Σ∂σᵤ²(ρ, σₐ², σᵤ², i, j, s, N) = χ(ρ, i, j, s, N)

derivatives = Dict("ρ" => ∂Σ∂ρ, "σₐ²" => ∂Σ∂σₐ², "σᵤ²" => ∂Σ∂σᵤ²)

# General derivative ∂Σ/∂y for y ∈ {"ρ", "σₐ²", "σᵤ²"} (quotes needed for y input)
∂Σ∂y(y, ρ, σₐ², σᵤ², i, j, s, N) = derivatives[y](ρ, σₐ², σᵤ², i, j, s, N)

"""
    Σ(ρ, σₐ², σᵤ², N, T)

Return the residuals' analytical N*TxN*T covariance matrix given values of parameters
@param ρ: float in [0,1], decay rate of AR1 emissions process
@param σₐ²: float >0, SD of year-specific emissions shock
@param σᵤ²: float >0, SD of region-year-specific emissions shock
@param N: int>1, number of units/regions
@param T: int>1, number of time periods
"""
function Σ(ρ, σₐ², σᵤ², N, T; verbose = false)
    # Initalize matrix of 0s
    V = zeros(N * T, N * T)

    # Fill in upper triangle
    idx = [(i, j) for i ∈ 1:N*T for j ∈ i:N*T]
    for (row, col) in idx
        t = Integer(ceil(row / N))
        i = row - (t - 1) * N

        τ = Integer(ceil(col / N))
        s = τ - t
        j = col - (τ - 1) * N

        V[row, col] = Evᵢₜvⱼₜ₊ₛ(ρ, σₐ², σᵤ², i, j, s, N)
        a = [i, t, j, τ, row, col]
    end

    # Fill in lower triangle by symmetry
    V = Symmetric(V)
    if verbose
        println("\nFull Covariance Matirx:")
        for r in 1:size(V)[1]
            println(round.(V[r, :]; sigdigits = 4))
        end
    end
    return (V)
end

"""
Function to fill the NT x NT matrix with elements defined by ijs_func
"""
function fill_matrix(ijs_func, ρ, σₐ², σᵤ², N, T; verbose = false)
    # Initalize matrix of 0s
    mat = zeros(N * T, N * T)

    # Fill in upper triangle
    idx = [(i, j) for i ∈ 1:N*T for j ∈ i:N*T]
    for (row, col) in idx
        t = Integer(ceil(row / N))
        i = row - (t - 1) * N

        τ = Integer(ceil(col / N))
        s = τ - t
        j = col - (τ - 1) * N

        mat[row, col] = ijs_func(ρ, σₐ², σᵤ², i, j, s, N)
    end

    # Fill in lower triangle by symmetry
    mat = Symmetric(mat)
    if verbose
        println("\nFull NT x NT Matirx:")
        for r in 1:size(mat)[1]
            println(round.(mat[r, :]; sigdigits = 4))
        end
    end
    return (mat)
end

# Negative Log Likelihood (want to minimize)
function nLL(ρ, σₐ², σᵤ², v, N, T)
    V = Σ(ρ,σₐ²,σᵤ²,N,T)
    nLL = (1/2) * ( v'*V^-1*v + log(det(V)) )
    return(nLL)
end

function ∂nLL∂y(y, ρ, σₐ², σᵤ², v, N, T)
    ijs_fun(ρ, σₐ², σᵤ², i, j, s, N) = ∂Σ∂y(y, ρ, σₐ², σᵤ², i, j, s, N)
    ∂V∂y = fill_matrix(ijs_fun, ρ, σₐ², σᵤ², N, T)

    V = Σ(ρ,σₐ²,σᵤ²,N,T)
    vV⁻¹∂V∂yV⁻¹v = v'*V^-1*∂V∂y*V^-1*v
    V⁻¹∂V∂y = V^-1*∂V∂y

    ∂LL∂y = -1/2 * (-vV⁻¹∂V∂yV⁻¹v + tr(V⁻¹∂V∂y))
    return(-∂LL∂y)
end

# Define the gradient of the negative log likelihood
function nLL_grad(gradient_vec, params, v, N, T)
    ρ = params[1]; σₐ² = params[2]; σᵤ² = params[3]
    gradient_vec[1] = ∂nLL∂y("ρ", ρ, σₐ², σᵤ², v, N, T)
    gradient_vec[2] = ∂nLL∂y("σₐ²", ρ, σₐ², σᵤ², v, N, T)
    gradient_vec[3] = ∂nLL∂y("σᵤ²", ρ, σₐ², σᵤ², v, N, T)
end

"""
    mymle(ρstart, σₐ²start, σᵤ²start, v)

Return likelihood-maximizing values of ρ, σₐ², and σᵤ².
Given vector of residuals, minimize the
@param ρstart: float, starting value of ρ
@param σₐ²start: float, starting value of σₐ²
@param σᵤ²start: float, starting value of σᵤ²
@param v: Nx1 vector of residuals
@param N: int>1, number of units/regions
@param T: int>1, number of time periods
@param lower: float 3-vector, lower bounds for parameters ρ, σₐ², σᵤ²
@param upper: float 3-vector, upper bounds for parameters ρ, σₐ², σᵤ²
"""
function mymle(ρstart, σₐ²start, σᵤ²start, v, N, T;
    lower = [1e-4, 1e-4, 1e-4], upper = [1 - 1e-4, Inf, Inf], analytical = true)

    # Starting paramerter values
    params0 = [ρstart, σₐ²start, σᵤ²start]

    # Function of only parameters (given residuals v)
    objective(ρσₐ²σᵤ²) = nLL(ρσₐ²σᵤ²[1], ρσₐ²σᵤ²[2], ρσₐ²σᵤ²[3], v, N, T)

    # Minimize over parameters
    grad!(storage, x) = nLL_grad(storage, x, v, N, T)

    # Calculates gradient with no parameter bounds
    # optimum = optimize(objective, params0, LBFGS())
    # Results in a domain error in the log function because the gradient pushes
    #   it to a negative value inside the log

    # Analytical gradient with no parameter bounds
    # optimum = optimize(objective, grad!, params0, LBFGS())
    # Results in a domain error in the log function because the gradient pushes
    #   it to a negative value inside the log

    if analytical
        # Analytical gradient with parameter bounds
        optimum = optimize(objective, grad!, lower, upper, params0, Fminbox(GradientDescent()))
    else
        # Calculates gradient with parameter bounds
        optimum = optimize(objective, lower, upper, params0, Fminbox(GradientDescent()))
    end

    # Return the values
    LL = -optimum.minimum
    ρ, σₐ², σᵤ² = optimum.minimizer
    return (ρ, σₐ², σᵤ², LL)
end

"""
    dgp(ρ, σₐ², σᵤ², β, N, T; v₀, μ₀, σ₀)

Return simulated data from the data generating process given paramters.
@param ρ: float in [0,1], decay rate of AR1 emissions process
@param σₐ²: float >0, SD of year-specific emissions shock
@param σᵤ²: float >0, SD of region-year-specific emissions shock
@param β: 2-vector, linear and quadratic time trend parameters
@param v₀: float, initial emissions shock of AR1 process (v_t where t=0)
@param b₀: Nx1 array of floats, b₀ᵢ for different regions i in {1, ..., N}
    if not given, then μ₀ and σ₀ are used to pull b₀ᵢ from a random distribution
@param μ₀: float, mean of region-specific fixed effect distribution (b₀ᵢ)
@param σ₀: float >0, SD of region-specific fixed effect distribution (b₀ᵢ)
"""
function dgp(ρ, σₐ², σᵤ², β, N, T;
    v₀ = 0.0, b₀ = nothing, μ₀ = 0, σ₀ = 10, random_seed::Integer = 1234)
    # Initial conditions
    Random.seed!(random_seed)
    b = 1  # unidentified scale parameter

    # Get region-specific fixed effects if not given
    if b₀ === nothing
        b₀ = rand(Distributions.Normal(μ₀, σ₀), N)
    end

    # Random shocks
    αₜ = rand(Distributions.Normal(0, σₐ²^0.5), T)
    μᵢₜ = rand(Distributions.Normal(0, σᵤ²^0.5), N * T)
    # assume μᵢₜ is stacked region first:
    # i.e. (i,t)=(1,1) (i,t)=(2,1) ... (i,t)=(N-1,T)  (i,t)=(N,T)

    # Fill in the aggregate shocks
    vₜ = [v₀]
    vᵢₜ = []
    for t in 1:T
        s = 0  # this period's sum of shocks
        for i in 1:N
            # Note that vₜ index starts at 1 instead of 0, so it's ahead of the others
            append!(vᵢₜ, ρ * vₜ[t] + αₜ[t] + μᵢₜ[(t-1)*N+i])
            s += last(vᵢₜ)
        end
        append!(vₜ, s / N)  # Aggregate shock = average of this period's shocks
    end

    data = DataFrame(t = repeat(1:T, inner = N),
        i = string.(repeat(1:N, outer = T)),
        b₀ᵢ = repeat(b₀, outer = T),
        αₜ = repeat(αₜ, inner = N),
        μᵢₜ = μᵢₜ,
        vᵢₜ = vᵢₜ,
        vₜ = repeat(vₜ[2:(T+1)], inner = N))

    # Generate the resulting emissions
    data.eᵢₜ = (1 / b) * (data.b₀ᵢ + β[1] * data.t + β[2] * data.t .^ 2 + data.vᵢₜ)

    return (data)
end;

"""
    mygls(formula, df, W, N)

@param formula: StatsModels formula for the linear regression model
@param df: dataframe with columns corresponding to the variables in formula
@param W: nxn weighting matrix, n = length(df)
@param N: number of units
@returns :β, :βse, :yhat, :resid, :βvar, :HC0, :sandwich

# References
- Bruce Hansen Econometrics (2021) section 17.15 Feasible GLS
- Greene Ed 7
- rsusmel lecture notes: bauer.uh.edu/rsusmel/phd/ec1-11.pdf
"""
function mygls(formula, df, W, N)
    X = StatsModels.modelmatrix(formula.rhs, df)
    y = StatsModels.modelmatrix(formula.lhs, df)
    XX = inv(X'X)
    XWX = X' * inv(W) * X
    NT, k = size(X)

    β = inv(XWX) * (X' * inv(W) * y)
    yhat = X * β
    resid = y - yhat
    XeX = X' * Diagonal(vec(resid .^ 2)) * X

    # Variance estimators
    # Hansen eq 4.16 (sandwich estimator)
    s = XX * XWX * XX
    s_root = diag(s) .^ 0.5
    # White estimator (rsusmel pg 11)
    HC0 = XX * XeX * XX
    HC0_root = diag(HC0) .^ 0.5
    # Greene (9-13)
    βvar = inv(XWX) / (NT - N - k)
    βse = diag(βvar) .^ 0.5

    # TODO: Add Newey-West HAC SE estimator.
    #   pg. 20 of bauer.uh.edu/rsusmel/phd/ec1-11.pdf

    # Go with βse for now
    return(Dict(:β => β,
        :βse => βse,
        :yhat => yhat,
        :resid => resid,
        :βvar => βvar,
        :HC0 => HC0_root,
        :sandwich => s_root
    ))
end

function estimate_rho()
    filepath = "ts_allYears_nation.1751_2014.csv"
    # See r-scripts/reproducing_andy.R result1 and result2
end

#######################################################################
#           Helper functions
#######################################################################

create_dataframe() = DataFrame(
    seed = [], type = [], N = [], T = [],
    b₀ᵢ = [], β = [], ρ = [], σₐ² = [], σᵤ² = [], iterations = [],
    LL = [], V_norm_finalstep = [], V_norm_true=[], V = [],
    runtime = [], notes=[]
)

function add_row(df, seed, type, N, T;
    b₀ᵢ=missing, β=missing, ρ=missing, σₐ²=missing, σᵤ²=missing, iterations=missing,
    LL=missing, V_norm_finalstep=missing, V_norm_true=missing, V=missing,
    runtime=missing, notes=missing)
    push!(df, Dict(
    :seed=>seed, :type=>type, :N=>N, :T=>T,
    :b₀ᵢ=>b₀ᵢ, :β=>β, :ρ=>ρ, :σₐ²=>σₐ², :σᵤ²=>σᵤ², :LL=>LL, :iterations=>iterations,
    :V_norm_finalstep=>V_norm_finalstep, :V_norm_true=>V_norm_true, :V=>V,
    :runtime=>runtime, :notes=>notes
    ))
end

function write_simulation_df(df; N=2)
    # Save dataframe
    filepath = "../../data/temp/simulation_results_log(N=$N).csv"
    if isfile(filepath)
        CSV.write(filepath, df,append=true)
    else
        CSV.write(filepath, df)
    end
end

function write_estimation_df(df; N=2)
    # Save dataframe
    filepath = "../../data/temp/realdata_estimation_results_log(N=$N).csv"
    if isfile(filepath)
        CSV.write(filepath, df, append=true)
    else
        CSV.write(filepath, df)
    end
end

function read_data(N, T)
    filepath = "../../data/sunny/clean/grouped_nation.1751_2014.csv"
    df = @chain CSV.read(filepath, DataFrame) begin
        @rsubset :Year >= 1945 && :Year < 1945+T && :group ∈ ["USA","EU","BRIC","Other"][1:N]
        @select(:t = :time .- 44,
                :i = categorical(:group),
                :eᵢₜ = :CO2)
        @orderby(:t, :i)
    end

    return(df)
end

function read_global_data(T)
    filepath = "../../data/sunny/clean/ts_allYears_nation.1751_2014.csv"
    df = @chain CSV.read(filepath, DataFrame) begin
        @rsubset :Year >= 1945 && :Year < 1945+T && :group ∈ ["USA","EU","BRIC","Other"][1:N]
        @select(:t = :time .- 44,
                :i = categorical(:group),
                :eᵢₜ = :CO2)
        @orderby(:t, :i)
    end

    return(df)
end

function get_covar_names(data)
    ilevels = unique(data.i); len = length(ilevels);
    inames = repeat(["i"], len-1) .* string.(ilevels[2:len])
    varnames = vcat(["intercept"], inames, ["t", "t^2"])
    return(varnames)
end

function get_sample_variances(data)
    σₐ² = sum((data.αₜ .- mean(data.αₜ)) .^ 2) / 2 / (T - 1)
    σᵤ² = sum((data.μᵢₜ .- mean(data.μᵢₜ)) .^ 2) / (N * T - 1)
    return(σₐ², σᵤ²)
end

function save_reformatted_data()
    N=2; T=60;
    data = read_data(N, T)
    filepath = "../../data/temp/realdata_reformated(N=$N).csv"
    CSV.write(filepath, data)
end

#######################################################################
#           Iterate to convergence!
#######################################################################
"""
Examine convergence behavior of covariance matrix
"""
function estimate_sigmas(N, T, starting_params;
        n_seeds = 10, iteration_max = 500, convergence_threshold = 1e-16,
        analytical = true)
    println()
    display([N,T,starting_params])
    @info "Starting estimation with N,T:" N T starting_params n_seeds iteration_max convergence_threshold
    df = create_dataframe()
    # Generate different data to see how convergence behaves
    println("\n", "_"^n_seeds)
    for seed in 1:n_seeds
        print("*")
        @info "seed: " seed
        Random.seed!(seed)
        # Starting parameter values
        ρ = starting_params[:ρ];    β = starting_params[:β]
        σₐ² = starting_params[:σₐ²];  σᵤ² = starting_params[:σᵤ²]
        LL = -Inf
        # Simuated data
        data = dgp(ρ, σₐ², σᵤ², β, N, T; random_seed = seed)
        # True values from DGP
        b₀ᵢ = unique(data.b₀ᵢ)
        V_true = Σ(ρ, σₐ², σᵤ², N, T)
        add_row(df, seed, "True", N, T; b₀ᵢ=b₀ᵢ[:], β=β, ρ=ρ, σₐ²=σₐ², σᵤ²=σᵤ², LL=LL)
        # Observed sample variance from DGP
        σ_observed = get_sample_variances(data)
        add_row(df, seed, "Sample", N, T; σₐ²=σ_observed[1], σᵤ²=σ_observed[2])
        # Initialize the variance matrix (identity matrix = OLS in first iteration)
        V = Diagonal(ones(length(data.eᵢₜ)))
        # Linear GLS formula
        linear_formula = @formula(eᵢₜ ~ 1 + i + t + t^2)

        Vnorm = Inf; i = 1;
        while Vnorm > convergence_threshold && i <= iteration_max
            # GLS to get residuals v
            global gls = mygls(linear_formula, data, V, N)
            v = vec(gls[:resid])
            # Uncomment the below to remove detrending/demeaning step
            # v = vec(data.vᵢₜ)

            # ML to get parameter estimates ρ, σₐ², σᵤ²
            ρ, σₐ², σᵤ², LL = mymle(ρ, σₐ², σᵤ², v, N, T; analytical)
            Vold = V
            V = Σ(ρ, σₐ², σᵤ², N, T)
            Vnorm = norm(Vold - V)
            i += 1
        end

        β = gls[:β][3:4]; b₀ᵢ[1] = gls[:β][1]; b₀ᵢ[2] = sum(gls[:β][1:2]);
        Vnorm_true = norm(V_true - V)
        add_row(df, seed, "Estimated (calc grad, bounds)", N, T;
            b₀ᵢ=b₀ᵢ, β=β, ρ=ρ, σₐ²=σₐ², σᵤ²=σᵤ², LL=LL, iterations=i-1,
            V_norm_finalstep=Vnorm, V_norm_true=Vnorm_true,
            runtime=Dates.format(Dates.now(), "yyyy-mm-dd HH:MM:SS")
        )
    end

    # Save dataframe of simulation results
    # write_simulation_df(df; N=N)
end

"""
Roughly estimating the model parameters to put into DGP for testing.
Also used to do final estimation...
@param starting_params: dict with ρ, σₐ², σᵤ² starting values
"""
function estimate_dgp_params(N, T, starting_params;
        iteration_max = 500, convergence_threshold = 1e-4,
        params_lower_bound = [1e-4, 1e-4, 1e-4],
        params_upper_bound = [1 - 1e-4, Inf, Inf],
        data = missing)
    println("Estimating parameters of real data using iteration method.")
    # dataframe to save to CSV
    df = create_dataframe()
    # Load data, or use provided simulated data
    if ismissing(data)
        data = read_data(N, T)
    end
    # Initialize the variance matrix (identity matrix = OLS in first iteration)
    V = Diagonal(ones(length(data.eᵢₜ)))
    # Linear GLS formula
    linear_formula = @formula(eᵢₜ ~ 1 + i + t + t^2)

    # Starting parameter values
    ρ = starting_params[:ρ]; σₐ² = starting_params[:σₐ²];  σᵤ² = starting_params[:σᵤ²]
    Vnorm = Inf; i = 1; LL = -Inf
    while Vnorm > convergence_threshold && i <= iteration_max
        # GLS to estimate b₀ᵢ and β; get residuals v
        global gls = mygls(linear_formula, data, V, N)
        v = vec(gls[:resid])

        # ML to get parameter estimates ρ, σₐ², σᵤ²
        ρ, σₐ², σᵤ², LL = mymle(ρ, σₐ², σᵤ², v, N, T;
                                lower=params_lower_bound, upper=params_upper_bound)
        Vold = V
        V = Σ(ρ, σₐ², σᵤ², N, T)
        Vnorm = norm(Vold - V)
        println(Vnorm)
        i += 1
    end

    β = gls[:β][3:4];   b₀ᵢ = [gls[:β][1], sum(gls[:β][1:2])];
    note = join(["lower param bounds ", join(params_lower_bound, ", "), " :: ",
                 "upper param bounds ", join(params_upper_bound, ", "), " :: ",
                 "starting params ", join(collect(keys(starting_params)), ", "), ": ",
                 join(collect(values(starting_params)), ", ")])
    add_row(df, "real data", "Estimated", N, T;
        b₀ᵢ=b₀ᵢ, β=β, ρ=ρ, σₐ²=σₐ², σᵤ²=σᵤ², LL=LL, iterations=i-1,
        V_norm_finalstep=Vnorm,
        runtime=Dates.format(Dates.now(), "yyyy-mm-dd HH:MM:SS"),
        notes=note
    )

    # Save dataframe of results
    write_estimation_df(df; N=N)
    println(note)
    println("DONE estimating parameters of real data.")
end

#######################################################################
#           Testing convergence properties of different dimensions
#######################################################################
N = 2;
T = 3;
starting_params = Dict(:ρ => 0.85,
    :σₐ² => 10,
    :σᵤ² => 10,
    :β => [1, 0.1]
)

param(plist) = Dict(
    :ρ => plist[1],
    :σₐ² => plist[2],
    :σᵤ² => plist[3],
)

save_reformatted_data()

function test_dgp_params_from_real()
    N = 2; T = 60  # 1945-2005
    # parameters estimated from data
    b₀ = [200000, 560000]  # from test_starting_real_estimation()
    β = [22000, -100]      # from test_starting_real_estimation()
    ρ = 0.8785             # from r-scripts/reproducting_andy.R  -> line 70 result2
    σₐ² = 40               # from test_starting_real_estimation()
    σᵤ² = 265              # from test_starting_real_estimation()
    params = Dict(:b₀ => b₀, :β => β, :ρ => 0.8785, :σₐ² => 40, :σᵤ² => 265)

    # Generate simulated data from realistic parameter estimates
    simulated_data = dgp(params[:ρ], params[:σₐ²], params[:σᵤ²], params[:β], N, T;
        v₀ = 0.0, b₀ = params[:b₀])
    filepath = "../../data/temp/simulated_data rho$ρ sig_a$σₐ² sig_u$σᵤ².csv"
    CSV.write(filepath, simulated_data)

    # See how close the estimated parameters are to the parameters used to generate the data
    # Test various starting parameter values
    # Lock in ρ around 0.87
    # Starting parameters for the search ρ, σₐ², σᵤ²
    plist1 = [0.8785, 1, 1]
    plist2 = [0.8785, 0.1, 0.1]
    plist3 = [0.8785, 10, 10]
    plist4 = [0.8785, 100, 100]
    plist5 = [0.8785, 1000, 1000]
    lower_bounds1 = [0, 1e-4, 1e-4]
    upper_bounds1 = [1, Inf, Inf]

    # @time estimate_dgp_params(N, T, param(plist1); data=simulated_data,
    #     params_lower_bound=lower_bounds1,
    #     params_upper_bound=upper_bounds1)

    # @time estimate_dgp_params(N, T, param(plist2); data=simulated_data,
    #     params_lower_bound=lower_bounds1,
    #     params_upper_bound=upper_bounds1)

    @time estimate_dgp_params(N, T, param(plist3); data=simulated_data,
        params_lower_bound=lower_bounds1,
        params_upper_bound=upper_bounds1)

    @time estimate_dgp_params(N, T, param(plist4); data=simulated_data,
        params_lower_bound=lower_bounds1,
        params_upper_bound=upper_bounds1)

    # @time estimate_dgp_params(N, T, param(plist5); data=simulated_data,
    #     params_lower_bound=lower_bounds1,
    #     params_upper_bound=upper_bounds1)

end
# test_dgp_params_from_real()
# Takes about 1.5 hours to run and drives sigma_a to 0

function test_starting_real_estimation()
    # Starting parameters for the search ρ, σₐ², σᵤ²
    plist1 = [0.85, 1, 1]
    plist2 = [0.1, 0.1, 0.1]
    plist3 = [0.1, 10, 10]
    plist4 = [0.1, 100, 100]
    plist5 = [0.1, 1000, 1000]
    plist6 = [0.99, 0.1, 0.1]
    plist7 = [0.99, 10, 10]
    plist8 = [0.99, 100, 100]
    plist9 = [0.99, 1000, 1000]
    N = 2;
    T = 60;
    @time estimate_dgp_params(N, T, param(plist1))

    @time estimate_dgp_params(N, T, param(plist2))

    @time estimate_dgp_params(N, T, param(plist3))

    @time estimate_dgp_params(N, T, param(plist4))

    @time estimate_dgp_params(N, T, param(plist5))

    @time estimate_dgp_params(N, T, param(plist6))

    @time estimate_dgp_params(N, T, param(plist7))

    @time estimate_dgp_params(N, T, param(plist8))

    @time estimate_dgp_params(N, T, param(plist9))
end
# test_starting_real_estimation()
# about 2 minutes to finish

function test_bounds_real_estimation()
    # lower bounds on params: ρ, σₐ², σᵤ²
    lower_bounds1 = [1e-4, 1e-4, 1e-4]
    lower_bounds2 = [0.1, 1, 1]
    upper_bounds1 = [1, Inf, Inf]
    upper_bounds2 = [0.85, Inf, Inf]
    N = 2;
    T = 60;
    @time estimate_dgp_params(N, T, starting_params;
        params_lower_bound=lower_bounds1,
        params_upper_bound=upper_bounds1)

    @time estimate_dgp_params(N, T, starting_params;
        params_lower_bound=lower_bounds1,
        params_upper_bound=upper_bounds2)

    @time estimate_dgp_params(N, T, starting_params;
        params_lower_bound=lower_bounds2,
        params_upper_bound=upper_bounds1)

    @time estimate_dgp_params(N, T, starting_params;
        params_lower_bound=lower_bounds2,
        params_upper_bound=upper_bounds2)
    # T=3:   53.165765 seconds (136.57 M allocations: 8.168 GiB,  4.27% gc time, 90.99% compilation time)
    # T=60:  92.230013 seconds (354.54 M allocations: 24.525 GiB, 4.01% gc time, 51.16% compilation time)
end

# Timing test between analytical gradient and calcualted gradient
# RESULT: analytical gradients take about 5/6 of the time as calculating gradients,
#       more "allocations" but less memory is written to
function time_gradient_methods()
    @time estimate_sigmas(2, 50, starting_params; n_seeds=10, analytical=true)
    # 232.415821 seconds (1.20 G allocations: 101.854 GiB, 3.80% gc time, 7.80% compilation time)
    @time estimate_sigmas(2, 50, starting_params; n_seeds=10, analytical=true)
    # 267.077041 seconds (1.15 G allocations: 99.125 GiB, 3.25% gc time)
    @time estimate_sigmas(2, 50, starting_params; n_seeds=10, analytical=false)
    # 316.247974 seconds (655.02 M allocations: 139.433 GiB, 3.79% gc time, 0.08% compilation time)
    @time estimate_sigmas(2, 50, starting_params; n_seeds=10, analytical=false)
    # 315.076701 seconds (654.42 M allocations: 139.398 GiB, 4.04% gc time)
end

function test_estimation()
    # Do the estimates get closer as T gets larger? (consistency)
    println("\nTesting T")
    for T ∈ [5,10,25,50,100]
        estimate_sigmas(2, T, starting_params; n_seeds=10)
    end

    # Do the estimates get closer as N gets larger? (consistency)
    println("\nTesting N")
    for N ∈ [3,4]
        estimate_sigmas(N, 50, starting_params; n_seeds=10)
    end

    # Is the quality of the estimates sensitive to the relative size of σₐ²
    println("\nTesting σₐ²")
    for σₐ² ∈ [0.1, 1, 5, 10]
        starting_params = Dict(:ρ => 0.85,
            :σₐ² => σₐ²,
            :σᵤ² => 1,
            :β => [1, 0.1]
        )
        estimate_sigmas(2, 50, starting_params; n_seeds=10)
    end

    # Is the quality of the estimates sensitive to the relative size of σᵤ²
    println("\nTesting σᵤ²")
    for σᵤ² ∈ [0.1, 1, 5, 10]
        starting_params = Dict(:ρ => 0.85,
            :σₐ² => 1,
            :σᵤ² => σᵤ²,
            :β => [1, 0.1]
        )
        estimate_sigmas(2, 50, starting_params; n_seeds=10)
    end
end

# Close logging file
close(io)

"""
Developer Notes:
- Try simple case with rho = 1 or zero or other cases Larry mentions

Using the residuals from the true beta and FE, we get a little bit closer
to small Vnorm between the trueV norm and the estV norm (24.2 instead of 27.4)

# Testing glm function (can't take non-diagonal matrix)
gls = glm(linear_formula, data, Normal(), IdentityLink(), wts=ones(length(data.eᵢₜ)))

# Creating a more analytical gradient
objective(ρ, σₐ², σᵤ², v) = norm(maximand(ρ, σₐ², σᵤ², v));
ρ, σₐ², σᵤ² = ρstart, σₐ²start, σᵤ²start
for k in 1:iteration_max
    grad = gradient(objective(ρ, σₐ², σᵤ², v))
# Could also reverse-mode auto-differentiate to get complete analytical gradient

# Testing symbolic gradient
@variables ρ σₐ² σᵤ²
V = Σ(2, 3; verbose=false)
Vinv = inv(V)
Vinv = substitute(Vinv, Dict(true => 1))
println(Vinv, "\n")
v = ones(N*T)
nLL(v, N, T) = (1/2)*( v' * inv(Σ(N, T)) * v + log(det(Σ(N, T))) );
nLL = nLL(v, 2, 3)
println(nLL, "\n")
nLLgrad = Symbolics.gradient(nLL, [ρ,σₐ²,σᵤ²])
println(nLLgrad)

Gls: could decompose using Cholskey and augment the data than feed it into glm
linear model and use CovarianceMatrices from https://gragusa.org/code/ to get HAC SEs
M.diagonalize()

Iterate many times creating new data and save the vnorm from trueV, see the
dist of vnorm and distance from true parameter values. Is it just the
randomness of small sampling?

Also create general V function building function from N, T that returns a
    function to generate V given params (general function should be used
        once only for given def of N,T)

Want to compare N=2, small T to N=4, large T. Large T should have better
convergence? Since we can get more info about variance from more obs

REFERENCES:
If you use Symbolics.jl, please cite this paper
@article{gowda2021high,
  title={High-performance symbolic-numerics via multiple dispatch},
  author={Gowda, Shashi and Ma, Yingbo and Cheli, Alessandro and Gwozdz, Maja and Shah, Viral B and Edelman, Alan and Rackauckas, Christopher},
  journal={arXiv preprint arXiv:2105.03949},
  year={2021}
}

"""


text:
# PDEs
abstract type AbstractPDEs end

mutable struct VarElliptic <: AbstractPDEs
    # domain is fixed [0,1]*[0,1]
    a::Function
    rhs::Function
    bdy_type::Function # 1 Diri, 2 Neum, 3 Robin
    bdy_Diri::Function
    bdy_Neum::Function
    bdy_Robin::Function
end


text:
# AS clause.

mutable struct AsClause <: AbstractSQLClause
    over::Union{SQLClause, Nothing}
    name::Symbol

    AsClause(;
             over = nothing,
             name::Union{Symbol, AbstractString}) =
        new(over, Symbol(name))
end

AsClause(name; over = nothing) =
    AsClause(over = over, name = name)

"""
    AS(; over = nothing, name)
    AS(name; over = nothing)

An `AS` clause.

# Examples

```jldoctest
julia> c = ID(:person) |> AS(:p);

julia> print(render(c))
"person" AS "p"
```
"""
AS(args...; kws...) =
    AsClause(args...; kws...) |> SQLClause

dissect(scr::Symbol, ::typeof(AS), pats::Vector{Any}) =
    dissect(scr, AsClause, pats)

Base.convert(::Type{AbstractSQLClause}, p::Pair{<:Union{Symbol, AbstractString}}) =
    AsClause(name = first(p), over = convert(SQLClause, last(p)))

function PrettyPrinting.quoteof(c::AsClause, ctx::QuoteContext)
    ex = Expr(:call, nameof(AS), quoteof(c.name))
    if c.over !== nothing
        ex = Expr(:call, :|>, quoteof(c.over, ctx), ex)
    end
    ex
end

rebase(c::AsClause, c′) =
    AsClause(over = rebase(c.over, c′), name = c.name)


text:
module Three

if startswith(msg, "say:")
    println("EMITTING MESSAGE: $msg")
    for (k,v) in connections
        if k != client.id
            write(v, (usernames[client.id] * ": " * msg[5: end]))
        end
    end
end

function makePoint(x::Float32, y::Float32)
   # this is a useless Function to make a Point
   point = Point(x,y)
   return point
end

struct Foo
          bar
          baz::Int
          qux::Float64
end
# package code goes here

if something
 c = a + b # Some coment
elseif somethingelse
 c = a * b
else
 c = a - b
end

end


text:
## Makie wrapper for implicit equations
## How to integrate in....
function implicit_equation(f, axes...; iteration::Int=4, constraint=nothing)

    axes = [axes...]

    if constraint == nothing
        prob = MDBM_Problem(f, axes)
    else
        prob = MDBM_Problem(f, axes, constraint=constraint)
    end

    solve!(prob, iteration)

    prob
end

#==

AbstractPlotting.plot(m::MDBM_Problem; kwargs...) = plot!(Scene(), m; kwargs...)
AbstractPlotting.plot!(m::MDBM_Problem; kwargs...) = plot!(AbstractPlotting.current_scene(), m; kwargs...)
AbstractPlotting.plot!(scene::AbstractPlotting.Scene, m::MDBM_Problem; kwargs...) =
    plot!(Val(_dim(m)), scene, m; kwargs...)

_dim(m::MDBM.MDBM_Problem{a,N,b,c,d,e,f,g,h}) where {a,N,b,c,d,e,f,g,h} = N

# 2D plot
function AbstractPlotting.plot!(::Val{2}, scene::AbstractPlotting.Scene,
    m::MDBM_Problem; color=:black, kwargs...)

    mdt=MDBM.connect(m)
    for i in 1:length(mdt)
        dt=mdt[i]
        P1=getinterpolatedsolution(m.ncubes[dt[1]], m)
        P2=getinterpolatedsolution(m.ncubes[dt[2]], m)
        lines!(scene, [P1[1],P2[1]],[P1[2],P2[2]], color=color, kwargs...)
    end

    scene
end

# 3D plot
function AbstractPlotting.plot!(::Val{3}, scene::AbstractPlotting.Scene,
    m::MDBM_Problem; color=:black, kwargs...)

    positions = Point{3, Float32}[]
    scales = Vec3[]

    mdt=MDBM.connect(m)
    for i in 1:length(mdt)
        dt=mdt[i]
        P1=getinterpolatedsolution(m.ncubes[dt[1]], m)
        P2=getinterpolatedsolution(m.ncubes[dt[2]], m)

        a, b = Vec3(P1), Vec3(P2)
        push!(positions, Point3(P1))
        push!(scales, b-a)
    end

    cube = Rect{3, Float32}(Vec3(-0.5, -0.5, -0.5), Vec3(1, 1, 1))
    meshscatter!(scene, positions, marker=cube, scale = scales, color=color, transparency=true, kwargs...)

    scene
end

==#