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

# ## Initial imports

from scm.plams import *

from ase import Atoms
from pymatgen.core.structure import Structure
from pymatgen.analysis.diffraction.xrd import XRDCalculator


# ## Create ASE atoms object for BaTiO3

at = Atoms(
    symbols=[
        "Ba",
        "Ti",
        "O",
        "O",
        "O",
    ],
    scaled_positions=[
        [
            0.0,
            0.0,
            0.0,
        ],
        [0.5, 0.5, 0.5],
        [0.0, 0.0, 0.5],
        [0.0, 0.5, 0.0],
        [0.5, 0.0, 0.0],
    ],
    cell=[4.01, 4.01, 4.01],
    pbc=(True, True, True),
)
view(fromASE(at), fixed_atom_size=False, direction="large_tilt_z", width=200, height=200, picture_path="picture1.png")


# ## Save ASE Atoms to .cif format

fname = "batio3.cif"
at.write(fname)


# ## Load .cif in pymatgen and calculate XRD

# Available radiation sources:

print(f"Available radiation sources: {XRDCalculator.AVAILABLE_RADIATION}")


# Let's choose Cu K-alpha (default):

structure = Structure.from_file(fname)
xrd_calc = XRDCalculator(wavelength="CuKa")
xrd_calc.show_plot(structure)


pattern = xrd_calc.get_pattern(structure)
print("2*Theta Intensity hkl d_hkl(angstrom)")
for two_theta, intensity, hkls, d_hkl in zip(pattern.x, pattern.y, pattern.hkls, pattern.d_hkls):
    hkl_tuples = [hkl["hkl"] for hkl in hkls]
    for hkl in hkl_tuples:
        label = ", ".join(map(str, hkl))
        print(f"{two_theta:.2f} {intensity:.2f} {hkl} {d_hkl:.3f}")