import Base.Threads
import LinearAlgebra

"""
    get_frame(frame_idx::Int64, trajectory::Trajectory)
    ::Frame

Return Frame object which correspond to `framd_idx` frame from `trajectory`.
"""
function get_frame(frame_idx::IntT, trj::Trajectory{RealT}
    )::Frame where {IntT <: Integer, RealT <: Real}

    Frame(trj.coordinates[:,frame_idx], trj.attributes)
end

"""
    get_atom(query::Int64, trajectory::Trajectory)
    ::Vector{Atom}

Return Vector of `query`th Atom object from `trajectory`.
"""
function get_atom(query::IntT, trj::Trajectory)::Vector{Atom} where IntT <: Integer
    map(xyz -> Atom(xyz, trj.attributes[query]), trj.coordinates[query, :])
end

"""
    get_atom(query::Int64, frame::Frame)
    ::Atom

Return `query`th Atom object from `frame`.
"""
function get_atom(query::IntT, frame::Frame)::Atom where IntT <: Integer
    Atom(frame.coordinates[query], frame.attributes[query])
end

"""
    clip_trajectory(query::Union{Integer, Vector{IntT}, OrdinalRange}, trj::Trajectory;
                    query_key::Symbol = :frame)
    ::Trajecotory where IntT <: Integer

Clip a part of trajectory you specified.
Expected query key is one of the following

- `:frame` : in this case, query specify the frame index.
- `:atom`  : in this case, query specify the atom index.
"""
function clip_trajectory(query::IntT, trj::Trajectory{RealT};
    query_key = query_key::Symbol = :frame
    )::Trajectory{RealT} where {IntT <: Integer, RealT <: Real}

    if query_key == :frame
        Trajectory(reshape(trj.coordinates[:, query], (trj.natom, 1)), trj.attributes)
    elseif query_key == :atom
        Trajectory(reshape(trj.coordinates[query, :], (1, trj.nframe)), [trj.attributes[query]])
    else
        throw(ArgumentError("""
                            clip_trajectory: invalid query key :$(query_key) here
                            expected query key is one of the following.
                            - :frame
                            - :atom
                            """))
    end
end

function clip_trajectory(query::Union{Vector{IntT}, OrdinalRange}, trj::Trajectory{RealT};
    query_key::Symbol = :frame
    )::Trajectory{RealT} where {IntT <: Integer, RealT <: Real}

    if query_key == :frame
        Trajectory(trj.coordinates[:, query], trj.attributes)
    elseif query_key == :atom
        Trajectory(trj.coordinates[query, :], trj.attributes[query])
    else
        throw(ArgumentError("""
                            clip_trajectory: invalid query key :$(query_key) here
                            expected query key is one of the following.
                            - :frame
                            - :atom
                            """))
    end
end

"""
    geometric_center_of_mass(coordinates::Vector{Coordinate};
                             atom_indices::Union{Vector, OrdinalRange, Colon} = :)
    ::Coordinate

Calculate the geometric center of mass of Coordinate vector for specified atom indices.

> For trajectory case, geometric center of mass is calculated by `center_of_mass` method with geometric flag.
"""
function geometric_center_of_mass(coordinates::Vector{<:Coordinate{RealT}};
    atom_indices::Union{Vector, OrdinalRange, Colon} = :)::Coordinate{RealT} where RealT <: Real

    selected_coord = view(coordinates, atom_indices)
    sum(selected_coord) / length(selected_coord)
end

"""
    center_of_mass(coordinates::Vector{Coordiante}, mass_vec{Real};
                   atom_indices::Union{Vector, OrdinalRange, Colon} = :)
    ::Coordinate{RealT}

Calculate the (non-geometric) center of mass of Coordinate vector for specified atom indices.
You have to pass mass information vector which have same length to coordinates vector.
"""
function center_of_mass(coordinates::Vector{<:Coordinate{RealT}}, mass_vec::Vector{<:Real};
    atom_indices::Union{Vector, OrdinalRange, Colon} = :)::Coordinate{RealT} where RealT <: Real

    if length(coordinates) != length(mass_vec)
        throw(AssertionError("""
                             coordinate vector and mass vector should be same length.
                             """))
    end

    selected_coord = view(coordinates, atom_indices)
    selected_mass  = view(mass_vec,    atom_indices)
    sum(selected_coord .* selected_mass) / sum(selected_mass)
end

"""
    center_of_mass(query::Trajectory;
                   frame_indices::Union{Vector, OrdinalRange, Colon} = :,
                   atom_indices::Union{Vector, OrdinalRange, Colon} = :,
                   geometric::Bool = false)
    ::Vector{Coordinate{Real}}

Calculate the center of mass of trajectory for specified frame and atom indices.
If you set `geometric` is `true`, this function calculate geometric center of mass.
"""
function center_of_mass(trj::Trajectory{RealT};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    geometric::Bool = false)::Vector{Coordinate{RealT}} where RealT <: Real

    if !geometric
        mass_vec = map(attributes -> attributes.mass, trj.attributes)
        selected_mass_vec = view(mass_vec, atom_indices)
        if nothing ∈ selected_mass_vec || length(selected_mass_vec) == 0
            throw(ArgumentError("""
                                center_of_mass: geometric flag is false but attributes of trajectory have Nothing value.
                                you can not use mass to calculation.
                                """))
        end
        total_mass = sum(selected_mass_vec)
        weited_sum = reshape(selected_mass_vec, (1, length(selected_mass_vec))) * view(trj.coordinates, atom_indices, frame_indices)
        vec(weited_sum / total_mass)
    else
        selected_coord = view(trj.coordinates, atom_indices, frame_indices)
        vec(sum(selected_coord, dims = 1) / size(selected_coord, 1))
    end
end

"""
    pair_length_matrix(corrds1::Vector{Coordinate}, coords2::Vector{Coordinate})
    ::Matrix{Real}

Calculate distance matrix for all combination between `coords1` and `coord2`.
The row of returned matrix coresspond to `coords1`, and the column correspond to `coords2`.
"""
function pair_length_matrix(first_coords::ArrayT1, second_coords::ArrayT2
    )::Matrix{<:Real} where {ArrayT1 <: AbstractArray{<:Coordinate{<:Real}, 1},
                             ArrayT2 <: AbstractArray{<:Coordinate{<:Real}, 1}}

    distance.(first_coords, reshape(second_coords, (1, length(second_coords))))
end

"""
     pair_length_matrix_pbc(corrds1::Vector{Coordinate}, coords2::Vector{Coordinate},
                            box_size::Coordinate)
    ::Matrix{Real}

Calculate distance matrix for all combination between `coords1` and `coord2` considering periodic boundary condition.
The box size information used for calculation is specified by `box_size` argument.
The row of returned matrix coresspond to `coords1`, and the column correspond to `coords2`.
"""
function pair_length_matrix_pbc(first_coords::ArrayT1, second_coords::ArrayT2,
    box_size::Coordinate{<:Real}
    )::Matrix{<:Real} where {ArrayT1 <: AbstractArray{<:Coordinate{<:Real}, 1},
                             ArrayT2 <: AbstractArray{<:Coordinate{<:Real}, 1}}

    distance_pbc.(first_coords, reshape(second_coords, (1, length(second_coords))), box_size)
end

