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


#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)


## 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


# 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


"""
unfold_particle_filter(num_particles::Int, objects_observed::Matrix{Array{Detection2D}}, camera_trajectories::Matrix{Camera_Params}, num_receptive_fields::Int64)

Performs inference procedure.

# Arguments
- v_matrix::Union{Matrix{Float64}, Nothing} V matrix to condition on. If nothing, do inference over V matrix.
- num_particles::Int
- objects_observed: indexed by ``(\\mathrm{scene, frame, receptive field, detection})``
- camera_trajectories: Indexed by ``(\\mathrm{scene, frame})``
- num_receptive_fields::Int64

array for each frame and array for each receptive field and array for those detections

    # Returns
    - A trace containing the inferred visual system and world state.

    """
    function unfold_particle_filter(v_matrix::Union{Matrix{Float64}, Nothing},
        num_particles::Int64,
        mcmc_steps_outer::Int64,
        mcmc_steps_inner::Int64,
        objects_observed::Matrix{Array{Detection2D}},
        camera_trajectories::Matrix{Camera_Params},
        params::Video_Params,
        V_file::IOStream, ws_file::IOStream, order::Vector{Int64})

        lesioned = !isnothing(v_matrix)

        init_obs = Gen.choicemap()

        if lesioned
            for j = 1:params.n_possible_objects
                init_obs[:init_v_matrix => :lambda_fa => j => :fa] = v_matrix[j,1]
                init_obs[:init_v_matrix => :miss_rate => j => :miss] = v_matrix[j,2]
            end
        end

        #no video, no frames
        println("initialize pf")
        state = initialize_particle_filter(main, (lesioned, 0, 0, params), init_obs, num_particles)

        num_videos, num_frames = size(objects_observed)

        inferred_realities = Array{Any}(undef, num_videos)
        avg_v = zeros(params.n_possible_objects, 2)

        global acceptance_counter = 0
        global total = 0

        for v=1:num_videos
            @time begin

                println("v ", v)

                maybe_resample!(state)
                ess = effective_sample_size(normalize_weights(state.log_weights)[2])
                println("ess after resample ", ess)

                #want to do this for every frame at once. put a : instead of f
                #rf is for receptive field
                obs = Gen.choicemap()

                for f = 1:num_frames
                    camera_params = camera_trajectories[v, f]
                    obs[:videos => v => :frame_chain => f => :camera => :camera_location_x] = camera_params.camera_location.x
                    obs[:videos => v => :frame_chain => f => :camera => :camera_location_y] = camera_params.camera_location.y
                    obs[:videos => v => :frame_chain => f => :camera => :camera_location_z] = camera_params.camera_location.z
                    obs[:videos => v => :frame_chain => f => :camera => :camera_focus_x] = camera_params.camera_focus.x
                    obs[:videos => v => :frame_chain => f => :camera => :camera_focus_y] = camera_params.camera_focus.y
                    obs[:videos => v => :frame_chain => f => :camera => :camera_focus_z] = camera_params.camera_focus.z

                    #for rf = 1:num_receptive_fields
                    #println("objects_observed[v, f][rf] ", objects_observed[v, f][rf])
                    #println("type ", typeof(objects_observed[v, f][rf]))
                    #println(typeof(objects_observed[v, f]))
                    #println(typeof(convert(Array{Any, 1}, objects_observed[v, f])))
                    obs[:videos => v => :frame_chain => f => :observations_2D] = convert(Array{Any, 1}, objects_observed[v, f])
                    #end
                end
                #def should be using map to replace for loops here
                #point is, condition on the camera trajectory and on the observations

                if lesioned
                    for j = 1:params.n_possible_objects
                        obs[:videos => v => :v_matrix => :lambda_fa => j => :fa] = v_matrix[j,1]
                        obs[:videos => v => :v_matrix => :miss_rate => j => :miss] = v_matrix[j,2]
                    end
                end

                particle_filter_step!(state, (lesioned, v, num_frames, params), (UnknownChange(),), obs)

                max_particle_weight = maximum(map(i -> get_score(state.traces[i]), collect(1:num_particles)))
                #println("max log score ", max_particle_weight)

                min_particle_weight = minimum(map(i -> get_score(state.traces[i]), collect(1:num_particles)))
                #println("min log score ", min_particle_weight)

                #println("particles")
                # for i = 1:num_particles
                #      println("particle ", i)
                #      trace = state.traces[i]
                #     choice = get_choices(trace)
                #     open("$i.txt", "w") do f
                #         Base.show(f, "text/plain", get_choices(trace))
                #     end
                #
                #      println("score ", get_score(trace))
                # #
                #      if v==2
                #          selection = select(:videos => 2 => :frame_chain => 111 => :observations_2D)
                #          args = get_args(trace)
                #          println("args ", args)
                #          trace2, weight = regenerate(trace, args, (), selection)
                #          println("after regenerate score ", get_score(trace2))
                #          #println("after regenerate score ", weight)
                #          println("after regenerate ", trace2[:videos => 2 => :frame_chain => 111 => :observations_2D])
                #      end
                #
                #
                #     #println("choices")
                #     #println(get_choices(trace))
                #     println("init v ", trace[:init_v_matrix])
                #     println("scene 1 ", trace[:videos => 1 => :init_scene])
                #     println("at v 1 ", trace[:videos => 1 => :v_matrix])
                #     if v==2
                #         println("scene 2 ", trace[:videos => 2 => :init_scene])
                #         println("at v 2 ", trace[:videos => 2 => :v_matrix])
                #     end
                #
                #     #println(trace)
                # end

                ess = effective_sample_size(normalize_weights(state.log_weights)[2])
                println("ess after pf step ", ess)

                #optional rejuvination
                (perturb_params, n_objects_per_category) = get_probs_categories(objects_observed, params, v, num_frames)
                #println("perturb_params ", perturb_params)
                line_segments_per_category = get_line_segments_per_category(params, objects_observed, camera_trajectories, v, num_frames)
                #line_segments = get_line_segments(objects_observed, camera_trajectories, params, v, num_frames, num_receptive_fields, total_n_objects)
                #to-do: threads here
                #seeds = collect(1:num_particles)
                Threads.@threads for i = 1:num_particles
                    #Random.seed!(seeds[i])
                    #for i = 1:num_particles
                    println("score particle i ", i, " before perturbation", get_score(state.traces[i]))
                    println("trace ", state.traces[i][:videos => v => :init_scene])
                    #println("perturb particle i ", i)
                    state.traces[i] = perturb(lesioned, state.traces[i], v, perturb_params, mcmc_steps_outer, mcmc_steps_inner, line_segments_per_category, params)
                    #println("done perturbing i ", i)
                    println("score particle i ", i, " after perturbation", get_score(state.traces[i]))
                    println("trace ", state.traces[i][:videos => v => :init_scene])
                    #println("trace ", state.traces[i][:videos => v => :init_scene])
                    #println("log score of this trace ", get_score(state.traces[i]))
                    #visualize_trace(state.traces, i, camera_trajectories, v, 1, params)

                    #visualize_trace(state.traces, i, camera_trajectories, v, 2, params)
                end

                #visualize_trace_with_heatmap(state.traces, camera_trajectories, v, 1, params )

                min_particle_weight = minimum(map(i -> get_score(state.traces[i]), collect(1:num_particles)))
                #println("min log score ", min_particle_weight)

                max_particle_weight = maximum(map(i -> get_score(state.traces[i]), collect(1:num_particles)))
                #println("max log score ", max_particle_weight)

                ess = effective_sample_size(normalize_weights(state.log_weights)[2])
                println("ess after rejuvination ", ess)

                inferred_realities[v], avg_v = print_Vs_and_Rs_to_file(V_file, ws_file, state.traces, num_particles, params, v, order, v==num_videos)
                #println("avg_v ", avg_v)
                #println("time of v ", v)
            end #end timer
        end

        println(acceptance_counter)
        println(total)

        return (sample_unweighted_traces(state, num_particles), inferred_realities, avg_v)

    end

    """
    perturb(trace, v::Int64, perturb_params::Perturb_Params, line_segments_per_category::Array{Array{Line_Segment,1},1})
    Does 500 MCMC steps (with different proposal functions) on the scene and on the v matrix.

    """

    @gen function perturb(lesioned::Bool,
        trace,
        v::Int64,
        perturb_params::Perturb_Params,
        mcmc_steps_outer::Int,
        mcmc_steps_inner::Int,
        line_segments_per_category::Array{Array{Line_Segment,1},1},
        params::Video_Params)
        #acceptance_counter = 0
        #proposal_counter = 0

        #println("before perturbations")
        #println("v ", v)
        #
        # println("alpha ", get_retval(trace)[end][2][5,2])
        # println("beta ", get_retval(trace)[end][3][5,2])
        #
        # #println("miss 2 ", trace[:init_v_matrix => :miss_rate => 2 => :miss])
        # println("init miss 5 ", trace[:init_v_matrix => :miss_rate => 5 => :miss])
        # #println("miss 2 ", trace[:videos => v => :v_matrix => :miss_rate => 2 => :miss])
        # println("miss rate 5 ", trace[:videos => v => :v_matrix => :miss_rate => 5 => :miss])
        #
        # println("scene at v ", trace[:videos => v => :init_scene])
        n = length(perturb_params.probs_possible_objects)

        for iter=1:mcmc_steps_outer #try 100 MH moves
            #println("iter ", iter)
            #println("trace ", trace[:videos => v => :init_scene])

            trace = perturb_scene(trace, v, perturb_params, line_segments_per_category, params)
            if lesioned == false
                for iter2=1:mcmc_steps_inner
                    trace = perturb_whole_v_matrix_mh(trace, v, n)
                end
            end
            # println("lambda_fa 2 ", trace[:v_matrix => (:lambda_fa, 2)])
            # println("miss 2 ", trace[:v_matrix => (:miss_rate, 2)])
            # println("lambda_fa 5 ", trace[:v_matrix => (:lambda_fa, 5)])
            # println("miss 5 ", trace[:v_matrix => (:miss_rate, 5)])
        end
        #println("acceptance_counter $(acceptance_counter/proposal_counter)")

        #println("after perturbations")
        #println("v ", v)

        # println("alpha ", get_retval(trace)[end][2][5,2])
        # println("beta ", get_retval(trace)[end][3][5,2])
        #
        # #println("miss 2 ", trace[:init_v_matrix => :miss_rate => 2 => :miss])
        # println("init miss 5 ", trace[:init_v_matrix => :miss_rate => 5 => :miss])
        # #println("miss 2 ", trace[:videos => v => :v_matrix => :miss_rate => 2 => :miss])
        # println("miss rate 5 ", trace[:videos => v => :v_matrix => :miss_rate => 5 => :miss])
        #
        # println("scene at v ", trace[:videos => v => :init_scene])
        return trace
    end

    """
    perturb_scene(trace, v::Int64, perturb_params::Perturb_Params, line_segments_per_category::Array{Array{Line_Segment,1},1})
    Does 3 MCMC steps (with different proposal functions) on the scene.

    """

    function perturb_scene(trace, v::Int64, perturb_params::Perturb_Params, line_segments_per_category::Array{Array{Line_Segment,1},1}, params::Video_Params)
        #println("trace ", trace[:videos => v => :init_scene])
        trace, accepted = add_remove_kernel(trace, v, line_segments_per_category, perturb_params)

        #println("accepted? ", accepted)
        #println("trace ", trace[:videos => v => :init_scene])

        #println("trace ", trace[:videos => v => :init_scene])

        #only try changing location or category if there's at least one object in the scene
        if length(trace[:videos => v => :init_scene]) > 0
            trace, accepted = change_location_kernel(trace, v, 0.1, params, line_segments_per_category)

            #global acceptance_counter = acceptance_counter + accepted
            #global total = total + 1

            #println("accepted? ", accepted)
            #println("trace ", trace[:videos => v => :init_scene])

            #trace, accepted = change_category_kernel(trace, v, perturb_params)
            #println("accepted? ", accepted)
            #println("trace ", trace[:videos => v => :init_scene])
        end
        return trace
    end

    #perturb the whole v matrix
    function perturb_whole_v_matrix_mh(trace, v::Int64, n::Int64)
        selection = get_selection_whole(v, n)
        trace, accepted = mh(trace, selection) #proposes new trace from the prior
        return trace
    end

    """
    perturb_v_matrix(trace, perturb_params::Perturb_Params)
    Picks one element of the v matrix to perturb using metropolis hastings.

    """
    #just pick an element of the matrix to perturb
    function perturb_individual_element_v_matrix_mh(trace, v::Int64, n::Int64)
        i = categorical([0.5, 0.5])
        j = categorical(fill(1/n, n))

        selection = get_selection(v, j, i)
        trace, accepted = mh(trace, selection) #proposes new trace from the prior

        # if i == 1 #if perturbing fa
        #     trace, accepted = mh(trace, proposal_for_v_matrix_fa, (v, j))
        # else
        #     trace, accepted = mh(trace, proposal_for_v_matrix_miss, (v, j))
        # end

        #println("accepted? ", accepted)
        #println("trace ", trace[:v_matrix])
        return trace
    end

    #return a selection for one element of the matrix
    function get_selection(v::Int64, j::Int64, i::Int64)
        if i == 1 #change lambda_fa
            if v == 1
                selection = select(:videos => v => :v_matrix => :lambda_fa => j => :fa,
                :init_v_matrix => :lambda_fa => j => :fa)
            else
                selection = select(:videos => v => :v_matrix => :lambda_fa => j => :fa,
                :videos => v-1 => :v_matrix => :lambda_fa => j => :fa)
            end
        else
            if v == 1
                selection = select(:videos => v => :v_matrix => :miss_rate => j => :miss,
                :init_v_matrix => :miss_rate => j => :miss)
            else
                selection = select(:videos => v => :v_matrix => :miss_rate => j => :miss,
                :videos => v-1 => :v_matrix => :miss_rate => j => :miss)
            end
        end
        return selection
    end

    #return a selection for the whole v matrix
    function get_selection_whole(v::Int64, n::Int64)
        selection = select()

        for j = 1:n
            if v == 1
                push!(selection, :videos => v => :v_matrix => :lambda_fa => j => :fa)
                push!(selection, :init_v_matrix => :lambda_fa => j => :fa)
                push!(selection, :v_matrix => :miss_rate => j => :miss)
                push!(selection, :init_v_matrix => :miss_rate => j => :miss)
            else
                push!(selection, :videos => v => :v_matrix => :lambda_fa => j => :fa)
                push!(selection, :videos => v-1 => :v_matrix => :lambda_fa => j => :fa)
                push!(selection, :videos => v => :v_matrix => :miss_rate => j => :miss)
                push!(selection, :videos => v-1 => :v_matrix => :miss_rate => j => :miss)
            end
        end

        return selection
    end

    """
    perturb_v_matrix_hmc(trace, perturb_params::Perturb_Params)
    Perturbs an element of the v_matrix using Hamiltonian MC

    """

    #just pick an element of the matrix to perturb
    function perturb_v_matrix_hmc(trace, v::Int64, perturb_params::Perturb_Params)
        n = length(perturb_params.probs_possible_objects)
        i = categorical([0.5, 0.5])
        j = categorical(fill(1/n, n))

        if i == 1 #if perturbing fa
            trace, accepted = hmc(trace, select(:videos => v => :v_matrix => :lambda_fa => j => :fa))
        else
            trace, accepted = hmc(trace, select(:videos => v => :v_matrix => :miss_rate => j => :miss))
        end

        #println("accepted? ", accepted)
        #println("trace ", trace[:v_matrix])

        return trace
    end

    """
    proposal_for_v_matrix_fa(trace, j::Int64)
    Performs on MH step on the false alarm rate for object of category j.

        """

        @gen function proposal_for_v_matrix_fa(trace, v::Int64, j::Int64)
            std = 0.01
            choices = get_choices(trace)
            #centered on previous value
            @trace(trunc_normal(choices[:videos => v => :v_matrix => :lambda_fa => j => :fa], std, 0.0, 10000.0), :videos => v => :v_matrix => :lambda_fa => j => :fa) #had to make trunc_normal because of error otherwise
        end

        """
        proposal_for_v_matrix_miss(trace, j::Int64)
        Performs on MH step on the miss rate for object of category j.

            """

            @gen function proposal_for_v_matrix_miss(trace, v::Int64, j::Int64)
                std = 0.1
                choices = get_choices(trace)
                #centered on previous value
                @trace(trunc_normal(choices[:videos => v => :v_matrix => :miss_rate => j => :miss], std, 0.0, 1.0), :videos => v => :v_matrix => :miss_rate => j => :miss)
            end

            """
            get_probs_categories(objects_observed::Matrix{Array{Detection2D}}, params::Video_Params, v::Int64, num_frames::Int64, num_receptive_fields::Int64)

            Returns a probability distribution over the object categories to be used
            in the proposal functions.

            Derived from the objects that the visual system observes.
            """

            function get_probs_categories(objects_observed::Matrix{Array{Detection2D}}, params::Video_Params, v::Int64, num_frames::Int64)
                track_categories = zeros(length(params.possible_objects)) #each element will be the number of times that category was detected. adding 1
                probabilities = zeros(length(params.possible_objects))
                for f = 1:num_frames
                    #for rf = 1:num_receptive_fields
                    for (index, value) in enumerate(objects_observed[v, f])
                        track_categories[value[3]] = track_categories[value[3]]+1#category
                    end
                    #end
                end
                #make it so that 10% of the weight gets divided evenly among all the categories
                other_ten_percent = sum(track_categories)/9
                each = other_ten_percent/length(params.possible_objects)
                probabilities = copy(track_categories)
                probabilities = probabilities .+ each

                #in case of no observations, make uniform
                if sum(probabilities) == 0
                    probabilities .= 1
                end

                return (Perturb_Params(probs_possible_objects = (probabilities)./sum(probabilities)), track_categories)
            end

            """
            get_line_segments_per_category(params::Video_Params, objects_observed::Matrix{Array{Detection2D}}, camera_trajectories::Matrix{Camera_Params}, v::Int64, num_frames::Int64, num_receptive_fields::Int64)

            Gets all line segments in 3d space that correspond to each 2D detection
            in every frame of a scene.
            """
            function get_line_segments_per_category(params::Video_Params, objects_observed::Matrix{Array{Detection2D}}, camera_trajectories::Matrix{Camera_Params}, v::Int64, num_frames::Int64)
                line_segments = Array{Array{Line_Segment, 1}}(undef, length(params.possible_objects))
                for j = 1:length(params.possible_objects)
                    line_segments[j] = []
                end
                for f = 1:num_frames
                    camera_params = camera_trajectories[v, f]
                    #for rf = 1:num_receptive_fields
                    for (index, value) in enumerate(objects_observed[v, f])
                        #println("value ", value)
                        line_segment = get_line_segment(camera_params, params, value)
                        #println("line_segment in pf_inference ", line_segment)
                        push!(line_segments[value[3]], line_segment)
                    end
                    #end
                end
                return line_segments
            end

            """Duplicated from Gen library"""
            function effective_sample_size(log_normalized_weights::Vector{Float64})
                log_ess = -logsumexp(2. * log_normalized_weights)
                return exp(log_ess)
            end

            """Duplicated from Gen Library"""
            function normalize_weights(log_weights::Vector{Float64})
                log_total_weight = logsumexp(log_weights)
                log_normalized_weights = log_weights .- log_total_weight
                return (log_total_weight, log_normalized_weights)
            end

            export unfold_particle_filter