### A Pluto.jl notebook ###
# v0.17.1

using Markdown
using InteractiveUtils

# ╔═╡ dc3c2030-144f-11eb-08ab-0b68dc2698ac
begin
	using Pkg
	cd(joinpath(dirname(@__FILE__),".."))
    Pkg.activate(pwd())
	using BenchmarkTools
end

# ╔═╡ 615d2b78-9a47-41bf-bae4-f0f7fb875956
# Make cells wider
html"""<style>
/*              screen size more than:                     and  less than:                     */
@media screen and (max-width: 699px) { /* Tablet */
  /* Nest everything into here */
    main { /* Same as before */
        max-width: 1200px !important; /* Same as before */
        margin-right: 100px !important; /* Same as before */
    } /* Same as before*/

}

@media screen and (min-width: 700px) and (max-width: 1199px) { /* Laptop*/
  /* Nest everything into here */
    main { /* Same as before */
        max-width: 1200px !important; /* Same as before */
        margin-right: 100px !important; /* Same as before */
    } /* Same as before*/
}

@media screen and (min-width:1000px) and (max-width: 1920px) { /* Desktop */
  /* Nest everything into here */
    main { /* Same as before */
        max-width: 1000px !important; /* Same as before */
        margin-right: 100px !important; /* Same as before */
    } /* Same as before*/
}

@media screen and (min-width:1921px) { /* Stadium */
  /* Nest everything into here */
    main { /* Same as before */
        max-width: 1200px !important; /* Same as before */
        margin-right: 100px !important; /* Same as before */
    } /* Same as before*/
}
</style>
"""

# ╔═╡ 1c49d8ae-144e-11eb-024a-137460f78b15
md"""
# Performance
When building more complex programs, performance can be an issue. The [performance tips](https://docs.julialang.org/en/v1/manual/performance-tips/) website already contains a lot tips. We will illustrate potential performance gains for a function by providing one or more alternatives. In order to make sure the the functions produce the same result, the results are verified to be identical before starting the benchmarks.

The actual performance is measured using [BenchmarkTools](https://github.com/JuliaCI/BenchmarkTools.jl). More advanced methods (that we cannot cover due to time constraints) included profiling your code. In Visual Studio Code you can get a detailed overview with the [`@profview`](https://discourse.julialang.org/t/julia-vs-code-extension-version-v0-17-released/42865) macro. For a REPL based method you could use [Traceur](https://github.com/JunoLab/Traceur.jl).

## Tools
Below you can find a set of very generic functions that we will be using to evaluate the functions that we will be benchmarking.

### Notes
At some point we try to make use of multithreading. You can see the number of threads that are currently being used with
```julia
Threads.nthreads()
```

Should this be only 1, you can set the number of threads to use by setting the environment variable `JULIA_NUM_THREADS` to the required before starting julia.

In a windows terminal:
```
$env:JULIA_NUM_THREADS = 4
julia
```
In a Linux/Mac OS terminal:
```
export JULIA_NUM_THREADS=4
julia
```

both examples above suppose "julia" is known in your terminal. If this is not the case, simply replace it by the path to the julia executable.

*Remark*: starting from julia 1.5 you can use command line arguments to specify the number of threads e.g. `julia --threads 4`
"""

# ╔═╡ a4ed00b8-144f-11eb-0c66-932242a1767c
begin
	"""
    testfuns(F, args...; kwargs...)

Benchmark the functions in F using `args` and `kwargs`.
"""
function testfuns(F::Union{Array{Function,1}, Array{Type,1}}, args...; kwargs...)
    # set up benchmark suite
    suite = BenchmarkGroup()
    for f in F
        suite[string(f)] = @benchmarkable $(f)($(args)...; $(kwargs)...)
    end
    # tune and run suite
    tune!(suite)
    results = run(suite, verbose=get(kwargs, :verbose, false))
    # show results
    for f in F
        println("\nResult for $(f):\n$("-"^length("Result for $(f):"))")
        display(results[string(f)])
    end
    medians = sort!([(fun, median(trial.times)) for (fun, trial) in results], by=x->x[2])
    speedup = round(Int,medians[end][2] / medians[1][2])
    println("\n$(speedup)x speedup (fastest: $(medians[1][1]), slowest: $(medians[end][1]))\n")
    return

end
	nothing
end

# ╔═╡ e32569b0-144f-11eb-33e9-8fd8755c208f
begin
	"""
		equality(F::Array{Function,1}, args...; kwargs...)

	Verify the result of an array of functions F is the same using `args` and `kwargs`.

	### Notes
	The first function in the array is used as the reference
	"""
	function equality(F::Array{Function,1}, args...; kwargs...)
		return all([equality(F[1], F[i], args...; kwargs...) for i in 2:length(F)])
	end

	"""
		equality(f::Function, g::Function, args...; kwargs...)

	Verify the result of function f and g is the same (element-wise) using `args` and `kwargs`.
	"""
	function equality(f::Function, g::Function, args...; kwargs...)
		return all(isequal.(map(f->f(args...; kwargs...), [f, g])...))
	end

	nothing
end

# ╔═╡ ff480c88-144f-11eb-2839-598b8452a945
md"""
## Hands-on
"""

# ╔═╡ 1e0ef46a-1450-11eb-2676-d94d5cc1f073
begin

	"""
		renew(char="", n::Int=0)

	Slow implementation of string joining
	"""
	function renew(char="", n::Int=0; kwargs...)
		res = ""
		for i in 1:n
			res *= char
		end

		return res
	end

	"""
		buffer(char="", n::Int=0)

	Fast implementation of string joining
	"""
	function buffer(char="", n::Int=0; kwargs...)
		res = IOBuffer()
		for i in 1:n
			print(res, char)
		end

		return String(take!(res))
	end

	"""
		demo1

	- Illustration of implementation issues

	speedup ∝ length
	"""
	function demo1()
		F = [renew, buffer];
		args=("a", 10000)
		# test validity
		@info "Testing demo 1"
		@assert equality(F, args...)
		# Run benchmark
		@info "Benching demo 1"
		testfuns(F, args...)
	end

	demo1()

end

# ╔═╡ 41ef249a-1450-11eb-2c52-0930531fa46e
begin
	function f_alloc(x::Vector{T}, y::Array{T,2}; props::Vector{T}, kwargs...) where T
		# read properties
		prop_A = props[1]
		prop_B = props[2]
		prop_C = props[3]
		prop_D = props[4]
		# initialise
		result = zeros(T, size(x))
		result .+= x
		# update result
		result = result .+ prop_A .* y[:,1]
		result = result .+ prop_B .* y[:,2]
		result = result .+ prop_C .* y[:,3]
		result = result .+ prop_D .* y[:,4]
		return result
	end

	struct Props{T<:AbstractFloat}
		prop_A::T
		prop_B::T
		prop_C::T
		prop_D::T
	end

	@views function f_noalloc(x::Vector{T}, y::Array{T,2}; p::Props, kwargs...) where T
		return x .+ p.prop_A .* y[:,1] .+ p.prop_B .* y[:,2] + p.prop_C .* y[:,3] + p.prop_D .* y[:,4]
	end

	"""
		demo2()

	- Avoiding unnecessary allocations
	- Using views
	- Avoiding fields with abstract types
	"""
	function demo2()
		F = [f_alloc, f_noalloc]
		N = 10000
		x = rand(N)
		y = rand(N,4)
		args = [x, y]
		props = Float64[1,2,3,4]
		p = Props(props...)
		kwargs = Dict(:props => props, :p => p)
		# test validity
		@info "Testing demo 2"
		@assert equality(F, args...; kwargs...)
		# run benchmark
		@info "Benching demo 2"
		testfuns(F, args...; kwargs...)
	end

	demo2()

end

# ╔═╡ 6332ef7e-1450-11eb-0a36-41818207b78f
begin
	# avoid changing type of a variable
	function f_typechange(x)
		res = 0
		for val in x
			res += val
		end
		return res
	end

	function f_notypechange(x::Array{T,1}) where T <: Number
		res = zero(T)
		for val in x
			res += val
		end
		return res
	end

	"""
		demo3()

	- avoid changing type of a variable
	"""
	function demo3()
		F = [f_typechange, f_notypechange]
		N = 2
		args = [rand(N)]
		# test validity
		@info "Testing demo 3"
		@assert equality(F, args...)
		# run benchmark
		@info "Benching demo 3"
		testfuns(F, args...)
	end

	demo3()
end

# ╔═╡ 75948388-1450-11eb-0a0e-431ad5929504
begin
	f_norm(x) = 3x.^2 + 4x + 7x.^3;
	f_vec(x) = @. 3x.^2 + 4x + 7x.^3;

	"""
		demo 4

	- use vectorisation where possible
	"""
	function demo4()
		F = [f_norm, f_vec]
		N::Int = 1e6
		args = [rand(N)]
		# test validity
		@info "Testing demo 4"
		@assert equality(F, args...)
		# run benchmark
		@info "Benching demo 4"
		testfuns(F, args...)
	end

	demo4()

end

