# Copyright (c) lobsterpy development team
# Distributed under the terms of a BSD 3-Clause "New" or "Revised" License

"""Script to analyze Lobster outputs from the command line."""

from __future__ import annotations

import argparse
import json
from math import log, sqrt
from pathlib import Path

import matplotlib.style
from monty.json import jsanitize
from monty.os.path import zpath
from monty.serialization import dumpfn
from pymatgen.core.structure import Structure
from pymatgen.electronic_structure.cohp import Cohp, CompleteCohp
from pymatgen.io.lobster import Doscar, Icohplist, Lobsterin

from lobsterpy.cohp.analyze import Analysis
from lobsterpy.cohp.describe import Description
from lobsterpy.featurize.core import FeaturizeIcoxxlist
from lobsterpy.featurize.utils import get_file_paths
from lobsterpy.plotting import (
    BWDFPlotter,
    IcohpDistancePlotter,
    InteractiveCohpPlotter,
    PlainCohpPlotter,
    PlainDosPlotter,
    get_style_list,
)


def main() -> None:
    """Entry point for setup.py installer."""
    args = get_parser().parse_args()
    run(args)


def get_parser() -> argparse.ArgumentParser:
    """Construct argumentparser with subcommands and sections."""
    parser = argparse.ArgumentParser(description="Analyze and plot results from Lobster runs.")

    # Arguments that are needed by different actions, but not always

    incar_file = argparse.ArgumentParser(add_help=False)
    incar_file.add_argument(
        "-fincar",
        "--file-incar",
        default="INCAR",
        dest="incar",
        type=Path,
        help='path to INCAR. Default is "INCAR".',
    )

    charge_file = argparse.ArgumentParser(add_help=False)
    charge_file.add_argument(
        "-fcharge",
        "--file-charge",
        default="CHARGE.lobster",
        dest="charge",
        type=Path,
        help='path to CHARGE.lobster. Default is "CHARGE.lobster"',
    )

    structure_file = argparse.ArgumentParser(add_help=False)
    structure_file.add_argument(
        "-fstruct",
        "--file-structure",
        default="CONTCAR",
        dest="structure",
        type=Path,
        help='path to structure file. Default is "CONTCAR". '
        'Can also read "POSCAR.lobster" file or any '
        'suitable file format supported by pymatgen "Structure.from_file" method.',
    )

    potcar_file = argparse.ArgumentParser(add_help=False)
    potcar_file.add_argument(
        "-fpotcar",
        "--file-potcar",
        default="POTCAR",
        dest="potcar",
        type=Path,
        help='path to POTCAR. Default is "POTCAR".',
    )

    icohplist_file = argparse.ArgumentParser(add_help=False)
    icohplist_file.add_argument(
        "-ficohp",
        "--file-icohplist",
        default="ICOHPLIST.lobster",
        dest="icohplist",
        type=Path,
        help='path to ICOHPLIST.lobster. Default is "ICOHPLIST.lobster"',
    )

    coxxcar_file = argparse.ArgumentParser(add_help=False)
    coxxcar_file.add_argument(
        "-fcohp",
        "--file-cohpcar",
        default="COHPCAR.lobster",
        dest="cohpcar",
        type=Path,
        help='path to COHPCAR.lobster. Default is "COHPCAR.lobster". This argument '
        "can also read COBICARs or COOPCARs. One needs to use appropriate --cobis or "
        "--coops options along with this argument when plotting",
    )

    doscar_file = argparse.ArgumentParser(add_help=False)
    doscar_file.add_argument(
        "-fdos",
        "--file-doscar",
        default="DOSCAR.lobster",
        dest="doscar",
        type=Path,
        help='path to DOSCAR.lobster. Default is "DOSCAR.lobster".',
    )

    user_basis_arg = argparse.ArgumentParser(add_help=False)
    user_basis_arg.add_argument(
        "--userbasis",
        "--user-basis",
        default=None,
        type=_element_basis,
        nargs="+",
        help="Use the specific basis provided by the user to "
        "generate the inputs (e.g.,  --userbasis Cr.3d.3p.4s N.2s.2p). "
        "Default is None.",
    )

    # groups of arguments for specific actions

    calc_quality_description_file_parent = argparse.ArgumentParser(add_help=False)

    # Input args for specifically needed for calculation quality description
    calc_quality_description_file_group = calc_quality_description_file_parent.add_argument_group("")
    calc_quality_description_file_group.add_argument(
        "-fvasprun",
        "--file-vasprun",
        default="vasprun.xml",
        dest="vasprun",
        type=Path,
        help='path to vasprun.xml. Default is "vasprun.xml".',
    )
    calc_quality_description_file_group.add_argument(
        "-fbandoverlaps",
        "--file-bandoverlaps",
        default="bandOverlaps.lobster",
        dest="bandoverlaps",
        type=Path,
        help='path to bandOverlaps.lobster. Default is "bandOverlaps.lobster".',
    )
    calc_quality_description_file_group.add_argument(
        "-flobsterin",
        "--file-lobsterin",
        default="lobsterin",
        dest="lobsterin",
        type=Path,
        help='path to lobsterin. Default is "lobsterin".',
    )
    calc_quality_description_file_group.add_argument(
        "-flobsterout",
        "--file-lobsterout",
        default="lobsterout",
        dest="lobsterout",
        type=Path,
        help='path to lobsterout. Default is "lobsterout".',
    )
    calc_quality_description_file_group.add_argument(
        "-potsymbols",
        "--potcar-symbols",
        dest="potcarsymbols",
        type=_potcar_symbols,
        # nargs="+",
        help="List of potcar symbols",
    )

    # group of arguments related to writing LOBSTER calcs inputs
    output_parent = argparse.ArgumentParser(add_help=False)
    output_file_group = output_parent.add_argument_group("Output files")
    output_file_group.add_argument(
        "-fincarout",
        "--file-incar-out",
        default="INCAR.lobsterpy",
        dest="incarout",
        type=Path,
        help='path to INCAR that lobsterpy generates. Default is "INCAR.lobsterpy"',
    )
    output_file_group.add_argument(
        "-flobsterin",
        "--file-lobsterin",
        default="lobsterin.lobsterpy",
        dest="lobsterinout",
        type=Path,
        help='Base for path to lobsterins that lobsterpy generates. Default is "lobsterin.lobsterpy"',
    )
    output_file_group.add_argument(
        "--overwrite",
        "--overwrite-files",
        dest="overwrite",
        default=False,
        action="store_true",
        help="overwrites already created INCARs and lobsterins with the give name.",
    )
    # TODO: Add some output arguments: options to supply your own basis
    # General matplotlib plotting arguments common to all kinds of plots

    plotting_parent = argparse.ArgumentParser(add_help=False)
    plotting_group = plotting_parent.add_argument_group("Plotting")
    plotting_group.add_argument(
        "--fontsize",
        "--font-size",
        type=float,
        dest="fontsize",
        default=None,
        help="Base font size",
    )
    plotting_group.add_argument(
        "-ht",
        "--height",
        type=float,
        dest="height",
        default=None,
        help="Plot height in inches",
    )
    plotting_group.add_argument(
        "--hideplot",
        "--hide-plot",
        dest="hideplot",
        action="store_true",
        help="Hide plot output. Especially relevant when plots are saved.",
    )
    plotting_group.add_argument(
        "-nbs",
        "--no-base-style",
        "--nobasestyle",
        action="store_true",
        dest="no_base_style",
        help=(
            "Disable inbuilt style entirely. This may prevent interference with external "
            "stylesheets when using --style."
        ),
    )
    plotting_group.add_argument(
        "-sty",
        "--style",
        type=str,
        nargs="+",
        dest="style",
        default=None,
        help="Matplotlib style sheet(s) for plot appearance",
    )
    plotting_group.add_argument(
        "--save-plot",
        "--saveplot",
        "-s",
        type=Path,
        metavar="FILENAME",
        default=None,
        dest="save_plot",
        help="Save plot to file",
    )
    plotting_group.add_argument("--title", type=str, default="", help="Plot title")
    plotting_group.add_argument(
        "-wd",
        "--width",
        type=float,
        dest="width",
        default=None,
        help="Plot width in inches",
    )
    plotting_group.add_argument(
        "-x",
        "--xlim",
        dest="xlim",
        nargs=2,
        default=None,
        type=float,
        help="Set x-axis limits for the plots",
    )
    plotting_group.add_argument(
        "-y",
        "--ylim",
        dest="ylim",
        nargs=2,
        default=None,
        type=float,
        help="Set y-axis limits for the plots",
    )
    broadening_group = plotting_group.add_mutually_exclusive_group()
    broadening_group.add_argument_group("Broadening")
    broadening_group.add_argument(
        "--fwhm",
        type=float,
        default=None,
        help="Full-width-half-maximum of Gaussian broadening.",
    )
    broadening_group.add_argument(
        "--sigma",
        type=float,
        default=None,
        help="Standard deviation of Gaussian broadening.",
    )
    # Group pf arguments specific to dos plotter
    dos_plotting_parent = argparse.ArgumentParser(add_help=False)
    dos_plotting_group = dos_plotting_parent.add_argument_group("Plotting")
    dos_plotting_group.add_argument(
        "-addtdos",
        "--addtotal",
        "--add-total",
        action="store_true",
        help="Add total dos to DOS plot.",
    )
    dos_plotting_group.add_argument(
        "-el",
        "--element",
        type=str,
        dest="element",
        nargs="+",
        default=None,
        help="Add spd DOS projections for requested element to the DOS plot",
    )
    dos_plotting_group.add_argument(
        "-orb",
        "--orbital",
        type=str,
        dest="orbital",
        nargs="+",
        default=None,
        help="Orbital name for the site for which DOS are to be added",
    )
    dos_plotting_group.add_argument(
        "--site",
        type=int,
        dest="site",
        nargs="+",
        default=None,
        help="Site index in the crystal structure for which DOS need to be added",
    )
    dos_plotting_group.add_argument(
        "-sspin",
        "--summedspins",
        "--summed-spins",
        dest="summedspins",
        action="store_true",
        help="Plot summed spins DOS",
    )

    group = dos_plotting_parent.add_mutually_exclusive_group()
    group.add_argument(
        "-eldos",
        "--elementdos",
        dest="elementdos",
        action="store_true",
        help="Add DOS projections for each element to the DOS plot",
    )
    group.add_argument(
        "--spddos",
        "--spd-dos",
        dest="spddos",
        action="store_true",
        help="Add spd projected dos to the DOS plot",
    )
    # Argument common to COHPs / COOPs / COBIs or DOS plots
    advanced_plotting_args = argparse.ArgumentParser(add_help=False)
    advanced_plotting_args.add_argument(
        "--invertaxis",
        "--invert-axis",
        dest="invertaxis",
        action="store_true",
        help="Invert plot axis of DOS or COOPs COHPs or COBIS",
    )
    # Argument specific to COHPs / COOPs / COBIs plots
    coxx_plotting_args = argparse.ArgumentParser(add_help=False)
    coxx_plotting_args.add_argument(
        "-integr",
        "--integrated",
        dest="integrated",
        action="store_true",
        help="Show integrated cohp/cobi/coop plots.",
    )
    # Arguments specific to lobsterpy.cohp.analyze.Analysis and
    # lobsterpy.cohp.describe.Description class
    auto_parent = argparse.ArgumentParser(add_help=False)
    auto_group = auto_parent.add_argument_group("Adjustable automatic analysis parameters")
    auto_group.add_argument(
        "-allb",
        "--allbonds",
        "--all-bonds",
        action="store_true",
        default=False,
        help="Consider all bonds during the automatic analysis, not only cation-anion bonds (default) ",
    )
    auto_group.add_argument(
        "--cutofficohp",
        type=float,
        default=0.1,
        help="Consider only bonds that are stronger than cutoff_icoxx * strongest ICOXX "
        " (ICOHP or ICOBI or ICOOP) for automatic analysis.",
    )
    auto_group.add_argument(
        "-fjson",
        "--file-json",
        nargs="?",
        type=Path,
        default=None,
        metavar="FILENAME",
        const=Path("lobsterpy.json"),
        help="Write a JSON file with the most important information",
    )
    auto_group.add_argument(
        "--noisecutoff",
        type=float,
        default=0.1,
        help="Sets the lower limit of icohps or icoops or icobis considered in automatic analysis",
    )
    auto_group.add_argument(
        "-orbresol",
        "--orbitalresolved",
        "--orbital-resolved",
        action="store_true",
        default=False,
        help="Switch on orbital resolved analysis of (I)COHPs or (I)COBIs or (I)COOPs with all relevant orbitals.",
    )
    auto_group.add_argument(
        "-orbcutoff",
        "--orbitalcutoff",
        "--orbital-cutoff",
        type=float,
        default=0.05,
        help="Consider only orbital interactions that are stronger than orbitalintcutoff * 100 of "
        "relevant bonds (ICOHP or ICOBI or ICOOP). Can only be used in combination "
        "with orbital-resolved automatic analysis (--orbitalresolved).",
    )

    auto_group.add_argument(
        "-sspins",
        "--summedspins",
        "--summed-spins",
        action="store_true",
        default=False,
        help="Sum COHP `Spin.up` and `Spin.down` populations for automatic analysis",
    )
    auto_group.add_argument(
        "--save-plot-json",
        "--saveplotjson",
        "-spj",
        nargs="?",
        type=Path,
        default=None,
        metavar="FILENAME",
        help="Write a JSON file with the plot data. "
        "`monty.serialization.loadfn` or `json.load` can be used to load the data.",
    )
    auto_group.add_argument(
        "-tchg",
        "--typecharge",
        type=str,
        dest="type_charge",
        default="Mulliken",
        help="Switch between different types of charges for automatic analysis. "
        "Possible values are: 'Mulliken', 'Loewdin' or 'Valences'. ",
    )

    # Argument that will help to switch automatic analysis
    analysis_switch = argparse.ArgumentParser(add_help=False)
    analysis_group = analysis_switch.add_argument_group(
        "Arguments to switch type of files analyzed during automatic analysis"
        " (Also indicates file to be read for 'plot/plot-icohp-distance/plot-bwdf' action in cli)"
    )
    analysis_group.add_argument(
        "--cobis",
        "--cobis",
        action="store_true",
        help="Setting this option starts automatic bonding analysis using COBIs"
        " (Also indicates plotter that input file contains COBI data)",
    )
    analysis_group.add_argument(
        "--coops",
        "--coops",
        action="store_true",
        help="Setting this option starts automatic bonding analysis using COOPs"
        " (Also indicates plotter that input file contains COOP data)",
    )

    # Specific to interactive plotter args
    interactive_plotter_args = argparse.ArgumentParser(add_help=False)
    interactive_plotter_group = interactive_plotter_args.add_argument_group("Options specific to interactive plotter")
    interactive_plotter_group.add_argument(
        "-labresol",
        "--labelresolved",
        "--label-resolved",
        action="store_true",
        help="Create automatic interactive plots with all relevant bond labels. "
        "If not set, plots consist of summed cohps.",
    )

    # Args specific to icohp distance plotter args
    icohp_distance_plotter_args = argparse.ArgumentParser(add_help=False)
    icohp_distance_plotter_group = icohp_distance_plotter_args.add_argument_group(
        "Options specific to ICOHP distance plotter"
    )
    icohp_distance_plotter_group.add_argument(
        "-alpha",
        type=float,
        default=0.4,
        dest="alpha",
        help="Transparency of the markers in the plot.",
    )
    icohp_distance_plotter_group.add_argument(
        "-cbonds",
        "--colorbonds",
        "--color-bonds",
        default=False,
        dest="colorbonds",
        action="store_true",
        help="If set, will color ICOHPs based on atom types",
    )
    icohp_distance_plotter_group.add_argument(
        "-c",
        "--colors",
        dest="colors",
        nargs="+",
        type=str,
        default=None,
        help="Custom list of colors for the plot. Should be equal to the number of unique atom pairs in the structure.",
    )
    icohp_distance_plotter_group.add_argument(
        "-lprefix",
        "--legendprefix",
        "--legend-prefix",
        type=str,
        dest="legendprefix",
        default="",
        help="Prefix for the legend in the plot.",
    )
    icohp_distance_plotter_group.add_argument(
        "-msize",
        "--markersize",
        "--marker-size",
        type=float,
        default=50,
        dest="markersize",
        help="Size of the markers in the plot.",
    )
    icohp_distance_plotter_group.add_argument(
        "-mstyle",
        "--markerstyle",
        "--marker-style",
        type=str,
        default="o",
        dest="markerstyle",
        help="Marker style for the plot.",
    )

    # Args specific to plotting BWDFs
    bwdf_plotting_args = argparse.ArgumentParser(add_help=False)
    bwdf_plotting_group = bwdf_plotting_args.add_argument_group("Options specific to plotting BWDFs")
    bwdf_plotting_group.add_argument(
        "-atompairs",
        "--atom-pairs",
        dest="atompairs",
        action="store_true",
        default=False,
        help="If True, will plot the BWDFs for all unique atom pairs. Default: False.",
    )
    bwdf_plotting_group.add_argument(
        "-binwidth",
        "--bin-width",
        type=float,
        dest="binwidth",
        default=0.02,
        help="Bin width used for computing the BWDFs. Default: 0.02.",
    )
    bwdf_plotting_group.add_argument(
        "-interacttol",
        "--interactions-tolerance",
        type=float,
        dest="interacttol",
        default=1e-3,
        help="Numerical tolerance considered for interactions to be insignificant. Default: 1e-3.",
    )
    bwdf_plotting_group.add_argument(
        "-maxlen",
        "--maxlength",
        "--max-length",
        type=float,
        dest="maxlen",
        default=6,
        help="Maximum bond length for the BWDFs in Angstroms. Default: 6.",
    )
    bwdf_plotting_group.add_argument(
        "-minlen",
        "--minlength",
        "--min-length",
        type=float,
        dest="minlen",
        default=0,
        help="Minimum bond length for the BWDFs in Angstroms. Default: 0.",
    )
    bwdf_plotting_group.add_argument(
        "-norm",
        "--normalization",
        type=str,
        default="formula_units",
        dest="norm",
        help="Normalization of the BWDFs. Options: 'formula_units', 'area', 'counts' and 'none'. "
        "Default: 'formula_units'.",
    )
    bwdf_plotting_group.add_argument(
        "-plotneg",
        "--plot-negative",
        dest="plotneg",
        action="store_true",
        default=False,
        help="If True, will plot -1*ICOHPs. Works only for ICOHPs. Default: True.",
    )
    bwdf_plotting_group.add_argument(
        "-siteindex",
        "--site-index",
        type=int,
        default=None,
        dest="siteindex",
        help="Site index for which the BWDFs are to be plotted. Default: None. If None, all sites are considered.",
    )

    # Args specific to calc quality description dict and texts
    calc_quality_args = argparse.ArgumentParser(add_help=False)
    calc_quality_args_group = calc_quality_args.add_argument_group(
        "Options to change default output of quality description"
    )
    calc_quality_args_group.add_argument(
        "--bvacomp",
        "--bva-comp",
        action="store_true",
        default=False,
        help="Enable the BVA charge comparison in automatic LOBSTER calc quality analysis ",
    )
    calc_quality_args_group.add_argument(
        "--doscomp",
        "--dos-comp",
        action="store_true",
        default=False,
        help="Enable the DOS comparison in automatic LOBSTER calc quality analysis ",
    )
    calc_quality_args_group.add_argument(
        "--erange",
        dest="erange",
        nargs=2,
        default=[-5, 0],
        type=int,
        help="Energy range for DOS comparisons",
    )
    calc_quality_args_group.add_argument(
        "-fcalcqualjson",
        "--file-calc-quality-json",
        nargs="?",
        type=Path,
        default=None,
        metavar="FILENAME",
        const=Path("calc_quality_json.json"),
        help="Write a JSON file with the LOBSTER calc quality analysis",
    )
    calc_quality_args_group.add_argument(
        "--nbins",
        dest="nbins",
        default=None,
        type=int,
        help="Number of bins for DOS comparisons",
    )
    # Build the actions using arguments defined earlier
    subparsers = parser.add_subparsers(
        dest="action",
        required=True,
        help="Use -h/--help after the chosen subcommand to see further options.",
    )
    subparsers.add_parser(
        "create-inputs",
        aliases=["createinputs"],
        parents=[
            incar_file,
            structure_file,
            potcar_file,
            user_basis_arg,
            output_parent,
        ],
        help="Create inputs for lobster computation. It works only with PBE POTCARs, "
        "as, currently, only the pbeVASPfit2015 basis fitted to PBE POTCARs in LOBSTER includes "
        "additional orbitals relevant for the solid-state materials. Please check out our "
        "publication https://doi.org/10.1002/cplu.202200123 and LOBSTER program manual for more information.",
    )
    subparsers.add_parser(
        "description",
        parents=[
            charge_file,
            coxxcar_file,
            icohplist_file,
            structure_file,
            auto_parent,
            analysis_switch,
        ],
        help=("Deliver a text description from automatic analysis of COHPs or COBIS or COOP results from Lobster run"),
    )
    subparsers.add_parser(
        "description-quality",
        parents=[
            calc_quality_description_file_parent,
            charge_file,
            doscar_file,
            potcar_file,
            structure_file,
            calc_quality_args,
        ],
        help=(
            "Deliver a text description of the LOBSTER calc quality analysis. "
            "Mandatory required files: structure file compatible with --file-structure (preferably CONTCAR), "
            "POTCAR, lobsterout, lobsterin. "
            "Optional files (BVA comparison): CHARGE.lobster, "
            "(DOS comparison): DOSCAR.lobster/ DOSCAR.LSO.lobster, Vasprun.xml."
        ),
    )
    subparsers.add_parser(
        "plot-automatic",
        aliases=[
            "plot-auto",
            "automatic-plot",
            "automaticplot",
            "auto-plot",
            "autoplot",
        ],
        parents=[
            charge_file,
            coxxcar_file,
            icohplist_file,
            structure_file,
            auto_parent,
            plotting_parent,
            coxx_plotting_args,
            analysis_switch,
        ],
        help=(
            "Plot most important COHPs or COBIs or COOPs automatically."
            " This option also includes an automatic description."
        ),
    )
    subparsers.add_parser(
        "plot-automatic-ia",
        aliases=[
            "plot-auto-ia",
            "automatic-plot-ia",
            "automaticplotia",
            "autoplotia",
            "auto-plot-ia",
        ],
        parents=[
            charge_file,
            coxxcar_file,
            icohplist_file,
            structure_file,
            auto_parent,
            plotting_parent,
            coxx_plotting_args,
            interactive_plotter_args,
            analysis_switch,
        ],
        help=("Creates an interactive plot of most important COHPs or COBIs or COOPs automatically."),
    )
    subparsers.add_parser(
        "plot-bwdf",
        aliases=["plotbwdf"],
        parents=[
            icohplist_file,
            structure_file,
            bwdf_plotting_args,
            plotting_parent,
            analysis_switch,
        ],
        help="Plot bond-weighted distribution functions (BWDFs) from ICOXXLIST.lobster.",
    )
    subparsers.add_parser(
        "plot-dos",
        aliases=["plotdos"],
        parents=[
            doscar_file,
            structure_file,
            dos_plotting_parent,
            advanced_plotting_args,
            plotting_parent,
        ],
        help="Plots DOS from lobster computation.",
    )
    subparsers.add_parser(
        "plot-icohp-distance",
        aliases=["ploticohpdistance"],
        parents=[icohplist_file, plotting_parent, analysis_switch, icohp_distance_plotter_args],
        help="Plot ICOHPs or ICOOPs or ICOBIs with respect to bond lengths",
    )
    # Mode for normal plotting (without automatic detection of relevant COHPs)
    plot_parser = subparsers.add_parser(
        "plot",
        parents=[
            coxxcar_file,
            structure_file,
            plotting_parent,
            advanced_plotting_args,
            coxx_plotting_args,
            analysis_switch,
        ],
        help="Plot specific COHPs/COBIs/COOPs based on bond numbers.",
    )
    plot_parser.add_argument(
        "bond_numbers",
        nargs="+",
        type=int,
        help="List of bond numbers, determining COHPs/COBIs/COOPs to include in plot.",
    )
    plot_grouping = plot_parser.add_mutually_exclusive_group()
    plot_grouping.add_argument("--multi_cobis", action="store_true", help="Plot COBICARs including multi-center bonds.")

    plot_grouping.add_argument(
        "--summed",
        action="store_true",
        help="Show a summed COHP",
    )
    plot_grouping.add_argument(
        "-orbwise",
        "--orbitalwise",
        dest="orbitalwise",
        nargs="+",
        default=None,
        type=str,
        help=(
            "Plot cohps of specific orbitals. e.g. to plot 2s-2s interaction of "
            'bond with label 1, use "lobsterpy plot 1 --orbitalwise 2s-2s". '
            'To plot all orbitalwise cohps of one bond, you can use "all" instead of "2s-2s". '
            "To plot orbitalwise interactions of more than one bond, use, for example, "
            '"lobsterpy plot 1 1 --orbitalwise "3s-3s" "2px-3s"'
        ),
    )
    return parser


