#!/usr/bin/env python

import IMP.em


def main():
    IMP.set_log_level(IMP.SILENT)
    desc = """Samples a protein into a simulated 3D density map."""
    p = IMP.ArgumentParser(description=desc)
    p.add_argument("pdb", help="input PDB file name")
    p.add_argument("density", help="output EM density file name")
    p.add_argument("resolution", type=float, help="resolution")
    p.add_argument("voxel_size", type=float, help="voxel size")
    args = p.parse_args()
    pdb_fn = args.pdb
    mrc_fn = args.density
    resolution = args.resolution
    apix = args.voxel_size
    # read the protein
    mdl = IMP.Model()
    mh = IMP.atom.read_pdb(
        pdb_fn,
        mdl,
        IMP.atom.NonWaterNonHydrogenPDBSelector())
    IMP.atom.add_radii(mh)
    # sample the density
    dmap = IMP.em.particles2density(IMP.core.get_leaves(mh), resolution, apix)
    IMP.em.write_map(dmap, mrc_fn, IMP.em.MRCReaderWriter())

if __name__ == "__main__":
    main()
