#!/usr/bin/env python
"""
Interactive tool to open and analyze files produced by Abinit and other atomistic codes.

This script supports a wide range of file formats, including NetCDF files (GSR, DDB, PHBST, etc.),
Abinit input and output files, structure files (CIF, POSCAR, XYZ), and more.

By default, the script starts an interactive IPython session with the file object
loaded as `abifile` or `obj`. This allows for direct interaction with the data
and execution of class methods.

Alternatively, it can generate Jupyter notebooks, serve interactive Dashboards
using the `panel` package, or automatically generate Matplotlib/Plotly figures.

Examples:
    Open a GSR file in an interactive shell:
        $ abiopen.py out_GSR.nc

    Print basic information about a structure file:
        $ abiopen.py crystal.cif -p

    Automatically generate and show plots for a phonon band structure:
        $ abiopen.py out_PHBST.nc -e

    Serve an interactive Dashboard for a DDB file:
        $ abiopen.py out_DDB -pn

    Generate a Jupyter-lab notebook to analyze a file:
        $ abiopen.py out_SIGRES.nc -nb
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from shutil import which

from monty.termcolor import cprint

import abipy.tools.cli_parsers as cli
from abipy import abilab
from abipy.tools.plotting import Exposer


def make_and_open_notebook(options):
    """
    Generate a Jupyter notebook and open it in the browser.

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

    Returns:
        int: System exit code (0 for success).

    Raises:
        RuntimeError: If `jupyter` is not found in the system PATH.
    """
    import os

    import nbformat

    nbf = nbformat.v4
    nb = nbf.new_notebook()

    nb.cells.extend(
        [
            nbf.new_markdown_cell("## This is an auto-generated notebook for %s" % os.path.relpath(options.filepath)),
            nbf.new_code_cell("""\

%matplotlib notebook
import numpy as np
#import seaborn as sns
#sns.set(context='notebook', style='darkgrid', palette='deep',
#        font='sans-serif', font_scale=1, color_codes=False, rc=None)
from abipy import abilab