def _element_basis(string: str):
    """
    Parse element and basis from string.

    :param string: string to parse

    Returns:
            element, basis
    """
    cut_list = string.split(".")
    element = cut_list[0]
    basis = " ".join(cut_list[1:])
    return element, basis


def _potcar_symbols(string: str):
    """
    Parse string of potcar symbols and return a list.

    :param string: string of potcar symbols

    Returns:
        list of potcar symbols
    """
    return string.split(" ")


def _user_figsize(width, height, aspect=None):
    """Get figsize options from user input, if any.

    If only width x or height is provided, use a target aspect ratio to derive
    the other one.

    Returns a dict which can be merged into style kwargs
    """
    if width is None and height is None:
        return {}
    if width is not None and height is not None:
        return {"figure.figsize": (width, height)}

    if aspect is None:
        aspect = (sqrt(5) + 1) / 2  # Golden ratio
    if width is None:
        return {"figure.figsize": (height * aspect, height)}
    return {"figure.figsize": (width, width / aspect)}


def run(args):
    """
    Run actions based on args.

    :param args: args for cli

    """
    if args.action in [
        "plot-auto",
        "automatic-plot",
        "automaticplot",
        "auto-plot",
        "autoplot",
    ]:
        args.action = "plot-automatic"

    if args.action in [
        "plot-auto-ia",
        "automatic-plot-ia",
        "automaticplotia",
        "autoplotia",
        "auto-plot-ia",
    ]:
        args.action = "plot-automatic-ia"

    if args.action in [
        "description",
        "plot-automatic",
        "plot-automatic-ia",
    ]:
        req_files = get_file_paths(
            path_to_lobster_calc=Path.cwd(), requested_files=["structure", "charge", "icohplist", "cohpcar"]
        )

        if args.coops:
            req_files_coops = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["icooplist", "coopcar"])

            req_files["icohplist"] = req_files_coops["icooplist"]
            req_files["cohpcar"] = req_files_coops["coopcar"]

        if args.cobis:
            req_files_cobis = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["icobilist", "cobicar"])

            req_files["icohplist"] = req_files_cobis["icobilist"]
            req_files["cohpcar"] = req_files_cobis["cobicar"]

        for arg_name in req_files:
            setattr(args, arg_name, req_files[arg_name])

    if args.action in ["description", "plot-automatic", "plot-automatic-ia"]:
        which_bonds = "all" if args.allbonds else "cation-anion"

        analyse = Analysis(
            path_to_poscar=args.structure,
            path_to_charge=args.charge,
            path_to_cohpcar=args.cohpcar,
            path_to_icohplist=args.icohplist,
            which_bonds=which_bonds,
            are_coops=args.coops,
            are_cobis=args.cobis,
            noise_cutoff=args.noisecutoff,
            cutoff_icohp=args.cutofficohp,
            orbital_cutoff=args.orbitalcutoff,
            orbital_resolved=args.orbitalresolved,
            summed_spins=args.summedspins,
            type_charge=args.type_charge,
        )

        describe = Description(analysis_object=analyse)
        describe.write_description()

        if args.file_json is not None:
            analysedict = analyse.condensed_bonding_analysis
            with open(args.file_json, "w") as fd:
                json.dump(analysedict, fd)

        if args.save_plot_json is not None:
            ic_plotter = InteractiveCohpPlotter(are_coops=args.coops, are_cobis=args.cobis)
            if args.action in ["plot-automatic-ia"]:
                ic_plotter.add_all_relevant_cohps(
                    analyse=analyse, label_resolved=args.labelresolved, orbital_resolved=args.orbitalresolved
                )
            else:
                ic_plotter.add_all_relevant_cohps(
                    analyse=analyse, label_resolved=False, orbital_resolved=args.orbitalresolved
                )
            monty_encoded_doc = jsanitize(ic_plotter._cohps, allow_bson=True, strict=True, enum_values=True)
            for bond_label, cohp_data in monty_encoded_doc["All"].items():
                cohp_data.update({"are_coops": args.coops, "are_cobis": args.cobis})
                monty_encoded_doc["All"][bond_label] = Cohp.from_dict(cohp_data)
            dumpfn(monty_encoded_doc["All"], args.save_plot_json)

    if args.action in [
        "plot",
        "plot-automatic",
        "plot-automatic-ia",
        "plot-bwdf",
        "plotbwdf",
        "plot-dos",
        "plotdos",
        "plot-icohp-distance",
        "ploticohpdistance",
    ]:
        style_kwargs = {}
        style_kwargs.update(_user_figsize(args.width, args.height))
        if args.fontsize:
            style_kwargs.update({"font.size": args.fontsize})

        style_list = get_style_list(no_base_style=args.no_base_style, styles=args.style, **style_kwargs)
        matplotlib.style.use(style_list)

        if args.sigma:
            sigma = args.sigma
        elif args.fwhm:
            sigma = args.fwhm / (2 * sqrt(2 * log(2)))
        else:
            sigma = None

    if args.action in ["plot-automatic"]:
        describe.plot_cohps(
            ylim=args.ylim,
            xlim=args.xlim,
            integrated=args.integrated,
            save=args.save_plot is not None,
            filename=args.save_plot,
            title=args.title,
            sigma=sigma,
            hide=args.hideplot,
        )

    if args.action in ["plot-automatic-ia"]:
        describe.plot_interactive_cohps(
            ylim=args.ylim,
            xlim=args.xlim,
            integrated=args.integrated,
            save_as_html=args.save_plot,
            filename=args.save_plot,
            title=args.title,
            sigma=sigma,
            hide=args.hideplot,
            label_resolved=args.labelresolved,
            orbital_resolved=args.orbitalresolved,
        )

    if args.action in ["plot-bwdf", "plotbwdf"]:
        if args.cobis:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["icobilist"]).get("icobilist")
            options = {"are_cobis": True, "are_coops": False}
        elif args.coops:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["icooplist"]).get("icooplist")
            options = {"are_cobis": False, "are_coops": True}
        else:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["icohplist"]).get("icohplist")
            options = {"are_cobis": False, "are_coops": False}

        structure_filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["structure"]).get(
            "structure"
        )

        feat_icoxx = FeaturizeIcoxxlist(
            path_to_icoxxlist=filename,
            path_to_structure=structure_filename,
            interactions_tol=args.interacttol,
            normalization=args.norm,
            bin_width=args.binwidth,
            min_length=args.minlen,
            max_length=args.maxlen,
            **options,
        )

        bwdf = feat_icoxx.calc_site_bwdf(site_index=args.siteindex) if args.siteindex else feat_icoxx.calc_bwdf()

        # Set label for plot (legend)
        if args.siteindex:
            site = feat_icoxx.structure.sites[args.siteindex]
            label = site.species_string
        else:
            label = feat_icoxx.structure.composition.get_reduced_formula_and_factor()[0]

        bwdf_plotter = BWDFPlotter(are_coops=args.coops, are_cobis=args.cobis)

        # Assimilate data to get separate plots for each pair
        plot_data = []
        for pair, value in bwdf.items():
            if args.atompairs and pair not in ["summed", "centers", "edges", "bin_width", "wasserstein_dist_to_rdf"]:
                data = {pair: value}
                data.update(
                    {
                        "centers": bwdf["centers"],
                        "edges": bwdf["edges"],
                        "bin_width": bwdf["bin_width"],
                        "wasserstein_dist_to_rdf": bwdf["wasserstein_dist_to_rdf"],
                    }
                )
                plot_data.append(data)
            elif not args.atompairs and pair == "summed":  # summed only
                data = {pair: value}
                data.update(
                    {
                        "centers": bwdf["centers"],
                        "edges": bwdf["edges"],
                        "bin_width": bwdf["bin_width"],
                        "wasserstein_dist_to_rdf": bwdf["wasserstein_dist_to_rdf"],
                    }
                )
                plot_data.append(data)
            elif not args.atompairs and pair == str(args.siteindex):  # specific site
                data = {pair: value}
                data.update(
                    {
                        "centers": bwdf["centers"],
                        "edges": bwdf["edges"],
                        "bin_width": bwdf["bin_width"],
                        "wasserstein_dist_to_rdf": bwdf["wasserstein_dist_to_rdf"],
                    }
                )
                plot_data.append(data)

        # Iterate over assimilated data and plot
        for bwdf_dict in plot_data:
            pair = next(iter(bwdf_dict.keys()))
            bwdf_plotter.add_bwdf(bwdf=bwdf_dict, label=label)
            plot = bwdf_plotter.get_plot(sigma=args.sigma, xlim=args.xlim, ylim=args.ylim, plot_negative=args.plotneg)
            title = f"{args.title} : {pair}" if args.title else ""
            plot.title(title)

            if args.save_plot:
                filename = Path(args.save_plot)
                plot_filename = filename.parent / f"{pair.replace('-', '_')}_{filename.stem}{filename.suffix}"

            if not args.hideplot and not args.save_plot:
                plot.show()
            elif args.save_plot and not args.hideplot:
                fig = plot.gcf()
                fig.savefig(plot_filename)
                plot.show()
            if args.save_plot and args.hideplot:
                plot.savefig(plot_filename)

    if args.action == "plot":
        if args.cobis:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["cobicar"]).get("cobicar")
            options = {"are_cobis": True, "are_coops": False}
        elif args.multi_cobis:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["cobicar"]).get("cobicar")
            options = {"are_cobis": False, "are_coops": False, "are_multi_center_cobis": True}
        elif args.coops:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["coopcar"]).get("coopcar")
            options = {"are_cobis": False, "are_coops": True}
        else:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["cohpcar"]).get("cohpcar")
            options = {"are_cobis": False, "are_coops": False}

        structure_filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["structure"]).get(
            "structure"
        )

        completecohp = CompleteCohp.from_file(
            fmt="LOBSTER",
            filename=filename,
            structure_file=structure_filename,
            **options,
        )
        if args.multi_cobis:
            options = {"are_cobis": True, "are_coops": False}
        cp = PlainCohpPlotter(**options)

        if not args.summed:
            # TODO: add checks for label in allowed labels -> print all labels
            # TODO: add check if args.oribtalwise is exactly as long as labels
            # TODO: add check if orbital is in args.orbitalwise

            for label in args.bond_numbers:
                if str(label) not in completecohp.bonds:
                    raise IndexError(
                        "The provided bond label " + str(label) + " is not available in ICO**LIST.lobster.\n "
                        "Allowed options are in this list: \n"
                        + str([int(listi) for listi in list(completecohp.bonds.keys())])
                    )

            if not args.orbitalwise:
                for label in args.bond_numbers:
                    cp.add_cohp(label, completecohp.get_cohp_by_label(label=str(label)))
            else:
                if len(args.bond_numbers) != len(args.orbitalwise):
                    raise IndexError(
                        "Please provide as mainy orbitals as bond labels,"
                        " e.g., lobsterpy plot 1 1 --orbitalwise '2s-2s' '2s-2px'"
                    )

                for ilabel, label in enumerate(args.bond_numbers):
                    orbitals = args.orbitalwise[ilabel]

                    availableorbitals = list(completecohp.orb_res_cohp[str(label)].keys())
                    orbitaloptions = [*availableorbitals, "all"]

                    if orbitals not in orbitaloptions:
                        raise IndexError(
                            "Orbital in not available for current bond. \n"
                            "For bond "
                            + str(label)
                            + " only the following orbital options are available: \n"
                            + str(orbitaloptions)
                        )

                    if orbitals != "all":
                        cp.add_cohp(
                            str(label) + ": " + orbitals,
                            completecohp.get_orbital_resolved_cohp(label=str(label), orbitals=orbitals),
                        )
                    else:
                        for orbitals in availableorbitals:
                            cp.add_cohp(
                                str(label) + ": " + orbitals,
                                completecohp.get_orbital_resolved_cohp(label=str(label), orbitals=orbitals),
                            )
        else:
            cp.add_cohp(
                str(args.bond_numbers),
                completecohp.get_summed_cohp_by_label_list(label_list=[str(label) for label in args.bond_numbers]),
            )

        plt = cp.get_plot(integrated=args.integrated, xlim=args.xlim, ylim=args.ylim, sigma=sigma)

        ax = plt.gca()
        ax.set_title(args.title)

        if not args.hideplot and not args.save_plot:
            plt.show()
        elif args.save_plot and not args.hideplot:
            fig = plt.gcf()
            fig.savefig(args.save_plot)
            plt.show()
        if args.save_plot and args.hideplot:
            plt.savefig(args.save_plot)

    if args.action in ["create-inputs", "createinputs"]:
        # Check for .gz files exist for default values and update accordingly
        req_file_arg_names = ["structure", "potcar", "incar"]

        for arg_name in req_file_arg_names:
            file_path = getattr(args, arg_name)
            if not file_path.exists():
                gz_file_path = file_path.with_name(zpath(file_path.name))
                if gz_file_path.exists():
                    setattr(args, arg_name, gz_file_path)
                else:
                    raise ValueError(
                        "Files necessary for creating inputs for LOBSTER calcs not found in the current directory."
                    )

        if args.userbasis is None:
            # This will rely on standard basis files as stored in pymatgen

            potcar_names = Lobsterin._get_potcar_symbols(POTCAR_input=args.potcar)

            list_basis_dict = Lobsterin.get_all_possible_basis_functions(
                structure=Structure.from_file(args.structure),
                potcar_symbols=potcar_names,
            )

            for ibasis, basis_dict in enumerate(list_basis_dict):
                lobsterinput = Lobsterin.standard_calculations_from_vasp_files(
                    args.structure,
                    args.incar,
                    None,
                    option="standard",
                    dict_for_basis=basis_dict,
                )

                lobsterin_path = Path(str(args.lobsterinout) + "-" + str(ibasis))
                incar_path = Path(str(args.incarout) + "-" + str(ibasis))

                if (not lobsterin_path.is_file() and not incar_path.is_file()) or (args.overwrite):
                    lobsterinput.write_lobsterin(lobsterin_path)
                    lobsterinput.write_INCAR(
                        incar_input=args.incar,
                        incar_output=incar_path,
                        poscar_input=args.structure,
                        isym=0,
                    )
                else:
                    raise ValueError('please use "--overwrite" if you would like to overwrite existing lobster inputs')
        else:
            # convert list userbasis to dict
            userbasis = {}
            for userbasis_single in args.userbasis:
                userbasis[userbasis_single[0]] = userbasis_single[1]

            lobsterinput = Lobsterin.standard_calculations_from_vasp_files(
                args.structure,
                args.incar,
                None,
                option="standard",
                dict_for_basis=userbasis,
            )

            lobsterin_path = Path(str(args.lobsterinout) + "-" + str(0))
            incar_path = Path(str(args.incarout) + "-" + str(0))

            if (not lobsterin_path.is_file() and not incar_path.is_file()) or (args.overwrite):
                lobsterinput.write_lobsterin(lobsterin_path)
                lobsterinput.write_INCAR(
                    incar_input=args.incar,
                    incar_output=incar_path,
                    poscar_input=args.structure,
                    isym=0,
                )
            else:
                raise ValueError('please use "--overwrite" if you would like to overwrite existing lobster inputs')

    if args.action in ["description-quality"]:
        # Check for .gz files exist for default values and update accordingly
        req_files = get_file_paths(
            path_to_lobster_calc=Path.cwd(), requested_files=["structure", "lobsterin", "lobsterout"]
        )
        for arg_name in req_files:
            setattr(args, arg_name, req_files[arg_name])

        optional_files = {
            "bandoverlaps": "bandOverlaps.lobster",
            "potcar": "POTCAR",
            "vasprun": "vasprun.xml",
        }

        for arg_name in optional_files:
            file_path = getattr(args, arg_name)
            if not file_path.exists():
                gz_file_path = file_path.with_name(zpath(file_path.name))
                if gz_file_path.exists():
                    setattr(args, arg_name, gz_file_path)

        bva_comp = args.bvacomp

        if bva_comp:
            bva_files = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["charge"])
            for arg_name in bva_files:
                setattr(args, arg_name, bva_files[arg_name])

        dos_comparison = args.doscomp

        if dos_comparison:
            if "DOSCAR.LSO.lobster" in args.doscar.name:
                dos_files = get_file_paths(
                    path_to_lobster_calc=Path.cwd(), requested_files=["vasprun", "doscar"], use_lso_dos=True
                )
            else:
                dos_files = get_file_paths(
                    path_to_lobster_calc=Path.cwd(), requested_files=["vasprun", "doscar"], use_lso_dos=False
                )
            for arg_name in dos_files:
                setattr(args, arg_name, dos_files[arg_name])

        potcar_file_path = args.potcar

        quality_dict = Analysis.get_lobster_calc_quality_summary(
            path_to_poscar=args.structure,
            path_to_charge=args.charge,
            path_to_lobsterout=args.lobsterout,
            path_to_lobsterin=args.lobsterin,
            path_to_potcar=None if not potcar_file_path.exists() else potcar_file_path,
            potcar_symbols=args.potcarsymbols,
            path_to_bandoverlaps=args.bandoverlaps,
            dos_comparison=dos_comparison,
            bva_comp=bva_comp,
            path_to_doscar=args.doscar,
            e_range=args.erange,
            n_bins=args.nbins,
            path_to_vasprun=args.vasprun,
        )

        quality_text = Description.get_calc_quality_description(quality_dict)
        Description.write_calc_quality_description(quality_text)

        if args.file_calc_quality_json is not None:
            with open(args.file_calc_quality_json, "w") as fd:
                json.dump(quality_dict, fd)

    if args.action in ["plot-dos", "plotdos"]:
        if "DOSCAR.LSO.lobster" in args.doscar.name:
            req_files = get_file_paths(
                path_to_lobster_calc=Path.cwd(), requested_files=["structure", "doscar"], use_lso_dos=True
            )
        else:
            req_files = get_file_paths(
                path_to_lobster_calc=Path.cwd(), requested_files=["structure", "doscar"], use_lso_dos=False
            )

        for arg_name in req_files:
            setattr(args, arg_name, req_files[arg_name])

        lobs_dos = Doscar(doscar=args.doscar, structure_file=args.structure).completedos

        dos_plotter = PlainDosPlotter(summed=args.summedspins, sigma=args.sigma)
        if args.addtotal:
            dos_plotter.add_dos(dos=lobs_dos, label="Total DOS")
        if args.spddos:
            dos_plotter.add_dos_dict(dos_dict=lobs_dos.get_spd_dos())

        if args.elementdos:
            dos_plotter.add_dos_dict(dos_dict=lobs_dos.get_element_dos())

        if args.element:
            for element in args.element:
                element_spddos = lobs_dos.get_element_spd_dos(el=element)
                for orbital, dos in element_spddos.items():
                    label = f"{element}: {orbital.name}"
                    dos_plotter.add_dos_dict(dos_dict={label: dos})

        if args.site is not None and args.orbital:
            if len(args.site) > len(args.orbital):
                for site in args.site:
                    for orbital in args.orbital:
                        dos_plotter.add_site_orbital_dos(site_index=site, orbital=orbital, dos=lobs_dos)
            elif len(args.orbital) > len(args.site):
                for orbital in args.orbital:
                    for site in args.site:
                        dos_plotter.add_site_orbital_dos(site_index=site, orbital=orbital, dos=lobs_dos)
            else:
                for site, orbital in zip(args.site, args.orbital):
                    dos_plotter.add_site_orbital_dos(site_index=site, orbital=orbital, dos=lobs_dos)
        elif (
            args.site is None
            and not args.orbital
            and not args.element
            and not args.spddos
            and not args.elementdos
            and not args.addtotal
        ):
            dos_plotter.add_dos(dos=lobs_dos, label="Total DOS")
            dos_plotter.add_dos_dict(dos_dict=lobs_dos.get_element_dos())

        elif (args.site is None or not args.orbital) and (not args.element and not args.spddos and not args.elementdos):
            raise ValueError("Please set both args i.e site and orbital to generate the plot")

        plt = dos_plotter.get_plot(
            xlim=args.xlim,
            ylim=args.ylim,
            beta_dashed=True,
            invert_axes=args.invertaxis,
        )

        ax = plt.gca()
        ax.set_title(args.title)

        if not args.hideplot and not args.save_plot:
            plt.show()
        elif args.save_plot and not args.hideplot:
            fig = plt.gcf()
            fig.savefig(args.save_plot)
            plt.show()
        if args.save_plot and args.hideplot:
            plt.savefig(args.save_plot)

    if args.action in ["plot-icohp-distance", "ploticohpdistance"]:
        if args.cobis:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["icobilist"]).get("icobilist")
            args.icohplist = filename
            options = {"are_cobis": True, "are_coops": False}
        elif args.coops:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["icooplist"]).get("icooplist")
            args.icohplist = filename
            options = {"are_cobis": False, "are_coops": True}
        else:
            filename = get_file_paths(path_to_lobster_calc=Path.cwd(), requested_files=["icohplist"]).get("icohplist")
            args.icohplist = filename
            options = {"are_cobis": False, "are_coops": False}

        icohpcollection = Icohplist(filename=args.icohplist, **options).icohpcollection
        icohp_plotter = IcohpDistancePlotter(**options)

        icohp_plotter.add_icohps(icohpcollection=icohpcollection, label=args.legendprefix)

        get_plot_kwargs = {
            "alpha": args.alpha,
            "color_interactions": args.colorbonds,
            "colors": args.colors,
            "marker_size": args.markersize,
            "marker_style": args.markerstyle,
            "xlim": args.xlim,
            "ylim": args.ylim,
        }

        plt = icohp_plotter.get_plot(**get_plot_kwargs)

        ax = plt.gca()
        ax.set_title(args.title)

        if not args.hideplot and not args.save_plot:
            plt.show()
        elif args.save_plot and not args.hideplot:
            fig = plt.gcf()
            fig.savefig(args.save_plot)
            plt.show()
        if args.save_plot and args.hideplot:
            plt.savefig(args.save_plot)


if __name__ == "__main__":
    main()
