text:
function set_acescore(ta)
    atom_type = Array{String}(undef, ta.natom)
    ace_score = Array{Float64}(undef, ta.natom)

    for iatom = 1:ta.natom
        # ATOM TYPE "N"
        if ta.atomname[iatom] == "N"
            atom_type[iatom] = "N"
            ace_score[iatom] = -0.495

            # ATOM TYPE "C"
        elseif ta.atomname[iatom] == "C"
            atom_type[iatom] = "C"
            ace_score[iatom] = -0.464

            # ATOM TYPE "O"
        elseif ta.atomname[iatom] == "O" || ta.atomname[iatom] == "OXT"
            atom_type[iatom] = "O"
            ace_score[iatom] = -0.079

            # ATOM TYPE "GCA"
        elseif ta.resname[iatom] == "GLY" && ta.atomname[iatom] == "CA"
            atom_type[iatom] = "GCA"
            ace_score[iatom] = -0.008

            # ATOM TYPE "CA"
        elseif ta.atomname[iatom] == "CA"
            atom_type[iatom] = "CA"
            ace_score[iatom] = -0.553

            # ATOM TYPE "CB"
        elseif ta.resname[iatom] ==  "ALA" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "ASN" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "ASP" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "CYS" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "GLN" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "GLU" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "HIS" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "ILE" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "LEU" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "LYS" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "MET" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "PHE" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "PRO" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "PRO" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "PRO" && ta.atomname[iatom] == "CD"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "THR" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "TYR" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353
        elseif ta.resname[iatom] ==  "VAL" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "CB"
            ace_score[iatom] = -0.353

            # ATOM TYPE "KNZ"
        elseif ta.resname[iatom] ==  "LYS" && ta.atomname[iatom] == "CE"
            atom_type[iatom] = "KNZ"
            ace_score[iatom] = 1.334
        elseif ta.resname[iatom] ==  "LYS" && ta.atomname[iatom] == "NZ"
            atom_type[iatom] = "KNZ"
            ace_score[iatom] = 1.334

            # ATOM TYPE "KCD"
        elseif ta.resname[iatom] ==  "LYS" && ta.atomname[iatom] == "CD"
            atom_type[iatom] = "KCD"
            ace_score[iatom] = 1.046

            # ATOM TYPE "DOD"
        elseif ta.resname[iatom] ==  "ASP" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "DOD"
            ace_score[iatom] = 0.933
        elseif ta.resname[iatom] ==  "ASP" && ta.atomname[iatom] == "OD1"
            atom_type[iatom] = "DOD"
            ace_score[iatom] = 0.933
        elseif ta.resname[iatom] ==  "ASP" && ta.atomname[iatom] == "OD2"
            atom_type[iatom] = "DOD"
            ace_score[iatom] = 0.933
        elseif ta.resname[iatom] ==  "GLU" && ta.atomname[iatom] == "CD"
            atom_type[iatom] = "DOD"
            ace_score[iatom] = 0.933
        elseif ta.resname[iatom] ==  "GLU" && ta.atomname[iatom] == "OE1"
            atom_type[iatom] = "DOD"
            ace_score[iatom] = 0.933
        elseif ta.resname[iatom] ==  "GLU" && ta.atomname[iatom] == "OE2"
            atom_type[iatom] = "DOD"
            ace_score[iatom] = 0.933

            # ATOM TYPE "RNH"
        elseif ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "CZ"
            atom_type[iatom] = "RNH"
            ace_score[iatom] = 0.726
        elseif ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "NH1"
            atom_type[iatom] = "RNH"
            ace_score[iatom] = 0.726
        elseif ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "NH2"
            atom_type[iatom] = "RNH"
            ace_score[iatom] = 0.726

            # ATOM TYPE "NND"
        elseif ta.resname[iatom] ==  "ASN" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "NND"
            ace_score[iatom] = 0.693
        elseif ta.resname[iatom] ==  "ASN" && ta.atomname[iatom] == "OD1"
            atom_type[iatom] = "NND"
            ace_score[iatom] = 0.693
        elseif ta.resname[iatom] ==  "ASN" && ta.atomname[iatom] == "ND2"
            atom_type[iatom] = "NND"
            ace_score[iatom] = 0.693
        elseif ta.resname[iatom] ==  "GLN" && ta.atomname[iatom] == "CD"
            atom_type[iatom] = "NND"
            ace_score[iatom] = 0.693
        elseif ta.resname[iatom] ==  "GLN" && ta.atomname[iatom] == "OE1"
            atom_type[iatom] = "NND"
            ace_score[iatom] = 0.693
        elseif ta.resname[iatom] ==  "GLN" && ta.atomname[iatom] == "NE2"
            atom_type[iatom] = "NND"
            ace_score[iatom] = 0.693

            # ATOM TYPE "RNE"
        elseif ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "CD"
            atom_type[iatom] = "RNE"
            ace_score[iatom] = 0.606
        elseif ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "NE"
            atom_type[iatom] = "RNE"
            ace_score[iatom] = 0.606

            # ATOM TYPE "SOG"
        elseif ta.resname[iatom] ==  "SER" && ta.atomname[iatom] == "CB"
            atom_type[iatom] = "SOG"
            ace_score[iatom] = 0.232
        elseif ta.resname[iatom] ==  "SER" && ta.atomname[iatom] == "OG"
            atom_type[iatom] = "SOG"
            ace_score[iatom] = 0.232
        elseif ta.resname[iatom] ==  "THR" && ta.atomname[iatom] == "OG1"
            atom_type[iatom] = "SOG"
            ace_score[iatom] = 0.232
        elseif ta.resname[iatom] ==  "TYR" && ta.atomname[iatom] == "OH"
            atom_type[iatom] = "SOG"
            ace_score[iatom] = 0.232

            # ATOM TYPE "HNE"
        elseif ta.resname[iatom] ==  "HIS" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "HNE"
            ace_score[iatom] = 0.061
        elseif ta.resname[iatom] ==  "HIS" && ta.atomname[iatom] == "ND1"
            atom_type[iatom] = "HNE"
            ace_score[iatom] = 0.061
        elseif ta.resname[iatom] ==  "HIS" && ta.atomname[iatom] == "CD2"
            atom_type[iatom] = "HNE"
            ace_score[iatom] = 0.061
        elseif ta.resname[iatom] ==  "HIS" && ta.atomname[iatom] == "CE1"
            atom_type[iatom] = "HNE"
            ace_score[iatom] = 0.061
        elseif ta.resname[iatom] ==  "HIS" && ta.atomname[iatom] == "NE2"
            atom_type[iatom] = "HNE"
            ace_score[iatom] = 0.061
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "NE1"
            atom_type[iatom] = "HNE"
            ace_score[iatom] = 0.061

            # ATOM TYPE "YCZ"
        elseif ta.resname[iatom] ==  "TYR" && ta.atomname[iatom] == "CE1"
            atom_type[iatom] = "YCZ"
            ace_score[iatom] = -0.289
        elseif ta.resname[iatom] ==  "TYR" && ta.atomname[iatom] == "CE2"
            atom_type[iatom] = "YCZ"
            ace_score[iatom] = -0.289
        elseif ta.resname[iatom] ==  "TYR" && ta.atomname[iatom] == "CZ"
            atom_type[iatom] = "YCZ"
            ace_score[iatom] = -0.289

            # ATOM TYPE "FCZ"
        elseif ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "GLN" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "GLU" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "ILE" && ta.atomname[iatom] == "CG1"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "LEU" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "LYS" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "MET" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "MET" && ta.atomname[iatom] == "SD"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "PHE" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "PHE" && ta.atomname[iatom] == "CD1"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "PHE" && ta.atomname[iatom] == "CD2"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "PHE" && ta.atomname[iatom] == "CE1"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "PHE" && ta.atomname[iatom] == "CE2"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "PHE" && ta.atomname[iatom] == "CZ"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "THR" && ta.atomname[iatom] == "CG2"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CD1"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CD2"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CE2"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CE3"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CZ2"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CZ3"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TRP" && ta.atomname[iatom] == "CH2"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TYR" && ta.atomname[iatom] == "CG"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TYR" && ta.atomname[iatom] == "CD1"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432
        elseif ta.resname[iatom] ==  "TYR" && ta.atomname[iatom] == "CD2"
            atom_type[iatom] = "FCZ"
            ace_score[iatom] = -0.432

            # ATOM TYPE "LCD"
        elseif ta.resname[iatom] ==  "ILE" && ta.atomname[iatom] == "CG2"
            atom_type[iatom] = "LCD"
            ace_score[iatom] = -0.987
        elseif ta.resname[iatom] ==  "ILE" && ta.atomname[iatom] == "CD"
            atom_type[iatom] = "LCD"
            ace_score[iatom] = -0.987
        elseif ta.resname[iatom] ==  "ILE" && ta.atomname[iatom] == "CD1"
            atom_type[iatom] = "LCD"
            ace_score[iatom] = -0.987
        elseif ta.resname[iatom] ==  "LEU" && ta.atomname[iatom] == "CD1"
            atom_type[iatom] = "LCD"
            ace_score[iatom] = -0.987
        elseif ta.resname[iatom] ==  "LEU" && ta.atomname[iatom] == "CD2"
            atom_type[iatom] = "LCD"
            ace_score[iatom] = -0.987
        elseif ta.resname[iatom] ==  "MET" && ta.atomname[iatom] == "CE"
            atom_type[iatom] = "LCD"
            ace_score[iatom] = -0.987
        elseif ta.resname[iatom] ==  "VAL" && ta.atomname[iatom] == "CG1"
            atom_type[iatom] = "LCD"
            ace_score[iatom] = -0.987
        elseif ta.resname[iatom] ==  "VAL" && ta.atomname[iatom] == "CG2"
            atom_type[iatom] = "LCD"
            ace_score[iatom] = -0.987

            # ATOM TYPE "CSG"
        elseif ta.resname[iatom] ==  "CYS" && ta.atomname[iatom] == "SG"
            atom_type[iatom] = "CSG"
            ace_score[iatom] = -1.827
        else
            println("error: faled to assign atom type " * ta.atomname[iatom] * "-" * ta.resname[iatom])
        end
    end

    return TrjArray(ta, mass=ace_score)
end

function golden_section_spiral(n)
    points = zeros(Float64, n, 3)
    inc = pi * (3.0 - sqrt(5.0))
    offset = 2.0 / Float64(n)
    for k = 1:n
        y = (k-1) * offset - 1.0 + (offset / 2.0)
        r = sqrt(1.0 - y*y)
        phi = (k-1) * inc
        points[k, 1] = cos(phi) * r
        points[k, 2] = y
        points[k, 3] = sin(phi) * r
    end
    return points
end

function set_radius(ta::TrjArray{T, U}) where {T, U}
    radius = Array{T}(undef, ta.natom)
    element = Array{String}(undef, ta.natom)
    for iatom = 1:ta.natom
        if !isnothing(match(r"H.*", ta.atomname[iatom]))
            element[iatom] = "H"
        elseif !isnothing(match(r"C.*", ta.atomname[iatom]))
            element[iatom] = "C"
        elseif !isnothing(match(r"N.*", ta.atomname[iatom]))
            element[iatom] = "N"
        elseif !isnothing(match(r"O.*", ta.atomname[iatom]))
            element[iatom] = "O"
        elseif !isnothing(match(r"S.*", ta.atomname[iatom]))
            element[iatom] = "S"
        else
            error("failed to assign element: " * ta.atomname[iatom])
        end
    end

    radius_dict = Dict("H" => 1.20,
    "C" => 1.70,
    "N" => 1.55,
    "O" => 1.52,
    "S" => 1.80)

    for iatom = 1:ta.natom
        radius[iatom] = radius_dict[element[iatom]]
    end

    return TrjArray(ta, radius=radius)
end

function set_charge(ta)
    charge = Array{Float64}(undef, ta.natom)

    is_first = true
    for iatom = 1:ta.natom
        # ATOM TYPE "N"
        if ta.atomname[iatom] == "N"
            if is_first
                charge[iatom] = 1.0
                is_first = false
            else
                charge[iatom] = 0.5
            end

            # ATOM TYPE "O"
        elseif ta.atomname[iatom] == "O"
            charge[iatom] = -0.5
        elseif  ta.atomname[iatom] == "OXT"
            charge[iatom] = -1.0

        elseif  ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "NH1"
            charge[iatom] = 0.5
        elseif ta.resname[iatom] ==  "ARG" && ta.atomname[iatom] == "NH2"
            charge[iatom] = 0.5
        elseif ta.resname[iatom] ==  "GLU" && ta.atomname[iatom] == "OE1"
            charge[iatom] = -0.5
        elseif ta.resname[iatom] ==  "GLU" && ta.atomname[iatom] == "OE2"
            charge[iatom] = -0.5
        elseif ta.resname[iatom] ==  "ASP" && ta.atomname[iatom] == "OD1"
            charge[iatom] = -0.5
        elseif ta.resname[iatom] ==  "ASP" && ta.atomname[iatom] == "OD2"
            charge[iatom] = -0.5
        elseif ta.resname[iatom] ==  "LYS" && ta.atomname[iatom] == "NZ"
            charge[iatom] = 1.0
        elseif ta.resname[iatom] ==  "PRO" && ta.atomname[iatom] == "N"
            charge[iatom] = -0.1

        else
            charge[iatom] = 0.0
        end
    end

    return TrjArray(ta, charge=charge)
end

