# fmt: off

import numpy as np

import ase  # Annotations
from ase.calculators.calculator import PropertyNotImplementedError
from ase.utils import jsonable


def calculate_band_structure(atoms, path=None, scf_kwargs=None,
                             bs_kwargs=None, kpts_tol=1e-6, cell_tol=1e-6):
    """Calculate band structure.

    The purpose of this function is to abstract a band structure calculation
    so the workflow does not depend on the calculator.

    First trigger SCF calculation if necessary, then set arguments
    on the calculator for band structure calculation, then return
    calculated band structure.

    The difference from get_band_structure() is that the latter
    expects the calculation to already have been done."""
    if path is None:
        path = atoms.cell.bandpath()

    from ase.lattice import celldiff  # Should this be a method on cell?
    if any(path.cell.any(1) != atoms.pbc):
        raise ValueError('The band path\'s cell, {}, does not match the '
                         'periodicity {} of the atoms'
                         .format(path.cell, atoms.pbc))
    cell_err = celldiff(path.cell, atoms.cell.uncomplete(atoms.pbc))
    if cell_err > cell_tol:
        raise ValueError('Atoms and band path have different unit cells.  '
                         'Please reduce atoms to standard form.  '
                         'Cell lengths and angles are {} vs {}'
                         .format(atoms.cell.cellpar(), path.cell.cellpar()))

    calc = atoms.calc
    if calc is None:
        raise ValueError('Atoms have no calculator')

    if scf_kwargs is not None:
        calc.set(**scf_kwargs)

    # Proposed standard mechanism for calculators to advertise that they
    # use the bandpath keyword to handle band structures rather than
    # a double (SCF + BS) run.
    use_bandpath_kw = getattr(calc, 'accepts_bandpath_keyword', False)
    if use_bandpath_kw:
        calc.set(bandpath=path)
        atoms.get_potential_energy()
        return calc.band_structure()

    atoms.get_potential_energy()

    if hasattr(calc, 'get_fermi_level'):
        # What is the protocol for a calculator to tell whether
        # it has fermi_energy?
        eref = calc.get_fermi_level()
    else:
        eref = 0.0

    if bs_kwargs is None:
        bs_kwargs = {}

    calc.set(kpts=path, **bs_kwargs)
    calc.results.clear()  # XXX get rid of me

    # Calculators are too inconsistent here:
    # * atoms.get_potential_energy() will fail when total energy is
    #   not in results after BS calculation (Espresso)
    # * calc.calculate(atoms) doesn't ask for any quantity, so some
    #   calculators may not calculate anything at all
    # * 'bandstructure' is not a recognized property we can ask for
    try:
        atoms.get_potential_energy()
    except PropertyNotImplementedError:
        pass

    ibzkpts = calc.get_ibz_k_points()
    kpts_err = np.abs(path.kpts - ibzkpts).max()
    if kpts_err > kpts_tol:
        raise RuntimeError('Kpoints of calculator differ from those '
                           'of the band path we just used; '
                           'err={} > tol={}'.format(kpts_err, kpts_tol))

    bs = get_band_structure(atoms, path=path, reference=eref)
    return bs


def get_band_structure(atoms=None, calc=None, path=None, reference=None):
    """Create band structure object from Atoms or calculator."""
    # path and reference are used internally at the moment, but
    # the exact implementation will probably change.  WIP.
    #
    # XXX We throw away info about the bandpath when we create the calculator.
    # If we have kept the bandpath, we can provide it as an argument here.
    # It would be wise to check that the bandpath kpoints are the same as
    # those stored in the calculator.
    atoms = atoms if atoms is not None else calc.atoms
    calc = calc if calc is not None else atoms.calc

    kpts = calc.get_ibz_k_points()

    energies = []
    for s in range(calc.get_number_of_spins()):
        energies.append([calc.get_eigenvalues(kpt=k, spin=s)
                         for k in range(len(kpts))])
    energies = np.array(energies)

    if path is None:
        from ase.dft.kpoints import (
            BandPath,
            find_bandpath_kinks,
            resolve_custom_points,
        )
        standard_path = atoms.cell.bandpath(npoints=0)
        # Kpoints are already evaluated, we just need to put them into
        # the path (whether they fit our idea of what the path is, or not).
        #
        # Depending on how the path was established, the kpoints might
        # be valid high-symmetry points, but since there are multiple
        # high-symmetry points of each type, they may not coincide
        # with ours if the bandpath was generated by another code.
        #
        # Here we hack it so the BandPath has proper points even if they
        # come from some weird source.
        #
        # This operation (manually hacking the bandpath) is liable to break.
        # TODO: Make it available as a proper (documented) bandpath method.
        kinks = find_bandpath_kinks(atoms.cell, kpts, eps=1e-5)
        pathspec, special_points = resolve_custom_points(
            kpts[kinks], standard_path.special_points, eps=1e-5)
        path = BandPath(standard_path.cell,
                        kpts=kpts,
                        path=pathspec,
                        special_points=special_points)

    # XXX If we *did* get the path, now would be a good time to check
    # that it matches the cell!  Although the path can only be passed
    # because we internally want to not re-evaluate the Bravais
    # lattice type.  (We actually need an eps parameter, too.)

    if reference is None:
        # Fermi level should come from the GS calculation, not the BS one!
        reference = calc.get_fermi_level()

    if reference is None:
        # Fermi level may not be available, e.g., with non-Fermi smearing.
        # XXX Actually get_fermi_level() should raise an error when Fermi
        # level wasn't available, so we should fix that.
        reference = 0.0

    return BandStructure(path=path,
                         energies=energies,
                         reference=reference)


