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

# # Search Structures and Prepare Slabs With CatKit

# ## Preparations

# Import `requests` library for http request and define API endpoints

# In[30]:


import requests
import pprint
import ase.io
import ase.visualize
try:# catch module change from Python2->3
    import io
except:
    import StringIO as io

API = 'http://api.catalysis-hub.org'
PROTOTYPE_ENDPOINT = API + '/apps/prototypeSearch/prototype/'
STRUCTURE_ENDPOINT = API + '/apps/prototypeSearch/get_structure/'
CATKIT_SLAB_ENDPOINT = API + '/apps/catKitDemo/generate_slab_cif/'
CATKIT_SITE_ENDPOINT = API + '/apps/catKitDemo/get_adsorption_sites'


# ## Search

# Ok, we use the prototype search API to find us some Calcium Titanate structures that are Perovkites.

# In[31]:


r = requests.post(PROTOTYPE_ENDPOINT, json={"prototype":"ABC3_1_a_b_c_221"}).json()


# In[32]:


#print(r)
prototypes = r['prototypes']
#print(str(len(prototypes)) + ' Prototypes')


# Let's just print one prototype to 'see' what it looks like. It will tell use where the prototype came from (repository and handle) and gives us all parameters (Spacegroup, Wyckoff Sites, and parameters), needed to create the structure but not the actual structure.

# In[33]:


for i, prototype in enumerate(prototypes):
    pprint.pprint(prototype)
    if i == 2:
        break
#pprint.pprint(r)    


# For actually getting the structure we can use a different endpoint.

# In[34]:


print(STRUCTURE_ENDPOINT)
for prototype in prototypes:
    structure_cif = requests.post(STRUCTURE_ENDPOINT, json=prototype)
    pprint.pprint(structure_cif)
    
    pprint.pprint(structure_cif.json()['structure'])
    break


# ### Cut Slabs

# So we have a structure. Would be nice if we could use the CatKit API to cut some slabs from it and place adsorbates on it.

# In[63]:


for prototype in prototypes:
    bulk_cif = requests.post(STRUCTURE_ENDPOINT, json=prototype).json()['structure']
    params = {
        "bulkParams": {
            "format":"vasp"
            
        },
        
        "slabParams":{
        "millerX": 1,
        "millerY": 1,
        "millerZ": 1,
        "layers": 12,
        "fixed": 6,
        "unitCellSize": 4,
        "vacuum":22,
        "format": "vasp",
       
    },
        "adsorbateParams":{
            "adsorbate":"S",
            "format":"vasp"
        },
    "bulk_cif": bulk_cif,
    }
    
    r = requests.post(CATKIT_SLAB_ENDPOINT, json=params).json()
    #pprint.pprint(r['input'])

    r = requests.post(CATKIT_SITE_ENDPOINT, json=params).json()
    #pprint.pprint(r)
    with io.StringIO() as tmp_file:
        tmp_file.write(r['inputImages'][6])
        tmp_file.seek(0)
        atoms = ase.io.read(tmp_file, format='vasp')
    break


# In[64]:


ase.visualize.view(atoms * [3, 3, 1], viewer='ngl')


# In[ ]: