# --- INITIAL DEFINITIONS ---
import numpy, math, random
from visualise import view
from ase import Atoms
import sys
sys.path.insert(0, './data/descriptor_codes/')
sys.path.insert(0, './data/descriptor_codes/src')
from dscribe.descriptors import SOAP

"""
Explanation: Smooth Overlap of Atomic Positions
SOAP is a local descriptor, that maps the local environment around a point very accurately. It eliminates rotational, and permutation redundancies by integrating the overlap of smoothed out atomic positions, by gaussian smearing, and mapping them into coefficients of orthornormal basis functions.
This is done by the following steps:


Smooth out the atomic positions:
The atomic positions are point objects in space. Integrating them would need a lot of basis functions. Thus, the atoms' positions are smeared as gaussian functions. 
$$ \rho(r) = \sum_i e^{-(r-r_i)^2}$$
However, this also makes all the elements indistinguishable. Thus, SOAP for individual elements, in molecule/unit-cell, is calculated, and then the values are concantenated at end. 
 Image courtesy Jäger Marc


Generate orthonormal basis set:
The obtained smeared atomic position, or atomic density, if you will, is decomposed using Laplace Spherical Harmonics -- spherical harmonics in real space -- and orthogonal basis set: $\Upsilon_{lm}(\theta, \phi)$ and  $g_n(r) $. 
Basis function for s orbital of hydrogen:  
Laplace spherical harmonics $\Upsilon_{ℓm}$ for l = 0, …, 4 (top to bottom) and m = 0, …, l (left to right). The negative order harmonics $\Upsilon_{ℓ-m}$ would be shown rotated about the z axis by $90^o$ with respect to the positive order ones.
  Image courtsey wikipedia.org/wiki/User:Cyp  


Integrate for all coefficients:
    $$c_{nlm} = \left< \rho | g_n(r)\Upsilon_{lm} \right>  = \int_V g_n(r)\Upsilon_{lm}(\theta, \phi)\rho(r, \theta, \phi)dV$$
    Further, a power spectrum, or a density matrix, per se, is made out of these parameters and summed for all m's for rotational invarience.
    $$P_{nn'l} = \sum_m c_{nlm}c^*_{n'lm}$$


For more info see:
Bartók, Albert P., Risi Kondor, and Gábor Csányi. <i>Physical Review B 87.18</i> (2013): <b>184115</b>
For calculating SOAP, we use the DScribe package as developed by Surfaces and Interfaces at the Nanoscale, Aalto
Example
We are going to see SOAP in action for a simple NaCl system.
End of explanation
"""


# Define the system under study: NaCl in a conventional cell.
NaCl_conv = Atoms(
    cell=[
        [5.6402, 0.0, 0.0],
        [0.0, 5.6402, 0.0],
        [0.0, 0.0, 5.6402]
    ],
    scaled_positions=[
        [0.0, 0.5, 0.0],
        [0.0, 0.5, 0.5],
        [0.0, 0.0, 0.5],
        [0.0, 0.0, 0.0],
        [0.5, 0.5, 0.5],
        [0.5, 0.5, 0.0],
        [0.5, 0.0, 0.0],
        [0.5, 0.0, 0.5]
    ],
    symbols=["Na", "Cl", "Na", "Cl", "Na", "Cl", "Na", "Cl"],
)
view(NaCl_conv)

"""
Explanation: Atom description
We'll make an ase.Atoms class for NaCl:
End of explanation
"""


# Computing SOAP
calcpos = [0, 0, 0]

soaper = SOAP(
    rcut=8,
    nmax=5,
    lmax=5,
    species=['Na', 'Cl'],
    sparse=False
)

"""
Explanation: Setting SOAP hyper-parameters
Next we set the hyper-parameters to SOAP. 
1. calcpos, center of SOAP calculation
1. rcut, sets the cutoff for atoms whose gaussian densities will be included in the integral.
2. nmax, sets the number of orthogonal radial basis functions to use.
3. lmax, sets the number of angular momentum terms, so l = 0, 1, ...., lmax
**Note: even when giving one SOAP calculation position, it should be wrapped in a list, as shown in example below**
End of explanation
"""


#calculation
soap1 = soaper.create(NaCl_conv, positions=[calcpos])
print("Size of descriptor: {}".format(soap1.shape[1]))
print("First five values, for position {}: {}".format(calcpos, soap1[0,:5]))

"""
Explanation: Calculation
Now we call the soap function, and pass all the parameters
End of explanation
"""


#Rotation of positions

print("Original positions:\n {}".format(NaCl_conv.positions))

NaCl_conv.rotate(90, [0,1,1], center=calcpos)

print("Rotated positions:\n {}".format(NaCl_conv.positions))

view(NaCl_conv)

"""
Explanation: Rotational invariance
End of explanation
"""


soap2 = soaper.create(NaCl_conv, positions=[calcpos])
print(numpy.linalg.norm(soap1 - soap2))

"""
Explanation: Recompute SOAP for the same atom, after rotation and show the difference in descriptors:
End of explanation
"""


# DIY...

"""
Explanation: Remark
The power spectrum at a desired position x is the fingerprint of the local chemical environment at this specific position. Thus, it can be used to:
1. Compare the similarity of two local chemical environments by comparing their SOAP descriptors.
2. Machine learn local properties, like charges, adsorption energies, etc.
Exercises
1. Smoothness
Verify that the SOAP is smooth under translations of point of interest.
End of explanation
"""


# atomic positions as matrix
molxyz = numpy.load("./data/molecule.coords.npy")
# atom types
moltyp = numpy.load("./data/molecule.types.npy")

atoms_sys = Atoms(positions=molxyz, numbers=moltyp)
view(atoms_sys)

# build SOAP at each atom location
# ...
# compute average soap for each specie
# ...
# concatenate the soaps to the the overall global one
# ...

"""
Explanation: 2. Construct a global SOAP
Use the atomic environments to construct an average SOAP descriptor for molecules
End of explanation
"""