function compute_sasa(ta::TrjArray{T, U}, probe_radius=1.4::T; npoint=960::Int, iframe=1::Int, candicate = 10) where {T, U}
    # construct neighbor rist
    max_radius = 2.0 * maximum(ta.radius) + 2.0 * probe_radius ############# TODO
    pairlist = compute_pairlist(ta[iframe, :], max_radius)
    neighbor_list = []
    for iatom in 1:ta.natom
        push!(neighbor_list, Array{U}(undef, 0))
    end
    for ipair in 1:size(pairlist.pair, 1)
        i = pairlist.pair[ipair, 1]
        j = pairlist.pair[ipair, 2]
        push!(neighbor_list[i], j)
        push!(neighbor_list[j], i)
    end

    # generate uniform points on a unit sphere
    points = golden_section_spiral(npoint)

    # compute the ratio of exposed area for each sphere
    sasa = Array{T}(undef, ta.natom)
    for iatom = 1:ta.natom
        n_accessible_point = 0
        neighbor_list_iatom = neighbor_list[iatom]
        for ipoint in 1:npoint
            is_accessible = true
            point = points[ipoint, :] .* (ta.radius[iatom] + probe_radius)
            point[1] += ta.xyz[iframe, 3*(iatom-1)+1]
            point[2] += ta.xyz[iframe, 3*(iatom-1)+2]
            point[3] += ta.xyz[iframe, 3*(iatom-1)+3]
            for j in 1:length(neighbor_list_iatom)
                jatom = neighbor_list_iatom[j]
                dx = point[1] - ta.xyz[iframe, 3*(jatom-1)+1]
                dy = point[2] - ta.xyz[iframe, 3*(jatom-1)+2]
                dz = point[3] - ta.xyz[iframe, 3*(jatom-1)+3]
                if !isempty(ta.boxsize)
                    dx = dx - round(dx/ta.boxsize[iframe, 1])*ta.boxsize[iframe, 1]
                    dy = dy - round(dy/ta.boxsize[iframe, 2])*ta.boxsize[iframe, 2]
                    dz = dz - round(dz/ta.boxsize[iframe, 3])*ta.boxsize[iframe, 3]
                end
                d = sqrt(dx^2 + dy^2 + dz^2)
                if d < (ta.radius[jatom] + probe_radius)
                    is_accessible = false
                    break
                end
            end
            if is_accessible
                n_accessible_point += 1
            end
        end
        sasa[iatom] = 4.0 * pi * (ta.radius[iatom] + probe_radius)^2 * n_accessible_point / npoint
    end

    return TrjArray(ta, sasa=sasa)
end

function assign_shape_complementarity!(grid, ta::TrjArray{T, U}, grid_space,
    rcut1, rcut2, x_grid, y_grid, z_grid, iframe) where {T,U}
    grid .= 0.0 + 0.0im
    nx, ny, nz = size(grid)

    for iatom = 1:ta.natom
        rcut = rcut1[iatom]
        if rcut < 0.0
            continue
        end

        x = ta.xyz[iframe, 3*(iatom-1)+1]
        y = ta.xyz[iframe, 3*(iatom-1)+2]
        z = ta.xyz[iframe, 3*(iatom-1)+3]

        dx = x - x_grid[1]
        ix_min = floor(U, (dx - rcut)/grid_space) + 1
        ix_min = max(ix_min, 1)
        ix_max = floor(U, (dx + rcut)/grid_space) + 2
        ix_max = min(ix_max, nx)

        dy = y - y_grid[1]
        iy_min = floor(U, (dy - rcut)/grid_space) + 1
        iy_min = max(iy_min, 1)
        iy_max = floor(U, (dy + rcut)/grid_space) + 2
        iy_max = min(iy_max, ny)

        dz = z - z_grid[1]
        iz_min = floor(U, (dz - rcut)/grid_space) + 1
        iz_min = max(iz_min, 1)
        iz_max = floor(U, (dz + rcut)/grid_space) + 2
        iz_max = min(iz_max, nz)

        for ix = ix_min:ix_max
            for iy = iy_min:iy_max
                for iz = iz_min:iz_max
                    if imag(grid[ix, iy, iz]) < 0.0001
                        dist = sqrt((x - x_grid[ix])^2 + (y - y_grid[iy])^2 + (z - z_grid[iz])^2)
                        if dist < rcut
                            grid[ix, iy, iz] = 0.0 + 9.0im
                        end
                    end
                end
            end
        end
    end

    for iatom = 1:ta.natom
        rcut = rcut2[iatom]
        if rcut < 0.0
            continue
        end

        x = ta.xyz[iframe, 3*(iatom-1)+1]
        y = ta.xyz[iframe, 3*(iatom-1)+2]
        z = ta.xyz[iframe, 3*(iatom-1)+3]

        dx = x - x_grid[1]
        ix_min = floor(U, (dx - rcut)/grid_space) + 1
        ix_min = max(ix_min, 1)
        ix_max = floor(U, (dx + rcut)/grid_space) + 2
        ix_max = min(ix_max, nx)

        dy = y - y_grid[1]
        iy_min = floor(U, (dy - rcut)/grid_space) + 1
        iy_min = max(iy_min, 1)
        iy_max = floor(U, (dy + rcut)/grid_space) + 2
        iy_max = min(iy_max, ny)

        dz = z - z_grid[1]
        iz_min = floor(U, (dz - rcut)/grid_space) + 1
        iz_min = max(iz_min, 1)
        iz_max = floor(U, (dz + rcut)/grid_space) + 2
        iz_max = min(iz_max, nz)

        for ix = ix_min:ix_max
            for iy = iy_min:iy_max
                for iz = iz_min:iz_max
                    if imag(grid[ix, iy, iz]) < 0.0001
                        dist = sqrt((x - x_grid[ix])^2 + (y - y_grid[iy])^2 + (z - z_grid[iz])^2)
                        if dist < rcut
                            grid[ix, iy, iz] = 1.0 + 0.0im
                        end
                    end
                end
            end
        end
    end
end

function compute_docking_score_with_fft(quaternion, grid_RSC, grid_LSC, ligand2, grid_space, rcut1, rcut2, x_grid, y_grid, z_grid, iframe, tops, iq)
    ligand2_rotated = rotate(ligand2, quaternion)
    assign_shape_complementarity!(grid_LSC, ligand2_rotated, grid_space, rcut1, rcut2, x_grid, y_grid, z_grid, iframe)
    #grid_LSC .= grid_LSC[end:-1:1, end:-1:1, end:-1:1]

    if CUDA.functional()
        grid_RSC_gpu = cu(grid_RSC)
        grid_LSC_gpu = cu(grid_LSC)
        t_gpu = ifft(fft(grid_RSC_gpu) .* conj.(fft(conj.(grid_LSC_gpu))))
        #score_gpu = real(t_gpu) .- imag(t_gpu)
        score_gpu = real(t_gpu)
        score = Array(score_gpu)
    else
        t = ifft(fft(grid_RSC) .* conj.(fft(conj.(grid_LSC))))
        #score = real(t) .- imag(t)
        score = real(t)
    end

    ret = []
    for t in 1:tops
        id = argmax(score)
        dx_estimated = id[1]
        dy_estimated = id[2]
        dz_estimated = id[3]
        push!(ret, (score[id], dx_estimated, dy_estimated, dz_estimated, iq))
        score[id] = -Inf
    end

    return ret
end

function compute_docking_score_on_xyplane(grid_RSC, grid_LSC)
    if CUDA.functional()
        grid_RSC_gpu = cu(grid_RSC)
        grid_LSC_gpu = cu(grid_LSC)
        t_gpu = ifft(fft(grid_RSC_gpu) .* conj.(fft(conj.(grid_LSC_gpu))))
        score_gpu = real(t_gpu)
        score = Array(score_gpu)
    else
        t = ifft(fft(grid_RSC) .* conj.(fft(conj.(grid_LSC))))
        score = real(t)
    end

    return score[:, :, 1]
end

function spread_nearest!(grid, x, y, z, grid_x, grid_y, grid_z)
    T = eltype(grid)
    natom = length(x)
    grid .= zero(T)

    for iatom = 1:natom
        dx = x[iatom] .- grid_x
        dy = y[iatom] .- grid_y
        dz = z[iatom] .- grid_z
        ix = argmin(abs.(dx))
        iy = argmin(abs.(dy))
        iz = argmin(abs.(dz))
        grid[ix, iy, iz] = T(1.0*im)
    end

    return
end

function spread_nearest_gpu!(grid, x, y, z, grid_x, grid_y, grid_z)
    T = eltype(grid)
    natom = length(x)
    ngid = gridDim().x*blockDim().x
    gid = (blockIdx().x - 1)*blockDim().x + threadIdx().x - 1

    grid_x_min = grid_x[1]
    grid_y_min = grid_y[1]
    grid_z_min = grid_z[1]

    grid_x_delta = grid_x[2] - grid_x[1]
    grid_y_delta = grid_y[2] - grid_y[1]
    grid_z_delta = grid_z[2] - grid_z[1]

    grid_x_delta_half = grid_x_delta*0.5
    grid_y_delta_half = grid_y_delta*0.5
    grid_z_delta_half = grid_z_delta*0.5

    ntile = ceil(Int, natom/ngid)

    for itile = 1:ntile
        iatom = ngid*(itile-1) + gid + 1
        if iatom > natom
            return nothing
        end
        atom_x = x[iatom]
        atom_y = y[iatom]
        atom_z = z[iatom]

        ix = ceil(Int, (atom_x - grid_x_min)/grid_x_delta + grid_x_delta_half) + 1
        iy = ceil(Int, (atom_y - grid_y_min)/grid_y_delta + grid_y_delta_half) + 1
        iz = ceil(Int, (atom_z - grid_z_min)/grid_z_delta + grid_z_delta_half) + 1
        grid[ix,iy,iz] = T(1.0)
    end

    return nothing
end

function spread_neighbors_add!(grid, x, y, z, grid_x, grid_y, grid_z, rcut, charge)
    T = eltype(grid)
    natom = length(x)
    nx, ny, nz = size(grid)

    for iatom = 1:natom
        for ix = 1:nx
            dx = x[iatom] - grid_x[ix]
            if abs(dx) > rcut[iatom]
                continue
            end
            for iy = 1:ny
                dy = y[iatom] - grid_y[iy]
                if abs(dy) > rcut[iatom]
                    continue
                end
                for iz = 1:nz
                    dz = z[iatom] - grid_z[iz]
                    if abs(dz) > rcut[iatom]
                        continue
                    end
                    d = dx*dx + dy*dy + dz*dz
                    if d < rcut[iatom]*rcut[iatom]
                        grid[ix, iy, iz] += charge[iatom]
                    end
                end
            end
        end
    end

    return
end

function spread_neighbors_add_gpu!(grid, x, y, z, grid_x, grid_y, grid_z, rcut, charge)
    T = eltype(grid)
    natom = length(x)
    ngid = gridDim().x*blockDim().x
    gid = (blockIdx().x - 1)*blockDim().x + threadIdx().x - 1

    grid_x_min = grid_x[1]
    grid_y_min = grid_y[1]
    grid_z_min = grid_z[1]

    grid_x_delta = grid_x[2] - grid_x[1]
    grid_y_delta = grid_y[2] - grid_y[1]
    grid_z_delta = grid_z[2] - grid_z[1]

    nx, ny, nz = size(grid)

    ntile = ceil(Int, natom/ngid)

    for itile = 1:ntile
        iatom = ngid*(itile-1) + gid + 1
        if iatom > natom
            return
        end
        atom_x = x[iatom]
        atom_y = y[iatom]
        atom_z = z[iatom]
        r = rcut[iatom]
        r2 = r*r
        c = charge[iatom]

        ix_min = max(floor(Int, (atom_x - r - grid_x_min)/grid_x_delta) - 1, 1)
        iy_min = max(floor(Int, (atom_y - r - grid_y_min)/grid_y_delta) - 1, 1)
        iz_min = max(floor(Int, (atom_z - r - grid_z_min)/grid_z_delta) - 1, 1)

        ix_max = min(floor(Int, (atom_x + r - grid_x_min)/grid_x_delta) + 2, nx)
        iy_max = min(floor(Int, (atom_y + r - grid_y_min)/grid_y_delta) + 2, ny)
        iz_max = min(floor(Int, (atom_z + r - grid_z_min)/grid_z_delta) + 2, nz)

        for ix = ix_min:ix_max
            for iy = iy_min:iy_max
                for iz = iz_min:iz_max
                    xx = grid_x_min + grid_x_delta*(ix-1)
                    yy = grid_y_min + grid_y_delta*(iy-1)
                    zz = grid_z_min + grid_z_delta*(iz-1)
                    d = (atom_x - xx)^2 + (atom_y - yy)^2 + (atom_z - zz)^2
                    if d < r2
                        #@cuprintln("called1")
                        #@cuprintln("c = $(c)")
                        #@cuprintln("called2")
                        #@cuprintln("grid[ix,iy,iz] = $(grid[ix,iy,iz])")
                        CUDA.@atomic grid[ix,iy,iz] += c
                    end
                end
            end
        end
    end

    return
end

function spread_neighbors_substitute!(grid, x, y, z, grid_x, grid_y, grid_z, rcut, charge)
    T = eltype(grid)
    natom = length(x)
    nx, ny, nz = size(grid)

    for iatom = 1:natom
        for ix = 1:nx
            dx = x[iatom] - grid_x[ix]
            if abs(dx) > rcut[iatom]
                continue
            end
            for iy = 1:ny
                dy = y[iatom] - grid_y[iy]
                if abs(dy) > rcut[iatom]
                    continue
                end
                for iz = 1:nz
                    dz = z[iatom] - grid_z[iz]
                    if abs(dz) > rcut[iatom]
                        continue
                    end
                    d = dx*dx + dy*dy + dz*dz
                    if d < rcut[iatom]*rcut[iatom]
                        grid[ix, iy, iz] = charge[iatom]
                    end
                end
            end
        end
    end

    return
end

