#! /usr/bin/env 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
from libcrp import crp_python_wrapper

class CRPCalculator:
	"""CRP 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, PES):
		self.atoms = None
		self.results = {}

		self.PES = PES

		crp_python_wrapper.initialize()

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

		return

	# For CRP only the atomic positions are important for calculation
	all_changes = ['positions', 'numbers']

	implemented_properties = ['energy', 'forces']

	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 skewing( u, v, latu, latv, skewing ):
		"""
		Convert from cartesian coordinates to the skewed coordinate system (in fractional coordinates to make our lives more convenient...)
		There are some implicit assumptions here.
		"""
		x = u * latu + v * latv * np.cos( np.radians( skewing ) )
		y = v * latv * np.sin( np.radians( skewing ) )
		return [ x, y ]

	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=None):
		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 set_PES(self, PES):
		"""
		If one wants to change the CRP PES.
		"""
		self.PES = PES
		return

	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()

		if len(atoms) != 2:
			raise CalculatorError('This calculator is build for a diatomic CRP PES, something went wrong with number of atoms')

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

		crp_python_wrapper.setup_crp_potential( self.PES )
		Energy, Forces = crp_python_wrapper.calculate_crp_potential( [self.atoms[0].position[2], self.atoms[0].position[0], self.atoms[0].position[1],
										self.atoms[1].position[2], self.atoms[1].position[0], self.atoms[1].position[1]] )
		crp_python_wrapper.destroy_crp_potential()

		self.results['energy'] = Energy
		self.results['forces'] = np.array( [[Forces[1], Forces[2], Forces[0]], [Forces[4], Forces[5], Forces[3]]] )

		return