# ╔═╡ 98e356f4-1450-11eb-0229-eb57f262ea4d
begin
	function f_nopre(x::Vector; kwargs...)
		y = similar(x)
		for i in eachindex(x)
			y[i] = sum(x[1:i])
		end
		return y
	end

	function f_pre!(x::Vector{T}; y::Vector{T}, kwargs...) where {T}
		for i in eachindex(x)
			y[i] = sum(x[1:i])
		end
		return y
	end

	function f_preboost!(x::Vector{T}; y::Vector{T}, kwargs...) where {T}
		Threads.@threads for i in eachindex(x)
			@inbounds y[i] = sum(x[1:i])
		end
		return y
	end

	@views function f_preturboboost!(x::Vector{T}; y::Vector{T}, kwargs...) where {T}
		Threads.@threads for i in eachindex(x)
			@inbounds y[i] = sum(x[1:i])
		end
		return y
	end

	"""
		demo5()

	- preallocate outputs if possible
	- make use of threads if possible
	- make use of views if possible
	- skip index check when sure of dimensions
	"""
	function demo5()
		F = [f_nopre, f_pre!, f_preboost!, f_preturboboost!]
		@info "Currently running on $(Threads.nthreads()) threads"
		N::Int = 1e4
		x::Vector = rand(N)
		y_pre = similar(x)
		args = [x]
		kwargs = Dict( :y => y_pre)
		# test validity
		@info "Testing demo 5"
		@assert equality(F, args...; kwargs...)
		# run benchmark
		@info "Benching demo 5"
		testfuns(F, args...; kwargs...)
	end

	demo5()
end

# ╔═╡ b341b4d2-1450-11eb-0549-7db5bd6756ba
begin
	function grow(N::Int; kwargs...)
		res = []
		for _ in 1:N
			push!(res, 1)
		end
		return res
	end

	function nogrow(N::Int; x::Array{T,1}) where{T}
		res = similar(x)
		for i in 1:N
			res[i] = 1
		end
		return res
	end

	"""
		demo6()

	- growing vs. allocating for known dimensions
	"""
	function demo6()
		N = 10000
		x = rand(N)
		F = [grow, nogrow]
		args = [N]
		kwargs = Dict(:x => x)
		# test validity
		@info "Testing demo 6"
		@assert equality(F, args...; kwargs...)
		# run benchmark
		@info "Benching demo 6"
		testfuns(F, args...; kwargs...)
	end

	demo6()
end

# ╔═╡ c6702048-1450-11eb-1f1d-a54089545949
begin
	struct mystruct
		a
		b
		c
	end

	struct mytypedstruct{T}
		a::T
		b::T
		c::T
	end

	mutable struct mymutablestruct
		a
		b
		c
	end

	"""
		demo7

	- creation time of mutable vs unmutable structs and the parameter version
	"""
	function demo7()
		structs = [mystruct, mytypedstruct, mymutablestruct]
		@info "Benching demo 7"
		testfuns(structs, [1;2;3]...)
	end

	demo7()

end

# ╔═╡ df4563da-1450-11eb-08a7-25ed365f578c
begin
	function reader(s::mystruct)
		return (s.a, s.b, s.c)
	end

	function reader(s::mytypedstruct)
		return (s.a, s.b, s.c)
	end

	function reader(sm::mymutablestruct)
		return (sm.a, sm.b, sm.c)
	end

	"""
		demo8

	- read times of fields of mutable vs unmutable structs and the parameter version
	"""
	function demo8()
		vals = [1;2;3]
		@info "Benching demo8"
		for arg in [mystruct, mytypedstruct, mymutablestruct]
			println("\nResult for $(arg):\n$("-"^length("Result for $(arg):"))")
			x = arg(vals...)
			display(@benchmark reader($(x)))
		end
	end

	demo8()
end

# ╔═╡ Cell order:
# ╟─615d2b78-9a47-41bf-bae4-f0f7fb875956
# ╟─1c49d8ae-144e-11eb-024a-137460f78b15
# ╠═dc3c2030-144f-11eb-08ab-0b68dc2698ac
# ╠═a4ed00b8-144f-11eb-0c66-932242a1767c
# ╠═e32569b0-144f-11eb-33e9-8fd8755c208f
# ╟─ff480c88-144f-11eb-2839-598b8452a945
# ╠═1e0ef46a-1450-11eb-2676-d94d5cc1f073
# ╠═41ef249a-1450-11eb-2c52-0930531fa46e
# ╠═6332ef7e-1450-11eb-0a36-41818207b78f
# ╠═75948388-1450-11eb-0a0e-431ad5929504
# ╠═98e356f4-1450-11eb-0229-eb57f262ea4d
# ╠═b341b4d2-1450-11eb-0549-7db5bd6756ba
# ╠═c6702048-1450-11eb-1f1d-a54089545949
# ╠═df4563da-1450-11eb-08a7-25ed365f578c


"""
    Calculators

This module exists to bridge the gap between the `Models` and the `Dynamics`.

Here, we provide functions and types for evaluating and storing quantities obtained from the
`Models`.
In addition any further manipulation of those quantities, such as computing eigenvalues,
is included here.

This module is largely needed to facilitate integration of both ring polymer and
classical dynamics to allow using the same models and functions for both.
Specific ring polymer types are provided that have the extra fields and methods needed
to evaluate the quantities for each bead.
"""
module Calculators

using LinearAlgebra: LinearAlgebra, Hermitian, I, Eigen, tr
using StaticArrays: SMatrix, SVector
using RingPolymerArrays: get_centroid!

using NQCModels: NQCModels, Model, nstates
using NQCModels.AdiabaticModels: AdiabaticModel
using NQCModels.DiabaticModels: DiabaticModel, DiabaticFrictionModel
using NQCModels.FrictionModels: AdiabaticFrictionModel

using NQCDynamics: ndofs

"""
    AbstractCalculator{M<:Model}

Top-level type for all calculators.

Each concrete calculator contains the `Model` and the fields to store the quantities
obtained from the model.
"""
abstract type AbstractCalculator{T,M<:Model} end
abstract type AbstractAdiabaticCalculator{T,M<:AdiabaticModel} <: AbstractCalculator{T,M} end
abstract type AbstractDiabaticCalculator{T,M<:Union{DiabaticFrictionModel,DiabaticModel}} <: AbstractCalculator{T,M} end
abstract type AbstractStaticDiabaticCalculator{T,M} <: AbstractDiabaticCalculator{T,M} end
abstract type AbstractFrictionCalculator{T,M<:AdiabaticFrictionModel} <: AbstractCalculator{T,M} end

NQCModels.nstates(calc::AbstractCalculator) = NQCModels.nstates(calc.model)
NQCModels.eachstate(calc::AbstractCalculator) = NQCModels.eachstate(calc.model)
NQCModels.nelectrons(calc::AbstractCalculator) = NQCModels.nelectrons(calc.model)
NQCModels.eachelectron(calc::AbstractCalculator) = NQCModels.eachelectron(calc.model)
NQCModels.mobileatoms(calc::AbstractCalculator) = NQCModels.mobileatoms(calc.model, size(calc.derivative, 2))
NQCModels.dofs(calc::AbstractCalculator) = NQCModels.dofs(calc.model)

Base.eltype(::AbstractCalculator{T}) where {T} = T

mutable struct DependentField{V,R}
    value::V
    position::R
end
needsupdate(field::DependentField, r) = field.position != r

function Base.getproperty(calc::AbstractCalculator, name::Symbol)
    field = getfield(calc, name)
    if field isa DependentField
        return field.value
    else
        return field
    end
end

function Base.setproperty!(calc::AbstractCalculator, name::Symbol, x)
    field = getfield(calc, name)
    if field isa DependentField
        setfield!(field, :value, x)
    else
        setfield!(calc, name, x)
    end
end

"""
Each of the quantities specified here has functions:
`get_quantity(calculator, r)`
`evaluate_quantity!(calculator, r)!`

The user should access only the former.
This will ensure quantities are correctly evaluated and cached accordingly.

The latter is called by the former and is where the details required to calculate the quantity are found.
"""
const quantities = [
    :potential,
    :derivative,
    :eigen,
    :adiabatic_derivative,
    :nonadiabatic_coupling,

    :traceless_potential,
    :V̄,
    :traceless_derivative,
    :D̄,
    :traceless_adiabatic_derivative,

    :centroid,
    :centroid_potential,
    :centroid_derivative,
    :centroid_eigen,
    :centroid_adiabatic_derivative,
    :centroid_nonadiabatic_coupling,

    :friction,
]

for quantity in quantities
    get_quantity = Symbol(:get_, quantity)
    evaluate_quantity! = Symbol(:evaluate_, quantity, :!)
    field = Expr(:call, :getfield, :calculator, QuoteNode(quantity))

    @eval function $(get_quantity)(calculator, r)

        if needsupdate($field, r)
            copyto!($field.position, r)
            $(evaluate_quantity!)(calculator, r)
        end
        return $field.value
    end
end

struct AdiabaticCalculator{T,M} <: AbstractAdiabaticCalculator{T,M}
    model::M
    potential::DependentField{T,Matrix{T}}
    derivative::DependentField{Matrix{T},Matrix{T}}
    stats::Dict{Symbol,Int}
    function AdiabaticCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        potential = zero(T)
        derivative = zeros(T, ndofs(model), atoms)
        position = fill(NaN, ndofs(model), atoms)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
        )

        new{T,M}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            stats
        )
    end
end

struct RingPolymerAdiabaticCalculator{T,M} <: AbstractAdiabaticCalculator{T,M}
    model::M
    potential::DependentField{Vector{T},Array{T,3}}
    derivative::DependentField{Array{T,3},Array{T,3}}
    stats::Dict{Symbol,Int}
    function RingPolymerAdiabaticCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        potential = zeros(T, beads)
        derivative = zeros(T, ndofs(model), atoms, beads)
        position = fill(NaN, ndofs(model), atoms, beads)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
        )

        new{T,M}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            stats
        )
    end
end