"""),
            nbf.new_code_cell("abifile = abilab.abiopen('%s')" % options.filepath),
        ]
    )

    import tempfile

    _, nbpath = tempfile.mkstemp(prefix="abinb_", suffix=".ipynb", dir=os.getcwd(), text=True)

    with open(nbpath, "w", encoding="utf8") as f:
        nbformat.write(nb, f)

    if which("jupyter") is None:
        raise RuntimeError("Cannot find jupyter in PATH. Install it with `pip install`")

    if not options.classic_notebook:
        # Use jupyter-lab instead of classic notebook
        has_jupyterlab = which("jupyter-lab") is not None
        appname = "jupyter-lab" if has_jupyterlab else "jupyter notebook"
    else:
        appname = "jupyter notebook"

    if options.foreground:
        return os.system("%s %s" % (appname, nbpath))
    fd, tmpname = tempfile.mkstemp(text=True)
    print(tmpname)
    cmd = "%s %s" % (appname, nbpath)
    print("Executing:", cmd, "\nstdout and stderr redirected to %s" % tmpname)
    process = subprocess.Popen(cmd.split(), shell=False, stdout=fd, stderr=fd)
    cprint("pid: %s" % str(process.pid), "yellow")
    return 0


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

    abiopen.py FILE          => Open file in ipython shell.
    abiopen.py FILE -p       => Print info on object to terminal.
    abiopen.py FILE -e       => Generate matplotlib figures automatically.
                                Use -sns to activate seaborn settings.
    abiopen.py FILE -eweb    => Generate matplotlib figures, show them in the $BROWSER.
    abiopen.py FILE -ply     => Generate plotly figures automatically. Show them in the $BROWSER.
                                Note that not all FILEs support plotly.
    abiopen.py FILE -pn      => Generate GUI in web BROWSER to interact with FILE. Requires panel package.
    abiopen.py FILE -nb      => Generate jupyter-lab notebook.
    abiopen.py FILE -cnb     => Generate classic jupyter notebook.

where `FILE` is any file supported by abipy/pymatgen e.g. Netcdf files, Abinit input, POSCAR, xsf.
File extensions supported (including zipped files with extension in ".bz2", ".gz", ".z"):
Use `-v` to increase verbosity level (can be supplied multiple times e.g -vv).

JSON file are supported as well. In this case, abiopen.py tries to reconstruct python objects
assuming JSON document in MSONable format and then invokes ipython with the `data` object.
Use `-e` or `--notebook` or `--panel` to print the JSON dictionary without reconstructing python objects.
======================================================================================================

Table mapping file extension to AbiPy object:

"""
    return s + abilab.abiopen_ext2class_table()


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

    Args:
        with_epilog: If True, include the epilog in the parser.
    """
    parser = argparse.ArgumentParser(
        epilog=get_epilog() if with_epilog else "", formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument(
        "--loglevel",
        default="ERROR",
        type=str,
        help="Set the loglevel. Possible values: CRITICAL, ERROR (default), WARNING, INFO, DEBUG",
    )
    parser.add_argument("-V", "--version", action="version", version=abilab.__version__)

    parser.add_argument(
        "-v",
        "--verbose",
        default=0,
        action="count",  # -vv --> verbose=2
        help="verbose, can be supplied multiple times to increase verbosity",
    )

    parser.add_argument("filepath", help="File to open. See table below for the list of supported extensions.")

    # notebook options.
    parser.add_argument("-nb", "--notebook", action="store_true", default=False, help="Open file in jupyter notebook")
    parser.add_argument(
        "--classic-notebook",
        "-cnb",
        action="store_true",
        default=False,
        help="Use classic jupyter notebook instead of jupyterlab.",
    )
    parser.add_argument(
        "--no-browser",
        action="store_true",
        default=False,
        help=(
            "Start the jupyter server to serve the notebook "
            "but don't open the notebook in the browser.\n"
            "Use this option to connect remotely from localhost to the machine running the kernel"
        ),
    )
    parser.add_argument(
        "--foreground", action="store_true", default=False, help="Run jupyter notebook in the foreground."
    )

    # print option
    parser.add_argument("-p", "--print", action="store_true", default=False, help="Print python object and return.")

    # panel options
    parser.add_argument(
        "-pn",
        "--panel",
        action="store_true",
        default=False,
        help="Open Dashboard in web browser, requires panel package.",
    )
    parser.add_argument(
        "-pnt",
        "--panel-template",
        default="FastList",
        type=str,
        help="Specify template for panel dashboard."
        "Possible values are: FastList, FastGrid, Golden, Bootstrap, Material, React, Vanilla."
        "Default: FastList",
    )
    parser.add_argument("--port", default=0, type=int, help="Allows specifying a specific port when serving panel app.")
    # add_expose_options_to_parser(parser)

    # Expose option.
    parser.add_argument(
        "-e",
        "--expose",
        action="store_true",
        default=False,
        help="Open file and generate matplotlib figures automatically by calling expose method.",
    )
    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.",
    )
    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.",
    )
    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',
    )
    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."
        ),
    )
    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.",
    )
    parser.add_argument(
        "-ply",
        "--plotly",
        default=False,
        action="store_true",
        help="Generate plotly plots in $BROWSER instead of matplotlib. WARNING: Not all the features are supported.",
    )

    return parser


def serve_kwargs_from_options(options):
    """
    Build and return a dictionary with keyword arguments for serving the panel app.

    Args:
        options: Namespace object containing command-line options.
    """
    if options.no_browser:
        print("""
Use:

    ssh -N -f -L localhost:{port}:localhost:{port} username@your_remote_cluster

