"""
Objects to analyze the screening files in netcdf format produced by the GW code (optdriver 3).
"""

from __future__ import annotations

from functools import cached_property

import numpy as np
import pymatgen.core.units as pmgu
from monty.bisect import index as bs_index
from monty.collections import AttrDict
from monty.inspect import all_subclasses
from monty.string import marquee
from monty.termcolor import cprint

from abipy.core.func1d import Function1D
from abipy.core.gsphere import GSphere
from abipy.core.kpoints import KpointList
from abipy.core.mixins import AbinitNcFile, Has_ElectronBands, Has_Header, Has_Structure, NotebookWriter
from abipy.core.structure import Structure
from abipy.electrons.ebands import ElectronBands
from abipy.iotools import ETSF_Reader
from abipy.tools import duck
from abipy.tools.plotting import (
    ArrayPlotter,
    add_fig_kwargs,
    data_from_cplx_mode,
    get_ax_fig_plt,
    get_axarray_fig_plt,
    set_axlims,
)
from abipy.tools.typing import Figure

_COLOR_CMODE = dict(re="red", im="blue", abs="black", angle="green")


def _latex_symbol_cplxmode(symbol: str, cplx_mode: str) -> str:
    """Latex label to be used to plot ``symbol`` in ``cplx_mode``."""
    return {
        "re": r"$\Re(" + symbol + ")$",
        "im": r"$\Im(" + symbol + ")$",
        "abs": r"$||" + symbol + "||$",
        "angle": r"$Phase(" + symbol + ")$",
    }[cplx_mode]


