#!/usr/bin/env python
"""
This script provides a simplified interface to the AbiPy factory functions.
For a more flexible interface, please use the AbiPy objects
to generate input files and workflows.
"""
from __future__ import annotations

import sys
import os
import argparse
import abipy.tools.cli_parsers as cli

from monty.termcolor import cprint
from monty.functools import prof_main
from abipy import abilab
from abipy.abio import factories
from abipy.abio.inputs import AnaddbInput
from abipy.dfpt.ddb import DdbFile


def get_structure(options):
    """Return structure object either from file or from the material project database."""
    if os.path.exists(options.filepath):
        return abilab.Structure.from_file(options.filepath)

    elif options.filepath.startswith("mp-"):
        return abilab.Structure.from_mpid(options.filepath, final=True,
                                          api_key=options.mapi_key, endpoint=options.endpoint)

    raise TypeError("Don't know how to extract structure object from %s" % options.filepath)


def get_pseudotable(options):
    """Return PseudoTable object."""
    if options.pseudos is not None:
        from abipy.flowtk import PseudoTable
        return PseudoTable.as_table(options.pseudos)

    try:
        from pseudo_dojo import OfficialTables

        dojo_tables = OfficialTables()
        if options.usepaw:
            raise NotImplementedError("PAW table is missing")
            #pseudos = dojo_tables["ONCVPSP-PBE-PDv0.2-accuracy"]
        else:
            pseudos = dojo_tables["ONCVPSP-PBE-PDv0.2-accuracy"]

        print("Using pseudos from PseudoDojo table", repr(pseudos))

    except ImportError as exc:
        from abipy.data.hgh_pseudos import HGH_TABLE
        pseudos = HGH_TABLE
        print("PseudoDojo package not installed. Please install it with `pip install pseudo_dojo`")
        print("or use `--pseudos FILE_LIST` to specify the pseudopotentials to use.")
        print("Using internal HGH_TABLE!!!!")
        #raise exc

    return pseudos


def finalize(obj, options):
    if options.mnemonics: obj.set_mnemonics(True)
    print(obj)
    print("\n")
    print("# This input file template has been automatically generated by the abinp.py script.")
    print("# Several input parameters have default values that might not be suitable for you particular calculation.")
    print("# Please check the input file and modify the template according to your needs.")
    return 0


def build_abinit_input_from_file(options, **abivars):
    """
    Build and return an AbinitInput instance from filepath.

    abivars are optional variables that will be added to the input.
    """
    from abipy.abio.abivars import AbinitInputFile
    abifile = AbinitInputFile(options.filepath)
    pseudos = get_pseudotable(options)
    jdtset = options.jdtset
    # Get vars from input
    abi_kwargs = abifile.datasets[jdtset - 1].get_vars()
    if abifile.ndtset != 1:
        cprint("# Input file contains %s datasets, will select jdtset index %s:" %
               (abifile.ndtset, jdtset), "yellow")
        abi_kwargs["jdtset"] = jdtset

    # Add input abivars (if any).
    abi_kwargs.update(abivars)

    return abilab.AbinitInput(abifile.structure, pseudos, pseudo_dir=None, comment=None, decorators=None, abi_args=None,
                              abi_kwargs=abi_kwargs, tags=None)


def abinp_validate(options):
    """Validate Abinit input file."""
    inp = build_abinit_input_from_file(options)
    r = inp.abivalidate()
    if r.retcode == 0:
        print("Validation completed succesfully.")
    else:
        print(r.log_file, r.stderr_file)

    return r.retcode


def abinp_autoparal(options):
    """Compute autoparal configurations."""
    inp = build_abinit_input_from_file(options)
    pconfs = inp.abiget_autoparal_pconfs(options.max_ncpus, autoparal=1, workdir=None, manager=None, verbose=options.verbose)
    print(pconfs)
    return 0


def abinp_abispg(options):
    """Call Abinit with chkprim = 0 to find space group."""
    inp = build_abinit_input_from_file(options, chkprim=0, mem_test=0)
    r = inp.abivalidate()
    if r.retcode != 0:
        print(r.log_file, r.stderr_file)
        return r.retcode

    try:
        out = abilab.abiopen(r.output_file.path)
    except Exception as exc:
        print("Error while trying to parse output file:", r.output_file.path)
        print("Exception:\n", exc)
        return 1

    #print(out)
    structure = out.initial_structure

    # Call spglib to get spacegroup if Abinit spacegroup is not available.
    # Return string with full information about crystalline structure i.e.
    # space group, point group, wyckoff positions, equivalent sites.
    print(structure.spget_summary(verbose=options.verbose))
    if options.verbose:
        print(structure.abi_spacegroup.to_string(verbose=options.verbose))

    return 0