function spread_neighbors_substitute_gpu!(grid, x, y, z, grid_x, grid_y, grid_z, rcut, charge)
    T = eltype(grid)
    natom = length(x)
    ngid = gridDim().x*blockDim().x
    gid = (blockIdx().x - 1)*blockDim().x + threadIdx().x - 1

    grid_x_min = grid_x[1]
    grid_y_min = grid_y[1]
    grid_z_min = grid_z[1]

    grid_x_delta = grid_x[2] - grid_x[1]
    grid_y_delta = grid_y[2] - grid_y[1]
    grid_z_delta = grid_z[2] - grid_z[1]

    nx, ny, nz = size(grid)

    ntile = ceil(Int, natom/ngid)

    for itile = 1:ntile
        iatom = ngid*(itile-1) + gid + 1
        if iatom > natom
            return
        end
        atom_x = x[iatom]
        atom_y = y[iatom]
        atom_z = z[iatom]
        r = rcut[iatom]
        r2 = r*r
        c = charge[iatom]

        ix_min = max(floor(Int, (atom_x - r - grid_x_min)/grid_x_delta) - 1, 1)
        iy_min = max(floor(Int, (atom_y - r - grid_y_min)/grid_y_delta) - 1, 1)
        iz_min = max(floor(Int, (atom_z - r - grid_z_min)/grid_z_delta) - 1, 1)

        ix_max = min(floor(Int, (atom_x + r - grid_x_min)/grid_x_delta) + 2, nx)
        iy_max = min(floor(Int, (atom_y + r - grid_y_min)/grid_y_delta) + 2, ny)
        iz_max = min(floor(Int, (atom_z + r - grid_z_min)/grid_z_delta) + 2, nz)

        for ix = ix_min:ix_max
            for iy = iy_min:iy_max
                for iz = iz_min:iz_max
                    xx = grid_x_min + grid_x_delta*(ix-1)
                    yy = grid_y_min + grid_y_delta*(iy-1)
                    zz = grid_z_min + grid_z_delta*(iz-1)
                    d = (atom_x - xx)^2 + (atom_y - yy)^2 + (atom_z - zz)^2
                    if d < r2
                        #@cuprintln("called")
                        #@cuprintln("grid[ix,iy,iz] = $(grid[ix,iy,iz])")
                        #@cuprintln("c = $(c)")
                        grid[ix,iy,iz] = c
                    end
                end
            end
        end
    end

    return
end

function rotate_with_matrix!(x, y, z, R)
    natom = length(x)
    for iatom = 1:natom
        x_curr = x[iatom]
        y_curr = y[iatom]
        z_curr = z[iatom]
        x_new = R[1, 1]*x_curr + R[1, 2]*y_curr + R[1, 3]*z_curr
        y_new = R[2, 1]*x_curr + R[2, 2]*y_curr + R[2, 3]*z_curr
        z_new = R[3, 1]*x_curr + R[3, 2]*y_curr + R[3, 3]*z_curr
        x[iatom] = x_new
        y[iatom] = y_new
        z[iatom] = z_new
    end
    return
end

#function rotate_with_matrix_gpu!(x, y, z, R)
#    natom = length(x)
#    ngid = gridDim().x*blockDim().x
#    gid = (blockIdx().x - 1)*blockDim().x + threadIdx().x - 1
#
#    ntile = ceil(Int, natom/ngid)
#    for itile = 1:ntile
#        iatom = ngid*(itile-1) + gid + 1
#        if iatom > natom
#            return
#        end
#        x_curr = x[iatom]
#        y_curr = y[iatom]
#        z_curr = z[iatom]
#        x_new = R[1, 1]*x_curr + R[1, 2]*y_curr + R[1, 3]*z_curr
#        y_new = R[2, 1]*x_curr + R[2, 2]*y_curr + R[2, 3]*z_curr
#        z_new = R[3, 1]*x_curr + R[3, 2]*y_curr + R[3, 3]*z_curr
#        x[iatom] = x_new
#        y[iatom] = y_new
#        z[iatom] = z_new
#    end
#    return
#end

function rotate!(x, y, z, q::AbstractVector{T}) where {T, U}
    natom = length(x)
    r1 = 1.0 - 2.0 * q[2] * q[2] - 2.0 * q[3] * q[3]
    r2 = 2.0 * (q[1] * q[2] + q[3] * q[4])
    r3 = 2.0 * (q[1] * q[3] - q[2] * q[4])
    r4 = 2.0 * (q[1] * q[2] - q[3] * q[4])
    r5 = 1.0 - 2.0 * q[1] * q[1] - 2.0 * q[3] * q[3]
    r6 = 2.0 * (q[2] * q[3] + q[1] * q[4])
    r7 = 2.0 * (q[1] * q[3] + q[2] * q[4])
    r8 = 2.0 * (q[2] * q[3] - q[1] * q[4])
    r9 = 1.0 - 2.0 * q[1] * q[1] - 2.0 * q[2] * q[2]
    for iatom = 1:natom
        x_new = r1 * x[iatom] + r2 * y[iatom] + r3 * z[iatom]
        y_new = r4 * x[iatom] + r5 * y[iatom] + r6 * z[iatom]
        z_new = r7 * x[iatom] + r8 * y[iatom] + r9 * z[iatom]
        x[iatom] = x_new
        y[iatom] = y_new
        z[iatom] = z_new
    end
    return nothing
end

function rotate_with_matrix!(x, y, z, R::AbstractMatrix{T}) where {T, U}
    natom = length(x)
    r1 = R[1, 1]
    r2 = R[1, 2]
    r3 = R[1, 3]
    r4 = R[2, 1]
    r5 = R[2, 2]
    r6 = R[2, 3]
    r7 = R[3, 1]
    r8 = R[3, 2]
    r9 = R[3, 3]
    for iatom = 1:natom
        x_new = r1 * x[iatom] + r2 * y[iatom] + r3 * z[iatom]
        y_new = r4 * x[iatom] + r5 * y[iatom] + r6 * z[iatom]
        z_new = r7 * x[iatom] + r8 * y[iatom] + r9 * z[iatom]
        x[iatom] = x_new
        y[iatom] = y_new
        z[iatom] = z_new
    end
    return nothing
end

function rotate_gpu!(x, y, z, q)
    natom = length(x)
    ngid = gridDim().x*blockDim().x
    gid = (blockIdx().x - 1)*blockDim().x + threadIdx().x - 1

    r1 = 1.0 - 2.0 * q[2] * q[2] - 2.0 * q[3] * q[3]
    r2 = 2.0 * (q[1] * q[2] + q[3] * q[4])
    r3 = 2.0 * (q[1] * q[3] - q[2] * q[4])
    r4 = 2.0 * (q[1] * q[2] - q[3] * q[4])
    r5 = 1.0 - 2.0 * q[1] * q[1] - 2.0 * q[3] * q[3]
    r6 = 2.0 * (q[2] * q[3] + q[1] * q[4])
    r7 = 2.0 * (q[1] * q[3] + q[2] * q[4])
    r8 = 2.0 * (q[2] * q[3] - q[1] * q[4])
    r9 = 1.0 - 2.0 * q[1] * q[1] - 2.0 * q[2] * q[2]

    ntile = ceil(Int, natom/ngid)
    for itile = 1:ntile
        iatom = ngid*(itile-1) + gid + 1
        if iatom > natom
            return
        end
        x_new = r1 * x[iatom] + r2 * y[iatom] + r3 * z[iatom]
        y_new = r4 * x[iatom] + r5 * y[iatom] + r6 * z[iatom]
        z_new = r7 * x[iatom] + r8 * y[iatom] + r9 * z[iatom]
        x[iatom] = x_new
        y[iatom] = y_new
        z[iatom] = z_new
    end
    return nothing
end

function rotate_with_matrix_gpu!(x, y, z, q)
    natom = length(x)
    ngid = gridDim().x*blockDim().x
    gid = (blockIdx().x - 1)*blockDim().x + threadIdx().x - 1

    r1 = R[1, 1]
    r2 = R[1, 2]
    r3 = R[1, 3]
    r4 = R[2, 1]
    r5 = R[2, 2]
    r6 = R[2, 3]
    r7 = R[3, 1]
    r8 = R[3, 2]
    r9 = R[3, 3]

    ntile = ceil(Int, natom/ngid)
    for itile = 1:ntile
        iatom = ngid*(itile-1) + gid + 1
        if iatom > natom
            return
        end
        x_new = r1 * x[iatom] + r2 * y[iatom] + r3 * z[iatom]
        y_new = r4 * x[iatom] + r5 * y[iatom] + r6 * z[iatom]
        z_new = r7 * x[iatom] + r8 * y[iatom] + r9 * z[iatom]
        x[iatom] = x_new
        y[iatom] = y_new
        z[iatom] = z_new
    end
    return nothing
end

function filter_tops!(score_tops, cartesian_tops, iq_tops, score, iq, tops)
    if any(score .< score_tops[tops+1])
        c = findall(score .< score_tops[tops+1])
        s = score[c]
        if isempty(s)
            return
        end
        nrows = length(s)
        nrows = min(nrows, tops)
        score_tops[(tops+1):(tops+nrows)] .= Array(s[1:nrows])
        cartesian_tops[(tops+1):(tops+nrows)] .= Array(c[1:nrows])
        iq_tops[(tops+1):(tops+nrows)] .= iq
        id = sortperm(score_tops)
        score_tops .= score_tops[id]
        cartesian_tops .= cartesian_tops[id]
        iq_tops .= iq_tops[id]
    end
    return
end

function filter_tops!(score_tops, cartesian_tops, rotz_tops, rotx_tops, score, rotz, rotx, tops)
    nz = size(score, 3)
    iz_center = ceil(Int, (nz/2.0)+1.0)
    score_xy = score[:, :, iz_center]
    if any(score_xy .< score_tops[tops+1])
        c = findall(score_xy .< score_tops[tops+1])
        s = score_xy[c]
        if isempty(s)
            return
        end
        nrows = length(s)
        nrows = min(nrows, tops)
        score_tops[(tops+1):(tops+nrows)] .= Array(s[1:nrows])
        cartesian_tops[(tops+1):(tops+nrows)] .= Array(c[1:nrows])
        rotz_tops[(tops+1):(tops+nrows)] .= rotz
        rotx_tops[(tops+1):(tops+nrows)] .= rotx
        id = sortperm(score_tops)
        score_tops .= score_tops[id]
        cartesian_tops .= cartesian_tops[id]
        rotz_tops .= rotz_tops[id]
        rotx_tops .= rotx_tops[id]
    end
    return
end