class BandStructurePlot:
    def __init__(self, bs):
        self.bs = bs
        self.ax = None
        self.xcoords = None

    def plot(self, ax=None, *, spin=None, emin=-10, emax=5, filename=None,
             show=False, ylabel=None, colors=None, point_colors=None,
             label=None, loc=None,
             cmap=None, cmin=-1.0, cmax=1.0, sortcolors=False,
             colorbar=True, clabel='$s_z$', cax=None,
             **plotkwargs):
        """Plot band-structure.

        ax: Axes
            MatPlotLib Axes object.  Will be created if not supplied.
        spin: int or None
            If given, only plot the specified spin channel.
            If None, plot all spins.
            Default: None, i.e., plot all spins.
        emin, emax: float
            Minimum and maximum energy above reference.
        filename: str
            If given, write image to a file.
        show: bool
            Show the image (not needed in notebooks).
        ylabel: str
            The label along the y-axis.  Defaults to 'energies [eV]'
        colors: sequence of str
            A sequence of one or two color specifications, depending on
            whether there is spin.
            Default: green if no spin, yellow and blue if spin is present.
        point_colors: ndarray
            An array of numbers of the shape (nspins, n_kpts, nbands) which
            are then mapped onto colors by the colormap (see ``cmap``).
            ``colors`` and ``point_colors`` are mutually exclusive
        label: str or list of str
            Label for the curves on the legend.  A string if one spin is
            present, a list of two strings if two spins are present.
            Default: If no spin is given, no legend is made; if spin is
            present default labels 'spin up' and 'spin down' are used, but
            can be suppressed by setting ``label=False``.
        loc: str
            Location of the legend.

        If ``point_colors`` is given, the following arguments can be specified.

        cmap:
            Only used if colors is an array of numbers.  A matplotlib
            colormap object, or a string naming a standard colormap.
            Default: The matplotlib default, typically 'viridis'.
        cmin, cmax: float
            Minimal and maximal values used for colormap translation.
            Default: -1.0 and 1.0
        colorbar: bool
            Whether to make a colorbar.
        clabel: str
            Label for the colorbar (default 's_z', set to None to suppress.
        cax: Axes
            Axes object used for plotting colorbar.  Default: split off a
            new one.
        sortcolors (bool or callable):
            Sort points so highest color values are in front.  If a callable is
            given, then it is called on the color values to determine the sort
            order.

        Any additional keyword arguments are passed directly to matplotlib's
        plot() or scatter() methods, depending on whether point_colors is
        given.
        """
        import matplotlib.pyplot as plt

        if colors is not None and point_colors is not None:
            raise ValueError("Don't give both 'color' and 'point_color'")

        if self.ax is None:
            ax = self.prepare_plot(ax, emin, emax, ylabel)

        if spin is None:
            e_skn = self.bs.energies
        elif spin not in [0, 1]:
            raise ValueError(f"spin should be 0 or 1, not {spin}")
        else:
            # Select only one spin channel.
            e_skn = self.bs.energies[spin, np.newaxis]

        nspins = len(e_skn)

        if point_colors is None:
            # Normal band structure plot
            if colors is None:
                if len(e_skn) == 1:
                    colors = 'g'
                else:
                    colors = 'yb'
            elif (len(colors) != nspins):
                raise ValueError(
                    "colors should be a sequence of {nspin} colors"
                )

            # Default values for label
            if label is None and nspins == 2:
                label = ['spin up', 'spin down']

            if label:
                if nspins == 1 and isinstance(label, str):
                    label = [label]
                elif len(label) != nspins:
                    raise ValueError(
                        f'label should be a list of {nspins} strings'
                    )

            for spin, e_kn in enumerate(e_skn):
                kwargs = dict(color=colors[spin])
                kwargs.update(plotkwargs)
                lbl = None   # Retain lbl=None if label=False
                if label:
                    lbl = label[spin]
                ax.plot(self.xcoords, e_kn[:, 0], label=lbl, **kwargs)

                for e_k in e_kn.T[1:]:
                    ax.plot(self.xcoords, e_k, **kwargs)
            show_legend = label is not None or nspins == 2

        else:
            # A color per datapoint.
            kwargs = dict(vmin=cmin, vmax=cmax, cmap=cmap, s=1)
            kwargs.update(plotkwargs)
            shape = e_skn.shape
            xcoords = np.zeros(shape)
            xcoords += self.xcoords[np.newaxis, :, np.newaxis]
            if sortcolors:
                if callable(sortcolors):
                    perm = sortcolors(point_colors).argsort(axis=None)
                else:
                    perm = point_colors.argsort(axis=None)
                e_skn = e_skn.ravel()[perm].reshape(shape)
                point_colors = point_colors.ravel()[perm].reshape(shape)
                xcoords = xcoords.ravel()[perm].reshape(shape)

            things = ax.scatter(xcoords, e_skn, c=point_colors, **kwargs)
            if colorbar:
                cbar = plt.colorbar(things, cax=cax)
                if clabel:
                    cbar.set_label(clabel)
            show_legend = False

        self.finish_plot(filename, show, loc, show_legend)

        return ax

    def prepare_plot(self, ax=None, emin=-10, emax=5, ylabel=None):
        import matplotlib.pyplot as plt
        if ax is None:
            ax = plt.figure().add_subplot(111)

        def pretty(kpt):
            if kpt == 'G':
                kpt = r'$\Gamma$'
            elif len(kpt) == 2:
                kpt = kpt[0] + '$_' + kpt[1] + '$'
            return kpt

        self.xcoords, label_xcoords, orig_labels = self.bs.get_labels()
        label_xcoords = list(label_xcoords)
        labels = [pretty(name) for name in orig_labels]

        i = 1
        while i < len(labels):
            if label_xcoords[i - 1] == label_xcoords[i]:
                labels[i - 1] = labels[i - 1] + ',' + labels[i]
                labels.pop(i)
                label_xcoords.pop(i)
            else:
                i += 1

        for x in label_xcoords[1:-1]:
            ax.axvline(x, color='0.5')

        ylabel = ylabel if ylabel is not None else 'energies [eV]'

        ax.set_xticks(label_xcoords)
        ax.set_xticklabels(labels)
        ax.set_ylabel(ylabel)
        ax.axhline(self.bs.reference, color='k', ls=':')
        ax.axis(xmin=0, xmax=self.xcoords[-1], ymin=emin, ymax=emax)
        self.ax = ax
        return ax

    def finish_plot(self, filename, show, loc, show_legend=False):
        import matplotlib.pyplot as plt

        if show_legend:
            leg = plt.legend(loc=loc)
            leg.get_frame().set_alpha(1)

        if filename:
            plt.savefig(filename)

        if show:
            plt.show()


