#!/usr/bin/env python
import argparse
from dfttools.parsers.generic import ParseError
from dfttools.simple import parse

from numericalunits import eV

parser = argparse.ArgumentParser(description="Plots a band structure from a supported file")
parser.add_argument("-o", "--output", nargs=1, help="plotted band structure file name", metavar="FILENAME")
parser.add_argument("-d", "--density", nargs='?', help="plot density of states", metavar="FILENAME", const=True)
parser.add_argument("-f", "--fermi", help="Fermi level in eV", metavar="VALUE", type=float)
parser.add_argument("-0", "--fermi-origin", nargs='?', help="Set Fermi level to zero", const=True)
parser.add_argument("-e", "--range", help="Energy range in eV", metavar="FROM,TO")
parser.add_argument("-c", "--color", nargs=1, help="Color bands according to a scalar", metavar="FILENAME")
parser.add_argument("-m", "--cmap", nargs=1, help="Color map used for coloring bands", metavar="VALUE")
parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output")
parser.add_argument("-l", "--labels", help="Labels for bands path")
parser.add_argument("file", help="file with the band structure", metavar="FILENAME")

options = parser.parse_args()


def report(_bands):
    fermi_present = _bands.fermi is not None
    print("Fermi present: {}".format("yes" if fermi_present else "no"))

    if fermi_present:

        print("Fermi value: {:.6f}".format(_bands.fermi / eV))

        if _bands.gapped:

            try:
                print("Gap size: {:.6f} eV".format(_bands.gap / eV))
            except ValueError:
                print("Gap size unknown (either occupied ot virtual bands are absent)")

        else:
            print("No band gap")


if options.fermi is not None:
    options.fermi = options.fermi * eV

if options.range is not None:
    options.range = list(map(float, options.range.split(",")))

if options.labels is not None:
    options.labels = options.labels.split(",")

data = {}

for file_name, destination in (
(options.file, "bands"), (options.density[0] if "__getitem__" in dir(options.density) else None, "density")):

    if file_name is not None:

        try:

            with open(file_name, 'r') as f:

                try:
                    if options.verbose:
                        print("Parsing {} ...".format(file_name))

                    bands = parse(f, 'band-structure')
                except ParseError:
                    print("Could not parse file '{}'\n".format(file_name))
                    raise

                if options.fermi is not None:
                    bands.fermi = options.fermi

                if options.fermi_origin:
                    bands.canonize_fermi()

                data[destination] = bands

        except IOError:

            print("Could no read file '{}'\n".format(options.file))
            raise

        if options.verbose:
            report(bands)

if options.color is not None:
    with open(options.color[0], 'r') as f:
        weights = parse(f, "basis-dependent", data["bands"]).values

else:
    weights = None

if options.cmap is not None:
    options.cmap = options.cmap[0]

try:
    from matplotlib import pyplot, gridspec
    from dfttools.presentation import matplotlib_bands, matplotlib_bands_density
except ImportError:
    print("Could not import matplotlib\n")
    raise

if "density" not in data:

    if options.density:

        if options.verbose:
            print("Plotting density only ...")

        matplotlib_bands_density(
            data["bands"],
            pyplot.gca(),
            200,
            orientation="landscape",
            energy_range=options.range,
            method='optimal',
        )

    else:

        if options.verbose:
            print("Plotting bands only ...")

        plot = matplotlib_bands(
            data["bands"],
            pyplot.gca(),
            energy_range=options.range,
            weights_color=weights,
            cmap=options.cmap,
            edge_names=options.labels,
        )

        if options.color is not None:
            pyplot.colorbar(plot)

else:

    fig = pyplot.figure()
    gs = gridspec.GridSpec(1, 2, width_ratios=[3, 1])
    ax_left = pyplot.subplot(gs[0])
    ax_right = pyplot.subplot(gs[1], sharey=ax_left)

    if options.verbose:
        print("Plotting bands (next: density) ...")

    plot = matplotlib_bands(
        data["bands"],
        ax_left,
        energy_range=options.range,
        weights_color=weights,
        cmap=options.cmap,
        edge_names=options.labels,
    )

    if options.verbose:
        print("Plotting density ...")

    matplotlib_bands_density(
        data["density"],
        ax_right,
        200,
        energy_range=ax_left.get_ylim(),
        orientation="portrait",
        method='optimal',
    )

    if options.color is not None:
        pyplot.colorbar(plot)

if options.output is not None:
    pyplot.savefig(options.output[0])
else:
    pyplot.show()