function dock!(receptor::TrjArray{T, U}, ligand::TrjArray{T, U}, quaternions::Matrix{T}; deg=15.0, grid_space=1.2, iframe=1, tops=100, alpha=0.01, beta=0.06) where {T, U}
    decenter!(receptor)
    decenter!(ligand)
    orient!(receptor)
    orient!(ligand)

    # Assign atom radius
    println("step1: assigning atom radius")
    receptor = set_radius(receptor)
    ligand = set_radius(ligand)

    # Solvent accessible surface
    println("step2: computing SASA")
    receptor = compute_sasa(receptor, 1.4)
    ligand = compute_sasa(ligand, 1.4)

    # ACE scores
    println("step3: assigning ACE scores")
    receptor = set_acescore(receptor)
    ligand = set_acescore(ligand)

    # Determine grid size and coordinates
    println("step4: computing grid size and coordinates")
    size_ligand = maximum(ligand.xyz[iframe, 1:3:end]) - minimum(ligand.xyz[iframe, 1:3:end])
    #size_ligand = size_ligand*2

    x_min = minimum(receptor.xyz[iframe, 1:3:end]) - size_ligand - grid_space
    y_min = minimum(receptor.xyz[iframe, 2:3:end]) - size_ligand - grid_space
    z_min = minimum(receptor.xyz[iframe, 3:3:end]) - size_ligand - grid_space

    x_max = maximum(receptor.xyz[iframe, 1:3:end]) + size_ligand + grid_space
    y_max = maximum(receptor.xyz[iframe, 2:3:end]) + size_ligand + grid_space
    z_max = maximum(receptor.xyz[iframe, 3:3:end]) + size_ligand + grid_space

    x_grid = collect(x_min:grid_space:x_max)
    y_grid = collect(y_min:grid_space:y_max)
    z_grid = collect(z_min:grid_space:z_max)

    nx, ny, nz = length(x_grid), length(y_grid), length(z_grid)
    nxyz = nx*ny*nz

    println("step5: assigning values to grid points")
    # receptor grid: shape complementarity
    grid_RSC = zeros(Complex{T}, (nx, ny, nz))

    id_surface = receptor.sasa .> 1.0
    id_core = .!id_surface

    x = receptor.xyz[iframe, 1:3:end]
    y = receptor.xyz[iframe, 2:3:end]
    z = receptor.xyz[iframe, 3:3:end]
    rcut = zeros(T, receptor.natom)
    rcut[id_core] .= receptor.radius[id_core] .* sqrt(1.5)
    rcut[id_surface] .= receptor.radius[id_surface] .* sqrt(0.8)
    value_core = fill(Complex{T}(9im), receptor.natom)

    x_surface = x[id_surface]
    y_surface = y[id_surface]
    z_surface = z[id_surface]
    rcut_surface = receptor.radius[id_surface] .+ 3.4
    value_surface = fill(Complex{T}(1.0), length(rcut_surface))

    spread_neighbors_substitute!(grid_RSC, x_surface, y_surface, z_surface, x_grid, y_grid, z_grid, rcut_surface, value_surface)
    spread_neighbors_substitute!(grid_RSC, x, y, z, x_grid, y_grid, z_grid, rcut, value_core)

    # receptor grid: desolvation free energy
    grid_RDS = zeros(Complex{T}, (nx, ny, nz))
    rcut_ds = fill(T(6.0), receptor.natom)
    spread_nearest!(grid_RDS, x, y, z, x_grid, y_grid, z_grid)
    spread_neighbors_add!(grid_RDS, x, y, z, x_grid, y_grid, z_grid, rcut_ds, receptor.mass)

    # ligand grid: shape complementarity
    grid_LSC = zeros(Complex{T}, (nx, ny, nz))

    id_surface = ligand.sasa .> 1.0
    id_core = .!id_surface

    x = ligand.xyz[iframe, 1:3:end]
    y = ligand.xyz[iframe, 2:3:end]
    z = ligand.xyz[iframe, 3:3:end]
    rcut = zeros(T, ligand.natom)
    rcut[id_core] .= ligand.radius[id_core] .* sqrt(1.5)
    rcut[id_surface] .= ligand.radius[id_surface] .* sqrt(0.8)
    value_core = fill(Complex{T}(9im), ligand.natom)
    rcut_surface = ligand.radius[id_surface]
    value_surface = fill(Complex{T}(1.0), length(rcut_surface))

    x_surface = x[id_surface]
    y_surface = y[id_surface]
    z_surface = z[id_surface]
    spread_neighbors_substitute!(grid_LSC, x, y, z, x_grid, y_grid, z_grid, rcut, value_core)
    spread_neighbors_substitute!(grid_LSC, x_surface, y_surface, z_surface, x_grid, y_grid, z_grid, rcut_surface, value_surface)

    # ligand grid: desolvation free energy
    grid_LDS = zeros(Complex{T}, (nx, ny, nz))
    rcut_ds = fill(T(6.0), ligand.natom)
    spread_nearest!(grid_LDS, x, y, z, x_grid, y_grid, z_grid)
    spread_neighbors_add!(grid_LDS, x, y, z, x_grid, y_grid, z_grid, rcut_ds, ligand.mass)

    println("step7: docking")
    score_sc = similar(grid_RSC, T)
    score_ds = similar(grid_RSC, T)
    score_tops = fill(typemax(eltype(Float32)), 2*tops)
    cartesian_tops = Vector{CartesianIndex{3}}(undef, 2*tops)
    iq_tops = fill(-1, 2*tops);
    if CUDA.functional()
        grid_RSC_d = cu(grid_RSC)
        grid_LSC_d = cu(grid_LSC)
        grid_RDS_d = cu(grid_RDS)
        grid_LDS_d = cu(grid_LDS)
        grid_LDS_real = real.(grid_LDS_d)
        grid_LDS_imag = real.(grid_LDS_d)
        x_grid_d = cu(x_grid)
        y_grid_d = cu(y_grid)
        z_grid_d = cu(z_grid)
        x_org = cu(x)
        y_org = cu(y)
        z_org = cu(z)
        x_d = cu(x)
        y_d = cu(y)
        z_d = cu(z)

        id_surface_d = cu(id_surface)
        rcut_d = cu(rcut)
        rcut_surface_d = cu(rcut_surface)
        value_core_d = cu(value_core)
        value_surface_d = cu(value_surface)

        rcut_ds_d = cu(rcut_ds)
        ligand_mass_d = cu(ligand.mass)

        t_d = similar(grid_RSC_d)
        score_sc_d = real(t_d)
        score_ds_d = real(t_d)
        score_d = real(t_d)

        quaternions_d = cu(quaternions)
        nblocks = ceil(Int, length(x)/256)
        @time @showprogress for iq = 1:size(quaternions, 1)
            # rotate ligand
            x_d .= x_org
            y_d .= y_org
            z_d .= z_org
            #@time CUDA.@sync @cuda threads=256 blocks=nblocks rotate_with_matrix_gpu!(x_d, y_d, z_d, R_d)
            @cuda threads=256 blocks=nblocks rotate_gpu!(x_d, y_d, z_d, quaternions_d[iq, :])

            # shape complementarity
            grid_LSC_d .= zero(eltype(grid_LSC_d))
            @cuda threads=256 blocks=nblocks spread_neighbors_substitute_gpu!(grid_LSC_d, x_d, y_d, z_d, x_grid_d, y_grid_d, z_grid_d, rcut_d, value_core_d)
            @cuda threads=256 blocks=nblocks spread_neighbors_substitute_gpu!(grid_LSC_d, x_d[id_surface_d], y_d[id_surface_d], z_d[id_surface_d], x_grid_d, y_grid_d, z_grid_d, rcut_surface_d, value_surface_d)
            #t_d .= ifftshift(ifft(ifft(grid_RSC_d) .* fft(grid_LSC_d))) .* nxyz
            #score_sc_d .= - real(t_d) .+ imag(t_d)
            t_d .= ifftshift(ifft(ifft(grid_RSC_d) .* fft(grid_LSC_d))) .* nxyz
            score_sc_d .= - real(t_d) .+ imag(t_d)

            # desolvation free energy
            grid_LDS_d .= zero(eltype(grid_LDS_d))
            grid_LDS_real .= zero(eltype(grid_LDS_real))
            grid_LDS_imag .= zero(eltype(grid_LDS_imag))
            @cuda threads=256 blocks=nblocks spread_nearest_gpu!(grid_LDS_imag, x_d, y_d, z_d, x_grid_d, y_grid_d, z_grid_d)
            #@show typeof(grid_LDS_real)
            #@show typeof(ligand_mass_d)
            @cuda threads=256 blocks=nblocks spread_neighbors_add_gpu!(grid_LDS_real, x_d, y_d, z_d, x_grid_d, y_grid_d, z_grid_d, rcut_ds_d, ligand_mass_d)
            grid_LDS_d .= grid_LDS_real .+ grid_LDS_imag .* im
            t_d .= 0.5 .* ifftshift(ifft(ifft(grid_RDS_d) .* fft(grid_LDS_d))) .* nxyz
            score_ds_d .= - imag(t_d)

            # filter top scores
            #score_d .= alpha .* score_sc_d
            score_d .= alpha .* score_sc_d .+ score_ds_d
            filter_tops!(score_tops, cartesian_tops, iq_tops, score_d, iq, tops)
        end
        grid_RSC .= Array(grid_RSC_d)
        grid_LSC .= Array(grid_LSC_d)
        grid_RDS .= Array(grid_RDS_d)
        grid_LDS .= Array(grid_LDS_d)
        score_sc .= Array(score_sc_d)
        score_ds .= Array(score_ds_d)
    else
        x_org = deepcopy(x)
        y_org = deepcopy(y)
        z_org = deepcopy(z)
        t = similar(grid_LSC)
        score = real(t)

        @time @showprogress for iq = 1:size(quaternions, 1)
            # rotate ligand
            x .= x_org
            y .= y_org
            z .= z_org
            rotate!(x, y, z, quaternions[iq, :])

            # shape complementarity
            grid_LSC .= zero(eltype(grid_LSC))
            spread_neighbors_substitute!(grid_LSC, x, y, z, x_grid, y_grid, z_grid, rcut, value_core)
            spread_neighbors_substitute!(grid_LSC, x[id_surface], y[id_surface], z[id_surface], x_grid, y_grid, z_grid, rcut_surface, value_surface)
            #t .= ifft(fft(grid_RSC) .* conj.(fft(conj.(grid_LSC))))
            #score_sc .= real(t)
            t .= ifftshift(ifft(ifft(grid_RSC) .* fft(grid_LSC))) .* nxyz
            score_sc .= - real(t) .+ imag(t)

            # desolvation free energy
            grid_LDS .= zero(eltype(grid_LDS))
            spread_nearest!(grid_LDS, x, y, z, x_grid, y_grid, z_grid)
            spread_neighbors_add!(grid_LDS, x, y, z, x_grid, y_grid, z_grid, rcut_ds, ligand.mass)
            t .= 0.5 .* ifftshift(ifft(ifft(grid_RDS) .* fft(grid_LDS))) .* nxyz
            score_ds .= - imag(t)

            # filter top scores
            #score .= score_sc
            score .= alpha .* score_sc .+ score_ds
            filter_tops!(score_tops, cartesian_tops, iq_tops, score, iq, tops)
        end
    end

    ligand_init = deepcopy(ligand[iframe, :])
    ligand_return = deepcopy(ligand[0, :])
    ix_center = ceil(Int, (nx/2.0)+1.0)
    iy_center = ceil(Int, (ny/2.0)+1.0)
    iz_center = ceil(Int, (nz/2.0)+1.0)
    for itop = 1:tops
        iq = iq_tops[itop]
        ligand_tmp = rotate(ligand_init, quaternions[iq, :])
        dx = (cartesian_tops[itop][1]-ix_center) * grid_space
        dy = (cartesian_tops[itop][2]-iy_center) * grid_space
        dz = (cartesian_tops[itop][3]-iz_center) * grid_space
        ligand_tmp.xyz[1, 1:3:end] .-= dx
        ligand_tmp.xyz[1, 2:3:end] .-= dy
        ligand_tmp.xyz[1, 3:3:end] .-= dz
        ligand_return = [ligand_return; ligand_tmp]
    end

    return (receptor=receptor, ligand=ligand_return, score=score_tops, iq=iq_tops, cartesian=cartesian_tops,
            grid_RSC=grid_RSC, grid_LSC=grid_LSC,
            grid_RDS=grid_RDS, grid_LDS=grid_LDS,
            score_sc, score_ds)
end

function dock_multimer!(receptor::TrjArray{T, U}; deg=15.0, grid_space=1.2, iframe=1, tops=100, alpha=0.01, beta=0.06, nfold=3, rot_space=10.0) where {T, U}
    decenter!(receptor)
    orient!(receptor)

    # Assign atom radius
    println("step1: assigning atom radius")
    receptor = set_radius(receptor)

    # Solvent accessible surface
    println("step2: computing SASA")
    receptor = compute_sasa(receptor, 1.4)

    # ACE scores
    println("step3: assigning ACE scores")
    receptor = set_acescore(receptor)

    # Determine grid size and coordinates
    println("step5: computing grid size and coordinates")
    size_ligand = maximum(receptor.xyz[iframe, 1:3:end]) - minimum(receptor.xyz[iframe, 1:3:end])

    x_min = minimum(receptor.xyz[iframe, 1:3:end]) - size_ligand - grid_space
    y_min = minimum(receptor.xyz[iframe, 2:3:end]) - size_ligand - grid_space
    z_min = minimum(receptor.xyz[iframe, 3:3:end]) - size_ligand - grid_space

    x_max = maximum(receptor.xyz[iframe, 1:3:end]) + size_ligand + grid_space
    y_max = maximum(receptor.xyz[iframe, 2:3:end]) + size_ligand + grid_space
    z_max = maximum(receptor.xyz[iframe, 3:3:end]) + size_ligand + grid_space

    x_grid = collect(x_min:grid_space:x_max)
    y_grid = collect(y_min:grid_space:y_max)
    z_grid = collect(z_min:grid_space:z_max)

    nx, ny, nz = length(x_grid), length(y_grid), length(z_grid)
    nxyz = nx*ny*nz

    println("step6: assigning values to grid points")
    # receptor grid: shape complementarity
    id_surface = receptor.sasa .> 1.0
    id_core = .!id_surface

    x = receptor.xyz[iframe, 1:3:end]
    y = receptor.xyz[iframe, 2:3:end]
    z = receptor.xyz[iframe, 3:3:end]
    rcut = zeros(T, receptor.natom)
    rcut[id_core] .= receptor.radius[id_core] .* sqrt(1.5)
    rcut[id_surface] .= receptor.radius[id_surface] .* sqrt(0.8)
    value_core = fill(Complex{T}(9im), receptor.natom)

    x_surface = x[id_surface]
    y_surface = y[id_surface]
    z_surface = z[id_surface]
    rcut_surface = receptor.radius[id_surface] .+ 3.4
    value_surface = fill(Complex{T}(1.0), length(rcut_surface))

    # receptor grid: desolvation free energy
    grid_RSC = zeros(Complex{T}, (nx, ny, nz))
    grid_LSC = zeros(Complex{T}, (nx, ny, nz))
    grid_RDS = zeros(Complex{T}, (nx, ny, nz))
    grid_LDS = zeros(Complex{T}, (nx, ny, nz))
    rcut_ds = fill(T(6.0), receptor.natom)
    #spread_nearest!(grid_RDS, x, y, z, x_grid, y_grid, z_grid)
    #spread_neighbors_add!(grid_RDS, x, y, z, x_grid, y_grid, z_grid, rcut_ds, receptor.mass)

    println("step7: docking")
    score_sc = similar(grid_RSC, T)
    score_ds = similar(grid_RSC, T)
    score_tops = fill(typemax(eltype(Float32)), 2*tops)
    cartesian_tops = Vector{CartesianIndex{2}}(undef, 2*tops)
    rotz_tops = fill(-1.0, 2*tops);
    rotx_tops = fill(-1.0, 2*tops);
    rot_z_pi = collect((0.0:rot_space:360.0)./ 360.0 .* (2.0*pi))
    rot_x_pi = collect((0.0:rot_space:90.0)./ 360.0 .* (2.0*pi))
    #if CUDA.functional()
    if 1 == 0
    else
        x_org = deepcopy(x)
        y_org = deepcopy(y)
        z_org = deepcopy(z)
        t = similar(grid_LSC)
        score = real(t)
        @showprogress for rotz in rot_z_pi
            R = [cos(rotz) -sin(rotz) 0.0; sin(rotz) cos(rotz) 0.0; 0.0 0.0 1.0]
            receptor_rot = rotate_with_matrix(receptor, R)
            for rotx in rot_x_pi
                R = [1.0 0.0 0.0; 0.0 cos(rotx) -sin(rotx); 0.0 sin(rotx) cos(rotx)]
                receptor_rot = rotate_with_matrix(receptor_rot, R)
                R = [cos(2.0*pi/nfold) -sin(2.0*pi/nfold) 0.0; sin(2.0*pi/nfold) cos(2.0*pi/nfold) 0.0; 0.0 0.0 1.0]
                ligand_rot = rotate_with_matrix(receptor_rot, R)

                x = receptor_rot.xyz[iframe, 1:3:end]
                y = receptor_rot.xyz[iframe, 2:3:end]
                z = receptor_rot.xyz[iframe, 3:3:end]
                grid_RSC .= zero(eltype(grid_RSC))
                spread_neighbors_substitute!(grid_RSC, x[id_surface], y[id_surface], z[id_surface], x_grid, y_grid, z_grid, rcut_surface, value_surface)
                spread_neighbors_substitute!(grid_RSC, x, y, z, x_grid, y_grid, z_grid, rcut, value_core)

                x = ligand_rot.xyz[iframe, 1:3:end]
                y = ligand_rot.xyz[iframe, 2:3:end]
                z = ligand_rot.xyz[iframe, 3:3:end]
                grid_LSC .= zero(eltype(grid_LSC))
                spread_neighbors_substitute!(grid_LSC, x, y, z, x_grid, y_grid, z_grid, rcut, value_core)
                spread_neighbors_substitute!(grid_LSC, x[id_surface], y[id_surface], z[id_surface], x_grid, y_grid, z_grid, rcut_surface, value_surface)

                t .= ifftshift(ifft(ifft(grid_RSC) .* fft(grid_LSC))) .* nxyz
                score_sc .= - real(t) .+ imag(t)

                x = receptor_rot.xyz[iframe, 1:3:end]
                y = receptor_rot.xyz[iframe, 2:3:end]
                z = receptor_rot.xyz[iframe, 3:3:end]
                grid_RDS .= zero(eltype(grid_RDS))
                spread_nearest!(grid_RDS, x, y, z, x_grid, y_grid, z_grid)
                spread_neighbors_add!(grid_RDS, x, y, z, x_grid, y_grid, z_grid, rcut_ds, receptor.mass)

                x = ligand_rot.xyz[iframe, 1:3:end]
                y = ligand_rot.xyz[iframe, 2:3:end]
                z = ligand_rot.xyz[iframe, 3:3:end]
                grid_LDS .= zero(eltype(grid_LDS))
                spread_nearest!(grid_LDS, x, y, z, x_grid, y_grid, z_grid)
                spread_neighbors_add!(grid_LDS, x, y, z, x_grid, y_grid, z_grid, rcut_ds, receptor.mass)

                t .= 0.5 .* ifftshift(ifft(ifft(grid_RDS) .* fft(grid_LDS))) .* nxyz
                score_ds .= - imag(t)

                score .= alpha .* score_sc .+ score_ds
                filter_tops!(score_tops, cartesian_tops, rotz_tops, rotx_tops, score, rotz, rotx, tops)
            end
        end
    end

    receptor_init = deepcopy(receptor[iframe, :])
    receptor_return = deepcopy(receptor[0, :])
    ligand_return = deepcopy(receptor[0, :])
    ix_center = ceil(Int, (nx/2.0)+1.0)
    iy_center = ceil(Int, (ny/2.0)+1.0)
    iz_center = ceil(Int, (nz/2.0)+1.0)
    for itop = 1:tops
        rotz = rotz_tops[itop]
        R = [cos(rotz) -sin(rotz) 0.0; sin(rotz) cos(rotz) 0.0; 0.0 0.0 1.0]
        receptor_rot = rotate_with_matrix(receptor_init, R)
        rotx = rotx_tops[itop]
        R = [1.0 0.0 0.0; 0.0 cos(rotx) -sin(rotx); 0.0 sin(rotx) cos(rotx)]
        receptor_rot = rotate_with_matrix(receptor_rot, R)
        receptor_return = [receptor_return; receptor_rot]
        R = [cos(2.0*pi/nfold) -sin(2.0*pi/nfold) 0.0; sin(2.0*pi/nfold) cos(2.0*pi/nfold) 0.0; 0.0 0.0 1.0]
        ligand_rot = rotate_with_matrix(receptor_rot, R)
        dx = (cartesian_tops[itop][1]-ix_center) * grid_space
        dy = (cartesian_tops[itop][2]-iy_center) * grid_space
        #dz = (cartesian_tops[itop][3]-iz_center) * grid_space
        ligand_rot.xyz[1, 1:3:end] .-= dx
        ligand_rot.xyz[1, 2:3:end] .-= dy
        #ligand_rot.xyz[1, 3:3:end] .-= dz
        ligand_return = [ligand_return; ligand_rot]
    end

    natom = receptor_return.natom
    cc = TrjArray(receptor_return[1, :], chainname=fill(string(1), natom), chainid=fill(1, natom))
    for ifold = 2:nfold
        cc = [cc TrjArray(receptor_return[1, :], chainname=fill(string(ifold), natom), chainid=fill(ifold, natom))]
    end
    @show cc
    complex_return = deepcopy(cc[0, :])
    for itop = 1:tops
        com1 = centerofmass(receptor_return[itop, :])
        com2 = centerofmass(ligand_return[itop, :])
        L = [com2.xyz[1, 1] - com1.xyz[1, 1], com2.xyz[1, 2] - com1.xyz[1, 2]]
        L_norm = sqrt(L[1]^2 + L[2]^2)
        alpha = (180.0 - (360.0/nfold))/2.0
        alpha = (alpha / 360.0) * 2.0 * pi
        d_norm = L_norm / (2.0*cos(alpha))

        cc = TrjArray(receptor_return[itop, :], chainname=fill(string(1), natom), chainid=fill(1, natom))
        cc.xyz[:, 1:3:end] .= cc.xyz[:, 1:3:end] .+ d_norm
        tmp = deepcopy(cc)
        for ifold = 2:nfold
            beta = (2.0 * pi / nfold) * (ifold-1)
            R = [cos(beta) -sin(beta) 0.0; sin(beta) cos(beta) 0.0; 0.0 0.0 1.0]
            r = rotate_with_matrix(tmp, R)
            cc = [cc TrjArray(r, chainname=fill(string(ifold), natom), chainid=fill(ifold, natom))]
        end
        complex_return = [complex_return; cc]
    end

    return (receptor=receptor, ligand=ligand_return, multimer=complex_return, score=score_tops, rotz=rotz_tops, rotx=rotx_tops, cartesian=cartesian_tops,
            grid_RSC=grid_RSC, grid_LSC=grid_LSC,
            grid_RDS=grid_RDS, grid_LDS=grid_LDS,
            score_sc, score_ds)
