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

# ## Set up imports and helper functions
#
# This example runs a VASP geometry optimization through AMS and PLAMS and then inspects the generated VASP input files.
#
# For a **real** VASP calculation, you must replace the `VASPExec` and `POTCARLibrary` values below with paths for your own VASP installation and POTCAR library. In this notebook they are intentionally set to bundled **dummy** example values so that you can focus on the workflow and inspect the generated VASP files.
#

import os
from pathlib import Path
from typing import Optional

import numpy as np
from scm.base import ChemicalSystem, Lattice
from scm.plams import AMSJob, Settings, view

AMSHOME = Path(os.environ["AMSHOME"])
AMSBIN = Path(os.environ["AMSBIN"])


def print_tree(root: Path, max_depth: Optional[int] = None, depth: int = 0) -> None:
    """Print a small directory tree rooted at ``root``."""
    if depth == 0:
        print(root)
    if max_depth is not None and depth >= max_depth:
        return

    entries = sorted(root.iterdir(), key=lambda p: (p.is_file(), p.name.lower()))
    for entry in entries:
        indent = "    " * depth
        suffix = "/" if entry.is_dir() else ""
        print(f"{indent}- {entry.name}{suffix}")
        if entry.is_dir():
            print_tree(entry, max_depth=max_depth, depth=depth + 1)


# ## Build the starting `ChemicalSystem`
#
# We start from `ChemicalSystem.from_smiles("COCO")`, add a cubic `Lattice`, and shift the atoms into the middle of the unit cell. We explicitly enable the `vasp` atom attributes on the `ChemicalSystem` and set `atom.vasp.label = "O_h"` for every oxygen atom - this will use the `O_h` (hard) POTCAR files from the POTCAR library.
#

system = ChemicalSystem.from_smiles("COCO")
system.lattice = Lattice([[8.0, 0.0, 0.0], [0.0, 8.0, 0.0], [0.0, 0.0, 8.0]])
system.perturb_coordinates(0.1)
system.translate([4.0, 4.0, 4.0])
system.bonds.clear_bonds()

system.enable_atom_attributes("vasp")
for atom in system:
    if atom.symbol == "O":
        atom.vasp.label = "O_h"


print(system)


view(system, width=400, height=300, guess_bonds=True, direction="tilt_z", picture_path="picture1.png")


# ## Show the dummy VASPExec and POTCAR library used in this example
#
# The current `POTCARLibrary` path is only meant as an example value. For your own calculations, point `POTCARLibrary` to your real VASP pseudopotential library, and point `VASPExec` to the command that launches your VASP executable.
#

VASP_EXEC_DUMMY = f"{AMSBIN}/amspython {AMSHOME}/scripting/scm/external_engines/backends/_vasp/fakevasp.py 6.4.0"
POTCAR_LIBRARY_DUMMY = AMSHOME / "atomicdata" / "VASP" / "test_potcars"


# ### Command to execute VASP

print("Dummy VASPExec:")
print(VASP_EXEC_DUMMY)
print()


# ### POTCAR Library

print("Dummy POTCARLibrary:")
print(POTCAR_LIBRARY_DUMMY)
print()
print("Current POTCAR library structure:")
print_tree(POTCAR_LIBRARY_DUMMY, max_depth=2)


# ## Build the PLAMS `Settings` object
#
# The `VASPExec` and `POTCARLibrary` settings are essential for the VASP engine. In this notebook they use the dummy paths shown above, so replace them before using this workflow with a real VASP installation.
#

settings = Settings()

settings.input.ams.Task = "GeometryOptimization"

settings.input.VASP.EnergyCutoff = 400
settings.input.VASP.VASPExec = VASP_EXEC_DUMMY
settings.input.VASP.POTCARLibrary = str(POTCAR_LIBRARY_DUMMY)
settings.input.VASP.KPOINTSOrigin = "Gamma-centered"
settings.input.VASP.nk1 = 1
settings.input.VASP.nk2 = 1
settings.input.VASP.nk3 = 1
settings.input.VASP.Occupation = "Gaussian"
settings.input.VASP.Smearing = 0.01

job = AMSJob(settings=settings, molecule=system, name="vasp_geoopt")
print(job.get_input())


# ## Run the geometry optimization
#
# The AMS driver performs the geometry optimization, while the VASP engine writes the VASP-style files in the worker directory. In this documentation example, the engine command is the bundled fake-VASP executable.
#

results = job.run()
print(f"Job status ok: {results.ok()}")
print(f"Job directory: {job.path}")


# ## Inspect the job directory
#
# The generated VASP files live inside the `worker.0.0.0` subdirectory of the PLAMS job directory.
#

job_dir = Path(job.path)
worker_dir = next(job_dir.glob("worker.*"))

print_tree(job_dir, max_depth=2)


# ## Inspect `INCAR`, `KPOINTS`, and `POTCAR`
#
# The files below are generated from the `Engine VASP` block and the atom labels on the `ChemicalSystem`:
#
# - `INCAR` reflects keywords such as `EnergyCutoff`, `Occupation`, and `Smearing`.
# - `KPOINTS` reflects `KPOINTSOrigin`, `nk1`, `nk2`, and `nk3`.
# - `POTCAR` is assembled from `POTCARLibrary` together with the atom labels, so the oxygen atoms labeled `O_h` pick up the `O_h/POTCAR` entry.
#

# ### INCAR

print((worker_dir / "INCAR").read_text())


# ### KPOINTS

print((worker_dir / "KPOINTS").read_text())


# ### POTCAR
#
# Note: each dummy POTCAR file is just a single line. Real POTCAR files are much bigger and have real data.

print((worker_dir / "POTCAR").read_text())


# ## Avoid looking at `POSCAR`
#
# When running VASP via AMS, the `POSCAR` in the worker directory does **not** necessarily correspond to the original input structure. This is because sometimes VASP needs to be restarted (in case of changing composition, lattice, ...).
#
# Here is the POSCAR file for reference:
#
# ### POSCAR

print((worker_dir / "POSCAR").read_text())


# In general, **prefer the normal AMS methods** to get the input or output systems from a finished job:

# ### Input system

input_system = results.get_input_system()
print(input_system)


# ### Output system

optimized_system = results.get_main_system()
print(optimized_system)


# ## Appendix: sort system by vasp.label
#
# It is common in VASP to sort the input structure by element (or, more precisely, by elemental POTCAR file). This reduces the size of the overall POTCAR file. You can do it like this:

sorted_system = system.copy()
sorted_system.enable_atom_attributes("vasp")
sorted_system.sort_atoms(lambda A, B: (A.vasp.label or A.symbol) < (B.vasp.label or B.symbol))
print(sorted_system)