struct DiabaticCalculator{T,M,S,L} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::DependentField{Hermitian{T,SMatrix{S,S,T,L}},Matrix{T}}
    derivative::DependentField{Matrix{Hermitian{T,SMatrix{S,S,T,L}}},Matrix{T}}
    eigen::DependentField{LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}},Matrix{T}}
    adiabatic_derivative::DependentField{Matrix{SMatrix{S,S,T,L}},Matrix{T}}
    nonadiabatic_coupling::DependentField{Matrix{SMatrix{S,S,T,L}},Matrix{T}}
    stats::Dict{Symbol,Int}
    function DiabaticCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        n = nstates(model)
        mat = NQCModels.DiabaticModels.matrix_template(model, T)
        vec = NQCModels.DiabaticModels.vector_template(model, T)

        potential = Hermitian(zero(mat))
        derivative = [Hermitian(zero(mat)) for _=1:ndofs(model), _=1:atoms]
        eigen = Eigen(zero(vec), zero(mat) + I)
        adiabatic_derivative = [zero(mat) for _ in CartesianIndices(derivative)]
        nonadiabatic_coupling = [zero(mat) for _ in CartesianIndices(derivative)]
        position = fill(NaN, ndofs(model), atoms)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
            :eigen=>0,
            :adiabatic_derivative=>0,
            :nonadiabatic_coupling=>0
        )

        new{T,M,n,n^2}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            DependentField(eigen, copy(position)),
            DependentField(adiabatic_derivative, copy(position)),
            DependentField(nonadiabatic_coupling, copy(position)),
            stats
        )
    end
end

struct RingPolymerDiabaticCalculator{T,M,S,L} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::DependentField{Vector{Hermitian{T,SMatrix{S,S,T,L}}},Array{T,3}}
    derivative::DependentField{Array{Hermitian{T,SMatrix{S,S,T,L}},3},Array{T,3}}
    eigen::DependentField{Vector{LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}}},Array{T,3}}
    adiabatic_derivative::DependentField{Array{SMatrix{S,S,T,L},3},Array{T,3}}
    nonadiabatic_coupling::DependentField{Array{SMatrix{S,S,T,L},3},Array{T,3}}

    traceless_potential::DependentField{Vector{Hermitian{T,SMatrix{S,S,T,L}}},Array{T,3}}
    V̄::DependentField{Vector{T},Array{T,3}}
    traceless_derivative::DependentField{Array{Hermitian{T,SMatrix{S,S,T,L}},3},Array{T,3}}
    D̄::DependentField{Array{T,3},Array{T,3}}
    traceless_adiabatic_derivative::DependentField{Array{SMatrix{S,S,T,L},3},Array{T,3}}

    centroid::DependentField{Matrix{T},Array{T,3}}
    centroid_potential::DependentField{Hermitian{T,SMatrix{S,S,T,L}},Array{T,3}}
    centroid_derivative::DependentField{Matrix{Hermitian{T,SMatrix{S,S,T,L}}},Array{T,3}}
    centroid_eigen::DependentField{LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}},Array{T,3}}
    centroid_adiabatic_derivative::DependentField{Matrix{SMatrix{S,S,T,L}},Array{T,3}}
    centroid_nonadiabatic_coupling::DependentField{Matrix{SMatrix{S,S,T,L}},Array{T,3}}

    stats::Dict{Symbol,Int}
    function RingPolymerDiabaticCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        n = nstates(model)
        mat = NQCModels.DiabaticModels.matrix_template(model, T)
        vec = NQCModels.DiabaticModels.vector_template(model, T)

        potential = [Hermitian(zero(mat)) for _=1:beads]
        derivative = [Hermitian(zero(mat)) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        adiabatic_derivative = [zero(mat) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        eigen = [Eigen(zero(vec), zero(mat) + I) for _=1:beads]
        nonadiabatic_coupling = [zero(mat) for _=1:ndofs(model), _=1:atoms, _=1:beads]

        traceless_potential = [Hermitian(zero(mat)) for _=1:beads]
        V̄ = zeros(beads)
        traceless_derivative = [Hermitian(zero(mat)) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        D̄ = zeros(ndofs(model), atoms, beads)
        traceless_adiabatic_derivative = [zero(mat) for _=1:ndofs(model), _=1:atoms, _=1:beads]

        centroid = zeros(T, ndofs(model), atoms)
        centroid_potential = Hermitian(zero(mat))
        centroid_derivative = [Hermitian(zero(mat)) for _=1:ndofs(model), _=1:atoms]
        centroid_eigen = Eigen(zero(vec), zero(mat) + I)
        centroid_adiabatic_derivative = [zero(mat) for _ in CartesianIndices(centroid_derivative)]
        centroid_nonadiabatic_coupling = [zero(mat) for _ in CartesianIndices(centroid_derivative)]

        position = fill(NaN, ndofs(model), atoms, beads)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
            :eigen=>0,
            :adiabatic_derivative=>0,
            :nonadiabatic_coupling=>0,
            :traceless_potential=>0,
            :V̄=>0,
            :traceless_derivative=>0,
            :D̄=>0,
            :traceless_adiabatic_derivative=>0,
            :centroid=>0,
            :centroid_potential=>0,
            :centroid_derivative=>0,
            :centroid_eigen=>0,
            :centroid_adiabatic_derivative=>0,
            :centroid_nonadiabatic_coupling=>0,
        )

        new{T,M,n,n^2}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            DependentField(eigen, copy(position)),
            DependentField(adiabatic_derivative, copy(position)),
            DependentField(nonadiabatic_coupling, copy(position)),

            DependentField(traceless_potential, copy(position)),
            DependentField(V̄, copy(position)),
            DependentField(traceless_derivative, copy(position)),
            DependentField(D̄, copy(position)),
            DependentField(traceless_adiabatic_derivative, copy(position)),

            DependentField(centroid, copy(position)),
            DependentField(centroid_potential, copy(position)),
            DependentField(centroid_derivative, copy(position)),
            DependentField(centroid_eigen, copy(position)),
            DependentField(centroid_adiabatic_derivative, copy(position)),
            DependentField(centroid_nonadiabatic_coupling, copy(position)),
            stats
        )
    end
end

function Calculator(model::DiabaticModel, atoms::Integer, t::Type{T}) where {T}
    DiabaticCalculator{t}(model, atoms)
end
function Calculator(model::AdiabaticModel, atoms::Integer, t::Type{T}) where {T}
    AdiabaticCalculator{t}(model, atoms)
end
function Calculator(model::DiabaticModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerDiabaticCalculator{t}(model, atoms, beads)
end
function Calculator(model::AdiabaticModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerAdiabaticCalculator{t}(model, atoms, beads)
end

function evaluate_potential!(calc::AbstractCalculator, R)
    calc.stats[:potential] += 1
    calc.potential = NQCModels.potential(calc.model, R)
    return nothing
end

function evaluate_potential!(calc::AbstractCalculator, R::AbstractArray{T,3}) where {T}
    calc.stats[:potential] += 1
    @views for i in axes(R, 3)
        calc.potential[i] = NQCModels.potential(calc.model, R[:,:,i])
    end
    return nothing
end

function evaluate_V̄!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:V̄] += 1
    potential = get_potential(calc, r)
    for i in 1:length(calc.V̄)
        calc.V̄[i] = tr(potential[i]) / nstates(calc.model)
    end
end

function evaluate_traceless_potential!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:traceless_potential] += 1
    n = nstates(calc.model)
    potential = get_potential(calc, r)
    V̄ = get_V̄(calc, r)
    for I in eachindex(potential)
        calc.traceless_potential[I] = Hermitian(SMatrix{n,n}(
            i != j ? potential[I][j,i] : potential[I][j,i] - V̄[I] for j=1:n, i=1:n
        ))
    end
end

function evaluate_centroid!(calc::AbstractCalculator, r::AbstractArray{T,3}) where {T}
    calc.stats[:centroid] += 1
    get_centroid!(calc.centroid, r)
end

function evaluate_centroid_potential!(calc::AbstractCalculator, r::AbstractArray{T,3}) where {T}
    calc.stats[:centroid_potential] += 1
    centroid = get_centroid(calc, r)
    calc.centroid_potential = NQCModels.potential(calc.model, centroid)
end

function evaluate_derivative!(calc::AbstractCalculator, R)
    calc.stats[:derivative] += 1
    NQCModels.derivative!(calc.model, calc.derivative, R)
end

function evaluate_derivative!(calc::AbstractCalculator, R::AbstractArray{T,3}) where {T}
    calc.stats[:derivative] += 1
    @views for i in axes(R, 3)
        NQCModels.derivative!(calc.model, calc.derivative[:,:,i], R[:,:,i])
    end
end

function evaluate_D̄!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:D̄] += 1
    derivative = get_derivative(calc, r)
    for I in eachindex(derivative)
        calc.D̄[I] = tr(derivative[I]) / nstates(calc.model)
    end
end

function evaluate_traceless_derivative!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:traceless_derivative] += 1
    n = nstates(calc.model)
    derivative = get_derivative(calc, r)
    D̄ = get_D̄(calc, r)
    for I in eachindex(derivative)
        calc.traceless_derivative[I] = Hermitian(SMatrix{n,n}(
            i != j ? derivative[I][j,i] : derivative[I][j,i] - D̄[I] for j=1:n, i=1:n
        ))
    end
end

function evaluate_traceless_adiabatic_derivative!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:traceless_adiabatic_derivative] += 1
    n = nstates(calc.model)
    adiabatic_derivative = get_adiabatic_derivative(calc, r)
    D̄ = get_D̄(calc, r)
    for I in eachindex(D̄)
        calc.traceless_adiabatic_derivative[I] = SMatrix{n,n}(
            i != j ? adiabatic_derivative[I][j,i] : adiabatic_derivative[I][j,i] - D̄[I] for j=1:n, i=1:n
        )
    end
end

function evaluate_centroid_derivative!(calc::AbstractCalculator, r::AbstractArray{T,3}) where {T}
    calc.stats[:centroid_derivative] += 1
    centroid = get_centroid(calc, r)
    NQCModels.derivative!(calc.model, calc.centroid_derivative, centroid)