class ScrFile(AbinitNcFile, Has_Header, Has_Structure, Has_ElectronBands, NotebookWriter):
    """
    This object provides an interface to the ``SCR.nc`` file produced by the GW code.

    Usage example:

    .. code-block:: python

        with ScrFile("out_SCR.nc") as scr:
            print(scr)
            scr.plot_emacro()

    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: ScrFile
    """

    @classmethod
    def from_file(cls, filepath: str) -> ScrFile:
        """Initialize the object from a Netcdf file"""
        return cls(filepath)

    def __init__(self, filepath: str):
        """Initialize the object from a file path."""
        super().__init__(filepath)
        self.r = ScrReader(filepath)

    def close(self) -> None:
        """Close the file."""
        self.r.close()

    def __str__(self) -> str:
        return self.to_string()

    def to_string(self, verbose=0) -> str:
        """String representation."""
        lines = []
        app = lines.append

        app(marquee("File Info", mark="="))
        app(self.filestat(as_string=True))
        app("")
        app(self.structure.to_string(verbose=verbose, title="Structure"))
        app("")
        # TODO: Fix problem with efermi
        # app(self.ebands.to_string(with_structure=False, title="Electronic Bands"))
        app(self.kpoints.to_string(verbose=verbose, title="K-points for screening function"))
        app("")
        app("Number of G-vectors in screening matrices: %d" % self.ng)
        app("Number of frequencies: %d (real: %d, imaginary: %d)" % (self.nw, self.nrew, self.nimw))

        if verbose:
            app(str(self.params))

        if verbose > 1:
            app("")
            app(self.hdr.to_string(verbose=verbose, title="Abinit Header"))

        return "\n".join(lines)

    @property
    def structure(self) -> Structure:
        """|Structure| object."""
        return self.r.structure

    @property
    def kpoints(self):
        """List of k-points in the dielectric function."""
        return self.r.kpoints

    @cached_property
    def ebands(self) -> ElectronBands:
        """
        |ElectronBands| object with the single-particle energies used to compute the screening.
        """
        ebands = ElectronBands.from_file(self.filepath)
        # FIXME
        cprint("Setting Fermi energy to zero since `fermie_energy` is not initialized in Abinit v8.2", "yellow")
        ebands.fermie = 0
        return ebands

    @property
    def ng(self) -> int:
        """Number of G-vectors in screening matrices."""
        return self.r.ng

    @property
    def wpoints(self) -> np.ndarray:
        """
        Array of complex numbers with the frequencies of the dielectric function in Hartree.
        """
        return self.r.wpoints

    @property
    def nw(self) -> int:
        """Total number of frequencies."""
        return self.r.nw

    @property
    def nrew(self) -> int:
        """Number of real frequencies."""
        return self.r.nrew

    @property
    def nimw(self) -> int:
        """Number of imaginary frequencies."""
        return self.r.nimw

    @property
    def netcdf_name(self) -> str:
        """The netcdf_ name associated to the data on disk."""
        return self.r.netcdf_name

    @cached_property
    def params(self) -> dict:
        """
        Dict with the most important parameters used to compute the screening
        keys can be accessed with the dot notation i.e. ``params.zcut``.
        """
        # od = self.get_ebands_params()
        return self.r.read_params()

    @add_fig_kwargs
    def plot_emacro(self, cplx_mode="re-im", ax=None, xlims=None, fontsize=8, **kwargs) -> Figure:
        r"""
        Plot the macroscopic dielectric function including local-field effects.

        Args:
            cplx_mode: string defining the data to print.
                Possible choices are (case-insensitive):
                "re" for the real part. "im" for the imaginary part.
                "abs" for the absolute value.
                "angle" will display the phase of the complex number in radians.
                Options can be concatenated with ``-`` e.g. ``re-im``.
            xlims: Set the data limits for the x-axis in eV. Accept tuple e.g. (left, right)
                or scalar e.g. left. If left (right) is None, default values are used.
            ax: |matplotlib-Axes| or None if a new figure should be created.
            fontsize: Legend and title fontsize.

        Returns:
            |matplotlib-Figure|
        """
        emlf = self.r.read_emacro_lf()
        xx, yy = emlf.mesh * pmgu.Ha_to_eV, emlf.values

        ax, fig, plt = get_ax_fig_plt(ax=ax)
        for c in cplx_mode.lower().split("-"):
            ax.plot(
                xx,
                data_from_cplx_mode(c, yy),
                color=_COLOR_CMODE[c],
                linewidth=kwargs.get("linewidth", 2),
                linestyle=kwargs.get("linestyle", "solid"),
                label=_latex_symbol_cplxmode(r"\varepsilon_{M}", c),
            )

        set_axlims(ax, xlims, "x")
        ax.grid(True)
        ax.set_xlabel(r"$\omega$ (eV)")
        ax.legend(loc="best", shadow=True, fontsize=fontsize)

        return fig

    @add_fig_kwargs
    def plot_eelf(self, ax=None, xlims=None, fontsize=8, **kwargs) -> Figure:
        r"""
        Plot electron energy loss function.

        Args:
            ax: |matplotlib-Axes| or None if a new figure should be created.
            xlims: Set the data limits for the x-axis in eV. Accept tuple e.g. ``(left, right)``
                   or scalar e.g. ``left``. If left (right) is None, default values are used
            fontsize: Legend and label fontsize:
        """
        eelf = self.r.read_eelf()
        xx, yy = eelf.mesh * pmgu.Ha_to_eV, eelf.values

        ax, fig, plt = get_ax_fig_plt(ax=ax)
        ax.plot(xx, yy, linewidth=kwargs.get("linewidth", 2), linestyle=kwargs.get("linestyle", "solid"), label="EELF")

        set_axlims(ax, xlims, "x")
        ax.grid(True)
        ax.set_xlabel(r"$\omega$ (eV)")
        ax.legend(loc="best", fontsize=fontsize, shadow=True)

        return fig

    def plot_freqs_at_ggp(
        self, gvec1, gvec2=None, waxis="real", cplx_modes=("re", "im"), colormap="viridis", fontsize=8, **kwargs
    ) -> Figure:
        """
        Plot the real and the imaginary part of the response function as a function of omega along
        eithe the real or the imaginary axis. Each subplot, shows all the q-points stored on file

        Args:
            gvec1, gvec2: G-vectors.
            waxis: ``real`` to plot along the real axis, ``imag`` for the imaginary axis.
            cplx_modes:
            colormap: Color map. Have a look at the colormaps here and decide which one you like:
                http://matplotlib.sourceforge.net/examples/pylab_examples/show_colormaps.html
        """
        ax_list, fig, plt = get_axarray_fig_plt(
            None, nrows=len(cplx_modes), ncols=1, sharex=True, sharey=False, squeeze=False
        )
        ax_list = ax_list.ravel()
        cmap = plt.get_cmap(colormap)

        nqpt = len(self.r.sorted_kpoints)
        for iq, qpoint in enumerate(self.r.sorted_kpoints):
            wggmat = self.r.read_wggmat(qpoint)
            for ax, cplx_mode in zip(ax_list, cplx_modes, strict=False):
                wggmat.plot_freq(
                    gvec1,
                    gvec2=gvec2,
                    waxis=waxis,
                    color=cmap(float(iq) / nqpt),
                    label=f"{qpoint}",
                    cplx_mode=cplx_mode,
                    ax=ax,
                    show=False,
                )

        return fig

    def yield_figs(self, **kwargs):  # pragma: no cover
        """
        This function *generates* a predefined list of matplotlib figures with minimal input from the user.
        Used in abiview.py to get a quick look at the results.
        """
        edos = self.ebands.get_edos()
        yield self.ebands.plot_with_edos(edos, show=False)
        # Plot spectra if there are enough frequencies.
        if self.nrew > 2:
            yield self.plot_emacro(show=False)
            yield self.plot_eelf(show=False)

    def write_notebook(self, nbpath=None) -> str:
        """
        Write a jupyter_ notebook to nbpath. If ``nbpath`` is None, a temporay file in the current
        working directory is created. Return path to the notebook.
        """
        nbformat, nbv, nb = self.get_nbformat_nbv_nb(title=None)

        nb.cells.extend(
            [
                nbv.new_code_cell("ncfile = abilab.abiopen('%s')" % self.filepath),
                nbv.new_code_cell("print(ncfile)"),
                nbv.new_code_cell("print(ncfile.params)"),
                # nbv.new_code_cell("ncfile.ebands.plot();"),
                nbv.new_code_cell("edos = ncfile.ebands.get_edos()\nncfile.ebands.plot_with_edos(edos);"),
            ]
        )

        if self.nrew > 2:
            # Plot optical properties and EELF
            nb.cells.extend(
                [
                    nbv.new_code_cell("ncfile.plot_emacro();"),
                    nbv.new_code_cell("ncfile.plot_eelf();"),
                ]
            )

        return self._write_nb_nbpath(nb, nbpath)


