"""
spin21.jl - derivative robustness for the empty problem with simultaneous errors
"""

using Altro
using HDF5
using LinearAlgebra
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const TO = TrajectoryOptimization

WDIR = get(ENV, "ROBUST_QOC_PATH", "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin21"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# Define the optimization.
const CONTROL_COUNT = 1
const PN_STEPS = 2
const MAX_COST_VALUE = 1e12
const ILQR_DJ_TOL = 1e-4

# Define the problem.
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const ASTATE_SIZE_BASE = STATE_COUNT * STATE_SIZE_ISO + 1 * CONTROL_COUNT
const INITIAL_STATE1 = [1., 0, 0, 0]
const INITIAL_STATE2 = [0., 1, 0, 0]
# state indices
const STATE1_IDX = 1:STATE_SIZE_ISO
const STATE2_IDX = STATE1_IDX[end] + 1:STATE1_IDX[end] + STATE_SIZE_ISO
const CONTROLS_IDX = STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT
# const DCONTROLS_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT
const D1STATE1_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + STATE_SIZE_ISO
const D1STATE2_IDX = D1STATE1_IDX[end] + 1:D1STATE1_IDX[end] + STATE_SIZE_ISO
const D2STATE1_IDX = D1STATE2_IDX[end] + 1:D1STATE2_IDX[end] + STATE_SIZE_ISO
const D2STATE2_IDX = D2STATE1_IDX[end] + 1:D2STATE1_IDX[end] + STATE_SIZE_ISO
# control indices
const DCONTROLS_IDX = 1:CONTROL_COUNT

# Specify logging.
const VERBOSE = true
const SAVE = true

# Misc. constants
EMPTY_V = []

# Define the dynamics.
struct Model{DO} <: AbstractModel
    Model(DO::Int64=0) = new{DO}()
end
RobotDynamics.state_dim(::Model{DO}) where DO = (
    ASTATE_SIZE_BASE + DO * STATE_COUNT * STATE_SIZE_ISO
)
RobotDynamics.control_dim(::Model{DO}) where DO = CONTROL_COUNT

function RobotDynamics.dynamics(model::Model{DO}, astate::StaticVector,
                                acontrols::StaticVector, time::Real) where DO
    negi_h = (
        astate[CONTROLS_IDX][1] * NEGI_H1_ISO
    )
    delta_state1 = negi_h * astate[STATE1_IDX]
    delta_state2 = negi_h * astate[STATE2_IDX]
    delta_control = acontrols[DCONTROLS_IDX]

    if DO == 1
        delta_d11state1 = NEGI_H0_ISO * astate[STATE1_IDX] + negi_h * astate[D1STATE1_IDX]
        delta_d11state2 = NEGI_H0_ISO * astate[STATE2_IDX] + negi_h * astate[D1STATE2_IDX]
        delta_d12state1 = NEGI_H0_ISO * astate[STATE1_IDX] + negi_h * astate[D1STATE1_IDX]
        delta_d12state2 = NEGI_H0_ISO * astate[STATE2_IDX] + negi_h * astate[D1STATE2_IDX]
        delta_astate = [
            delta_state1;
            delta_state2;
            delta_control;
            delta_d11state1;
            delta_d11state2;
            delta_d12state1;
            delta_d12state2;
        ]
    elseif DO == 2
        delta_d1state1 = NEGI_H0_ISO * astate[STATE1_IDX] + negi_h * astate[D1STATE1_IDX]
        delta_d1state2 = NEGI_H0_ISO * astate[STATE2_IDX] + negi_h * astate[D1STATE2_IDX]
        delta_d2state1 = 2 * NEGI_H0_ISO * astate[D1STATE1_IDX] + negi_h * astate[D2STATE1_IDX]
        delta_d2state2 = 2 * NEGI_H0_ISO * astate[D1STATE2_IDX] + negi_h * astate[D2STATE2_IDX]
        delta_astate = [
            delta_state1;
            delta_state2;
            delta_control;
            delta_d1state1;
            delta_d1state2;
            delta_d2state1;
            delta_d2state2;
        ]
    else
        delta_astate = [
            delta_state1;
            delta_state2;
            delta_control;
        ]
    end

    return delta_astate
end

function run_traj(;gate_type=xpi, evolution_time=10., solver_type=alilqr,
                  postsample=false, initial_save_file_path=nothing,
                  initial_save_type=jl, sqrtbp=false, derivative_order=0,
                  integrator_type=rk6, max_penalty=1e11, qs=nothing,
                  smoke_test=false, dt=2.5e-3, al_kickout_tol=1e-9,
                  constraint_tol=1e-10)
    model = Model(derivative_order)
    n = state_dim(model)
    m = control_dim(model)
    t0 = 0.
    x0 = SVector{n}([
        INITIAL_STATE1;
        INITIAL_STATE2;
        zeros(1 * CONTROL_COUNT);
        zeros(derivative_order * STATE_COUNT * STATE_SIZE_ISO);
    ])
    if gate_type == xpi
        target_state1 = Array(XPI_ISO_1)
        target_state2 = Array(XPI_ISO_2)
    end
    xf = SVector{n}([
        target_state1;
        target_state2;
        zeros(1 * CONTROL_COUNT);
        zeros(derivative_order * STATE_COUNT * STATE_SIZE_ISO);
    ])

    # control amplitude constraint
    x_max = SVector{n}([
        fill(Inf, STATE_COUNT * STATE_SIZE_ISO);
        fill(MAX_CONTROL_NORM_0, 1); # control
        fill(Inf, derivative_order * STATE_COUNT * STATE_SIZE_ISO)
    ])
    x_min = SVector{n}([
        fill(-Inf, STATE_COUNT * STATE_SIZE_ISO);
        fill(-MAX_CONTROL_NORM_0, 1); # control
        fill(-Inf, derivative_order * STATE_COUNT * STATE_SIZE_ISO)
    ])

    N = Int(floor(evolution_time / dt)) + 1
    U0 = [SVector{m}(
        fill(1e-4, CONTROL_COUNT)
    ) for k = 1:N-1]
    X0 = [SVector{n}([
        fill(NaN, n);
    ]) for k = 1:N]
    Z = Traj(X0, U0, dt * ones(N))

    if isnothing(qs)
        qs = ones(6)
    end
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * STATE_SIZE_ISO); # state1, state2
        fill(qs[2], 1); # control
        fill(qs[3], eval(:($derivative_order >= 1 ? $STATE_COUNT * $STATE_SIZE_ISO : 0))); # d1state<1,2>
        fill(qs[4], eval(:($derivative_order >= 2 ? $STATE_COUNT * $STATE_SIZE_ISO : 0))); # d2state<1,2>
        fill(qs[5], eval(:($derivative_order >= 3 ? $STATE_COUNT * $STATE_SIZE_ISO : 0))); # d3state<1,2>
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[6], CONTROL_COUNT);
    ]))
    obj = LQRObjective(Q, R, Qf, xf, N)

    # Must satisfy control amplitude bound.
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # Must reach target state. Must have zero net flux.
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX])

    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, target_astate_constraint, N:N);

    # Instantiate problem and solve.
    prob = Problem{IT_RDI[integrator_type]}(model, obj, constraints, x0, xf, Z, N, t0, evolution_time)
    opts = SolverOptions(verbose=VERBOSE)
    solver = AugmentedLagrangianSolver(prob, opts)
    if solver_type == alilqr
        solver = AugmentedLagrangianSolver(prob, opts)
        solver.solver_uncon.opts.square_root = sqrtbp
        solver.solver_uncon.opts.max_cost_value = MAX_COST_VALUE
        solver.opts.constraint_tolerance = al_kickout_tol
        solver.opts.constraint_tolerance_intermediate = al_kickout_tol
        solver.opts.cost_tolerance_intermediate = ILQR_DJ_TOL
        solver.opts.penalty_max = max_penalty
        if smoke_test
            solver.opts.iterations = 1
            solver.solver_uncon.opts.iterations = 1
        end
    elseif solver_type == altro
        solver = ALTROSolver(prob, opts)
        solver.opts.constraint_tolerance = constraint_tol
        solver.solver_al.solver_uncon.opts.square_root = sqrtbp
        solver.solver_al.solver_uncon.opts.max_cost_value = MAX_COST_VALUE
        solver.solver_al.opts.constraint_tolerance = al_kickout_tol
        solver.solver_al.opts.constraint_tolerance_intermediate = al_kickout_tol
        solver.solver_al.opts.cost_tolerance_intermediate = ILQR_DJ_TOL
        solver.solver_al.opts.penalty_max = max_penalty
        solver.solver_pn.opts.constraint_tolerance = constraint_tol
        solver.solver_pn.opts.n_steps = PN_STEPS
        if smoke_test
            solver.solver_al.opts.iterations = 1
            solver.solver_al.solver_uncon.opts.iterations = 1
            solver.solver_pn.opts.n_steps = 1
        end
    end
    Altro.solve!(solver)

    # Post-process.
    acontrols_raw = controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    Q_raw = Array(Q)
    Q_arr = [Q_raw[i, i] for i in 1:size(Q_raw)[1]]
    Qf_raw = Array(Qf)
    Qf_arr = [Qf_raw[i, i] for i in 1:size(Qf_raw)[1]]
    R_raw = Array(R)
    R_arr = [R_raw[i, i] for i in 1:size(R_raw)[1]]
    cidx_arr = Array(CONTROLS_IDX)
    cmax = TrajectoryOptimization.max_violation(solver)
    cmax_info = TrajectoryOptimization.findmax_violation(get_constraints(solver))
    iterations_ = iterations(solver)

    # Save.
    if SAVE
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            write(save_file, "acontrols", acontrols_arr)
            write(save_file, "controls_idx", cidx_arr)
            write(save_file, "evolution_time", evolution_time)
            write(save_file, "astates", astates_arr)
            write(save_file, "Q", Q_arr)
            write(save_file, "Qf", Qf_arr)
            write(save_file, "R", R_arr)
            write(save_file, "cmax", cmax)
            write(save_file, "cmax_info", cmax_info)
            write(save_file, "dt", dt)
            write(save_file, "derivative_order", derivative_order)
            write(save_file, "solver_type", Integer(solver_type))
            write(save_file, "sqrtbp", Integer(sqrtbp))
            write(save_file, "max_penalty", max_penalty)
            write(save_file, "constraint_tol", constraint_tol)
            write(save_file, "al_kickout_tol", al_kickout_tol)
            write(save_file, "ilqr_dj_tol", ILQR_DJ_TOL)
            write(save_file, "integrator_type", Integer(integrator_type))
            write(save_file, "gate_type", Integer(gate_type))
            write(save_file, "save_type", Integer(jl))
            write(save_file, "iterations", iterations_)
            write(save_file, "max_cost_value", MAX_COST_VALUE)
        end
    end
end


"""
    Hopper <: Environment

    planar hopping robot, based on https://gym.openai.com/envs/Hopper-v2/
"""
struct Hopper end

