# Source code for pymatgen.io.abinit.abiinspect

    # coding: utf-8
    # Copyright (c) Pymatgen Development Team.
    # Distributed under the terms of the MIT License.
    """
    This module provides objects to inspect the status of the Abinit tasks at run-time.
    by extracting information from the main output file (text format).
    """

    import os
    import numpy as np
    import ruamel.yaml as yaml

    from collections.abc import Mapping, Iterable, Iterator
    from collections import OrderedDict
    from tabulate import tabulate
    from monty.functools import lazy_property
    from monty.collections import AttrDict
    from pymatgen.util.plotting import add_fig_kwargs, get_axarray_fig_plt

    def straceback():
        """Returns a string with the traceback."""
        import traceback
        return traceback.format_exc()

    def _magic_parser(stream, magic):
        """
        Parse the section with the SCF cycle

        Returns:
            dict where the key are the name of columns and
            the values are list of numbers. Note if no section was found.

        .. warning::

            The parser is very fragile and should be replaced by YAML.
        """
        # Example (SCF cycle, similar format is used for phonons):
        #
        #  iter   Etot(hartree)      deltaE(h)  residm     vres2
        #  ETOT  1  -8.8604027880849    -8.860E+00 2.458E-02 3.748E+00

        #  At SCF step    5       vres2   =  3.53E-08 < tolvrs=  1.00E-06 =>converged.
        in_doc, fields = 0, None

        for line in stream:
            line = line.strip()

            if line.startswith(magic):
                keys = line.split()
                fields = OrderedDict((k, []) for k in keys)

            if fields is not None:
                # print(line)
                in_doc += 1
                if in_doc == 1:
                    continue

                # End of the section.
                if not line:
                    break

                tokens = list(map(float, line.split()[1:]))
                assert len(tokens) == len(keys)
                for l, v in zip(fields.values(), tokens):
                    l.append(v)

        # Convert values to numpy arrays.
        if fields:
            return OrderedDict([(k, np.array(v)) for k, v in fields.items()])
        else:
            return None

    def plottable_from_outfile(filepath):
        """
        Factory function that returns a plottable object by inspecting the main output file of abinit
        Returns None if it is not able to detect the class to instantiate.
        """
        # TODO
        # Figure out how to detect the type of calculations
        # without having to parse the input. Possible approach: YAML doc
        # with YamlTokenizer(filepath) as r:
        #    doc = r.next_doc_with_tag("!CalculationType")
        #    d = yaml.safe_load(doc.text_notag)
        #    calc_type = d["calculation_type"]

        # ctype2class = {
        #    "Ground State": GroundStateScfCycle,
        #    "Phonon": PhononScfCycle,
        #    "Relaxation": Relaxation,
        # }
        # obj = ctype2class.get(calc_type, None)

        obj = GroundStateScfCycle
        if obj is not None:
            return obj.from_file(filepath)
        else:
            return None

    # Use log scale for these variables.
    _VARS_SUPPORTING_LOGSCALE = set(["residm", "vres2", "nres2"])

    # Hard-coded y-range for selected variables.
    _VARS_WITH_YRANGE = {
        "deltaE(h)": (-1e-3, +1e-3),
        "deltaE(Ha)": (-1e-3, +1e-3),
    }

    class ScfCycle(Mapping):
        """
        It essentially consists of a dictionary mapping string
        to list of floats containing the data at the different iterations.

        .. attributes::

            num_iterations: Number of iterations performed.
        """

        def __init__(self, fields):
            self.fields = fields
            all_lens = [len(lst) for lst in self.values()]
            self.num_iterations = all_lens[0]
            assert all(n == self.num_iterations for n in all_lens)

        def __getitem__(self, slice):
            return self.fields.__getitem__(slice)

        def __iter__(self):
            return self.fields.__iter__()

        def __len__(self):
            return len(self.fields)

        def __str__(self):
            return self.to_string()

        def to_string(self, verbose=0):
            """String representation."""
            rows = [[it + 1] + list(map(str, (self[k][it] for k in self.keys())))
                    for it in range(self.num_iterations)]

            return tabulate(rows, headers=["Iter"] + list(self.keys()))

        @property
        def last_iteration(self):
            """Returns a dictionary with the values of the last iteration."""
            return {k: v[-1] for k, v in self.items()}

        @classmethod
        def from_file(cls, filepath):
            """Read the first occurrence of ScfCycle from file."""
            with open(filepath, "rt") as stream:
                return cls.from_stream(stream)

        @classmethod
        def from_stream(cls, stream):
            """
            Read the first occurrence of ScfCycle from stream.

            Returns:
                None if no `ScfCycle` entry is found.
            """
            fields = _magic_parser(stream, magic=cls.MAGIC)

            if fields:
                fields.pop("iter")
                return cls(fields)
            else:
                return None

        @add_fig_kwargs
        def plot(self, ax_list=None, fontsize=12, **kwargs):
            """
            Uses matplotlib to plot the evolution of the SCF cycle.

            Args:
                ax_list: List of axes. If None a new figure is produced.
                fontsize: legend fontsize.
                kwargs: keyword arguments are passed to ax.plot

            Returns: matplotlib figure
            """
            # Build grid of plots.
            num_plots, ncols, nrows = len(self), 1, 1
            if num_plots > 1:
                ncols = 2
                nrows = num_plots // ncols + num_plots % ncols

            ax_list, fig, plot = get_axarray_fig_plt(ax_list, nrows=nrows, ncols=ncols,
                                                     sharex=True, sharey=False, squeeze=False)
            ax_list = np.array(ax_list).ravel()

            iter_num = np.array(list(range(self.num_iterations))) + 1
            label = kwargs.pop("label", None)

            for i, ((key, values), ax) in enumerate(zip(self.items(), ax_list)):
                ax.grid(True)
                ax.set_xlabel('Iteration Step')
                ax.set_xticks(iter_num, minor=False)
                ax.set_ylabel(key)

                xx, yy = iter_num, values
                if self.num_iterations > 1:
                    # Don't show the first iteration since it's not very useful.
                    xx, yy = xx[1:], values[1:]

                if not kwargs and label is None:
                    ax.plot(xx, yy, "-o", lw=2.0)
                else:
                    ax.plot(xx, yy, label=label if i == 0 else None, **kwargs)

                if key in _VARS_SUPPORTING_LOGSCALE and np.all(yy > 1e-22):
                    ax.set_yscale("log")

                if key in _VARS_WITH_YRANGE:
                    ymin, ymax = _VARS_WITH_YRANGE[key]
                    val_min, val_max = np.min(yy), np.max(yy)
                    if abs(val_max - val_min) > abs(ymax - ymin):
                        ax.set_ylim(ymin, ymax)

                if label is not None:
                    ax.legend(loc="best", fontsize=fontsize, shadow=True)

            # Get around a bug in matplotlib.
            if num_plots % ncols != 0:
                ax_list[-1].plot(xx, yy, lw=0.0)
                ax_list[-1].axis('off')

            return fig

    class GroundStateScfCycle(ScfCycle):
        """Result of the Ground State self-consistent cycle."""
        MAGIC = "iter   Etot(hartree)"

        @property
        def last_etotal(self):
            """The total energy at the last iteration."""
            return self["Etot(hartree)"][-1]

    class D2DEScfCycle(ScfCycle):
        """Result of the Phonon self-consistent cycle."""
        MAGIC = "iter   2DEtotal(Ha)"

        @property
        def last_etotal(self):
            """The 2-nd order derivative of the energy at the last iteration."""
            return self["2DEtotal(Ha)"][-1]

    class PhononScfCycle(D2DEScfCycle):
        """Iterations of the DFPT SCF cycle for phonons."""

    class CyclesPlotter:
        """Relies on the plot method of cycle objects to build multiple subfigures."""

        def __init__(self):
            self.labels = []
            self.cycles = []

        def items(self):
            """To iterate over (label, cycle)."""
            return zip(self.labels, self.cycles)

        def add_label_cycle(self, label, cycle):
            """Add new cycle to the plotter with label `label`."""
            self.labels.append(label)
            self.cycles.append(cycle)

        @add_fig_kwargs
        def combiplot(self, fontsize=8, **kwargs):
            """
            Compare multiple cycels on a grid: one subplot per quantity,
            all cycles on the same subplot.

            Args:
                fontsize: Legend fontsize.
            """
            ax_list = None
            for i, (label, cycle) in enumerate(self.items()):
                fig = cycle.plot(ax_list=ax_list, label=label, fontsize=fontsize,
                                 lw=2.0, marker="o", linestyle="-", show=False)
                ax_list = fig.axes

            return fig

        def slideshow(self, **kwargs):
            """
            Produce slides show of the different cycles. One plot per cycle.
            """
            for label, cycle in self.items():
                cycle.plot(title=label, tight_layout=True)

    class Relaxation(Iterable):
        """
        A list of :class:`GroundStateScfCycle` objects.

        .. attributes::

            num_iterations: Number of iterations performed.

        .. note::

            Forces, stresses  and crystal structures are missing.
            This object is mainly used to analyze the behavior of the Scf cycles
            during the structural relaxation. A more powerful and detailed analysis
            can be obtained by using the HIST.nc file.
        """

        def __init__(self, cycles):
            self.cycles = cycles
            self.num_iterations = len(self.cycles)

        def __iter__(self):
            return self.cycles.__iter__()

        def __len__(self):
            return self.cycles.__len__()

        def __getitem__(self, slice):
            return self.cycles[slice]

        def __str__(self):
            return self.to_string()

        def to_string(self, verbose=0):
            """String representation."""
            lines = []
            app = lines.append
            for i, cycle in enumerate(self):
                app("")
                app("RELAXATION STEP: %d" % (i + 1))
                app(cycle.to_string(verbose=verbose))

            return "\n".join(lines)

        @classmethod
        def from_file(cls, filepath):
            """Initialize the object from the Abinit main output file."""
            with open(filepath, "rt") as stream:
                return cls.from_stream(stream)

        @classmethod
        def from_stream(cls, stream):
            """
            Extract data from stream. Returns None if some error occurred.
            """
            cycles = []
            while True:
                scf_cycle = GroundStateScfCycle.from_stream(stream)
                if scf_cycle is None:
                    break
                cycles.append(scf_cycle)

            return cls(cycles) if cycles else None

        @lazy_property
        def history(self):
            """
            Ordered Dictionary of lists with the evolution of
            the data as function of the relaxation step.
            """
            history = OrderedDict()
            for cycle in self:
                d = cycle.last_iteration
                for k, v in d.items():
                    if k in history:
                        history[k].append(v)
                    else:
                        history[k] = [v]

            # Convert to numpy arrays.
            for k, v in history.items():
                history[k] = np.array(v)

            return history

        def slideshow(self, **kwargs):
            """
            Uses matplotlib to plot the evolution of the structural relaxation.

            Args:
                ax_list: List of axes. If None a new figure is produced.

            Returns:
                `matplotlib` figure
            """
            for i, cycle in enumerate(self.cycles):
                cycle.plot(title="Relaxation step %s" % (i + 1),
                           tight_layout=kwargs.pop("tight_layout", True),
                           show=kwargs.pop("show", True))

        @add_fig_kwargs
        def plot(self, ax_list=None, fontsize=12, **kwargs):
            """
            Plot relaxation history i.e. the results of the last iteration of each SCF cycle.

            Args:
                ax_list: List of axes. If None a new figure is produced.
                fontsize: legend fontsize.
                kwargs: keyword arguments are passed to ax.plot

            Returns: matplotlib figure
            """
            history = self.history

            # Build grid of plots.
            num_plots, ncols, nrows = len(history), 1, 1
            if num_plots > 1:
                ncols = 2
                nrows = num_plots // ncols + num_plots % ncols

            ax_list, fig, plot = get_axarray_fig_plt(ax_list, nrows=nrows, ncols=ncols,
                                                     sharex=True, sharey=False, squeeze=False)
            ax_list = np.array(ax_list).ravel()

            iter_num = np.array(list(range(self.num_iterations))) + 1
            label = kwargs.pop("label", None)

            for i, ((key, values), ax) in enumerate(zip(history.items(), ax_list)):
                ax.grid(True)
                ax.set_xlabel('Relaxation Step')
                ax.set_xticks(iter_num, minor=False)
                ax.set_ylabel(key)

                xx, yy = iter_num, values
                if not kwargs and label is None:
                    ax.plot(xx, yy, "-o", lw=2.0)
                else:
                    ax.plot(xx, yy, label=label if i == 0 else None, **kwargs)

                if key in _VARS_SUPPORTING_LOGSCALE and np.all(yy > 1e-22):
                    ax.set_yscale("log")

                if key in _VARS_WITH_YRANGE:
                    ymin, ymax = _VARS_WITH_YRANGE[key]
                    val_min, val_max = np.min(yy), np.max(yy)
                    if abs(val_max - val_min) > abs(ymax - ymin):
                        ax.set_ylim(ymin, ymax)

                if label is not None:
                    ax.legend(loc="best", fontsize=fontsize, shadow=True)

            # Get around a bug in matplotlib.
            if num_plots % ncols != 0:
                ax_list[-1].plot(xx, yy, lw=0.0)
                ax_list[-1].axis('off')

            return fig

    # TODO
    # class HaydockIterations(Iterable):
    #    """This object collects info on the different steps of the Haydock technique used in the Bethe-Salpeter code"""
    #    @classmethod
    #    def from_file(cls, filepath):
    #        """Initialize the object from file."""
    #        with open(filepath, "rt") as stream:
    #            return cls.from_stream(stream)
    #
    #    @classmethod
    #    def from_stream(cls, stream):
    #        """Extract data from stream. Returns None if some error occurred."""
    #        cycles = []
    #        while True:
    #            scf_cycle = GroundStateScfCycle.from_stream(stream)
    #            if scf_cycle is None: break
    #            cycles.append(scf_cycle)
    #
    #        return cls(cycles) if cycles else None
    #
    #    #def __init__(self):
    #
    #    def plot(self, **kwargs):
    #        """
    #        Uses matplotlib to plot the evolution of the structural relaxation.
    #        ==============  ==============================================================
    #        kwargs          Meaning
    #        ==============  ==============================================================
    #        title           Title of the plot (Default: None).
    #        how            True to show the figure (Default).
    #        savefig         'abc.png' or 'abc.eps'* to save the figure to a file.
    #        ==============  ==============================================================
    #        Returns:
    #            `matplotlib` figure
    #        """
    #        import matplotlib.pyplot as plt
    #        title = kwargs.pop("title", None)
    #        show = kwargs.pop("show", True)
    #        savefig = kwargs.pop("savefig", None)
    #        if title: fig.suptitle(title)
    #        if savefig is not None: fig.savefig(savefig)
    #        if show: plt.show()
    #        return fig

    ##################
    # Yaml parsers.
    ##################

    class YamlTokenizerError(Exception):
        """Exceptions raised by :class:`YamlTokenizer`."""

    class YamlTokenizer(Iterator):
        """
        Provides context-manager support so you can use it in a with statement.
        """
        Error = YamlTokenizerError

        def __init__(self, filename):
            # The position inside the file.
            self.linepos = 0
            self.filename = filename

            try:
                self.stream = open(filename, "rt")
            except IOError as exc:
                # Look for associated error file.
                root, ext = os.path.splitext(self.filename)
                errfile = root + ".err"
                if os.path.exists(errfile) and errfile != self.filename:
                    print("Found error file: %s" % errfile)
                    with open(errfile, "rt") as fh:
                        print(fh.read())
                raise exc

        def __iter__(self):
            return self

        def __enter__(self):
            return self

        def __exit__(self, type, value, traceback):
            self.close()

        def __del__(self):
            self.close()

        def close(self):
            try:
                self.stream.close()
            except Exception:
                print("Exception in YAMLTokenizer.close()")
                print("Python traceback:")
                print(straceback())

        def seek(self, offset, whence=0):
            """
            seek(offset[, whence]) -> None.  Move to new file position.

            Argument offset is a byte count.  Optional argument whence defaults to
            0 (offset from start of file, offset should be >= 0); other values are 1
            (move relative to current position, positive or negative), and 2 (move
            relative to end of file, usually negative, although many platforms allow
            seeking beyond the end of a file).  If the file is opened in text mode,
            only offsets returned by tell() are legal.  Use of other offsets causes
            undefined behavior.
            Note that not all file objects are seekable.
            """
            assert offset == 0
            self.linepos = 0
            return self.stream.seek(offset, whence)

        # Python 3 compatibility
        def __next__(self):
            return self.next()

        def next(self):
            """
            Returns the first YAML document in stream.

            .. warning::

                Assume that the YAML document are closed explicitely with the sentinel '...'
            """
            in_doc, lines, doc_tag = None, [], None

            for i, line in enumerate(self.stream):
                self.linepos += 1
                # print(i, line)

                if line.startswith("---"):
                    # Include only lines in the form:
                    #  "--- !tag"
                    #  "---"
                    # Other lines are spurious.
                    in_doc = False
                    l = line[3:].strip().lstrip()

                    if l.startswith("!"):
                        # "--- !tag"
                        doc_tag = l
                        in_doc = True
                    elif not l:
                        #  "---"
                        in_doc = True
                        doc_tag = None

                    if in_doc:
                        lineno = self.linepos

                if in_doc:
                    lines.append(line)

                if in_doc and line.startswith("..."):
                    return YamlDoc(text="".join(lines), lineno=lineno, tag=doc_tag)

            raise StopIteration("Cannot find next YAML document in %s" % self.filename)

        def all_yaml_docs(self):
            """
            Returns a list with all the YAML docs found in stream.
            Seek the stream before returning.

            .. warning::

                Assume that all the YAML docs (with the exception of the last one)
                are closed explicitely with the sentinel '...'
            """
            docs = [doc for doc in self]
            self.seek(0)
            return docs

        def next_doc_with_tag(self, doc_tag):
            """
            Returns the next document with the specified tag. Empty string is no doc is found.
            """
            while True:
                try:
                    doc = next(self)
                    if doc.tag == doc_tag:
                        return doc

                except StopIteration:
                    raise

        def all_docs_with_tag(self, doc_tag):
            """
            Returns all the documents with the specified tag.
            """
            docs = []

            while True:
                try:
                    doc = self.next_doc_with(doc_tag)
                    docs.append(doc)

                except StopIteration:
                    break

            self.seek(0)
            return docs

    def yaml_read_kpoints(filename, doc_tag="!Kpoints"):
        """Read the K-points from file."""
        with YamlTokenizer(filename) as r:
            doc = r.next_doc_with_tag(doc_tag)
            d = yaml.safe_load(doc.text_notag)

            return np.array(d["reduced_coordinates_of_qpoints"])

    def yaml_read_irred_perts(filename, doc_tag="!IrredPerts"):
        """Read the list of irreducible perturbations from file."""
        with YamlTokenizer(filename) as r:
            doc = r.next_doc_with_tag(doc_tag)
            d = yaml.safe_load(doc.text_notag)

            return [AttrDict(**pert) for pert in d["irred_perts"]]
            # return d["irred_perts"]

    class YamlDoc:
        """
        Handy object that stores that YAML document, its main tag and the
        position inside the file.
        """
        __slots__ = [
            "text",
            "lineno",
            "tag",
        ]

        def __init__(self, text, lineno, tag=None):
            """
            Args:
                text: String with the YAML document.
                lineno: The line number where the document is located.
                tag: The YAML tag associate to the document.
            """
            # Sanitize strings: use "ignore" to skip invalid characters in .encode/.decode like
            if isinstance(text, bytes):
                text = text.decode("utf-8", "ignore")
            text = text.rstrip().lstrip()
            self.text = text
            self.lineno = lineno
            if isinstance(tag, bytes):
                tag = tag.decode("utf-8", "ignore")
            self.tag = tag

        def __str__(self):
            return self.text

        def __eq__(self, other):
            if other is None:
                return False
            return (self.text == other.text and
                    self.lineno == other.lineno and
                    self.tag == other.tag)

        def __ne__(self, other):
            return not self == other

        @property
        def text_notag(self):
            """
            Returns the YAML text without the tag.
            Useful if we don't have any constructor registered for the tag
            (we used the tag just to locate the document).
            """
            if self.tag is not None:
                return self.text.replace(self.tag, "")
            else:
                return self.text

        def as_dict(self):
            """Use Yaml to parse the text (without the tag) and returns a dictionary."""
            return yaml.safe_load(self.text_notag)