class SusFile(ScrFile):
    """File object for SUS.nc files."""


class ScrReader(ETSF_Reader):
    """Class to read screening files."""
    """
    This object reads the results stored in the SCR (Screening) file produced by ABINIT.
    It provides helper functions to access the most important quantities.

    double inverse_dielectric_function(number_of_qpoints_dielectric_function,
    number_of_frequencies_dielectric_function, number_of_spins, number_of_spins,
    number_of_coefficients_dielectric_function, number_of_coefficients_dielectric_function, complex)

    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: ScrReader
    """

    def __init__(self, filepath: str):
        """
        Args:
            filepath: Path to the netcdf file.
        """
        super().__init__(filepath)

        # Read and store important quantities.
        self.structure = self.read_structure()
        qfrac_coords = self.read_value("qpoints_dielectric_function")
        self.kpoints = KpointList(self.structure.reciprocal_lattice, qfrac_coords)
        self.wpoints = self.read_value("frequencies_dielectric_function", cmode="c")
        self.ng = self.read_dimvalue("number_of_coefficients_dielectric_function")

        # Find number of real/imaginary frequencies.
        self.nw = len(self.wpoints)
        self.nrew = self.nw
        self.nimw = 0
        for i, w in enumerate(self.wpoints):
            if np.iscomplex(w):
                self.nrew = i
                break

        self.nimw = self.nw - self.nrew
        if self.nimw and not np.all(np.iscomplex(self.wpoints[self.nrew + 1 :])):
            raise ValueError(
                "wpoints should contained real points packed in the first positions\n"
                "followed by imaginary points but got: %s" % str(self.wpoints)
            )

        # Define self.netcdf_name from the data available on file.
        nfound = 0
        netcdf_names = ["polarizability", "dielectric_function", "inverse_dielectric_function"]
        for try_name in netcdf_names:
            if try_name in self.rootgrp.variables:
                self.netcdf_name = try_name
                nfound += 1

        if nfound == 0:
            raise RuntimeError("Cannot find `%s` in netcdf file" % str(netcdf_names))
        if nfound > 1:
            raise RuntimeError("Find multiple netcdf arrays (%s) in netcdf file!" % str(netcdf_names))

    @cached_property
    def sorted_kpoints(self) -> KpointList:
        """
        Build list of k-points sorted by norm (ascending order).
        """
        norms = [(kpt.norm, kpt) for kpt in self.kpoints]
        sorted_frac_coords = [t[1].frac_coords for t in sorted(norms, key=lambda t: t[0], reverse=False)]
        return KpointList(self.structure.reciprocal_lattice, sorted_frac_coords)

    def read_params(self) -> AttrDict:
        """
        Read the most important parameters used to compute the screening i.e.
        the parameters that may be subject to convergence studies.

        Returns:
            |AttrDict| a dictionary whose keys can be accessed with the dot notation i.e. ``d.key``.
        """
        # TODO: ecuteps is missing!
        keys = [
            "ikxc",
            "inclvkb",
            "gwcalctyp",
            "nbands_used",
            "npwwfn_used",
            "spmeth",
            "test_type",
            "tordering",
            "awtr",
            "icutcoul",
            "gwcomp",
            "gwgamma",
            "mbpt_sciss",
            "spsmear",
            "zcut",
            "gwencomp",
        ]

        def convert(arr):
            """Convert to scalar if size == 1"""
            # return np.asscalar(arr) if arr.size == 1 else arr
            return np.asarray(arr).item() if arr.size == 1 else arr

        return AttrDict({k: convert(self.read_value(k)) for k in keys})

    def read_emacro_lf(self, kpoint=(0, 0, 0)) -> Function1D:
        """
        Read the macroscopic dielectric function *with* local field effects 1 / em1_{0,0)(kpoint, omega).
        """
        if self.netcdf_name == "inverse_dielectric_function":
            em1 = self.read_wslice(kpoint, ig1=0, ig2=0)
            emacro = 1 / em1[: self.nrew]
        else:
            raise NotImplementedError("emacro_lf with netcdf != InverseDielectricFunction")

        return Function1D(np.real(self.wpoints[: self.nrew]).copy(), emacro)

    def read_emacro_nlf(self, kpoint=(0, 0, 0)) -> Function1D:
        """
        Read the macroscopic dielectric function *without* local field effects e_{0,0)(kpoint, omega).

        .. warning::

            This function performs the inversion of e-1 to get e.
            that can be quite expensive and memory demanding for large matrices!
        """
        if self.netcdf_name == "inverse_dielectric_function":
            em1 = self.read_wggmat(kpoint)
            e = np.linalg.inv(em1.wggmat[: self.nrew, :, :])
        else:
            raise NotImplementedError("emacro_nlf with netcdf != InverseDielectricFunction")

        return Function1D(np.real(self.wpoints[: self.nrew]).copy(), e[:, 0, 0])

    def read_eelf(self, kpoint=(0, 0, 0)) -> Function1D:
        """
        Read electron energy loss function: - Im(1/ emacro)
        """
        # eelf = -Im(1 / eM)
        emacro_lf = self.read_emacro_lf(kpoint=kpoint)
        # emacro_lf = self.read_emacro_nlf(kpoint=kpoint)
        values = (-1 / emacro_lf.values).imag

        return Function1D(emacro_lf.mesh.copy(), values)

    def read_wggmat(self, kpoint, spin1=0, spin2=0, cls=None):
        """
        Read data at the given k-point and return an instance of ``cls`` where
        ``cls`` is a subclass of :class:`_AwggMatrix`
        """
        cls = _AwggMatrix.class_from_netcdf_name(self.netcdf_name) if cls is None else cls

        var = self.rootgrp.variables["reduced_coordinates_plane_waves_dielectric_function"]
        # Use ik = 0 because the basis set is not k-dependent.
        ik0 = 0
        gvecs = var[ik0, :]
        # print("gvecs", gvecs)

        kpoint, ik = self.find_kpoint_fileindex(kpoint)

        # FIXME ecuteps is missing
        # TODO: Gpshere.find is very slow if we don't take advantage of shells
        ecuteps = 2
        gsphere = GSphere(ecuteps, self.structure.reciprocal_lattice, kpoint, gvecs)

        # Exchange spin due to F --> C
        values = self.rootgrp.variables[self.netcdf_name][ik, :, spin2, spin1, :, :, :]
        wggmat = values[:, :, :, 0] + 1j * values[:, :, :, 1]

        return cls(self.wpoints, gsphere, wggmat, inord="F")

    def find_kpoint_fileindex(self, kpoint):
        """
        Returns the k-point and the index of the k-point in the netcdf file.
        Accepts |Kpoint| instance or integer.
        """
        if duck.is_intlike(kpoint):
            ik = int(kpoint)
        else:
            ik = self.kpoints.index(kpoint)

        return self.kpoints[ik], ik

    def read_wslice(self, kpoint, ig1=0, ig2=0, spin1=0, spin2=0):
        """Read slice along the frequency dimension."""
        kpoint, ik = self.find_kpoint_fileindex(kpoint)
        var = self.rootgrp.variables[self.netcdf_name]
        values = var[ik, :, spin2, spin1, ig2, ig1, :]  # Exchange G indices F --> C

        return values[:, 0] + 1j * values[:, 1]


