#!/usr/bin/env python
"""
Command-line tool to visualize and analyze Abinit results.

This script provides specialized visualization tools for various Abinit output
formats (GSR, DDB, PHBST, HIST, DEN/POT, etc.). It supports generating plots
using Matplotlib or Plotly, and can export data to external applications like
VESTA, XCrySDen, and OVITO.

Main Categories:
    Structure: Visualize crystal structures and relaxation/MD trajectories.
    Electrons: Plot band structures, DOS, and Fermi surfaces.
    Phonons:   Analyze vibrational properties, IR spectra, and elastic tensors from DDBs.
    Fields:    Visualize volumetric data like electronic density and potentials.
    Timing:    Analyze Abinit timer data and memory logs.

Examples:
    Plot electron bands from a GSR file:
        $ abiview.py ebands out_GSR.nc

    Visualize a structural relaxation trajectory in HIST.nc:
        $ abiview.py hist out_HIST.nc

    Compute and plot phonon bands and DOS from a DDB file:
        $ abiview.py ddb out_DDB

    Visualize the electronic density (DEN.nc) in VESTA:
        $ abiview.py denpot out_DEN.nc -a vesta

    Analyze the Abinit timer from an output file:
        $ abiview.py timer out_ABO
"""

from __future__ import annotations

import argparse
import os
import sys
from pprint import pformat

from monty.termcolor import cprint

import abipy.tools.cli_parsers as cli
from abipy import abilab
from abipy.abio.outputs import AbinitOutputFile
from abipy.iotools.visualizer import Visualizer
from abipy.tools.plotting import GenericDataFilePlotter, MplExposer, PanelExposer

# , plotlyfigs_to_browser, push_to_chart_studio


def handle_overwrite(path, options):
    """Exit 1 if file ``path`` exists and not options.force else return path."""
    name_parts = os.path.splitext(path)
    print("Writing %s file:" % name_parts[-1].replace(".", "").upper())
    if os.path.exists(path) and not options.force:
        cprint("Cannot overwrite pre-existent file. Use `-f` options.", "red")
        sys.exit(1)
    return path


def df_to_clipboard(options, df) -> None:
    """Copy dataframe to clipboard if options.clipboard."""
    if getattr(options, "clipboard", False):
        cprint("Copying dataframe to the system clipboard.", "green")
        cprint("This can be pasted into Excel, for example", "green")
        df.to_clipboard()


class NegateAction(argparse.Action):
    """Argparse action to negate a boolean option."""
    def __call__(self, parser, ns, values, option):
        """Toggle the boolean value of the option."""
        setattr(ns, self.dest, option[2:4] != "no")


def abiview_structure(options):
    """
    Visualize a crystal structure with an external application.

    Args:
        options: Namespace object containing command-line options (must have `filepath` and `appname`).

    Returns:
        int: System exit code.
    """
    structure = abilab.Structure.from_file(options.filepath)
    print(structure.to_string(verbose=options.verbose))
    print("Visualizing structure with:", options.appname)
    structure.visualize(appname=options.appname)
    return 0


def abiview_input(options):
    """Read input file from netcdf file and print it to terminal."""
    from abipy.iotools import ETSF_Reader

    with ETSF_Reader(options.filepath) as r:
        input_str = r.read_string("input_string")
        print(input_str)
    return 0


def abiview_hist(options):
    """
    Visualize a structural relaxation or molecular dynamics run.

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

    Returns:
        int: System exit code.
    """
    with abilab.abiopen(options.filepath) as hist:
        print(hist.to_string(verbose=options.verbose))
        if options.appname is not None:
            hist.visualize(appname=options.appname, to_unit_cell=options.to_unit_cell)
        elif options.xdatcar:
            xpath = options.filepath + ".XDATCAR"
            handle_overwrite(xpath, options)
            hist.write_xdatcar(filepath=xpath, groupby_type=True, overwrite=True, to_unit_cell=options.to_unit_cell)
        else:
            hist.plot()

        return 0


def abiview_data(options):
    """
    Extract data from a generic text file with results
    in tabular format and plot data with matplotlib.
    Multiple datasets are supported.
    No attempt is made to handle metadata (e.g. column name)
    """
    plotter = GenericDataFilePlotter(options.filepath)
    print(plotter.to_string(verbose=options.verbose))
    plotter.plot(use_index=options.use_index)

    return 0


# def abiview_abo(options):
#    """
#    Plot self-consistent iterations extracted from Abinit output file
#    as well as timer data (if present)
#    """
#    with abilab.abiopen(options.filepath) as abo:
#        print(abo.to_string(verbose=options.verbose))
#        abo.plot()
#        #timer = abo.get_timer()
#    return 0


def abiview_timer(options):
    """Plot the ABINIT timer reported in the main output file. Requires timopt != 0."""
    key = "wall_time"
    nmax = 5
    from abipy.tools.plotting import MplExposer

    with AbinitOutputFile(options.filepath) as abo:
        timer = abo.get_timer()
        # print(timer.to_string(verbose=options.verbose))

        with MplExposer(
            slide_mode=options.slide_mode,
            slide_timeout=options.slide_timeout,
            expose_web=options.expose_web,
            verbose=options.verbose,
        ) as e:
            # e(timer.plot_efficiency(key=key, what="good+bad", nmax=nmax))
            e(timer.plot_pie(key=key, minfract=0.05, show=False))
            e(timer.plot_stacked_hist(key=key, nmax=nmax, show=False))

    return 0


def abiview_memlog(options):
    """Analyze the ABINIT log file to extract virtual memory info (requires Linux)."""
    from abipy.flowtk.events import MemLogParser

    parser = MemLogParser(options.filepath)
    print(parser.to_string(verbose=options.verbose))
    # return 0

    with MplExposer(
        slide_mode=options.slide_mode,
        slide_timeout=options.slide_timeout,
        expose_web=options.expose_web,
        verbose=options.verbose,
    ) as e:
        e(parser.plot(show=False))
        e(parser.plot_by_file(show=False))

    return 0


