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

# ## Charge transfer integrals with ADF
#
# This example computes charge-transfer integrals for a benzene dimer with an ADF fragment calculation.
#

# ### Initial Imports

from pathlib import Path
from typing import Dict

from scm.plams import ADFFragmentJob, ADFFragmentResults, Molecule, Settings, view, plot_image_grid


# ### Set up ADF Calculations

# Set common ADF settings and enable transfer-integral analysis for the full dimer job.
# A minimal basis (`SZ`) is used here for speed. For quantitative work, use a larger basis.

common = Settings()
common.input.ams.Task = "SinglePoint"
common.input.adf.Basis.Type = "SZ"
common.input.adf.Basis.Core = "None"
common.input.adf.Symmetry = "NoSym"

full = Settings()
full.input.adf.transferintegrals = True


# ### Load Benzene Dimer and Create Fragments

# Load a benzene dimer, split it into two fragments, and visualize the full dimer and fragments.
#

xyz_path = Path("BenzeneDimer.xyz")
xyz_path.write_text("""24

C      -1.90007930      -0.01180491      -1.63051319
C       0.87349469      -0.01023939      -1.76821915
C      -1.20564674      -1.21381033      -1.65931829
C      -1.20758387       1.19098295      -1.67103029
C       0.17915587       1.19167022      -1.73987307
C       0.18110787      -1.21293268      -1.72799219
H      -2.98312785      -0.01239662      -1.57479944
H      -1.74970769       2.12993064      -1.64644374
H       0.72018058       2.13112890      -1.76828739
H       1.95680577      -0.00959486      -1.81797251
H       0.72362449      -2.15176759      -1.74701794
H      -1.74626146      -2.15334648      -1.62565114
C      -1.04276513       0.00660438       1.20777829
C       1.73079327       0.00816985       1.07007850
C      -0.34843214      -1.19529353       1.17943874
C      -0.35038422       1.20928613       1.16755800
C       1.03636300       1.21016618       1.09888633
C       1.03830008      -1.19460875       1.11059824
H      -2.12607520       0.00595978       1.25753750
H      -0.89290110       2.14811978       1.18659394
H       1.57697669       2.14970259       1.06522938
H       2.81384166       0.00876150       1.01437093
H       1.58042281      -2.13355664       1.08602198
H      -0.88945704      -2.13475088       1.20786337
""")

dimer = Molecule(str(xyz_path))
dimer.guess_bonds()
fragments = dimer.separate()
if len(fragments) != 2:
    raise RuntimeError(f"Expected 2 fragments, got {len(fragments)}")

fragment1, fragment2 = fragments
imgs = {
    name: view(mol, width=250, height=250)
    for name, mol in [("Dimer", dimer), ("Fragment 1", fragment1), ("Fragment 2", fragment2)]
}
plot_image_grid(imgs, rows=1, save_path="picture1.png")
# ### Run the Fragment Calculation and Extract Results

job = ADFFragmentJob(
    name="ADFTI",
    fragment1=fragment1,
    fragment2=fragment2,
    settings=common,
    full_settings=full,
)
results = job.run()


transfer_integrals = results.job.full.results.read_rkf_section("TransferIntegrals", file="adf")
print("== Results ==")
for key, value in sorted(transfer_integrals.items()):
    print(f"{key:<28}: {float(value):>12.6f}")