class _AwggMatrix:
    r"""
    Base class for two-point functions expressed in reciprocal space
    i.e. a complex matrix :math:`A_{G,G'}(\omega)` where G, G' are reciprocal
    lattice vectors defined inside the G-sphere.

    This class is not supposed to be instantiated directly.

    .. attributes:

        wpoints:
        gpshere:
        nrew:
        nwim:
    """

    netcdf_name = "_AwggMatrix"

    latex_name = "Unknown"

    @classmethod
    def class_from_netcdf_name(cls, netcdf_name: str):
        """Return the subclass associated to the given netcdf name."""
        nfound = 0
        for subclass in all_subclasses(cls):
            if subclass.netcdf_name == netcdf_name:
                nfound += 1
                out_cls = subclass

        if nfound == 0:
            raise ValueError(f"Cannot find subclass associated to {netcdf_name=}")
        if nfound > 1:
            raise ValueError(f"Find multiple subclasses associated to {netcdf_name=}")

        return out_cls

    def __init__(self, wpoints, gsphere, wggmat, inord="C"):
        """
        Initialize the matrix with frequencies, G-sphere and data.

        Args:
            gsphere: |GSphere| with G-vectors and k-point object.
            wpoints: Complex frequency points in Hartree.
            wggmat: [nw, ng, ng] complex array.
            inord: storage order of ``wggmat``. If inord == "F", ``wggmat`` is in
                in Fortran column-major order. Default: "C" i.e. C row-major order.
        """
        self.wpoints = np.array(wpoints, dtype=complex)
        self.gsphere = gsphere
        self.wggmat = np.reshape(wggmat, (self.nw, self.ng, self.ng))

        if inord.lower() == "f":
            # Fortran to C.
            for iw, _ in enumerate(wpoints):
                self.wggmat[iw] = self.wggmat[iw].T.copy()

        for i in (1, 2):
            if len(gsphere) != wggmat.shape[-i]:
                raise ValueError(f"{len(gsphere)=} != {wggmat.shape[-i]=}")

        if len(self.wpoints) != len(self.wggmat):
            raise ValueError(f"{len(self.wpoints)=} != {len(self.wggmat)=}")

        # Find number of real/imaginary frequencies.
        self.nrew = self.nw
        self.nimw = 0
        for i, w in enumerate(self.wpoints):
            if np.iscomplex(w):
                self.nrew = i
                break

        self.nimw = self.nw - self.nrew
        if self.nimw and not np.all(np.iscomplex(self.wpoints[self.nrew + 1 :])):
            raise ValueError(
                "wpoints should contained real points packed in the first positions\n"
                "followed by imaginary points but got: %s" % str(self.wpoints)
            )

    def __str__(self) -> str:
        return self.to_string()

    def to_string(self, verbose: int = 0) -> str:
        """String representation."""
        lines = []
        app = lines.append

        app(marquee(self.netcdf_name, mark="="))
        app("  K-point: %s" % self.kpoint)
        app("  Number of G-vectors: %d" % self.ng)
        app("  Total number of frequencies: %d (real: %s, imaginary: %s)" % (self.nw, self.nrew, self.nimw))
        if self.nrew:
            app("  Real frequencies up to %.2f (eV)" % self.real_wpoints[-1].real)
        if self.nimw:
            app("  Imaginary frequencies up to %.2f (eV)" % self.imag_wpoints[-1].imag)

        return "\n".join(lines)

    @property
    def kpoint(self):
        """|Kpoint| object."""
        return self.gsphere.kpoint

    @property
    def ng(self) -> int:
        """Number of G-vectors."""
        return len(self.gsphere)

    @property
    def nw(self) -> int:
        """Total number of frequencies (real + imag)."""
        return len(self.wpoints)

    @property
    def real_wpoints(self):
        """
        Real frequencies in Hartree. Empty list if not available.
        """
        if self.nrew > 0:
            return np.real(self.wpoints[: self.nrew])
        return []

    @property
    def imag_wpoints(self):
        """
        Imaginary frequencies in Hartree. Empty list if not available.
        """
        if self.nimw > 0:
            return self.wpoints[self.nrew :]
        return []

    @property
    def wggmat_realw(self):
        """The slice of wggmat along the real axis."""
        return self.wggmat[: self.nrew, :, :]

    @property
    def wggmat_imagw(self):
        """The slice of wggmat along the imaginary axis."""
        return self.wggmat[self.nrew :, :, :]

    def windex(self, w, atol: float = 0.001) -> int:
        """
        Find the index of the **closest** frequency in ``wpoints``.
        """
        if np.iscomplex(w):
            iw = bs_index(self.imag_wpoints.imag, w.imag, atol=atol)
            iw += self.nrew
        else:
            iw = bs_index(self.real_wpoints.real, w, atol=atol)

        return iw

    def gindex(self, gvec):
        """
        Find the index of gvec. If ``gvec`` is an integer, gvec is returned.

        Raises:
            `ValueError` if gvec is not found.
        """
        if duck.is_intlike(gvec):
            return int(gvec)
        return self.gsphere.index(gvec)

    def latex_label(self, cplx_mode: str) -> str:
        """Return a latex string that can be used in matplotlib plots."""
        return _latex_symbol_cplxmode(self.latex_name, cplx_mode)

    @add_fig_kwargs
    def plot_freq(
        self,
        gvec1,
        gvec2=None,
        waxis="real",
        cplx_mode="re-im",
        ax=None,
        color: str = None,
        label: str = None,
        fontsize=8,
        **kwargs,
    ) -> Figure:
        r"""
        Plot the frequency dependence of :math:`W_{G1, G2}(\omega)`

        Args:
            gvec1, gvec2: G-vectors.
            waxis: ``real`` to plot along the real axis, ``imag`` for the imaginary axis.
            cplx_mode: string defining the data to print.
                Possible choices are (case-insensitive): ``re`` for the real part
                ``im`` for the imaginary part, ``abs`` for the absolute value.
                ``angle`` will display the phase of the complex number in radians.
                Options can be concatenated with ``-`` e.g. ``re-im``.
            ax: |matplotlib-Axes| or None if a new figure should be created.
            color: Optional color passed to ax.plot.
            label: Optional string for label.
            fontsize: legend and label fontsize.

        Returns: |matplotlib-Figure|
        """
        # Select data to plot
        ig1 = self.gindex(gvec1)
        ig2 = ig1 if gvec2 is None else self.gindex(gvec2)

        ax, fig, plt = get_ax_fig_plt(ax=ax)

        if waxis == "real":
            if self.nrew == 0:
                cprint("File does not contain frequencies along the real axis!", color="orange")
                return fig

            xx = self.real_wpoints.real * pmgu.Ha_to_eV
            yy = self.wggmat_realw[:, ig1, ig2]

        elif waxis == "imag":
            if self.nimw == 0:
                cprint("File does not contain frequencies along the imaginary axis!", color="orange")
                return fig

            xx = self.imag_wpoints.imag * pmgu.Ha_to_eV
            yy = self.wggmat_imagw[:, ig1, ig2]

        else:
            raise ValueError(f"Wrong value for {waxis=}")

        plt_kwargs = dict(
            linewidth=kwargs.pop("linewidth", 2),
            linestyle=kwargs.pop("linestyle", "solid"),
        )
        if color is not None:
            plt_kwargs["color"] = color

        for c in cplx_mode.lower().split("-"):
            if "color" not in plt_kwargs:
                plt_kwargs["color"] = _COLOR_CMODE[c]

            ax.plot(xx, data_from_cplx_mode(c, yy), label=self.latex_label(c) if label is None else label, **plt_kwargs)

        ax.grid(True)
        ax.set_xlabel(r"$\omega$ (eV)")
        ax.set_title("%s, kpoint: %s" % (self.netcdf_name, self.kpoint), fontsize=fontsize)
        ax.legend(loc="best", fontsize=fontsize, shadow=True)

        return fig

    @add_fig_kwargs
    def plot_gg(self, cplx_mode="abs", wpos=None, **kwargs) -> Figure:
        r"""
        Use matplotlib imshow to plot :math:`W_{GG'}` matrix

        Args:
            cplx_mode: string defining the data to print.
                Possible choices are (case-insensitive): ``re`` for the real part
                ``im`` for the imaginary part, ``abs`` for the absolute value.
                ``angle`` will display the phase of the complex number in radians.
            wpos: List of frequency indices to plot. If None, the first frequency is used (usually w=0).
                If wpos == "all" all frequencies are shown (use it carefully)
                Other possible values: "real" if only real frequencies are wanted.
                "imag" for imaginary frequencies only.

        Returns: |matplotlib-Figure|
        """
        # Get wpos indices.
        choice_wpos = {None: [0], "all": range(self.nw), "real": range(self.nrew), "imag": range(self.nrew, self.nw)}

        if any(wpos == k for k in choice_wpos):
            wpos = choice_wpos[wpos]
        else:
            if duck.is_intlike(wpos):
                wpos = [int(wpos)]
            wpos = np.array(wpos)

        # Build plotter.
        plotter = ArrayPlotter()
        for iw in wpos:
            data = data_from_cplx_mode(cplx_mode, self.wggmat[iw])
            label = r"%s $\omega=%s$" % (self.latex_label(cplx_mode), self.wpoints[iw])
            plotter.add_array(label, data)

        return plotter.plot(show=False, **kwargs)