def abiview_dirviz(options) -> int:
    """Visualize directory tree with graphviz."""
    import tempfile

    from abipy.flowtk.utils import Dirviz

    graph = Dirviz(options.filepath).get_cluster_graph(engine=options.engine)
    directory = tempfile.mkdtemp()
    print("Producing source files in:", directory)
    graph.view(directory=directory)

    return 0


def abiview_ebands(options) -> int:
    """
    Plot electronic bands if file contains high-symmetry k-path or DOS if k-sampling.
    Accept any file with ElectronBands e.g. GSR.nc, WFK.nc, ...
    """
    with abilab.abiopen(options.filepath) as abifile:
        if options.xmgrace:
            outpath = options.filepath + ".agr"
            abifile.ebands.to_xmgrace(handle_overwrite(outpath, options))
        elif options.bxsf:
            outpath = options.filepath + ".bxsf"
            abifile.ebands.to_bxsf(handle_overwrite(outpath, options))
        # elif options.plotly:
        #    print(abifile.to_string(verbose=options.verbose))
        else:
            print(abifile.to_string(verbose=options.verbose))
            abifile.expose_ebands(
                slide_mode=options.slide_mode,
                slide_timeout=options.slide_timeout,
                expose_web=options.expose_web,
                verbose=options.verbose,
            )

        return 0


def abiview_effmass(options) -> int:
    """
    Compute electron effective masses with finite differences using energies along a k-path
    Accept any file with ElectronBands e.g. GSR.nc, WFK.nc, ...
    """
    from abipy.electrons.effmass_analyzer import EffMassAnalyzer

    emana = EffMassAnalyzer.from_file(options.filepath)
    print(emana)
    emana.select_band_edges()
    emana.summarize()
    emana.plot_emass()

    return 0


def abiview_skw(options) -> int:
    """
    Interpolate energies in k-space along a k-path with the star-function method.
    Note that the interpolation will likely fail if there are symmetric k-points in the input set of k-points
    so it is highly recommended to call this method with energies obtained in the IBZ.
    Accept any file with ElectronBands e.g. GSR.nc, WFK.nc, ...
    """
    ebands = abilab.ElectronBands.as_ebands(options.filepath)
    if not ebands.kpoints.is_ibz:
        cprint("SKW interpolator should be called with energies in the IBZ", "yellow")
    r = ebands.interpolate(lpratio=options.lpratio, line_density=options.line_density, verbose=options.verbose)
    r.ebands_kpath.plot()
    return 0


def abiview_fs(options) -> int:
    """
    Extract eigenvalues in the IBZ from file and visualize Fermi surface with
    the external application specified via --appname
    """
    with abilab.abiopen(options.filepath) as abifile:
        eb3d = abifile.ebands.get_ebands3d()

        if options.appname == "mpl":
            eb3d.plot_isosurfaces()
        elif options.appname == "xsf":
            eb3d.xcrysden_view()
        elif options.appname == "mayavi":
            eb3d.mvplot_isosurfaces()
        else:
            raise ValueError("Unsupported value for --appname: %s" % str(options.appname))

        return 0


def abiview_ifermi_fs(options) -> int:
    """
    Use ifermi package to visualize the Fermi surface. Requires netcdf file with energies in the IBZ.
    See <https://fermisurfaces.github.io/IFermi/>
    """
    with abilab.abiopen(options.filepath) as abifile:
        kwargs = dict(
            interpolation_factor=options.interpolation_factor,
            mu=options.mu,
            eref=options.eref,
            wigner_seitz=options.wigner,
            with_velocities=options.with_velocities,
        )
        print(abifile.ebands, "\n")
        print("Building Fermi surface with options:\n\n", pformat(kwargs, indent=4), end=2 * "\n")

        r = abifile.ebands.get_ifermi_fs(**kwargs)
        r.fs_plotter.get_plot(plot_type=options.plot_type).show()
        # abifile.ebands.get_ifermi_slices(**kwargs)

        return 0