end

function evaluate_eigen!(calc::AbstractDiabaticCalculator, r)
    calc.stats[:eigen] += 1
    potential = get_potential(calc, r)
    eig = LinearAlgebra.eigen(potential)
    corrected_vectors = correct_phase(eig.vectors, calc.eigen.vectors)
    calc.eigen = Eigen(eig.values, corrected_vectors)
    return nothing
end

function evaluate_centroid_eigen!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:centroid_eigen] += 1
    potential = get_centroid_potential(calc, r)
    eig = LinearAlgebra.eigen(potential)
    corrected_vectors = correct_phase(eig.vectors, calc.centroid_eigen.vectors)
    calc.centroid_eigen = Eigen(eig.values, corrected_vectors)
    return nothing
end

function evaluate_eigen!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:eigen] += 1
    potential = get_potential(calc, r)
    for i=1:length(potential)
        eig = LinearAlgebra.eigen(potential[i])
        corrected_vectors = correct_phase(eig.vectors, calc.eigen[i].vectors)
        calc.eigen[i] = Eigen(eig.values, corrected_vectors)
    end
    return nothing
end

function correct_phase(new_vectors::SMatrix, old_vectors::SMatrix)
    n = size(new_vectors, 1)
    vect = SVector{n}(sign(LinearAlgebra.dot(new_vectors[:,i], old_vectors[:,i])) for i=1:n)
    return new_vectors .* vect'
end

function evaluate_adiabatic_derivative!(calc::AbstractDiabaticCalculator, r)
    calc.stats[:adiabatic_derivative] += 1
    U = get_eigen(calc, r).vectors
    diabatic_derivative = get_derivative(calc, r)
    for I in eachindex(diabatic_derivative)
        calc.adiabatic_derivative[I] = U' * diabatic_derivative[I] * U
    end
end

function evaluate_centroid_adiabatic_derivative!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:centroid_adiabatic_derivative] += 1
    centroid_derivative = get_centroid_derivative(calc, r)
    centroid_eigen = get_centroid_eigen(calc, r)
    for I in eachindex(centroid_derivative)
        calc.centroid_adiabatic_derivative[I] = centroid_eigen.vectors' * centroid_derivative[I] * centroid_eigen.vectors
    end
end

function evaluate_adiabatic_derivative!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:adiabatic_derivative] += 1
    derivative = get_derivative(calc, r)
    eigen = get_eigen(calc, r)
    for i in axes(derivative, 3) # Beads
        for j in axes(derivative, 2) # Atoms
            for k in axes(derivative, 1) # DoFs
                calc.adiabatic_derivative[k,j,i] = eigen[i].vectors' * derivative[k,j,i] * eigen[i].vectors
            end
        end
    end
end

function evaluate_nonadiabatic_coupling!(calc::AbstractDiabaticCalculator, r)
    calc.stats[:nonadiabatic_coupling] += 1
    adiabatic_derivative = get_adiabatic_derivative(calc, r)
    eigen = get_eigen(calc, r)
    for I in eachindex(calc.adiabatic_derivative)
        calc.nonadiabatic_coupling[I] = evaluate_nonadiabatic_coupling(adiabatic_derivative[I], eigen.values)
    end
end

function evaluate_centroid_nonadiabatic_coupling!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:centroid_nonadiabatic_coupling] += 1
    centroid_adiabatic_derivative = get_centroid_adiabatic_derivative(calc, r)
    centroid_eigen = get_centroid_eigen(calc, r)
    for I in eachindex(centroid_adiabatic_derivative)
        calc.centroid_nonadiabatic_coupling[I] = evaluate_nonadiabatic_coupling(centroid_adiabatic_derivative[I], centroid_eigen.values)
    end
end

function evaluate_nonadiabatic_coupling!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:nonadiabatic_coupling] += 1
    adiabatic_derivative = get_adiabatic_derivative(calc, r)
    eigen = get_eigen(calc, r)
    for i in eachindex(eigen) # Beads
        for I in CartesianIndices(size(adiabatic_derivative)[1:2])
            calc.nonadiabatic_coupling[I,i] = evaluate_nonadiabatic_coupling(adiabatic_derivative[I,i], eigen[i].values)
        end
    end
end

"""
# References

- HammesSchifferTully_JChemPhys_101_4657_1994 Eq. (32)
- SubotnikBellonzi_AnnuRevPhyschem_67_387_2016, section 2.3
"""
function evaluate_nonadiabatic_coupling(adiabatic_derivative::SMatrix, eigenvalues::SVector)
    n = length(eigenvalues)
    SMatrix{n,n}(
        (i != j ? adiabatic_derivative[j,i] / (eigenvalues[i] - eigenvalues[j]) : 0
        for j=1:n, i=1:n))
end

include("friction.jl")
include("large_diabatic.jl")

"""
Evaluates all electronic properties for the current position `r`.
# Properties evaluated:
- Diabatic potential
- Diabatic derivative
- Eigenvalues and eigenvectors
- Adiabatic derivative
- Nonadiabatic coupling

This should no longer be used, instead access the quantities directly with `get_quantity(calc, r)`.
"""
function update_electronics!(calculator::AbstractDiabaticCalculator, r::AbstractArray)
    get_nonadiabatic_coupling(calculator, r)
    return nothing
end

function update_electronics!(calculator::RingPolymerDiabaticCalculator, r::AbstractArray{T,3}) where {T}
    get_nonadiabatic_coupling(calculator, r)
    update_centroid_electronics!(calculator, r)
    return nothing
end

function update_centroid_electronics!(calculator::RingPolymerDiabaticCalculator, r::AbstractArray{T,3}) where {T}
    get_centroid_nonadiabatic_coupling(calculator, r)
    return nothing
end

end # module


using TimetableSolver
using JuMP
using ConstraintSolver
using DataStructures: OrderedDict

# TODO: add return type for functions

const CS = ConstraintSolver

"""
    VarRefOrExpr

Represents either a model variable reference or expression reference.

Alias for union of VariableRef and AffExpr.
"""
VarRefOrExpr = Union{VariableRef,AffExpr}

#=
Each variable is a string of the form:
    `{div}_{row}_{period}_{type}`,
where div is the str representation of division,
row is the number of row, period is the number of period,
and type is either 'subject' or 'teacher'.

Each domain value represents a subject or teacher, so its value
should be a str of the form: `{subject}` or `{teacher}`.

However, all the strings are hidden and ints are used for solving directly.
Mappings are generated and stored in VariableData struct.
=#

"""
    VariableData

Immutable struct to store data related to divisions, variables and values,
and their mappings from ints to strings, because the solver only works on ints.

Fields:
- `divvarval_maps::OrderedDict{Symbol,OrderedDict}`: Keys are a symbol `:subject` or `:teacher`.
  Values are map of ints (div) to 2-tuple of vecs of ints (var, val).
- `divvarval_mapstrs::OrderedDict{Symbol,OrderedDict}`: Same as `div_var_val_map` but with strings instead of ints.

- `alldivs::Vector{Int}`: Vector of all divs as ints.
- `allvars::Vector{Int}`: Vector of all vars as ints. (no `:subject` `:teacher` bifurcation)
- `allvals::Vector{Int}`: Vector of all vals as ints. (no `:subject` `:teacher` bifurcation)

- `divmap::OrderedDict{Int,String}`: Map from div ints to strings.
- `varmap::OrderedDict{Int,String}`: Map from variable ints to strings.
- `valmap::OrderedDict{Int,String}`: Map from values ints to strings.
- `inversedivmap::OrderedDict{String,Int}`: Inverted `divmap`.
- `inversevarmap::OrderedDict{String,Int}`: Inverted `varmap`.
- `inversevalmap::OrderedDict{String,Int}`: Inverted `valmap`.

- `schedule::Schedule`: Schedule type instance.

"""
struct VariableData
    divvarval_maps::OrderedDict{Symbol,OrderedDict}
    divvarval_mapstrs::OrderedDict{Symbol,OrderedDict}

    alldivs::Vector{Int}
    allvars::Vector{Int}
    allvals::Vector{Int}

    divmap::OrderedDict{Int,String}
    varmap::OrderedDict{Int,String}
    valmap::OrderedDict{Int,String}
    inversedivmap::OrderedDict{String,Int}
    inversevarmap::OrderedDict{String,Int}
    inversevalmap::OrderedDict{String,Int}

    schedule::Schedule
end