"""
    pair_length_matrix(trj::Trajectory;
                       frame_indices::Union{Vector, OrdinalRange, Colon}       = :,
                       first_atom_indices::Union{Vector, OrdinalRange, Colon}  = :,
                       second_atom_indices::Union{Vector, OrdinalRange, Colon} = atom_indices1)
    ::Vector{Matrix{Coordinate}}

Calculate distance matrix for all combination between `first_atom_indices` and `second_atom_indices`.
This cauculation apply to each frame of trajectory and the result matrices are stored in Vector.
The target frame can be restricted by pass indeces vector or range to `frame_indices`.
"""
function pair_length_matrix(trj::Trajectory{<:Real};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Vector{<:Matrix{<:Real}}

    zip_iterate4frame = zip(eachcol(view(trj.coordinates, first_atom_indices, frame_indices)),
                            eachcol(view(trj.coordinates, second_atom_indices, frame_indices)))
    map(coords_vec_pair -> pair_length_matrix(coords_vec_pair...), zip_iterate4frame)
end

"""
    pair_length_matrix_pbc(trj::Trajectory, box_size::Coordinate;
                           frame_indices::Union{Vector, OrdinalRange, Colon}       = :,
                           first_atom_indices::Union{Vector, OrdinalRange, Colon}  = :,
                           second_atom_indices::Union{Vector, OrdinalRange, Colon} = atom_indices1)
    ::Vector{Matrix{Coordinate}}

Calculate distance matrix for all combinations between `first_atom_indices` and `second_atom_indices` considering periodic boundary condition.
The box size information used for calculation is specified by `box_size` argument.
This cauculation apply to each frame of trajectory and the result matrices are stored in Vector.
The target frame can be restricted by pass indeces vector or range to `frame_indices`.
"""
function pair_length_matrix_pbc(trj::Trajectory{<:Real}, box_size::Coordinate{<:Real};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Vector{<:Matrix{<:Real}}

    zip_iterate4frame = zip(eachcol(view(trj.coordinates, first_atom_indices, frame_indices)),
                            eachcol(view(trj.coordinates, second_atom_indices, frame_indices)))
    map(coords_vec_pair -> pair_length_matrix_pbc(coords_vec_pair..., box_size),
        zip_iterate4frame)
end

"""
    pair_length_matrix_parallel(trj::Trajectory;
                                frame_indices::Union{Vector, OrdinalRange, Colon}       = :,
                                first_atom_indices::Union{Vector, OrdinalRange, Colon}  = :,
                                second_atom_indices::Union{Vector, OrdinalRange, Colon} = atom_indices1)
    ::Vector{Matrix{Coordinate}}

Multi-threads version of pair_length_matrix. If you set available threads number to `Threads.nthreads()`, this function would faster than non-parallel version.
"""
function pair_length_matrix_parallel(trj::Trajectory{RealT};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Vector{Matrix{RealT}} where RealT <: Real

    first_target_coords  = view(trj.coordinates, first_atom_indices,  frame_indices)
    second_target_coords = view(trj.coordinates, second_atom_indices, frame_indices)
    target_frame_num = size(first_target_coords, 2)
    return_vec = Vector{Matrix{Float32}}(undef, target_frame_num)
    Threads.@threads for frame_idx in 1:target_frame_num
        return_vec[frame_idx] =
            pair_length_matrix(view(first_target_coords, :, frame_idx), view(second_target_coords, :, frame_idx))
    end
    return_vec
end

"""
    contact_bool_matrix(threshold::Real, trj::Trajectory;
                        frame_indices::Union{Array, OrdinalRange, Colon}       = :,
                        first_atom_indices::Union{Array, OrdinalRange, Colon}  = :,
                        second_atom_indices::Union{Array, OrdinalRange, Colon} = first_atom_indices)
    ::Vector{Matrix{Bool}}

Judge contact is formed or not. If the distance between two coordinate is shorter than threshold, contact is considered to be formed. In returned vector of matrices, each matrix correspond to contact matrix of each frame.
You can specify the target frames or atoms by `frame_indices`, `first_atom_indices` or `second_atom_indices`. When you specify the target atoms, the row of matrices corresponds to first_atom_indices and column of matrices corresponds to second_atom_indices.
"""
function contact_bool_matrix(threshold::RealT,
    first_coords::AbstractArray{<:Coordinate, 1},
    second_coords::AbstractArray{<:Coordinate, 1}
    )::Matrix{Bool} where RealT <: Real

    length_matrix = pair_length_matrix(first_coords, second_coords)
    map(length -> length < threshold, length_matrix)
end

function contact_bool_matrix(threshold::RealT1, trj::Trajectory{RealT2};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Vector{Matrix{Bool}} where {RealT1 <: Real, RealT2 <: Real}

    length_mat_arr = pair_length_matrix(trj, frame_indices = frame_indices,
                                        first_atom_indices = first_atom_indices,
                                        second_atom_indices = second_atom_indices)
    map(length_mat_arr) do length_matrix
        map(length -> length < threshold, length_matrix)
    end
end

"""
    contact_bool_matrix_pbc(threshold::Real, trj::Trajectory, box_size::Coordinate;
                            frame_indices::Union{Array, OrdinalRange, Colon}       = :,
                            first_atom_indices::Union{Array, OrdinalRange, Colon}  = :,
                            second_atom_indices::Union{Array, OrdinalRange, Colon} = first_atom_indices)
    ::Vector{Matrix{Bool}}

Judge contact is formed or not considering periodic boundary condition. If the distance between two coordinate is shorter than threshold, contact is considered to be formed. In returned vector of matrices, each matrix correspond to contact matrix of each frame.
The box size information used for calculation is specified `box_size` argument.
You can specify the target frames or atoms by `frame_indices`, `first_atom_indices` or `second_atom_indices`. When you specify the target atoms, the row of matrices corresponds to first_atom_indices and column of matrices corresponds to second_atom_indices.
"""
function contact_bool_matrix_pbc(threshold::RealT, trj::Trajectory{<:Real},
    box_size::Coordinate{<:Real};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Vector{Matrix{Bool}} where RealT <: Real

    length_mat_arr = pair_length_matrix_pbc(trj, box_size,
                                            frame_indices = frame_indices,
                                            first_atom_indices = first_atom_indices,
                                            second_atom_indices = second_atom_indices)
    map(length_mat_arr) do length_matrix
        map(length -> length < threshold, length_matrix)
    end
end

"""
    contact_bool_matrix_parallel(threshold::Real, trj::Trajectory;
                                 frame_indices::Union{Vector, OrdinalRange, Colon}       = :,
                                 first_atom_indices::Union{Vector, OrdinalRange, Colon}  = :,
                                 second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices)
    ::Vector{Matrix{Bool}}

Multi-threads version of contact_bool_matrix. If you set available threads number to `Threads.nthreads()`, this function would faster than non-parallel version.
"""
function contact_bool_matrix_parallel(threshold::RealT1, trj::Trajectory{RealT2};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Vector{Matrix{Bool}} where {RealT1 <: Real, RealT2 <: Real}

    length_mat_arr = pair_length_matrix_parallel(trj, frame_indices = frame_indices,
                                                 first_atom_indices = first_atom_indices,
                                                 second_atom_indices = second_atom_indices)
    target_frame_num = length(length_mat_arr)
    return_vec = Vector{Matrix{Float32}}(undef, target_frame_num)
    threads_num = Threads.nthreads()
    iteration_in_thread = Int32(floor(target_frame_num / threads_num))
    Threads.@threads for frame_idx in 1:target_frame_num
        return_vec[frame_idx] = map(length -> length < threshold, length_mat_arr[frame_idx])
    end
    return_vec
end

"""
    contact_count_matrix(threshold::Real, trj::Trajectory;
                         frame_indices::Union{Vector, OrdinalRange, Colon}       = :,
                         first_atom_indices::Union{Vector, OrdinalRange, Colon}  = :,
                         second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices)
    ::Matrix{Real}

Calculate contact formation count over the trajectory. If the distance between two coordinate is shorter than threshold, contact is considered to be formed.
You can specify the target frames or atoms by `frame_indices`, `first_atom_indices` or `second_atom_indices`. When you specify the target atoms, the row of matrices corresponds to first_atom_indices and column of matrices corresponds to second_atom_indices.
"""
function contact_count_matrix(threshold::RealT1, trj::Trajectory{RealT2};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Matrix{RealT2} where {RealT1 <: Real, RealT2 <: Real}

    first_target_coords  = view(trj.coordinates, first_atom_indices,  frame_indices)
    second_target_coords = view(trj.coordinates, second_atom_indices, frame_indices)
    target_frame_num = size(first_target_coords, 2)
    count_mat = zeros(RealT2, (size(first_target_coords, 1), size(second_target_coords, 1)))
    for frame_idx in 1:target_frame_num
        count_mat +=
            contact_bool_matrix(threshold, view(first_target_coords, :, frame_idx), view(second_target_coords, :, frame_idx))
    end
    count_mat
end

"""
    contact_count_matrix_parallel(threshold::Real, trj::Trajectory;
                                  frame_indices::Union{Vector, OrdinalRange, Colon}       = :,
                                  first_atom_indices::Union{Vector, OrdinalRange, Colon}  = :,
                                  second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices)
    ::Matrix{Real}

Multi-threads version of contact_count_matrix. If you set available threads number to `Threads.nthreads()`, this function would faster than non-parallel version.
"""
function contact_count_matrix_parallel(threshold::RealT1, trj::Trajectory{RealT2};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Matrix{RealT2} where {RealT1 <: Real, RealT2 <: Real}

    first_target_coords  = view(trj.coordinates, first_atom_indices,  frame_indices)
    second_target_coords = view(trj.coordinates, second_atom_indices, frame_indices)
    target_frame_num = size(first_target_coords, 2)
    first_atoms_len = size(first_target_coords, 1)
    second_atoms_len = size(second_target_coords, 1)
    count_mat_arr =
        [zeros(RealT2, first_atoms_len, second_atoms_len) for thread in 1:Threads.nthreads()]
    Threads.@threads for frame_idx in 1:target_frame_num
        count_mat_arr[Threads.threadid()] +=
            contact_bool_matrix(threshold,
                                view(first_target_coords, :, frame_idx), view(second_target_coords, :, frame_idx))
    end
    sum(count_mat_arr)
end

"""
    contact_probability_matrix(threshold::Real, trj::Trajectory;
                               frame_indices::Union{Vector, OrdinalRange, Colon}       = :,
                               first_atom_indices::Union{Vector, OrdinalRange, Colon}  = :,
                               second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices)
    ::Matrix{Real}

Calculate contact formation probability over the trajectory. If the distance between two coordinate is shorter than threshold, contact is considered to be formed.
You can specify the target frames or atoms by `frame_indices`, `first_atom_indices` or `second_atom_indices`. When you specify the target atoms, the row of matrices corresponds to first_atom_indices and column of matrices corresponds to second_atom_indices.
"""
function contact_probability_matrix(threshold::RealT1, trj::Trajectory{RealT2};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Matrix{RealT2} where {RealT1 <: Real, RealT2 <: Real}

    count_mat= contact_count_matrix(threshold, trj,
                                    frame_indices       = frame_indices,
                                    first_atom_indices  = first_atom_indices,
                                    second_atom_indices = second_atom_indices)
    count_mat / size(view(trj.coordinates, :, frame_indices), 2)
end

"""
    contact_probability_matrix_parallel(threshold::Real, trj::Trajectory;
                                        frame_indices::Union{Vector, OrdinalRange, Colon}       = :,
                                        first_atom_indices::Union{Vector, OrdinalRange, Colon}  = :,
                                        second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices)
    ::Matrix{Real}

Multi-threads version of contact_probability_matrix. If you set available threads number to `Threads.nthreads()`, this function would faster than non-parallel version.
"""
function contact_probability_matrix_parallel(threshold::RealT1, trj::Trajectory{RealT2};
    frame_indices::Union{Vector, OrdinalRange, Colon} = :,
    first_atom_indices::Union{Vector, OrdinalRange, Colon} = :,
    second_atom_indices::Union{Vector, OrdinalRange, Colon} = first_atom_indices
    )::Matrix{RealT2} where {RealT1 <: Real, RealT2 <: Real}

    count_mat_arr = contact_count_matrix_parallel(threshold, trj,
                                                  frame_indices       = frame_indices,
                                                  first_atom_indices  = first_atom_indices,
                                                  second_atom_indices = second_atom_indices)
    count_mat_arr / size(view(trj.coordinates, :,frame_indices), 2)
end

"""
    radius_of_gyration(trj::Trajectory;
                       frame_indices::Union{Vector, OrdinalRange, Colon} = :,
                       atom_indices ::Union{Vector, OrdinalRange, Colon} = :,
                       geometric::Bool = false)
    ::Vector{Real}

Calculate radius of gyration over the trajectory.
If you set `geometric` is `true`, this function calculate radius of gyration without mass info.
"""
function radius_of_gyration(trj::Trajectory{RealT};
             frame_indices::Union{Vector, OrdinalRange, Colon} = :,
             atom_indices ::Union{Vector, OrdinalRange, Colon} = :,
             geometric::Bool = false
             )::Vector{RealT} where RealT <: Real

    com = center_of_mass(trj,
                         frame_indices = frame_indices,
                         atom_indices  = atom_indices,
                         geometric     = geometric)

    target_coordinates = view(trj.coordinates, atom_indices, frame_indices)
    target_frame_num = length(com)
    return_vec = Vector{RealT}(undef, target_frame_num)

    dist_from_com_vec = norm.(target_coordinates .- permutedims(com))

    if !geometric
        mass_vec = map(attributes -> attributes.mass, view(trj.attributes, atom_indices))
        mass_sum = sum(mass_vec)
        for frame_idx in 1:target_frame_num
            return_vec[frame_idx] = sqrt(sum(view(dist_from_com_vec, :, frame_idx).^2 .* mass_vec) / mass_sum)
        end
    else target_atom_num = size(target_coordinates)[1]
        for frame_idx in 1:target_frame_num
            return_vec[frame_idx] = sqrt(sum(view(dist_from_com_vec, :, frame_idx).^2) / target_atom_num)
        end
    end
    return_vec
end

"""
    atom_mass(name::AbstractString)
    ::Float32

Return mass of the `name` atom.
The kind of each atom is judged by PDB format atom names rule.
https://cdn.rcsb.org/wwpdb/docs/documentation/file-format/PDB_format_1992.pdf
"""
function atom_mass(name::AbstractString)::Float32
    if occursin(r"^(H|[1-4]H)", name)
        ATOMNAME2MASS["H"]
    elseif occursin(r"^C($|[^L]+)", name)
        ATOMNAME2MASS["C"]
    elseif occursin(r"^O", name)
        ATOMNAME2MASS["O"]
    elseif occursin(r"^N($|[^A]+)", name)
        ATOMNAME2MASS["N"]
    elseif occursin(r"^S", name)
        ATOMNAME2MASS["S"]
    elseif occursin(r"^P", name)
        ATOMNAME2MASS["P"]
    else
        ATOMNAME2MASS[name]
    end
end

"""
    residue_mass(name::AbstractString)
    ::Float32

Return mass of the `name` residue.
The kind of each residue is judged by amino acid 3-letter abbreviation.
"""
function residue_mass(name::AbstractString)::Float32
    RESNAME2MASS[name]
end

"""
    distance_pdc(first_atom::Coordinate,  second_atom::Coordinate,
                 box_size::Coordinate)
    ::Real

Calculate distance `first_atom` and `second_atom` considering periodic boundary condition.
The box size information used calculatin is specified by `box_size` argument.
"""
function distance_pbc(first_coord::Coordinate{RealT},  second_coord::Coordinate{RealT},
    box_size::Coordinate{<:Real}
    )::RealT where RealT <: Real

    half_box_size = box_size * 0.5
    dist_vec = first_coord  - second_coord
    x = abs(dist_vec.x) < half_box_size.x ? dist_vec.x : box_size.x - abs(dist_vec.x)
    y = abs(dist_vec.y) < half_box_size.y ? dist_vec.y : box_size.y - abs(dist_vec.y)
    z = abs(dist_vec.z) < half_box_size.z ? dist_vec.z : box_size.z - abs(dist_vec.z)
    sqrt(x^2 + y^2 + z^2)
end

# This function get half_box_size, because if process will apply to all trajectory,
# half box size calculation will occur every time frame, this is overhead.
function fix_pbc(coordinates::Vector{<:Coordinate{RealT}}, groupid_vec::Vector{<:Integer},
    box_size::Coordinate{<:Real}, half_box_size::Coordinate{<:Real};
    x::Bool= true, y::Bool = true, z::Bool = true
    )::Vector{<:Coordinate{RealT}} where RealT <: Real

    if length(coordinates) != length(groupid_vec)
        throw(AssertionError("""
                             coordinate vector and groupid vector should be same length.
                             """))
    end

    new_coords = deepcopy(coordinates)
    unique_groupid_vec = unique(groupid_vec)

    for groupid in unique_groupid_vec
        same_group_ids = findall(id->id==groupid, groupid_vec)
        @views sbj_coords = new_coords[same_group_ids[2:end]]
        @views dist2first = sbj_coords .- coordinates[same_group_ids[1]]
        sbj_dist_zip      = zip(sbj_coords, dist2first)

        if x
            for (coord, dist) in sbj_dist_zip
                coord.x = abs(dist.x) < half_box_size.x ? coord.x : coord.x - sign(dist.x) * box_size.x
            end
        end

        if y
            for (coord, dist) in sbj_dist_zip
                coord.y = abs(dist.y) < half_box_size.y ? coord.y : coord.y - sign(dist.y) * box_size.y
            end
        end

        if z
            for (coord, dist) in sbj_dist_zip
                coord.z = abs(dist.z) < half_box_size.z ? coord.z : coord.z - sign(dist.z) * box_size.z
            end
        end
    end
    new_coords
end

"""
    fix_pbc(coordinates::Vector{Coordinate}, groupid_vec::Vector{Integer},
    box_size::Coordinate{<:Real};
    x::Bool = true, y::Bool = true, z::Bool = true)::Vector{Coordinate}

Fix the atom group splited by periodic boundary box. This group is defined by `groupid_vec` which contain each particles group id.
For example, if your system have 3 atom, and atom 1 and 2 are group 1, and atom 3 is group 2, this `groupid_vec` is [1, 1, 2].
If the group over the boundary of the box, the atoms in the group which separated from first atom of that move to the position where the atom locate without periodic boundary condition.
You can specify the axies which applied fixing by x,y and option. If you set `z = false`, only x and y coordinate fixed and z don't change.
"""
function fix_pbc(coordinates::Vector{<:Coordinate{RealT}}, groupid_vec::Vector{<:Integer},
    box_size::Coordinate{<:Real};
    x::Bool = true, y::Bool = true, z::Bool = true
    )::Vector{<:Coordinate{RealT}} where RealT <: Real
    half_box_size = box_size * 0.5
    fix_pbc(coordinates, groupid_vec, box_size, half_box_size, x=x, y=y, z=z)
end

"""
    fix_pbc(trj::Trajectory, groupid_vec::Vector{Integer},
    box_size::Coordinate;
    x::Bool = true, y::Bool = true, z::Bool = true)::Trajectory

Fix residues splited by periodic boundary condition. This is more specific version of `fix_pbc` function for trajectory handling.
For example, if your system have 3 atom, and atom 1 and 2 are group 1, and atom 3 is group 2, this `groupid_vec` is [1, 1, 2].
If the group over the boundary of the box, the atoms in the group which separated from first atom of that move to the position where the atom locate without periodic boundary condition.
You can specify the axies which applied fixing by x,y and option. If you set `z = false`, only x and y coordinate fixed and z don't change.
"""
function fix_pbc(trj::Trajectory{RealT}, groupid_vec::Vector{<:Integer},
    box_size::Coordinate{<:Real};
    x::Bool = true, y::Bool = true, z::Bool = true
    )::Trajectory{RealT} where RealT <: Real

    new_trj = deepcopy(trj)
    half_box_size = box_size * 0.5

    unique_groupid_vec = unique(groupid_vec)
    coordinates        = new_trj.coordinates
    for groupid in unique_groupid_vec
        same_mol_indices = findall(id->id==groupid, groupid_vec)
        @views sbj_coords = coordinates[same_mol_indices[2:end], :]
        @views dist2first_mat = sbj_coords .- permutedims(coordinates[same_mol_indices[1], :])
        if x
            for (coord, dist) in zip(sbj_coords, dist2first_mat)
                coord.x = abs(dist.x) < half_box_size.x ? coord.x : coord.x - sign(dist.x) * box_size.x
            end
        end

        if y
            for (coord, dist) in zip(sbj_coords, dist2first_mat)
                coord.y = abs(dist.y) < half_box_size.y ? coord.y : coord.y - sign(dist.y) * box_size.y
            end
        end

        if z
            for (coord, dist) in zip(sbj_coords, dist2first_mat)
                coord.z = abs(dist.z) < half_box_size.z ? coord.z : coord.z - sign(dist.z) * box_size.z
            end
        end
    end
    new_trj
end

function fix_pbc!(trj::Trajectory{RealT}, box_size::Coordinate{<:Real};
    x::Bool = true, y::Bool = true, z::Bool = true
    ) where RealT <: Real

    # fix atom in over pbc residue based on first atom of the residue
    half_box_size = box_size * 0.5

    resid_vec  = map(attr->attr.resid, trj.attributes)
    unique_resid_vec = unique(resid_vec)
    coordinates = trj.coordinates
    attributes  = trj.attributes
    for resid in unique_resid_vec
        same_mol_indices = findall(id->id==resid, resid_vec)
        @views sbj_coords = coordinates[same_mol_indices[2:end], :]
        @views dist2first_mat = sbj_coords .- permutedims(coordinates[same_mol_indices[1], :])
        if x
            for (coord, dist) in zip(sbj_coords, dist2first_mat)
                coord.x = abs(dist.x) < half_box_size.x ? coord.x : coord.x - sign(dist.x) * box_size.x
            end
        end
        if y
            for (coord, dist) in zip(sbj_coords, dist2first_mat)
                coord.y = abs(dist.y) < half_box_size.y ? coord.y : coord.y - sign(dist.y) * box_size.y
            end
        end
        if z
            for (coord, dist) in zip(sbj_coords, dist2first_mat)
                coord.z = abs(dist.z) < half_box_size.z ? coord.z : coord.z - sign(dist.z) * box_size.z
            end
        end
    end
end

"""
    move_pbc_center(coordinates::Vector{Coordinate},
                    new_center::Coordinate, box_size::Coordinate)
                    )::Vector{Coordinate}
Fix coordinate to consistent with new box center and box size.
Box size should match to original periodic boundary box size.
"""
function move_pbc_center(coordinates::Vector{<:Coordinate{RealT}},
    center::Coordinate{<:Real}, box_size::Coordinate{<:Real}
    )::Vector{Coordinate{RealT}} where RealT <: Real

    new_coords = deepcopy(coordinates)
    half_box = box_size * 0.5
    for coord in new_coords
        dist = coord - center
        coord.x = abs(dist.x) < half_box.x ? coord.x : coord.x - sign(dist.x) * box_size.x
        coord.y = abs(dist.y) < half_box.y ? coord.y : coord.y - sign(dist.y) * box_size.y
        coord.z = abs(dist.z) < half_box.z ? coord.z : coord.z - sign(dist.z) * box_size.z
    end
    new_coords
end


using Cxx

# Adding path to tetgen folder
const path_to_tetgen = "/home/olli/programming/forks/tetgen_src"

# Adding source files
tetgen_source = "/home/olli/programming/forks/tetgen_src/tetgen.cxx"
predicates_source = "/home/olli/programming/forks/tetgen_src/predicates.cxx"

# Adding header dir path
addHeaderDir(path_to_tetgen, kind=C_System)

# include source files
cxxinclude(predicates_source)
cxxinclude(tetgen_source)

# Creating tetgen objects
b_ = @cxxnew tetgenbehavior()
in_ = @cxxnew tetgenio()
out = @cxxnew tetgenio()
addin_ = @cxxnew tetgenio()
bgmin_ = @cxxnew tetgenio()

# including polynomial file
# It's necessary to procide 2 arguments for parse_commandline. "test" is only for that purpose
icxx"""
char *file_name[] = {"test", "example.poly"};
$b_->parse_commandline(2, file_name);
"""

# checking input
icxx"""
if ($b_->refine) {
    if (!$in_->load_tetmesh($b_->infilename)) {
      std::cout << "Something went wrong in refine section?"<< std::endl;
    }
} else {
if (!$in_->load_plc($b_->infilename, (int) $b_->object)) {
      std::cout << "Something went wrong in load plc?"<< std::endl;
    }
}
"""

# Adding points
icxx"""
  if ($b_->insertaddpoints) {
    if (!$addin_->load_node($b_->addinfilename)) {
      $addin_->numberofpoints = 0l;
    }
  }
  if ($b_->metric) {
    if (!$bgmin_->load_tetmesh($b_->bgmeshfilename)) {
      $bgmin_->numberoftetrahedra = 0l;
    }
  }

  if ($bgmin_->numberoftetrahedra > 0l) {
    tetrahedralize($b_, $in_, NULL, $addin_, $bgmin_);
  } else {
    tetrahedralize($b_, $in_, NULL, $addin_, NULL);
  }
"""


using Oceananigans.Operators: Δx, Δy, ΔzC, div_xyᶜᶜᵃ, div_xzᶜᵃᶜ, div_yzᵃᶜᶜ

#####
##### Outer functions for setting velocity on boundary and filling halo beyond boundary.
#####

@inline function fill_west_halo!(u, bc::NFBC, arch, grid, clock, state)
    @launch(device(arch), config=launch_config(grid, :yz), set_east_west_u_velocity!(u, 1, bc, grid, clock, state))
    #@launch(device(arch), config=launch_config(grid, :yz), _fill_west_halo!(u, bc, grid, state.velocities.v, state.velocities.w))
    return nothing
end

@inline function fill_east_halo!(u, bc::NFBC, arch, grid, clock, state)
    @launch(device(arch), config=launch_config(grid, :yz), set_east_west_u_velocity!(u, grid.Nx + 1, bc, grid, clock, state))
    #@launch(device(arch), config=launch_config(grid, :yz), _fill_east_halo!(u, bc, grid, state.velocities.v, state.velocities.w))
    return nothing
end

@inline function fill_south_halo!(v, bc::NFBC, arch, grid, clock, state)
    @launch(device(arch), config=launch_config(grid, :xz), set_north_south_v_velocity!(v, 1, bc, grid, clock, state))
    #@launch(device(arch), config=launch_config(grid, :xz), _fill_south_halo!(v, bc, grid, state.velocities.u, state.velocities.w))
    return nothing
end

@inline function fill_north_halo!(v, bc::NFBC, arch, grid, clock, state)
    @launch(device(arch), config=launch_config(grid, :xz), set_north_south_v_velocity!(v, grid.Ny + 1, bc, grid, clock, state))
    #@launch(device(arch), config=launch_config(grid, :xz), _fill_north_halo!(v, bc, grid, state.velocities.u, state.velocities.v))
    return nothing
end

@inline function fill_bottom_halo!(w, bc::NFBC, arch, grid, clock, state)
    @launch(device(arch), config=launch_config(grid, :xy), set_top_bottom_w_velocity!(w, 1, bc, grid, clock, state))
    #@launch(device(arch), config=launch_config(grid, :xy), _fill_bottom_halo!(w, bc, grid, state.velocities.u, state.velocities.v))
    return nothing
end

@inline function fill_top_halo!(w, bc::NFBC, arch, grid, clock, state)
    @launch(device(arch), config=launch_config(grid, :xy), set_top_bottom_w_velocity!(w, grid.Nz + 1, bc, grid, clock, state))
    #@launch(device(arch), config=launch_config(grid, :xy), _fill_top_halo!(w, bc, grid, state.velocities.u, state.velocities.v))
    return nothing
end

#####
##### Kernels that use the continuity equation to fill interior points
#####

# No penetration for x-boundaries: fill in x, loop over yz

function _fill_west_halo!(u, ::NFBC, grid, v, w)
    @loop_yz j k grid begin
        @unroll for i in 0 : -1 : (1 - grid.Hx) # integrate westward
            @inbounds u[i, j, k] = u[i+1, j, k] + Δx(0, j, k, grid) * div_yzᵃᶜᶜ(i, j, k, grid, v, w)
        end
    end
    return nothing
end

function _fill_east_halo!(u, ::NFBC, grid, v, w)
    @loop_yz j k grid begin
        @unroll for i in (grid.Nx + 2) : (grid.Nx + 1 + grid.Hx) # integrate eastward, using face spacing
            @inbounds u[i, j, k] = u[i-1, j, k] - Δx(grid.Nx+1, j, k, grid) * div_yzᵃᶜᶜ(i-1, j, k, grid, v, w)
        end
    end
    return nothing
end

# No penetration for y-boundaries: fill in y, loop over xz

function _fill_south_halo!(v, ::NFBC, grid, u, w)
    @loop_xz i k grid begin
        @unroll for j in 0 : -1 : (1 - grid.Hy) # integrate southward
            @inbounds v[i, j, k] = v[i, j+1, k] + Δy(i, 0, k, grid) * div_xzᶜᵃᶜ(i, j, k, grid, u, w)
        end
    end
    return nothing
end

function _fill_north_halo!(v, ::NFBC, grid, u, w)
    @loop_xz i k grid begin
        @unroll for j in (grid.Ny + 2) : (grid.Ny + 1 + grid.Hy) # integrate northward
            @inbounds v[i, j, k] = v[i, j-1, k] - Δy(i, grid.Ny+1, k, grid) * div_xzᶜᵃᶜ(i, j-1, k, grid, u, w)
        end
    end
    return nothing
end

# No penetration for z-boundaries: fill in z, loop over xy

function _fill_bottom_halo!(w, ::NFBC, grid, u, v)
    @loop_xy i j grid begin
        @unroll for k in 0 : -1 : (1 - grid.Hz) # integrate downwards
            @inbounds w[i, j, k] = w[i, j, k+1] + ΔzF(i, j, 0, grid) * div_xyᶜᶜᵃ(i, j, k, grid, u, v)
        end
    end
    return nothing
end

function _fill_top_halo!(w, ::NFBC, grid, u, v)
    @loop_xy i j grid begin
        @unroll for k in (grid.Nz + 2) : (grid.Nz + 1 + grid.Hz) # integrate upwards
            @inbounds w[i, j, k] = w[i, j, k-1] - ΔzF(i, j, grid.Nz+1, grid) * div_xyᶜᶜᵃ(i, j, k-1, grid, u, v)
        end
    end
    return nothing
end

function set_east_west_u_velocity!(u, i_boundary, bc, grid, clock, state)
    @loop_yz j k grid begin
        @inbounds u[i_boundary, j, k] = getbc(bc, j, k, grid, clock, state)
    end
    return nothing
end

function set_north_south_v_velocity!(v, j_boundary, bc, grid, clock, state)
    @loop_xz i k grid begin
        @inbounds v[i, j_boundary, k] = getbc(bc, i, k, grid, clock, state)
    end
    return nothing
end

function set_top_bottom_w_velocity!(w, k_boundary, bc, grid, clock, state)
    @loop_xy i j grid begin
        @inbounds w[i, j, k_boundary] = getbc(bc, i, j, grid, clock, state)
    end
    return nothing
end


# wrapper to convert LPQP solver into Polyhedra solver

# To enable Polyhedra support from an LPQP solver, define, e.g.,
# PolyhedraModel(s::GurobiSolver) = LPQPtoPolyhedraBridge(LinearQuadraticModel(s))

# Cheating a bit here
#PolyhedraModel(s::MathProgbase.AbstractMathProgSolver) = LPQPtoPolyhedraBridge(LinearQuadraticModel(s))
PolyhedraModel(s) = LPQPtoPolyhedraBridge(MPBSI.LinearQuadraticModel(s))

struct LPQPtoPolyhedraBridge <: AbstractPolyhedraModel
    lpqpmodel::MPB.AbstractLinearQuadraticModel
    rep
    c
    sense
end

LPQPtoPolyhedraBridge(m::MPB.AbstractLinearQuadraticModel) = LPQPtoPolyhedraBridge(m, nothing, nothing, nothing)

export LPQPtoPolyhedraBridge

# To transform Polyhedra problems into LinearQuadratic problems
function MPBSI.loadproblem!(m::LPQPtoPolyhedraBridge, rep::HRep, c, sense)
    lp = LPHRepresentation(rep)
    MPBSI.loadproblem!(m.lpqpmodel, lp.A, lp.l, lp.u, c, lp.lb, lp.ub, sense)
end

for f in [:optimize!, :status, :getsolution, :getobjval, :getreducedcosts, :getconstrduals]
    @eval MPBSI.$f(model::LPQPtoPolyhedraBridge) = MPBSI.$f(model.lpqpmodel)
end


#include("Crystal.jl")

"""
    module DFToutMod

Deal with energy, force, stress, band structure, etc, from a DFT calculation

Generic to different DFT codes.
"""
module DFToutMod

#using LinearAlgebra
using ..CrystalMod:crystal
using ..CrystalMod:makecrys
using ..CrystalMod:orbital_index
using Printf

using ..ThreeBodyTB:convert_energy
using ..ThreeBodyTB:convert_force
using ..ThreeBodyTB:convert_stress
using ..ThreeBodyTB:global_energy_units

##using Formatting

export dftout
export makedftout
#export test
export test2
#export crystal
#export makecrys
#export generate_supercell
#export generate_random
#export write_poscar
#export write_efs

#holds data
"""
    mutable struct bandstructure

Band structure. Has

- `nbnd::Int` Number of bands
- `nks::Int` Number of k-points
- `nelec::Float64` Number of electrons
- `efermi::Float64` Fermi energy
- `kpts::Array{Float64,2}` List of k-points, `nks × 3` array in BZ crystal units.
- `kweights::Array{Float64,1}` Weights of k-points (`nks`).
- `kgrid::Array{Int,1}` Equivalent gamma centered Monkhorst-Pack k-grid dimensions, if applies.
- `eigs::Array{Float64,2}` Eigenvalues. `nks × nbnd`

"""
mutable struct bandstructure
    nbnd::Int
    nks::Int
    nelec::Float64
    efermi::Float64
    kpts::Array{Float64,2}
    kweights::Array{Float64,1}
    kgrid::Array{Int,1}
    eigs::Array{Float64,2}
end

"""
    mutable struct dftout

DFT output struct. Has

- `crys::crystal`  crystal structure
- `energy::Float64` the actual DFT energy, depends on pseudopotentials.
- `energy_smear::Float64` smearing energy
- `forces::Array{Float64,2}` forces Ryd / a.u.
- `stress::Array{Float64,2}` stress  Ryd / (a.u.)^3
- `bandstruct::bandstructure` See bandstructure struct
- `hasband::Bool` does this object have a band structure, usually `true`
- `hasham::Bool` not used, always `false`
- `prefix::String`  A string has a name used to find output files.
- `outdir::String` Directly loaded from
- `tot_charge::Float64` If charge of unit cell is nonzero
- `atomize_energy::Float64` Atomization energy, relative to non-spin-polarized atoms.
"""
mutable struct dftout

    crys::crystal
    energy::Float64
    energy_smear::Float64
    forces::Array{Float64,2}
    stress::Array{Float64,2}
    bandstruct::bandstructure
    hasband::Bool
    hasham::Bool
    prefix::String
    outdir::String
    tot_charge::Float64
    atomize_energy::Float64

end

"""
    function get_atomize_energy(d::dftout)

Return atomization energy in current energy units
"""
function get_atomize_energy(d::dftout)

    return convert_energy(d.atomize_energy)

end

#printing
Base.show(io::IO, d::dftout) = begin
#    println(io)
#    for i in 1:3
#        @printf(io, "A%.1i=     % .5f  % .5f  % .5f\n", i, d.crys.A[i,1],  d.crys.A[i,2],  d.crys.A[i,3])
#    end
    println(io, d.crys)

    println(io)
    println("Energy (direct from dft) ", d.energy, " Ryd")
    println(io)
    println(io,"STRUCTURE =====================| FORCES =================")
    println(io)

    forces = convert_force(d.forces)

    for i in 1:d.crys.nat
        @printf(io, "%-3s  % .5f  % .5f  % .5f | % .5f % .5f % .5f\n", d.crys.types[i], d.crys.coords[i,1], d.crys.coords[i,2], d.crys.coords[i,3], forces[i,1], forces[i,2], forces[i,3])
    end
    println(io,)
    println(io,"=========================================================")
    println(io)
    println(io,"STRESS ==================")

    stress = convert_stress(d.stress)

    @printf(io, "\n")
    for i in 1:3
        @printf(io, "% .5f  % .5f  % .5f\n", stress[i,1], stress[i,2], stress[i,3])
    end
    println(io,)
    println(io, "Atomization energy (no spin): ", convert_energy(d.atomize_energy), " $global_energy_units ")
    println(io,"=========================")
    println(io,"Has bandstructure: ", d.hasband,"; Has hamiltonian: ", d.hasham)
    println(io)

end

#print bandstructure
Base.show(io::IO, d::bandstructure) = begin
    println(io,"nbnd = ", d.nbnd, "; nkpts = ", d.nks)
    println(io,"nelec = ", d.nelec, "; efermi = ", d.efermi)

    println(io)

#    pst = "{:d} {: f} {: f} {: f} , {: f},  "
#    pst = "%-3s % .5f % .5f % .5f , % .5f, "
#    println(pst)

    for i = 1:min(d.nks, 200)
        t = @sprintf("%-3s", i)
        k = @sprintf(" [ % .5f % .5f % .5f ]", d.kpts[i,1],d.kpts[i,2],d.kpts[i,3]  )
        kw = @sprintf(" % .5f ",  d.kweights[i])
        v = " "
        for j = 1:d.nbnd
            v = v * @sprintf(" % .5f",d.eigs[i,j])
        end
        println(io,t,kw,k,v)

    end
    if d.nks > 200
        println(io, "... [truncated]")
    end
    println(io)

end

"""
    function makebs(nelec::Number, efermi::Number,  kpoints, kweights,kgrid, vals)

Constructor for `bandstructure`.
"""
function makebs(nelec::Number, efermi::Number,  kpoints, kweights,kgrid, vals)
    try
        kpoints = convert(Array{Float64,2},kpoints)
        kweights = convert(Array{Float64,1}, kweights)
        vals = convert(Array{Float64,2}, vals)
    catch
        println(typeof(kpoints), typeof(kweights), typeof(vals))
        error("error make bs type conversion")
    end
    nks = size(kpoints,1)
    if nks != size(kweights,1) || nks != size(vals,1)
        error("kpoints kweights eigs don't match")
    end
    nbnd = size(vals,2)

    K2 = zeros(Float64,size(kpoints))
    for i in 1:nks
        K2[i,:] = Int64.(round.(kpoints[i,:] * 100000000))/100000000
    end

    return bandstructure(nbnd, nks, nelec, efermi, K2, kweights, kgrid, vals)

end

function make_empty_bs()

    return bandstructure(1, 1, 1, 0.0, zeros(1,3), zeros(3), [1, 1, 1], zeros(1,1))

end

"""
    function makedftout(crys::crystal, energy::Number, energy_smear::Number,  forces, stress, bandstruct=missing; prefix="PREFIX", outdir="TMPDIR", tot_charge=0.0)

Constructor for dftout. Usually called by function that loads DFT output files, not called directly.
"""
function makedftout(crys::crystal, energy::Number, energy_smear::Number,  forces, stress, bandstruct=missing; prefix="PREFIX", outdir="TMPDIR", tot_charge=0.0)
"""
Creates a struct with the desired data
"""
    try
        forces = convert(Array{Float64,2}, forces)
        stress = convert(Array{Float64,2}, stress)
        energy = convert(Float64, energy)
        energy_smear = convert(Float64, energy_smear)
    catch
        println(typeof(energy), typeof(forces), typeof(stress))
        error("Error converting types")
        return
    end

    nat = crys.nat
    if (nat,3) != size(forces)
        println( forces)
        error("Error forces")

        return
    end
    if size(stress) != (3,3)
        println( stress)
        error("Error stress", size(stress))
        return
    end

    ind2orb, orb2ind, etotal_atoms, nval =  orbital_index(crys)
    atomize_energy = energy - etotal_atoms

    if ismissing(bandstruct)
        return dftout(crys, energy, energy_smear, forces, stress, make_empty_bs() , false, false, prefix, outdir, tot_charge, atomize_energy) #, missing, False, False)
    else
        return dftout(crys, energy, energy_smear,forces, stress, bandstruct, true, false, prefix, outdir, tot_charge, atomize_energy) #, missing, False, False)
    end
end

#alterate version, will make the crys for you.
"""
    function makedftout(A, pos, types, energy::Number,energy_smear::Number,  forces, stress, bandstruct=missing; prefix="PREFIX", outdir="TMPDIR", tot_charge=0.0)
"""
function makedftout(A, pos, types, energy::Number,energy_smear::Number,  forces, stress, bandstruct=missing; prefix="PREFIX", outdir="TMPDIR", tot_charge=0.0)
    c = makecrys(A,pos,types, units="Bohr")
    if ismissing(bandstruct)
        return makedftout(c, energy, energy_smear, forces,stress, prefix=prefix, outdir=outdir, tot_charge=0.0)
    else
        return makedftout(c, energy, energy_smear, forces,stress, bandstruct,prefix=prefix, outdir=outdir, tot_charge=0.0)
    end
end

end #end module

#using .DFToutMod
#Base.:(==)(x::dftout, y::dftout) = (x.crys == y.crys && x.energy == y.energy && x.forces == y.forces && x.stress ==y.stress)


""" Very coarse solution for gm to check for feasibility and an initial starting point. """
function surveysolve(gm::GlobalModel)
    bbcs = [bbl for bbl in gm.bbls if bbl isa BlackBoxClassifier]
    if !isempty(bbcs)
        for bbc in bbcs
            learn_constraint!(bbc)
            update_tree_constraints!(gm, bbc)
        end
    end
    bbrs = [bbl for bbl in gm.bbls if bbl isa BlackBoxRegressor]
    if !isempty(bbrs)
        for bbr in bbrs
            learn_constraint!(bbr, regression_sparsity = 0, max_depth = 6, minbucket = 2*length(bbr.vars))
            update_tree_constraints!(gm, bbr)
        end
    end
    optimize!(gm)
    return
end

"""
    function add_relaxation_variables!(gm::GlobalModel, bbl::Union{BlackBoxLearner, LinkedLearner})
    function add_relaxation_variables!(gm::GlobalModel, bbls::Array)

Populates relax_var attributes of all substructs.
"""
function add_relaxation_variables!(gm::GlobalModel, bbl::Union{BlackBoxLearner, LinkedLearner})
    if bbl.relax_var isa Real
        bbl.relax_var = @variable(gm.model)
        set_lower_bound(bbl.relax_var, 0)
    end
    if !isempty(bbl.mi_constraints)
        clear_tree_constraints!(gm, bbl)
        @info "Clearing MI constraints from $(bbl.name) due to relaxation."
    end
    if bbl isa BlackBoxLearner
        for ll in bbl.lls
            add_relaxation_variables!(gm, ll)
        end
    end
end

function add_relaxation_variables!(gm::GlobalModel, bbls::Array)
    for bbl in bbls
        add_relaxation_variables!(gm, bbl)
    end
end

add_relaxation_variables!(gm::GlobalModel) = add_relaxation_variables!(gm, gm.bbls)

"""
    $(TYPEDSIGNATURES)

Clears slack variables in GlobalModel.
"""
function clear_relaxation_variables!(gm::GlobalModel, bbl::Union{BlackBoxLearner, LinkedLearner})
    if bbl.relax_var isa JuMP.VariableRef
        is_valid(gm.model, bbl.relax_var) ||
            throw(OCTHaGOnException("$(bbl.relax_var) is not a valid variable of  $(gm.model)."))
        delete_lower_bound(bbl.relax_var)
        delete(gm.model, bbl.relax_var)
        bbl.relax_var = 0
    end
    if bbl isa BlackBoxLearner
        for ll in bbl.lls
            clear_relaxation_variables!(gm, ll)
        end
    end
end

function clear_relaxation_variables!(gm::GlobalModel, bbls::Array)
    for bbl in bbls
        clear_relaxation_variables!(gm, bbl)
    end
end

clear_relaxation_variables!(gm::GlobalModel) = clear_relaxation_variables!(gm, gm.bbls)

"""
    $(TYPEDSIGNATURES)

Includes relaxation variables in as a linear penalty in the objective function.
"""
function relax_objective!(gm::GlobalModel, M::Real = 1e8)
    no = 0
    for bbl in gm.bbls
        no += M*bbl.relax_var
        for ll in bbl.lls
            no += M*ll.relax_var
        end
    end
    @objective(gm.model, Min, gm.objective + no)
    return
end

"""
Removes relaxation variables from objective.
"""
function tight_objective!(gm::GlobalModel)
    @objective(gm.model, Min, gm.objective)
end

"""
    $(TYPEDSIGNATURES)

Performs gradient descent on the last optimal solution in gm.solution_history.
In case of infeasibility, first projects the feasible point using the local
constraint gradients.

Optional kwargs:
max_iterations: maximum number of gradient steps or projections.
step_penalty: magnitude of penalty on step size during projections.
step_size: Size of 0-1 normalized Euclidian ball we can step.
decay_rate: Exponential coefficient of step size reduction.
"""
function descend!(gm::GlobalModel; kwargs...)
    # Initialization
    clear_tree_constraints!(gm)
    bbls = gm.bbls
    vars = gm.vars
    bounds = get_bounds(vars)

    # Update descent algorithm parameters
    if !isempty(kwargs)
        for item in kwargs
            set_param(gm, item.first, item.second)
        end
    end

    # Checking for a nonlinear objective
    obj_bbl = nothing
    if gm.objective isa VariableRef
        obj_bbl = filter(x -> x.dependent_var == gm.objective,
                            [bbl for bbl in gm.bbls if bbl isa BlackBoxRegressor])
        if isempty(obj_bbl)
            obj_bbl = nothing
        elseif length(obj_bbl) == 1
            obj_bbl = obj_bbl[1]
            vars = [var for var in vars if var != obj_bbl.dependent_var]
            bbls = [bbl for bbl in gm.bbls if bbl != obj_bbl]
        elseif length(obj_bbl) > 1
            throw(OCTHaGOnException("GlobaModel $(gm.name) has more than one BlackBoxRegressor" *
                               " on the objective variable."))
        end
    end

    grad_shell = DataFrame(string.(vars) .=> [Float64[] for i=1:length(vars)])

    # Add relaxation variables
    add_relaxation_variables!(gm)

    # Final checks
    obj_gradient = copy(grad_shell)
    if isnothing(obj_bbl)
        if gm.objective isa VariableRef
            append!(obj_gradient, DataFrame(string.(gm.objective) => 1), cols = :subset)
        elseif gm.objective isa JuMP.GenericAffExpr
            append!(obj_gradient, DataFrame(Dict(string(key) => value for (key,value) in gm.objective.terms)), cols = :subset)
        else
            @warn "Type of objective $(gm.objective) is unsupported."
        end
        obj_gradient = coalesce.(obj_gradient, 0)
    end
    if isempty(obj_gradient)
        obj_gradient = DataFrame(string.(vars) .=> zeros(length(vars)))
    end

    # x0 initialization, and actual objective computation
    x0 = DataFrame(gm.solution_history[end, :])
    sol_vals = x0[:,string.(vars)]
    if !isnothing(obj_bbl)
        actual_cost = obj_bbl(x0)
        push!(gm.cost, actual_cost[1])
        x0[:, string(gm.objective)] = actual_cost
        sol_vals = x0[:,string.(vars)]
        feas_gap(gm, x0) # Checking feasibility gaps with updated objective
        append!(gm.solution_history, x0) # Pushing last solution
    end

    # Descent direction, counting and book-keeping
    d = @variable(gm.model, [1:length(vars)])
    var_diff = []
    max_diff = 0
    for key in vars
        push!(var_diff, maximum(bounds[key]) - minimum(bounds[key]))
        if !isinf(var_diff[end])
            max_diff = maximum([var_diff[end], max_diff])
        end
    end
    if any(isinf.(var_diff))
        @warn "Unbounded variables detected. " *
              "PGD may fail to converge to a local minimum."
        replace!(var_diff, Inf => max_diff)
    end
    replace!(var_diff, 0 => 1) # Sanitizing fixed values.

    ct = 0
    d_improv = 1e5
    prev_feas = false
    feas = is_feasible(gm)
    abstol = get_param(gm, :abstol)
    tighttol = get_param(gm, :tighttol)
    max_iterations = get_param(gm, :max_iterations)
    step_penalty = get_param(gm, :step_penalty)
    equality_penalty = get_param(gm, :equality_penalty)
    step_size = get_param(gm, :step_size)
    decay_rate = get_param(gm, :decay_rate)

    # WHILE LOOP
    @info("Starting projected gradient descent...")
    while (!feas || !prev_feas ||
        (prev_feas && feas && d_improv >= get_param(gm, :abstol)) ||
        ct == 0) && ct < max_iterations
        prev_obj = gm.cost[end]
        constrs = []
        ct += 1
        if step_penalty == 0
            step_penalty = maximum([1e3, 1e4*abs(prev_obj)])
        end
        if equality_penalty == 0
            equality_penalty = 100*step_penalty
        end

        # Linear objective gradient and constraints
        if !isnothing(obj_bbl)
            update_gradients(obj_bbl, [length(obj_bbl.Y)])
            obj_gradient = copy(grad_shell)
            append!(obj_gradient,
                DataFrame(obj_bbl.gradients[end,:]), cols = :subset)
            obj_gradient = coalesce.(obj_gradient, 0)
            # Update objective constraints
            push!(constrs, @constraint(gm.model, sum(Array(obj_gradient[end,:]) .* d) +
                              obj_bbl.dependent_var >= obj_bbl.Y[end]))
        end

        # Constraint evaluation
        errors = 0
        for bbl in bbls
            update_gradients(bbl, [length(bbl.Y)])
            constr_gradient = copy(grad_shell)
            append!(constr_gradient, DataFrame(bbl.gradients[end,:]), cols = :subset)
            constr_gradient = coalesce.(constr_gradient, 0)
            if bbl isa BlackBoxClassifier
                error_diff = maximum([tighttol, abs(bbl.Y[end])])
                if bbl.equality
                    push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d)  + bbl.Y[end] + error_diff * bbl.relax_var >= 0))
                    push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d)  + bbl.Y[end] <= error_diff * bbl.relax_var))
                    push!(constrs, @constraint(gm.model, bbl.relax_var <= 1))
                    errors += bbl.relax_var^2
                elseif !is_feasible(bbl) || (is_feasible(bbl) && bbl.feas_gap[end] <= 0)
                    push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d)  + bbl.Y[end] + error_diff * bbl.relax_var >= 0))
                    push!(constrs, @constraint(gm.model, bbl.relax_var <= 1))
                    errors += bbl.relax_var^2
                else # feasible to zero tolerance
                    push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d)  + bbl.Y[end] >= 0))
                end
            elseif bbl isa BlackBoxRegressor
                error_diff = maximum([tighttol, abs(bbl.optima[end] - bbl.actuals[end])])
                if bbl.equality
                    push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + bbl.dependent_var + error_diff * bbl.relax_var >= bbl.Y[end]))
                    push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + bbl.dependent_var <= bbl.Y[end] + error_diff * bbl.relax_var))
                    push!(constrs, @constraint(gm.model, bbl.relax_var <= 1))
                    errors += bbl.relax_var^2
                elseif !is_feasible(bbl) || (is_feasible(bbl) && bbl.feas_gap[end] <= 0)
                    push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d)  + bbl.Y[end] + error_diff * bbl.relax_var >= 0))
                    push!(constrs, @constraint(gm.model, bbl.relax_var <= 1))
                    errors += bbl.relax_var^2
                else # feasible to zero tolerance
                    push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + bbl.dependent_var >= bbl.Y[end]))
                end
            end
            if !isempty(bbl.lls)
                n_lls = length(bbl.lls)
                update_gradients(bbl, collect(length(bbl.Y) - n_lls:length(bbl.Y)-1))
                for i = 1:n_lls
                    ll = bbl.lls[i]
                    constr_gradient = copy(grad_shell)
                    append!(constr_gradient, # Changing the headers of the gradient
                        DataFrame(string.(ll.vars) .=> Array(bbl.gradients[end-n_lls-1+i,:])), cols = :subset)
                    constr_gradient = coalesce.(constr_gradient, 0)
                    Y_val = bbl.Y[end-n_lls-1+i]
                    if bbl isa BlackBoxClassifier
                        error_diff = maximum([tighttol, abs(Y_val)])
                        if bbl.equality
                            push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + Y_val + error_diff * ll.relax_var >= 0))
                            push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + Y_val <= error_diff * ll.relax_var))
                            push!(constrs, @constraint(gm.model, ll.relax_var <= 1))
                            errors += ll.relax_var^2
                        elseif !is_feasible(ll) || (is_feasible(ll) && ll.feas_gap[end] <= 0)
                            push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d)  + Y_val + error_diff * ll.relax_var >= 0))
                            push!(constrs, @constraint(gm.model, ll.relax_var <= 1))
                            errors += ll.relax_var^2
                        else # feasible to zero tolerance
                            push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d)  + Y_val >= 0))
                        end
                    elseif bbl isa BlackBoxRegressor
                        error_diff = maximum([tighttol, abs(ll.optima[end] - ll.actuals[end])])
                        if bbl.equality
                            push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + ll.dependent_var + error_diff * ll.relax_var >= Y_val))
                            push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + ll.dependent_var <= Y_val + error_diff * ll.relax_var))
                            push!(constrs, @constraint(gm.model, ll.relax_var <= 1))
                            errors += ll.relax_var^2
                        elseif !is_feasible(ll) || (is_feasible(ll) && ll.feas_gap[end] <= 0)
                            push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + ll.dependent_var + error_diff * ll.relax_var >= Y_val))
                            push!(constrs, @constraint(gm.model, ll.relax_var <= 1))
                            errors += ll.relax_var^2
                        else # feasible to zero tolerance
                            push!(constrs, @constraint(gm.model, sum(Array(constr_gradient[end,:]) .* d) + ll.dependent_var >= Y_val))
                        end
                    end
                end
            end
        end

        # Initializing descent direction
        push!(constrs, @constraint(gm.model, d .== vars .- Array(sol_vals[end,:])))
        if feas # Small gradient step
            push!(constrs, @constraint(gm.model, sum((d ./ var_diff).^2) <=
                step_size/exp(decay_rate*(ct-1)/max_iterations)))
            @objective(gm.model, Min, sum(Array(obj_gradient[end,:]) .* d) + equality_penalty*errors)
        else    # Projection step
            @objective(gm.model, Min, sum(Array(obj_gradient[end,:]) .* d) + step_penalty*sum((d ./ var_diff).^2) + equality_penalty*errors)
        end

        # Optimizing the step, and finding next x0
        optimize!(gm.model)
        x0 = DataFrame(string.(vars) .=> getvalue.(vars))
        if !isnothing(obj_bbl)
            x0[!, string(gm.objective)] = obj_bbl(x0)
        end
        sol_vals = x0[:,string.(vars)]
        feas_gap(gm, x0)
        append!(gm.solution_history, x0)

        # Measuring improvements
        if !isnothing(obj_bbl)
            push!(gm.cost, gm.solution_history[end, string(gm.objective)])
        else
            push!(gm.cost,  JuMP.getvalue(gm.objective))
        end
        d_improv = gm.cost[end-1] - gm.cost[end]

        # Feasibility check
        prev_feas = feas
        feas = is_feasible(gm)

        # Saving solution dict for JuMP-style recovery.
        # TODO: do not regenerate soldict unless final solution.
        gm.soldict = Dict(key => JuMP.getvalue.(gm.model[key]) for (key, value) in gm.model.obj_dict)

        # Delete gradient constraints (TODO: perhaps add constraints to avoid cycling?)
        for con in constrs
            delete(gm.model, con)
        end
    end

    fincost = round(gm.cost[end], digits = -Int(round(log10(abstol))))
    # Termination criteria
    if ct >= max_iterations
        @info "Max iterations ($(ct)) reached."
        if prev_feas && feas
            @info "Solution is not converged to tolerance $(abstol)!"
        elseif (feas && !prev_feas) && (abs(d_improv) <= 100*abstol)
            @info "Solution is feasible and likely cycling, but the solution is close. Reduce step size and descend again. "
        elseif (!feas && prev_feas) && (abs(d_improv) <= 100*abstol)
            @info "Solution is infeasible and likely cycling, but the solution is close. Reduce step size and descend again. "
        elseif ((feas && !prev_feas) || (prev_feas && !feas))
            @info "Solution is likely cycling, with > $(abstol) changes in cost. Reduce step size and descend again."
        elseif !feas && !prev_feas
            @info "Solution is infeasible to tolerance $(tighttol)."
        end
        @info("Final cost is $(fincost).")
    else
        @info("PGD converged in $(ct) iterations!")
        @info("The optimal cost is $(fincost).")
    end

    # Reverting objective, and deleting vars
    @objective(gm.model, Min, gm.objective)
    delete(gm.model, d)
    return
