#!/usr/bin/env python
# coding: utf-8

# # Gradients

# In[1]:


import pandas as pd
import numpy as np
import chemcoord as cc
import sympy
sympy.init_printing()


# In[2]:


molecule = cc.Cartesian.read_xyz('MIL53_beta.xyz', start_index=1)
r, theta = sympy.symbols('r, theta', real=True)


# Let's build the construction table in order to bend one of the terephtalic acid ligands.

# In[3]:


fragment = molecule.get_fragment([(12, 17), (55, 60)])
connection = np.array([[3, 99, 1, 12], [17, 3, 99, 12], [60, 3, 17, 12]])
connection = pd.DataFrame(connection[:, 1:], index=connection[:, 0], columns=['b', 'a', 'd'])
c_table = molecule.get_construction_table([(fragment, connection)])
molecule = molecule.loc[c_table.index]
zmolecule = molecule.get_zmat(c_table)


# This gives the following movement:

# In[4]:


zmolecule_symb = zmolecule.copy()
zmolecule_symb.safe_loc[3, 'angle'] += theta

cc.xyz_functions.view([zmolecule_symb.subs(theta, a).get_cartesian() for a in [-30, 0, 30]])


# # Gradient for Zmat to Cartesian

# For the gradients it is very illustrating to compare:
# $$
# f(x + h) \approx f(x) + f'(x) h
# $$
# 
# $f(x + h)$ will be ``zmolecule2``
# 
# and
# $h$ will be dist_zmol
# 
# The boolean ``chain`` argument denotes if the movement should be chained or not.

# ##### Bond

# In[5]:


dist_zmol1 = zmolecule.copy()

r = 3

dist_zmol1.unsafe_loc[:, ['bond', 'angle', 'dihedral']] = 0
dist_zmol1.unsafe_loc[3, 'bond'] = r


# In[6]:


cc.xyz_functions.view([molecule,
                       molecule + zmolecule.get_grad_cartesian(chain=False)(dist_zmol1),
                       molecule + zmolecule.get_grad_cartesian()(dist_zmol1),
                       (zmolecule + dist_zmol1).get_cartesian()])


# ##### Angle

# In[7]:


angle = 30


# In[8]:


dist_zmol2 = zmolecule.copy()
dist_zmol2.unsafe_loc[:, ['bond', 'angle', 'dihedral']] = 0
dist_zmol2.unsafe_loc[3, 'angle'] = angle


# In[9]:


cc.xyz_functions.view([molecule,
                       molecule + zmolecule.get_grad_cartesian(chain=False)(dist_zmol2),
                       molecule + zmolecule.get_grad_cartesian()(dist_zmol2),
                       (zmolecule + dist_zmol2).get_cartesian()])


# Note that the deviation between $f(x + h)$ and $f(x) + h f'(x)$ is not an error in the implementation but a visualisation of the [small angle approximation](https://en.wikipedia.org/wiki/Small-angle_approximation).
# 
# The smaller the angle the better is the linearisation.

# # Gradient for Cartesian to Zmat

# In[10]:


x_dist = 2

dist_mol = molecule.copy()
dist_mol.loc[:, ['x', 'y', 'z']] = 0.
dist_mol.loc[13, 'x'] = x_dist


# In[11]:


zmat_dist = molecule.get_grad_zmat(c_table)(dist_mol)


# It is immediately obvious, that only the ``['bond', 'angle', 'dihedral']`` of those atoms change,
# which are either moved themselves in cartesian space or use moved references.

# In[12]:


zmat_dist[(zmat_dist.loc[:, ['bond', 'angle', 'dihedral']] != 0).any(axis=1)]