function hopper(;
    representation=:minimal,
    timestep=0.05,
    gravity=[0.0; 0.0; -9.81],
    friction_coefficient=1.0,
    spring=10.0,
    damper=50.0,
    seed=1,
    contact_foot=true,
    contact_body=true,
    info=nothing,
    vis=Visualizer(),
    name=:robot,
    opts_step=SolverOptions(rtol=3.0e-4, btol=3.0e-4, undercut=1.5),
    opts_grad=SolverOptions(rtol=3.0e-4, btol=3.0e-4, undercut=1.5),
    T=Float64)

    mechanism = get_hopper(
        timestep=timestep,
        gravity=gravity,
        friction_coefficient=friction_coefficient,
        spring=spring,
        damper=damper,
        contact_foot=contact_foot,
        contact_body=contact_body)

    initialize_hopper!(mechanism)

    if representation == :minimal
        nx = minimal_dimension(mechanism)
    elseif representation == :maximal
        nx = maximal_dimension(mechanism)
    end
    nu = 3
    no = nx - 1 # full_state is not being returned by default

    # values taken from Mujoco's model, combining the control range -1, 1 and the motor gears.
    aspace = BoxSpace(nu,
        low=(-Inf * ones(nu)),
        high=(Inf * ones(nu)))
    ospace = BoxSpace(no,
        low=(-Inf * ones(no)),
        high=(Inf * ones(no)))

    rng = MersenneTwister(seed)

    z = get_maximal_state(mechanism)
    x = representation == :minimal ? maximal_to_minimal(mechanism, z) : z

    fx = zeros(nx, nx)
    fu = zeros(nx, nu)

    u_prev = zeros(nu)
    control_mask = [zeros(nu, 3) I(nu)]
    motor_gear = [200, 200, 200.]
    control_scaling = Diagonal(timestep * motor_gear)

    build_robot(mechanism, vis=vis, name=name)

    TYPES = [Hopper, T, typeof(mechanism), typeof(aspace), typeof(ospace), typeof(info)]
    env = Environment{TYPES...}(mechanism, representation, aspace, ospace,
        x, fx, fu,
        u_prev, control_mask' * control_scaling,
        nx, nu, no,
        info,
        [rng], vis,
        opts_step, opts_grad)

    return env
end

function Base.reset(env::Environment{Hopper};
    x=nothing, reset_noise_scale=0.005)

    if x != nothing
        env.state .= x
    else
        # initialize above the ground to make sure that with random initialization we do not violate the ground constraint.
        initialize!(env.mechanism, :hopper,
            z=0.25)
        x0 = get_minimal_state(env.mechanism)
        nx = minimal_dimension(env.mechanism)

        low = -reset_noise_scale
        high = reset_noise_scale
        x = x0 + (high - low) .* rand(env.rng[1], nx) .+ low # we ignored the normal distribution on the velocities
        z = minimal_to_maximal(env.mechanism, x)
        set_maximal_state!(env.mechanism, z)
        if env.representation == :minimal
            env.state .= get_minimal_state(env.mechanism)
        elseif env.representation == :maximal
            env.state .= get_maximal_state(env.mechanism)
        end
        env.input_previous .= 0.0
    end
    return get_observation(env)
end

function get_observation(env::Environment{Hopper};
    full_state=false)

    full_state && (return env.state)

    nx = minimal_dimension(env.mechanism)
    if env.representation == :minimal
        o = env.state
    elseif env.representation == :maximal
        o = maximal_to_minimal(env.mechanism, env.state)
    end

    # clamp velocities and remove x position
    ind = velocity_index(env.mechanism)
    vel = clamp.(o[ind], -10, 10)
    o[ind] .= vel
    o = o[[1; 3:nx]]
    return o
end

function cost(env::Environment{Hopper}, x, u;
        alive_bonus=0.1)

    if env.representation == :minimal
        x_velocity = -x[5]
    else
        i_torso = findfirst(body -> body.name == "torso", collect(env.mechanism.bodies))
        z_torso = x[(i_torso-1)*13 .+ (1:13)]
        x_velocity = -z_torso[4]
    end
    # c = -x_velocity/10 # Gym reward
    c = -x_velocity
    c -= alive_bonus
    c += 1e-4 * u'*u
    return c
end

function is_done(env::Environment{Hopper}, x)
    nx = minimal_dimension(env.mechanism)
    if env.representation == :minimal
        x0 = x
    elseif env.representation == :maximal
        x0 = maximal_to_minimal(env.mechanism, x)
    end
    height = x0[1]
    ang = x0[3]
    done = !(
        all(isfinite.(x0)) &&
        all(abs.(x0[[1;3:nx]]) .< 100) &&
        (height > 0.7) &&
        (abs(ang) < 0.2)
        )
    return done
end


################################################################################
#
#   TYPE DEFINITIONS OF OBJECT CLASSES IN JULIA
#
#   STRUCTURE OF THE FILE:
#
#   1) Definition of UNITCELL OBJECT
#       - Saving
#       - Loading
#
#   2) Definition of LATTICE OBJECT
#       - Saving
#       - Loading
#
#   3) CONVERSION LATTICE --> UNITCELL
#
#   4) OBTAINING GENERAL INFORMATION
#       - Testing a Unitcell object
#       - Print Information (printInfo(..))
#       - obtain the connection strengths of an object
#
#   5) OBTAINING CONNECTIVITY INFORMATION
#       - getConnectivityList(..)
#       - getConnectionList(..)
#
################################################################################

#-----------------------------------------------------------------------------------------------------------------------------
#
#   THE UNIT CELL OBJECT CLASS IN JULIA
#
#   FIELDS:
#   - lattice_vectors   lattice vectors of the underlying bravais lattice
#   - basis             basis sites of the unit cell
#   - connections       the connetions between sites of the unit cells
#   - filename          The filename that the unitcell should be stored into
#
#-----------------------------------------------------------------------------------------------------------------------------
"""
    mutable struct Unitcell

The type that contains information on a unitcell. Fields are

    lattice_vectors :: Array{Array{Float64, 1}, 1}
    basis           :: Array{Array{Float64, 1}, 1}
    connections     :: Array{Array{Any, 1}, 1}
    filename        :: String

Note that connections have the format

    [index_from::Int, index_to::Int, strength::Any, unitcell_offset::Tuple]

New objects can be created with

    Unitcell(filename::String)

to load from a given file with name `filename` or with

    Unitcell()

if the constructed object serves as a placeholder. In this case, the filename
is set to `UNITCELL_DUMMY_FILENAME`.

# Examples

```julia-repl
julia> Unitcell("myunitcellfile.jld")
```
"""
mutable struct Unitcell

    # FIELDS

    # basis vectors of the Bravais lattice
    lattice_vectors::Array{Array{Float64, 1}, 1}

    # basis sites within the unitcell
    basis::Array{Array{Float64, 1}, 1}

    # list of connections of the format
    # [index_from, index_to, strength, unitcell_offset]
    connections::Array{Array{Any, 1}, 1}

    # filename if saved
    filename::String

    # the general constructor
    function Unitcell(lattice_vectors::Array{Array{Float64, 1}, 1}, basis::Array{Array{Float64, 1}, 1}, connections::Array{Array{Any,1},1}, filename::String)
        return new(lattice_vectors, basis, connections, filename)
    end
    # the custom constructor when loading a unitcell from a file
    function Unitcell(filename::String)
        # define an empty unitcell
        uc = new(Array{Float64, 1}[],Array{Float64, 1}[],Array{Any, 1}[],filename)
        # load the uc
        loadUnitcell(uc)
        # return the uc
        return uc
    end
    # the dummy constructor (for lattices that do not need a UC)
    function Unitcell()
        # just initialize everything
        return new(Array{Float64, 1}[],Array{Float64, 1}[],Array{Any, 1}[],UNITCELL_DUMMY_FILENAME)
    end

end

#-----------------------------------------
#
#   METHODS FOR SAVING UNITCELLS
#
#-----------------------------------------

"""
    saveUnitcell(uc::Unitcell [, filename::String])

Save the `Unitcell` object `unitcell` to a file by using the module `JLD`.
A new filename can be given by passing an additional argument.

Note that the filename has to end with `.jld` or the ending will be added automatically.

Folders that lead to the specified file will be created on the way.

# Examples

```julia-repl
julia> saveUnitcell(unitcell)
"filename.jld"

julia> saveUnitcell(unitcell, "testfile.jld")
"testfile.jld"
```
"""
function saveUnitcell(uc::Unitcell, filename::String="NONE")
    # maybe set the new filename
    if filename != "NONE"
        uc.filename = filename
    end
    # make sure that the file ends with .jld
    if uc.filename[end-3:end] != ".jld"
        uc.filename = "$(uc.filename).jld"
    end
    # ensure default path is build
    if contains(uc.filename, "/")
		# get the containing folder
		folder = uc.filename[1:findlast(uc.filename, '/')]
		# build the path to that folder
		mkpath(folder)
	end
    # open the file and write the fields into it
    save(uc.filename,
        "basis",            uc.basis,
        "connections",      uc.connections,
        "lattice vectors",  uc.lattice_vectors,
        "filename",         uc.filename
    )
    # return the filename
    return uc.filename
end

#-----------------------------------------
#
#   METHODS FOR LOADING UNITCELLS
#
#-----------------------------------------

"""
    loadUnitcell(uc::Unitcell [, filename::String])
    loadUnitcell(filename::String)

Loads a `Unitcell` object `unitcell` from its file by using the module `JLD`.
A new filename can be given by passing an additional argument.
Moreover, if there is no `Unitcell` object which should be loaded but one wants to obtain a new
object, one can only pass the filename of the unitcell file.

# Examples

```julia-repl
julia> loadUnitcell(unitcell)
"filename.jld"

julia> loadUnitcell(unitcell, "testfile.jld")
"testfile.jld"

julia> loadUnitcell("testfile.jld")
LatticePhysics.Unitcell[...]
```
"""
function loadUnitcell(uc::Unitcell, filename::String="NONE")
    # maybe set the new filename
    if filename != "NONE"
        uc.filename = filename
    end
    # open the file and load all fields from it
    uc.basis            = load(uc.filename, "basis")
    uc.connections      = load(uc.filename, "connections")
    uc.lattice_vectors  = load(uc.filename, "lattice vectors")
    # return the filename
    return uc.filename
end
function loadUnitcell(filename::String)
    # return the filename
    return Unitcell(filename)
end

# make the type public accessible
export Unitcell
export loadUnitcell, saveUnitcell

#-----------------------------------------------------------------------------------------------------------------------------
#
#   THE LATTICE OBJECT CLASS IN JULIA
#
#   FIELDS:
#   - unitcell              The unit cell of the lattice (if present)
#   - unitcellRepetitions   Number of repetitions for each lattice direction (integer array with length of lattice vectors of UC)
#   - lattice_vectors       all periodic lattice vectors that span the complete lattice (empty list if finite lattice)
#   - positions             all sites of the lattice
#   - connections           the connetions between sites
#   - filename              The filename that the lattice should be stored into
#
#-----------------------------------------------------------------------------------------------------------------------------
"""
    mutable struct Lattice

The type that contains information on a lattice. Fields are

	unitcell            :: Unitcell
	unitcellRepetitions :: Array{Int64, 1}
    lattice_vectors     :: Array{Array{Float64, 1}, 1}
    positions           :: Array{Array{Float64, 1}, 1}
	positions_indices   :: Array{Int64, 1}
    connections         :: Array{Array{Any, 1}, 1}
    filename            :: String

Note that connections have the format

    [index_from::Int, index_to::Int, strength::Any, lattice_offset::Tuple]

where `lattice_offset` only describes the offset along the periodic directions

New objects can be created with

    Lattice(filename::String)

to load from a given file with name `filename` or with

    Lattice()

if the constructed object serves as a placeholder. In this case, the filename
is set to `LATTICE_DUMMY_FILENAME`.

For a complete list of all constructors please consult the documentation.

# Examples

```julia-repl
julia> Lattice("mylatticefile.jld")
```
"""
type Lattice

    # FIELDS

	# The unitcell of the lattice if applicable
    unitcell::Unitcell
	# how often the unitcell is repeated in the elementary directions
    unitcellRepetitions::Array{Int64, 1}

	# The lattice vectors that span the LATTICE in periodic dimensions
    lattice_vectors::Array{Array{Float64, 1}, 1}

	# Array with the positions of all sites within the lattice
    positions::Array{Array{Float64, 1}, 1}
	# array of ints that give the index of the site in the original unit cell
    positions_indices::Array{Int64, 1}

	# Array with connections between sites (same format as in Unitcell)
    connections::Array{Array{Any, 1}, 1}

	# Filename of the jld file
    filename::String

	# general Constructor if everything is given
    function Lattice(
				unitcell::Unitcell,
				unitcellRepetitions::Array{Int64, 1},
				lattice_vectors::Array{Array{Float64, 1}, 1},
				positions::Array{Array{Float64, 1}, 1},
				positions_indices::Array{Int64, 1},
				connections::Array{Array{Any, 1}, 1},
				filename::String
			)
        # just initialize everything and set the position indices to 1
        return new(unitcell, unitcellRepetitions, lattice_vectors, positions, positions_indices, connections, filename)
    end

	# Constructor if no position indices should be given
    function Lattice(
				unitcell::Unitcell,
				unitcellRepetitions::Array{Int64, 1},
				lattice_vectors::Array{Array{Float64, 1}, 1},
				positions::Array{Array{Float64, 1}, 1},
				connections::Array{Array{Any, 1}, 1},
				filename::String
			)
        # just initialize everything and set the position indices to 1
        return new(unitcell, unitcellRepetitions, lattice_vectors, positions,ones(size(positions,1)), connections, filename)
    end

    # Constructor if the lattice is not constructed from a unitcell
    function Lattice(
				lattice_vectors::Array{Array{Float64, 1}, 1},
				positions::Array{Array{Float64, 1}, 1},
				connections::Array{Array{Any, 1}, 1},
				filename::String
			)
        # just initialize everything
        return new(Unitcell(), [], lattice_vectors, positions,ones(size(positions,1)), connections, filename)
    end

    # Constructor when loading from a file
    function Lattice(filename::String)
        # define an empty unitcell
        lattice = new(Unitcell(),Array{Int64, 1}[], Array{Float64, 1}[],Array{Float64, 1}[],Array{Int64, 1}[],Array{Any, 1}[],filename)
        # load the uc
        loadLattice(lattice)
        # return the uc
        return lattice
    end

    # Dummy constructor
    function Lattice()
        # define an empty unitcell
        lattice = new(Unitcell(),Array{Int64, 1}[], Array{Float64, 1}[],Array{Float64, 1}[],Array{Int64, 1}[],Array{Any, 1}[],LATTICE_DUMMY_FILENAME)
        # return the uc
        return lattice
    end