end

""" Complete solution procedure for GlobalModel. """
function globalsolve!(gm::GlobalModel)
    @info "GlobalModel " * gm.name * " solution in progress..."
    for bbl in gm.bbls
        if !is_sampled(bbl)
            uniform_sample_and_eval!(bbl)
        end
    end
    set_param(gm, :ignore_accuracy, true)
    set_param(gm, :ignore_feasibility, true)
    @info "Training OptimalTreeLearners..."
    for bbl in gm.bbls
        if isempty(bbl.learners)
            learn_constraint!(bbl)
        end
        if isempty(bbl.mi_constraints)
            add_tree_constraints!(gm, bbl)
        end
    end
    @info "Solving MIP..."
    optimize!(gm)
    descend!(gm)
end

function globalsolve_and_time!(m::Union{JuMP.Model, GlobalModel})
    t1 = time()
    if m isa JuMP.Model
        optimize!(m);
    else
        globalsolve!(m)
    end
    println("Model solution time: " * string(time()-t1) * " seconds.")
end


## Main function to run Bayesian OPTIMISATION
function settingsBayesOptMCBU(oedmc_def)

    # Manage inputs order
    inpnam = oedmc_def["Model"]["inpName"];

    # Optimisation section
    steps = length(oedmc_def["switchT"])-1;
    induc = oedmc_def["Model"]["nInp"];
    fiii = length(oedmc_def["fixedInp"]);

    doof = (steps*(induc-fiii))+fiii;

    # THE THIRD MOUNTAINS!
    uppe = Array{Any,1}(undef, doof); # Vector of upper bounds. Length set with doof but deppending on the user settings from fixedStep and equalStep this can change.
    lowe = Array{Any,1}(undef, doof); # Vector of lower bounds. Same ^
    coun = 1; # Index for the entry of ustr that will advance only if certain conditions are meet (that that position has been already checked and a variable or number has been introduced)
                # Tecnichally this indicated the first entry of uppe and lowwe for each inducer variables.
    fs = [oedmc_def["fixedStep"][j][1] for j in 1:length(oedmc_def["fixedStep"])]; # Fixed Steps (the actual step number)
    us = [oedmc_def["fixedStep"][j][2] for j in 1:length(oedmc_def["fixedStep"])]; # Input values (or Any entry) for each step specified in fixedStep
    fixin = 0; # Index to account for fixedInp. If the fixedInp is the first one in the inpam list this is the best way (that I have found) to account for it. This is because of the use of i (input index). This is because first I consider the case where there is a fisedInp and then check the rest.

    for i in 1:length(inpnam) # Loop over the number of inducers
        if inpnam[i] in oedmc_def["fixedInp"] # If the current inducer is fixed then add only one bound and move index (coun). Fixin works as described in the other MOUNTAINS
            uppe[coun] = oedmc_def["uUpper"][i];
            lowe[coun] = oedmc_def["uLower"][i];
            coun += 1;
            fixin += 1;
        else
            try # Using coun you take all the indexes corresponding to the current inducer. The you fill it with the specified uper and lower bounds. Added the try just in case someoen runns it in a machien that does not support Float 64? (I am not sure)
                uppe[coun:coun+(steps-1)] = repeat([convert(Float64, oedmc_def["uUpper"][i])], outer = steps);
                lowe[coun:coun+(steps-1)] = repeat([convert(Float64, oedmc_def["uLower"][i])], outer = steps);
            catch
                uppe[coun:coun+(steps-1)] = repeat([convert(Float32, oedmc_def["uUpper"][i])], outer = steps);
                lowe[coun:coun+(steps-1)] = repeat([convert(Float32, oedmc_def["uLower"][i])], outer = steps);
            end

            for j in 1:length(uppe[coun:coun+(steps-1)]) # Now we loop over each entry that we filled to check if we need to delete it according to the users fixedStep and equalStep
                if j in fs # Since this is not a fixed input j also indicates the step index, so check if this is in the list of fixedStep
                    for w in 1:length(fs) # As commentetd in other loops, this for and if is to consider the case the user does not introduce fixedStep entries sorted by step.
                        if j == fs[w]
                            if us[w][i-fixin] != Any # Check if current input contained in a fixedStep is defined as Any (to be optimised) or not (has been fixed. )
                                uppe[coun-1+j] = "";# Since it was fixed (different from Any) ad an empty string there to remove later in both bounds.
                                lowe[coun-1+j] = "";
                            end
                        end
                    end
                end
                if oedmc_def["equalStep"] != [] # If equalStep is empty then we do not need to do anything else, if it has something then there are entries to remove from the bounds
                    for k in 1:length(oedmc_def["equalStep"]) # Loop over each entry of equalStep
                        tmpk = zeros(length(oedmc_def["equalStep"][k])); # Temporal variable that will contain the same values as the equalStep entry except the ones that have been fixed for that step (this will happen if the user set an Any for the current inducer)
                        for h in 1:length(oedmc_def["equalStep"][k]) # Loop over each value in the current entry
                            if oedmc_def["equalStep"][k][h] in fs # Check if the current step considered is also in fixedStep
                                for w in 1:length(fs) # Account for ht euser not introducng steps sorted
                                    if oedmc_def["equalStep"][k][h] == fs[w]
                                        if us[w][i-fixin] == Any # If the current step considered was defined as Any for the current inducer add the index, if it was fixed to a value ignore (this has been taken into account before)
                                            tmpk[h] = oedmc_def["equalStep"][k][h];
                                        end
                                    end
                                end
                            else
                                tmpk[h] = oedmc_def["equalStep"][k][h] # If the current step was not in fixedStep just add it
                            end
                        end
                        tmpk = convert.(Int, filter!(e->e∉[0],tmpk)); # Remove 0 entries (fixed values)
                        if j in tmpk && j != tmpk[1] # Now, check if the current step is in our temporary index. If it is the first entry in the vector then do nothing, if not remove (it will be the same as the first, so no bound for it)
                            uppe[coun-1+j] = "";
                            lowe[coun-1+j] = "";
                        end
                    end
                end
            end
            coun += (steps); # Add the number of steps to the varable coun so we can move to the next inducer
        end
    end

    uppe = [uppe[i] for i in 1:length(uppe) if uppe[i] != ""]; # Remove empty strings so we are left with a vector of numbers with the bounds for the inducers.
    lowe = [lowe[i] for i in 1:length(lowe) if lowe[i] != ""];

    #################################### OPTIMISER SETINGS THAT CAN BE MODIFIED START HERE ####################################

    model = ElasticGPE(convert(Int, length(uppe)),                            #
           mean = MeanConst(0.),
           kernel = SEArd(zeros(length(uppe)), 5.),
           logNoise = 0.,
           capacity = 3000);

    modeloptimizer = MAPGPOptimizer(every = 50, noisebounds = [-4, 3],       # bounds of the logNoise
                   kernbounds = [[-1*ones(length(uppe)); 0], [4*ones(length(uppe)); 10]],  # bounds of the 3 parameters GaussianProcesses.get_param_names(model.kernel)
                   maxeval = 40);

    sut = Symbol(string(oedmc_def["Model"]["NameF"],"Utility"));
    opt = BOpt((@eval $sut),
          model,
          ExpectedImprovement(),                   # type of acquisition
          modeloptimizer,
          lowe,
          uppe,                                     # lowerbounds, upperbounds
          repetitions = 1,                          # evaluate the function for each input 1 times
          maxiterations = oedmc_def["maxiter"],     # evaluate at 50 input positions
          sense = Max,                              # maximise the function
          verbosity = Progress);

    #################################### OPTIMISER SETINGS THAT CAN BE MODIFIED STOP HERE ####################################

    return opt

