#!/usr/bin/python
import ase
import numpy as np
from subprocess import call
from copy import deepcopy as dcopy
from scipy.constants import physical_constants as phc

atmas = phc['atomic mass constant'][0]
eV2J  = phc['electron volt-joule relationship'][0]
Bc    = phc['Boltzmann constant'][0]
hbar  = phc['Planck constant'][0] / (2*np.pi)
H2J   = phc['Hartree energy'][0]
A2B   = 1e-10 / phc['Bohr radius'][0]
H2eV  = phc['Hartree energy in eV'][0]

class RuNNerCalculator:
	"""ASE calculator.

	A calculator should store a copy of the atoms object used for the
	last calculation.  When one of the *get_potential_energy*,
	*get_forces*, or *get_stress* methods is called, the calculator
	should check if anything has changed since the last calculation
	and only do the calculation if it's really needed.  Two sets of
	atoms are considered identical if they have the same positions,
	atomic numbers, unit cell and periodic boundary conditions."""

	def __init__(self):
		self.atoms = None
		self.results = {}

		if not hasattr(self, 'name'):
			self.name = self.__class__.__name__.lower()

		return

	# For RuNNer only the atomic positions, cell size and periodic boundary conditions are important for calculation
	all_changes = ['positions', 'numbers', 'cell', 'pbc']

	implemented_properties = ['energy', 'forces']

	def setup_input(self, atoms=None):
		"""The input file for RuNNer is generated here"""

		if not atoms is None:
			if not self.atoms==atoms:
				self.atoms=atoms

		data_file = open('input.data', 'w')

		Type = atoms.get_chemical_symbols()

		data_file.write( "begin \n" )
		for i  in range(3):
			data_file.write( "lattice {0:15.8f} {1:15.8f} {2:15.8f}\n" .format(atoms.cell[i,0]*A2B, atoms.cell[i,1]*A2B, atoms.cell[i,2]*A2B)  )

		for i  in range( len( atoms ) ):
			data_file.write( "atom {0:15.8f} {1:15.8f} {2:15.8f} {3:10s} 0.0  0.0  0.0  0.0  0.0\n" .format(atoms.positions[i,0]*A2B, atoms.positions[i,1]*A2B, atoms.positions[i,2]*A2B, Type[i])  )
		data_file.write( "energy 0.0\n")
		data_file.write( "charge 0.0\n")
		data_file.write( "end\n")
		data_file.close()

		return

	def reset(self):
		"""Clear all information from old calculation."""

		self.atoms = None
		self.results = {}

	def check_state(self, atoms, tol=1e-15):
		"""Check for any atomic position changes since last calculation."""
		state = ase.calculators.calculator.compare_atoms(self.atoms, atoms, tol)
		return state

	def get_property(self, name, atoms=None, allow_calculation=True):
		if name not in self.implemented_properties:
			raise ase.calculators.calculator.PropertyNotImplementedError('{0:s} property not implemented'.format(name))

		if atoms is None:
			atoms = self.atoms
			system_changes = []
		else:
			system_changes = self.check_state(atoms)
			if system_changes:
				self.reset()
		if name not in self.results:
			if not allow_calculation:
				return None
			self.calculate(atoms, [name], system_changes)

		if name not in self.results:
			# For some reason the calculator was not able to do what we want,
			# and that is not OK.
			raise ase.calculators.calculator.PropertyNotImplementedError('{0:s} not present in this calculation'.format(name))

		result = self.results[name]
		if isinstance(result, np.ndarray):
			result = result.copy()

		return result

	def get_potential_energy(self, atoms=None, force_consistent=False):
		if force_consistent:
			return self.get_property('free_energy', atoms)
		else:
			return self.get_property('energy', atoms)

	def get_potential_energies(self, atoms=None):
		return self.get_property('energies', atoms)

	def get_forces(self, atoms=None):
		return self.get_property('forces', atoms)

	def get_stress(self, atoms):
		return self.get_property('stress', atoms)

	def get_stresses(self, atoms=None):
		return self.get_property('stresses', atoms)

	def get_dipole_moment(self, atoms=None):
		return self.get_property('dipole', atoms)

	def get_charges(self, atoms=None):
		return self.get_property('charges', atoms)

	def get_magnetic_moment(self, atoms=None):
		return self.get_property('magmom', atoms)

	def get_magnetic_moments(self, atoms=None):
		"""Calculate magnetic moments projected onto atoms."""
		return self.get_property('magmoms', atoms)

	def Econverged(self, forces=None):
		"""Did the optimization converge?
		It checks not only for forces, but also energies."""
		if forces is None:
			forces = self.atoms.get_forces()
		if hasattr(self.atoms, 'get_curvature'):
			conv = ((forces**2).sum(axis=1).max() < self.fmax**2 and
				self.atoms.get_curvature() < 0.0)
		else:
			conv = (forces**2).sum(axis=1).max() < self.fmax**2

		if conv == False:
			self.Energy = self.atoms.calc.Energy
			if hasattr(self, 'Energy_old'):
				if abs(self.Energy_old - self.Energy) < 10**-8:
					self.conv_count += 1
				else:
					self.conv_count = 0
				conv = self.conv_count > 9
				self.Energy_old = self.Energy
			else:
				self.conv_count = 0
				self.Energy_old = self.Energy

		return conv

	def calculate(self, atoms=None, properties=['energy', 'forces'], system_changes=all_changes):
		"""Do the calculation.

		properties: list of str
			List of what needs to be calculated.  Can be any combination
			of 'energy', 'forces', 'stress', 'dipole', 'charges', 'magmom'
			and 'magmoms'.
		system_changes: list of str
			List of what has changed since last calculation.  Can be
			any combination of these six: 'positions', 'numbers', 'cell',
			'pbc', 'initial_charges' and 'initial_magmoms'.

		Subclasses need to implement this, but can ignore properties
		and system_changes if they want.  Calculated properties should
		be inserted into results dictionary like shown in this dummy
		example::

			self.results = {'energy': 0.0,
					'forces': np.zeros((len(atoms), 3)),
					'stress': np.zeros(6),
					'dipole': np.zeros(3),
					'charges': np.zeros(len(atoms)),
					'magmom': 0.0,
					'magmoms': np.zeros(len(atoms))}

		The subclass implementation should first call this
		implementation to set the atoms attribute.

		Current implementation expects RuNNer to be in the same folder as from which the script is run.
		"""

		if atoms is not None:
			self.atoms = atoms.copy()
#			self.atoms=atoms

		self.results = {'energy': 0.0,
				'forces': np.zeros((len(atoms), 3))}

		self.setup_input(atoms)

		myf = open("mytmp", "w")
		call(["./RuNNer.serial.x"],stdout=myf)

		energy_file = open('energy.out', 'r').readlines()
		Energy = float( energy_file[-1].split()[-2] ) * H2eV
		forces_file = open('nnforces.out', 'r').readlines()
		Forces = []
		for i in range( 1, len(forces_file) ):
			Forces.append([ float(x)*A2B*H2eV for x in forces_file[i].split()[-3:] ])

		self.results['energy'] = Energy
		self.results['forces'] = np.array( Forces )

		call(["rm", "nnatoms.out", "nnforces.out", "output.data", "debug.out", "mytmp", "energy.out"])

		return