def abiview_ddb(options) -> int:
    """Invoke anaddb to compute phonon bands and DOS from the DDB, plot the results."""
    with abilab.abiopen(options.filepath) as ddb:
        print(ddb.to_string(verbose=options.verbose))

        # Don't need PHDOS if phononwebsite
        nqsmall = 0 if options.phononwebsite else 10
        ndivsm = 20
        asr = 2
        chneut = 1
        dipdip = 1
        dos_method = "tetra"
        lo_to_splitting = "automatic"
        # chneut = 0
        print(f"""
Computing phonon bands and DOS from DDB file with:

    {nqsmall=}, {ndivsm=},
    {asr=}, {chneut=}, {dipdip=}, {lo_to_splitting=}, {dos_method=}
""")

        print("Invoking anaddb ...  ", end="")
        phbst_file, phdos_file = ddb.anaget_phbst_and_phdos_files(
            nqsmall=nqsmall,
            ndivsm=ndivsm,
            asr=asr,
            chneut=chneut,
            dipdip=dipdip,
            dos_method=dos_method,
            lo_to_splitting=lo_to_splitting,
            with_phonopy_obj=options.with_phonopy,
            verbose=options.verbose,
            mpi_procs=1,
        )
        print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

        phbands = phbst_file.phbands
        units = "mev"

        if options.xmgrace:
            outpath = options.filepath + ".agr"
            phbands.to_xmgrace(handle_overwrite(outpath, options))
            return 0

        # elif options.bxsf:
        #    outpath = options.filepath + ".bxsf"
        #    phbands.to_bxsf(handle_overwrite(outpath, options))
        #    return 0

        if options.phononwebsite:
            return phbands.view_phononwebsite(browser=options.browser, verbose=options.verbose)

        if options.plotly:
            # Plotly + Panel version.
            phdos = phdos_file.phdos
            with PanelExposer(title=f"Vibrational properties of {phdos_file.structure.formula}") as e:
                e(phbands.qpoints.plotly(show=False))
                e(phbands.plotly_with_phdos(phdos, units=units, show=False))
                e(phdos_file.plotly_pjdos_type(units=units, show=False))
                e(
                    phdos.plotly_harmonic_thermo(
                        tstart=5, tstop=300, num=50, units="eV", formula_units=1, quantities="all", show=False
                    )
                )
                e(phdos.plotly(units=units, show=False))
                e(phbands.plot_colored_matched(units=units, show=False))
                e(phbands.plot_fatbands(units=units, phdos_file=phdos_file, show=False))
                # push_to_chart_studio(figs) if options.chart_studio else plotlyfigs_to_browser(figs)

        else:
            phdos = phdos_file.phdos

            if not options.expose_web:
                # matplotlib figure and X-server.
                e = MplExposer(slide_mode=options.slide_mode, slide_timeout=options.slide_timeout)
            else:
                # panel version with matplotlib.
                e = PanelExposer(title=f"Vibrational properties of {phdos_file.structure.formula}")

            with e:
                e(phbands.qpoints.plot(show=False))
                e(phbands.plot_with_phdos(phdos, units=units, show=False))
                e(phdos_file.plot_pjdos_type(units=units, show=False))
                e(phbands.plot_colored_matched(units=units, show=False))
                e(
                    phdos.plot_harmonic_thermo(
                        tstart=5, tstop=300, num=50, units="eV", formula_units=1, quantities="all", show=False
                    )
                )
                e(phdos.plot(units=units, show=False))
                e(phbands.plot_fatbands(units=units, phdos_file=phdos_file, show=False))

        phbst_file.close()
        phdos_file.close()

        if options.with_phonopy:
            phonon = phbands.phonopy_obj
            # Save phonopy object in Yaml format.
            filename = "phonopy_params.yaml"
            print(f"Saving phonopy object to {filename}")
            phonon.save(filename="phonopy_params.yaml", settings={"force_constants": True})

    return 0


def abiview_ddb_becs(options) -> int:
    """Invoke anaddb to compute eps_inf and Born effective charges."""
    with abilab.abiopen(options.filepath) as ddb:
        print(ddb.to_string(verbose=options.verbose))

        chneut = 1
        print(f"""
Calling anaddb to compute the macroscopic electronic dielectric tensor (e_inf)
and the Born effective charges in Cartesian coordinates.

    chneut: {chneut}
""")
        r = ddb.anaget_epsinf_and_becs(chneut=chneut, verbose=options.verbose)
        print("epsilon_inf:\n", r.epsinf, end=2 * "\n")
        print(r.becs)
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_ddb_vs(options) -> int:
    """Compute speed of sound by fitting phonon frequencies along selected directions."""
    num_points = 20
    asr = 2
    chneut = 1
    dipdip = 1
    print(f"""
Computing phonon frequencies for linear least-squares with:

    {num_points=}, {asr=}, {chneut=}, {dipdip=}
""")
    print("Invoking anaddb ...  ")
    sv = abilab.SoundVelocity.from_ddb(
        options.filepath, num_points=num_points, asr=asr, chneut=chneut, dipdip=dipdip, verbose=options.verbose
    )
    # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    df = sv.get_dataframe()
    abilab.print_dataframe(df, title="Speed of sound for different directions:")
    df_to_clipboard(options, df)
    sv.plotly if options.plotly else sv.plot()

    return 0