end


# In the .mat file, the following code is used:
#  [chosen_juice, licking_result] = [1,1]   # animal chose A
#                                 = [2,1]   # animal chose B
#                                 = [0,1]   # animal made wrong choice during forced choice trial
#                                 = [0,0]   # animal made no licking responses
function mat_trialresult(nA, nB, choice, licked)
    leftA = nA < 0 || nB > 0
    nA, nB = abs(nA), abs(nB)
    licked == 0 && return TrialResult(nA, nB, leftA, missing)
    if choice == 0
        nA == 0 && nB == 0 && error("double-forced unexpected")
        nA == 0 || nB == 0 || error("expected this to be a forced-choice")
        choiceA = nA == 0
    else
        choiceA = choice == 1
    end
    return TrialResult(nA, nB, leftA, choiceA)
end

"""
    parsemat(filename)
    parsemat(readfunction, filename)

Parse the trial data in `filename`. Optionally pass a `readfunction` to extract additional elements;
`readfunction` should have syntax

```julia
function readfunction(data, args...)
    moredata = data["matlab_variable_name"]
    # maybe do some processing/validation
    return args..., moredata
end
```

`data` is the result of `MAT.matread(filename)`.

See also: [`positive_cells`](@ref).
"""
parsemat(filename::AbstractString) = parsemat(returntail, filename)