end

#-----------------------------------------
#
#   METHODS FOR SAVING LATTICES
#
#-----------------------------------------

"""
    saveLattice(lattice::Lattice [, filename::String])

Save the `Lattice` object `lattice` to a file by using the module `JLD`.
A new filename can be given by passing an additional argument.

Note that the filename has to end with `.jld` or the ending will be added automatically.

Folders that lead to the specified file will be created on the way.

# Examples

```julia-repl
julia> saveLattice(lattice)
"filename.jld"

julia> saveLattice(lattice, "testfile.jld")
"testfile.jld"
```
"""
function saveLattice(lattice::Lattice, filename::String="NONE")
    # maybe set the new filename
    if filename != "NONE"
        lattice.filename = filename
    end
    # make sure that the file ends with .jld
    if lattice.filename[end-3:end] != ".jld"
        lattice.filename = "$(lattice.filename).jld"
    end
    # ensure default path is build
    if contains(lattice.filename, "/")
		# get the containing folder
		folder = lattice.filename[1:findlast(lattice.filename, '/')]
		# build the path to that folder
		mkpath(folder)
	end
    # open the file and write the fields into it
    save(lattice.filename,
        "unitcell",             lattice.unitcell,
        "unitcell repetitions", lattice.unitcellRepetitions,
        "positions",            lattice.positions,
        "positions indices",    lattice.positions_indices,
        "connections",          lattice.connections,
        "lattice vectors",      lattice.lattice_vectors,
        "filename",             lattice.filename
    )
    # return the filename
    return lattice.filename
end

#-----------------------------------------
#
#   METHODS FOR LOADING LATTICES
#
#-----------------------------------------

"""
    loadLattice(lattice::Lattice [, filename::String])
    loadLattice(filename::String)

Loads a `Lattice` object `lattice` from its file by using the module `JLD`.
A new filename can be given by passing an additional argument.
Moreover, if there is no `Lattice` object which should be loaded but one wants to obtain a new
object, one can only pass the filename of the unitcell file. (Note, that this can also be done
by using the constructor `Lattice(filename)`)

# Examples

```julia-repl
julia> loadLattice(lattice)
"filename.jld"

julia> loadLattice(lattice, "testfile.jld")
"testfile.jld"

julia> loadLattice("testfile.jld")
LatticePhysics.Lattice[...]
```
"""
function loadLattice(lattice::Lattice, filename::String="NONE")
    # maybe set the new filename
    if filename != "NONE"
        lattice.filename = filename
    end
    # open the file and load all fields from it
    lattice.unitcell            = load(lattice.filename, "unitcell")
    lattice.unitcellRepetitions = load(lattice.filename, "unitcell repetitions")
    lattice.positions           = load(lattice.filename, "positions")
    lattice.positions_indices   = load(lattice.filename, "positions indices")
    lattice.connections         = load(lattice.filename, "connections")
    lattice.lattice_vectors     = load(lattice.filename, "lattice vectors")
    # return the filename
    return lattice.filename
end
function loadLattice(filename::String)
    # just return the object that the constructor would return
	return Lattice(filename)
end

# EXPORT the relevant types and methods
export Lattice
export saveLattice, loadLattice

#-----------------------------------------------------------------------------
#
#   METHODS FOR CONVERTING LATTICE --> UNITCELL
#
#-----------------------------------------------------------------------------

"""
	toUnitcell(lattice::Lattice [, filename::String])

Converts a `Lattice` object `lattice` to a `Unitcell` object by keeping
connections and lattice vectors and setting all sites as basis sites.
Optionally, a further argument can specify a new filename for the newly
created unitcell.

# Examples

```julia-repl
julia> toUnitcell(lattice)
LatticePhysics.Unitcell[...]

julia> toUnitcell(lattice, "unitcell_test.jld")
LatticePhysics.Unitcell[...]
```
"""
function toUnitcell(lattice::Lattice, filename::String="AUTOMATIC")
    # maybe set the new filename
    if filename == "AUTOMATIC"
        filename = replace(lattice.filename, ".jld", "_to_unitcell.jld")
    end
	# return a newly built unitcell object
    return Unitcell(
        lattice.lattice_vectors,
        lattice.positions,
        lattice.connections,
		filename
            )
end

# Export the conversion method
export toUnitcell

#-----------------------------------------------------------------------------
#
#   METHODS FOR PRINTING OR OBTAINING GENERAL INFORMATION
#
#-----------------------------------------------------------------------------

"""
    testUnitcell(unitcell::Unitcell, N_lattice_vectors::Int64, N_space_dimension::Int64)

Performs a series of tests on a given `Unitcell` object so see if it is implemented correctly.
Checks include the following:
1) test if the number of `lattice_vectors` is as expected
2) test if the number of spatial dimensions is correct for every basis site
3) test if the filename does not contain `"FOLDER_UNITCELLS"` (as in this case a \$ would be missing somewhere)
4) test if the connectivity is okay (if every connection has a counterpart)

The return value is a `Bool` indicating if the test is passed or not. If the test fails, the
failure reason is printed as well.

# Examples

```julia-repl
julia> testUnitcell(unitcell, 2, 2)
...
```
"""
function testUnitcell(unitcell::Unitcell, N_lattice_vectors::Int64, N_space_dimension::Int64)
    # check if the unitcell has enough lattice vectors
    if size(unitcell.lattice_vectors,1) != N_lattice_vectors
        println("Uncorrect number of lattice vectors ($(size(unitcell.lattice_vectors,1)) instead of $(N_lattice_vectors))")
        return false
    end
    # check if the space dimension for all sites is okay
    for p in unitcell.basis
        if size(p,1) != N_space_dimension
            println("Site ($(p)) has uncorrect space dimension ($(size(p,1)) instead of $(N_space_dimension))")
            return false
        end
    end
    # check if the filename does not contain FOLDER_UNITCELLS (this would mean someone forgot a "$")
    if contains(unitcell.filename, "FOLDER_UNITCELLS")
        println("filename ($(unitcell.filename)) still contains \"FOLDER_UNITCELLS\"")
        return false
    end
    # check if connectivity okay
    for c1 in unitcell.connections
        # so far no counterpart found
        counterpart = false
        # search for a counterpart in all connections
        for c2 in unitcell.connections
            # skip identical
            if c1==c2
                continue
            end
            # skip if indices of sites are not correct
            if c1[1] != c2[2] || c1[2] != c2[1]
                continue
            end
            # skip if connection strength not equal
            if c1[3] != c2[3]
                continue
            end
            # skip if the unitcell wraps dont sum up
            if sum([abs(element) for element in collect(c1[4]) .+ collect(c2[4])]) > 1e-9
                continue
            end
            # if everything passed up to here, a counterpart has been found
            counterpart = true
            break
        end
        if counterpart == false
            println("connection $(c1) does not have counterpart, connectivity broken")
            return false
        end
    end
    # now, return true as all tests are passed
    return true
end
export testUnitcell

"""
    testLattice(lattice::Lattice, N_lattice_vectors::Int64, N_space_dimension::Int64)

Performs a series of tests on a given `Lattice` object so see if it is implemented correctly.
Checks include the following:
1) test if the number of `lattice_vectors` is as expected
2) test if the number of spatial dimensions is correct for every site in the lattice
3) test if the filename does not contain `"FOLDER_LATTICES"` (as in this case a \$ would be missing somewhere)
4) test if the connectivity is okay (if every connection has a counterpart)

The return value is a `Bool` indicating if the test is passed or not. If the test fails, the
failure reason is printed as well.

# Examples

```julia-repl
julia> testLattice(lattice, 2, 2)
...
```
"""
function testLattice(lattice::Lattice, N_lattice_vectors::Int64, N_space_dimension::Int64)
    # check if the unitcell has enough lattice vectors
    if size(lattice.lattice_vectors,1) != N_lattice_vectors
        println("Uncorrect number of lattice vectors ($(size(lattice.lattice_vectors,1)) instead of $(N_lattice_vectors))")
        return false
    end
    # check if the space dimension for all sites is okay
    for p in lattice.positions
        if size(p,1) != N_space_dimension
            println("Site ($(p)) has uncorrect space dimension ($(size(p,1)) instead of $(N_space_dimension))")
            return false
        end
    end
    # check if the filename does not contain FOLDER_UNITCELLS (this would mean someone forgot a "$")
    if contains(lattice.filename, "FOLDER_LATTICES")
        println("filename ($(lattice.filename)) still contains \"FOLDER_LATTICES\"")
        return false
    end
    # check if connectivity okay
    for c1 in lattice.connections
        # so far no counterpart found
        counterpart = false
        # search for a counterpart in all connections
        for c2 in lattice.connections
            # skip identical
            if c1==c2
                continue
            end
            # skip if indices of sites are not correct
            if c1[1] != c2[2] || c1[2] != c2[1]
                continue
            end
            # skip if connection strength not equal
            if c1[3] != c2[3]
                continue
            end
            # skip if the lattice wraps dont sum up
            if sum([abs(element) for element in collect(c1[4]) .+ collect(c2[4])]) > 1e-9
                continue
            end
            # if everything passed up to here, a counterpart has been found
            counterpart = true
            break
        end
        if counterpart == false
            println("connection $(c1) does not have counterpart, connectivity broken")
            return false
        end
    end
    # now, return true as all tests are passed
    return true
end
export testLattice

