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

# ## Initial imports

import scm.plams as plams
from scm.params import ResultsImporter
from scm.plams import Settings, AMSJob, log, Molecule

# common_ru_h.py must exist in the current working directory
from common_ru_h import (
    rotation,
    dft_settings,
    QEKPointsConfig,
    lattice_optimization_settings,
    plot_pesscan,
    check_installation,
)

check_installation()


# ## Initialize PLAMS working directory

plams.init()


# ## Bulk structure: hcp Ru

initial_bulk = plams.Molecule()
a = 2.7  # hexagonal lattice parameter, angstrom
c = 4.2768  # hexagonal lattice parameter, angstrom
initial_bulk.add_atom(plams.Atom(symbol="Ru", coords=(0.0, 0.0, 0.0)))
initial_bulk.add_atom(plams.Atom(symbol="Ru", coords=(0.0, a / 3**0.5, c / 2)))
initial_bulk.lattice = [[a, 0, 0], [-a / 2, a * 3**0.5 / 2, 0], [0, 0, c]]
log("Initial structure")
log(initial_bulk)


plams.plot_molecule(initial_bulk, rotation=rotation)


# ## Lattice optimization of bulk Ru with DFT

lattopt_job = plams.AMSJob(
    settings=dft_settings(QEKPointsConfig(11, 11, 11)) + lattice_optimization_settings(),
    name="hcp_lattopt_Ru_dft",
    molecule=initial_bulk,
)
lattopt_job.run()


optimized_bulk: Molecule = lattopt_job.results.get_main_molecule()  # type: ignore
log(optimized_bulk)
log(f"Density: {optimized_bulk.get_density():.2f} kg/m^3")


# ## Volume scan of bulk Ru with DFT

from common_ru_h import (
    dft_settings,
    QEKPointsConfig,
    pesscan_settings,
    CellVolumeScalingRangeScanCoordinate,
)

s = dft_settings(QEKPointsConfig(11, 11, 11))
s += pesscan_settings([CellVolumeScalingRangeScanCoordinate(0.85, 1.15)], n_points=7)
volume_scan_job = AMSJob(
    settings=s,
    molecule=optimized_bulk,
    name="bulk_hcp_Ru_volume_scan_dft",
)
volume_scan_job.run()


plot_pesscan(volume_scan_job)


# ## Bond scan of H2 with DFT

h2_mol = plams.from_smiles("[HH]")
h2_mol.lattice = [[5, 0, 0], [0, 5, 0], [0, 0, 5]]
plams.plot_molecule(h2_mol, rotation=rotation)


from common_ru_h import (
    dft_settings,
    QEKPointsConfig,
    pesscan_settings,
    DistanceScanCoordinate,
)

scan_coordinate = DistanceScanCoordinate(atom1=1, atom2=2, start=0.55, end=1.0)
s = dft_settings(QEKPointsConfig(1, 1, 1))
s += pesscan_settings([scan_coordinate], n_points=7)

h2_bond_scan_job = AMSJob(settings=s, molecule=h2_mol, name="h2_bond_scan_dft")

h2_bond_scan_job.run()


plot_pesscan(h2_bond_scan_job)


# ## Store results

ri = ResultsImporter.from_ase()
properties = ["energy", "forces"]
ri.add_pesscan_singlepoints(volume_scan_job, properties=properties)
ri.add_pesscan_singlepoints(h2_bond_scan_job, properties=properties)
ri.add_singlejob(lattopt_job, task="SinglePoint", properties=properties)

# Also add as PES Scans - these will not be used during training but
# will plot the energy-volume curve and bond-scan curve at the end
# of the training
ri.add_singlejob(volume_scan_job, task="PESScan", properties=["pes"])
ri.add_singlejob(h2_bond_scan_job, task="PESScan", properties=["pes"])

ri.store("reference_data_1")