#!/usr/bin/env python
"""
This example shows how to compute the SCR file by splitting the calculation of the SCR file
over q-points with the input variables nqptdm and qptdm.
"""
from __future__ import division, print_function, unicode_literals, absolute_import

import sys
import os
import abipy.abilab as abilab
import abipy.data as data

from abipy import flowtk


def all_inputs(paral_kgb=1):
    """
    Build the input files of the calculation.
    Returns: gs_input, nscf_input, scr_input, sigma_input
    """
    ecut = ecutwfn = 4
    global_vars = dict(
        ecut=ecut,
        istwfk="*1",
        paral_kgb=paral_kgb,
    )

    multi = abilab.MultiDataset(structure=data.cif_file("si.cif"),
                                pseudos=data.pseudos("14si.pspnc"), ndtset=4)
    multi.set_vars(global_vars)

    gs, nscf, scr, sigma = multi.split_datasets()

    # This grid is the most economical, but does not contain the Gamma point.
    gs_kmesh = dict(
        ngkpt=[2, 2, 2],
        shiftk=[0.5, 0.5, 0.5,
                0.5, 0.0, 0.0,
                0.0, 0.5, 0.0,
                0.0, 0.0, 0.5]
    )

    # This grid contains the Gamma point, which is the point at which
    # we will compute the (direct) band gap.
    gw_kmesh = dict(
        ngkpt=[2, 2, 2],
        shiftk=[0.0, 0.0, 0.0,
                0.0, 0.5, 0.5,
                0.5, 0.0, 0.5,
                0.5, 0.5, 0.0]
    )

    # Dataset 1 (GS run)
    gs.set_kmesh(**gs_kmesh)
    gs.set_vars(tolvrs=1e-6, nband=4)

    # Dataset 2 (NSCF run)
    # Here we select the second dataset directly with the syntax inp[2]
    nscf.set_kmesh(**gw_kmesh)
    nscf.set_vars(iscf=-2,
                  tolwfr=1e-10,
                  nband=15,
                  nbdbuf=5)

    # Dataset3: Calculation of the screening.
    scr.set_kmesh(**gw_kmesh)

    scr.set_vars(
        optdriver=3,
        nband=6,
        ecutwfn=ecutwfn,
        symchi=1,
        inclvkb=0,
        ecuteps=2.0)

    # Dataset4: Calculation of the Self-Energy matrix elements (GW corrections)
    sigma.set_kmesh(**gw_kmesh)

    sigma.set_vars(
            optdriver=4,
            nband=8,
            ecutwfn=ecutwfn,
            ecuteps=2.0,
            ecutsigx=2.0,
            symsigma=1,
            )

    kptgw = [
         -2.50000000E-01, -2.50000000E-01,  0.00000000E+00,
         -2.50000000E-01,  2.50000000E-01,  0.00000000E+00,
          5.00000000E-01,  5.00000000E-01,  0.00000000E+00,
         -2.50000000E-01,  5.00000000E-01,  2.50000000E-01,
          5.00000000E-01,  0.00000000E+00,  0.00000000E+00,
          0.00000000E+00,  0.00000000E+00,  0.00000000E+00,
      ]

    bdgw = [1, 8]

    sigma.set_kptgw(kptgw, bdgw)

    return multi.split_datasets()


def build_flow(options):
    """
    Construct the flow for G0W0 calculations.
    The calculation of the SCR file is split into nqptdm independent calculations.
    Partial SCR files are then merged with mrgddb before starting the sigma run.
    """
    # Working directory (default is the name of the script with '.py' removed and "run_" replaced by "flow_")
    workdir = options.workdir
    if not options.workdir:
        workdir = os.path.basename(__file__).replace(".py", "").replace("run_", "flow_")

    # Build the input files for GS, NSCF, SCR and SIGMA runs.
    gs, nscf, scr_input, sigma_input = all_inputs()

    # Construct the flow.
    return flowtk.G0W0WithQptdmFlow(workdir, gs, nscf, scr_input, sigma_input,
                                    manager=options.manager)


@abilab.flow_main
def main(options):
    flow = build_flow(options)
    flow.build_and_pickle_dump()
    return flow


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