for port forwarding.
""")

    import abipy.panels as mod

    assets_path = os.path.join(os.path.dirname(mod.__file__), "assets")

    return dict(
        debug=options.verbose > 0,
        show=not options.no_browser,
        port=options.port,
        static_dirs={"/assets": assets_path},
        # address=address,
        # websocket_origin="{address}:{port}",
    )


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

    cli.set_loglevel(options.loglevel)

    ##############################################################################################
    # Handle meta options i.e. options that set other options.
    # OK, it's not very clean but I haven't find any parse API to express this kind of dependency.
    ##############################################################################################
    if options.plotly:
        options.expose = True
    if options.expose_web:
        options.expose = True
    if options.classic_notebook:
        options.notebook = True

    if options.verbose > 2:
        print(options)

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

        matplotlib.use(options.mpl_backend)

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

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

    if not os.path.exists(options.filepath):
        raise RuntimeError("%s: no such file" % options.filepath)

    if options.filepath.endswith(".json"):
        return handle_json(options)

    if options.filepath.endswith(".csv"):
        return handle_csv(options)

    if options.filepath.endswith(".traj"):
        return handle_ase_traj(options)

    if options.filepath.endswith("md.aselog"):
        return handle_ase_md_log(options)

    if options.filepath.endswith(".pickle") and not options.filepath.endswith("__AbinitFlow__.pickle"):
        # Handle pickle file.
        import pickle

        with open(options.filepath, "rb") as fh:
            obj = pickle.load(fh)

        return handle_object(obj, options)

    if os.path.basename(options.filepath) == "flows.db":
        from abipy.flowtk.launcher import print_flowsdb_file

        return print_flowsdb_file(options.filepath)

    if os.path.basename(options.filepath) == "phonopy_params.yaml":
        import phonopy

        phonon = phonopy.load("phonopy_params.yaml")
        # Start ipython shell with namespace
        import IPython

        IPython.embed(
            header="""
The Phonopy object is associated to the `phonon` python variable.
Use `phonon.<TAB>` to list available methods.

#phonon.set_irreps(
#  (0.0, 0.0, 0),
#  is_little_cogroup=False
#  nac_q_direction=False,
#  #nac_q_direction=settings.nac_q_direction,
#  degeneracy_tolerance=
#)

#phonon.show_irreps(show_irreps=True)
#phonon.write_yaml_irreps(show_irreps=True)
"""
        )
        return 0

    if not options.notebook:
        abifile = abilab.abiopen(options.filepath)
        return handle_object(abifile, options)

    # Call specialized method if the object is a NotebookWriter
    # else generate simple notebook by calling `make_and_open_notebook`
    cls = abilab.abifile_subclass_from_filename(options.filepath)
    if hasattr(cls, "make_and_open_notebook"):
        if hasattr(cls, "__exit__"):
            with abilab.abiopen(options.filepath) as abifile:
                return abifile.make_and_open_notebook(
                    foreground=options.foreground,
                    classic_notebook=options.classic_notebook,
                    no_browser=options.no_browser,
                )
        else:
            abifile = abilab.abiopen(options.filepath)
            return abifile.make_and_open_notebook(
                foreground=options.foreground, classic_notebook=options.classic_notebook, no_browser=options.no_browser
            )
    else:
        return make_and_open_notebook(options)

    return 0


def handle_object(obj, options):
    """
    Postprocess or visualize the object according to command-line options.

    Args:
        obj: The AbiPy object to handle.
        options: Namespace object containing command-line options.

    Returns:
        int: System exit code.
    """
    if options.print:
        # Print object to terminal.
        if hasattr(obj, "to_string"):
            print(obj.to_string(verbose=options.verbose))
        else:
            print(obj)
        return 0

    if options.expose:
        # Print info to terminal
        if hasattr(obj, "to_string"):
            print(obj.to_string(verbose=options.verbose))
        else:
            print(obj)

        # Generate plots automatically.
        if options.plotly:
            # plotly version
            if hasattr(obj, "plotly_expose"):
                obj.plotly_expose(verbose=options.verbose)
            else:
                cprint("`%s` does not implement plotly_expose method" % type(obj), "red")

        elif hasattr(obj, "expose"):
            # matplotlib version
            obj.expose(
                slide_mode=options.slide_mode,
                slide_timeout=options.slide_timeout,
                use_web=options.expose_web,
                verbose=options.verbose,
            )
        else:
            if not hasattr(obj, "yield_figs"):
                raise TypeError("Object of type `%s` does not implement (expose or yield_figs methods" % type(obj))
            from abipy.tools.plotting import MplExposer

            with MplExposer(
                slide_mode=options.slide_mode, slide_timeout=options.slide_timeout, verbose=options.verbose
            ) as e:
                e(obj.yield_figs())

    elif options.panel:
        if not hasattr(obj, "get_panel"):
            raise TypeError("Object of type `%s` does not implement get_panel method" % type(obj))

        import matplotlib

        matplotlib.use("Agg")
        pn = abilab.abipanel()

        app = obj.get_panel(template=options.panel_template)
        serve_kwargs = serve_kwargs_from_options(options)

        return pn.serve(app, **serve_kwargs)

    else:
        # 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 AbiPy object is associated to the `obj` python variable.
Use `obj.<TAB>` to list all available methods.
Use e.g. `obj.plot?` to access docstring and `obj.plot??` to visualize source code.
Use `print(obj)` to print the object.
"""
        )