end

function dock_multimer(receptor::TrjArray{T, U}; rot_space=10.0, grid_space=1.2, iframe=1, tops=10, nfold=3) where {T, U}
    # generate grid coordinates for receptor
    receptor2, dummy = decenter(receptor)

    @show "step 1"
    x_min, x_max = minimum(receptor2.xyz[iframe, 1:3:end]), maximum(receptor2.xyz[iframe, 1:3:end])
    y_min, y_max = minimum(receptor2.xyz[iframe, 2:3:end]), maximum(receptor2.xyz[iframe, 2:3:end])
    z_min, z_max = minimum(receptor2.xyz[iframe, 3:3:end]), maximum(receptor2.xyz[iframe, 3:3:end])
    size_receptor = sqrt((x_max - x_min)^2 + (y_max - y_min)^2 + (z_max - z_min)^2)

    x_min = minimum(receptor2.xyz[iframe, 1:3:end]) - size_receptor - grid_space
    y_min = minimum(receptor2.xyz[iframe, 2:3:end]) - size_receptor - grid_space
    z_min = minimum(receptor2.xyz[iframe, 3:3:end]) - size_receptor - grid_space

    x_max = maximum(receptor2.xyz[iframe, 1:3:end]) + size_receptor + grid_space
    y_max = maximum(receptor2.xyz[iframe, 2:3:end]) + size_receptor + grid_space
    z_max = maximum(receptor2.xyz[iframe, 3:3:end]) + size_receptor + grid_space

    x_grid = collect(x_min:grid_space:x_max)
    y_grid = collect(y_min:grid_space:y_max)
    z_grid = collect(z_min:grid_space:z_max)

    @show "step 2"
    nx, ny, nz = length(x_grid), length(y_grid), length(z_grid)

    # shape complementarity of receptor
    @show "step 3"
    iatom_surface = receptor2.sasa .> 1.0
    iatom_core = .!iatom_surface
    rcut1 = zeros(T, receptor2.natom)
    rcut2 = zeros(T, receptor2.natom)

    rcut1[iatom_core] .= receptor2.radius[iatom_core] * sqrt(1.5)
    rcut2[iatom_core] .= -1.0
    rcut1[iatom_surface] .= receptor2.radius[iatom_surface] * sqrt(0.8)
    rcut2[iatom_surface] .= receptor2.radius[iatom_surface] .+ 3.4

    @show nx ny nz
    @show "step 4"
    grid_RSC = zeros(complex(T), nx, ny, nz)
    grid_LSC = zeros(complex(T), nx, ny, nz)
    #assign_shape_complementarity!(grid_RSC, receptor2, grid_space, rcut1, rcut2, x_grid, y_grid, z_grid, iframe)

    @show "step 4.1"
    # compute score with FFT
    #nq = size(quaternions, 1)
    #s = @showprogress pmap(q -> compute_docking_score_with_fft(quaternions[q, :], grid_RSC, grid_LSC, ligand2, grid_space, rcut1, rcut2, x_grid, y_grid, z_grid, iframe, tops, q), 1:nq)
    rot_z_pi = collect((0.0:rot_space:360.0)./ 360.0 .* (2.0*pi))
    rot_x_pi = collect((0.0:rot_space:360.0)./ 360.0 .* (2.0*pi))
    score_best = -Inf
    @show "step 4.2"
    receptor_best = deepcopy(receptor2)
    ligand_best = deepcopy(receptor2)
    dx_estimated = 0
    dy_estimated = 0
    @show "step 5"
    for z in rot_z_pi
        @printf "%f \n" z
        R = [cos(z) -sin(z) 0.0; sin(z) cos(z) 0.0; 0.0 0.0 1.0]
        receptor_rot = rotate_with_matrix(receptor2, R)
        for x in rot_x_pi
            R = [1.0 0.0 0.0; 0.0 cos(x) -sin(x); 0.0 sin(x) cos(x)]
            receptor_rot = rotate_with_matrix(receptor_rot, R)
            R = [cos(2.0*pi/nfold) -sin(2.0*pi/nfold) 0.0; sin(2.0*pi/nfold) cos(2.0*pi/nfold) 0.0; 0.0 0.0 1.0]
            ligand_rot = rotate_with_matrix(receptor_rot, R)
            assign_shape_complementarity!(grid_RSC, receptor_rot, grid_space, rcut1, rcut2, x_grid, y_grid, z_grid, iframe)
            assign_shape_complementarity!(grid_LSC, ligand_rot, grid_space, rcut1, rcut2, x_grid, y_grid, z_grid, iframe)
            score = compute_docking_score_on_xyplane(grid_RSC, grid_LSC)
            score_max = maximum(score)
            if score_best < score_max
                id = argmax(score)
                dx_estimated = id[1]
                dy_estimated = id[2]
                receptor_best = deepcopy(receptor_rot)
                ligand_best = deepcopy(ligand_rot)
                score_best = score_max
            end
        end
    end

    ligand_return = deepcopy(ligand_best)
    dx = (dx_estimated-1) * grid_space
    if dx > (nx*grid_space / 2.0)
        dx = dx - (nx*grid_space)
    end
    dy = (dy_estimated-1) * grid_space
    if dy > (ny*grid_space / 2.0)
        dy = dy - (ny*grid_space)
    end
    ligand_return.xyz[iframe, 1:3:end] .+= dx
    ligand_return.xyz[iframe, 2:3:end] .+= dy

    return (receptor=receptor_best, ligand=ligand_return, score=score_best)
end

function docking_by_electrostatic_energy(receptor::TrjArray{T, U}, ligand::TrjArray{T, U}, quaternions; grid_space=1.2, iframe=1, tops=10) where {T, U}

    # generate grid coordinates of receptor
    decenter!(receptor)
    decenter!(ligand)

    # determine size of ligand
    x_min = minimum(ligand.xyz[iframe, 1:3:end])
    y_min = minimum(ligand.xyz[iframe, 2:3:end])
    z_min = minimum(ligand.xyz[iframe, 3:3:end])
    x_max = maximum(ligand.xyz[iframe, 1:3:end])
    y_max = maximum(ligand.xyz[iframe, 2:3:end])
    z_max = maximum(ligand.xyz[iframe, 3:3:end])
    size_ligand = sqrt((x_max-x_min)^2 + (y_max-y_min)^2 + (z_max-z_min)^2)

    # extension grid of receptor using size of ligand
    x_min = minimum(receptor.xyz[iframe, 1:3:end]) - size_ligand - grid_space
    y_min = minimum(receptor.xyz[iframe, 2:3:end]) - size_ligand - grid_space
    z_min = minimum(receptor.xyz[iframe, 3:3:end]) - size_ligand - grid_space
    x_max = maximum(receptor.xyz[iframe, 1:3:end]) + size_ligand + grid_space
    y_max = maximum(receptor.xyz[iframe, 2:3:end]) + size_ligand + grid_space
    z_max = maximum(receptor.xyz[iframe, 3:3:end]) + size_ligand + grid_space

    x_grid = collect(x_min:grid_space:x_max)
    y_grid = collect(y_min:grid_space:y_max)
    z_grid = collect(z_min:grid_space:z_max)

    nx = length(x_grid)
    ny = length(y_grid)
    nz = length(z_grid)

    grid_elec = zeros(T, nx, ny, nz)
    # assign ace score for RDS
    for iatom = 1:receptor.natom
        # create atom coordinates of receptor
        x_atom = receptor.xyz[iframe, 3*(iatom-1)+1]
        y_atom = receptor.xyz[iframe, 3*(iatom-1)+2]
        z_atom = receptor.xyz[iframe, 3*(iatom-1)+3]

        # determin Real part of RDS
        rcut = 20.0
        ix_min = findfirst(abs.(x_atom .- x_grid) .< rcut)
        iy_min = findfirst(abs.(y_atom .- y_grid) .< rcut)
        iz_min = findfirst(abs.(z_atom .- z_grid) .< rcut)
        ix_max = findlast(abs.(x_atom .- x_grid) .< rcut)
        iy_max = findlast(abs.(y_atom .- y_grid) .< rcut)
        iz_max = findlast(abs.(z_atom .- z_grid) .< rcut)

        # ６Åの判別（まだ条件分岐をつくってないので注意）
        for ix = ix_min:ix_max
            for iy = iy_min:iy_max
                for iz = iz_min:iz_max
                    dist = sqrt((x_atom - x_grid[ix])^2 + (y_atom - y_grid[iy])^2 + (z_atom - z_grid[iz])^2)
                    if dist < 2.0
                        grid_elec[ix, iy, iz] += 0.0
                    elseif dist < 6.0
                        dielec = 4.0
                        grid_elec[ix, iy, iz] += receptor.charge[iatom] / (dielec * dist^2)
                    elseif dist < 8.0
                        dielec = 38.0*dist - 224.0
                        grid_elec[ix, iy, iz] += receptor.charge[iatom] / (dielec * dist^2)
                    else
                        dielec = 80.0
                        grid_elec[ix, iy, iz] += receptor.charge[iatom] / (dielec * dist^2)
                    end
                end
            end
        end
    end

    for iatom = 1:receptor.natom
        if receptor.radius[iatom] > 1.0
            rcut = receptor.radius[iatom] * sqrt(0.8)
        else
            rcut = receptor.radius[iatom] * sqrt(1.5)
        end

        x_atom = receptor.xyz[iframe, 3*(iatom-1)+1]
        y_atom = receptor.xyz[iframe, 3*(iatom-1)+2]
        z_atom = receptor.xyz[iframe, 3*(iatom-1)+3]

        ix_min = findfirst(abs.(x_atom .- x_grid) .< rcut)
        iy_min = findfirst(abs.(y_atom .- y_grid) .< rcut)
        iz_min = findfirst(abs.(z_atom .- z_grid) .< rcut)
        ix_max = findlast(abs.(x_atom .- x_grid) .< rcut)
        iy_max = findlast(abs.(y_atom .- y_grid) .< rcut)
        iz_max = findlast(abs.(z_atom .- z_grid) .< rcut)

        for ix = ix_min:ix_max
            for iy = iy_min:iy_max
                for iz = iz_min:iz_max
                    grid_elec[ix, iy, iz] = 0.0
                end
            end
        end
    end

    # generate grid coordinates of ligand at every rotation by quaternions
    grid_charge = zeros(T, nx, ny, nz)
    score = similar(grid_charge)
    nq = size(quaternions, 1)
    @showprogress for iq = 1:nq
        grid_charge .= 0.0

        # rotate ligand by quaternions
        ligand_rotated = rotate(ligand, quaternions[iq, :])

        # assign ace score for LDS
        for iatom = 1:ligand_rotated.natom

            # create atom coordinates of ligand
            x_atom = ligand_rotated.xyz[iframe, 3*(iatom-1)+1]
            y_atom = ligand_rotated.xyz[iframe, 3*(iatom-1)+2]
            z_atom = ligand_rotated.xyz[iframe, 3*(iatom-1)+3]

            # invert atoms coordinates into grid coordinates
            ix_min = floor(Int, (x_atom - x_grid[1]) / grid_space)
            iy_min = floor(Int, (y_atom - y_grid[1]) / grid_space)
            iz_min = floor(Int, (z_atom - z_grid[1]) / grid_space)

            a = x_atom - x_grid[ix_min] - (0.5*grid_space)
            b = y_atom - y_grid[iy_min] - (0.5*grid_space)
            c = z_atom - z_grid[iz_min] - (0.5*grid_space)

            for ix = 1:2
                for iy = 1:2
                    for iz = 1:2
                        X = (-1)^ix * (0.5*grid_space)
                        Y = (-1)^iy * (0.5*grid_space)
                        Z = (-1)^iz * (0.5*grid_space)
                        weight = (X + a)*(Y + b)*(Z + c) / (X*Y*Z) / 8.0
                        grid_charge[ix_min+ix-1, iy_min+iy-1, iz_min+iz-1] = weight * ligand.charge[iatom]
                    end
                end
            end
        end

        # compute elec socre with FFT
        t = ifft(fft(grid_elec) .* conj.(fft(conj.(grid_charge))))
        score = real(t)
    end

    return score