function parsemat(f, filename::AbstractString)
    data = matread(filename)
    celldata = data["celldata"]
    goodtrials = data["goodTrials"]
    psyphydata = data["psyphydata"]

    # parse celldata
    cts = OrderedDict{Int,CellsTrial{Float64}}()   # not all trial indexes are guaranteed to be present
    i0 = i = 1
    idx = Int(celldata[i,2])
    while i <= size(celldata, 1)
        i += 1
        idxnew = i <= size(celldata, 1) ? Int(celldata[i,2]) : idx+1
        if idxnew != idx
            rng = i0:i-1
            cts[idx] = CellsTrial{Float64}(celldata[rng,1]*ms, celldata[rng,3:end])
            i0 = i
            idx = idxnew
        end
    end

    # parse goodtrials (here we assume all trials are listed)
    tts = OrderedDict{Int,TrialResult}()
    for r in eachrow(goodtrials)
        tts[Int(r[1])] = mat_trialresult(r[2:end]...)
    end

    # parse psyphydata
    ets = OrderedDict{Int,EventTiming}()
    j = 0
    while j < size(psyphydata,1)
        trial_start = offer_on = offer_off = go = choice = trial_end = Tms(NaN)
        i = Int(psyphydata[j+1,2])
        while j < size(psyphydata,1) && psyphydata[j+1,2] == i
            j += 1
            t, code = psyphydata[j,1]*ms, psyphydata[j,3]
            if code == 21
                trial_start = t
            elseif code == 30
                offer_on = t
            elseif code == 31
                offer_off = t
            elseif code == 35
                go = t
            elseif code == 39
                choice = t
            elseif code == 43 || code == 44
                @assert t == choice
            elseif code == 60
                trial_end = t
            else
                error("code $code not recognized")
            end
        end
        ets[i] = EventTiming(trial_start, offer_on, offer_off, go, choice, trial_end)
    end

    keys(cts) == keys(tts) == keys(ets) || throw(DimensionMismatch("trial indices must agree"))

    return f(data, cts, tts, ets)