def abinp_ibz(options):
    """Get k-points in the irreducible weights."""
    inp = build_abinit_input_from_file(options)
    ibz = inp.abiget_ibz(ngkpt=None, shiftk=None, kptopt=None, workdir=None, manager=None)

    nkibz = len(ibz.points)
    print("kptopt 0 nkpt ", nkibz)
    print("kpts")
    for i, (k, w) in enumerate(zip(ibz.points, ibz.weights)):
        print("%12.8f  %12.8f  %12.8f  # index: %d, weigth: %10.8f" % (k[0], k[1], k[2], i + 1, w))

    print("\nwtk")
    for i, w in enumerate(ibz.weights):
        print("%10.8f  # index: %d" % (w, i + 1))

    return 0


def abinp_phperts(options):
    """Get list of phonon perturbations."""
    inp = build_abinit_input_from_file(options)
    qpt = None if "qpt" in inp else [0, 0, 0]
    perts = inp.abiget_irred_phperts(qpt=qpt)
    print(perts)

    return 0


def abinp_gs(options):
    """Build Abinit input for ground-state calculation."""
    structure = abilab.Structure.from_file(options.filepath)
    pseudos = get_pseudotable(options)
    gsinp = factories.gs_input(structure, pseudos,
                               kppa=options.kppa, ecut=None, pawecutdg=None, scf_nband=None,
                               accuracy="normal", spin_mode=options.spin_mode,
                               smearing=options.smearing, charge=0.0, scf_algorithm=None)

    return finalize(gsinp, options)


def abinp_ebands(options):
    """Build Abinit input for band structure calculations."""
    structure = get_structure(options)
    pseudos = get_pseudotable(options)
    multi = factories.ebands_input(structure, pseudos,
                 kppa=options.kppa, nscf_nband=None, ndivsm=15,
                 ecut=None, pawecutdg=None, scf_nband=None, accuracy="normal", spin_mode=options.spin_mode,
                 smearing=options.smearing, charge=0.0, scf_algorithm=None, dos_kppa=None)

    # Add getwfk variables.
    for inp in multi[1:]:
        inp["getwfk"] = 1

    return finalize(multi, options)


def abinp_phonons(options):
    """Build Abinit input for phonon calculations."""
    structure = get_structure(options)
    pseudos = get_pseudotable(options)

    gs_inp = factories.gs_input(structure, pseudos,
                               kppa=options.kppa, ecut=None, pawecutdg=None, scf_nband=None,
                               accuracy="normal", spin_mode=options.spin_mode,
                               smearing=options.smearing, charge=0.0, scf_algorithm=None)

    multi = factories.phonons_from_gsinput(gs_inp, ph_ngqpt=None, qpoints=None, with_ddk=True, with_dde=True, with_bec=False,
                                           ph_tol=None, ddk_tol=None, dde_tol=None, wfq_tol=None, qpoints_to_skip=None)

    # Add getwfk variables.
    for inp in multi[1:]:
        inp["getwfk"] = 1

    return finalize(multi, options)


def abinp_g0w0(options):
    """Generate input files for G0W0 calculations."""
    structure = get_structure(options)
    pseudos = get_pseudotable(options)
    nscf_nband, ecuteps, ecutsigx = 100, 4, 12

    multi = factories.g0w0_with_ppmodel_inputs(structure, pseudos,
                             options.kppa, nscf_nband, ecuteps, ecutsigx,
                             ecut=None, pawecutdg=None,
                             accuracy="normal", spin_mode=options.spin_mode, smearing=options.smearing,
                             ppmodel="godby", charge=0.0, scf_algorithm=None, inclvkb=2, scr_nband=None,
                             sigma_nband=None, gw_qprange=1)

    # Add getwfk and getscr variables.
    for inp in multi[1:]:
        inp["getwfk"] = -1
    multi[-1]["getscr"] = -1

    return finalize(multi, options)


def abinp_anaph(options):
    """Build Anaddb input file for the computation of phonon bands DOS."""
    ddb = DdbFile(options.filepath)
    nqsmall = 10
    inp = AnaddbInput.phbands_and_dos(ddb.structure, ddb.guessed_ngqpt, nqsmall, ndivsm=20, q1shft=(0, 0, 0),
        qptbounds=None, asr=2, chneut=0, dipdip=1, dos_method="tetra", lo_to_splitting=False,
        anaddb_args=None, anaddb_kwargs=None)

    return finalize(inp, options)