"""
    VariableData(schedule::Schedule)::VariableData

Return `VariableData` instance with all the fields filled based on the schedule.
"""
function VariableData(schedule::Schedule)::VariableData
    # get divvarval_mapstrs
    divvarval_mapstrs = get_divvarval_mapstrs(schedule)

    # extract str representations of all data
    alldiv_strs = Vector{String}([])
    allvar_strs = OrderedDict{Symbol,Vector{String}}(
        :subject => [],
        :teacher => []
    )
    allval_strs = OrderedDict{Symbol,Vector{String}}(
        :subject => [],
        :teacher => []
    )
    # symbol is `:subject` or `:teacher`
    # symbol_map is div_var_val_map_strs but for either subject or teacher only
    for (symbol, symbolmap) in divvarval_mapstrs
        for (div_str, varval_strs) in symbolmap
            push!(alldiv_strs, div_str)
            for var_str in varval_strs[1]
                # push the var to the proper key in all_var_strs
                push!(allvar_strs[symbol], var_str)
            end
            for val_str in varval_strs[2]
                # push the val to the proper key in all_var_strs
                push!(allval_strs[symbol], val_str)
            end
        end
    end

    # remove duplicate values from allval_strs
    allval_strs[:subject] = unique(allval_strs[:subject])
    allval_strs[:teacher] = unique(allval_strs[:teacher])

    # remove duplicate values from alldiv_strs
    alldiv_strs = unique(alldiv_strs)

    # define int representations of all data
    # and int to str mappings
    alldivs = Vector{Int}(1:length(alldiv_strs))
    divmap::OrderedDict{Int,String} = OrderedDict(
        enumerate(alldiv_strs)
    )
    inversedivmap::OrderedDict{String,Int} = OrderedDict(v => k for (k, v) in divmap)

    allvars = Vector{Int}(1:(length(allvar_strs[:subject]))+length(allvar_strs[:teacher]))
    # nested list comp
    # return var_str for var_str in symbol_var_strs which are in all_var_strs
    varmap::OrderedDict{Int,String} = OrderedDict(
        enumerate([v for (_, var_strs) in allvar_strs for v in var_strs])
    )
    inversevarmap::OrderedDict{String,Int} = OrderedDict(v => k for (k, v) in varmap)

    allvals = Vector{Int}(1:(length(allval_strs[:subject]))+length(allval_strs[:teacher]))
    # nested list comp (same as for `var_map`)
    valmap::OrderedDict{Int,String} = OrderedDict(
        enumerate([v for (_, val_strs) in allval_strs for v in val_strs])
    )
    inversevalmap::OrderedDict{String,Int} = OrderedDict(v => k for (k, v) in valmap)

    # dict to store ordereddict of all vars and vals mapped to divs
    # dict keys are :subject and :teacher
    divvarval_maps::OrderedDict{Symbol,OrderedDict} = OrderedDict(
        :subject => OrderedDict(),
        :teacher => OrderedDict()
    )

    # generate values for div_var_val_map
    # symbol is `:subject` or `:teacher`
    # symbol_map is divvarval_mapstrs but for either subject or teacher only
    for (symbol, symbolmap) in divvarval_mapstrs
        for (i, (div_str, (var_strs, val_strs))) in enumerate(symbolmap)
            # int representing current div
            div::Int = inversedivmap[div_str]

            # ints representing current vars
            vars = Vector{Int}([inversevarmap[v] for v in var_strs])
            vals = Vector{Int}([inversevalmap[v] for v in val_strs])

            # add key value pair to div_var_val_map
            divvarval_maps[symbol][div] = (vars, vals)
        end
    end

    return VariableData(
        divvarval_maps, divvarval_mapstrs,
        alldivs, allvars, allvals,
        divmap, varmap, valmap, inversedivmap, inversevarmap, inversevalmap,
        schedule
    )
end

"""
    get_divvarval_mapstrs(schedule::Schedule)::OrderedDict{Symbol,OrderedDict}

Return `OrderedDict` mapping symbol (either `:subject` or `:teacher`), to an `OrderedDict`.

The first-order values (`OrderedDict`), are mappings from division string, to a 2-tuple of variables and values.
"variables" and "values" are vectors of variable strings and value strings.
"""
function get_divvarval_mapstrs(schedule::Schedule)::OrderedDict{Symbol,OrderedDict}
    # dict to store ordereddict of all vars and vals mapped to divs
    # dict keys are :subject and :teacher
    divvarval_maps::OrderedDict{
        Symbol,OrderedDict{String,Tuple{Vector{String},Vector{String}}}
    } = OrderedDict(
        :subject => OrderedDict(),
        :teacher => OrderedDict()
    )

    # fill divvarval_maps with data
    for (div, tt) in schedule.data
        subjectvals::Vector{String} = []
        teachervals::Vector{String} = []

        # add dom vals for both teachers and subjects
        for (subject, teachers) in tt.subjectteachers
            push!(subjectvals, subject)
            for t in teachers
                push!(teachervals, t)
            end
        end

        subjectvars::Vector{String} = []
        teachervars::Vector{String} = []

        for (i, row) in enumerate(tt.data)
            for j in (1:length(row))
                # add subject variable
                subject = "$(div)_$(i)_$(j)_subject"
                push!(subjectvars, subject)
                # add teacher variable
                teacher = "$(div)_$(i)_$(j)_teacher"
                push!(teachervars, teacher)
            end
        end

        # only use unique vals
        subjectvals = unique(subjectvals)
        teachervals = unique(teachervals)

        # create a key for every div, in teacher and subject
        divvarval_maps[:subject][div] = (subjectvars, subjectvals)
        divvarval_maps[:teacher][div] = (teachervars, teachervals)
    end

    return divvarval_maps
end

# TODO: add kwargs arg to this and other function to pass kwargs to Model()
"""
    get_model(vardata::VariableData, timeout::Float64=Inf, all_solutions::Bool=false)::Tuple{Model,OrderedDict}

Return a JuMP model created using `vardata`, and an OrderedDict of the variables defined in the model.

# Arguments
- `vardata::VariableData`: VariableData object representing the variables and constraints to be added to the model.
- `timeout::Float64=Inf`: Float representing the maximum time to run the model.
- `all_solutions::Bool=false`: Bool representing whether to return all solutions or just the first.

# Notes
- Create a model with the correct variables, and adds constraints to the model.
- The variables defined are anonymous, so they are stored in an OrderedDict mapped to their int representation.
"""
function get_model(vardata::VariableData, timeout::Float64 = Inf, all_solutions::Bool = false
)::Tuple{Model,OrderedDict}
    # create a constraint solver model and set ConstraintSolver as the optimizer
    m = Model(optimizer_with_attributes(CS.Optimizer,
        # logging vector containing :Info and/or :Table or empty
        "logging" => [],

        # whether to return all possible solutions
        "all_solutions" => all_solutions,

        # seed for random operations given for reproducibility
        "seed" => 0,
        # default to :Auto
        # "traverse_strategy" => :BFS,
        # "traverse_strategy"=> :DFS,
        # "traverse_strategy"=> :DBFS,

        # default :Auto
        # "branch_split"=>:Smallest,
        # "branch_split"=>:Biggest,
        # "branch_split" => :InHalf,

        "branch_strategy" => :IMPS, # default

        # simplify alldifferent, etc. constraints
        "simplify" => true, # default

        # max seconds before exiting with status `MOI.TIME_LIMIT`
        "time_limit" => timeout,

        # backtracking options
        "backtrack" => true, # default
        "backtrack_sorting" => true, # default

        # "lp_optimizer" => cbc_optimizer,
        # "lp_optimizer" => glpk_optimizer,
        # "lp_optimizer" => ipopt_optimizer,
    ))

    # define variables
    modelvars = define_variables!(m, vardata)

    # define constraints
    define_constraints!(m, modelvars, vardata)

    return m, modelvars
end

"""
    define_variables!(m::Model, vardata::VariableData)::OrderedDict{Int,VarRefOrExpr}

Define variables in the model `m` using the `vardata` object.
Return an OrderedDict mapping the int representation of the variable to the variable reference/expression reference.

# Notes
- Modify the model `m` in place.
- See [`VarRefOrExpr`](@ref).
"""
function define_variables!(
    m::Model, vardata::VariableData)::OrderedDict{Int,VarRefOrExpr}
    divvarval_maps = vardata.divvarval_maps

    # mapping variable int repr to its model variableref
    modelvars = OrderedDict{Int,VarRefOrExpr}()

    for (symbol, symbolmap) in divvarval_maps
        for (div, (vars, vals)) in symbolmap
            # define a set of vars at a time
            currentmodelvars = @variable(m, [vars], CS.Integers(vals))
            # allocate them to the value of var in modelvars one by one
            for v in vars
                modelvars[v] = currentmodelvars[v]
            end
        end
    end

    return modelvars
end

"""
    define_subjectconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing

Define subject constraints in the model `m` using the `vardata` object and the model variables `modelvars`.

# Notes
- Modify the model `m` in place.
"""
function define_subjectconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing
    schedule = vardata.schedule
    divvarval_maps = vardata.divvarval_maps
    inversedivmap = vardata.inversedivmap
    inversevalmap = vardata.inversevalmap
    for (div_str, tt) in schedule.data
        # define terms
        div = inversedivmap[div_str]
        counts = OrderedDict{Int,Int}(
            inversevalmap[subject_str] => count
            for (subject_str, count) in tt.subjectcounts
        )
        vars = divvarval_maps[:subject][div][1]
        vals = collect(Int, keys(counts))
        # tl;dr (var::Int = 2) -> (var::Binary[] = [0, 1, 0]) index shows the value
        # create a matrix of binary variables:
        # `vars` as indices and `vals` as number of values
        # so instead of setting an int as value of each var
        # each var has `vals` corresponding binary variables
        binvars = @variable(m, [vars, vals], Bin)
        # add constraint to only choose one val
        @constraint(m, [i = vars], sum(binvars[i, :]) == 1)
        # add constraint to only choose that val `counts[val]` times
        @constraint(m, [j = vals], sum(binvars[:, j]) == counts[j])
        # array of expressions that actually store the value of var as int
        exprs = @expression(m, [i = vars], sum(j * binvars[i, j] for j in vals))
        # replace the corresponding vars in modelvars with exprs
        for var in vars
            # delete the var from the model
            delete(m, modelvars[var])
            # set the var to an expr defined above
            modelvars[var] = exprs[var]
        end
    end
end