@jsonable('bandstructure')
class BandStructure:
    """A band structure consists of an array of eigenvalues and a bandpath.

    BandStructure objects support JSON I/O.
    """

    def __init__(self, path, energies, reference=0.0):
        self._path = path
        self._energies = np.asarray(energies)
        assert self.energies.shape[0] in [1, 2]  # spins x kpts x bands
        assert self.energies.shape[1] == len(path.kpts)
        assert np.isscalar(reference)
        self._reference = reference

    @property
    def energies(self) -> np.ndarray:
        """The energies of this band structure.

        This is a numpy array of shape (nspins, nkpoints, nbands)."""
        return self._energies

    @property
    def path(self) -> 'ase.dft.kpoints.BandPath':
        """The :class:`~ase.dft.kpoints.BandPath` of this band structure."""
        return self._path

    @property
    def reference(self) -> float:
        """The reference energy.

        Semantics may vary; typically a Fermi energy or zero,
        depending on how the band structure was created."""
        return self._reference

    def subtract_reference(self) -> 'BandStructure':
        """Return new band structure with reference energy subtracted."""
        return BandStructure(self.path, self.energies - self.reference,
                             reference=0.0)

    def todict(self):
        return dict(path=self.path,
                    energies=self.energies,
                    reference=self.reference)

    def get_labels(self, eps=1e-5):
        """"See :func:`ase.dft.kpoints.labels_from_kpts`."""
        return self.path.get_linear_kpoint_axis(eps=eps)

    def plot(self, *args, **kwargs):
        """Plot this band structure."""
        bsp = BandStructurePlot(self)
        return bsp.plot(*args, **kwargs)

    def __repr__(self):
        return ('{}(path={!r}, energies=[{} values], reference={})'
                .format(self.__class__.__name__, self.path,
                        '{}x{}x{}'.format(*self.energies.shape),
                        self.reference))