def handle_ase_traj(obj, options):
    """Handle ASE trajectory file."""
    from abipy.ml.aseml import AseTrajectoryPlotter

    plotter = AseTrajectoryPlotter.from_file(options.filepath)

    print(plotter.to_string(verbose=options.verbose))
    if options.expose:
        if len(plotter.traj) > 1:
            plot_kws = dict(show=False)
            with Exposer.as_exposer("mpl") as e:
                e(plotter.plot(**plot_kws))
                e(plotter.plot_lattice(**plot_kws))

    return 0


def handle_ase_md_log(options):
    """Handle ASE MD log file."""
    from abipy.ml.aseml import AseMdLog

    md_log = AseMdLog(options.filepath)
    with Exposer.as_exposer("mpl") as e:
        e.add_obj_with_yield_figs(md_log)
    return 0


def handle_csv(options):
    """Handle CSV file."""
    import pandas as pd

    df = pd.read_csv(options.filepath)

    def print_df():
        print("=== Dataframe info ===")
        print(df.info())
        print("=== Dataframe describe ===")
        print(df.describe())

    if options.notebook:
        raise NotImplementedError("")
        # Visualize JSON document in jupyter
        # cmd = "jupyter-lab %s" % options.filepath
        # print("Executing:", cmd)
        # process = subprocess.Popen(cmd.split(), shell=False) #, stdout=fd, stderr=fd)
        # cprint("pid: %s" % str(process.pid), "yellow")
        return 0

    if options.panel:
        raise NotImplementedError("")
        # Visualize JSON document in panel dashboard.
        # pn = abilab.abipanel()
        # with abilab.abiopen(options.filepath) as json_file:
        #    app = json_file.get_panel()

        # serve_kwargs = serve_kwargs_from_options(options)
        # return pn.serve(app, **serve_kwargs)

    if options.print:
        # Print python object to terminal.
        print_df()
        return 0
    if options.expose:
        print_df()
        raise NotImplementedError("")
        return 0

    # Start ipython shell with namespace
    # Use embed because I don't know how to show a header with start_ipython.
    print_df()
    import IPython

    IPython.embed(
        header="""
The pandas DataFrame initialized from the csv file can be accessed via the `df` python variable.
"""
    )

    return 0


def handle_json(options):
    """
    Handle JSON files by loading them with MSONable support.

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

    Returns:
        int: System exit code.
    """
    if options.notebook:
        # Visualize JSON document in jupyter
        cmd = "jupyter-lab %s" % options.filepath
        print("Executing:", cmd)
        process = subprocess.Popen(cmd.split(), shell=False)  # , stdout=fd, stderr=fd)
        cprint("pid: %s" % str(process.pid), "yellow")
        return 0

    if options.panel:
        # Visualize JSON document in panel dashboard.
        pn = abilab.abipanel()
        with abilab.abiopen(options.filepath) as json_file:
            app = json_file.get_panel()

        serve_kwargs = serve_kwargs_from_options(options)
        return pn.serve(app, **serve_kwargs)

    obj = abilab.mjson_load(options.filepath)
    handle_object(obj, options)

    return 0


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