"""
    printInfo(unitcell::Unitcell [; detailed::Bool=false])
    printInfo(lattice::Lattice [; detailed::Bool=false])

Prints (detailed) information about a `Unitcell` object or a `Lattice` object.
Information Includes Bravais lattice vectors, number of sites and number of bonds.
When detailed output is enabled, furthermore all sites and bonds are printed.
Last but now least, statistics is printed about how many bonds per site are found.

# Examples

```julia-repl
julia> printInfo(unitcell)
...

julia> printInfo(unitcell, detailed=true)
...

```
"""
function printInfo(unitcell::Unitcell; detailed::Bool=false)

    # Header of the information
    println("Information on the unitcell stored in file \"$(unitcell.filename)\":")

    # General information on Bravais lattice vectors
    println(" - periodicity given by $(size(unitcell.lattice_vectors,1)) lattice vectors:")
    for l in unitcell.lattice_vectors
        println("     - $(l)")
    end

    # Information on basis sites
    if detailed
        println(" - $(size(unitcell.basis,1)) sites in unitcell of dimension $(length(unitcell.basis[1])):")
        for site in unitcell.basis
            println("     - $(site)")
        end
    else
        println(" - $(size(unitcell.basis,1)) sites in unitcell of dimension $(length(unitcell.basis[1]))")
    end

    # Information on bonds / connections
    if detailed
        println(" - $(size(unitcell.connections,1)) connections in the unitcell:")
        for c in unitcell.connections
            if typeof(c[3]) == String
                println("     - from $(c[1]) to $(c[2]) (with warp $(c[4])): \"$(c[3])\"")
            else
                println("     - from $(c[1]) to $(c[2]) (with warp $(c[4])): $(c[3])")
            end
        end
    else
        println(" - $(size(unitcell.connections,1)) connections in the unitcell")
    end

    # Information on bond / connection statistics
    println(" - $(size(unitcell.connections,1)/size(unitcell.basis,1)) connections per site")
    # check statistics of connections
    nc = zeros(Int64, size(unitcell.basis,1))
    for i in 1:length(nc)
        for c in unitcell.connections
            if Int(c[1]) == i
                nc[i] += 1
            end
        end
    end
    tc = 0
    c = 0
    print(" - statistics of connections per site:")
    while tc < size(unitcell.basis, 1)
        cc = 0
        for ncc in nc
            if ncc == c
                cc += 1
            end
        end
        if cc != 0
            print(" $(c)($(cc))")
        end
        c = c+1
        tc += cc
    end
    println("")
    broken = false
    for c1 in unitcell.connections
        counterpart = false
        for c2 in unitcell.connections
            if c1==c2
                continue
            end
            if c1[1] != c2[2] || c1[2] != c2[1]
                continue # site indices not correct
            end
            if c1[3] != c2[3]
                continue # connection strength not equal
            end
            if sum([abs(element) for element in collect(c1[4]) .+ collect(c2[4])]) > 1e-9
                continue # wrap not equal
            end
            counterpart = true
            break
        end
        if counterpart == false
            broken = true
        end
    end
    if broken
        println(" - connectivity of unitcell is broken (connections not vice versa)")
    else
        println(" - connectivity of unitcell is okay (including periodic BC)")
    end
end
function printInfo(lattice::Lattice; detailed=false)
	# Print the header informations
    println("Information on the lattice stored in file \"$(lattice.filename)\":")

	# Information on lattice vectors, print all of them
    println(" - periodicity given by $(size(lattice.lattice_vectors,1)) lattice vectors:")
    for l in lattice.lattice_vectors
        println("     - $(l)")
    end

	# Information on sites, depending on detailed print information
    if detailed
        println(" - $(size(lattice.positions,1)) sites in lattice of dimension $(length(lattice.positions[1])):")
        for site in lattice.positions
            println("     - $(site)")
        end
    else
        println(" - $(size(lattice.positions,1)) sites in lattice of dimension $(length(lattice.positions[1]))")
    end

	# Information on connectioins, depending on detailed print information
    if detailed
        println(" - $(size(lattice.connections,1)) connections in the lattice:")
        for c in lattice.connections
            if typeof(c[3]) == String
                println("     - from $(c[1]) to $(c[2]) (with warp $(c[4])): \"$(c[3])\"")
            else
                println("     - from $(c[1]) to $(c[2]) (with warp $(c[4])): $(c[3])")
            end
        end
    else
        println(" - $(size(lattice.connections,1)) connections in the lattice")
    end
	# Number of connections per site
    println(" - $(size(lattice.connections,1)/size(lattice.positions,1)) connections per site")
    # check statistics of connections
    nc = zeros(Int64, size(lattice.positions,1))
    for i in 1:length(nc)
        for c in lattice.connections
            if Int(c[1]) == i
                nc[i] += 1
            end
        end
    end
    tc = 0
    c = 0
	# print the statistics of the connections
    print(" - statistics of connections per site:")
    while tc < size(lattice.positions, 1)
        cc = 0
        for ncc in nc
            if ncc == c
                cc += 1
            end
        end
        if cc != 0
            print(" $(c)($(cc))")
        end
        c = c+1
        tc += cc
    end
    println("")
	# Test the connectivity of the lattice
    broken = false
    for c1 in lattice.connections
        counterpart = false
        for c2 in lattice.connections
            if c1==c2
                continue
            end
            if c1[1] != c2[2] || c1[2] != c2[1]
                continue # site indices not correct
            end
            if c1[3] != c2[3]
                continue # connection strength not equal
            end
            if sum([abs(element) for element in collect(c1[4]) .+ collect(c2[4])]) > 1e-9
                continue # wrap not equal
            end
            counterpart = true
            break
        end
        if counterpart == false
            broken = true
        end
    end
	# Print the connectivity of the lattice
    if broken
        println(" - connectivity of lattice is broken (connections not vice versa)")
    else
        println(" - connectivity of lattice is okay (including periodic BC)")
    end
end
export printInfo

"""
    getConnectionStrengthList(unitcell::Unitcell)
	getConnectionStrengthList(lattice::Lattice)

Obtains which connection strength values are taken in the definition of a given
`Lattice` object `lattice` or a `Unitcell` object `unitcell`.
The functions return a list which contains all connection strengths in the order
of appearence.

# Examples

```julia-repl
julia> connection_strength_list = getConnectionStrengthList(lattice)
["tx", "ty", "tz"]
```
"""
function getConnectionStrengthList(unitcell::Unitcell)
    # list of connetion strengths
    cs_list = []
    # iterate over all connections
    for c in unitcell.connections
        if !(c[3] in cs_list)
            push!(cs_list, c[3])
        end
    end
    # return the list
    return cs_list
end
function getConnectionStrengthList(lattice::Lattice)
    # list of connetion strengths
    cs_list = []
    # iterate over all connections
    for c in lattice.connections
        if !(c[3] in cs_list)
            push!(cs_list, c[3])
        end
    end
    # return the list
    return cs_list
end

# export the FUNCTIONS
export getConnectionStrengthList

#-----------------------------------------------------------------------------
#
#   METHODS FOR CONNECTIVITY INFORMATION
#
#-----------------------------------------------------------------------------

"""
	getConnectivityList(lattice::Lattice)
    getConnectivityList(unitcell::Unitcell)

Obtains connectivity information of a `Lattice` object `lattice` or a `Unitcell` object `unitcell`.
The functions return a list which contains a list of Tuples for every site in the
passed objects. Therefore, `connectivity_list[i][j]` can be used to access the connection `j`
from site `i` which will give a Tuple `(to_index, strength)`.

# Examples

```julia-repl
julia> connectivity_list = getConnectivityList(lattice)
Array{...}[...]

julia> connectivity_list[2][1]   # connection 1 outgoing from site 2
(12, "tx")
```
"""
function getConnectivityList(lattice::Lattice)
    # list of lists
    lists = Array[]
    for p in lattice.positions
        push!(lists, [])
    end
    for c in lattice.connections
        push!(lists[Int(c[1])], (c[2], c[3]) )
    end
    # return the lists
    return lists
end
function getConnectivityList(unitcell::Unitcell)
    # list of lists
    lists = Array[]
    for p in unitcell.basis
        push!(lists, [])
    end
    for c in unitcell.connections
        push!(lists[Int(c[1])], (c[2], c[3]) )
    end
    # return the lists
    return lists
end

# export the FUNCTIONS
export getConnectivityList

"""
	getConnectionList(lattice::Lattice)
    getConnectionList(unitcell::Unitcell)

Obtains connectivity information of a `Lattice` object `lattice` or a `Unitcell` object `unitcell`.
The functions return a list which contains a list of connection Arrays for every site in the
passed objects. Therefore, `connectivity_list[i][j]` can be used to access the connection `j`
from site `i` which will give an Array `[from_index, to_index, strength, wrap_tuple]`
like it is saved in the list of all connections of the respective object.

# Examples

```julia-repl
julia> connection_list = getConnectionList(lattice)
Array{...}[...]

julia> connection_list[2][1]   # connection 1 outgoing from site 2
[1, 12, "tx", (0,0)]
```
"""
function getConnectionList(lattice::Lattice)
    # list of lists
    lists = Array[]
    for p in lattice.positions
        push!(lists, Array[])
    end
    for c in lattice.connections
        push!(lists[Int(c[1])], c )
    end
    # return the lists
    return lists
end
function getConnectionList(unitcell::Unitcell)
    # list of lists
    lists = Array[]
    for p in unitcell.basis
        push!(lists, Array[])
    end
    for c in unitcell.connections
        push!(lists[Int(c[1])], c )
    end
    # return the lists
    return lists
end

# export the FUNCTIONS
export getConnectionList


mutable struct DDEIntegrator{algType,IIP,uType,tType,P,eigenType,absType,relType,
                             residType,tTypeNoUnits,tdirType,ksEltype,
                             SolType,F,CacheType,
                             IType,NType,O,tstopsType,
                             FSALType,EventErrorType} <: AbstractDDEIntegrator{algType,IIP,uType,tType}
    sol::SolType
    u::uType
    k::ksEltype
    t::tType
    dt::tType
    f::F
    p::P
    uprev::uType
    uprev2::uType
    tprev::tType
    prev_idx::Int
    prev2_idx::Int
    fixedpoint_abstol::absType
    fixedpoint_reltol::relType
    resid::residType # This would have to resize for resizing DDE to work
    fixedpoint_norm::NType
    max_fixedpoint_iters::Int
    saveat::tstopsType
    tracked_discontinuities::Vector{Discontinuity{tType}}
    alg::algType
    dtcache::tType
    dtchangeable::Bool
    dtpropose::tType
    tdir::tdirType
    eigen_est::eigenType
    EEst::tTypeNoUnits
    qold::tTypeNoUnits
    q11::tTypeNoUnits
    erracc::tTypeNoUnits
    dtacc::tType
    success_iter::Int
    iter::Int
    saveiter::Int
    saveiter_dense::Int
    cache::CacheType
    kshortsize::Int
    force_stepfail::Bool
    last_stepfail::Bool
    just_hit_tstop::Bool
    event_last_time::Int
    last_event_error::EventErrorType
    accept_step::Bool
    isout::Bool
    reeval_fsal::Bool
    u_modified::Bool
    opts::O
    integrator::IType
    fsalfirst::FSALType
    fsallast::FSALType

    # incomplete initialization without fsalfirst and fsallast
    function DDEIntegrator{algType,IIP,uType,tType,P,eigenType,absType,relType,
                           residType,tTypeNoUnits,
                           tdirType,ksEltype,SolType,F,CacheType,IType,
                           NType,O,tstopsType,FSALType,EventErrorType}(
                               sol,u,k,t,dt,f,p,uprev,uprev2,tprev,prev_idx,prev2_idx,
                               fixedpoint_abstol,fixedpoint_reltol,resid,fixedpoint_norm,
                               max_fixedpoint_iters,saveat,tracked_discontinuities,alg,
                               dtcache,dtchangeable,dtpropose,tdir,eigen_est,EEst,qold,
                               q11,erracc,dtacc,success_iter,iter,saveiter,saveiter_dense,
                               cache,kshortsize,force_stepfail,last_stepfail,
                               just_hit_tstop,event_last_time,last_event_error,
                               accept_step,isout,reeval_fsal,u_modified,opts,
                               integrator) where
        {algType,IIP,uType,tType,P,eigenType,absType,relType,residType,tTypeNoUnits,
        tdirType,ksEltype,
         SolType,F,CacheType,IType,NType,O,tstopsType,FSALType,EventErrorType}

        new{algType,IIP,uType,tType,P,eigenType,absType,relType,
                               residType,tTypeNoUnits,
                               tdirType,ksEltype,SolType,F,CacheType,IType,
                               NType,O,tstopsType,FSALType,EventErrorType}(
            sol,u,k,t,dt,f,p,uprev,uprev2,tprev,prev_idx,prev2_idx,fixedpoint_abstol,
            fixedpoint_reltol,resid,fixedpoint_norm,max_fixedpoint_iters,saveat,
            tracked_discontinuities,alg,dtcache,dtchangeable,dtpropose,tdir,eigen_est,
            EEst,qold,q11,erracc,dtacc,success_iter,iter,saveiter,saveiter_dense,cache,
            kshortsize,force_stepfail,last_stepfail,just_hit_tstop,accept_step,
            event_last_time,last_event_error,isout,reeval_fsal,u_modified,opts,integrator)
    end