end


text:
#-----------------------#
# GENERIC MODEL METHODS #
#-----------------------#

@doc """Generic constructor for a DynamicalModel's underlying ODEFunctions."""
function (T::Type{<:Abstract_ModelODEFunctions})(args...; kwargs...)
    _ode, deduplicate_terms = ODESystem(T, args...; kwargs...)

    # Reduce high-order system to 1st-order (re-ordering to keep the original equations first)
    # NOTE: need to expand the RHS derivatives before calling order-lowering
    eqs = expand_derivatives.(equations(_ode))
    num_orig_eqs = length(eqs)
    eqs, dvs = ode_order_lowering(eqs, independent_variable(_ode), ModelingToolkit.states(_ode))
    eqs = [eqs[end - num_orig_eqs + 1:end]..., eqs[1:num_orig_eqs]...]
    dvs = [dvs[end - num_orig_eqs + 1:end]..., dvs[1:num_orig_eqs]...]
    ode = ODESystem(simplify.(eqs), independent_variable(_ode), dvs, parameters(_ode))

    # Generate the functions
    # TODO: Add support for tgrad (need to define derivative(get_pos) for EphemerisNBP)
    ode_f = ODEFunction(ode; jac=true, tgrad=false, eval_expression=false, eval_module=@__MODULE__, deduplicate_terms)
    ode_stm_f = STM_ODEFunction(ode, ode_f; deduplicate_terms)
    T(ode, ode_f, ode_stm_f)
end

@doc """Generic model function."""
(model::Abstract_DynamicalModel)(args...; kwargs...) = model.ode.ode_f(args...; kwargs...)

@traitdef HasJacobian{X}
@traitimpl HasJacobian{X} <- has_jacobian(X)
has_jacobian(X::Type{<:DiffEqBase.ODEFunction}) = !isnothing(fieldtype(X, :jac))

@doc """
    Function to generate an ODE Function that computes the State Transition
    Matrix simultaneously with the given function f. This requires f to have
    a computable Jacobian function.
"""
@traitfn function STM_ODEFunction(ode::ModelingToolkit.AbstractODESystem, f::::HasJacobian; deduplicate_terms=[])
    # TODO: Output just the variational equations, join as a SplitODEProblem
    # TODO: Memoize this function! It's very slightly slow for EphemerisNBP

    # Convert Variable -> Operation
    iv  = independent_variable(ode)
    dvs = states(ode)
    params = parameters(ode)

    # NOTE: Depends on the Jacobian, corresponding to A(t) matrix (for the
    # State Transition Matrix) [Parker & Anderson 2014].
    @variables   ϕ[1:length(dvs),1:length(dvs)](iv)
    D = Differential(iv)

    # Get the Jacobian matrix (A(t))
    A = Base.invokelatest(f.jac, Num.(dvs), Num.(params), Num(iv))

    # The State Transition Matrix (STM) ODE function is defined as follows, including the N^2 Jacobian equations +
    # the N first-order equations of motion. [Koon 2011]
    # NOTE: the Differential is defined element-wise and flattened to a list.
    stm_eqs = simplify.(D.(ϕ) .~ A * ϕ)

    # Create the ODE system and generate its functions
    stm_ode = ODESystem(
        [equations(ode)..., stm_eqs...], # Append the ODE equations.
        iv,
        [dvs..., ϕ...],  # Append the STM and motion state variables
        params)
    stm_f = ODEFunction(stm_ode; sparse=true, eval_expression=false, eval_module=@__MODULE__, deduplicate_terms)
end

#---------#
# DISPLAY #
#---------#
Base.show(io::IO, x::Abstract_DynamicalModel) = show(io, typeof(x))
Base.show(io::IO, x::Type{<:Abstract_DynamicalModel}) = print(io, nameof(x))
ModelingToolkit.varmap_to_vars(model::Abstract_DynamicalModel, varmap) = ModelingToolkit.varmap_to_vars(varmap, parameters(model))
DiffEqBase.isinplace(f::Abstract_DynamicalModel, _) = true

# XXX: Need these due to new ModelingToolkit interface.
ModelingToolkit.get_systems(::Abstract_ModelODEFunctions) = []
ModelingToolkit.get_eqs(f::Abstract_ModelODEFunctions) = ModelingToolkit.get_eqs(f.ode_system)
ModelingToolkit.get_states(f::Abstract_ModelODEFunctions) = ModelingToolkit.get_states(f.ode_system)
ModelingToolkit.get_ps(f::Abstract_ModelODEFunctions) = ModelingToolkit.get_ps(f.ode_system)
Base.nameof(f::Abstract_ModelODEFunctions) = nameof(typeof(f))

#----------#
# INCLUDES #
#----------#
include("models/3bp.jl")
include("models/4bp.jl")
include("models/nbp.jl")

ext: jl


text:
is_given_macro(::Any) = false
is_given_macro(e::Expr) = e.head == :macrocall && e.args[1] == Symbol("@given")

given_parameters(::Any) = []
function given_parameters(e::Expr)
    if e.head == :call && e.args[1] == :∈
        # The name and associated iterator generator.
        return [(e.args[2], e.args[3])]
    elseif e.head == :tuple || is_given_macro(e)
        return vcat(map(given_parameters, e.args)...)
    end

    error("Cannot parse given parameters from expression $(e)")
end

function settingexpr(blocks::Expr...)
    call = Expr(:call, Expr(:., :ParameterTests, :(:ParamTestSettings)))
    params = Expr(:parameters, [Expr(:kw, block.args...) for block in blocks]...)
    push!(call.args, params)
    return call
end

paramkeys(params::Vector) = vcat(map(paramkeys, params)...)
paramkeys(param::Symbol) = [param]
function paramkeys(param::Expr)
    param.head == :tuple || error("param keys must be a tuple")
    param.args
end

macro paramtest(name, blocks::Expr...)
    block = blocks[end]
    block.head == :block || error("paramtest requires a block expression")
    givenlines = filter(is_given_macro, block.args)
    length(givenlines) > 0 || error("paramtest must have a @given parameters")
    # Drop the given lines from our block, we don't need them.
    filter!((!)∘is_given_macro, block.args)
    givenparams = vcat(map(given_parameters, givenlines)...)
    setparams = paramkeys(map(first, givenparams))

    settings = settingexpr(blocks[1:end-1]...)
    testblock = quote
        @testset $name begin
            @testset ParameterTestSet begin
                ParameterTests.withsettings($settings) do
                    for $(Expr(:tuple, map(first, givenparams)...)) in $(Expr(:call, :zip, map(last, givenparams)...))
                        @testset SingleParameterTestSet begin
                            local ts = Test.get_testset()
                            $(map(p -> :(ts.parameters[$(QuoteNode(p))] = $p), setparams)...)
                            $block
                        end
                    end
                end
            end
        end
    end
    return esc(testblock)
end


text:
@cenum hipError_t::UInt32 begin
    hipSuccess = 0
    hipErrorOutOfMemory = 2
    hipErrorNotInitialized = 3
    hipErrorDeinitialized = 4
    hipErrorProfilerDisabled = 5
    hipErrorProfilerNotInitialized = 6
    hipErrorProfilerAlreadyStarted = 7
    hipErrorProfilerAlreadyStopped = 8
    hipErrorInsufficientDriver = 35
    hipErrorInvalidImage = 200
    hipErrorInvalidContext = 201
    hipErrorContextAlreadyCurrent = 202
    hipErrorMapFailed = 205
    hipErrorUnmapFailed = 206
    hipErrorArrayIsMapped = 207
    hipErrorAlreadyMapped = 208
    hipErrorNoBinaryForGpu = 209
    hipErrorAlreadyAcquired = 210
    hipErrorNotMapped = 211
    hipErrorNotMappedAsArray = 212
    hipErrorNotMappedAsPointer = 213
    hipErrorECCNotCorrectable = 214
    hipErrorUnsupportedLimit = 215
    hipErrorContextAlreadyInUse = 216
    hipErrorPeerAccessUnsupported = 217
    hipErrorInvalidKernelFile = 218
    hipErrorInvalidGraphicsContext = 219
    hipErrorInvalidSource = 300
    hipErrorFileNotFound = 301
    hipErrorSharedObjectSymbolNotFound = 302
    hipErrorSharedObjectInitFailed = 303
    hipErrorOperatingSystem = 304
    hipErrorSetOnActiveProcess = 305
    hipErrorInvalidHandle = 400
    hipErrorNotFound = 500
    hipErrorIllegalAddress = 700
    hipErrorInvalidSymbol = 701
    # Runtime Error Codes start here.
    hipErrorMissingConfiguration = 1001
    hipErrorMemoryAllocation = 1002
    hipErrorInitializationError = 1003
    hipErrorLaunchFailure = 1004
    hipErrorPriorLaunchFailure = 1005
    hipErrorLaunchTimeOut = 1006
    hipErrorLaunchOutOfResources = 1007
    hipErrorInvalidDeviceFunction = 1008
    hipErrorInvalidConfiguration = 1009
    hipErrorInvalidDevice = 1010
    hipErrorInvalidValue = 1011
    hipErrorInvalidDevicePointer = 1017
    hipErrorInvalidMemcpyDirection = 1021
    hipErrorUnknown = 1030
    hipErrorInvalidResourceHandle = 1033
    hipErrorNotReady = 1034
    hipErrorNoDevice = 1038
    hipErrorPeerAccessAlreadyEnabled = 1050
    hipErrorPeerAccessNotEnabled = 1051
    hipErrorRuntimeMemory = 1052
    hipErrorRuntimeOther = 1053
    hipErrorHostMemoryAlreadyRegistered = 1061
    hipErrorHostMemoryNotRegistered = 1062
    hipErrorMapBufferObjectFailed = 1071
    hipErrorAssert = 1081
    hipErrorNotSupported = 1082
    hipErrorTbd
end

hipStream_t = Ptr{Cvoid}


text:
import GeoInterface, DBFTables, Tables

"Shapefile.Table represents both the geometries and associated fields"
struct Table{T}
    shp::Handle{T}
    dbf::DBFTables.Table
    # ensure matching sizes on construction
    function Table{T}(shp, dbf) where {T}
        if length(shp) != length(dbf)
            throw(ArgumentError("Shapefile and DBF file contain a different amount of rows"))
        end
        new{T}(shp, dbf)
    end
end

function Table(shp::Handle{T}, dbf::DBFTables.Table) where {T}
    Table{T}(shp, dbf)
end

"Read a file into a Shapefile.Table"
function Table(path::AbstractString)
    stempath, ext = splitext(path)
    if lowercase(ext) == ".shp"
        shp_path = path
        shx_path = string(stempath, ".shx")
        dbf_path = string(stempath, ".dbf")
    elseif ext == ""
        shp_path = string(stempath, ".shp")
        shx_path = string(stempath, ".shx")
        dbf_path = string(stempath, ".dbf")
    else
        throw(ArgumentError("Provide the shapefile with either .shp or no extension"))
    end
    isfile(shp_path) || throw(ArgumentError("File not found: $shp_path"))
    isfile(dbf_path) || throw(ArgumentError("File not found: $dbf_path"))

    shp = if isfile(shx_path)
        Shapefile.Handle(shp_path,shx_path)
    else
        Shapefile.Handle(shp_path)
    end
    dbf = DBFTables.Table(dbf_path)
    return Shapefile.Table(shp, dbf)
end

"Struct representing a singe record in a shapefile"
struct Row{T}
    geometry::T
    record::DBFTables.Row
end

getshp(t::Table) = getfield(t, :shp)
getdbf(t::Table) = getfield(t, :dbf)

Base.length(t::Table) = length(getshp(t).shapes)

