import pymatgen as pmg
si = pmg.Element('Si')
print('Si:')
print(si.atomic_mass)
print(si.common_oxidation_states)
print(si.number)

lattice = pmg.Lattice.cubic(4.2)
cscl = pmg.Structure(lattice, ["Cs", "Cl"], [[0,0,0],[0.5,0.5,0.5]])
print('\nCsCl structure')
print(cscl)
print('\n')
print(cscl.composition)
cscl.to(filename='POSCAR')
# ... 
# The examples are online at http://pymatgen.org/#quick-start

"""
Explanation: Assignment 3 Solution: Introduction to pymatgen
1. Install pymatgen (reccomended: install into a virtual environment. See resources)
If you need help, please send an email to whomever is helping you with this course. 
2. Open up an interactive Python interpreter (or Jupyter notebook) and go through the pymatgen quick start examples. Pay particular attention to the instance of the Structure class (rather than Element or Molecule). Since our group deals primarily with infinite crystalline lattices, we use pymatgen structures almost exclusively.
End of explanation
"""


from pymatgen import Structure, Lattice
# convert 335 pm to A:
a = 335*10**-2 # lattice parameter, a
latt = Lattice([[a, 0, 0], [0, a, 0], [0, 0, a]])
Po = Structure(latt, ['Po'], [(0, 0, 0)])
print(Po)
print(Po.volume)
print(Po.density)

"""
Explanation: 3. Create the following Structures from Lattices. Instead of using the convience class methods e.g. Lattice.cubic, create them from basis vectors. Print the volumes and densities of the structures after you create them to check yourself. The API documentation for these classes may be helpful (see resources). Pay attention to the units! You are encourage, but not required to use Structure.to('POSCAR') (or CIF) and then load the file into something like VESTA (free) or CrystalMaker (commercial).
i. Po. Simple cubic with lattice parameter of 335 pm
End of explanation
"""


a = 2.866 # lattice parameter, a
latt = Lattice([[a, 0, 0], [0, a, 0], [0, 0, a]])
Fe = Structure(latt, ['Fe', 'Fe'], [(0, 0, 0), (0.5, 0.5, 0.5)])
print(Fe)
print(Fe.volume)
print(Fe.density)

"""
Explanation: ii. Fe. Body centered cubic with lattice parameter of 2.866 Angstrom
End of explanation
"""


a = 65.89**(1.0/3.0) # lattice parameter, a
latt = Lattice([[a, 0, 0], [0, a, 0], [0, 0, a]])
Al = Structure(latt, ['Al', 'Al', 'Al', 'Al'], [(0.0,0.0,0.0), (0.5,0.5,0.0),
                                                (0.5,0.0,0.5), (0.0,0.5,0.5)])
print(Al)
print(Al.volume)
print(Al.density)

"""
Explanation: iii. Al. Face centered cubic with volume of 65.89 cubic Angstrom
End of explanation
"""


a = 5.6402 # lattice parameter, a
latt = Lattice.cubic(a)
NaCl = Structure(latt, ['Na', 'Na', 'Na', 'Na', 'Cl', 'Cl', 'Cl', 'Cl'], 
                 [(0.0,0.0,0.0), (0.5,0.5,0.0), (0.5,0.0,0.5), (0.0,0.5,0.5),
                  (0.5,0.0,0.0), (0.0,0.5,0.0), (0.0,0.0,0.5), (0.5,0.5,0.5)])
print(NaCl)
print(NaCl.volume)
print(NaCl.density)
NaCl.to(filename='POSCAR')

"""
Explanation: iv. NaCl. Rock salt structure. Use the Lattice.cubic staticmethod and a lattice parameter of 5.6402 Angstrom
End of explanation
"""


# found lattice parameters for P6_3/mmc at https://materialsproject.org/materials/mp-46/
# this i
a = 2.939
c = 4.641
latt = Lattice.hexagonal(a, c)
Ti = Structure(latt, ['Ti', 'Ti'], 
                 [(0.0,0.33333,0.0), (0.33333,0.0, 0.5)])
print(Ti)
print(Ti.volume)
print(Ti.density)
Ti.to(filename='POSCAR')

"""
Explanation: v. Ti. Hexagonal. For this, you can use the hexagonal staticmethod. Try to look up the correct lattice parameter by browsing the Materials Project. It's highly suggested to plot this and convince yourself that the answer is correct. For best results, turn on space filling and increase the boundary to see the close packed structure.
An interesting note here: first principles calculations find the most stable Ti phase to be omega-Ti, when we usually find beta-Ti (bcc) in nature. The reason for this is that bcc is dynamically stabilized (with temperature). See Kadkhodaei, Hong, and van de Walle. Phys. Rev. B 95 (2017) 064101.
Any answer for this question is really fine, but the idea was to use the hexagonal constructor and the true hcp structure is mp-46 (the omega phase is mp-72)
End of explanation
"""


