#!/usr/bin/env python
"""
Command-line interface to AbiPy's input factory functions.

This script provides a simplified way to generate input files for Abinit and other
supported codes (VASP, Wannier90, Lobster) from various structure file formats.
It also includes tools for validating Abinit inputs, finding space groups,
analyzing k-point grids, and more.

Examples:
    Generate a ground-state Abinit input from a CIF file:
        $ abinp.py gs crystal.cif > run.abi

    Generate Abinit inputs for an electron band structure using a structure from MP:
        $ abinp.py ebands mp-149

    Validate an existing Abinit input file:
        $ abinp.py validate run.abi

    Generate VASP input files (INCAR, KPOINTS, etc.) from a POSCAR:
        $ abinp.py vasp POSCAR

Note: For more complex workflows or fine-grained control, it is recommended
to use the AbiPy Python API directly.
"""

from __future__ import annotations

import argparse
import os
import sys

from monty.termcolor import cprint
from pymatgen.io.vasp.sets import VaspInputSet

import abipy.tools.cli_parsers as cli
from abipy import abilab
from abipy.abio import factories
from abipy.abio.inputs import AnaddbInput
from abipy.dfpt.ddb import DdbFile


def vasp_dict_set_cls(s: str | VaspInputSet) -> type | list[str]:
    """
    Return a subclass of DictSect from string `s`.
    If s == "__all__", return list with all VaspInputSet subclasses supported by pymatgen.
    """
    from inspect import isclass

    from pymatgen.io.vasp import sets

    def is_dict_set(key: str) -> bool:
        return isclass(obj := getattr(sets, key)) and issubclass(obj, VaspInputSet)

    valid_keys = [key for key in dir(sets) if is_dict_set(key)]

    if s == "__all__":
        return valid_keys

    if isinstance(s, VaspInputSet):
        return s

    if s not in valid_keys:
        raise ValueError(f"Unknown VaspInputSet {s}, must be one of {valid_keys}")

    return getattr(sets, s)


ALL_VASP_DICT_SETS = vasp_dict_set_cls("__all__")


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)

    if options.filepath.startswith("mp-"):
        return abilab.Structure.from_mpid(options.filepath)

    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"]
        pseudos = dojo_tables["ONCVPSP-PBE-PDv0.2-accuracy"]

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

    except ImportError:
        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!!!!")

    return pseudos


def finalize(obj, options):
    """
    Finalize the object and print the input file to stdout.

    Args:
        obj: The object to finalize (e.g. AbinitInput or MultiDataset).
        options: Namespace object containing command-line 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), color="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 an Abinit input file by calling Abinit in dry-run mode.

    Args:
        options: Namespace object containing command-line options.

    Returns:
        int: Abinit exit code (0 if valid).
    """
    inp = build_abinit_input_from_file(options)
    r = inp.abivalidate()
    if r.retcode == 0:
        print("Validation completed successfully.")
    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, strict=False)):
        print("%12.8f  %12.8f  %12.8f  # index: %d, weight: %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):
    """
    Generate an Abinit input for a ground-state (GS) calculation.

    Args:
        options: Namespace object containing command-line options.

    Returns:
        int: System exit code.
    """
    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_vasp(options):
    """Build VASP input files from a FILE defining the structure."""
    structure = abilab.Structure.from_file(options.filepath)
    cls = vasp_dict_set_cls(options.dict_set)
    cprint(f"Generating VASP input using {cls}. Use -d option to change settings.", color="yellow")
    cls(structure).write_input(".")
    return 0


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 abinp_slurm(options):
    """Print template for Slurm submmission script."""
    from abipy.flowtk.qutils import get_slurm_template

    body = "srun abinit run.abi > run.log 2> run.err"
    print(get_slurm_template(body))


def get_epilog():
    """
    Return the epilog string for the command-line parser.
    """
    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 vasp FILE              # Build and write Vasp input files starting from a FILE with structure.
    abinp.py wannier90 FILE         # Build and print wannier90 input file starting from a FILE with structure.
    abinp.py lobster .              # Build and print lobster input file from directory.
    abinp.py slurm                  # Print template for Slurm submission script


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(
        "-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 vasp command.
    p_vasp = subparsers.add_parser("vasp", parents=[path_selector], help=abinp_vasp.__doc__)
    p_vasp.add_argument(
        "--dict-set",
        default="MPStaticSet",
        type=str,
        help="VaspDictSet. Default: MPStaticSet. For further info see pymatgen.io.vasp.sets",
        choices=ALL_VASP_DICT_SETS,
    )

    # 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__)

    p_slurm = subparsers.add_parser("slurm", help=abinp_slurm.__doc__)

    return parser


@cli.prof_main
def main():
    """
    Main entry point for the script.
    """

    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())
