from smact.structure_prediction import prediction, database, mutation, probability_models, structure, utilities
import json
import itertools
from itertools import zip_longest
import smact

# An optional utility to display a progress bar
# for long-running loops. `pip install tqdm`.
from tqdm import tqdm
from ipywidgets import widgets

from pymatgen.ext.matproj import MPRester
from pprint import pprint
from pymatgen.analysis import structure_matcher
from pymatgen.util.plotting import pretty_plot
import pandas as pd

"""
Explanation: Create a database of structures
This notebook provides a short demo on how to use SMACT and the Materials Project to generate a database of materials which could be used for structure prediction.
In the example below, we generate a database of garnet materials
End of explanation
"""


#Use the API Key
m = MPRester()
#We use a csv file downloaded from the web version of MP with the MP ids
mp_df=pd.read_csv("_Materials Project .csv")

#Extracts a list of mp-ids
mp_ids=mp_df["Materials Id"].to_list()

print(f"We have {len(mp_ids)} entries with formula A3B2C3O12")
mp_df.head(10)

#Query materials project
data=m.query(criteria={"task_id": {"$in": mp_ids}}, properties=["pretty_formula","material_id","spacegroup.symbol","icsd_ids","e_above_hull","exp","structure","cif"])

"""
Explanation: Querying the MP for garnets
End of explanation
"""


#get structure of Ca3Fe2(SiO4)3
SM=structure_matcher.StructureMatcher(attempt_supercell=True)

known_garnet=m.query("mp-6672", properties=["pretty_formula","material_id","spacegroup.symbol","icsd_ids","e_above_hull","exp","structure","cif"])
known_garnet_structure=known_garnet[0]["structure"]

#Iterate over query data and verify if they have the correct structure
fitted_data=[]
for i in data:
    if SM.fit_anonymous(i['structure'], known_garnet_structure):
        fitted_data.append(i)
print(len(fitted_data))



"""
Explanation: Structure matching
Here, we use the structure of Ca3Fe2(SiO4)3 to filter our query data for only materials which form the garnet structure.
End of explanation
"""



experimental_list=[]
theoretical_list=[]
for i in fitted_data:
    if len(i["icsd_ids"])!=0:
        experimental_list.append(i)
    else:
        theoretical_list.append(i)
print(len(experimental_list))
print(len(theoretical_list))

"""
Explanation: Sorting out experimental data
Here, we determine which garnet structures are experimental (i.e. have a corresponding structure in the ICSD) or theoretical (i.e. has had electronic structure calculations done and submitted to the MP).
End of explanation
"""


#Query materials project for X3Y5O12
data_2=m.query("*3*5O12", properties=["pretty_formula","material_id","spacegroup.symbol","icsd_ids","e_above_hull","exp","structure","cif"])
print(f"There are {len(data_2)} materials with formula X3Y2Y3O12 in the MP")
print("")

#get structure of YAG
YAG=m.query("mp-3050", properties=["pretty_formula","material_id","spacegroup.symbol","icsd_ids","e_above_hull","exp","structure","cif"])
YAG_structure=YAG[0]["structure"]

#Iterate over query data and verify if they have the correct structure
fitted_data_2=[]
for i in data_2:
    if SM.fit_anonymous(i['structure'], YAG_structure):
        fitted_data_2.append(i)
print(f"Of the {len(data_2)} materials, {len(fitted_data_2)} match the structure of YAG")
print("")

#Find the number of experimental and theoretical materials
experimental_list_2=[]
theoretical_list_2=[]
for i in fitted_data_2:
    if len(i["icsd_ids"])!=0:
        experimental_list_2.append(i)
    else:
        theoretical_list_2.append(i)
print(f"There are {len(experimental_list_2)} experimental garnet structures with formula X3Y5O12")
print(f"There are {len(theoretical_list_2)} theoretical structures with formula X3Y5O12")

print(f"Considering formulas A3B2C3O12 and X3Y5O12, suggest there are {len(fitted_data)+len(fitted_data_2)} materials with the garnet structure in the Materials Project")
print("")
print(f"Of these {len(fitted_data)+len(fitted_data_2)} materials, {len(experimental_list)+len(experimental_list_2)} are experimental ")
print("")
print(f"Of these {len(fitted_data)+len(fitted_data_2)} materials, {len(theoretical_list)+len(theoretical_list_2)} are theoretical ")

new_fitted_data=fitted_data+fitted_data_2
new_theoretical_list=theoretical_list+theoretical_list_2
new_experimental_list=experimental_list+experimental_list_2
df_new=pd.DataFrame(new_fitted_data)

df_new.to_csv("MP_garnets.csv", index=False)
print(df_new.shape)
df_new.head()

theoretical=[]
for i in df_new["icsd_ids"]:
    if len(i)!=0:
        theoretical.append("No")
    else:
        theoretical.append("Yes")
df_new["theoretical?"]=theoretical

    

"""
Explanation: Other garnet structures
Garnet materials with formulae A3B2C3D12 and X3Y5Z12 both have the same crystal structure however, the pymatgen structure matcher is unable to match a garnet with the former formula to a garnet with the latter formula. As well, querying the MP for the quarternary formula will not produce results containing the ternary formula and vice versa.
Hence, for this particular example we also query for garnets with the ternary formula.
End of explanation
"""


#This creates the database object
DB=database.StructureDB("Garnets.db")

#These create tables within the database
DB.add_table("Garnets")
DB.add_table("Experimental")
DB.add_table("Theoretical")

#Create an iterable of the query data
structs=[]
for i in new_fitted_data:
    structs.append(database.parse_mprest(i))
    
exp_structs=[]
for i in new_experimental_list:
    exp_structs.append(database.parse_mprest(i))
    
theo_structs=[]
for i in new_theoretical_list:
    theo_structs.append(database.parse_mprest(i))
    
#Uncomment the line below the first time you run this notebook
DB.add_structs(structs, "Garnets")
DB.add_structs(exp_structs, "Experimental")
DB.add_structs(theo_structs, "Theoretical")

"""
Explanation: Storing in SMACT compatible database
Now that we have all filtered query data, we now proceed to store the results in a locally accessible database.
End of explanation
"""