"""
    define_subjectteacherconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing

Define subject-teacher constraints in the model `m` using the `vardata` object and the model variables `modelvars`.

# Notes
- Modify the model `m` in place.
"""
function define_subjectteacherconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing
    divvarval_maps = vardata.divvarval_maps
    inversevalmap = vardata.inversevalmap
    schedule = vardata.schedule

    # mapping div to corresponding subjectteachers
    subjectteacher_maps = OrderedDict{Int,OrderedDict}()
    for (div_str, tt) in schedule.data
        div = vardata.inversedivmap[div_str]
        # subject teacher mappings as ints
        subjectteachers = OrderedDict{Int,Vector{Int}}()
        # get the subjectteacher value for this timetable
        subjectteacher_strs = tt.subjectteachers
        # convert strs to ints, and add to `subjectteachers`
        for (subject, teachers) in subjectteacher_strs
            subject_int = inversevalmap[subject]
            teacher_ints = [inversevalmap[t] for t in teachers]
            subjectteachers[subject_int] = teacher_ints
        end
        subjectteacher_maps[div] = subjectteachers

    end

    # iterate over every set of subject variables (by div)
    for (div, (subjectvars, _)) in divvarval_maps[:subject]
        # generate pairs of all possible subject teacher combinations
        # and store them in a matrix
        numpairs = sum(length.(values(subjectteacher_maps[div])))
        subjectteacherpairs = Matrix{Int}(undef, numpairs, 2)
        i = 0
        for (subject, teachers) in subjectteacher_maps[div]
            for teacher in teachers
                i += 1
                subjectteacherpairs[i, 1:2] = [subject, teacher]
            end
        end
        # add using CS.TableSet(table)
        teachervars = divvarval_maps[:teacher][div][1]
        for (svar, tvar) in zip(subjectvars, teachervars)
            # model vars corresponding to this pair of subject and teacher vars
            (modelsvar, modeltvar) = modelvars[svar], modelvars[tvar]
            @constraint(m, [modelsvar, modeltvar] in CS.TableSet(subjectteacherpairs))
        end
    end

    return nothing
end

"""
    define_teacherconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing

Define teacher constraints in the model `m` using the `vardata` object and the model variables `modelvars`.

# Notes
- Modify the model `m` in place.
"""
function define_teacherconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing
    # variables with same row,period should have different teachers
    divvarval_maps = vardata.divvarval_maps

    # following block of code creates a list of vars
    # grouped by the same row and period
    i = 0
    while true
        i += 1
        # store var or 0 if var not found
        rowperiodvars = [get(vars, i, 0) for (_, (vars, _)) in divvarval_maps[:teacher]]
        # remove zeroes from rowperiodvars
        filter!(!iszero, rowperiodvars)
        # if no vars were pushed to list,
        # then we have exceeded the number of rows and periods
        if length(rowperiodvars) == 0
            break
        end
        # now get model vars from these vars
        rowperiodmodelvars = [modelvars[x] for x in rowperiodvars]
        # add alldifferent constraint
        @constraint(m, rowperiodmodelvars in CS.AllDifferent())
    end

    return nothing
end

"""
    define_constraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing

Define all constraints in the model `m` using the `vardata` object and the model variables `modelvars`.

# Notes
- Modify the model `m` in place.
- Call [`define_subjectconstraints!`](@ref), [`define_subjectteacherconstraints!`](@ref)
  and [`define_teacherconstraints!`](@ref) internally, in the respective order.
"""
function define_constraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing
    define_subjectconstraints!(m, modelvars, vardata)

    define_subjectteacherconstraints!(m, modelvars, vardata)

    define_teacherconstraints!(m, modelvars, vardata)

    return nothing
end

"""
    get_solution(m::Model, modelvars::OrderedDict, vardata::VariableData)::Tuple{OrderedDict{Int,Int},MOI.TerminationStatusCode}

Return a solution and status code after solving the given model `m` and using model variables `modelvars` and variable data `vardata` to access the solution values.

# Notes
- Model is optimally solved only if the status code is [`MOI.OPTIMAL`](https://jump.dev/JuMP.jl/stable/moi/reference/models/#MathOptInterface.TerminationStatusCode).
- Solution is a mapping from variable to its value (both as ints).
"""
function get_solution(
    m::Model, modelvars::OrderedDict, vardata::VariableData)::Tuple{OrderedDict{Int,Int},MOI.TerminationStatusCode}
    # solve model
    optimize!(m)

    # convert model values to var=>sol mapping
    solution = OrderedDict{Int,Int}()

    # model status
    status = JuMP.termination_status(m)

    if status != MOI.OPTIMAL
        error("Model unsolved!\nstatus = $status")
    end

    for (symbol, symbolmap) in vardata.divvarval_maps
        for (div, (vars, vals)) in symbolmap
            for var in vars
                # get the corresponding var defined in the model
                mvar = modelvars[var]
                # get its value
                mval = Int(JuMP.value(mvar))
                # add var val mapping to solution
                solution[var] = mval
            end
        end
    end

    # return solution
    return solution, status
end

"""
    convertsolution(rawsolution::OrderedDict{Int,Int}, vardata::VariableData)::OrderedDict{String,String}

Return a solution in the form of a mapping from variable to its value (both as strings).

# Notes
- Convert an `Int=>Int` solution to a `String=>String` one using the `inverse*` mappings defined in `vardata`.
"""
function convertsolution(
    rawsolution::OrderedDict{Int,Int}, vardata::VariableData)::OrderedDict{String,String}
    solution = OrderedDict{String,String}()
    # iterate over every var and val in solution
    for (var, val) in rawsolution
        # get var val strs using mappings
        var_str = vardata.varmap[var]
        val_str = vardata.valmap[val]
        # add it to solution
        solution[var_str] = val_str
    end
    return solution
end

"""
    applysolution!(schedule::Schedule, solution::OrderedDict{String,String})::Nothing

Apply a solution (of strings) to the given schedule.

# Notes
- Modify the schedule in place.
"""
function applysolution!(schedule::Schedule, solution::OrderedDict{String,String})::Nothing
    # iterate over every var and val in solution
    for (var, val) in solution
        # modify schedule with var and val
        modify!(schedule, var, val)
    end
end

"""
    solve!(schedule::Schedule, timeout::Float64 = Inf, all_solutions::Bool = false)::Tuple{OrderedDict{String,String},MOI.TerminationStatusCode}

Solve a given schedule in place and return a `String=>String` solution and status code.

# Arguments
- `schedule::Schedule`: Schedule object representing the problem.
- `timeout::Float64=Inf`: Float representing the maximum time to run the model.
- `all_solutions::Bool=false`: Bool representing whether to return all solutions or just the first.

# Notes
- Schedule is optimally solved only if the status code is [`MOI.OPTIMAL`](https://jump.dev/JuMP.jl/stable/moi/reference/models/#MathOptInterface.TerminationStatusCode).
"""
function solve!(
    schedule::Schedule, timeout::Float64 = Inf, all_solutions::Bool = false
)::Tuple{OrderedDict{String,String},MOI.TerminationStatusCode}
    # get vardata
    vardata = VariableData(schedule)

    # get model and model vars
    m, modelvars = get_model(vardata, timeout, all_solutions)

    # get solution
    rawsolution, status = get_solution(m, modelvars, vardata)

    # convert solution from ints to strs
    solution = convertsolution(rawsolution, vardata)

    # apply solution
    applysolution!(schedule, solution)

    # return solution
    return solution, status
end


# Atomic weights taken from:

# Atomic weights of the elements 2013 (IUPAC Technical  Report)
# http://doi.org/10.1515/pac-2015-0305

# The "Conventional" weights are used where a range is provided, shown in comments.

# hydrogen 1
@mol H 1.008 # [1.0078, 1.0082]
# helium 2
@mol He 4.0026
# lithium 3
@mol Li 6.94 # [6.938, 6.997]
# beryllium 4
@mol Be 9.0122
# boron 5
@mol B 10.81 # [10.806, 10.821]
# carbon 6
@mol C 12.011 # [12.009, 12.012]
# nitrogen 7
@mol N 14.007 # [14.006, 14.008]
# oxygen 8
@mol O 15.999 # [15.999, 16.000]
# fluorine 9
@mol F 18.998
# neon 10
@mol Ne 20.180
# sodium 11
@mol Na 22.990
# magnesium 12
@mol Mg 24.305 # [24.304, 24.307]
# aluminium (aluminum) 13
@mol Al 26.982
# silicon 14
@mol Si 28.085 # [28.084, 28.086]
# phosphorus 15
@mol P 30.974
# sulfur 16
@mol S 32.06 # [32.059, 32.076]
# chlorine 17
@mol Cl 35.45 # [35.446, 35.457]
# argon 18
@mol Ar 39.948
# potassium 19
@mol K 39.098
# calcium 20
@mol Ca 40.078 # (4)
# scandium 21
@mol Sc 44.956
# titanium 22
@mol Ti 47.867
# vanadium 23
@mol V 50.942
# chromium 24
@mol Cr 51.996
# manganese 25
@mol Mn 54.938
# iron 26
@mol Fe 55.845 # (2)
# cobalt 27
@mol Co 58.933
# nickel 28
@mol Ni 58.693
# copper 29
@mol Cu 63.546 # (3)
# zinc 30
@mol Zn 65.38 # (2)
# gallium 31
@mol Ga 69.723
# germanium 32
@mol Ge 72.630 # (8)
# arsenic 33
@mol As 74.922
# selenium 34
@mol Se 78.971 # (8)
# bromine 35
@mol Br 79.904 # [79.901, 79.907]
# krypton 36
@mol Kr 83.798 # (2)
# rubidium 37
@mol Rb 85.468
# strontium 38
@mol Sr 87.62
# yttrium 39
@mol Y 88.906
# zirconium 40
@mol Zr 91.224 # (2)
# niobium 41
@mol Nb 92.906
# molybdenum 42
@mol Mo 95.95
# technetium* 43
# @mol Tc –
# ruthenium 44
@mol Ru 101.07 # (2)
# rhodium 45
@mol Rh 102.91
# palladium 46
@mol Pd 106.42
# silver 47
@mol Ag 107.87
# cadmium 48
@mol Cd 112.41
# indium 49
@mol In 114.82
# tin 50
@mol Sn 118.71
# antimony 51
@mol Sb 121.76
# tellurium 52
@mol Te 127.60 # (3)
# iodine 53
@mol I 126.90
# xenon 54
@mol Xe 131.29
# caesium (cesium) 55
@mol Cs 132.91
# barium 56
@mol Ba 137.33
# lanthanum 57
@mol La 138.91
# cerium 58
@mol Ce 140.12
# praseodymium 59
@mol Pr 140.91
# neodymium 60
@mol Nd 144.24
# promethium* 61
# @mol Pm –
# samarium 62
@mol Sm 150.36 # (2)
# europium 63
@mol Eu 151.96
# gadolinium 64
@mol Gd 157.25 # (3)
# terbium 65
@mol Tb 158.93
# dysprosium 66
@mol Dy 162.50
# holmium 67
@mol Ho 164.93
# erbium 68
@mol Er 167.26
# thulium 69
@mol Tm 168.93
# ytterbium 70
@mol Yb 173.05
# lutetium 71
@mol Lu 174.97
# hafnium 72
@mol Hf 178.49 # (2)
# tantalum 73
@mol Ta 180.95
# tungsten 74
@mol W 183.84
# rhenium 75
@mol Re 186.21
# osmium 76
@mol Os 190.23 # (3)
# iridium 77
@mol Ir 192.22
# platinum 78
@mol Pt 195.08
# gold 79
@mol Au 196.97
# mercury 80
@mol Hg 200.59
# thallium 81
@mol Tl 204.38 # [204.38, 204.39]
# lead 82
@mol Pb 207.2
# bismuth* 83
@mol Bi 208.98


export segyhdr
# ============================================================================
# Utility functions not for export
function ibmfloat(x::UInt32)
  local fra = ntoh(x)
  local sgn = UInt8((fra & 0x80000000)>>31)
  fra <<= 1
  local exp = Int16(fra >> 25) - Int16(64)
  fra <<= 7
  y = (sgn == 0x00 ? 1.0 : -1.0) * 16.0^exp * signed(fra >> 8)/16777216
  return y
end
#=
  This is actual IBM hexadecimal float, and correctly parses the examples from
  https://en.wikipedia.org/wiki/IBM_hexadecimal_floating_point ; see tests.

  The version in JuliaSeis is wrong because IBM float has a range too wide for
  IEEE single float in Julia.

  The description in the SEG Y manual does something strange with the last bit.

  My last line is computationally expensive; working out the radix shift and
  using >> would be much better.
=#

function trid(i::Int16, fs::Float64, fc::Float64)
  S = ["DH", "HZ", "H1", "H2", "HZ", "HT", "HR"]
  return string(getbandcode(fs, fc=fc), S[i])
end

function mk_lc(lv::Int32)
  L1 = UInt8(div(lv, Int32(36)))
  L2 = UInt8(rem(lv, Int32(36)))
  c1 = 0x30 + L1 + (L1 > 0x09 ? 0x07 : 0x00)
  c2 = 0x30 + L2 + (L2 > 0x09 ? 0x07 : 0x00)
  return String([c1, c2])
end

function auto_coords(xy::Array{Int32, 1}, sc::Array{Int16, 1})
  xy == Int32[0,0] && return (0.0, 0.0)
  lon = Float64(xy[1])
  lat = Float64(xy[2])
  coord_scale = sc[1]
  coord_units = sc[2]
  (coord_scale < 0) && (coord_scale = -1 / coord_scale)
  lat *= coord_scale
  lon *= coord_scale

  #=
  Conversion formula for coord_units == 1 "borrowed" from PASSSOFT segy2sac.
  Assumes (possibly wrong) x-y coordinate origin & spherical Earth.
  coord_units == 3 and coord_units == 4 aren't supported; never seen them.
  =#
  if coord_units == 1
      iflg = lon < 0 ? -1 : 1
      c = 111194.6976
      D = sqrt(lat^2 + lon^2) / c
      lat /= c
      d = cosd(D)
      lon = iflg*acosd(d / cosd(lat))
  else
    lat /= 3600.0
    lon /= 3600.0
  end
  return lat, lon
end

function do_trace(f::IO,
                  passcal::Bool,
                  full::Bool,
                  ll::UInt8,
                  swap::Bool,
                  fh::Array{Int16,1}
                  )

  # First part of trace header is quite standard
  buf = BUF.buf
  ints = BUF.int32_buf
  shorts = BUF.int16_buf
  lat = 0.0
  lon = 0.0
  checkbuf_8!(buf, max(180, length(buf)))
  fast_readbytes!(f, buf, 180)

  intbuf = reinterpret(Int32, buf)
  copyto!(ints, 1, intbuf, 1, 7)
  copyto!(ints, 8, intbuf, 10, 8)
  copyto!(ints, 16, intbuf, 19, 4)

  shortbuf = reinterpret(Int16, buf)
  copyto!(shorts, 1, shortbuf, passcal ? 1 : 15, 4) # shorts[1:4]
  copyto!(shorts, 5, shortbuf, 35, 2)               # shorts[5:6]
  copyto!(shorts, 7, shortbuf, 45, 46)              # shorts[7:52]

  if passcal
    fast_readbytes!(f, buf, 40)

    scale_fac     = fastread(f, Float32)
    inst_no       = fastread(f, UInt16)
    shorts[62]    = fastread(f, Int16)
    ints[21]      = fastread(f, Int32)
    ints[22]      = fastread(f, Int32)
    ints[23]      = fastread(f, Int32)

    setindex!(shorts, shortbuf[10], 53)
    copyto!(shorts, 54, shortbuf, 13, 8)
    setindex!(ints, intbuf[6], 20)

    if swap
      shorts     .= ntoh.(shorts)
      ints       .= ntoh.(ints)
      scale_fac   = bswap(scale_fac)
      inst_no     = bswap(inst_no)
    end

    chars = buf[1:18]
    dt    = getindex(ints,20)
    n     = getindex(shorts, 20)
    fmt   = getindex(shorts, 54)
    nx    = (n == typemax(Int16) ? getindex(ints,21) : Int32(n))
    T     = (fmt == one(Int16) ? Int32 : Int16)
    nb    = checkbuf!(buf, nx, T)

    fast_readbytes!(f, buf, nb)

    # trace processing
    y = reinterpret(T, buf)
    x = Array{Float32,1}(undef, nx)
    if swap
      for i = 1:nx
        x[i] = bswap(y[i])
      end
    else
      copyto!(x, 1, y, 1, nx)
      # faster than reprocessing with fillx_ for long files
    end

    z = getindex(shorts, 5)
    δ = getindex(shorts, 21)
    fs        = sμ / Float64(δ == one(Int16) ? dt : δ)
    gain      = 1.0 / (Float64(scale_fac)  * 10.0^(-1.0*Float64(shorts[24]) / 10.0) / Float64(shorts[23]))
    lat, lon  = auto_coords(ints[18:19], shorts[6:7])
    if abs(lat) < 0.1 && abs(lon) < 0.1
      lat *= 3600.0
      lon *= 3600.0
    end
    el        = Float64(ints[9]) * Float64(abs(z))^(z < zero(Int16) ? -1.0 : 1.0)

    # Create ID
    id_arr = zeros(UInt8,8)
    i = one(Int8)
    o = one(Int8)
    while i < Int8(18)
      if chars[i] == 0x00
        chars[i] = 0x20
      end
      i = i+o
    end
    fill_id!(id_arr, chars, one(Int16), Int16(6), Int16(2), Int16(6))
    id_arr[1] = 0x2e
    id_arr[8] = 0x2e
    deleteat!(id_arr, id_arr.==0x00)

    # Channel string is tedious; use of one-char channel names like "Z"
    inds = Int64[]
    for i in 15:18
      if chars[i] > 0x20
        push!(inds, i)
      end
    end
    ch_arr = isempty(inds) ? UInt8[] : chars[inds]
    nc = length(ch_arr)
    if nc == 1
      c = uppercase(Char(ch_arr[1]))
      if c in ('Z','N','E')
        cha = string(getbandcode(fs), 'H', c)
      else
        cha = "YYY"
      end
    elseif nc == 0
      cha = "YYY"
    else
      cha = String(ch_arr[1:min(3, nc)])
    end
    id = String(id_arr) * cha

  else
    (dt, nx, fmt) = fh
    fast_readbytes!(f, buf, 38)
    ints[26]      = fastread(f, Int32)
    shorts[60]    = fastread(f, Int16)
    ints[27]      = fastread(f, Int32)
    shorts[61]    = fastread(f, Int16)
    shorts[62]    = fastread(f, Int16)
    ints[28]      = fastread(f, Int32)
    ints[29]      = fastread(f, Int32)

    copyto!(ints, 20, intbuf, 1, 5)
    copyto!(shorts, 53, shortbuf, 11, 2)          # shorts[53:54]
    ints[25] = getindex(intbuf, 7)
    copyto!(shorts, 55, shortbuf, 15, 5)          # shorts[55:59]

    T = getindex(segy_ftypes, fmt)
    if (T == Any)
      close(f)
      error(string("Trace data code = ", fmt, "  unsupported!"))
    end
    nb = checkbuf!(buf, nx, T)
    fast_readbytes!(f, buf, nb)

    # trace processing
    x = Array{Float32,1}(undef, nx)
    if T == Int16
      fillx_i16_be!(x, buf, nx, 0)
    elseif T == Int32
      fillx_i32_be!(x, buf, nx, 0)
    elseif T == Int8
      fillx_i8!(x, buf, nx, 0)
    elseif T == Float32
      x .= bswap.(reinterpret(Float32, buf))[1:nx]
    elseif T == UInt32
      y = Array{UInt32,1}(undef, nx)
      fillx_u32_le!(y, buf, nx, 0) # _le because ibmfloat bswaps
      x = ibmfloat.(y)
    end

    if swap
    shorts  .= ntoh.(shorts)
    ints    .= ntoh.(ints)
    end
    z       = getindex(shorts, 5)
    δ       = getindex(shorts, 21)
    fs      = sμ / Float64(δ)

    # not sure about meaning of "dB" in gain constants
    gain    = Float64(ints[25]) * 10.0^(shorts[55] + (shorts[23] + shorts[24])/10.0)
    el      = Float64(ints[9]) * Float64(abs(z))^(z<Int16(0) ? -1.0 : 1.0)

    # Create ID
    sta = string(reinterpret(UInt16, shorts[57]))
    lc = ll > 0x00 ? mk_lc(ints[ll]) : ""
    cha = Int16(10) < shorts[1] < Int16(18) ? trid(shorts[1]-Int16(10), fs, 1.0) : "YYY"
    id = string(".", sta[1:min(lastindex(sta),5)], ".", lc, ".", cha)

  end
  if full == true
    shorts_k = String["trace_id_code", "n_summed_z", "n_summed_h", "data_use",
      "z_sc", "h_sc", "h_units_code", "v_weather",
      "v_subweather", "t_src_uphole", "t_rec_uphole", "src_static_cor",
      "rec_static_cor", "total_static", "t_lag_a", "t_lag_b",
      "t_delay", "t_mute_st", "t_mute_en", "nx",
      "delta", "gain_type", "gain_const", "init_gain",
      "correlated", "sweep_st", "sweep_en", "sweep_len",
      "sweep_type", "sweep_tap_st", "sweep_tap_en", "tap_type",
      "f_alias", "slope_alias", "f_notch", "slope_notch",
      "f_low_cut", "f_high_cut", "slope_low_cut", "slope_high_cut",
      "year", "day", "hour", "minute",
      "second", "time_code", "trace_wt_fac", "geophone_roll_p1",
      "geophone_first_tr", "geophone_last_tr", "gap_size", "overtravel"]
    append!(shorts_k, passcal ?
        String["total_static_hi", "data_form", "ms", "trigyear",
        "trigday", "trighour", "trigminute", "trigsecond",
        "trigms", "not_to_be_used"] :
        String["shot_scalar", "trace_units_code", "trans_exp", "trans_units_code",
        "device_id", "trace_time_sc", "src_type_code", "src_energy_dir",
        "src_exp", "src_units_code"])

    ints_k = String["trace_seq_line", "trace_seq_file", "rec_no", "channel_no",
      "energy_src_pt", "ensemble_no", "trace_in_ensemble", "src-rec_dist",
      "rec_ele", "src_ele", "src_dep", "rec_datum_ele",
      "src_datum_ele", "src_water_dep", "rec_water_dep", "src_x",
      "src_y", "rec_x", "rec_y"]
    append!(ints_k, passcal ? String["samp_rate", "num_samps", "max", "min"] :
      String["cdp_x", "cdp_y", "inline_3d", "crossline_3d", "shot_point",
      "trans_mant", "unassigned_1", "unassigned_2"] )
    misc = Dict{String,Any}()
    [misc[shorts_k[i]] = shorts[i] for i in 1:length(shorts_k)]
    [misc[ints_k[i]] = ints[i] for i in 1:length(ints_k)]
    misc["scale_fac"] = gain
    if passcal
      sta = String(chars[1:6])
      misc["inst_no"] = inst_no
      misc["sensor_serial"] = String(chars[7:14])
    end
    misc["station_name"] = sta
    misc["channel_name"] = cha
  end

  # Trace info
  ts = mktime(shorts[41], shorts[42], shorts[43], shorts[45], shorts[45], zero(Int16)) +
        shorts[53] + 1000*sum(shorts[15:17])
  loc = GeoLoc()
  loc.lat = lat
  loc.lon = lon
  loc.el = el

  C = SeisChannel()
  setfield!(C, :name, id)
  setfield!(C, :id, id)
  setfield!(C, :loc, loc)
  setfield!(C, :gain, gain)
  setfield!(C, :fs, fs)
  mk_t!(C, length(x), ts)
  setfield!(C, :x, x)
  if passcal == false
    setfield!(C, :units, get(segy_units, Int16(shorts[56]), ""))
  end
  if full == true
    setfield!(C, :misc, misc)
  end
  return C