Tables.istable(::Type{<:Table}) = true
Tables.rows(t::Table) = t
Tables.columns(t::Table) = t
Tables.rowaccess(::Type{<:Table}) = true
Tables.columnaccess(::Type{<:Table}) = true

"Iterate over the rows of a Shapefile.Table, yielding a Shapefile.Row for each row"
function Base.iterate(t::Table, st = 1)
    st > length(t) && return nothing
    geom = @inbounds getshp(t).shapes[st]
    record = DBFTables.Row(getdbf(t), st)
    return Row(geom, record), st + 1
end

Base.getproperty(row::Row, name::Symbol) = getproperty(getfield(row, :record), name)

function Base.getproperty(t::Table, name::Symbol)
    if name === :geometry
        shapes(t)
    else
        getproperty(getdbf(t), name)
    end
end

Base.propertynames(row::Row) = propertynames(getfield(row, :record))

function Base.propertynames(t::Table)
    names = propertynames(getdbf(t))
    pushfirst!(names, :geometry)
    return names
end

function Tables.schema(t::Table)
    dbf = getdbf(t)
    dbf_schema = Tables.schema(dbf)
    names = (:geometry, dbf_schema.names...)
    types = (eltype(shapes(t)), dbf_schema.types...)
    return Tables.Schema(names, types)
end

function Base.show(io::IO, t::Table)
    tt = typeof(t)
    nr = length(t)
    nc = length(propertynames(t))
    println(io, "$tt with $nr rows and the following $nc columns:\n\t")
    join(io, propertynames(t), ", ")
    println(io)
end

# TODO generalize these with a future GeoInterface/GeoTables
# should probably be geometry/geometries but don't want to claim these names yet
"Get the geometry associated with a shapefile row"
shape(row::Row) = getfield(row, :geometry)
"Get a vector of the geometries in a shapefile, without any metadata"
shapes(t::Table) = getshp(t).shapes


text:
#####
##### `reshape`
#####

function rrule(::typeof(reshape), A::AbstractArray, dims::Tuple{Vararg{Union{Colon,Int}}})
    A_dims = size(A)
    function reshape_pullback(Ȳ)
        return (NoTangent(), reshape(Ȳ, A_dims), NoTangent())
    end
    return reshape(A, dims), reshape_pullback
end

function rrule(::typeof(reshape), A::AbstractArray, dims::Union{Colon,Int}...)
    A_dims = size(A)
    function reshape_pullback(Ȳ)
        ∂A = reshape(Ȳ, A_dims)
        ∂dims = broadcast(_ -> NoTangent(), dims)
        return (NoTangent(), ∂A, ∂dims...)
    end
    return reshape(A, dims...), reshape_pullback
end

#####
##### `hcat`
#####

function rrule(::typeof(hcat), Xs::Union{AbstractArray, Number}...)
    Y = hcat(Xs...)  # note that Y always has 1-based indexing, even if X isa OffsetArray
    ndimsY = Val(ndims(Y))  # this avoids closing over Y, Val() is essential for type-stability
    sizes = map(size, Xs)   # this avoids closing over Xs
    function hcat_pullback(ȳ)
        dY = unthunk(ȳ)
        hi = Ref(0)  # Ref avoids hi::Core.Box
        dXs = map(sizes) do sizeX
            ndimsX = length(sizeX)
            lo = hi[] + 1
            hi[] += get(sizeX, 2, 1)
            ind = ntuple(ndimsY) do d
                if d==2
                    d > ndimsX ? lo : lo:hi[]
                else
                    d > ndimsX ? 1 : (:)
                end
            end
            if ndimsX > 0
                # Here InplaceableThunk breaks @inferred, removed for now
                # InplaceableThunk(@thunk(dY[ind...]), dX -> dX .+= view(dY, ind...))
                dY[ind...]
            else
                # This is a hack to perhaps avoid GPU scalar indexing
                sum(view(dY, ind...))
            end
        end
        return (NoTangent(), dXs...)
    end
    return Y, hcat_pullback
end

function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVecOrMat})
    widths = map(A -> size(A,2), As)
    function reduce_hcat_pullback_2(dY)
        hi = Ref(0)
        dAs = map(widths) do w
            lo = hi[]+1
            hi[] += w
            dY[:, lo:hi[]]
        end
        return (NoTangent(), NoTangent(), dAs)
    end
    return reduce(hcat, As), reduce_hcat_pullback_2
end

function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVector})
    axe = axes(As,1)
    function reduce_hcat_pullback_1(dY)
        hi = Ref(0)
        dAs = map(_ -> dY[:, hi[]+=1], axe)
        return (NoTangent(), NoTangent(), dAs)
    end
    return reduce(hcat, As), reduce_hcat_pullback_1
end

#####
##### `vcat`
#####

function rrule(::typeof(vcat), Xs::Union{AbstractArray, Number}...)
    Y = vcat(Xs...)
    ndimsY = Val(ndims(Y))
    sizes = map(size, Xs)
    function vcat_pullback(ȳ)
        dY = unthunk(ȳ)
        hi = Ref(0)
        dXs = map(sizes) do sizeX
            ndimsX = length(sizeX)
            lo = hi[] + 1
            hi[] += get(sizeX, 1, 1)
            ind = ntuple(ndimsY) do d
                if d==1
                    d > ndimsX ? lo : lo:hi[]
                else
                    d > ndimsX ? 1 : (:)
                end
            end
            if ndimsX > 0
                # InplaceableThunk(@thunk(dY[ind...]), dX -> dX .+= view(dY, ind...))
                dY[ind...]
            else
                sum(view(dY, ind...))
            end
        end
        return (NoTangent(), dXs...)
    end
    return Y, vcat_pullback
end

function rrule(::typeof(reduce), ::typeof(vcat), As::AbstractVector{<:AbstractVecOrMat})
    Y = reduce(vcat, As)
    ndimsY = Val(ndims(Y))
    heights = map(A -> size(A,1), As)
    function reduce_vcat_pullback(dY)
        hi = Ref(0)
        dAs = map(heights) do z
            lo = hi[]+1
            hi[] += z
            ind = ntuple(d -> d==1 ? (lo:hi[]) : (:), ndimsY)
            dY[ind...]
        end
        return (NoTangent(), NoTangent(), dAs)
    end
    return Y, reduce_vcat_pullback
end

#####
##### `cat`
#####

_val(::Val{x}) where {x} = x

function rrule(::typeof(cat), Xs::Union{AbstractArray, Number}...; dims)
    Y = cat(Xs...; dims=dims)
    cdims = dims isa Val ? Int(_val(dims)) : dims isa Integer ? Int(dims) : Tuple(dims)
    ndimsY = Val(ndims(Y))
    sizes = map(size, Xs)
    function cat_pullback(ȳ)
        dY = unthunk(ȳ)
        prev = fill(0, _val(ndimsY))  # note that Y always has 1-based indexing, even if X isa OffsetArray
        dXs = map(sizes) do sizeX
            ndimsX = length(sizeX)
            index = ntuple(ndimsY) do d
                if d in cdims
                    d > ndimsX ? (prev[d]+1) : (prev[d]+1:prev[d]+sizeX[d])
                else
                    d > ndimsX ? 1 : (:)
                end
            end
            for d in cdims
                prev[d] += get(sizeX, d, 1)
            end
            if ndimsX > 0
                # InplaceableThunk(@thunk(dY[index...]), dX -> dX .+= view(dY, index...))
                dY[index...]
            else
                sum(view(dY, index...))
            end
        end
        return (NoTangent(), dXs...)
    end
    return Y, cat_pullback
end

#####
##### `hvcat`
#####

function rrule(::typeof(hvcat), rows, values::Union{AbstractArray, Number}...)
    Y = hvcat(rows, values...)
    cols = size(Y,2)
    ndimsY = Val(ndims(Y))
    sizes = map(size, values)
    function hvcat_pullback(dY)
        prev = fill(0, 2)
        dXs = map(sizes) do sizeX
            ndimsX = length(sizeX)
            index = ntuple(ndimsY) do d
                if d in (1, 2)
                    d > ndimsX ? (prev[d]+1) : (prev[d]+1:prev[d]+sizeX[d])
                else
                    d > ndimsX ? 1 : (:)
                end
            end
            prev[2] += get(sizeX, 2, 1)
            if prev[2] == cols
                prev[2] = 0
                prev[1] += get(sizeX, 1, 1)
            end
            dY[index...]
        end
        return (NoTangent(), NoTangent(), dXs...)
    end
    return Y, hvcat_pullback
end

#####
##### `fill`
#####

function rrule(::typeof(fill), value::Any, dims::Tuple{Vararg{Int}})
    function fill_pullback(Ȳ)
        return (NoTangent(), sum(Ȳ), NoTangent())
    end
    return fill(value, dims), fill_pullback
end

function rrule(::typeof(fill), value::Any, dims::Int...)
    function fill_pullback(Ȳ)
        return (NoTangent(), sum(Ȳ), ntuple(_->NoTangent(), length(dims))...)
    end
    return fill(value, dims), fill_pullback
end


text:
__precompile__(true)

module PaddedViews
using Base: OneTo, tail
using OffsetArrays

export PaddedView, paddedviews

"""
    datapadded = PaddedView(fillvalue, data, padded_indices)
    datapadded = PaddedView(fillvalue, data, padded_indices, data_indices)
    datapadded = PaddedView(fillvalue, data, sz)
    datapadded = PaddedView(fillvalue, data, sz, first_datum)

Create a padded version of the array `data`, where any elements within
the span of `padded_indices` not assigned in `data` will have value
`fillvalue`. If a second set of indices `data_indices` is not supplied
it is assumed the array `data` spans the indices from the first up until
`size(data)`, otherwise `data` spans the specified `data_indices`.
Alternately, the padded array size `sz` can be specified along with the
location of the first element of `data`, `first_datum`. If `first_datum`
is omitted, it is assumed to be the first element of the padded array.

# Example

```julia
julia> a = reshape(1:9, 3, 3)
3×3 Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}:
 1  4  7
 2  5  8
 3  6  9

julia> PaddedView(-1, a, (4, 5))
4×5 PaddedViews.PaddedView{Int64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}},Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}}:
  1   4   7  -1  -1
  2   5   8  -1  -1
  3   6   9  -1  -1
 -1  -1  -1  -1  -1

 julia> PaddedView(-1, a, (1:5,1:5), (2:4,2:4))
 PaddedViews.PaddedView{Int64,2,Tuple{UnitRange{Int64},UnitRange{Int64}},OffsetArrays.OffsetArray{Int64,2,Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}}} with indices 1:5×1:5:
  -1  -1  -1  -1  -1
  -1   1   4   7  -1
  -1   2   5   8  -1
  -1   3   6   9  -1
  -1  -1  -1  -1  -1

julia> PaddedView(-1, a, (5,5), (2,2))
5×5 PaddedViews.PaddedView{Int64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}},OffsetArrays.OffsetArray{Int64,2,Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}}}:
-1  -1  -1  -1  -1
-1   1   4   7  -1
-1   2   5   8  -1
-1   3   6   9  -1
-1  -1  -1  -1  -1
```
"""
struct PaddedView{T,N,I,A} <: AbstractArray{T,N}
    fillvalue::T
    data::A
    indices::I

    function PaddedView{T,N,I,A}(fillvalue::T,
                                 data::AbstractArray{T,N},
                                 indices::NTuple{N,AbstractUnitRange}) where {T,N,I,A}
        new{T,N,I,A}(fillvalue, data, indices)
    end
end

PaddedView(fillvalue, data::AbstractArray{T,N}, indices) where {T,N} =
    PaddedView{T,N,typeof(indices),typeof(data)}(convert(T, fillvalue), data, indices)
function PaddedView(fillvalue, data::AbstractArray{T,N}, sz::Tuple{Integer,Vararg{Integer}}) where {T,N}
    inds = map(OneTo, sz)
    PaddedView{T,N,typeof(inds),typeof(data)}(convert(T, fillvalue), data, inds)
end

# This method eliminates an ambiguity between the two below it
function PaddedView(fillvalue,
                    data::AbstractArray{T,0},
                    ::Tuple{},
                    ::Tuple{}) where T
    return PaddedView(fillvalue, data, ())
end

function PaddedView(fillvalue,
                    data::AbstractArray{T,N},
                    padded_inds::NTuple{N,AbstractUnitRange},
                    data_inds::NTuple{N,AbstractUnitRange}) where {T,N}

    off_data = OffsetArray(data, data_inds...)
    return PaddedView(fillvalue, off_data, padded_inds)
end

function PaddedView(fillvalue,
                    data::AbstractArray{T,N},
                    sz::NTuple{N,Integer},
                    first_datum::NTuple{N,Integer}) where {T,N}
    padded_inds = map(OneTo, sz)
    data_inds   = map(:, first_datum, size(data) .+ first_datum .- 1)
    return PaddedView(fillvalue, data, padded_inds, data_inds)
end

Base.axes(A::PaddedView) = A.indices
@inline Base.axes(A::PaddedView, d::Integer) = d <= ndims(A) ? A.indices[d] : default_axes(A.indices)
default_axes(::NTuple{N,I}) where {N,I<:AbstractUnitRange} = convert(I, OneTo(1))
default_axes(::Any) = OneTo(1)

Base.size(A::PaddedView) = map(length, axes(A))

@inline function Base.getindex(A::PaddedView{T,N}, i::Vararg{Int,N}) where {T,N}
    @boundscheck checkbounds(A, i...)
    if Base.checkbounds(Bool, A.data, i...)
        return A.data[i...]
    end
    return A.fillvalue
end

"""
    Aspad = paddedviews(fillvalue, A1, A2, ....)

Pad the arrays `A1`, `A2`, ..., to a common size or set of axes,
chosen as the span of axes enclosing all of the input arrays.

# Example:
```julia
julia> a1 = reshape([1,2], 2, 1)
2×1 Array{Int64,2}:
 1
 2

julia> a2 = [1.0,2.0]'
1×2 Array{Float64,2}:
 1.0  2.0

julia> a1p, a2p = paddedviews(0, a1, a2);

julia> a1p
2×2 PaddedViews.PaddedView{Int64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}},Array{Int64,2}}:
 1  0
 2  0

julia> a2p
2×2 PaddedViews.PaddedView{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}},Array{Float64,2}}:
 1.0  2.0
 0.0  0.0
```
"""
function paddedviews(fillvalue, As::AbstractArray...)
    inds = outerinds(As...)
    map(A->PaddedView(fillvalue, A, inds), As)