a = 5.6402 # lattice parameter, a
latt = Lattice.cubic(a)
NaCl_spg = Structure.from_spacegroup('Fm-3m', latt, ['Na', 'Cl'], [(0,0,0), (0.5,0.5,0.5)])
print(NaCl_spg)
print(NaCl_spg.volume)
print(NaCl_spg.density)

"""
Explanation: HCP Ti

HCP Ti with space filling

4. Now that you have some appreciation for the tedious creation of structures from lattices, even for pure elements, you might imaging that 'there must be a better way'. You may have noticed on the Materials Project that there are many different spacegroups. Spacegroups define different symmetry elements in infinite 2D or 3D lattices and are useful for both understanding and predicting properties and useful in a practical sense. Make the following structures from spacegroups (note, if not given, you're intended to look them up)
i. NaCl has the spacegroupm F-m3-m, where an Na can be found at (0,0,0) and Cl at (0.5,0.5,0.5). Make a Structure object directly from the spacegroup and atom positions (don't make a Lattice). Look up the API to do this in the pymatgen documentation. Take note that You only have to enter one of each element, unlike in the problem above. Explain why this is.
Note that the Structure.from_spacegroup classmethod can take either the symbol or the number. I've used both here. For NaCl, you should get exactly the same result as above (unless you use a different lattice or site coordinates)
End of explanation
"""


a, b, c = 4.803, 6.048, 10.323 # lattice parameter, a
latt = Lattice.orthorhombic(a, b, c)
forsterite = Structure.from_spacegroup(62, latt, ['Mg', 'Mg', 'Si', 'O', 'O', 'O'],
                                       [(0, 0, 0), 
                                        (0.222736, 0.750000, 0.491625),
                                        (0.093638, 0.250000, 0.425974),
                                        (0.162835, 0.033014, 0.276946),
                                        (0.053218, 0.750000, 0.722149),
                                        (0.091677, 0.250000, 0.765961),])
print(forsterite)
print(forsterite.volume)
print(forsterite.density)
forsterite.to(filename='POSCAR')

"""
Explanation: ii. Pick a relatively complex (3 or more component) metal or oxide structure from the Materials Project and create it from the spacegroup. Optionally, visualize it in VESTA to verify that you constructed it correctly.
For this I chose forsterite. After applying symmetry operations, there are 28 atoms in the cell but we only needed to enter 6 positions. Symmetry is great.
End of explanation
"""


from pymatgen import MPRester
with MPRester() as mpr:
    Al_api = mpr.get_structure_by_material_id('mp-134')
print(Al_api)    
print()
print(Al_api.lattice.matrix)
print(Al_api.volume)
print(Al_api.density)

"""
Explanation: 5. We've now streamlined Structure generation, but we can do better still. Using the Materials Project API, we can directly access the structures that we want and modify them from there. Read through the couple examples at the pymatgen REST API and find the actual API documentation and skim it. Then do the following using MPRester:
i. Get the Structure object for mp-134. Print it and compare it to the structure for the same material you made before. What is different? Why?
Keep in mind that we are using with MPRester() as mpr syntax. This is called a context manager.
This Al only has one site and the lattice is different, but the density is the same. The lattice vectors are printed below. You can see in the structure result that the angles are different and the lattice parameters are smaller as well. Tiling this structure in space gives the same structure as an fcc cell that is defined by 4 atoms.
End of explanation
"""


with MPRester() as mpr:
    result = mpr.get_entry_by_material_id('mp-134', property_data=['energy_per_atom'])
print(result.energy)
print(result.energy_per_atom)
# since we just saw that the structure only had one atom, this makes sense.


"""
Explanation: ii. Get the energy of mp-134.
Hint: since we specified the mp-id, we should use get_entry_by_material_id. If you wanted to use the formula, you could have used get_entry or get_data, but you would have had to find the right form of Al.
End of explanation
"""


with MPRester() as mpr:
    result = mpr.get_data('C', prop='energy')
print(result)

"""
Explanation: iii. Get all of the energies for C (carbon) structures as a dictionary
Here we actually do want to get the data for all structures for a formula, so we'll use get_data.
End of explanation
"""


def mp_dict_to_energies(d):
    """Return list of energies from a list of dictionaries `{'material_id':mp-id, 'energy': energy}`
    
    Parameters
    ----------
    d (list): List of dictionaries of the form `{'material_id':mp-id, 'energy': energy}`
    
    Returns
    -------
    list:
        List of energies from the input dictionary
    """
    energies = []
    for entry in d:
        energies.append(entry['energy'])
    return energies
print(mp_dict_to_energies(result))

