import re

import numpy as np

from ase.dft.kpoints import get_bandpath, special_paths, special_points


def isstring(obj):
    return hasattr(obj, 'isalnum')


# Actually this method should probably call get_special_points()
# since some of them depend on the cell.  (Then of course we must provide
# the cell)
def normalize_special_points(lattice, symbols=None):
    """Convert special k-point symbols (G, M, ...) into coordinates.

    See also get_band_path.

    symbols can be a single string (e.g. 'GXWK') or a list.
    If symbols is a list, elements may be strings or vectors (kx, ky, kx).
    If symbols is None, a standard path is chosen.

    Returns a list of symbols and an array of coordinates."""

    try:
        special = special_points[lattice]
    except KeyError:
        raise KeyError('No lattice "%s".  Lattices: %s'
                       % (lattice, list(special_points.keys())))

    if symbols is None:
        symbols = special_paths[lattice][0]  # Is this a good default?
    elif isstring(symbols):
        # Convert strings like 'GHXMH1' into ['G', 'H', ..., 'H1']
        # Also works if one of them is 'Gamma' which is used somewhere.
        symbols = [name for name in re.split(r'([A-Z][a-z0-9]*)', symbols)
                   if name]

    # Convert symbols into coordinates:
    coords = []
    for kpt in symbols:
        if isstring(kpt):
            if kpt not in special:
                raise KeyError('No special point "%s" for lattice %s.  '
                               'Special points: %s'
                               % (kpt, lattice, list(special.keys())))
            kpt = special[kpt]
        # Sym could also be a vector (kx, ky, kz)
        assert np.shape(kpt) == (3,)
        coords.append(kpt)

    symbols = [str(sym) for sym in symbols]
    coords = np.array(coords, float)
    return symbols, coords


def get_band_structure(atoms, lattice, symbols=None, npoints=50):
    symbols, special_kpts = normalize_special_points(lattice, symbols)
    kpts, xcoords, special_xcoords = get_bandpath(special_kpts, atoms.cell,
                                                  npoints)

    # Maybe change get_band_path so it returns an object?
    band_path = BandPath(lattice, symbols, special_kpts, special_xcoords,
                         kpts, xcoords)

    calc = atoms.get_calculator()
    # XXX how many calculators support set?
    calc.set(kpts=kpts)
    calc.calculate(atoms)
    # XXX spin polarized?

    # XXX Why does calc.get_eigenvalues() not just return
    # a [spin x kpts x bands] array by default?

    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)
    return BandStructure(band_path, energies)


class BandPath:
    def __init__(self, lattice, symbols, special_kpts, special_xcoords,
                 kpts, xcoords):
        self.lattice = lattice  # 'fcc' or similar

        self.symbols = symbols  # 'G', 'K', '[0.5, 0.5, 0.0]', 'M', ...
        self.special_kpts = special_kpts  # Special k-points
        self.special_xcoords = special_xcoords

        self.kpts = kpts
        self.xcoords = xcoords

    # Implement read(), write()


class BandStructure:
    # Should BandStructure and BandPath be the same object?
    # We could add an occupations array and do some more stuff with this
    # object, which would then be worth a separate class

    def __init__(self, band_path, energies):
        self.band_path = band_path
        self.energies = energies
        # Relevant to have Fermi energy?
        # Save occupations too for some reason?

    def plot(self, ax=None, spin=None):
        if ax is None:
            import matplotlib.pyplot as plt
            ax = plt.gca()

        path = self.band_path

        def pretty(kpt):
            return r'$\Gamma$' if kpt in ['G', 'Gamma'] else kpt

        if spin is None:
            e_skn = self.energies
        else:
            e_skn = self.energies[spin][None]

        for spin, e_kn in enumerate(self.energies):
            color = 'br'[spin]
            for e_k in e_kn.T:
                ax.plot(path.xcoords, e_k, marker='.', ls='-', color=color)

        ax.set_xticks(path.special_xcoords)
        ax.set_xticklabels([pretty(name) for name in path.symbols])
        ax.axis(xmin=0, xmax=path.xcoords[-1])
        return ax

    # Implement read/write

def main2():
    from gpaw import GPAW, PW, MixerSum
    from ase.build import bulk

    atoms = bulk('Fe', 'bcc')
    atoms.set_initial_magnetic_moments([1.] * len(atoms))
    calc = GPAW(mode=PW(300),
                mixer=MixerSum(),
                kpts=(4, 4, 4),
                nbands='200%')
    atoms.calc = calc
    atoms.get_potential_energy()

    calc.set(fixdensity=True, symmetry='off')

    bs = get_band_structure(atoms, 'bcc', npoints=30)
    bs.plot()

    import matplotlib.pyplot as plt
    plt.show()


def main():
    from gpaw import GPAW, PW
    from ase.build import bulk

    si = bulk('Si', crystalstructure='diamond', a=5.4)
    calc = GPAW(mode=PW(250), nbands='200%')
    si.calc = calc
    si.get_potential_energy()
    eFermi = calc.get_fermi_level()

    calc.set(fixdensity=True, symmetry='off')

    bs = get_band_structure(si, 'fcc', 'GXWKGLUWLK', npoints=60)
    # Also possible:
    #bs = get_band_structure(si, 'fcc', ['G', [0., .5, 0.], 'X'], npoints=20)
    ax = bs.plot()

    # Maybe: bs.get_dos()?  bs2 = bs.fold_bands(...)?

    # Should we save eFermi and other things in BS object?
    # How many decorations should the plot() function add?
    ax.axis(ymax=eFermi + 10)
    ax.grid(axis='x')
    ax.axhline(eFermi, color='k')

    import matplotlib.pyplot as plt
    plt.show()

if __name__ == '__main__':
    main()