end
paddedviews(fillvalue) = ()

@inline outerinds(A::AbstractArray, Bs...) = _outerinds(axes(A), Bs...)
@inline _outerinds(inds, A::AbstractArray, Bs...) =
    _outerinds(_outerinds(inds, axes(A)), Bs...)
_outerinds(inds) = inds
@inline _outerinds(inds1::NTuple{N,I}, inds2::NTuple{N,I}) where {N,I<:AbstractUnitRange} =
    map((i1, i2) -> convert(I, padrange(i1, i2)), inds1, inds2)
_outerinds(inds1::NTuple{N,AbstractUnitRange}, inds2::NTuple{N,AbstractUnitRange}) where {N} =
    map((i1, i2) -> convert(UnitRange{Int}, padrange(i1, i2)), inds1, inds2)

# This shouldn't be reached due to the `paddedviews(fillvalue)` method above,
# but it's here in case anyone extends `paddedviews` for types other than AbstractArrays.
# See https://github.com/JuliaImages/ImageCore.jl/pull/32#discussion_r111545756
outerinds() = error("must supply at least one array with concrete axes")

padrange(i1::OneTo, i2::OneTo) = OneTo(max(last(i1), last(i2)))
padrange(i1::AbstractUnitRange, i2::AbstractUnitRange) = min(first(i1),first(i2)):max(last(i1), last(i2))

end # module


text:
using HDF5

filename = "test.h5"

x = 0.:10:100 |> collect
y = 0.:2:10 |> collect
z = [ i^2 + j for i in x, j in y ]

h5open(filename, "w") do fid
    fid["x"] = x[:]
    fid["y"] = y[:]
    fid["z"] = z[:, :]
    HDF5.h5ds_set_scale(fid["x"], "x")
    HDF5.h5ds_set_scale(fid["y"], "y")
    HDF5.h5ds_attach_scale(fid["z"], fid["x"], 1)
    HDF5.h5ds_attach_scale(fid["z"], fid["y"], 0)
end


text:
using GPMethylation
using Test

@testset "GPMethylation.jl" begin
    # Write your tests here.
end


text:
#        _|                                                        _|                  _|
#    _|_|_|    _|_|      _|_|_|    _|_|_|  _|    _|        _|_|_|  _|_|_|      _|_|_|      _|_|_|
#  _|    _|  _|_|_|_|  _|        _|    _|  _|    _|      _|        _|    _|  _|    _|  _|  _|    _|
#  _|    _|  _|        _|        _|    _|  _|    _|      _|        _|    _|  _|    _|  _|  _|    _|
#    _|_|_|    _|_|_|    _|_|_|    _|_|_|    _|_|_|        _|_|_|  _|    _|    _|_|_|  _|  _|    _|
#                                                _|
#                                            _|_|

abstract type Recoupling end
@with_kw struct NoRecoupling <: Recoupling
    two_λa::Int
    two_λb::Int
end

amplitude(cs::NoRecoupling, two_λa, two_λb) =
    (cs.two_λa == two_λa) *
    (cs.two_λb == two_λb)

@with_kw struct ParityRecoupling <: Recoupling
    two_λa::Int
    two_λb::Int
    ηηηphaseisplus::Bool
end
ParityRecoupling(two_λa::Int, two_λb::Int, ηηηphasesign::Char) = ParityRecoupling(two_λa, two_λb, ηηηphasesign=='+')
function ParityRecoupling(two_λa::Int, two_λb::Int, (jp,(jp1,jp2))::Pair{jp, Tuple{jp, jp}})
    ηηη = jp1.p ⊗ jp2.p ⊗ jp.p
    ηηηphase = (2*(ηηη == '+') - 1) * x"-1"^(jp.j - jp1.j - jp2.j)
    return ParityRecoupling(two_λa, two_λb, ηηηphase == 1)
end

function amplitude(cs::ParityRecoupling, two_λa, two_λb)
    (cs.two_λa == two_λa) * (cs.two_λb == two_λb) && return 1
    (cs.two_λa == -two_λa) * (cs.two_λb == -two_λb) && return 2*cs.ηηηphaseisplus-1
    return 0
end

@with_kw struct RecouplingLS <: Recoupling
    two_j::Int
    two_ls::Tuple{Int,Int}
    two_ja::Int
    two_jb::Int
end
RecouplingLS(two_ls, (jp,(jpa,jpb))::Pair{jp, Tuple{jp, jp}}) =
    RecouplingLS(jp.j |> x2, two_ls, jpa.j |> x2, jpb.j |> x2)

amplitude(cs::RecouplingLS, two_λa, two_λb) =
    jls_coupling(cs.two_ja, two_λa, cs.two_jb, two_λb, cs.two_j, cs.two_ls[1], cs.two_ls[2])

@with_kw struct DecayChain{X, V1<:Recoupling, V2<:Recoupling, T}
    k::Int
    #
    two_s::Int # isobar spin
    #
    Xlineshape::X # lineshape
    #
    HRk::V1
    Hij::V2
    #
    tbs::T # the structure with masses and spins
end

function printable_l(two_l)
    l = div(two_l,2)
    waves = ['S', 'P', 'D', 'F', 'G', 'H'];
    return l < 6 ? waves[l+1] : string(l)[1]
end
printable_s(two_s) = two_s//2
printable_ls(two_ls) = (printable_s(two_ls[2]), printable_l(two_ls[1]))

"""
    DecayChainLS(k, Xlineshape;
        two_s, # give two_s, i.e. the spin of the isobar
        parity, #
        Ps, # need parities: [1,2,3,0]
        tbs) # give three-body-system structure

    Returns the decay chain with the smallest LS, ls
"""
function DecayChainLS(k, Xlineshape;
    tbs=error("give three-body-system structure"),
    two_s=0,
    parity::Char='+',
    Ps=SVector('+','+','+','+'))
    #
    lsLS = vcat(possible_lsLS(k, two_s, parity, tbs.two_js, Ps)...)
    length(lsLS) == 0 && error("there are no possible LS couplings")
    #
    lsLS_sorted = sort(lsLS, by=x->x.LS[1])
    @unpack ls, LS = lsLS_sorted[1]
    #
    i,j = ij_from_k(k)
    return DecayChain(; k, Xlineshape, tbs, two_s,
        Hij=RecouplingLS(two_s, Int.(2 .* ls), tbs.two_js[i], tbs.two_js[j]),
        HRk=RecouplingLS(tbs.two_js[4], Int.(2 .* LS), two_s, tbs.two_js[k]))
end

"""
    DecayChainsLS(k, Xlineshape;
        two_s, # give two_s, i.e. the spin of the isobar
        parity, #
        Ps, # need parities: [1,2,3,0]
        tbs) # give three-body-system structure

    Returns an array of the decay chains with all possible couplings
"""
function DecayChainsLS(k, Xlineshape;
    two_s = error("give two_s, i.e. the spin of the isobar, two_s=..."),
    parity::Char = error("give the parity, parity=..."),
    Ps = error("need parities: Ps=[P1,2,3,0]"),
    tbs = error("give three-body-system structure, tbs=..."))
    #
    LSlsv = possible_lsLS(k, two_s, parity, tbs.two_js, Ps)
    return [DecayChain(;
        k, Xlineshape, tbs, two_s,
            Hij=RecouplingLS(two_s, Int.(2 .* ls), tbs.two_js[i], tbs.two_js[j]),
            HRk=RecouplingLS(tbs.two_js[4], Int.(2 .* LS), two_s, tbs.two_js[k]))
        for x in LSlsv]
end

function amplitude(σs, two_λs, dc)
    k = dc.k; i,j = ij_from_k(k);
    tbs = dc.tbs
    s = tbs.ms.m0^2
    #
    two_s = dc.two_s
    two_js = tbs.two_js
    #
    itr_two_λs′ = itr(SVector{3}(tbs.two_js[1],tbs.two_js[2],tbs.two_js[3]))
    f = 0.0
    for two_τ = -two_s:2:two_s, two_λs′ in itr_two_λs′
        f += Zksτ(k,two_s,two_τ,two_λs,two_λs′,σs,tbs) *
            amplitude(dc.HRk, two_τ, two_λs′[k]) *
            amplitude(dc.Hij, two_λs′[i], two_λs′[j])
    end
    lineshape = dc.Xlineshape(σs[k])
    return f * lineshape
end
#
amplitude(dpp, dc) = amplitude(dpp.σs, dpp.two_λs, dc)
#
summed_over_polarization(fn, two_js) = σs->sum(fn(σs,two_λs) for two_λs in itr(two_js))
#
itr(two_js) = Iterators.ProductIterator(Tuple([-two_j:2:two_j for two_j in two_js]))

max_stars_repo_name: mmikhasenko/ThreeBodyDecay

max_issues_repo_name: mmikhasenko/ThreeBodyDecay

max_forks_repo_name: mmikhasenko/ThreeBodyDecay


text:
#=
    Examples file
=#

using Nemo
using SparseArrays: sparse

include("../src/ExactODEReduction.jl")
using .ExactODEReduction: many_invariant_subspaces, check_invariance!, polynormalize, @sparse

#------------------------------------------------------------------------------

# Example 1

A1 = @sparse [
    1 1 0 0;
    1 1 0 0;
    0 0 0 2;
    0 0 0 1;
]
A2 = @sparse[
    0 1 0 0;
    1 0 0 0;
    0 0 0 2;
    0 0 0 0;
]

# find an invariant subspaces of the matrices
invariants = many_invariant_subspaces([A1, A2])

println("--------------\n",
    join(map(string, invariants), "\n"),
    "\n--------------")

@assert all(V -> check_invariance!([A1, A2], V),
            invariants)

#------------------------------------------------------------------------------

# Example 2

B1 = @sparse [
    0 0 0 0 0;
    0 0 0 0 0;
    0 0 3 0 0;
    0 0 3 3 1;
    0 0 0 0 2;
]

# find an invariant subspaces of the matrices
invariants = many_invariant_subspaces([B1])

println("--------------\n",
    join(map(string, invariants), "\n"),
    "\n--------------")

@assert all(V -> check_invariance!([B1], V),
            invariants)

#------------------------------------------------------------------------------


text:
module FASTX

export
    FASTA,
    FASTQ,
    identifier,
    description,
    sequence,
    quality,
    transcribe

# Generic methods
function identifier end
function description end
function sequence end

include("fasta/fasta.jl")
include("fastq/fastq.jl")

import .FASTA
import .FASTQ
import .FASTQ: quality

function FASTA.Record(record::FASTQ.Record)
    FASTQ.checkfilled(record)
    slice = (first(record.identifier) - 1):last(record.sequence)
    newdata = @inbounds record.data[slice]
    newdata[1] = 0x3E
    FASTA.Record(newdata)
end

"""
    transcribe(in::FASTQ.Reader, out::FASTA.Writer)

Convert a FASTQ file to a FASTA file.
"""
function transcribe(in::FASTQ.Reader, out::FASTA.Writer)
    buff_record = FASTQ.Record()
    while !eof(in)
        read!(in, buff_record)
        write(out, FASTA.Record(buff_record))
    end
end

end # module


text:
function LoadTestSparseData(Oracles)

	features = sparse([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16],[1,2,3,4,1,2,3,4,1,2,3,4,1,2,3,4], [0.746453, -0.646761,-0.768557,-1.26793, -2.18751,-0.0529194, -3.98887,-0.30549,  1.10894, 0.111473, 0.639841, 0.337336,  0.348856, 2.40815, 0.573903, -1.27225])

	# featuresTP = sparse([1,2,3,4,1,2,3,4,1,2,3,4,1,2,3,4],[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16], [0.746453, -0.646761,-0.768557,-1.26793, -2.18751,-0.0529194, -3.98887,-0.30549,  1.10894, 0.111473, 0.639841, 0.337336,  0.348856, 2.40815, 0.573903, -1.27225])

	# featuresVec = SparseMatrixCSC{Float64,Int64}[]
	#
	# push!(featuresVec, sparsevec([1],[0.746453]))
	# push!(featuresVec, sparsevec([2],[-0.646761]))
	# push!(featuresVec, sparsevec([3],[-0.768557]))
	# push!(featuresVec, sparsevec([4],[-1.26793]))
	# push!(featuresVec, sparsevec([1],[ -2.18751]))
	# push!(featuresVec, sparsevec([2],[-0.0529194]))
	# push!(featuresVec, sparsevec([3],[-3.98887]))
	# push!(featuresVec, sparsevec([4],[-0.30549]))
	# push!(featuresVec, sparsevec([1],[ 1.10894]))
	# push!(featuresVec, sparsevec([2],[0.111473]))
	# push!(featuresVec, sparsevec([3],[0.639841]))
	# push!(featuresVec, sparsevec([4],[0.337336]))
	# push!(featuresVec, sparsevec([1],[ 0.348856]))
	# push!(featuresVec, sparsevec([2],[2.40815]))
	# push!(featuresVec, sparsevec([3],[0.573903]))
	# push!(featuresVec, sparsevec([4],[ -1.27225]))

	labels = [1.0,1.0,-1.0,-1.0,1.0,-1.0,1.0,-1.0,1.0,-1.0,1.0,-1.0,1.0,-1.0,1.0,-1.0]

	(classLabels, numClasses) = createClassLabels(labels)

	(trf, trl, trl2, numTP, tef, tel, tel2) = trainTestRandomSeparate(features, MinusPlusOneVector(labels, Set([1.0])), classLabels)

	L2regs=[
	false,
	true
	]
	for L2reg in L2regs

		(gradientOracle, numVars, numTP, csDataType, LossFunctionString, myOutputter, L2reg) = createSBLOracles(trf, trl, numTP, tef, tel, L2reg; outputLevel = 0)

		push!(Oracles, (gradientOracle, numVars, numTP, csDataType, LossFunctionString, myOutputter, L2reg, "TestSparse", (t)-> Problem(L2reg, "TestSparse", LossFunctionString,  numTP, t,(j)->getSampleFunctionAt(j,gradientOracle))))

	end
end

LoadTestSparseData(Oracles)