#! /usr/bin/env python3

import numpy as np
from CRPCalculator import CRPCalculator

from ase import Atoms
from ase.build import molecule

eV2kJ = 96.4853075

def set_positions_spherical(atoms, x, y, z, r, theta, phi, apply_constraint=True ):
	"""
	Set position of the two atoms with spherical coordinates instead of cartesian coordinates
	"""
	ctheta = np.cos( np.radians( theta ) )
	stheta = np.sin( np.radians( theta ) )
	cphi = np.cos( np.radians( phi ) )
	sphi = np.sin( np.radians( phi ) )

	atoms.set_positions([[x - 0.5 * r * stheta * cphi, y - 0.5 * r * stheta * sphi, z - 0.5 * r * ctheta], [x + 0.5 * r * stheta * cphi, y + 0.5 * r * stheta * sphi, z + 0.5 * r * ctheta]], apply_constraint=apply_constraint)

	return

atoms = molecule('O2')
#lattice vectors for O2 on Al(111) HSE03-1/3x@RPBE 
#
a 		= 2.844266999968277             # Bohr
#a               = 5.374885268357408             # Angstrom
SuperCellSize 	= 2 				# Number of Unit Cells in the SuperCell
UnitCell 	= np.array([
				[a, 0.000000000, 0.000000000],
				[-0.5*a, 0.5*np.sqrt(3)*a, 0.000000000],
				[0.000000000, 0.000000000, 0.000000000]])		# Create the UnitCell, FCC!
Cell	= SuperCellSize * UnitCell


atoms.set_cell( Cell )
atoms.set_pbc( [True, True, True])

#Select the PES
#Retrieved from: http://pubs.tc.lic.leidenuniv.nl/public/elham.ghassemi/2017-CPL-683-329/SPO-DVR-PES/H2Pt111PBEa057vdW2_C3v/
PES = 'O2-Al111_HSE03-13x@RPBE_CRP_data/CRP_Data/'

atoms.set_calculator( CRPCalculator(PES) )

# Remember that the direct coordinates are for the entire super cell, and not just a single unit cell (see for example how VASP POSCARs work)
u = 0.25
v = 0.
Z = 2.256
r = 1.268
theta = 90.
phi = 60.
pos = np.dot( [u, v, 0.], Cell )
set_positions_spherical( atoms, pos[0], pos[1], Z, r, theta, phi )

print( atoms.get_potential_energy()*eV2kJ )

# Example of how to extract information about the lattice vectors from the CRP
# In principle this could have been used to set up the lattice vectors, but this is not how we implemented it in this example.
from libcrp import crp_python_wrapper
crp_python_wrapper.setup_crp_potential( PES )
print(*crp_python_wrapper.get_potential_info())
