# Source code for pymatgen.cli.pmg

    #!/usr/bin/env python
    # coding: utf-8
    # Copyright (c) Pymatgen Development Team.
    # Distributed under the terms of the MIT License.

    """
    A master convenience script with many tools for vasp and structure analysis.
    """

    import argparse
    import sys
    import itertools

    from tabulate import tabulate, tabulate_formats
    from pymatgen import Structure, SETTINGS
    from pymatgen.io.vasp import Incar, Potcar

    from pymatgen.cli.pmg_analyze import analyze
    from pymatgen.cli.pmg_config import configure_pmg
    from pymatgen.cli.pmg_potcar import generate_potcar
    from pymatgen.cli.pmg_plot import plot
    from pymatgen.cli.pmg_structure import analyze_structures
    from pymatgen.cli.pmg_query import do_query
    from pymatgen import __version__

    def parse_view(args):
        """
        Handle view commands.

        :param args: Args from command.
        """
        from pymatgen.vis.structure_vtk import StructureVis
        excluded_bonding_elements = args.exclude_bonding[0].split(",") \
            if args.exclude_bonding else []
        s = Structure.from_file(args.filename[0])
        vis = StructureVis(excluded_bonding_elements=excluded_bonding_elements)
        vis.set_structure(s)
        vis.show()
        return 0

    def diff_incar(args):
        """
        Handle diff commands.

        :param args: Args from command.
        """
        filepath1 = args.incars[0]
        filepath2 = args.incars[1]
        incar1 = Incar.from_file(filepath1)
        incar2 = Incar.from_file(filepath2)

        def format_lists(v):
            if isinstance(v, (tuple, list)):
                return " ".join(["%d*%.2f" % (len(tuple(group)), i)
                                 for (i, group) in itertools.groupby(v)])
            return v

        d = incar1.diff(incar2)
        output = [['SAME PARAMS', '', ''], ['---------------', '', ''],
                  ['', '', ''], ['DIFFERENT PARAMS', '', ''],
                  ['----------------', '', '']]
        output.extend([(k, format_lists(d['Same'][k]), format_lists(d['Same'][k]))
                       for k in sorted(d['Same'].keys()) if k != "SYSTEM"])
        output.extend([(k, format_lists(d['Different'][k]['INCAR1']),
                        format_lists(d['Different'][k]['INCAR2']))
                       for k in sorted(d['Different'].keys()) if k != "SYSTEM"])
        print(tabulate(output, headers=['', filepath1, filepath2]))
        return 0

    def main():
        """
        Handle main.
        """
        parser = argparse.ArgumentParser(
            description="""
        pmg is a convenient script that uses pymatgen to perform many
        analyses, plotting and format conversions. This script works based on
        several sub-commands with their own options. To see the options for the
        sub-commands, type "pmg sub-command -h".""",
            epilog="""Version: {}""".format(__version__)
        )

        subparsers = parser.add_subparsers()

        parser_config = subparsers.add_parser(
            "config", help="Tools for configuring pymatgen, e.g., "
                           "potcar setup, modifying .pmgrc.yaml "
                           "configuration file.")
        groups = parser_config.add_mutually_exclusive_group(required=True)
        groups.add_argument("-p", "--potcar", dest="potcar_dirs",
                            metavar="dir_name",
                            nargs=2,
                            help="Initial directory where downloaded VASP "
                                 "POTCARs are extracted to, and the "
                                 "output directory where the reorganized "
                                 "potcars will be stored. The input "
                                 "directory should be "
                                 "the parent directory that contains the "
                                 "POT_GGA_PAW_PBE or potpaw_PBE type "
                                 "subdirectories.")
        groups.add_argument("-i", "--install", dest="install",
                            metavar="package_name",
                            choices=["enumlib", "bader"],
                            help="Install various optional command line "
                                 "tools needed for full functionality.")

        groups.add_argument("-a", "--add", dest="var_spec", nargs="+",
                            help="Variables to add in the form of space "
                                 "separated key value pairs. E.g., "
                                 "PMG_VASP_PSP_DIR ~/psps")
        parser_config.set_defaults(func=configure_pmg)

        parser_analyze = subparsers.add_parser(
            "analyze", help="Vasp calculation analysis tools.")
        parser_analyze.add_argument("directories", metavar="dir", default=".",
                                    type=str, nargs="*",
                                    help="directory to process (default to .)")
        parser_analyze.add_argument("-e", "--energies", dest="get_energies",
                                    action="store_true", help="Print energies")
        parser_analyze.add_argument(
            "-m", "--mag", dest="ion_list", type=str, nargs=1,
            help="Print magmoms. ION LIST can be a range "
                 "(e.g., 1-2) or the string 'All' for all ions.")
        parser_analyze.add_argument(
            "-r", "--reanalyze", dest="reanalyze", action="store_true",
            help="Force reanalysis. Typically, vasp_analyzer"
                 " will just reuse a vasp_analyzer_data.gz if "
                 "present. This forces the analyzer to reanalyze "
                 "the data.")
        parser_analyze.add_argument(
            "-f", "--format", dest="format", choices=tabulate_formats,
            default="simple",
            help="Format for table. Supports all options in tabulate package.")
        parser_analyze.add_argument(
            "-v", "--verbose", dest="verbose", action="store_true",
            help="Verbose mode. Provides detailed output on progress.")
        parser_analyze.add_argument(
            "-q", "--quick", dest="quick", action="store_true",
            help="Faster mode, but less detailed information. Parses individual vasp files.")
        parser_analyze.add_argument(
            "-s", "--sort", dest="sort", choices=["energy_per_atom", "filename"],
            default="energy_per_atom",
            help="Sort criteria. Defaults to energy / atom.")
        parser_analyze.set_defaults(func=analyze)

        parser_query = subparsers.add_parser(
            "query",
            help="Search for structures and data from the Materials Project.")
        parser_query.add_argument(
            "criteria", metavar="criteria",
            help="Search criteria. Supported formats in formulas, chemical "
                 "systems, Materials Project ids, etc.")
        group = parser_query.add_mutually_exclusive_group(required=True)
        group.add_argument(
            "-s", "--structure", dest="structure", metavar="format",
            choices=["poscar", "cif", "cssr"], type=str.lower,
            help="Get structures from Materials Project and write them to a "
                 "specified format.")
        group.add_argument(
            "-e", "--entries", dest="entries", metavar="filename",
            help="Get entries from Materials Project and write them to "
                 "serialization file. JSON and YAML supported.")
        group.add_argument(
            "-d", "--data", dest="data", metavar="fields", nargs="*",
            help="Print a summary of entries in the Materials Project satisfying "
                 "the criteria. Supply field names to include additional data. "
                 "By default, the Materials Project id, formula, spacegroup, "
                 "energy per atom, energy above hull are shown.")
        parser_query.set_defaults(func=do_query)

        parser_plot = subparsers.add_parser("plot", help="Plotting tool for "
                                                         "DOS, CHGCAR, XRD, etc.")
        group = parser_plot.add_mutually_exclusive_group(required=True)
        group.add_argument('-d', '--dos', dest="dos_file", metavar="vasprun.xml",
                           help="Plot DOS from a vasprun.xml")
        group.add_argument('-c', '--chgint', dest="chgcar_file", metavar="CHGCAR",
                           help="Generate charge integration plots from any "
                                "CHGCAR")
        group.add_argument('-x', '--xrd', dest="xrd_structure_file",
                           metavar="structure_file",
                           help="Generate XRD plots from any supported structure "
                                "file, e.g., CIF, POSCAR, vasprun.xml, etc.")

        parser_plot.add_argument("-s", "--site", dest="site", action="store_const",
                                 const=True, help="Plot site projected DOS")
        parser_plot.add_argument("-e", "--element", dest="element", type=str,
                                 nargs=1,
                                 help="List of elements to plot as comma-separated"
                                      " values e.g., Fe,Mn")
        parser_plot.add_argument("-o", "--orbital", dest="orbital",
                                 action="store_const", const=True,
                                 help="Plot orbital projected DOS")

        parser_plot.add_argument("-i", "--indices", dest="inds", type=str,
                                 nargs=1,
                                 help="Comma-separated list of indices to plot "
                                      "charge integration, e.g., 1,2,3,4. If not "
                                      "provided, the code will plot the chgint "
                                      "for all symmetrically distinct atoms "
                                      "detected.")
        parser_plot.add_argument("-r", "--radius", dest="radius", type=float,
                                 default=3,
                                 help="Radius of integration for charge "
                                      "integration plot.")
        parser_plot.add_argument("--out_file", dest="out_file", type=str,
                                 help="Save plot to file instead of displaying.")
        parser_plot.set_defaults(func=plot)

        parser_structure = subparsers.add_parser(
            "structure",
            help="Structure conversion and analysis tools.")

        parser_structure.add_argument(
            "-f", "--filenames", dest="filenames",
            metavar="filename", nargs="+",
            help="List of structure files.")

        groups = parser_structure.add_mutually_exclusive_group(required=True)
        groups.add_argument("-c", "--convert", dest="convert", action="store_true",
                            help="Convert from structure file 1 to structure "
                                 "file 2. Format determined from filename. "
                                 "Supported formats include POSCAR/CONTCAR, "
                                 "CIF, CSSR, etc. If the keyword'prim' is within "
                                 "the filename, the code will automatically attempt "
                                 "to find a primitive cell.")
        groups.add_argument("-s", "--symmetry", dest="symmetry",
                            metavar="tolerance", type=float,
                            help="Determine the spacegroup using the "
                                 "specified tolerance. 0.1 is usually a good "
                                 "value for DFT calculations.")
        groups.add_argument("-g", "--group", dest="group",
                            choices=["element", "species"],
                            metavar="mode",
                            help="Compare a set of structures for similarity. "
                                 "Element mode does not compare oxidation states. "
                                 "Species mode will take into account oxidations "
                                 "states.")
        groups.add_argument(
            "-l", "--localenv", dest="localenv", nargs="+",
            help="Local environment analysis. Provide bonds in the format of"
                 "Center Species-Ligand Species=max_dist, e.g., H-O=0.5.")

        parser_structure.set_defaults(func=analyze_structures)

        parser_view = subparsers.add_parser("view", help="Visualize structures")
        parser_view.add_argument("filename", metavar="filename", type=str,
                                 nargs=1, help="Filename")
        parser_view.add_argument("-e", "--exclude_bonding", dest="exclude_bonding",
                                 type=str, nargs=1,
                                 help="List of elements to exclude from bonding "
                                      "analysis. E.g., Li,Na")
        parser_view.set_defaults(func=parse_view)

        parser_diff = subparsers.add_parser(
            "diff", help="Diffing tool. For now, only INCAR supported.")
        parser_diff.add_argument("-i", "--incar", dest="incars",
                                 metavar="INCAR", required=True,
                                 nargs=2, help="List of INCARs to compare.")
        parser_diff.set_defaults(func=diff_incar)

        parser_potcar = subparsers.add_parser("potcar",
                                              help="Generate POTCARs")
        parser_potcar.add_argument("-f", "--functional", dest="functional",
                                   type=str,
                                   choices=sorted(Potcar.FUNCTIONAL_CHOICES),
                                   default=SETTINGS.get("PMG_DEFAULT_FUNCTIONAL",
                                                        "PBE"),
                                   help="Functional to use. Unless otherwise "
                                        "stated (e.g., US), "
                                        "refers to PAW psuedopotential.")
        group = parser_potcar.add_mutually_exclusive_group(required=True)

        group.add_argument("-s", "--symbols", dest="symbols",
                           type=str, nargs="+",
                           help="List of POTCAR symbols. Use -f to set "
                                "functional. Defaults to PBE.")
        group.add_argument("-r", "--recursive", dest="recursive",
                           type=str,
                           help="Dirname to find and generate from POTCAR.spec.")
        parser_potcar.set_defaults(func=generate_potcar)

        try:
            import argcomplete
            argcomplete.autocomplete(parser)
        except ImportError:
            # argcomplete not present.
            pass

        args = parser.parse_args()

        try:
            getattr(args, "func")
        except AttributeError:
            parser.print_help()
            sys.exit(-1)
        return args.func(args)

    if __name__ == "__main__":
        main()