def abinp_wannier90(options):
    """
    Build wannier90 template input file from Abinit input/output file.
    possibly with electron bands
    """
    from abipy.wannier90.win import Wannier90Input
    inp = Wannier90Input.from_abinit_file(options.filepath)
    return finalize(inp, options)


def abinp_lobster(options):
    """
    Build and print Lobster input file from directory with files file (for pseudos)
    and GSR.nc output file (multidatasets are not allowed).
    Supports also directory with vasprun.xml and POTCAR file.
    """
    lobinp = abilab.LobsterInput.from_dir(os.path.dirname(options.dirpath))
    print(lobinp)


def get_epilog():
    return r"""
Usage example:

######################
# Require Abinit Input
######################

    abinp.py validate run.abi       # Call abinit to validate run.abi input file
    abinp.py abispg run.abi         # Call abinit to get space group information.
    abinp.py autoparal run.abi      # Call abinit to get list of autoparal configurations.
    abinp.py ibz run.abi            # Call abinit to get list of k-points in the IBZ.
    abinp.py phperts run.abi        # Call abinit to get list of atomic perturbations for phonons.

########################
# Abinit Input Factories
########################

    abinp.py gs si.cif > run.abi    # Build input for GS run for structure read from CIF file.
                                    # Redirect output to run.abi.
    abinp.py ebands out_GSR.nc      # Build input for SCF + NSCF run with structure read from GSR.nc file.
    abinp.py ebands mp-149          # Build input for SCF+NSCF run with (relaxed) structure taken from the
                                    # materials project database. Requires internet connection and MAPI_KEY.
    abinp.py phonons POSCAR         # Build input for GS + DFPT calculation of phonons with DFPT.
    abinp.py phonons out_HIST.nc    # Build input for phonons run with (relaxed) structure read from HIST.nc file.
                                    # Use e.g. --kppa=100 --spin-mode=polarized --smearing="gaussian: 0.3 eV"
                                    # to specify the k-points sampling, the treatment of spin and the occupation scheme.

########################
# Anaddb Input Factories
########################

    abinp.py anaph out_DDB          # Build anaddb input file for phonon bands + DOS from DDB file.

#############
# Other Codes
#############

    abinp.py wannier90 FILE         # Build and print wannier90 input file from FILE with structure.
    abinp.py lobster .              # Build and print lobster input file from directory.


Note that one can use pass any file providing a pymatgen structure
e.g. Abinit netcdf files, CIF files, POSCAR, ...
Use `abinp.py --help` for help and `abinp.py COMMAND --help` to get the documentation for `COMMAND`.
Use `-v` to increase verbosity level (can be supplied multiple times e.g -vv).

CAVEAT: This script provides a simplified interface to the AbiPy factory functions.
For a more flexible interface, please use the AbiPy objects to generate input files and workflows.
"""