end

function (integrator::DDEIntegrator)(t, deriv::Type=Val{0}; idxs=nothing)
    OrdinaryDiffEq.current_interpolant(t, integrator, idxs, deriv)
end

function (integrator::DDEIntegrator)(val::AbstractArray, t::Union{Number,AbstractArray},
                                     deriv::Type=Val{0}; idxs=nothing)
    OrdinaryDiffEq.current_interpolant!(val, t, integrator, idxs, deriv)
end


@doc """  if pos_gene have this position, return gene; else return ""
          Foundational function
""" ->
function query_gene(chr::ASCIIString,pos::Int64)
    info("query_gene for a $chr,$pos")
    if !isdefined(GeneMisc,:chr_sted_gene)
        warn("chr_sted_gene is not defined,load_index starts")
        load_index()
    end
    sted_gene = chr_sted_gene[chr]

    steds = collect(keys(sted_gene))
    # sorted by acclerator speed
    starts = sort(map(x->x[1], steds))
    ends   = sort(map(x->x[2], steds))
    debug = true
    function check_sted(starts,ends)
        l = length(starts)
        @assert length(ends) == l
        for i in 1:l
            @assert starts[i] < ends[i]
        end
    end
    if debug check_sted(starts,ends) end
    idx_st = searchsortedfirst(starts, pos)
    idx_ed = searchsortedfirst(ends, pos)
    function result(st_idx,ed_idx)
        info("$st_idx, $ed_idx")
        st = starts[st_idx]
        ed = ends[ed_idx]
        sted_gene[(st,ed)]
    end
    if idx_st == 1
        if starts[1] < pos
            return ""
        end
        @assert starts[1] >= pos
        return result(1,1)
    end
    if idx_st == idx_ed
        #if ends[idx_ed] > pos
        return ""
        #end
        #return result(idx_st)
    else
        @assert idx_st > idx_ed
        return result(idx_st-1,idx_ed)
    end

    #=
    for (st,ed) in steds
        if st<=pos<=ed
            return sted_gene[(st,ed)]
        end
    end
    =#

    ""
end

@doc """ if pos_gene have this position, return gene; else return ""
""" ->
function query_gene(chr::ASCIIString,pos::ASCIIString)
    pos_int = parse(Int64,pos)
    query_gene(chr,pos_int)
end

@doc """ query a lot of positions
""" ->
function query_gene(data::Array{ASCIIString,2})
    info("query gene location for you ...")
    @assert size(data,2) == 2
    n = size(data,1)
    genes = Array{ASCIIString,1}(n)
    for i = 1:n
        chr,pos = data[i,:]
        genes[i] = query_gene(chr,pos)
    end
    genes
end

@doc """ Given a gene name, find its location: chr, st, ed
""" ->
function query_geneloc(genename::ASCIIString)
    if !isdefined(GeneMisc, :gene_chrsted)
        load_index()
    end
    gene_chrsted[genename]
end


using SeaPearl
using SeaPearlExtras
using ReinforcementLearning
const RL = ReinforcementLearning
using Flux
using GeometricFlux
using JSON
using BSON: @load, @save
using Random
using Dates

# -------------------
# Generator
# -------------------
n_city = 21
grid_size = 100
max_tw_gap = 0
max_tw = 100
tsptw_generator = SeaPearl.TsptwGenerator(n_city, grid_size, max_tw_gap, max_tw, true)
SR = SeaPearl.TsptwStateRepresentation{SeaPearl.TsptwFeaturization, SeaPearl.TsptwTrajectoryState}

# -------------------
# Internal variables
# -------------------
numInFeatures=SeaPearl.feature_length(SR)

# -------------------
# Experience variables
# -------------------
nbEpisodes = 500
evalFreq = 200
nbInstances = 10
nbRandomHeuristics = 1

# -------------------
# Agent definition
# -------------------
include("agents.jl")

# -------------------
# Value Heuristic definition
# -------------------
learnedHeuristic = SeaPearl.LearnedHeuristic{SR, SeaPearl.TsptwReward, SeaPearl.VariableOutput}(agent)
include("nearest_heuristic.jl")
nearest_heuristic = SeaPearl.BasicHeuristic(select_nearest_neighbor) # Basic value-selection heuristic

function select_random_value(x::SeaPearl.IntVar; cpmodel=nothing)
    selected_number = rand(1:length(x.domain))
    i = 1
    for value in x.domain
        if i == selected_number
            return value
        end
        i += 1
    end
    @assert false "This should not happen"
end

randomHeuristics = []
for i in 1:nbRandomHeuristics
    push!(randomHeuristics, SeaPearl.BasicHeuristic(select_random_value))
end

valueSelectionArray = [learnedHeuristic, nearest_heuristic]
append!(valueSelectionArray, randomHeuristics)

# -------------------
# Variable Heuristic definition
# -------------------
struct TsptwVariableSelection{TakeObjective} <: SeaPearl.AbstractVariableSelection{TakeObjective} end
TsptwVariableSelection(;take_objective=false) = TsptwVariableSelection{take_objective}()
function (::TsptwVariableSelection{false})(cpmodel::SeaPearl.CPModel; rng=nothing)
    for i in 1:length(keys(cpmodel.variables))
        if haskey(cpmodel.variables, "a_"*string(i)) && !SeaPearl.isbound(cpmodel.variables["a_"*string(i)])
            return cpmodel.variables["a_"*string(i)]
        end
    end
end
variableSelection = TsptwVariableSelection()

# -------------------
# -------------------
# Core function
# -------------------
# -------------------
function trytrain(nbEpisodes::Int)
    experienceTime = now()
    dir = mkdir(string("exp_",Base.replace("$(round(experienceTime, Dates.Second(3)))",":"=>"-")))
    expParameters = Dict(
        :nbEpisodes => nbEpisodes,
        :evalFreq => evalFreq,
        :nbInstances => nbInstances
    )
    open(dir*"/params.json", "w") do file
        JSON.print(file, expParameters)
    end

    metricsArray, eval_metricsArray = SeaPearl.train!(
    valueSelectionArray=valueSelectionArray,
    generator=tsptw_generator,
    nbEpisodes=nbEpisodes,
    strategy=SeaPearl.DFSearch(),
    variableHeuristic=variableSelection,
    out_solver=true,
    verbose = false,
    evaluator=SeaPearl.SameInstancesEvaluator(valueSelectionArray,tsptw_generator; evalFreq = evalFreq, nbInstances = nbInstances)
)

    trained_weights = params(agent.policy.learner.approximator.model)
    @save dir*"/model_weights_tsptw"*string(n_city)*".bson" trained_weights

    SeaPearlExtras.storedata(metricsArray[1]; filename=dir*"/tsptw_$(n_city)_training")
    SeaPearlExtras.storedata(eval_metricsArray[:,1]; filename=dir*"/tsptw_$(n_city)_trained")
    SeaPearlExtras.storedata(eval_metricsArray[:,2]; filename=dir*"/tsptw_$(n_city)_nearest")
    for i = 1:nbRandomHeuristics
        SeaPearlExtras.storedata(eval_metricsArray[:,i+2]; filename=dir*"/tsptw_$(n_city)_random$(i)")
    end

    return metricsArray, eval_metricsArray
end

metricsArray, eval_metricsArray = trytrain(nbEpisodes)
nothing


### Julia OpenStreetMap Package ###
### MIT License                 ###
### Copyright 2014              ###

type Highway
    @compat class::AbstractString       # Type of highway
    lanes::Int          # Number of lanes (1 if unspecified)
    oneway::Bool        # True if road is one-way
    @compat sidewalk::AbstractString    # Sidewalk classifier, if available
    @compat cycleway::AbstractString    # Cycleway classifier, if available
    @compat bicycle::AbstractString     # Bicycle classifier, if available
    @compat name::AbstractString        # Name, if available
    maxspeed::Float64       # Speed limit in kph
    nodes::Vector{Int}  # List of nodes
end

type Segment
    node0::Int          # Source node ID
    node1::Int          # Target node ID
    nodes::Vector{Int}  # List of nodes falling within node0 and node1
    dist::Real          # Length of the segment
    class::Int          # Class of the segment
    parent::Int         # ID of parent highway
    oneway::Bool        # True if road is one-way
end

type Feature
    @compat class::AbstractString       # Shop, amenity, crossing, etc.
    @compat detail::AbstractString      # Class qualifier
    @compat name::AbstractString        # Name
end

type Building
    @compat class::AbstractString       # Building type (usually "yes")
    @compat name::AbstractString        # Building name (usually unavailable)
    nodes::Vector{Int}  # List of nodes
end

type Intersection
    highways::Set{Int}  # Set of highway IDs
end
Intersection() = Intersection(Set{Int}())

type HighwaySet # Multiple highways representing a single "street" with a common name
    highways::Set{Int}
end

# Transporation network graph data and helpers to increase routing speed
type Network
    g                                   # Graph object
    v::Dict{Int,Graphs.KeyVertex{Int}}  # (node id) => (graph vertex)
    w::Vector{Float64}                  # Edge weights, indexed by edge id
    class::Vector{Int}                 # Road class of each edge
end

### Rendering style data
type Style
    @compat color::UInt32
    width::Real
    @compat spec::AbstractString
end
Style(x, y) = Style(x, y, "-")


"""
```
simulate_switching(m, scen::SwitchingScenario; verbose = :low)
```

Simulate switching in and out of a default scenario for the
SwitchingScenario scen. See `SwitchingScenario` for more
info. Returns a dictionary of results and the proportion
of times a switch actually occured.
"""
function simulate_switching(m::AbstractDSGEModel, scen::SwitchingScenario;
                            verbose::Symbol = :low)

    info_print(verbose, :low, "Simulating switching for " * string(scen.key) * "...")
    println(verbose, :low, "Start time: " * string(now()))
    println(verbose, :low, "Outputs will be saved in " * rawpath(m, "scenarios"))

    tic = time_ns()
    # Revert model alt policy to historical rule
    m <= Setting(:alternative_policy, AltPolicy(:historical, solve, eqcond), false, "apol",
                 "Alternative policy")

    # Get output file names
    original_output_files = get_scenario_output_files(m, scen.original,
                                                      [:forecastobs, :forecastpseudo])
    default_output_files = get_scenario_output_files(m, scen.default,
                                                     [:forecastobs, :forecastpseudo])

    results = Dict{Symbol, Array{Float64}}()
    switching_results = [0.0, 0.0]

    for (i, output_var) in enumerate([:forecastobs, :forecastpseudo])
        # Read in original and default draws
        original_draws = JLD2.jldopen(original_output_files[output_var], "r") do file
            read(file, "arr")
        end
        default_draws  = JLD2.jldopen(default_output_files[output_var], "r") do file
            read(file, "arr")
        end

        n_draws, n_vars, n_periods = size(original_draws)
        n_default_draws = size(default_draws, 1)

        @assert n_vars == size(default_draws, 2)
        @assert length(scen.probs_enter) == n_periods <= size(default_draws, 3)

        # Simulate switching
        results[output_var] = zeros(n_draws, n_vars, n_periods)
        proportion_switched = 0.0
        for j = 1:n_draws
            k = rand(1:n_default_draws)
            results[output_var][j, :, :], switched = switch(original_draws[j, :, :], default_draws[k, :, :],
                                                            scen.probs_enter, scen.probs_exit)
            proportion_switched += Float64(switched)
        end
        switching_results[i] = proportion_switched / n_draws
    end

    results[:proportion_switched] = [mean(switching_results)]

    # Write outputs
    output_files = get_scenario_output_files(m, scen, [:forecastobs, :forecastpseudo])
    write_scenario_forecasts(m, output_files, results, verbose = verbose)

    # Print
    switching_time = (time_ns() - tic)/1e9
    switching_time_min = switching_time/60
    println(verbose, :low, "\nTime elapsed: " * string(switching_time_min) * " minutes")
    println(verbose, :low, "Switching complete: " * string(now()))

    return results, switching_results