class Polarizability(_AwggMatrix):
    """
    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: Polarizability
    """

    netcdf_name = "polarizability"
    latex_name = r"\tilde chi"


class DielectricFunction(_AwggMatrix):
    """
    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: DielectricFunction
    """

    netcdf_name = "dielectric_function"
    latex_name = r"\epsilon"


class InverseDielectricFunction(_AwggMatrix):
    """
    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: InverseDielectricFunction
    """

    netcdf_name = "inverse_dielectric_function"
    latex_name = r"\epsilon^{-1}"

    # def _add_ppmodel(self, ppm):
    #    """
    #    Add a :class:`PPModel` object to the internal list. Return ppm.
    #    """
    #    if not hasattr(self, "ppmodels"): self.ppmodels = []
    #    self.ppmodels.append(ppm)
    #    return ppm

    # def build_godby_needs_ppmodel(self, wplasma):
    #    ppm = GodbyNeeds.from_em1(self, wplasma)
    #    return self._add_ppmodel(ppm)

    # @add_fig_kwargs
    # def plot_with_ppmodels(self, gvec1, gvec2=None, waxis="real", cplx_mode="re",
    #                       zcut=0.1/pmgu.Ha_to_eV, **kwargs):
    #    """
    #    Args:
    #        gvec1, gvec2:
    #        waxis: "real" to plot along the real axis, "imag" for the imaginary axis.
    #        cplx_mode: string defining the data to print.
    #           Possible choices are (case-insensitive): `re` for the real part
    #           "im" for the imaginary part, "abs" for the absolute value.
    #           "angle" will display the phase of the complex number in radians.
    #           Options can be concatenated with "-" e.g. "re-im"
    #        zcut: Small shift along the imaginary axis to avoid poles. 0.1 eV is the Abinit default.
    #        ax: matplotlib :class:`Axes` or None if a new figure should be created.

    #    Returns:
    #        matplotlib figure.
    #    """
    #    # Select the (G,G') indices to plot.
    #    ig1 = self.gindex(gvec1)
    #    ig2 = ig1 if gvec2 is None else self.gindex(gvec2)

    #    ax, fig, plt = get_ax_fig_plt(ax=None)

    #    self.plot_freq(gvec1, gvec2=gvec2, waxis=waxis, cplx_mode=cplx_mode, ax=ax, show=False)

    #    # Compute em1 from the ppmodel on the same grid used for self.
    #    omegas = {"real": self.real_wpoints, "imag": self.imag_wpoints}[waxis]

    #    # Get y-limits of the ab-initio em1 to zoom-in the interesting region
    #    ymin_em1, ymax_em1 = ax.get_ylim()

    #    for ppm in self.ppmodels:
    #        em1_ppm = ppm.eval_em1(omegas, zcut)
    #        em1_ppm.plot_freq(ig1, gvec2=ig2, waxis=waxis, cplx_mode=cplx_mode,
    #                          ax=ax, linestyle="--", show=False)

    #    ax.set_ylim(ymin_em1, ymax_em1)

    #    return fig