end

returntail(data, args...) = args

"""
    positive_cells(data, args...)

Extract the vector of `idx_positive_cells` from the .mat file.
"""
function positive_cells(data, args...)
    celldata = data["celldata"]
    poscells = data["idx_positive_cells"]
    size(poscells, 2) == 1 || throw(DimensionMismatch("expected idx_positive_cells to be a column vector"))
    poscells = vec(poscells)
    size(celldata, 2) == length(poscells) + 2 ||
        throw(DimensionMismatch("$filename: idx_positive_cells must match columns in celldata"))

    poscells = convert(AbstractVector{Bool}, poscells)

    return args..., poscells
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

"""
Problem u(X) = u₀ in Γ(d)
"""
type Dirichlet <: BoundaryProblem
    formulation :: Symbol
    variational :: Bool
    dual_basis :: Bool
    order :: Int
end

function Dirichlet()
    Dirichlet(:incremental, false, false, 1)
end

function get_unknown_field_name(::Type{Dirichlet})
    return "reaction force"
end

function get_formulation_type(problem::Problem{Dirichlet})
    return problem.properties.formulation
end

function assemble!(problem::Problem{Dirichlet}, time::Float64=0.0;
                   auto_initialize=true)
    # FIXME: boilerplate
    if !isempty(problem.assembly)
        warn("Assemble problem $(problem.name): problem.assembly is not empty and assembling, are you sure you know what are you doing?")
    end
    if isempty(problem.elements)
        warn("Assemble problem $(problem.name): problem.elements is empty, no elements in problem?")
    else
        first_element = first(problem.elements)
        unknown_field_name = get_unknown_field_name(problem)
        if !haskey(first_element, unknown_field_name)
            warn("Assemble problem $(problem.name): seems that problem is uninitialized.")
            if auto_initialize
                info("Initializing problem $(problem.name) at time $time automatically.")
                initialize!(problem, time)
            end
        end
    end

    if method_exists(assemble_prehook!, Tuple{typeof(problem), Float64})
        assemble_prehook!(problem, time)
    end

    if problem.properties.variational
        for element in get_elements(problem)
            assemble!(problem.assembly, problem, element, time)
        end
    else # nodal collocation
        field_vals = Dict{Int64, Float64}()
        field_name = get_parent_field_name(problem)
        field_dim = get_unknown_field_dimension(problem)
        for element in get_elements(problem)
            gdofs = get_gdofs(problem, element)
            for i=1:field_dim
                haskey(element, field_name*" $i") || continue
                ldofs = gdofs[i:field_dim:end]
                xis = get_reference_coordinates(typeof(element.properties))
                vals = Float64[]
                for xi in xis
                    g = element(field_name*" $i", xi, time)
                    # u = u_prev + Δu ⇒ Δu = u - u_prev
                    if haskey(element, field_name)
                        g_prev = element(field_name, xi, time)
                        g -= g_prev[i]
                    end
                    push!(vals, g)
                end
                for (dof, g) in zip(ldofs, vals)
                    field_vals[dof] = g
                end
            end
        end
        for (k, v) in field_vals
            push!(problem.assembly.C1, k, k, 1.0)
            push!(problem.assembly.C2, k, k, 1.0)
            push!(problem.assembly.g, k, 1, v)
        end
    end

    if method_exists(assemble_posthook!, Tuple{typeof(problem), Float64})
        assemble_posthook!(problem, time)
    end