def abiview_ddb_ir(options) -> int:
    """Compute infra-red spectrum from DDB. Plot results."""
    asr = 2
    chneut = 1
    dipdip = 1
    print(f"""
Computing phonon frequencies for infra-red spectrum with:

    {asr=}, {chneut=}, {dipdip=}
""")

    with abilab.abiopen(options.filepath) as ddb:
        tgen = ddb.anaget_dielectric_tensor_generator(asr=asr, chneut=chneut, dipdip=dipdip, verbose=options.verbose)
        print(tgen.to_string(verbose=options.verbose))

        gamma_ev = 1e-4
        print("Plotting dielectric tensor with constant phonon damping: %s (eV)" % gamma_ev)
        if options.plotly:
            tgen.plotly_all(gamma_ev=gamma_ev)
        else:
            tgen.plot_all(gamma_ev=gamma_ev)
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_ddb_asr(options) -> int:
    """Compute phonon band structure from DDB with/without acoustic sum rule. Plot results."""
    print("Computing phonon frequencies with/without ASR")

    with abilab.abiopen(options.filepath) as ddb:
        plotter = ddb.anacompare_asr(
            asr_list=(0, 2),
            chneut_list=(1,),
            dipdip=1,
            lo_to_splitting="automatic",
            nqsmall=10,
            ndivsm=20,
            dos_method="tetra",
            ngqpt=None,
            verbose=0,
            mpi_procs=1,
        )

        title = ddb.structure.formula
        renderer = "browser" if not options.chart_studio else "chart_studio"
        plotter.combiplotly(renderer=renderer, title=title) if options.plotly else plotter.plot(title=title)
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_ddb_dipdip(options) -> int:
    """Compute phonon band structure from DDB with/without dipole-dipole interaction. Plot results."""
    print("Computing phonon frequencies with/without dipdip.")

    with abilab.abiopen(options.filepath) as ddb:
        plotter = ddb.anacompare_dipdip(
            chneut_list=(1,),
            asr=2,
            lo_to_splitting="automatic",
            nqsmall=10,
            ndivsm=20,
            dos_method="tetra",
            ngqpt=None,
            verbose=0,
            mpi_procs=1,
        )

        title = ddb.structure.formula
        renderer = "browser" if not options.chart_studio else "chart_studio"
        plotter.combiplotly(renderer=renderer, title=title) if options.plotly else plotter.plot(title=title)
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_ddb_quad(options) -> int:
    """Compute phonon band structure from DDB with/without quadrupole terms. Plot results."""
    print("Computing phonon frequencies with/without dip-quad and quad-quad terms.")

    with abilab.abiopen(options.filepath) as ddb:
        plotter = ddb.anacompare_phbands_with_quad(
            asr=2,
            chneut=1,
            dipdip=1,
            lo_to_splitting="automatic",
            nqsmall=0,
            ndivsm=20,
            dos_method="tetra",
            ngqpt=None,
            verbose=0,
            mpi_procs=1,
        )

        title = ddb.structure.formula
        renderer = "browser" if not options.chart_studio else "chart_studio"
        plotter.combiplotly(renderer=renderer, title=title) if options.plotly else plotter.plot(title=title)
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_ddb_isodistort_ph(options) -> int:
    """
    Compute ph-freqs for given q-point (default: Gamma), produce CIF files for unperturbed and distorded structure
    that can be used with ISODISTORT (https://stokes.byu.edu/iso/isodistort.php) to analyze the symmetry of phonon modes.
    See README.me file produced in output directory.
    """
    with abilab.abiopen(options.filepath) as ddb:
        print(ddb.to_string(verbose=options.verbose))
        qpoint = options.qpoint
        asr = 2
        chneut = 1
        dipdip = 1
        lo_to_splitting = False
        print(f"""
Computing phonon frequencies and eigenvectors with:

    {asr=}, {chneut=}, {dipdip=}, {lo_to_splitting=}
    {qpoint=}
""")

        print("Invoking anaddb ...  ", end="")
        phbands = ddb.anaget_phmodes_at_qpoint(
            qpoint=qpoint, asr=asr, chneut=chneut, dipdip=dipdip, verbose=options.verbose, lo_to_splitting=False
        )

        phbands.make_isodistort_ph_dir(qpoint, select_modes=None, eta=1, workdir=None)
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_ddb_qpt(options) -> int:
    """
    Compute ph-frequencies at the selected q-point without passing through the
    Fourier interpolation of the interatomic force constants.
    """
    with abilab.abiopen(options.filepath) as ddb:
        # Execute anaddb to compute phbands without Fourier interpolation.
        phbands, inp = ddb.anaget_phmodes_at_qpoint(
            qpoint=options.qpoint, asr=2, chneut=1, dipdip=1, ifcflag=0, return_input=True
        )
        print(inp)
        df = phbands.get_dataframe()
        abilab.print_dataframe(df, title="Phonon frequencies:")
        df_to_clipboard(options, df)
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_ddb_ifc(options) -> int:
    """Visualize interatomic force constants in real space."""
    asr = 2
    chneut = 1
    dipdip = 1
    print(f"""
Computing interatomic force constants with:

    {asr=}, {chneut=}, {dipdip=}
""")

    with abilab.abiopen(options.filepath) as ddb:
        # Execute anaddb to compute the interatomic force constants.
        ifc = ddb.anaget_ifc(asr=asr, chneut=chneut, dipdip=dipdip)
        # print(ifc)

        if not options.expose_web:
            # matplotlib figure and X-server.
            e = MplExposer(slide_mode=options.slide_mode, slide_timeout=options.slide_timeout)
        else:
            # panel version
            e = PanelExposer(title=f"Interatomic force constants of {ddb.structure.formula}")

        with e:
            e(ifc.plot_longitudinal_ifc(title="Longitudinal IFCs", show=False))
            e(ifc.plot_longitudinal_ifc_short_range(title="Longitudinal IFCs short range", show=False))
            e(ifc.plot_longitudinal_ifc_ewald(title="Longitudinal IFCs Ewald", show=False))
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_ddb_elastic(options) -> int:
    """Invoke anaddb to compute elastic and piezoelectric tensors. Require DDB with strain terms."""
    with abilab.abiopen(options.filepath) as ddb:
        print(ddb.to_string(verbose=options.verbose))

        asr = 2
        chneut = 1
        relaxed_ion = "automatic"
        piezo = "automatic"
        dde = False
        stress_correction = False

        print(f"""
Calling anaddb to compute elastic and piezoelectric tensors. Require DDB with strain terms.

    {asr=}, {chneut=}
    {relaxed_ion=}, {piezo=}, {dde=}, {stress_correction=}
""")

        el_data = ddb.anaget_elastic(
            relaxed_ion=relaxed_ion,
            piezo=piezo,
            dde=dde,
            stress_correction=stress_correction,
            asr=asr,
            chneut=chneut,
            verbose=options.verbose,
            retpath=False,
            return_input=False,
        )
        print(el_data)
        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

    return 0


def abiview_phbands(options) -> int:
    """Plot phonon bands. Accept any file with PhononBands e.g. PHBST.nc, ..."""
    with abilab.abiopen(options.filepath) as abifile:
        if options.xmgrace:
            outpath = options.filepath + ".agr"
            abifile.phbands.to_xmgrace(handle_overwrite(outpath, options))
        # elif options.bxsf:
        #    outpath = options.filepath + ".bxsf"
        #    abifile.phbands.to_bxsf(handle_overwrite(outpath, options))
        #    return 0
        elif options.phononwebsite:
            return abifile.phbands.view_phononwebsite(browser=options.browser)
        else:
            print(abifile.to_string(verbose=options.verbose))
            abifile.expose_phbands(
                slide_mode=options.slide_mode, slide_timeout=options.slide_timeout, verbose=options.verbose, units="mev"
            )

        # print("Calculation completed.\nResults available in:", os.path.dirname(phbst_file.filepath))

        return 0