"""
Explanation: iv. Write a function that takes that dictionary and makes a list of energies (in any order)
End of explanation
"""


energies = [entry['energy'] for entry in result]
print(energies)

"""
Explanation: v. Do the same as above, except use a list comprehension instead of a function
End of explanation
"""


with MPRester() as mpr:
    struct = mpr.get_structure_by_material_id('mp-510624')
print(struct)
print(struct.formula)
print(struct.volume)
print(struct.density)

"""
Explanation: vi. What are the drawbacks of storing the energies in this way? Are there drawbacks for storing the data as two lists of ids and energies?
Variable answers. The idea was to think about differences between dicts and lists. 

A dictionary is unordered (though they are now ordered by default in Python 3.6 and later). 
Dictionaries excel when you need to drill down and look at all of the data for one structure. 
When you want to compare a certain property for all structures, you more or less have to convert to a list first. 
A database addresses both of these problems because you can query for specific entries or for a specific key/value for many entries. 

6. Finally, we will practice using some of the methods of Structure objects to transform them. You are expected to read the documentation to see what these do. If the documentation is unclear, try clicking the [Source] button on the right to see the actual Python code. Print the result of each transformation to prove to yourself that you did it correctly. Note that some transformations modify the Structure object and have no return value while others the return the new object, leaving the original unchanged. You should keep track of this and make sure to update the same structure the whole way through, applying all of the transformations to the same structure.
i. Create a Structure of strontium ferrite, mp-510624. What is the chemical formula? What is the volume of the structure? The density?
End of explanation
"""


struct.scale_lattice(56.965)
print(struct.volume)

"""
Explanation: ii. Look at the band structure and band gap for that structure on the Materials Project website. Is this a metal, semiconductor or an insulator?
There is a gap in the bandstructure at about 5 eV above the Fermi energy. The gap is 0.332 eV. That means SrFeO3 is predicted to be a semiconductor
iii. Scale the volume to 56.965 cubic Angstrom
End of explanation
"""


struct.replace_species({'Fe':'Co'})
print(struct)

"""
Explanation: iv. Replace all of the Fe species with Co in one step
End of explanation
"""


struct.make_supercell([2,2,2]) # or just struct.make_supercell(2) works for this
print(struct)

"""
Explanation: v. Make a 2x2x2 supercell
End of explanation
"""


print(struct.volume)
struct.scale_lattice(struct.volume*1.1)
print(struct.volume)

"""
Explanation: vi. Scale the volume up by 10%
End of explanation
"""


from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
sga = SpacegroupAnalyzer(struct)
print(sga.get_space_group_symbol())

"""
Explanation: vii. Import and create a SpacegroupAnalyzer from pymatgen.symmetry.analyzer with your structure. Print the spacegroup symbol
End of explanation
"""


struct.replace(24, 'Fe')
print(struct)
sga = SpacegroupAnalyzer(struct)
print(sga.get_space_group_symbol())

"""
Explanation: viii. Replace one of Co with an Fe. What is the chemical formula? What is the spacegroup symbol of this structure?
I chose the site manually from the site number printed above in the supercell.
End of explanation
"""


struct.perturb(0.1)
sga = SpacegroupAnalyzer(struct)
print(sga.get_space_group_symbol())

"""
Explanation: ix. Break the symmetry with a perturbation of 0.1 Angstrom. What is the spacegroup symbol of this structure?
Since we broke the symmetry, we see that all symmetry goes away (unless you had a 'lucky' perturbation) leaving you with a primitive Bravais and no symmetry.
End of explanation
"""


my_string = 'Python is a scientific programming language'

print('Original string:')
print(my_string)
print('Try to mutate the string:')
my_string.replace('scientific', 'fun')
print(my_string)
print('It wasn\'t replaced!')
my_string = my_string.replace('scientific', 'fun')
print(my_string)
print('Now it has been replaced.')


"""
Explanation: x. Explain what the difference is between mutable and immutable objects. What was the implication for these problems?
Mutable objects can be changed in place. They (usually) have a return value of None. In order for your structure to be updated, you run the method or function that updates the variable without you having to explictly set it again. The methods for the Structure class change the object's state so it is a mutable object. There is also an IStructure object, where I means immutable. You may have seen this in the documentation. Structure is a subclass of this object.
More examples of mutable objects are lists and dictionaries. Lists have append() and  insert() methods that update the object for you. Dictionaries are similar with the update() method.
Immutable objects cannot be changed and if you want the value of a variable with an immutable object to change, you must get a new object. Methods usually return  All of the methods for immutable objects return a new object. A main example is the string (str) object in Python. Strings cannot be changed so to update the variable, you have to set it again to the result of the method. An example of this is below, when we use the replace() method for a string, the return value is None.
End of explanation
"""


