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

# ## Initial imports, files, and helper function

import os
from pathlib import Path
from scm.plams import view

# Make sure to source amsbashrc.sh before launching this example so that the AMSHOME environment variable is set.

AMSHOME = os.environ["AMSHOME"]
cif_file = f"{AMSHOME}/atomicdata/Molecules/IZA-Zeolites/ABW.cif"
xyz_file = f"{AMSHOME}/scripting/scm/params/examples/benchmark/ISOL6/e_13.xyz"
badxyz_file = f"{AMSHOME}/scripting/scm/plams/unit_tests/xyz/reactant2.xyz"

assert Path(cif_file).exists(), f"{cif_file} does not exist."
assert Path(xyz_file).exists(), f"{xyz_file} does not exist."

output_dir = Path("molecule_formats_outputs")
output_dir.mkdir(exist_ok=True)


def head(filename, n: int = 4):
    """Print the first ``n`` lines of a file"""
    with open(filename, "r") as f:
        lines = f.readlines()
        lines = lines[: min(n, len(lines))]
    print("".join(lines))


# ## SMILES

# ### Load PLAMS Molecule from SMILES string

from scm.plams import from_smiles, Molecule, plot_molecule

mol = from_smiles("CCCCO")
print(f"{type(mol)=}")
view(mol, padding=-0.5, width=300, height=300, picture_path=output_dir / "example_mol.png")


# ### Convert PLAMS Molecule to SMILES string
#
# Note: This requires that bonds are defined in the PLAMS Molecule.

from scm.plams import to_smiles

smiles = to_smiles(mol)
print(smiles)


# ## .xyz files

# ### Load PLAMS Molecule from .xyz file

from scm.plams import Molecule, plot_molecule

mol = Molecule(xyz_file)
print(f"{type(mol)=}")
view(mol, guess_bonds=True, padding=-2, width=500, height=300, picture_path="picture1.png")


# ### Write PLAMS Molecule to .xyz file

mol.properties.comment = "The comment line (2nd line after the number of atoms)"
mol.write(output_dir / "out.xyz")


head(output_dir / "out.xyz")


# ## .cif files

# ### Load PLAMS Molecule from .cif file
#
# PLAMS cannot natively read .cif files. Instead, go through another library, for example ASE or pymatgen.

from ase.io import read
from scm.plams import fromASE

mol: Molecule = fromASE(read(cif_file))
print(f"{type(mol)=}")
view(mol, fixed_atom_size=False, padding=-2, width=400, height=300, picture_path="picture2.png")


# ### Write PLAMS Molecule to .cif file
#
# PLAMS cannot natively export to .cif files. Instead, go through another library, for example ASE or pymatgen.
#
# ASE can be used to write many file formats. See https://wiki.fysik.dtu.dk/ase/ase/io/io.html

from scm.plams import toASE

toASE(mol).write(output_dir / "out.cif")
head(output_dir / "out.cif")


# ## AMS .in system block format
#
# ### Write PLAMS Molecule to AMS .in system file

mol.write(output_dir / "ams_system_block.in")
head(output_dir / "ams_system_block.in")


# ### Load PLAMS Molecule from AMS .in system file

from scm.plams import Molecule

mol = Molecule(output_dir / "ams_system_block.in")
view(mol, fixed_atom_size=False, padding=-2, width=400, height=300, picture_path="picture3.png")


# ## POSCAR/CONTCAR (VASP input format)

# ### Write PLAMS Molecule to POSCAR/CONTCAR (VASP input format)
#
# ASE can be used to write many file formats. See https://wiki.fysik.dtu.dk/ase/ase/io/io.html

from scm.plams import toASE

toASE(mol).write(output_dir / "POSCAR")
head(output_dir / "POSCAR", 10)


# ### Load PLAMS Molecule from POSCAR/CONTCAR (VASP input format)

from scm.plams import fromASE
from ase.io import read

mol: Molecule = fromASE(read(output_dir / "POSCAR"))

print(f"{type(mol)=}")
view(mol, fixed_atom_size=False, padding=-2, width=400, height=300, picture_path="picture4.png")


# ## ASE Atoms Python class

# ### Convert PLAMS Molecule to ASE Atoms

from scm.plams import toASE
from ase import Atoms
from ase.visualize.plot import plot_atoms
import matplotlib.pyplot as plt

print(f"{type(mol)=}")
print(f"{mol.get_formula()=}")

ase_atoms: Atoms = toASE(mol)
print(f"{type(ase_atoms)=}")
print(f"{ase_atoms.get_chemical_formula()=}")

_, ax = plt.subplots(figsize=(2, 2))
plot_atoms(ase_atoms, rotation="-85x,5y,0z", ax=ax)
# ### Convert ASE Atoms to PLAMS Molecule

from scm.plams import fromASE, plot_molecule, Molecule

mol: Molecule = fromASE(ase_atoms)
print(f"{type(mol)=}")
view(
    mol,
    direction="tilt_z",
    show_lattice_vectors=True,
    padding=-1,
    fixed_atom_size=False,
    width=400,
    height=400,
    picture_path="picture5.png",
)


# ## RDKit Mol Python class

# ### Convert PLAMS Molecule to RDKit Mol

from scm.plams import to_rdmol, Molecule
from rdkit.Chem import Draw
from rdkit.Chem.Draw import IPythonConsole

IPythonConsole.ipython_useSVG = True
IPythonConsole.molSize = 250, 250