end

function read_segy_file!( S::GphysData,
                          fname::String,
                          ll::UInt8,
                          passcal::Bool,
                          memmap::Bool,
                          full::Bool,
                          swap::Bool,
                          strict::Bool)

  f = memmap ? IOBuffer(Mmap.mmap(fname)) : open(fname, "r")
  trace_fh = Array{Int16, 1}(undef, 3)
  if passcal == true
    C = do_trace(f, true, full, ll, swap, trace_fh)
    add_chan!(S, C, strict)
    close(f)
  else
    shorts  = getfield(BUF, :int16_buf)

    # File headers
    filehdr       = fastread(f, 3200)
    jobid         = bswap(fastread(f, Int32))
    lineid        = bswap(fastread(f, Int32))
    reelid        = bswap(fastread(f, Int32))
    fast_readbytes!(f, BUF.buf, 48)
    fillx_i16_be!(shorts, BUF.buf, 24, 0)
    fastskip(f, 240)

    # Some early sample files had these in little-endian byte order
    for i = 25:27
      shorts[i] = bswap(fastread(f, Int16))
    end
    fastskip(f, 94)

    # Process file header
    nh = max(zero(Int16), getindex(shorts, 27))
    if full == false
      fastskip(f, 3200*nh)
    else
      exthdr = Array{String,1}(undef, nh)
      [exthdr[i] = fastread(f, 3200) for i in 1:nh]
      fhd = Dict{String,Any}(
              zip(String["ntr", "naux", "filedt", "origdt", "filenx",
                         "orignx", "fmt", "cdpfold", "trasort", "vsum",
                         "swst", "swen0", "swlen", "swtyp", "tapnum",
                         "swtapst", "swtapen", "taptyp", "corrtra", "bgainrec",
                         "amprec", "msys", "zupdn", "vibpol", "segyver",
                         "isfixed", "n_exthdr"], shorts[1:27])
                            )
      fhd["jobid"]  = jobid
      fhd["lineid"] = lineid
      fhd["reelid"] = reelid
      fhd["filehdr"] = filehdr
      fhd["exthdr"] = exthdr
    end

    trace_fh[1] = getindex(shorts,3)
    trace_fh[2] = getindex(shorts,5)
    trace_fh[3] = getindex(shorts,7)

    # Channel headers
    nt = shorts[1]
    for i = 1:nt
      # "swap" is always true for valid SEG Y data
      C = do_trace(f, false, full, ll, true, trace_fh)
      j = add_chan!(S, C, strict)
      if full == true
        merge!(S.misc[j], fhd)
      end
    end
    close(f)
  end
  resize!(BUF.buf, 65535)
  return S
end

# ============================================================================

"""
    segyhdr(f[; passcal=false, ll=LL, swap=false])

Print formatted, sorted SEG-Y headers of file `f` to stdout. Use keyword
`passcal=true` for PASSCAL/NMT modified SEG Y; use `swap=true` for big-endian
PASSCAL. See SeisIO `read_data` documentation for `ll` codes.
"""
function segyhdr(fname::String; ll::UInt8=0x00, passcal::Bool=false, swap::Bool=false)
  if passcal
    seis = read_data("passcal", fname::String, full=true, ll=ll, swap=swap)
  else
    seis = read_data("segy", fname::String, ll=ll, full=true)
  end
  if passcal
    printstyled(stdout, @sprintf("%20s: %s\n", "PASSCAL SEG-Y FILE", realpath(fname)), color=:green, bold=true)
    D = getindex(getfield(seis, :misc),1)
    for k in sort(collect(keys(D)))
      @printf(stdout, "%20s: %s\n", k, string(get(D, k, "")))
    end
  else
    p = 1; w = 32; W = displaysize(stdout)[2]-2
    S = fill("", length(seis.misc[1])+1)
    printstyled(stdout, @sprintf("%20s: %s\n", "SEG-Y FILE", realpath(fname)), color=:green, bold=true)
    for i = 1:seis.n
      if p > 1
        s = @sprintf("%20i/%i", i, seis.n)
      else
        s = @sprintf(" Trace # %11i/%i", i, seis.n)
      end
      S[1] *= s
      S[1] *= " "^(w-length(s))
      D = getindex(getfield(seis, :misc),1)
      for (j,k) in enumerate(sort(collect(keys(D))))
        if k == "exthdr" || k == "filehdr"
          val = get(seis.misc[i], k, "")
          if isempty(val)
            s = "(empty)"
          else
            s = length(val) > 8 ? String(val[1:8])*"…" : String(val)
          end
        else
          s = string(get(seis.misc[i], k, ""))
        end
        filler = " "^max(0, 10-length(s))
        S[j+1] *= @sprintf("%20s: %s%s", k, s, filler)
      end
      if (p+2*w > W) || (i == seis.n)
        printstyled(stdout, S[1]*"\n", color=:yellow, bold=true)
        [println(S[j]) for j=2:length(S)]
        println("")
        S = fill("", length(seis.misc[i])+1)
        p = 1
      else
        p += w
      end
    end
    printstyled(stdout, @sprintf("%20s\n","END OF RECORD"), color=208, bold=true)
  end
  return nothing
end