#!/usr/bin/env python
r"""
Flow for v-ZSISA-QHA calculations
=================================

See [Phys. Rev. B 110, 014103](https://doi.org/10.1103/PhysRevB.110.014103)
"""

import os
import sys

import abipy.data as abidata
from abipy import abilab, flowtk
from abipy.flowtk.vzsisa import VzsisaFlow


def build_flow(options):
    """
    Create a `VzsisaFlow`.
    """
    # Working directory (default is the name of the script with '.py' removed and "run_" replaced by "flow_")
    if not options.workdir:
        __file__ = os.path.join(os.getcwd(), "run_qha_vzsisa.py")
        options.workdir = os.path.basename(__file__).replace(".py", "").replace("run_", "flow_")

    # Initialize structure from cif file.
    structure = abilab.Structure.from_file(abidata.cif_file("si.cif"))

    # Get NC pseudos from pseudodojo.
    from abipy.flowtk.psrepos import get_oncvpsp_pseudos

    pseudos = get_oncvpsp_pseudos(xc_name="PBEsol", version="0.4")

    # Select k-mesh for electrons and q-mesh for phonons. NB: ngqpt should be a divisor of ngkpt.
    # These values are clearly UNDERCONVERGED, just for testing.
    ngkpt = [2, 2, 2]
    ngqpt = [1, 1, 1]

    # BECS are not needed for Si.
    with_becs = True
    with_quad = True
    # with_quad = not structure.has_zero_dynamical_quadrupoles

    # List of volumetric scaling factors for the BO energies and the phonon part.
    # bo_vol_scales = [0.96, 0.98, 1.0, 1.02, 1.04, 1.06]
    # ph_vol_scales = [0.98, 1.0, 1.02, 1.04, 1.06] # EinfVib4(D)
    bo_vol_scales = [0.96, 0.98, 1, 1.02, 1.04]  # EinfVib4(S)
    ph_vol_scales = [1, 1.02, 1.04]  # EinfVib2(D)

    scf_input = abilab.AbinitInput(structure, pseudos)

    # Set other important variables assuming a spin unpolarized semiconductor (nsppol 1, nspinor 1).
    # All the other DFPT runs will inherit these parameters.
    scf_input.set_vars(
        nband=scf_input.num_valence_electrons // 2,
        nline=10,
        nbdbuf=0,
        nstep=100,
        ecutsm=1.0,
        tolvrs=1.0e-8,  # SCF stopping criterion (modify default)
        # tolvrs=1.0e-18, # SCF stopping criterion (modify default)
        paral_kgb=0,
    )

    scf_input.set_kmesh(ngkpt=ngkpt, shiftk=[0, 0, 0])
    print("scf_input\n", scf_input)

    return VzsisaFlow.from_scf_input(
        options.workdir, scf_input, bo_vol_scales, ph_vol_scales, ngqpt, with_becs, with_quad, edos_ngkpt=None
    )


# This block generates the thumbnails in the Abipy gallery.
# You can safely REMOVE this part if you are using this script for production runs.
if os.getenv("READTHEDOCS", False):
    __name__ = None
    import tempfile

    options = flowtk.build_flow_main_parser().parse_args(["-w", tempfile.mkdtemp()])
    build_flow(options).graphviz_imshow()


@flowtk.flow_main
def main(options):
    """
    This is our main function that will be invoked by the script.
    flow_main is a decorator implementing the command line interface.
    Command line args are stored in `options`.
    """
    return build_flow(options)


if __name__ == "__main__":
    sys.exit(main())