end

#for switching with multiple models
function simulate_switching(m_original::AbstractDSGEModel, m_default::AbstractDSGEModel, scen::SwitchingScenario;
                            verbose::Symbol = :low)

    info(verbose, :low, "Simulating switching for " * string(scen.key) * "...")
    println(verbose, :low, "Start time: " * string(now()))
    println(verbose, :low, "Outputs will be saved in " * rawpath(m_original, "scenarios"))
    tic()

    # Revert model alt policy to historical rule
    m_original <= Setting(:alternative_policy, AltPolicy(:historical, solve, eqcond), false, "apol",
                 "Alternative policy")
    m_default <= Setting(:alternative_policy, AltPolicy(:historical, solve, eqcond), false, "apol",
                 "Alternative policy")

    # Get output file names
    original_output_files = get_scenario_output_files(m_original, scen.original,
                                                      [:forecastobs, :forecastpseudo])
    default_output_files = get_scenario_output_files(m_default, scen.default,
                                                     [:forecastobs, :forecastpseudo])

    results = Dict{Symbol, Array{Float64}}()
    switching_results = [0.0, 0.0]

    for (i, output_var) in enumerate([:forecastobs, :forecastpseudo])
        if output_var==:forecastobs
            include_obs = intersect(keys(m_default.observables), keys(m_original.observables))
            inds_default = Vector(0)
            inds_original = Vector(0)
            for key in include_obs
                inds_default = vcat(inds_default, m_default.observables[key])
                inds_original = vcat(inds_original, m_original.observables[key])
            end
        elseif output_var==:forecastpseudo
            include_obs = intersect(keys(m_default.pseudo_observables), keys(m_original.pseudo_observables))
            inds_default = Vector(0)
            inds_original = Vector(0)
            for key in include_obs
                inds_default = vcat(inds_default, m_default.pseudo_observables[key])
                inds_original = vcat(inds_original, m_original.pseudo_observables[key])
            end
        end

        # Read in original and default draws
        original_draws = load(original_output_files[output_var], "arr")[:, inds_original, :]
        default_draws  = load(default_output_files[output_var], "arr")[:, inds_default, :]

        n_draws, n_vars, n_periods = size(original_draws)
        n_default_draws = size(default_draws, 1)

        @assert n_vars == size(default_draws, 2)
        @assert length(scen.probs_enter) == n_periods <= size(default_draws, 3)

        # Simulate switching
        results[output_var] = zeros(n_draws, n_vars, n_periods)
        proportion_switched = 0.0
        for j = 1:n_draws
            k = rand(1:n_default_draws)
            results[output_var][j, :, :], switched = switch(original_draws[j, :, :], default_draws[k, :, :],
                                                            scen.probs_enter, scen.probs_exit)
            proportion_switched += Float64(switched)
        end
        switching_results[i] = proportion_switched / n_draws
    end

    results[:proportion_switched] = [mean(switching_results)]

    # Write outputs
    output_files = get_scenario_output_files(m_original, scen, [:forecastobs, :forecastpseudo])
    write_scenario_forecasts(m_original, output_files, results, verbose = verbose)

    # Print
    switching_time = toq()
    switching_time_min = switching_time/60
    println(verbose, :low, "\nTime elapsed: " * string(switching_time_min) * " minutes")
    println(verbose, :low, "Switching complete: " * string(now()))

    return results, switching_results
end

"""
```
switch(original, default, probs_enter, probs_exit)
```

Simulate entry and exit from the original vector into the default one
according to the specified probabilities.
"""
function switch(original::Matrix{Float64}, default::Matrix{Float64},
                probs_enter::Vector{Float64}, probs_exit::Vector{Float64})

    n_vars, n_periods = size(original)
    output = zeros(n_vars, n_periods)

    # Find period of switch from default to original
    period_in = choose_switch_period(probs_enter, 1)

    # Record whether original is actually switched into
    switched = period_in < n_periods

    # Find period of switch from original to default
    # Start in period_in + 1 s.t. at least 1 period is spent in the original scenario
    period_out = choose_switch_period(probs_exit, period_in + 1)

    output[:, 1:(period_in-1)]          = default[:, 1:(period_in-1)]
    output[:, period_in:(period_out-1)] = original[:, 1:(period_out-period_in)]
    output[:, period_out:n_periods]     = default[:, period_out:n_periods]

    return output, switched
end

function choose_switch_period(probs::Vector{Float64}, first_period::Int)
    n_periods = length(probs)

    for t = first_period:n_periods
        if rand() < probs[t]
            return t
        end
    end

    return n_periods + 1
end


# extract estimates from jld file

using Caesar, Distributions, RoME

const KDE = KernelDensityEstimate

datetimestamp = "2018-08-09T00:04:36.373"
datadir = "/home/dehann/Pictures/racecarimgs/$(datetimestamp)/"
filename = datadir*"racecar_fg_x237.jld"

fg, = loadjld(file=filename)

# using RoMEPlotting
# drawPosesLandms(fg, spscale=0.1)

xx,ll = ls(fg)

fid = open(datadir*"results_$(datetimestamp).csv","w")
for x in xx
  # @show x,
  mx = KDE.getKDEMean(getBelief(getVariable(fg, x)))
  println(fid, "$x, $(mx[1]), $(mx[2]), $(mx[3]), $(mx[4]), $(mx[5])")
end
for x in ll
  # @show x,
  mx = KDE.getKDEMean(getBelief(getVariable(fg, x)))
  println(fid, "$x, $(mx[1]), $(mx[2])")
end
close(fid)

#


module DataGen

# Types
import RIOPA: TagBase, DataStream, PayloadRange, DataSet
import OrderedCollections: LittleDict
# Functions
import RIOPA: evolve_payload_range!
# Modules
import MPI, Random

abstract type DataGenTag <: TagBase end

function generate! end
function initialize_streams! end

tagmap = LittleDict{String,DataGenTag}()

function add(key::String, tag::DataGenTag)
    tagmap[key] = tag
end

function get_tag(key::String)
    try
        return tagmap[key]
    catch
        println("Invalid DataGen backend: ", key)
        rethrow()
    end
end

struct DefaultDataGenTag <: DataGenTag end

get_tag(::Nothing) = DefaultDataGenTag()

struct ProcessPayloadRatioError <: Exception
    msg::String
end

function check_payload_group_ratios(ds::DataSet)
    for stream_cfg in ds.cfg.streams
        sum = 0.0
        for grp in stream_cfg.payload_groups
            sum += grp.ratio
        end
        if !isapprox(sum, 1.0)
            throw(
                ProcessPayloadRatioError(
                    "Sum of payload group ratios ($sum) must equal 1" *
                    "; dataset: " *
                    ds.cfg.name *
                    ", stream: " *
                    stream_cfg.name,
                ),
            )
        end
    end
end

function initialize_streams!(::DefaultDataGenTag, ds::DataSet)
    check_payload_group_ratios(ds)
    ds.streams = map(stream_cfg -> DataStream(stream_cfg), ds.cfg.streams)
end

function generate_stream_data!(
    stream::DataStream,
    step::Integer,
    step_factor::Integer,
)
    newsize = rand((stream.range.a):(stream.range.b))
    resize!(stream.data.vec, newsize)
    Random.rand!(stream.data.vec)
    evolve_payload_range!(stream, step * step_factor)
end

function generate!(::DefaultDataGenTag, ds::DataSet)
    foreach(
        stream ->
            generate_stream_data!(stream, ds.curr_step, ds.cfg.step_factor),
        ds.streams,
    )
end

end


# FourierLayerState
immutable FourierLayerState <: AbstractScatteredLayerState
    blobs::Vector{Mocha.Blob}
    layer::FourierLayer
end

function FourierLayerState(
        backend::Mocha.CPUBackend,
        layer::FourierLayer,
        inputs::Vector{Mocha.Blob},
        diffs::Vector{Mocha.Blob})
    blobs = Vector{Mocha.Blob}(length(inputs))
    for idblob in eachindex(inputs)
        innodes = inputs[idblob].nodes
        # TODO: tighten typing
        outnodes = DataStructures.SortedDict{Path,RealFourierNode,
            Base.Order.ForwardOrdering}()
        for path in keys(innodes)
            outnodes[path] = RealFourierNode(
                innodes[path],
                findin(innodes[path].ranges, layer.pathkeys),
                layer.flags,
                layer.timelimit)
        end
        blobs[idblob] = ScatteredBlob(outnodes)
    end
    return FourierLayerState(blobs, layer)
end

function Base.findin{N}(
        ranges::NTuple{N,Pair{PathKey,StepRange{Int,Int}}},
        fourierkeys::Vector{PathKey})
    return find([any(r.first.==fourierkeys) for r in ranges])
end

function Mocha.setup(
        backend::Mocha.Backend,
        layer::FourierLayer,
        inputs::Vector{Mocha.Blob},
        diffs::Vector{Mocha.Blob})
    return FourierLayerState(backend, layer, inputs, diffs)
end


# Abstract path function
abstract type PathFunction{type <: FunctionType} end

# PathFunction contructors
function PathFunction(type::FunctionType, func!::Function, nFuncs::Int, nStates::Int, nControls::Int, nStatic::Int)
    return ADPathFunction(type, func!, nFuncs, nStates, nControls, nStatic)
end
function PathFunction(type::FunctionType, func!::Function, stateJac!::Union{Function,Nothing},
    controlJac!::Union{Function,Nothing}, staticJac!::Union{Function,Nothing}, timeJac!::Union{Function,Nothing},
    nFuncs::Int, nStates::Int, nControls::Int, nStatic::Int,
    stateSP::AbstractVecOrMat, controlSP::AbstractVecOrMat, staticSP::AbstractVecOrMat, timeSP::AbstractVecOrMat)
    return AnalyticPathFunction(type, func!, stateJac!, controlJac!, staticJac!, timeJac!,
        nFuncs, nStates, nControls, nStatic, stateSP, controlSP, staticSP, timeSP)
end

# Get the type of path function
function GetFunctionType(pf::PathFunction{type}) where {type}
    return type
end

# Evaluate functions method
function EvaluateFunction(pf::PathFunction,
                          out::AbstractVector,
                          state::AbstractVector,
                          control::AbstractVector,
                          static::AbstractVector,
                          time::AbstractFloat)
    pf.func!(out,state,control,static,time)
    return nothing
end

# Get number of parameters methods
function GetNumberOfFunctions(pf::PathFunction)
    return pf.nFuncs
end
function GetNumberOfStates(pf::PathFunction)
    return pf.nStates
end
function GetNumberOfControls(pf::PathFunction)
    return pf.nControls
end
function GetNumberOfStatics(pf::PathFunction)
    return pf.nStatic
end

# Function name methods
function SetFunctionNames!(pf::PathFunction, names::Vector{String})
    for i in 1:length(names)
        push!(pf.funcNames, name[i])
    end
    return nothing
end
function GetFunctionNames(pf::PathFunction)
    return pf.funcNames
end

# Get Jacobian Sparsity Methods
function GetJacobianSparsity(jacType::State, pf::PathFunction)
    return pf.stateSP
