#!/usr/bin/env python

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

from __future__ import annotations

import argparse
import itertools

from tabulate import tabulate, tabulate_formats

from pymatgen.cli.pmg_analyze import analyze
from pymatgen.cli.pmg_config import configure_pmg
from pymatgen.cli.pmg_plot import plot
from pymatgen.cli.pmg_potcar import generate_potcar
from pymatgen.cli.pmg_structure import analyze_structures
from pymatgen.core import SETTINGS
from pymatgen.core.structure import Structure
from pymatgen.io.vasp import Incar, Potcar


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

    Args:
        args: Args from command.
    """
    from pymatgen.vis.structure_vtk import StructureVis

    excluded_bonding_elements = args.exclude_bonding[0].split(",") if args.exclude_bonding else []
    struct = Structure.from_file(args.filename[0])
    vis = StructureVis(excluded_bonding_elements=excluded_bonding_elements)
    vis.set_structure(struct)
    vis.show()
    return 0


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

    Args:
        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(f"{len(tuple(group))}*{i:.2f}" for (i, group) in itertools.groupby(v))
        return v

    diff = incar1.diff(incar2)
    output = [
        ["SAME PARAMS", "", ""],
        ["---------------", "", ""],
        ["", "", ""],
        ["DIFFERENT PARAMS", "", ""],
        ["----------------", "", ""],
    ]
    output += [
        (k, format_lists(diff["Same"][k]), format_lists(diff["Same"][k])) for k in sorted(diff["Same"]) if k != "SYSTEM"
    ]
    output += [
        (
            k,
            format_lists(diff["Different"][k]["INCAR1"]),
            format_lists(diff["Different"][k]["INCAR2"]),
        )
        for k in sorted(diff["Different"])
        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="""Author: Pymatgen Development Team""",
    )

    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",
    )

    groups.add_argument(
        "--cp2k",
        dest="cp2k_data_dirs",
        metavar="dir_name",
        nargs=2,
        help="Initial directory where the CP2K data is located and the output directory where the "
        "CP2K YAML data files will be written",
    )

    parser_config.add_argument(
        "-b",
        "--backup",
        default=".bak",
        help="Suffix to append to a backup of .pmgrc.yaml when changing this file. "
        "Defaults to '.bak'. Set to '' to disable.",
    )
    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_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:
        _ = args.func
    except AttributeError as exc:
        parser.print_help()
        raise SystemExit("Please specify a command.") from exc
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())
