#

"""Dump energies and coordinates."""
__author__ = "Randall J. Radmer"
__version__ = "1.0"

import sys

import simtk.unit as unit

def writeEnergyAndCoords(fOut, step, eK, eP, coords,
                         simTime, modelFrameNumber, verbose):
    if verbose:
        sys.stdout.write("%s: Potential Energy = %s, Kinetic Energy = %s\n" %
                           (simTime.format("%7.2f"),
                            eP.format("%5.2f"),
                            eK.format("%5.2f")))
        sys.stdout.flush()

    # Add coords to PDB output file
    appendCoordsToPDB(fOut, step, eK, eP, coords, modelFrameNumber+1)

def appendCoordsToPDB(fOut, step, eK, eP, coords, modelFrameNumber):
    """Append one coord set to the open PDB file"""
    fOut.write("REMARK  step %6d: EK = %s   EP = %s\n"
                % (step, eK.format("%.3f"), eP.format("%.3f")))
    fOut.write("MODEL     %d\n" % modelFrameNumber)
    for ii in range(len(coords)):
        fOut.write("ATOM   %4d  AR  ARA  %4d    %8.3f%8.3f%8.3f\n"
                   % (ii, ii,
                      coords[ii][0].value_in_unit(unit.angstrom),
                      coords[ii][1].value_in_unit(unit.angstrom),
                      coords[ii][2].value_in_unit(unit.angstrom)) )
    fOut.write("ENDMDL\n")

def closePDB(fOut):
    # Close PDB file
    fOut.write("END\n")
    if fOut!=sys.stdout:
        fOut.close()