end

function assemble!(assembly::Assembly, problem::Problem{Dirichlet},
                   element::Element, time::Float64)

    # get dimension and name of PARENT field
    nnodes = length(element)
    field_dim = get_unknown_field_dimension(problem)
    field_name = get_parent_field_name(problem)
    gdofs = get_gdofs(element, field_dim)
    props = problem.properties

    if problem.properties.dual_basis
        De, Me, Ae = get_dualbasis(element, time)
    else
        Ae = eye(nnodes)
        De = zeros(nnodes, nnodes)
        for ip in get_integration_points(element, props.order)
            N = element(ip, time)
            detJ = element(ip, time, Val{:detJ})
            De += ip.weight*N'*N*detJ
        end
    end

    # left hand side
    for i=1:field_dim
        ldofs = gdofs[i:field_dim:end]
        if haskey(element, field_name*" $i")
            add!(assembly.C1, ldofs, ldofs, De)
            add!(assembly.C2, ldofs, ldofs, De)
        end
    end

    # right hand side
    for ip in get_integration_points(element, props.order)
        detJ = element(ip, time, Val{:detJ})
        w = ip.weight*detJ
        N = element(ip, time)

        for i=1:field_dim
            ldofs = gdofs[i:field_dim:end]
            if haskey(element, field_name*" $i")
                g = element(field_name*" $i", ip, time)
                # u = u_prev + Δu ⇒ Δu = u - u_prev
                if haskey(element, field_name)
                    g_prev = element(field_name, ip, time)
                    g -= g_prev[i]
                end
                add!(assembly.g, ldofs, w*g*Ae*N')
            end
        end

    end

end

#=

function assemble!(assembly::Assembly, problem::Problem{DirichletProblem}, element::Element, time::Real)

    # get dimension and name of PARENT field
    field_dim = problem.parent_field_dim
    field_name = problem.parent_field_name

    gdofs = get_gdofs(element, field_dim)
    for ip in get_integration_points(element, Val{2})
        w = ip.weight
        J = get_jacobian(element, ip, time)
        JT = transpose(J)
        if size(JT, 2) == 1  # plane problem
            w *= norm(JT)
        else
            w *= norm(cross(JT[:,1], JT[:,2]))
        end
        N = element(ip, time)
        A = w*N'*N

        if haskey(element, field_name)
            # add all dimensions at once if defined
            # element["blaa"] = 0.0
            # or
            # element["blaa"] = Vector{Float64}[[0.1, 0.2], [0.3, 0.4]]
            g = element(field_name, ip, time)
            if length(g) != length(N)
                g = g*ones(length(N))
            end
            for i=1:field_dim
                ldofs = gdofs[i:field_dim:end]
                add!(assembly.C1, ldofs, ldofs, A)
                add!(assembly.C2, ldofs, ldofs, A)
            end
            add!(assembly.g, gdofs, w*g*N)
        end

        for i=1:field_dim
            if haskey(element, field_name*" $i")
                g = element(field_name*" $i", ip, time)
                ldofs = gdofs[i:field_dim:end]
                add!(assembly.C1, ldofs, ldofs, A)
                add!(assembly.C2, ldofs, ldofs, A)
                add!(assembly.g, ldofs, w*g*N)
            end
        end
    end
end
=#


function initjumpmodel(options)
    @unpack solver, threads, showsolverlog = options
    if solver == :cplex
    # https://www.ibm.com/support/knowledgecenter/SSSA5P_12.6.0/ilog.odms.studio.help/CPLEX/ReleaseNotes/topics/releasenotes126/newParameterNames.html
    # https://www.ibm.com/support/knowledgecenter/SSSA5P_12.8.0/ilog.odms.cplex.help/CPLEX/Parameters/topics/introListAlpha.html
    # https://www.ibm.com/support/knowledgecenter/SSSA5P_12.8.0/ilog.odms.cplex.help/CPLEX/UsrMan/topics/cont_optim/barrier/19_tuning_title_synopsis.html
        m = Model(solver=CplexSolver(CPXPARAM_LPMethod=4, CPXPARAM_Threads=threads, CPXPARAM_Read_Scale=1,
                    CPXPARAM_Preprocessing_Dual=-1, CPXPARAM_SolutionType=2,
                    CPXPARAM_ScreenOutput=Int(showsolverlog)))
                    # CPXPARAM_Emphasis_Numerical=1,
                    # CPXPARAM_Barrier_ConvergeTol=1e-8 (default), CPXPARAM_Barrier_Display=2,
                    # CPXPARAM_Barrier_Crossover=-1 or 2, CPXPARAM_Emphasis_Memory=1,
                    # Tom: CPXPARAM_Barrier_ColNonzeros=200, CPXPARAM_Preprocessing_Fill=0 (doesn't help)

    elseif solver == :gurobi
        m = Model(solver=GurobiSolver(Method=2, Threads=threads, BarConvTol=1.5e-8, Crossover=0))
                # BarConvTol=1e-8, NumericFocus=2 (0-3), ScaleFlag=3 (-1-3), Crossover=0
    elseif solver == :glpk
        m = Model(solver=GLPKSolverLP(method=:InteriorPoint, msg_lev=GLPK.MSG_ON))
    elseif solver == :clp
        m = Model(solver=ClpSolver(SolveType=4, LogLevel=4))
    else
        println("No solver installed called $solver.")
    end
    return m
end

function makevariables(m, sets)
    @unpack REGION, FUEL, TECH, CLASS, STORAGECLASS, HOUR, techtype = sets

    storagetechs = [k for k in TECH if techtype[k] == :storage]

    @variables m begin
        Systemcost[r in REGION]                                                         # M€/year
        CO2emissions[r in REGION]                                                       # kton CO2/year
        FuelUse[r in REGION, f in FUEL] >= 0                                            # GWh fuel/year
        Electricity[r in REGION, k in TECH, c in CLASS[k], h in HOUR] >= 0              # GWh elec/period
        AnnualGeneration[r in REGION, k in TECH] >= 0                                   # GWh elec/year
        Charging[r in REGION, k in storagetechs, h in HOUR] >= 0       # GWh elec/period (electricity used to charge)
        StorageLevel[r in REGION, k in storagetechs, c in STORAGECLASS[k], h in HOUR] >= 0  # TWh elec (in storage)
        Transmission[r1 in REGION, r2 in REGION, h in HOUR] >= 0                        # GWh elec/period
        TransmissionCapacity[r1 in REGION, r2 in REGION] >= 0                           # GW elec
        Capacity[r in REGION, k in TECH, c in CLASS[k]] >= 0                            # GW elec
        SolarCapacity[r in REGION, k in [:pv, :csp], pv in CLASS[:pv], csp in CLASS[:csp]] >= 0     # GW elec
    end #variables

    return Vars(Systemcost, CO2emissions, FuelUse, Electricity, AnnualGeneration, Charging, StorageLevel,
                    Transmission, TransmissionCapacity, Capacity, SolarCapacity)
end

function setbounds(sets, params, vars, options)
    @unpack REGION, TECH, CLASS, techtype = sets
    @unpack Capacity, TransmissionCapacity = vars
    @unpack classlimits, hydrocapacity, transmissionislands, demand = params
    @unpack hydroinvestmentsallowed, nuclearallowed, transmissionallowed, disabletechs = options
    for r in REGION, k in TECH
        if techtype[k] == :vre || k == :csp
            for c in CLASS[k]
                setupperbound(Capacity[r,k,c], classlimits[r,k,c])
            end
        end
    end
    for r in REGION
        setlowerbound(Capacity[r,:hydro,:x0], hydrocapacity[r,:x0])
        for c in CLASS[:hydro]
            if hydroinvestmentsallowed
                setupperbound(Capacity[r,:hydro,c], hydrocapacity[r,c])
            else
                setupperbound(Capacity[r,:hydro,c], c == :x0 ? hydrocapacity[r,c] : 0.0)
            end
        end
    end
    if !nuclearallowed
        for r in REGION
            setupperbound(Capacity[r,:nuclear,:_], 0.0)
        end
    end
    for r1 in REGION, r2 in REGION
        if transmissionallowed == :none || (transmissionallowed == :islands && !transmissionislands[r1,r2])
            setupperbound(TransmissionCapacity[r1,r2], 0.0)
        end
    end
    for r in REGION, k in disabletechs, c in CLASS[k]
        setupperbound(Capacity[r,k,c], 0.0)
    end
    # for r in REGION
    #   setupperbound(AnnualGeneration[r,:csp], 0.50 * sum(demand[r,h] for h in HOUR) * hoursperperiod)
    # end
end

function makeconstraints(m, sets, params, vars, hourinfo, options)
    @unpack REGION, FUEL, TECH, CLASS, STORAGECLASS, HOUR, techtype, techfuel, reservoirclass = sets
    @unpack cf, transmissionlosses, demand, cfhydroinflow, efficiency, rampingrate, dischargetime, initialstoragelevel,
            minflow_existinghydro, emissionsCO2, fuelcost, variablecost, smalltransmissionpenalty, investcost, crf, fixedcost,
            transmissioninvestcost, transmissionfixedcost, hydroeleccost, solarcombinedarea,
            pv_density, csp_density, cspsolarmultiple = params
    @unpack Systemcost, CO2emissions, FuelUse, Electricity, AnnualGeneration, Charging, StorageLevel,
            Transmission, TransmissionCapacity, Capacity, SolarCapacity = vars
    @unpack hoursperperiod = hourinfo
    @unpack carbontax, carboncap, rampingconstraints, maxbioenergy, globalnuclearlimit = options

    storagetechs = [k for k in TECH if techtype[k] == :storage]

    @constraints m begin
        ElecCapacity[r in REGION, k in TECH, c in CLASS[k], h in HOUR],
            Electricity[r,k,c,h] <= Capacity[r,k,c] * (k == :csp ? 1 : cf[r,k,c,h]) * hoursperperiod

        ElecDemand[r in REGION, h in HOUR],
            sum(Electricity[r,k,c,h] for k in TECH, c in CLASS[k]) - sum(Charging[r,k,h] for k in TECH if techtype[k] == :storage) +
                + sum((1-transmissionlosses[r2,r])*Transmission[r2,r,h] - Transmission[r,r2,h] for r2 in REGION) >=
                    demand[r,h] * hoursperperiod

        # <= instead of == to avoid need of slack variable to deal with spillage during spring floods, etc
        StorageBalance[r in REGION, k in storagetechs, sc in STORAGECLASS[k], h in HOUR],
            (StorageLevel[r,k,sc,h] - StorageLevel[r,k,sc, (h>1) ? h-1 : length(HOUR)]) / 1 <=  # unit: energy diff per period (TWh/period)
                0.001 * Charging[r,k,h] +
                + (k == :hydro ? 0.001 * hoursperperiod * sum(cfhydroinflow[r,c,h] * Capacity[r,:hydro,c] for c in reservoirclass[sc])
                                : 0.0) +
                + (k == :csp ? 0.001 * hoursperperiod * sum(cf[r,:csp,c,h] * Capacity[r,:csp,c] for c in reservoirclass[sc])
                                : 0.0) +
                - 0.001 * sum(Electricity[r,k,c,h]/efficiency[k] for c in reservoirclass[sc])

        MaxStorageCapacity[r in REGION, k in storagetechs, sc in STORAGECLASS[k], h in HOUR],
            StorageLevel[r,k,sc,h] <= sum(Capacity[r,k,c] * dischargetime[r,k,c] for c in reservoirclass[sc]) / 1000

        InitialStorageLevel[r in REGION, k in storagetechs, sc in STORAGECLASS[k]],
            StorageLevel[r,k,sc,1] ==
                initialstoragelevel * sum(Capacity[r,k,c] * dischargetime[r,k,c] for c in reservoirclass[sc]) / 1000

        # consider turbine efficiencies later
        # POSSIBLE BUG (why only class x0 here?)
        MinHydroFlow[r in REGION, h in HOUR],
            Electricity[r,:hydro,:x0,h] >= minflow_existinghydro * hoursperperiod * cfhydroinflow[r,:x0,h] * Capacity[r,:hydro,:x0]

        # We'll add pumped hydro later
        NoCharging[r in REGION, h in HOUR, k in [:hydro, :csp]],
            Charging[r,k,h] == 0

        ChargingNeedsBattery[r in REGION, h in HOUR],
            Charging[r,:battery,h] <= Capacity[r,:battery, :_] * hoursperperiod

        MaxTransmissionCapacity[r1 in REGION, r2 in REGION, h in HOUR],
            Transmission[r1,r2,h] <= TransmissionCapacity[r1,r2] * hoursperperiod

        TwoWayStreet[r1 in REGION, r2 in REGION],
            TransmissionCapacity[r1,r2] == TransmissionCapacity[r2,r1]

        NoTransmission[r1 in REGION, r2 in REGION; transmissioninvestcost[r1,r2] == 0],
            TransmissionCapacity[r1,r2] == 0

        Calculate_AnnualGeneration[r in REGION, k in TECH],
            AnnualGeneration[r,k] == sum(Electricity[r,k,c,h] for c in CLASS[k], h in HOUR)

        Calculate_FuelUse[r in REGION, f in FUEL; f != :_],
            FuelUse[r,f] == sum(AnnualGeneration[r,k]/efficiency[k] for k in TECH if techfuel[k]==f)

        BioLimit[r in REGION],
            sum(AnnualGeneration[r,k] for k in [:bioGT, :bioCCGT]) <= maxbioenergy * sum(demand[r,h] for h in HOUR) * hoursperperiod

        # This does not quite make the variable bound redundant, because e.g. some pixels in PV class 1 are class "0" for CSP,
        # and are therefore unaffected by this constraint.
        SolarOverlap[r in REGION, pv in CLASS[:pv], csp in CLASS[:csp]],
            SolarCapacity[r,:pv,pv,csp]/pv_density + SolarCapacity[r,:csp,pv,csp]/csp_density/cspsolarmultiple <=
                solarcombinedarea[r,pv,csp]

        # There could be a minor bug here depending on the numerical class limits. If there are pixels in e.g. PV class a1
        # which do not belong to any CSP class, then we may miscount PV capacity. But this does not happen with our parameters.
        PVcapacity[r in REGION, pv in CLASS[:pv]],
            Capacity[r,:pv,pv] == sum(SolarCapacity[r,:pv,pv,csp] for csp in CLASS[:csp])

        CSPcapacity[r in REGION, csp in CLASS[:csp]],
            Capacity[r,:csp,csp] == sum(SolarCapacity[r,:csp,pv,csp] for pv in CLASS[:pv])

        TotalCO2[r in REGION],
            CO2emissions[r] == sum(FuelUse[r,f] * emissionsCO2[f] for f in FUEL)

        GlobalCO2Cap,
            sum(CO2emissions[r] for r in REGION) <= carboncap * sum(demand) * hoursperperiod

        # Storage costs included in ordinary Capacity costs
        # Transmission costs halved since they are counted twice (allocate half the cost to sending and receiving regions)
        # make this regional later
        Totalcosts[r in REGION],
            Systemcost[r] ==
                0.001 * sum(FuelUse[r,f] * fuelcost[f] for f in FUEL) +
                + 0.001 * carbontax * CO2emissions[r] +
                + 0.001 * sum(Electricity[r,k,c,h] * variablecost[k] for k in TECH, c in CLASS[k], h in HOUR) +
                + 0.001 * sum(Electricity[r,:hydro,c,h] * hydroeleccost[r,c] for c in CLASS[:hydro], h in HOUR) +
                + 0.001 * sum(Transmission[r,r2,h] * smalltransmissionpenalty for r2 in REGION, h in HOUR) +
                + sum(Capacity[r,k,c] * (investcost[k,c] * crf[k] + fixedcost[k]) for k in TECH, c in CLASS[k]) +
                + 0.5 * sum(TransmissionCapacity[r,r2] *
                            (transmissioninvestcost[r,r2] * crf[:transmission] + transmissionfixedcost[r,r2]) for r2 in REGION)
        # =#
    end #constraints

    if rampingconstraints
        @constraints m begin
            RampingDown[r in REGION, k in TECH, c in CLASS[k], h in HOUR; rampingrate[k] < 1],
                Electricity[r,k,c,h] - Electricity[r,k,c, (h>1) ? h-1 : length(HOUR)] >=
                    -rampingrate[k] * Capacity[r,k,c] * cf[r,k,c,h] * hoursperperiod

            RampingUp[r in REGION, k in TECH, c in CLASS[k], h in HOUR; rampingrate[k] < 1],
                Electricity[r,k,c,h] - Electricity[r,k,c, (h>1) ? h-1 : length(HOUR)] <=
                    rampingrate[k] * Capacity[r,k,c] * cf[r,k,c,h] * hoursperperiod
        end
    else
        RampingDown = RampingUp = nothing
    end

    if globalnuclearlimit != Inf
        @constraints m begin
            NuclearCapacityLimit,
                sum(Capacity[r,:nuclear,:_] for r in REGION) <= globalnuclearlimit
        end
    end

    return Constraints(ElecCapacity, ElecDemand, RampingDown, RampingUp, StorageBalance, MaxStorageCapacity, InitialStorageLevel,
                MaxTransmissionCapacity, TwoWayStreet, NoTransmission, NoCharging, ChargingNeedsBattery,
                Calculate_AnnualGeneration, Calculate_FuelUse, TotalCO2, Totalcosts)
end

function makeobjective(m, sets, vars)
    @unpack REGION = sets
    @unpack Systemcost = vars

    @objective m Min begin
        sum(Systemcost[r] for r in REGION)
    end
end


#####
# some common components for Prioritized Experience Replay based methods
#####

const PERLearners = Union{PrioritizedDQNLearner,RainbowLearner,IQNLearner}

function extract_experience(t::AbstractTrajectory, learner::PERLearners)
    s = learner.stack_size
    h = learner.update_horizon
    n = learner.batch_size
    γ = learner.γ

    # 1. sample indices based on priority
    valid_ind_range = isnothing(s) ? (1:(length(t)-h)) : (s:(length(t)-h))
    if t isa CircularCompactPSARTSATrajectory
        inds = Vector{Int}(undef, n)
        priorities = Vector{Float32}(undef, n)
        for i in 1:n
            ind, p = sample(learner.rng, get_trace(t, :priority))
            while ind ∉ valid_ind_range
                ind, p = sample(learner.rng, get_trace(t, :priority))
            end
            inds[i] = ind
            priorities[i] = p
        end
    else
        inds = rand(learner.rng, valid_ind_range, n)
        priorities = nothing
    end

    # 2. extract SARTS
    states = consecutive_view(get_trace(t, :state), inds; n_stack = s)
    actions = consecutive_view(get_trace(t, :action), inds)
    next_states = consecutive_view(get_trace(t, :state), inds .+ h; n_stack = s)
    consecutive_rewards = consecutive_view(get_trace(t, :reward), inds; n_horizon = h)
    consecutive_terminals = consecutive_view(get_trace(t, :terminal), inds; n_horizon = h)
    rewards, terminals = zeros(Float32, n), fill(false, n)

    rewards = discount_rewards_reduced(
        consecutive_rewards,
        γ;
        terminal = consecutive_terminals,
        dims = 1,
    )
    terminals = mapslices(any, consecutive_terminals; dims = 1) |> vec

    inds,
    (
        states = states,
        actions = actions,
        rewards = rewards,
        terminals = terminals,
        next_states = next_states,
        priorities = priorities,
    )
end

function RLBase.update!(p::QBasedPolicy{<:PERLearners}, t::AbstractTrajectory)
    learner = p.learner
    length(t) < learner.min_replay_history && return

    learner.update_step += 1

    if learner.update_step % learner.target_update_freq == 0
        copyto!(learner.target_approximator, learner.approximator)
    end

    learner.update_step % learner.update_freq == 0 || return

    inds, experience = extract_experience(t, p.learner)

    if t isa CircularCompactPSARTSATrajectory
        priorities = update!(p.learner, experience)
        get_trace(t, :priority)[inds] .= priorities
    else
        update!(p.learner, experience)
    end
end

function (agent::Agent{<:QBasedPolicy{<:PERLearners},<:CircularCompactPSARTSATrajectory})(
    ::RLCore.Training{PostActStage},
    obs,
)
    push!(
        agent.trajectory;
        reward = get_reward(obs),
        terminal = get_terminal(obs),
        priority = agent.policy.learner.default_priority,
    )
    nothing
end