end
function GetJacobianSparsity(jacType::Control, pf::PathFunction)
    return pf.controlSP
end
function GetJacobianSparsity(jacType::Static, pf::PathFunction)
    return pf.staticSP
end
function GetJacobianSparsity(jacType::Time, pf::PathFunction)
    return pf.timeSP
end


using Unitful, Statistics, StaticArrays, Rotations, CoordinateTransformations

function getdata(x)
    track = Track(x[:track])
    feeder = point(get(x, :feeder, x[:track].data[1,1:2]))
    pellet = pointcollection(get(x, :pellet, missing), x[:track].data[1,3])
    feeder, track, pellet
end
abstract type DungMethod end
struct ClosedNest <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
end
function ClosedNest(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    ClosedNest(feeder, nest, track, pellet)
end
struct Transfer <: DungMethod
    feeder::Point
    track::Track
    pellet::PointCollection
    south::Point
    north::Point
    nest2feeder::Float64
    azimuth::Float64
end
function _getvalueunit(txt, default)
    m = match(r"^(\d+)\s*(\D*)", txt)
    d, txtu = m.captures
    u = try
        _u = getfield(Unitful, Symbol(txtu))
        @assert _u isa Unitful.FreeUnits
        _u
    catch
        default
    end
    Float64(ustrip(uconvert(default, parse(Int, d)*u)))
end
function Transfer(x)
    feeder, track, pellet = getdata(x.data)
    south = point(x.data[:south])
    north = point(x.data[:north])
    nest2feeder = _getvalueunit(x.metadata.setup[:nest2feeder], u"cm")
    azimuth = _getvalueunit(x.metadata.setup[:azimuth], u"°")
    Transfer(feeder, track, pellet, south, north, nest2feeder, azimuth)
end
struct TransferNest <: DungMethod
    feeder::Point
    track::Track
    pellet::PointCollection
    south::Point
    north::Point
    nest2feeder::Float64
    azimuth::Float64
    originalnest::Point
end
function TransferNest(x)
    y = Transfer(x)
    originalnest = point(x.data[:originalnest])
    TransferNest(y.feeder, y.track, y.pellet, y.south, y.north, y.nest2feeder, y.azimuth, originalnest)
end
Transfer(x::TransferNest) = Transfer(x.feeder, x.track, x.pellet, x.south, x.north, x.nest2feeder, x.azimuth)
struct TransferNestBelen <: DungMethod
    feeder::Point
    track::Track
    pellet::PointCollection
    southbefore::Point
    northbefore::Point
    feederbefore::Point
    nestbefore::Point
    south::Point
    north::Point
end
function TransferNestBelen(x)
    feeder, track, pellet = getdata(x.data)
    southbefore = point(x.data[:southbefore])
    northbefore = point(x.data[:northbefore])
    feederbefore = point(x.data[:feederbefore])
    nestbefore = point(x.data[:nestbefore])
    south = point(x.data[:south])
    north = point(x.data[:north])
    TransferNestBelen(feeder, track, pellet, rightdowninitial, leftdowninitial, rightupinitial, leftupinitial, rightdownfinal, leftdownfinal, rightupfinal, leftupfinal)
end
struct DawaySandpaper <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    rightdowninitial::Point
    leftdowninitial::Point
    rightupinitial::Point
    leftupinitial::Point
    rightdownfinal::Point
    leftdownfinal::Point
    rightupfinal::Point
    leftupfinal::Point
end
function DawaySandpaper(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    rightdowninitial = point(x.data[:rightdowninitial])
    leftdowninitial = point(x.data[:leftdowninitial])
    rightupinitial = point(x.data[:rightupinitial])
    leftupinitial = point(x.data[:leftupinitial])
    rightdownfinal = point(x.data[:rightdownfinal])
    leftdownfinal = point(x.data[:leftdownfinal])
    rightupfinal = point(x.data[:rightupfinal])
    leftupfinal = point(x.data[:leftupfinal])
    DawaySandpaper(feeder, nest, track, pellet, rightdowninitial, leftdowninitial, rightupinitial, leftupinitial, rightdownfinal, leftdownfinal, rightupfinal, leftupfinal)
end
struct DawayNest <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    pickup::Point
end
function DawayNest(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    pickup = point(x.data[:pickup])
    DawayNest(feeder, nest, track, pellet, pickup)
end
struct Daway <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    initialfeeder::Point
end
function Daway(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    initialfeeder = point(x.data[:initialfeeder])
    Daway(feeder, nest, track, pellet, initialfeeder)
end
struct Displacement <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    pickup::Point
    dropoff::Point
end
function Displacement(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    pickup = point(x.data[:pickup])
    dropoff = point(x.data[:dropoff])
    Displacement(feeder, nest, track, pellet, pickup, dropoff)
end

######################### DungMethod methods ###########

DungMethod(x, displace_location, displace_direction::Missing, transfer::Missing, nest_coverage::Missing) = error("unidentified experimental setup")
DungMethod(x, displace_location::Missing, displace_direction::Missing, transfer::Missing, nest_coverage) = ClosedNest(x)
function DungMethod(x, displace_location::Missing, displace_direction::Missing, transfer, nest_coverage)
    if x.metadata.setup[:person] == "belen"
        TransferNestBelen(x)
    else
        if transfer == "back"
            TransferNest(x)
        elseif transfer ∈ ("far", "50m")
            Transfer(x)
        else
            error("unidentified experimental setup")
        end
    end
end
function DungMethod(x, displace_location, displace_direction, transfer::Missing, nest_coverage)
    if displace_location == "feeder"
        if x.metadata.setup[:person] == "belen"
            DawaySandpaper(x)
        else
            Daway(x)
        end
    elseif displace_location == "nest"
        DawayNest(x)
    elseif displace_location == "halfway"
        Displacement(x)
    else
        error("unidentified experimental setup")
    end
end

######################### Common methods ###########

#=mutable struct Common{N}
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    originalnest::N
end=#
common(x::ClosedNest) = Common(x.feeder, x.nest, x.track, x.pellet, x.nest)
# nest(x::Common) = x.nest
# turning(x::Common) = x.track.homing[end]
# originalnest(x::Common) = x.originalnest

function common(x::TransferNestBelen)
    v = x.northbefore - x.southbefore
    u = x.nestbefore - x.feederbefore
    azimuth = atan(v[2], v[1]) - atan(u[2], u[1])
    nest2feeder = norm(x.nestbefore - x.feederbefore)
    v = x.north - x.south
    α = atan(v[2], v[1]) + azimuth
    u = Point(cos(α), sin(α))
    nest = x.feeder + u*nest2feeder
    Common(x.feeder, nest, x.track, x.pellet, missing)
end

function common(x::Transfer)
    v = x.north - x.south
    α = atan(v[2], v[1]) + x.azimuth - π
    u = Point(cos(α), sin(α))
    nest = x.feeder + u*x.nest2feeder
    Common(x.feeder, nest, x.track, x.pellet, missing)
end

function common(x::TransferNest)
    y = common(Transfer(x))
    Common(y.feeder, y.nest, y.track, y.pellet, x.originalnest)
end

function common(x::DawaySandpaper)
    originalnest = x.nest
    initial = mean(getproperty(x, k) for k in [:rightdowninitial, :leftdowninitial, :rightupinitial, :leftupinitial])
    final = mean(getproperty(x, k) for k in [:rightdownfinal, :leftdownfinal, :rightupfinal, :leftupfinal])
    v = final - initial
    nest = originalnest + v
    _feeder = x.feeder
    feeder = _feeder + v
    Common(feeder, nest, x.track, x.pellet, originalnest)
end

function common(x::DawayNest)
    originalnest = x.nest
    feeder = x.feeder
    v = originalnest - x.pickup
    nest = feeder + v
    Common(feeder, nest, x.track, x.pellet, originalnest)
end

function common(x::Daway)
    originalnest = x.nest
    v = x.feeder - x.initialfeeder
    nest = originalnest + v
    Common(x.feeder, nest, x.track, x.pellet, originalnest)
end

function common(x::Displacement)
    v = x.dropoff - x.pickup
    originalnest = x.nest
    nest = originalnest + v
    Common(x.feeder, nest, x.track, x.pellet, originalnest)
end
######################### END ######################

common(x) = common(DungMethod(x, get(x.metadata.setup, :displace_location, missing), get(x.metadata.setup, :displace_direction, missing), get(x.metadata.setup, :transfer, missing), get(x.metadata.setup, :nest_coverage, missing)))

function get_rotation(nest, feeder)
    v = feeder - nest
    α = -atan(v[2], v[1])# - π/2
    rot = LinearMap(Angle2d(α))
    rot, α
end

(p::LinearMap{Angle2d{Float64}})(x::Missing) = missing

# LinearMap{Angle2d{Float64}}(x::Missing) = missing
function rotate!(x)
    rot, α = get_rotation(x.nest, x.feeder)
    x.nest = rot(x.nest)
    x.feeder = rot(x.feeder)
    x.track.coords .= rot.(x.track.coords)
    # x.track.direction .+= α
    x.pellet.xy .= rot.(x.pellet.xy)
    x.originalnest = rot(x.originalnest)
end

function rotate2!(x)
    rot, α = get_rotation(c, x.feeder)
    x.nest = rot(x.nest)
    x.feeder = rot(x.feeder)
    x.track.coords .= rot.(x.track.coords)
    # x.track.direction .+= α
    x.pellet.xy .= rot.(x.pellet.xy)
    x.originalnest = rot(x.originalnest)
end

function center2!(x)
    x.nest -= c
    x.feeder -= c
    for i in eachindex(x.track.coords)
        x.track.coords[i] -= c
    end
    for i in eachindex(x.pellet)
        x.pellet[i].xy -= c
    end
    if !ismissing(x.originalnest)
        x.originalnest -= c
    end
end


# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

using Distributed
using Random

function _run_benchmark_sample(;
    case::String,
    method::SolutionMethod,
    formulation::Formulation,
    solution_filename::String,
    optimizer,
)::Nothing
    total_time = @elapsed begin
        @info "Reading: $case"
        time_read = @elapsed begin
            instance = read_benchmark(case)
        end
        @info @sprintf("Read problem in %.2f seconds", time_read)
        BLAS.set_num_threads(Threads.nthreads())
        model = build_model(
            instance = instance,
            formulation = formulation,
            optimizer = optimizer,
            variable_names = true,
        )
        @info "Optimizing..."
        BLAS.set_num_threads(1)
        optimize!(model, method)
    end
    @info @sprintf("Total time was %.2f seconds", total_time)

    @info "Writing solution: $solution_filename"
    solution = UnitCommitment.solution(model)
    write("$solution_filename", solution)

    @info "Verifying solution..."
    validate(instance, solution)
    return
end

function _run_benchmark_combination(
    case::String,
    optimizer_name::String,
    optimizer,
    method_name::String,
    method::SolutionMethod,
    formulation_name::String,
    formulation::Formulation,
    trial,
)
    dirname = "results/$optimizer_name/$method_name/$formulation_name/$case"
    function info(msg)
        @info @sprintf(
            "%-8s %-16s %-16s %-16s %-8s %s",
            msg,
            optimizer_name,
            method_name,
            formulation_name,
            trial,
            case
        )
    end
    mkpath(dirname)
    trial_filename = @sprintf("%s/%03d.json", dirname, trial)
    if isfile(trial_filename)
        info("skip")
        return
    end
    info("run")
    open("$trial_filename.log", "w") do file
        redirect_stdout(file) do
            redirect_stderr(file) do
                return _run_benchmark_sample(
                    case = case,
                    method = method,
                    formulation = formulation,
                    solution_filename = trial_filename,
                    optimizer = optimizer,
                )
            end
        end
    end
    return info("done")
end

function _run_benchmarks(;
    cases::Vector{String},
    optimizers::Dict,
    formulations::Dict,
    methods::Dict,
    trials,
)
    combinations = [
        (c, s.first, s.second, m.first, m.second, f.first, f.second, t) for
        c in cases for s in optimizers for f in formulations for
        m in methods for t in trials
    ]
    shuffle!(combinations)
    if nworkers() > 1
        @printf("%24s", "")
    end
    @info @sprintf(
        "%-8s %-16s %-16s %-16s %-8s %s",
        "STATUS",
        "SOLVER",
        "METHOD",
        "FORMULATION",
        "TRIAL",
        "CASE"
    )
    @sync @distributed for c in combinations
        _run_benchmark_combination(c...)
    end
end


using Minibatch
using Revtok
using DataStructures
using GPUArrays
using NNlib
using Flux

english = """This is an account of how magical thinking made us modern.
When people talk about magical thinking, it is usually as a cognitive feature of children, uneducated people, the mushy-minded, or the mentally ill.
If we notice magical thinking in ourselves, it is with a pang of shame: literate adults are supposed to be more sophisticated than that.
At the same time, magical thinking is obviously rampant in the world.
It’s hard not to be fascinated, even if it’s a horrified fascination.
I think that the reason that it has been so difficult to precisely define “magical thinking” is that what we call “magical thinking” is a collection of stigmatized examples of a more general, and generally useful, cognitive capacity.
This is the ability to think in “as if” mode: “as if” inanimate objects had minds, “as if” thoughts could affect reality, “as if” symbols had power over their referents."""

spanish = """Este es un relato de cómo el pensamiento mágico nos hizo modernos.
Cuando la gente habla de pensamiento mágico, generalmente es como un rasgo cognitivo de los niños, de las personas sin educación, de los que tienen la mente blanda o de los enfermos mentales.
Si nos damos cuenta del pensamiento mágico en nosotros mismos, es con una espiral de vergüenza: se supone que los adultos alfabetizados son más sofisticados que eso.
Al mismo tiempo, el pensamiento mágico es obviamente desenfrenado en el mundo.
Es difícil no estar fascinado, aunque sea una fascinación horrorizada.
Creo que la razón por la que ha sido tan difícil definir con precisión el “pensamiento mágico” es que lo que llamamos "pensamiento mágico" es una colección de ejemplos estigmatizados de una capacidad cognitiva más general, y en general útil.
Esta es la capacidad de pensar en modo “como si”: “como si” los objetos inanimados tuvieran mente, “como si” los pensamientos pudieran afectar a la realidad, “como si” los símbolos tuvieran poder sobre sus referentes."""

struct Vocab
    itos::Vector
    stoi::Associative
end
Base.length(vocab::Vocab) = length(vocab.itos)

function process(corpus, vocabsize=200, batchsize=4)
    corpus = tokenize.(split(corpus, '\n'))
    segments = buildvocab(counter(Iterators.flatten(corpus)), vocabsize)
    corpus = segment.(corpus, segments)
    itos = keys(segments)
    stoi = Dict((reverse(p) for p in enumerate(itos)))
    vocab = Vocab(collect(itos), stoi)
    corpus = [[vocab.stoi[x] for x in ex] for ex in corpus]
    corpus = Iterators.partition(corpus, batchsize)
    corpus = map(b -> MaskedBatch(b, (true,)), corpus)
    return corpus, vocab
end

# first step to patching issues that come up when we try to use Flux.Tracker
function Base.getindex(a::Flux.TrackedArray{T, 2} where T, ::Colon, b::MaskedBatch{<:AbstractVector})
    return invoke(getindex, Tuple{AbstractMatrix, Colon, MaskedBatch{<:AbstractVector}}, a, :, b)
end

function posenc(timestep::Integer, channel::Integer, nchannels::Integer)
    if iseven(channel)
        return sin(timestep/(10000^(channel/nchannels)))
    else
        return cos(timestep/(10000^((channel-1)/nchannels)))
    end
end

function posenc(idx::CartesianIndex, nchannels::Integer)
    return posenc(idx[2], idx[1], nchannels)
end

function posenc!(A::GPUArray, state, nchannels::Integer)
    idx = @cartesianidx A state
    @inbounds A[idx] += posenc(idx, nchannels)
    return A
end

function posenc!(A::GPUArray)
    nchannels = size(A, 1)
    gpucall(posenc, A, (nchannels,))
    return A
end

function posenc!(A::AbstractArray)
    nchannels = size(A, 1)
    for idx in CartesianRange(size(A))
        # @inbounds
        A[idx] += posenc(idx, nchannels)
    end
    return A
end

function posenc!(B::MaskedBatch)
    posenc!(B.data)
    B.data = B.data .* B.mask # TODO in-place
    return B
end

function posenc!(A::Flux.TrackedArray)
    return posenc!(Flux.Tracker.data(A))
end

struct Embedding
    W
end
Embedding(vocabsize, embedsize) = Embedding(param(randn(embedsize, vocabsize)))
(embedding::Embedding)(x) = embedding.W[:, x]
Flux.treelike(Embedding)

# struct LayerNorm
#     γ
#     β
# end
# LayerNorm(nchannels) = LayerNorm(param(ones(nchannels)), param(zeros(nchannels)))
# (l::LayerNorm)(x) = l.γ .* (x .- mean(x, 1)) ./ (std(x, 1) .+ ϵ) .+ l.β

FeedForward(d_model, d_hidden) = Chain(Dense(d_model, d_hidden, relu), Dense(d_hidden, d_model))

struct ResidualBlock
    l
    norm
    # TODO dropout
end
ResidualBlock(l, d_model::Integer) = ResidualBlock(l, LayerNorm(d_model))
(l::ResidualBlock)(x...) = x[1] .+ l.norm(l.l(x...))
Flux.treelike(ResidualBlock)

struct Attention
    scale
    # TODO dropout
    # TODO causal
end
Attention(d_key::Integer, causal=false) = Attention(sqrt(d_key))
function (l::Attention)(q, k, v)
    alpha = k'*q
    # TODO causal mask
    return v*softmax(alpha, 1)
end
Flux.mapchildren(f, l::Attention) = Attention(l.scale)

# TODO multihead

struct EncoderLayer
    selfattn
    feedforward
end
EncoderLayer(d_model::Integer, d_hidden::Integer) = EncoderLayer(
    ResidualBlock(Attention(d_model), d_model),
    ResidualBlock(FeedForward(d_model, d_hidden), d_model))
(l::EncoderLayer)(x) = l.feedforward(l.selfattn(x, x, x))
Flux.treelike(EncoderLayer)

struct Encoder
    embed
    layers
end
Encoder(d_model, d_hidden, n_layers, vocabsize) = Encoder(
    Embedding(vocabsize, d_model),
    Chain((EncoderLayer(d_model, d_hidden) for _ in 1:n_layers)...))
function (l::Encoder)(x)
    x = posenc!(l.embed(x))
    encoding = [x = l(x) for l in l.layers.layers]
    return encoding
end
Flux.treelike(Encoder)

#corpus = zip(english, spanish)
en, vocab_en = process(english)
es, vocab_es = process(spanish)

# d = 4
# embed = Embedding(length(vocab_en), d)
# layernorm = LayerNorm(d)
# linear = Dense(d, d)
# feedforward = FeedForward(d, 2d)
# residual = ResidualBlock(FeedForward(d, 2d), d)
#
# x = en[1]
# x = embed(x)
# x = posenc!(x)
# x = layernorm(x)
# x = linear(x)
# x = feedforward(x)
# x = residual(x)

encoder = Encoder(4, 8, 3, length(vocab_en))
encoder = Flux.mapleaves(Flux.Tracker.data, encoder)
encoder(en[1])


using DataStructures
params = OrderedDict(
    "VISCR"  => 1:4:9,
    "NSUB"   => 1:3:7,
    "REMULT" => 0.25:0.25:1.5
)
static_params = Dict(
    "USE_GPU"     => true,
    "GPU_ID"      => 1,
    "DO_VIZ"      => false,
    "DO_SAVE"     => true,
    "DO_SAVE_VIZ" => true,
    "RESOL"       => 2047,
    "EBG"         => 1.0,
    "RHOG0"       => 0.0,
    "RHOGI"       => 0.0,
    "SIMNAME"     => "shear_weak"
)
par_names = Iterators.flatten([typeof(par)<:Tuple ? [par...] : [par] for par ∈ keys(params)])
for par in Iterators.product(values(params)...)
    par_values = Iterators.flatten(par)
    println(collect(par_values))
    for (p,v) ∈ zip(par_names,par_values) ENV[p] = v end
    for (p,v) ∈ static_params ENV[p] = v end
    run(`julia --project -O3 --check-bounds=no Stokes2D_param.jl`)
end


function ParAmulB!(nthreads::Int64, yt::Array{Float64,1}, y::Array{Float64,1}, A::SparseMatrixCSC{Float64,Int64}, x::Array{Float64,1})
    (m,n) = size(A)
    length(yt) == nthreads*m || throw(DimensionMismatch("yt"))
    if length(x) != n || length(y) != m
       throw(DimensionMismatch("length(x) != n || length(y) != m"))
    end
    p = ccall((:a_mul_b_prr_, spmatveclib), Int64, (Ptr{Int64}, Ptr{Int64}, Ptr{Int64}, Ptr{Float64}, Ptr{Int64}, Ptr{Int64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}),
                                                    &nthreads,  &m,         &n,         A.nzval,      A.rowval,   A.colptr,   x,            y,            yt           )
    return nothing
end

function AmulB!(y::Array{Float64,1}, A::SparseMatrixCSC{Float64,Int64}, x::Array{Float64,1})
    (m,n) = size(A)
    if length(x) != n || length(y) != m
       throw(DimensionMismatch("length(x) != n || length(y) != m"))
    end
    p = ccall((:a_mul_b_rr_, spmatveclib), Int64, (Ptr{Int64}, Ptr{Int64}, Ptr{Float64}, Ptr{Int64}, Ptr{Int64}, Ptr{Float64}, Ptr{Float64}),
                                                   &m,         &n,         A.nzval,      A.rowval,   A.colptr,   x,            y             )
    return nothing
end

function AmulB!(y::Array{Complex128,1}, A::SparseMatrixCSC{Float64,Int64}, x::Array{Complex128,1})
    (m,n) = size(A)
    if length(x) != n || length(y) != m
       throw(DimensionMismatch("length(x) != n || length(y) != m"))
    end
    p = ccall((:a_mul_b_rc_, spmatveclib), Int64, (Ptr{Int64}, Ptr{Int64}, Ptr{Float64}, Ptr{Int64}, Ptr{Int64}, Ptr{Complex128}, Ptr{Complex128}),
                                                   &m,         &n,         A.nzval,      A.rowval,   A.colptr,   x,               y                )
    return nothing
end

function AmulB!(y::Array{Complex128,1}, A::SparseMatrixCSC{Complex128,Int64}, x::Array{Complex128,1})
    (m,n) = size(A)
    if length(x) != n || length(y) != m
       throw(DimensionMismatch("length(x) != n || length(y) != m"))
    end
    p = ccall((:a_mul_b_cc_, spmatveclib), Int64, (Ptr{Int64}, Ptr{Int64}, Ptr{Complex128}, Ptr{Int64}, Ptr{Int64}, Ptr{Complex128}, Ptr{Complex128}),
                                                      &m,         &n,         A.nzval,         A.rowval,   A.colptr,   x,               y                )
    return nothing
end


function read_VolumeImage(fname::AbstractString)

    Logging.info("Creating volume image from file $fname")

    if contains(fname, ".dat")
        x, y, z, s, t = read_dat(fname)
        method = "CLARA"
        header = Dict()
        units = "nAm/cm^3"
        x = x/1000
        y = y/1000
        z = z/1000
        t = t/1000
    else
        Logging.warn("Unknown file type")
    end

    header["FileName"] = fname

    coord_system = "?"

    VolumeImage(s, units, collect(x), collect(y), collect(z), collect(t), method, header, coord_system)
end