plams_mol = Molecule(xyz_file)
# guess bonds, the bonds will be included in the RDKit molecule
plams_mol.guess_bonds()

rdkit_mol = to_rdmol(plams_mol)
print(f"{type(rdkit_mol)=}")
rdkit_mol


# ### Convert RDKit Mol to PLAMS Molecule

from scm.plams import from_rdmol, plot_molecule, Molecule

mol: Molecule = from_rdmol(rdkit_mol)

print(f"{type(rdkit_mol)=}")
print(f"{type(mol)=}")
view(mol, padding=-2, width=500, height=300, picture_path="picture6.png")


# ### Convert problematic PLAMS Molecule to RDKit Mol

mol = Molecule(badxyz_file)
mol.guess_bonds()
view(mol, padding=-0.5, width=300, height=300, picture_path="picture7.png")


# This molecule will fail to convert to an RDKit Mol object, because RDKit does not like the AMS assignment of double bonds.

try:
    rdkit_mol = to_rdmol(mol)
except ValueError as exc:
    print("Failed to convert")


# The problem can be fixed by passing the argument `presanitize` to the `to_rdmol` function.

rdkit_mol = to_rdmol(mol, presanitize=True)
rdkit_mol


# ## SCM Base ChemicalSystem Python class
#
# ### Convert PLAMS Molecule to ChemicalSystem

from scm.utils.conversions import plams_molecule_to_chemsys, chemsys_to_plams_molecule
from scm.plams import Molecule
from scm.base import ChemicalSystem

mol = Molecule(xyz_file)
chemsys = plams_molecule_to_chemsys(mol)
print(f"{type(chemsys)=}")
print(chemsys)


# ### Convert ChemicalSystem to PLAMS Molecule

from scm.utils.conversions import plams_molecule_to_chemsys, chemsys_to_plams_molecule
from scm.plams import Molecule
from scm.base import ChemicalSystem

mol = chemsys_to_plams_molecule(chemsys)
print(f"{type(chemsys)=}")
print(f"{type(mol)=}")
view(mol, guess_bonds=True, padding=-2, width=500, height=300, picture_path="picture8.png")


# ## pymatgen Structure and Molecule Python classes

# ### Convert PLAMS Molecule to pymatgen Structure (periodic)
#
# There is no builtin converter between PLAMS Molecule and pymatgen Structure (periodic crystal). Instead, you need to go through the ASE interface to both packages:

from pymatgen.core.structure import Structure
from pymatgen.io.ase import AseAtomsAdaptor
import scm.plams
from scm.plams import fromASE, toASE, Molecule
from ase.io import read


def convert_plams_molecule_to_pymatgen_structure(mol: Molecule) -> Structure:
    return AseAtomsAdaptor().get_structure(toASE(mol))


mol: scm.plams.Molecule = fromASE(read(cif_file))

pymatgen_structure: Structure = convert_plams_molecule_to_pymatgen_structure(mol)

print(f"{type(mol)=}")
print(f"{type(pymatgen_structure)=}")
print(pymatgen_structure)


# ### Convert pymatgen Structure (periodic) to PLAMS Molecule
#
# Go through the ASE interface:

from pymatgen.io.ase import AseAtomsAdaptor
from pymatgen.core.structure import Structure
from scm.plams import fromASE
from scm.plams import Molecule


def pymatgen_structure_to_plams_molecule(pymatgen_structure: Structure) -> Molecule:
    return fromASE(AseAtomsAdaptor().get_atoms(pymatgen_structure))


print(f"{type(pymatgen_structure)=}")

mol = pymatgen_structure_to_plams_molecule(pymatgen_structure)
print(f"{type(mol)=}")


# ### Convert PLAMS Molecule to pymatgen Molecule (non-periodic)
#
# pymatgen has a special ``Molecule`` class for non-periodic systems. In PLAMS, the ``Molecule`` class is used for both periodic and non-periodic systems.

import pymatgen.core.structure
import scm.plams
from pymatgen.io.ase import AseAtomsAdaptor
from scm.plams import toASE


def convert_plams_molecule_to_pymatgen_molecule(
    mol: scm.plams.Molecule,
) -> pymatgen.core.structure.Molecule:
    return AseAtomsAdaptor().get_molecule(toASE(mol))


plams_molecule = scm.plams.Molecule(xyz_file)

pymatgen_molecule: pymatgen.core.structure.Molecule = convert_plams_molecule_to_pymatgen_molecule(plams_molecule)

print(f"{type(plams_molecule)=}")
print(f"{type(pymatgen_molecule)=}")
print(pymatgen_molecule)


# ### Convert pymatgen Molecule (non-periodic) to PLAMS Molecule

from pymatgen.io.ase import AseAtomsAdaptor
import pymatgen.core.structure
from scm.plams import fromASE
from scm.plams import Molecule


def pymatgen_molecule_to_plams_molecule(
    pymatgen_molecule: pymatgen.core.structure.Molecule,
) -> scm.plams.Molecule:
    return fromASE(AseAtomsAdaptor().get_atoms(pymatgen_molecule))


print(f"{type(pymatgen_molecule)=}")

mol = pymatgen_molecule_to_plams_molecule(pymatgen_molecule)
print(f"{type(mol)=}")
view(
    mol,
    guess_bonds=True,
    padding=-2,
    width=500,
    height=300,
    picture_path=output_dir / "example_mol.png",
)