# pdb2reaction/dft.py

"""
Single-point DFT calculation with GPU acceleration (GPU4PySCF or CPU PySCF).

Example:
    pdb2reaction dft -i input.pdb -q 0 -m 1 --func-basis 'wb97m-v/6-31g**'

For detailed documentation, see: docs/dft.md
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Dict, Optional, Sequence, Tuple, List, Union

import sys
import traceback
import textwrap
import time
from functools import reduce

import click
import yaml
import numpy as np

from pysisyphus.helpers import geom_loader
from pysisyphus.constants import AU2KCALPERMOL

from .utils import (
    apply_yaml_overrides,
    pretty_block,
    format_geom_for_echo,
    format_elapsed,
    prepared_cli_input,
    set_convert_file_enabled,
    YamlFlowList,
    cli_param_overridden,
)
from .cli_utils import resolve_yaml_sources, load_merged_yaml_cfg
from .defaults import GEOM_KW_DEFAULT

logger = logging.getLogger(__name__)


# -----------------------------------------------
# Defaults (override via CLI / YAML)
# -----------------------------------------------

DFT_DEFAULT_FUNC = "wb97m-v"
DFT_DEFAULT_BASIS = "def2-tzvpd"

DFT_KW: Dict[str, Any] = {
    "conv_tol": 1e-9,          # SCF convergence tolerance (Eh)
    "max_cycle": 100,          # Maximum number of SCF iterations
    "grid_level": 3,           # Numerical integration grid level (PySCF grids.level)
    "verbose": 0,              # PySCF verbosity (0..9)
    "out_dir": "./result_dft/",# Output directory
    "func": DFT_DEFAULT_FUNC,  # XC functional (can be overridden via YAML)
    "basis": DFT_DEFAULT_BASIS,# Basis set (can be overridden via YAML)
}


# -----------------------------------------------
# Utilities
# -----------------------------------------------

def _parse_func_basis(s: str) -> Tuple[str, str]:
    """
    Parse 'FUNC/BASIS' into (xc, basis).
    Mixed case is accepted (PySCF is case-insensitive for common names).
    """
    if not s or "/" not in s:
        raise click.BadParameter("Expected 'FUNC/BASIS' (e.g., 'wb97m-v/def2-tzvpd').")
    func, basis = s.split("/", 1)
    func = func.strip()
    basis = basis.strip()
    if not func or not basis:
        raise click.BadParameter("Functional or basis is empty. Example: --func-basis 'wb97m-v/6-31g**'")
    return func, basis


def _geometry_to_pyscf_atoms_string(geometry) -> Tuple[str, Sequence[Tuple[str, Tuple[float, float, float]]]]:
    """
    Convert a pysisyphus Geometry to (xyz_string, PySCF atom list).
    The atom list is [(symbol, (x, y, z)), ...] in Angstrom.
    """
    s = geometry.as_xyz()
    lines = s.splitlines()
    atoms: list[Tuple[str, Tuple[float, float, float]]] = []
    for ln in lines[2:]:
        parts = ln.split()
        if len(parts) >= 4:
            sym = parts[0]
            x, y, z = map(float, parts[1:4])
            atoms.append((sym, (x, y, z)))
    return s, atoms


def _AU2KCALPERMOL(Eh: float) -> float:
    return float(Eh * AU2KCALPERMOL)


def _configure_scf_object(mf, dft_cfg: Dict[str, Any], xc: str):
    """Apply common SCF settings (XC, DF, tolerances, grids) to an SCF object."""
    mf.xc = xc
    mf.max_cycle = int(dft_cfg["max_cycle"])
    mf.conv_tol = float(dft_cfg["conv_tol"])
    mf.grids.level = int(dft_cfg["grid_level"])
    mf.chkfile = None
    mf = mf.density_fit()

    return mf


def _format_row_for_echo(row: List[Union[int, str, float, None]]) -> str:
    """Format a row like: [0, H, 0.0, 0.0, 0.0]."""
    return "[" + ", ".join(_format_row_value_for_echo(v) for v in row) + "]"


def _format_row_value_for_echo(x: Union[int, str, float, None]) -> str:
    if x is None:
        return "null"
    if isinstance(x, float):
        return f"{x:.10g}"
    return str(x)


# This function is based on https://pyscf.org/_modules/pyscf/lo/iao.html
def fast_iao_mullikan_spin_pop(mol, dm, iaos, verbose=None):
    """
    IAO-basis Mulliken spin population analysis.

    Args:
        mol : Mole or Cell object
        dm  : AO density matrix; for UKS/UHF, a (2, nao, nao) array
        iaos: 2D array of IAO orbitals (orthogonal or non-orthogonal)
        verbose: PySCF logger level

    Returns:
        (spin_pop_ao, Ms_by_atom)
        spin_pop_ao : Mulliken spin population on each IAO
        Ms_by_atom  : Mulliken spin density per atom (sum over IAOs on atom)
    """
    import numpy
    from pyscf.lib import logger as pyscf_logger
    from pyscf.lo.iao import reference_mol
    from pyscf.scf import uhf as scf_uhf

    if verbose is None:
        verbose = pyscf_logger.DEBUG

    pmol = reference_mol(mol)
    # Overlap in the large basis
    if getattr(mol, 'pbc_intor', None):  # cell?
        ovlpS = mol.pbc_intor('int1e_ovlp')
    else:
        ovlpS = mol.intor_symmetric('int1e_ovlp')

    # Transform DM in big basis to IAO basis
    # |IAO> = |big> C
    # DM_IAO = C^{-1} DM (C^{-1})^T = S_IAO^{-1} C^T S DM S C S_IAO^{-1}
    cs = numpy.dot(iaos.T.conj(), ovlpS)
    s_iao = numpy.dot(cs, iaos)
    iao_inv = numpy.linalg.solve(s_iao, cs)

    # Restricted case: spin density is identically zero
    if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
        spin_pop_ao = numpy.zeros(s_iao.shape[0], dtype=float)
        Ms = numpy.zeros(pmol.natm, dtype=float)
        return spin_pop_ao, Ms

    # Unrestricted: transform alpha/beta DM to IAO basis
    dm_a = reduce(numpy.dot, (iao_inv, dm[0], iao_inv.conj().T))
    dm_b = reduce(numpy.dot, (iao_inv, dm[1], iao_inv.conj().T))

    return scf_uhf.mulliken_spin_pop(pmol, [dm_a, dm_b], s_iao, verbose)


def _get_occupied_orbitals(mf) -> np.ndarray:
    """Return occupied MO coefficients (AO→MO) for RKS or UKS/UHF."""
    mo = mf.mo_coeff
    mo_occ = mf.mo_occ
    if isinstance(mo, np.ndarray) and mo.ndim == 2:
        occ_idx = np.asarray(mo_occ) > 0
        return mo[:, occ_idx]
    else:
        occ_idx = np.asarray(mo_occ[0]) > 0
        return mo[0][:, occ_idx]


def _safe_array(label: str, what: str, func):
    try:
        vals = func()
        return np.asarray(vals, dtype=float).tolist()
    except Exception as e:
        click.echo(f"{label} WARNING: Failed to compute {what}: {e}", err=True)
        return None


def _compute_atomic_charges(mol, mf) -> Dict[str, Optional[List[float]]]:
    """
    Compute atomic charges by three schemes:
      - 'mulliken' : scf.hf.mulliken_pop (failure → None)
      - 'lowdin'   : scf.hf.mulliken_pop_meta_lowdin_ao (failure → None)
      - 'iao'      : lo.iao.fast_iao_mullikan_pop (failure → None)
    """
    from pyscf.scf import hf as scf_hf
    from pyscf.lo import iao as lo_iao

    dm = mf.make_rdm1()
    S = mf.get_ovlp()
    # Total density (for charges)
    dm_tot = dm[0] + dm[1] if (isinstance(dm, np.ndarray) and dm.ndim == 3) else dm

    mull_q: Optional[List[float]] = _safe_array(
        "[Mulliken]",
        "charges",
        lambda: scf_hf.mulliken_pop(mol, dm_tot, s=S, verbose=0)[1],
    )
    low_q: Optional[List[float]] = _safe_array(
        "[Löwdin]",
        "charges",
        lambda: scf_hf.mulliken_pop_meta_lowdin_ao(mol, dm_tot, verbose=0, s=S)[1],
    )
    iao_q: Optional[List[float]] = _safe_array(
        "[IAO]",
        "charges",
        lambda: lo_iao.fast_iao_mullikan_pop(
            mol,
            dm,
            lo_iao.iao(mol, _get_occupied_orbitals(mf), minao="minao"),
            verbose=0,
        )[1],
    )

    return {
        "mulliken": mull_q,
        "lowdin": low_q,
        "iao": iao_q,
    }


def _compute_atomic_spin_densities(mol, mf) -> Dict[str, Optional[List[float]]]:
    """
    Compute atomic spin densities (Ms) by three schemes:
      - 'mulliken' : scf.uhf.mulliken_spin_pop (RKS → zeros; failure → None)
      - 'lowdin'   : scf.uhf.mulliken_spin_pop_meta_lowdin_ao (RKS → zeros; failure → None)
      - 'iao'      : fast_iao_mullikan_spin_pop (RKS → zeros; failure → None)
    """
    from pyscf.scf import uhf as scf_uhf
    from pyscf.lo import iao as lo_iao

    dm = mf.make_rdm1()
    S = mf.get_ovlp()
    nat = mol.natm

    # RKS (restricted) → spin densities are zero
    if not (isinstance(dm, np.ndarray) and dm.ndim == 3):
        zeros = [0.0] * nat
        return {"mulliken": zeros, "lowdin": zeros, "iao": zeros}

    mull: Optional[List[float]] = _safe_array(
        "[Spin Mulliken]",
        "spin densities",
        lambda: scf_uhf.mulliken_spin_pop(mol, dm, s=S, verbose=0)[1],
    )
    low: Optional[List[float]] = _safe_array(
        "[Spin Löwdin]",
        "spin densities",
        lambda: scf_uhf.mulliken_spin_pop_meta_lowdin_ao(mol, dm, verbose=0, s=S)[1],
    )
    iao_ms: Optional[List[float]] = _safe_array(
        "[Spin IAO]",
        "spin densities",
        lambda: fast_iao_mullikan_spin_pop(
            mol,
            dm,
            lo_iao.iao(mol, _get_occupied_orbitals(mf), minao="minao"),
            verbose=0,
        )[1],
    )

    return {"mulliken": mull, "lowdin": low, "iao": iao_ms}


# -----------------------------------------------
# CLI
# -----------------------------------------------

@click.command(
    help="Single-point DFT using GPU4PySCF (CPU PySCF backend).",
    context_settings={"help_option_names": ["-h", "--help"]},
)
@click.option(
    "-i", "--input",
    "input_path",
    type=click.Path(path_type=Path, exists=True, dir_okay=False),
    required=True,
    help="Input structure file (.pdb, .xyz, _trj.xyz, etc.; loaded via pysisyphus.helpers.geom_loader).",
)
@click.option(
    "-q",
    "--charge",
    type=int,
    required=False,
    help=(
        "Total charge. Required for non-.gjf inputs unless --ligand-charge is provided "
        "(PDB inputs or XYZ/GJF with --ref-pdb)."
    ),
)
@click.option(
    "-l",
    "--ligand-charge",
    type=str,
    default=None,
    show_default=False,
    help=(
        "Total charge or per-resname mapping (e.g., GPP:-3,SAM:1) used to derive charge "
        "when -q is omitted (requires PDB input or --ref-pdb)."
    ),
)
@click.option("-m", "--multiplicity", "spin", type=int, default=None, show_default=False, help="Spin multiplicity (2S+1; inherits from .gjf when available; otherwise defaults to 1).")
@click.option(
    "--convert-files/--no-convert-files",
    "convert_files",
    default=True,
    show_default=True,
    help="Accepted for interface consistency; dft does not emit PDB/GJF outputs.",
)
@click.option(
    "--ref-pdb",
    type=click.Path(path_type=Path, exists=True, dir_okay=False),
    default=None,
    help="Reference PDB topology to use when the input is XYZ/GJF (keeps XYZ coordinates).",
)
@click.option(
    "--func-basis",
    "func_basis",
    type=str,
    default=f"{DFT_DEFAULT_FUNC}/{DFT_DEFAULT_BASIS}",
    show_default=True,
    help="Exchange–correlation functional and basis set as 'FUNC/BASIS' (e.g., 'wb97m-v/6-31g**', 'wb97m-v/def2-tzvpd').",
)
@click.option("--max-cycle", type=int, default=DFT_KW["max_cycle"], show_default=True, help="Maximum SCF iterations.")
@click.option("--conv-tol", type=float, default=DFT_KW["conv_tol"], show_default=True, help="SCF convergence tolerance (Eh).")
@click.option("--grid-level", type=int, default=DFT_KW["grid_level"], show_default=True, help="Numerical integration grid level (PySCF grids.level).")
@click.option("-o", "--out-dir", type=str, default=DFT_KW["out_dir"], show_default=True, help="Output directory.")
@click.option(
    "--engine",
    type=click.Choice(["gpu", "cpu", "auto"], case_sensitive=False),
    default="gpu",
    show_default=True,
    help="Preferred SCF backend: GPU (strict), CPU, or auto (try GPU then CPU if unavailable).",
)
@click.option(
    "--config",
    "config_yaml",
    type=click.Path(path_type=Path, exists=True, dir_okay=False),
    default=None,
    help="Base YAML configuration file applied before explicit CLI options.",
)
@click.option(
    "--show-config/--no-show-config",
    "show_config",
    default=False,
    show_default=True,
    help="Print resolved configuration and continue execution.",
)
@click.option(
    "--dry-run/--no-dry-run",
    "dry_run",
    default=False,
    show_default=True,
    help="Validate options and print the execution plan without running DFT.",
)
@click.pass_context
def cli(
    ctx: click.Context,
    input_path: Path,
    charge: Optional[int],
    ligand_charge: Optional[str],
    spin: Optional[int],
    convert_files: bool,
    ref_pdb: Optional[Path],
    func_basis: str,
    max_cycle: int,
    conv_tol: float,
    grid_level: int,
    out_dir: str,
    engine: str,
    config_yaml: Optional[Path],
    show_config: bool,
    dry_run: bool,
) -> None:
    config_yaml, override_yaml, used_legacy_yaml = resolve_yaml_sources(
        config_yaml=config_yaml,
        override_yaml=None,
        args_yaml_legacy=None,
    )
    merged_yaml_cfg, _, _ = load_merged_yaml_cfg(
        config_yaml=config_yaml,
        override_yaml=None,
    )

    set_convert_file_enabled(convert_files)
    with prepared_cli_input(
        input_path,
        ref_pdb=ref_pdb,
        charge=charge,
        spin=spin,
        ligand_charge=ligand_charge,
        prefix="[dft]",
    ) as (prepared_input, resolved_charge, resolved_spin):
        geom_input_path = prepared_input.geom_path
        try:
            time_start = time.perf_counter()
            # --------------------------
            # 1) Assemble configuration (defaults < config < CLI(explicit) < override)
            # --------------------------
            geom_cfg = dict(GEOM_KW_DEFAULT)
            dft_cfg = dict(DFT_KW)

            apply_yaml_overrides(
                merged_yaml_cfg,
                [
                    (geom_cfg, (("geom",),)),
                    (dft_cfg, (("dft",),)),
                ],
            )

            if cli_param_overridden(ctx, "conv_tol"):
                dft_cfg["conv_tol"] = float(conv_tol)
            if cli_param_overridden(ctx, "max_cycle"):
                dft_cfg["max_cycle"] = int(max_cycle)
            if cli_param_overridden(ctx, "grid_level"):
                dft_cfg["grid_level"] = int(grid_level)
            if cli_param_overridden(ctx, "out_dir"):
                dft_cfg["out_dir"] = out_dir

            func_basis_value = str(
                dft_cfg.get(
                    "func_basis",
                    f"{DFT_DEFAULT_FUNC}/{DFT_DEFAULT_BASIS}",
                )
            )
            if cli_param_overridden(ctx, "func_basis"):
                func_basis_value = func_basis
            if func_basis_value:
                cfg_func, cfg_basis = _parse_func_basis(func_basis_value)
                dft_cfg["func"] = cfg_func
                dft_cfg["basis"] = cfg_basis

            xc = str(dft_cfg.get("func", "")).strip()
            basis = str(dft_cfg.get("basis", "")).strip()
            if not xc or not basis:
                raise click.BadParameter("Functional and basis must be non-empty (set via --func-basis or YAML dft.func/basis)")
            multiplicity = int(resolved_spin)
            if multiplicity < 1:
                raise click.BadParameter("Multiplicity (spin) must be >= 1.")
            spin2s = multiplicity - 1  # PySCF expects 2S

            # Echo resolved config
            engine_name = str(dft_cfg.get("engine", engine if engine else "gpu")).strip().lower()
            if cli_param_overridden(ctx, "engine"):
                engine_name = (engine or "gpu").strip().lower()
            out_dir_path = Path(dft_cfg["out_dir"]).resolve()
            echo_cfg = {
                "charge": int(resolved_charge),
                "multiplicity": multiplicity,
                "spin (PySCF expects 2S)": spin2s,
                "xc": xc,
                "basis": basis,
                "conv_tol": dft_cfg["conv_tol"],
                "max_cycle": dft_cfg["max_cycle"],
                "grid_level": dft_cfg["grid_level"],
                "out_dir": str(out_dir_path),
                "engine": engine_name,
            }
            click.echo(pretty_block("geom", format_geom_for_echo(geom_cfg)))
            click.echo(pretty_block("dft", echo_cfg))
            if show_config:
                click.echo(
                    pretty_block(
                        "yaml_layers",
                        {
                            "config": None if config_yaml is None else str(config_yaml),
                            "override_yaml": None if override_yaml is None else str(override_yaml),
                            "merged_keys": sorted(merged_yaml_cfg.keys()),
                        },
                    )
                )
            if dry_run:
                click.echo(
                    pretty_block(
                        "dry_run_plan",
                        {
                            "input_geometry": str(geom_input_path),
                            "output_dir": str(out_dir_path),
                            "engine": engine_name,
                            "convert_files": bool(convert_files),
                            "will_run_scf": True,
                            "will_write_result_yaml": True,
                            "will_run_population_analysis": True,
                        },
                    )
                )
                click.echo("[dry-run] Validation complete. DFT execution was skipped.")
                return

            # --------------------------
            # 2) Load geometry
            # --------------------------
            coord_type = geom_cfg.get("coord_type", GEOM_KW_DEFAULT["coord_type"])
            coord_kwargs = dict(geom_cfg)
            coord_kwargs.pop("coord_type", None)
            geometry = geom_loader(geom_input_path, coord_type=coord_type, **coord_kwargs)
            xyz_s, atoms_list = _geometry_to_pyscf_atoms_string(geometry)

            out_dir_path.mkdir(parents=True, exist_ok=True)
            # Write a provenance snapshot of the input geometry
            input_xyz = out_dir_path / "input_geometry.xyz"
            input_xyz.write_text(xyz_s if xyz_s.endswith("\n") else (xyz_s + "\n"))
            click.echo(f"[write] Wrote '{input_xyz}'.")

            # --------------------------
            # 3) Build PySCF molecule
            # --------------------------
            try:
                from pyscf import gto
            except (ModuleNotFoundError, ImportError) as e:
                click.echo(f"ERROR: PySCF import failed: {e}", err=True)
                sys.exit(2)

            mol = gto.Mole()
            mol.verbose = int(dft_cfg.get("verbose", 4))
            mol.build(
                atom=atoms_list,
                unit="Angstrom",
                charge=int(resolved_charge),
                spin=int(spin2s),
                basis=basis,
            )

            # --------------------------
            # 4) Activate GPU & build SCF object
            # --------------------------
            engine = engine_name
            using_gpu = False
            engine_label = "pyscf(cpu)"
            make_ks = (lambda mod: mod.RKS(mol) if spin2s == 0 else mod.UKS(mol))


            # --- Detect Blackwell GPU and emit warning ---
            is_blackwell_gpu = False
            try:
                import cupy as cp
                dev_id = cp.cuda.runtime.getDevice()
                props = cp.cuda.runtime.getDeviceProperties(dev_id)
                name = props["name"]
                if isinstance(name, bytes):
                    name = name.decode()
                if ("rtx 50" in name.lower()) or ("nvidia b" in name.lower()):
                    is_blackwell_gpu = True
            except Exception:
                is_blackwell_gpu = False

            if is_blackwell_gpu:
                click.echo("[gpu] WARNING: Detected a Blackwell GPU; GPU4PySCF may be unsupported.", err=True)
            # --------------------------------------------------


            if engine in ("gpu", "auto"):
                try:
                    from gpu4pyscf import dft as gdf
                    mf = make_ks(gdf)
                    using_gpu = True
                    engine_label = "gpu4pyscf"
                    mf = _configure_scf_object(mf, dft_cfg, xc)
                    e_tot = mf.kernel()

                except Exception as e:
                    if engine == "gpu":
                        raise click.ClickException(f"[gpu] GPU backend unavailable: {e}")
                    click.echo(
                        f"[gpu] WARNING: GPU backend unavailable ({e}); falling back to CPU.",
                    )
                    using_gpu = False
                    engine_label = "pyscf(cpu)"
                    engine = "cpu"

            if not using_gpu:
                from pyscf import dft as pdft
                mf = make_ks(pdft)
                mf = _configure_scf_object(mf, dft_cfg, xc)
                e_tot = mf.kernel()

            # --------------------------
            # 5) Run SCF
            # --------------------------


            converged = bool(getattr(mf, "converged", False))
            if e_tot is None:
                e_tot = float(getattr(mf, "e_tot", np.nan))

            e_h = float(e_tot)
            e_kcal = _AU2KCALPERMOL(e_h)

            # --------------------------
            # 6) Charges (Mulliken / meta-Löwdin-AO / IAO) & Spin densities
            # --------------------------
            try:
                mf_for_analysis = mf.to_cpu()  # GPU → CPU (no-op on CPU backend)
            except Exception:
                mf_for_analysis = mf

            charges = _compute_atomic_charges(mol, mf_for_analysis)
            spins   = _compute_atomic_spin_densities(mol, mf_for_analysis)

            def _round_list(xs, tol=1e-10):
                return [0.0 if (x == x) and abs(x) < tol else float(x) for x in xs]  # keep NaN as-is

            # Round tiny numbers (None-safe)
            for dct in (charges, spins):
                for key in ("mulliken", "lowdin", "iao"):
                    dct[key] = None if dct[key] is None else _round_list(dct[key])

            # Build per-atom tables
            charges_table: List[List[Any]] = []
            spins_table:   List[List[Any]] = []
            for i in range(mol.natm):
                elem = mol.atom_symbol(i)
                q_mull = None if charges["mulliken"] is None else charges["mulliken"][i]
                q_low  = None if charges["lowdin"]   is None else charges["lowdin"][i]
                q_iao  = None if charges["iao"]      is None else charges["iao"][i]
                charges_table.append([i, elem, q_mull, q_low, q_iao])

                s_mull = None if spins["mulliken"] is None else spins["mulliken"][i]
                s_low  = None if spins["lowdin"]   is None else spins["lowdin"][i]
                s_iao  = None if spins["iao"]      is None else spins["iao"][i]
                spins_table.append([i, elem, s_mull, s_low, s_iao])

            # ---- Echo charges/spins to stdout in flow style lines ----
            click.echo("")
            click.echo("charges [index, element, mulliken, lowdin, iao]:")
            for row in charges_table:
                click.echo(f"- {_format_row_for_echo(row)}")

            click.echo("")
            click.echo("spin_densities [index, element, mulliken, lowdin, iao]:")
            for row in spins_table:
                click.echo(f"- {_format_row_for_echo(row)}")

            # --------------------------
            # 7) Save result.yaml (flow style rows for readability)
            # --------------------------
            charges_rows_flow = [YamlFlowList(r) for r in charges_table]
            spins_rows_flow = [YamlFlowList(r) for r in spins_table]

            result_yaml = {
                "input": dict(echo_cfg),  # configuration snapshot
                "energy": {
                    "hartree": e_h,
                    "kcal_per_mol": e_kcal,
                    "converged": converged,
                    "engine": engine_label,
                    "used_gpu": bool(using_gpu),
                },
                # Table-style outputs (flow lists)
                "charges [index, element, mulliken, lowdin, iao]": charges_rows_flow,
                "spin_densities [index, element, mulliken, lowdin, iao]": spins_rows_flow,
            }
            (out_dir_path / "result.yaml").write_text(
                yaml.safe_dump(result_yaml, sort_keys=False, allow_unicode=True)
            )
            click.echo(f"[write] Wrote '{out_dir_path / 'result.yaml'}'.")
            # summary.md and key_* outputs are disabled.
            # --------------------------
            # 8) Final print: energies
            # --------------------------
            click.echo("")
            click.echo(f"E_total (Hartree): {e_h:.12f}")
            click.echo(f"E_total (kcal/mol): {e_kcal:.6f}")

            # Exit codes: 0 if converged, 3 otherwise
            if not converged:
                click.echo("WARNING: SCF did not converge to the requested tolerance.", err=True)
                sys.exit(3)

            click.echo(format_elapsed("[time] Elapsed Time for DFT", time_start))

        except KeyboardInterrupt:
            click.echo("Interrupted by user.", err=True)
            sys.exit(130)
        except click.ClickException:
            raise
        except Exception as e:
            tb = "".join(traceback.format_exception(type(e), e, e.__traceback__))
            click.echo("Unhandled error during DFT single-point:\n" + textwrap.indent(tb, "  "), err=True)
            sys.exit(1)