# class PPModel(metaclass=abc.ABCMeta):
#    """
#    Abstract base class for Plasmonpole models.
#    """
#
#    #@abc.abstractmethod
#    #def from_em1(cls, em1):
#    #    """Compute the plasmon-pole parameters from the inverse dielectric function."""
#
#    @abc.abstractmethod
#    def eval_em1(self, omegas, zcut):
#        """Compute the plasmon-pole model at frequency omega (Ha units)."""
#
#
# class GodbyNeeds(PPModel):
#
#    def __init__(self, gsphere, omegatw, bigomegatwsq):
#        r"""
#        bigomegatwsq(:)
#        Plasmon pole parameters $\tilde\Omega^2_{G Gp}(q)$.
#
#        omegatw(:)
#        omegatw(nqibz)%value(npwc,dm2_otq)
#        Plasmon pole parameters $\tilde\omega_{G Gp}(q)$.
#        """
#        self.gsphere = gsphere
#        self.kpoint = gsphere.kpoint
#        self.omegatw = omegatw
#        self.bigomegatwsq = bigomegatwsq
#
#        self.ng = len(gsphere)
#        assert len(omegatw) == len(bigomegatwsq)
#        assert len(omegatw) == len(gsphere)
#
#    @classmethod
#    def from_em1(cls, em1, wplasma):
#        # Find omega=0 and the second imaginary frequency to fit the ppm parameters.
#        iw0 = -1; iw1 = -1
#        for i, w in enumerate(em1.wpoints):
#            if np.abs(w) <= 1e-6: iw0 = i
#            if np.abs(w - 1j*wplasma) <= 1e-6: iw1 = i
#        if iw0 == -1:
#            raise ValueError("Cannot find omega=0 in em1")
#        if iw1 == -1:
#            raise ValueError("Cannot find second imaginary frequency at %s in em1!" % wplasma)
#
#        w0gg = em1.wggmat[iw0, :, :]
#        w1gg = em1.wggmat[iw1, :, :]
#
#        aa = w0gg - np.eye(em1.ng)
#        diff = w0gg - w1gg
#        ratio = aa / diff
#        omegatwsq = (ratio - 1.0) * (wplasma ** 2)
#
#        # If omega-twiddle-squared is negative,set omega-twiddle-squared to 1.0 (a reasonable way of treating
#        # such terms, in which epsilon**-1 was originally increasing along this part of the imaginary axis)
#        # (note: originally these terms were ignored in Sigma; this was changed on 6 March 1990.)
#        #if (REAL(omegatwsq) <= 0.0) omegatwsq=one
#        omegatwsq[np.where(omegatwsq.real <= 0.0)] = 1.0
#        #
#        # Get omega-twiddle
#        # * Neglect the imag part (if any) in omega-twiddle-squared
#        #omegatw(ig,igp)=SQRT(REAL(omegatwsq))
#        #omegatw = np.sqrt(omegatwsq)
#        omegatw = np.sqrt(omegatwsq.real)
#        #omegatw = omegatw + 7j * pmgu.eV_to_Ha
#
#        bigomegatwsq = -aa * omegatw**2
#
#        return cls(em1.gsphere, omegatw, bigomegatwsq)
#
#    def eval_em1(self, omegas, zcut):
#        omegas = np.array(omegas)
#
#        wggmat = np.empty((len(omegas), self.ng, self.ng), dtype=complex)
#        for i, w in enumerate(omegas):
#            # Add shift but only along the real axis.
#            delta = 0.0 if w.imag != 0 else 1j * zcut
#            den = w**2 - np.real((self.omegatw - delta)**2)
#            em1gg = np.eye(self.ng) + self.bigomegatwsq / den
#            #arg = (w - self.omegatw.real) / (0.01 * pmgu.eV_to_Ha)
#            #em1gg = em1gg * (1 - np.exp(-arg**2))
#            wggmat[i] = em1gg
#
#        return InverseDielectricFunction(self.kpoint, omegas, self.gsphere, wggmat)
#
#    @add_fig_kwargs
#    def plot_ggparams(self, **kwargs):
#        plotter = ArrayPlotter(*[
#            (r"$\tilde\omega_{G G'}$", self.omegatw),
#            (r"$\tilde\Omega^2_{G, G'}$", self.bigomegatwsq)])
#
#        return plotter.plot(show=False, **kwargs)