def get_parser(with_epilog=False):
    """Build parser."""
    # Parent parser for common options.
    copts_parser = argparse.ArgumentParser(add_help=False)
    copts_parser.add_argument('-v', '--verbose', default=0, action='count', # -vv --> verbose=2
        help='verbose, can be supplied multiple times to increase verbosity')
    copts_parser.add_argument('--loglevel', default="ERROR", type=str,
        help="Set the loglevel. Possible values: CRITICAL, ERROR (default), WARNING, INFO, DEBUG")
    copts_parser.add_argument("--mapi-key", default=None,
        help="Pymatgen MAPI_KEY used if mp identifier is used to select structure.\n"
             "Use value in .pmgrc.yaml if not specified.")
    copts_parser.add_argument("--endpoint", help="Pymatgen database.", default="https://www.materialsproject.org/rest/v2")

    copts_parser.add_argument("-m", '--mnemonics', default=False, action="store_true",
        help="Print brief description of input variables in the input file.")
    copts_parser.add_argument('--usepaw', default=False, action="store_true",
        help="Use PAW pseudos instead of norm-conserving.")

    # Parent parser for command options operating on Abinit input files.
    abiinput_parser = argparse.ArgumentParser(add_help=False)
    abiinput_parser.add_argument('--jdtset', default=1, type=int,
        help="jdtset index. Used to select the dataset index when the input file " +
             "contains more than one dataset.")
    abiinput_parser.add_argument("-p", '--pseudos', nargs="+", default=None, help="List of pseudopotentials")

    abiinput_parser.add_argument('--kppa', type=int, default=None,
        help="Number of k-points per reciprocal atom. Use default value (1000) if not specified.")
    abiinput_parser.add_argument('--spin-mode', type=str, default="unpolarized",
        help="Spin polarization. Default: unpolarized. Allowed values in: " +
             "[polarized, unpolarized, afm (anti-ferromagnetic), spinor (non-collinear magnetism) " +
             " spinor_nomag (non-collinear, no magnetism)].")
    abiinput_parser.add_argument('--smearing', type=str, default="fermi_dirac:0.1 eV",
        help="Smearing scheme. Defaults to Fermi-Dirac.")

    # Parent parser for commands requiring a file.
    path_selector = argparse.ArgumentParser(add_help=False)
    path_selector.add_argument('filepath', type=str,
        help="File with the crystalline structure (netcdf, cif, POSCAR, input files ...)")

    # Parent parser for commands requiring a path to a directory.
    dir_selector = argparse.ArgumentParser(add_help=False)
    dir_selector.add_argument('dirpath', type=str,
        help="Directory with input files required by COMMAND.")

    parser = argparse.ArgumentParser(epilog=get_epilog() if with_epilog else "",
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--loglevel', default="ERROR", type=str,
        help="Set the loglevel. Possible values: CRITICAL, ERROR (default), WARNING, INFO, DEBUG")
    parser.add_argument('-V', '--version', action='version', version=abilab.__version__)
    parser.add_argument('-v', '--verbose', default=0, action='count', # -vv --> verbose=2
        help='verbose, can be supplied multiple times to increase verbosity')

    # Create the parsers for the sub-commands
    subparsers = parser.add_subparsers(dest='command', help='sub-command help',
                                       description="Valid subcommands, use command --help for help")

    abifile_parsers = [copts_parser, path_selector, abiinput_parser]

    # Subparser for validate command.
    p_validate = subparsers.add_parser('validate', parents=abifile_parsers, help=abinp_validate.__doc__)

    # Subparser for autoparal command.
    p_autoparal = subparsers.add_parser('autoparal', parents=abifile_parsers, help=abinp_autoparal.__doc__)
    p_autoparal.add_argument("-n", '--max-ncpus', default=50, type=int, help="Maximum number of CPUs")

    # Subparser for abispg command.
    p_abispg = subparsers.add_parser('abispg', parents=abifile_parsers, help=abinp_abispg.__doc__)

    # Subparser for ibz command.
    p_ibz = subparsers.add_parser('ibz', parents=abifile_parsers, help=abinp_ibz.__doc__)

    # Subparser for phperts command.
    p_phperts = subparsers.add_parser('phperts', parents=abifile_parsers, help=abinp_phperts.__doc__)

    inpgen_parsers = [copts_parser, path_selector, abiinput_parser]

    # Subparser for gs command.
    p_gs = subparsers.add_parser('gs', parents=inpgen_parsers, help=abinp_gs.__doc__)

    # Subparser for ebands command.
    p_ebands = subparsers.add_parser('ebands', parents=inpgen_parsers, help=abinp_ebands.__doc__)

    # Subparser for phonons command.
    p_phonons = subparsers.add_parser('phonons', parents=inpgen_parsers, help=abinp_phonons.__doc__)

    # Subparser for g0w0 command.
    p_g0w0 = subparsers.add_parser('g0w0', parents=inpgen_parsers, help=abinp_g0w0.__doc__)

    # Subparser for anaph command.
    p_anaph = subparsers.add_parser('anaph', parents=inpgen_parsers, help=abinp_anaph.__doc__)

    # Subparser for wannier90 command.
    p_wannier90 = subparsers.add_parser('wannier90', parents=[path_selector], help=abinp_wannier90.__doc__)

    # Subparser for lobster command.
    p_lobster = subparsers.add_parser('lobster', parents=[dir_selector], help=abinp_lobster.__doc__)

    return parser


@prof_main
def main():

    def show_examples_and_exit(err_msg=None, error_code=1):
        """Display the usage of the script."""
        sys.stderr.write(get_epilog())
        if err_msg:
            sys.stderr.write("Fatal Error\n" + err_msg + "\n")
        sys.exit(error_code)

    parser = get_parser(with_epilog=True)

    # Parse the command line.
    try:
        options = parser.parse_args()
    except Exception:
        show_examples_and_exit(error_code=1)

    if not options.command:
        show_examples_and_exit(error_code=1)

    cli.set_loglevel(options.loglevel)

    # Dispatch
    return globals()["abinp_" + options.command](options)


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