def abiview_denpot(options) -> int:
    """
    Visualize netcdf DEN/POT files with --appname (default: Vesta).
    NB: Application must be installed and in $PATH.
    """
    with abilab.abiopen(options.filepath) as abifile:
        print(abifile.to_string(verbose=options.verbose))
        field = abifile.field
        if options.chgcar:
            if not field.is_density_like:
                cprint("Warning: expecting DENSITY file but received %s" % field.__class__.__name__, "yellow")
            chg_path = options.filepath + ".CHGCAR"
            field.to_chgcar(filename=handle_overwrite(chg_path, options))

        elif options.cube:
            outpath = options.filepath + ".cube"
            field.export_to_cube(filename=handle_overwrite(outpath, options), spin="total")

        else:
            abifile.field.visualize(appname=options.appname)
    return 0


def abiview_lobster(options) -> int:
    """Analyze lobster output files in directory."""
    from abipy.electrons.lobster import LobsterAnalyzer

    lobana = LobsterAnalyzer.from_dir(os.path.dirname(options.filepath), prefix=options.prefix)
    print(lobana.to_string(verbose=options.verbose))

    if options.ipython:
        # Start ipython shell with namespace
        # Use embed because I don't know how to show a header with start_ipython.
        import IPython

        IPython.embed(header="The LobsterAnalyzer is bound to the `lobana` variable.\nTry `print(lobana)`")

    elif options.notebook:
        return lobana.make_and_open_notebook(foreground=options.foreground)

    else:
        lobana.plot()
        # lobana.plot_coxp_with_dos(from_site_index=[0, 1])

    return 0


def abiview_xrd_traj(options) -> int:
    """Compare XRD spectra using the first and the last structure read from a trajectory file."""
    from abipy.core.structure import get_first_and_last_structure_from_file

    structures = get_first_and_last_structure_from_file(options.filepath)

    dfs = abilab.dataframes_from_structures(structures, index=["first", "last"])
    abilab.print_dataframe(dfs.lattice, title="Lattice parameters:")
    if options.verbose:
        abilab.print_dataframe(dfs.coords, title="Atomic positions (columns give the site index):")

    from pymatgen.analysis.diffraction.xrd import XRDCalculator

    xrd = XRDCalculator(wavelength=options.wavelength, symprec=options.symprec)
    two_theta_range = tuple(float(t) for t in options.two_theta_range)
    xrd.plot_structures(
        structures,
        two_theta_range=two_theta_range,
        fontsize=6,
        annotate_peaks=not options.no_annotate_peaks,
        tight_layout=True,
    )
    return 0


def get_epilog() -> str:
    """
    Return the epilog string for the command-line parser.
    """
    return """\
Usage example:

###########
# Structure
###########

    abiview.py structure FILE                ==> Visualize structure with Vesta (default)
    abiview.py structure FILE -a xcrysden    ==> Visualize structure with Xcrysden (default)
    abiview.py hist out_HIST.nc              ==> Plot relaxation/molecular dynamics results with matplotlib.
    abiview.py hist out_HIST.nc -a ovito     ==> Visualize relaxation/molecular dynamics results with ovito.
    abiview.py hist out_HIST.nc --xdatcar    ==> Convert HIST.nc into XDATCAR format (caveat: assume fixed unit cell!)

############
# Text files
############

    abiview.py data FILE     ==> Parse text FILE with data in tabular format and plot arrays.

###########
# Electrons
###########

    abiview.py ebands out_WFK.nc              ==>  Plot electrons bands (or DOS) with matplotlib.
    abiview.py ebands out_GSR.nc --xmgrace    ==>  Generate xmgrace file with electron bands.
    abiview.py effmass out_GSR.nc
    abiview.py fs FILE_WITH_KMESH.nc          ==>  Visualize Fermi surface from netcdf file with electron energies
                                                   on a k-mesh. Use `-a xsf` to change application e.g. Xcrysden.
    abiview.py skw out_GSR.nc                 ==> Interpolate IBZ energies with star-functions and plot interpolated bands.
    abiview.py denpot out_DEN.nc              ==> Visualize DEN/POT file with Vesta. Use `-a xcrysden` to change app.
    abiview.py denpot out_DEN.nc --chgcar     ==> Convert DEN file to CHGCAR fileformat.

#########
# Phonons
#########

    abiview.py ddb in_DDB                 ==>  Compute ph-bands and DOS from DDB, plot results.
    abiview.py ddb_becs in_DDB            ==>  Ccompute eps_inf and Born effective charges. Print results.
    abiview.py ddb_vs in_DDB              ==>  Compute speed of sound from DDB by fitting phonon frequencies.
    abiview.py ddb_ir in_DDB              ==>  Compute infra-red spectrum from DDB. Plot results.
    abiview.py ddb_asr in_DDB             ==>  Compute ph-bands from DDB with/wo acoustic rule.
                                               Plot results with matplotlib (default) or plotly (--plotly)
    abiview.py ddb_dipdip in_DDB          ==>  Compute ph-bands from DDB with/wo dipole-dipole treatment.
                                               Plot results with matplotlib (default) or plotly (--plotly)
    abiview.py ddb_quad in_DDB            ==>  Compute ph-bands from DDB with/wo dipole-quadrupole terms.
                                               Plot results with matplotlib (default) or plotly (--plotly)
    abiview.py ddb_qpt in_DDB -q 0 0.5 0  ==>  Compute ph-frequencies at the selected q-point without passing
                                               through the Fourier interpolation of the interatomic force constants.
    abiview.py ddb_ifc in_DDB             ==>  Visualize interatomic force constants in real space.
    abiview.py phbands out_PHBST.nc -web  ==>  Visualize ph-bands and displacements with phononwebsite.

#########
# Elastic
#########
    abiview.py ddb_elastic in_DDB        ==> Compute elastic and piezoelectric tensors.
                                             Require DDB with strain terms.

###############
# Miscelleanous
###############

  abiview.py dirviz DIRECTORY            ==> Visualize directory tree with graphviz.
  abiview.py lobster DIRECTORY           ==> Visualize Lobster results.

Use `abiview.py --help` for help and `abiview.py COMMAND --help` to get the documentation for `COMMAND`.
Use `-v` to increase verbosity level (can be supplied multiple times e.g -vv).
"""


