"""
efieldreporter.py: Outputs the electric field and position of a particular atom.
"""
__author__ = "Kyle Beauchamp"
__version__ = "1.0"

import simtk.openmm as mm
from simtk.unit import *
from numpy import array
import tables
Filter=tables.Filters(complevel=9,complib='blosc',shuffle=True)
    
class EFieldReporter(object):
    """Report the electric field and position of a single atom.
    
    Notes:
    This reporter requires multiple memory copies to the GPU.  It will be slow when done at high frequency.
    """
    
    def __init__(self, filename, reportInterval,other_simulation,atom_index):
        """Create a EFieldReporter.
    
        Parameters:
         - file (string) The file to write to
         - reportInterval (int) The interval (in time steps) at which to write frames
         - other_simulation (simulation): another simulation object used to calculate field
         - atom_index (int): the index of the atom at which to calculate field.
        """
        self._reportInterval = reportInterval
	self._out = tables.File(filename,'a')
	self._data=self._out.createEArray("/","Data",tables.Float32Atom(),(0,7),filters=Filter)
        #self._out.write("#Timestep [ps] x y z [m] Ex Ey Ez [N/C]\n")
        self._topology = None
        self._nextModel = 0
        self._other_simulation=other_simulation
        self._atom_index=atom_index
        self._adjust_charges()

    def _adjust_charges(self):
        """Sets the charge of atom_index to 0 for calculation of electric field."""
        f=self._other_simulation.system.getForce(3)
        self._q,sig,eps=f.getParticleParameters(self._atom_index)
        f.setParticleParameters(0,0.0,sig,eps)
    
    def describeNextReport(self, simulation):
        """Get information about the next report this object will generate.
        
        Parameters:
         - simulation (Simulation) The Simulation to generate a report for
        Returns: A five element tuple.  The first element is the number of steps until the
        next report.  The remaining elements specify whether that report will require
        positions, velocities, forces, and energies respectively.
        """
        steps = self._reportInterval - simulation.currentStep%self._reportInterval
        return (steps, True, False, True, False)
    
    def report(self, simulation, state):
        """Generate a report.
        
        Parameters:
         - simulation (Simulation) The Simulation to generate a report for
         - state (State) The current state of the simulation
        """

        current_time=simulation.integrator.getStepSize().value_in_unit(picosecond)
        current_time*=simulation.currentStep
        
        xyz=state.getPositions()
        f0=state.getForces(asNumpy=True)[self._atom_index]

        self._other_simulation.context.setPositions(xyz)
        other_state=self._other_simulation.context.getState(getForces=True)
        f=other_state.getForces(asNumpy=True)[self._atom_index]

        efield=(f0-f)
        efield/=constants.AVOGADRO_CONSTANT_NA
        efield/=self._q

        efield=efield.value_in_unit(newton/coulomb)

        xyz=xyz[self._atom_index].value_in_unit(meter)

        self._write(efield,xyz,current_time)
        self._nextModel += 1

    def _write(self,efield,xyz,current_time):
        """Save results to disk using pure python text IO."""
	self._data.append([[current_time,xyz[0],xyz[1],xyz[2],efield[0],efield[1],efield[2]]])
        self._out.flush()
    
    def __del__(self):
        self._out.close()