def get_parser(with_epilog=False):
    """
    Return the ArgumentParser object for the script.

    Args:
        with_epilog: If True, include the epilog in the parser.
    """
    # Parent parser for common options.
    copts_parser = argparse.ArgumentParser(add_help=False)
    copts_parser.add_argument("filepath", type=str, help="File to visualize.")

    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(
        "-v",
        "--verbose",
        default=0,
        action="count",  # -vv --> verbose=2
        help="verbose, can be supplied multiple times to increase verbosity.",
    )
    copts_parser.add_argument(
        "-sns",
        "--seaborn",
        const="paper",
        default=None,
        action="store",
        nargs="?",
        type=str,
        help='Use seaborn settings. Accept value defining context in ("paper", "notebook", "talk", "poster"). Default: paper',
    )
    copts_parser.add_argument(
        "-mpl",
        "--mpl-backend",
        default=None,
        help=(
            "Set matplotlib interactive backend. "
            "Possible values: GTKAgg, GTK3Agg, GTK, GTKCairo, GTK3Cairo, WXAgg, WX, TkAgg, Qt4Agg, Qt5Agg, macosx."
            "See also: https://matplotlib.org/faq/usage_faq.html#what-is-a-backend."
        ),
    )

    # Parent parser for commands supporting MplExposer.
    slide_parser = argparse.ArgumentParser(add_help=False)
    slide_parser.add_argument(
        "-s",
        "--slide-mode",
        default=False,
        action="store_true",
        help="Iterate over figures. Expose all figures at once if not given on the CLI.",
    )
    slide_parser.add_argument(
        "-t",
        "--slide-timeout",
        type=int,
        default=None,
        help="Close figure after slide-timeout seconds (only if slide-mode). Block if not specified.",
    )
    slide_parser.add_argument(
        "-ew",
        "--expose-web",
        default=False,
        action="store_true",
        help="Generate matplotlib plots in $BROWSER instead of X-server. WARNING: Not all the features are supported.",
    )

    # Parent parser for commands that operating on pandas dataframes
    pandas_parser = argparse.ArgumentParser(add_help=False)
    pandas_parser.add_argument(
        "-c",
        "--clipboard",
        default=False,
        action="store_true",
        help="Copy dataframe to the system clipboard. This can be pasted into Excel, for example",
    )

    # Parent parser for commands supporting ipython
    ipy_parser = argparse.ArgumentParser(add_help=False)
    ipy_parser.add_argument("-ipy", "--ipython", default=False, action="store_true", help="Invoke ipython terminal.")

    # Parent parser for commands supporting (jupyter notebooks)
    nb_parser = argparse.ArgumentParser(add_help=False)
    nb_parser.add_argument("-nb", "--notebook", default=False, action="store_true", help="Generate jupyter notebook.")
    nb_parser.add_argument(
        "--foreground", action="store_true", default=False, help="Run jupyter notebook in the foreground."
    )

    # Parent parser for commands supporting plotly plots
    plotly_parser = argparse.ArgumentParser(add_help=False)
    plotly_parser.add_argument(
        "-ply",
        "--plotly",
        default=False,
        action="store_true",
        help="Generate plotly plots in the browser instead of matplotlib.",
    )
    plotly_parser.add_argument(
        "-cs",
        "--chart-studio",
        default=False,
        action="store_true",
        help="Push figure to plotly chart studio. "
        "Requires --plotly and user account at https://chart-studio.plotly.com.",
    )

    # Build the main parser.
    parser = argparse.ArgumentParser(
        epilog=get_epilog() if with_epilog else "", formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("-V", "--version", action="version", version=abilab.__version__)

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

    def add_args(p, *args):
        """Add arguments to subparser `p`."""
        for arg in args:
            if arg == "xmgrace":
                p.add_argument(
                    "--xmgrace",
                    default=False,
                    action="store_true",
                    help="Print bands in xmgrace format to stdout and exit.",
                )
            elif arg == "bxsf":
                p.add_argument(
                    "--bxsf",
                    default=False,
                    action="store_true",
                    help=(
                        "Generate BXSF file suitable for the visualization of isosurfaces with Xcrysden"
                        "(xcrysden --bxsf FILE).\n Requires k-points in IBZ. Print to stdout and exit."
                    ),
                )
            elif arg == "phononweb":
                p.add_argument(
                    "-web",
                    "--phononwebsite",
                    default=False,
                    action="store_true",
                    help=(
                        "Visualize phonon band structure on the phononwebsite. "
                        "http://henriquemiranda.github.io/phononwebsite/"
                    ),
                )
            elif arg == "browser":
                p.add_argument(
                    "-b",
                    "--browser",
                    default=None,
                    help="Define browser used by python webbrowser. "
                    "See https://docs.python.org/2/library/webbrowser.html#webbrowser.register",
                )
            elif arg == "force":
                p.add_argument(
                    "-f",
                    "--force",
                    default=False,
                    action="store_true",
                    help="Overwrite pre-existent files without prompting for confirmation.",
                )
            else:
                raise ValueError("Invalid arg: %s" % arg)

    # Subparser for structure command.
    p_structure = subparsers.add_parser("structure", parents=[copts_parser], help=abiview_structure.__doc__)
    p_structure.add_argument(
        "-a",
        "--appname",
        nargs="?",
        type=str,
        default="vesta",
        help=(
            "Application name. Default: vesta. "
            "Possible options: %s, mayavi, vtk" % ", ".join(Visualizer.all_visunames())
        ),
    )

    # Subparser for input command.
    p_input = subparsers.add_parser("input", parents=[copts_parser], help=abiview_input.__doc__)

    # Subparser for hist command.
    p_hist = subparsers.add_parser("hist", parents=[copts_parser], help=abiview_hist.__doc__)
    p_hist.add_argument(
        "-a",
        "--appname",
        nargs="?",
        default=None,
        const="ovito",
        help=(
            "Application name. Default: ovito. "
            "Possible options: `%s`, `mpl` (matplotlib) `mayavi`, `vtk`" % ", ".join(Visualizer.all_visunames())
        ),
    )
    p_hist.add_argument("--xdatcar", default=False, action="store_true", help="Convert HIST file into XDATCAR format.")
    p_hist.add_argument(
        "--to-unit-cell", default=False, action="store_true", help="Whether to translate sites into the unit cell."
    )
    add_args(p_hist, "force")

    # Subparser for data command.
    p_data = subparsers.add_parser("data", parents=[copts_parser], help=abiview_data.__doc__)
    p_data.add_argument(
        "-i",
        "--use-index",
        default=False,
        action="store_true",
        help="Use the row index as x-value in the plot. By default the plotter uses the first column as x-values",
    )

    # Subparser for abo command.
    # p_abo = subparsers.add_parser('abo', parents=[copts_parser], help=abiview_abo.__doc__)

    p_timer = subparsers.add_parser("timer", parents=[copts_parser, slide_parser], help=abiview_timer.__doc__)

    p_memlog = subparsers.add_parser("memlog", parents=[copts_parser, slide_parser], help=abiview_memlog.__doc__)

    # Subparser for dirviz command.
    p_dirviz = subparsers.add_parser("dirviz", parents=[copts_parser], help=abiview_dirviz.__doc__)
    p_dirviz.add_argument(
        "-e",
        "--engine",
        type=str,
        default="fdp",
        help=(
            "graphviz engine: ['dot', 'neato', 'twopi', 'circo', 'fdp', 'sfdp', 'patchwork', 'osage']. "
            "See http://www.graphviz.org/pdf/dot.1.pdf "
            "Use `conda install python-graphviz` or `pip install graphviz` to install the python package."
        ),
    )

    # Subparser for ebands command.
    p_ebands = subparsers.add_parser("ebands", parents=[copts_parser, slide_parser], help=abiview_ebands.__doc__)
    add_args(p_ebands, "xmgrace", "bxsf", "force")

    # Subparser for effmass command.
    p_effmass = subparsers.add_parser("effmass", parents=[copts_parser, slide_parser], help=abiview_effmass.__doc__)
    # add_args(p_ebands, "xmgrace", "bxsf", "force")

    # Subparser for skw command.
    p_skw = subparsers.add_parser("skw", parents=[copts_parser], help=abiview_skw.__doc__)
    p_skw.add_argument(
        "-lp",
        "--lpratio",
        type=int,
        default=5,
        help=(
            "Ratio between the number of star functions and the number of ab-initio k-points. "
            "The default should be OK in many systems, larger values may be required for accurate derivatives."
        ),
    )
    p_skw.add_argument(
        "-ld", "--line-density", type=int, default=20, help="Number of points in the smallest segment of the k-path."
    )

    # Subparser for fs command.
    p_fs = subparsers.add_parser("fs", parents=[copts_parser], help=abiview_fs.__doc__)
    p_fs.add_argument(
        "-a",
        "--appname",
        type=str,
        default="mpl",
        help="Application name. Possible options: mpl (matplotlib, default), xsf (xcrysden), mayavi.",
    )

    # Parent parser for ifermi commands
    ifermi_parser = argparse.ArgumentParser(add_help=False)
    ifermi_parser.add_argument(
        "-i",
        "--interpolation-factor",
        default=8,
        type=float,
        help="interpolation factor for band structure  [default: 8.0]",
    )
    ifermi_parser.add_argument(
        "--eref",
        default="fermie",
        type=str,
        choices=["fermie", "cbm", "vbm"],
        help="Energy reference for isosurface. `eref` and `mu` define the energy level: isoe = eref + mu"
        "Use `fermie` for metals, `cbm` for the conduction band minimum, and `vbm` for the valence band maximum, "
        "[default: `fermie`]",
    )
    ifermi_parser.add_argument(
        "-t",
        "--plot-type",
        default="plotly",
        choices=["plotly", "matplotlib", "mayavi"],
        help="Plot type. Possible options: plotly (default), matplotlib, mayavi.",
    )
    ifermi_parser.add_argument(
        "--wigner",
        "--no-wigner",
        dest="wigner",
        default=True,
        action=NegateAction,
        nargs=0,
        help="Use the Wigner-Seitz cell or the reciprocal lattice parallelepiped. Default is: wigner",
    )
    ifermi_parser.add_argument(
        "-vel",
        "--with-velocities",
        default=False,
        action="store_true",
        help="Show color map with interpolated group velocities. Default is False.",
    )

    # Subparser for ifermi_fs command.
    p_ifermi_fs = subparsers.add_parser(
        "ifermi_fs", parents=[ifermi_parser, copts_parser], help=abiview_ifermi_fs.__doc__
    )
    p_ifermi_fs.add_argument(
        "-m",
        "--mu",
        default=0.0,
        type=float,
        help="Offset in eV from the energy reference, eref, at which to calculate the isosurface. Default 0 i.e. use `eref`",
    )

    # Subparser for ddb command.
    p_ddb = subparsers.add_parser("ddb", parents=[copts_parser, slide_parser, plotly_parser], help=abiview_ddb.__doc__)
    p_ddb.add_argument(
        "--with-phonopy",
        default=False,
        action="store_true",
        help="Produce phonopy.yaml file. Use e.g. `abiopen.py  phonopy.yaml` to load the phonon object in the ipython terminal.",
    )
    add_args(p_ddb, "xmgrace", "phononweb", "browser", "force")

    # Subparser for ddb_vs command.
    p_ddb_vs = subparsers.add_parser(
        "ddb_vs", parents=[copts_parser, pandas_parser, slide_parser], help=abiview_ddb_vs.__doc__
    )

    # Subparser for ddb_ir command.
    p_ddb_becs = subparsers.add_parser("ddb_becs", parents=[copts_parser, pandas_parser], help=abiview_ddb_becs.__doc__)

    # Subparser for ddb_ir command.
    p_ddb_ir = subparsers.add_parser(
        "ddb_ir", parents=[copts_parser, pandas_parser, slide_parser], help=abiview_ddb_ir.__doc__
    )

    # Subparser for ddb_asr command.
    p_ddb_asr = subparsers.add_parser(
        "ddb_asr", parents=[copts_parser, pandas_parser, slide_parser, plotly_parser], help=abiview_ddb_asr.__doc__
    )

    # Subparser for ddb_dipdip command.
    p_ddb_dipdip = subparsers.add_parser(
        "ddb_dipdip",
        parents=[copts_parser, pandas_parser, slide_parser, plotly_parser],
        help=abiview_ddb_dipdip.__doc__,
    )

    # Subparser for ddb_quad command.
    p_ddb_quad = subparsers.add_parser(
        "ddb_quad", parents=[copts_parser, pandas_parser, slide_parser, plotly_parser], help=abiview_ddb_quad.__doc__
    )

    # Subparser for ddb_ph_isodistort command.
    p_ddb_isodistort_ph = subparsers.add_parser(
        "ddb_isodistort_ph", parents=[copts_parser], help=abiview_ddb_isodistort_ph.__doc__
    )
    p_ddb_isodistort_ph.add_argument(
        "-q",
        "--qpoint",
        nargs=3,
        type=float,
        help="q-point in reduced coordinates e.g. 0.25 0 0. Default: 0, 0, 0",
        default=[0, 0, 0],
    )

    # Subparser for ddb_ifc command.
    p_ddb_ifc = subparsers.add_parser(
        "ddb_ifc", parents=[copts_parser, pandas_parser, slide_parser], help=abiview_ddb_ifc.__doc__
    )

    # Subparser for ddb_ifc command.
    p_ddb_elastic = subparsers.add_parser(
        "ddb_elastic", parents=[copts_parser, pandas_parser, slide_parser], help=abiview_ddb_elastic.__doc__
    )

    # Subparser for ddb_qpt command.
    p_ddb_qpt = subparsers.add_parser(
        "ddb_qpt", parents=[copts_parser, pandas_parser, slide_parser], help=abiview_ddb_qpt.__doc__
    )

    p_ddb_qpt.add_argument(
        "-q",
        "--qpoint",
        nargs=3,
        type=float,
        help="q-point in reduced coordinates e.g. `-q 0.25 0 0`. Default: 0, 0, 0",
        default=[0, 0, 0],
    )

    # Subparser for phbands command.
    p_phbands = subparsers.add_parser("phbands", parents=[copts_parser, slide_parser], help=abiview_phbands.__doc__)
    add_args(p_phbands, "xmgrace", "phononweb", "browser", "force")

    # Subparser for lobster command.
    p_lobster = subparsers.add_parser(
        "lobster", parents=[copts_parser, ipy_parser, nb_parser], help=abiview_lobster.__doc__
    )
    p_lobster.add_argument("--prefix", type=str, default="", help="Prefix for lobster output files. Default: ''")

    # Subparser for xrd.
    p_xrd = subparsers.add_parser(
        "xrd_traj",
        parents=[copts_parser],
        help="Compare X-ray diffraction for the first and the last structure in a trajectory file.",
    )
    p_xrd.add_argument(
        "-w",
        "--wavelength",
        default="CuKa",
        type=str,
        help=(
            "The wavelength can be specified as a string. It must be one of the "
            "supported definitions in the WAVELENGTHS dict declared in pymatgen/analysis/diffraction/xrd.py."
            "Defaults to 'CuKa', i.e, Cu K_alpha radiation."
        ),
    )
    p_xrd.add_argument(
        "-s",
        "--symprec",
        default=0,
        type=float,
        help=(
            "Symmetry precision for structure refinement. "
            "If set to 0, no refinement is done. Otherwise, refinement is performed using spglib with provided precision."
        ),
    )
    p_xrd.add_argument(
        "-t",
        "--two-theta-range",
        default=(0, 90),
        nargs=2,
        help=("Tuple for range of two_thetas to calculate in degrees. Defaults to (0, 90)."),
    )
    p_xrd.add_argument(
        "-nap",
        "--no-annotate-peaks",
        default=False,
        action="store_true",
        help="Whether to annotate the peaks with plane information.",
    )

    # Subparser for denpot command.
    p_denpot = subparsers.add_parser("denpot", parents=[copts_parser], help=abiview_denpot.__doc__)
    p_denpot.add_argument(
        "-a",
        "--appname",
        type=str,
        default="vesta",
        help=(
            "Application name. Default: vesta. "
            + "Possible options: `%s`, `mayavi`, `vtk`" % ", ".join(Visualizer.all_visunames())
        ),
    )
    p_denpot.add_argument("--chgcar", default=False, action="store_true", help="Convert Density to CHGCAR format.")
    p_denpot.add_argument(
        "--cube", default=False, action="store_true", help="Convert Density/Potential to CUBE format."
    )

    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)

    if getattr(options, "plotly", False):
        options.expose = True
    if options.verbose > 2:
        print(options)

    if options.mpl_backend is not None:
        # Set matplotlib backend
        import matplotlib

        matplotlib.use(options.mpl_backend)

    if options.seaborn:
        # Use seaborn settings.
        import seaborn as sns

        sns.set(
            context=options.seaborn,
            style="darkgrid",
            palette="deep",
            font="sans-serif",
            font_scale=1,
            color_codes=False,
            rc=None,
        )

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


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