#!/usr/bin/env python
"""
Parse the json file generated by the magnon band calculation. 
And plot the magnon band structure from the json file. 
You can modify this file to plot the magnon band structure, all plot 
multiple magnon band structures in one figure.
"""
import numpy as np
import matplotlib.pyplot as plt
import json


def load_magnon_json(filename):
    """
    Load the magnon band structure from the json file.
    args:
        filename: the json file generated by the magnon band calculation.
    return:
        kptlist: all the kpoints
        xlist: the x coordinates of the kpoints in the plot
        knames: names of the high symmetry points
        Xs: The x coordinates of the high symmetry points in the plot
        allevals: eigenvalues, in meV
            Note that the eigenvalues are splitted into groups of kpoints.
            the first index is the index of the group of kpoints
            the second index is the index of the kpoint in the group
            the third index is the index of the band
    """
    with open(filename) as f:
        data = json.load(f)
    kptlist = data["kptlist"]
    xlist = data["xlist"]
    knames = data["knames"]
    Xs = data["X_for_highsym_kpoints"]
    allevals = np.array(data["evals"])
    return kptlist, xlist, knames, Xs, allevals


def plot_magnon_json(filename, ax, color="k", show=True):
    """
    Plot the magnon band structure from the json file.
    args:
        filename: the json file generated by the magnon band calculation.
        ax: the matplotlib ax to plot the band structure.
            If None, a new figure will be created.
        color: the color of the band structure.
        show: whether to show the plot.
    """
    if ax is None:
        fig, ax = plt.subplots()

    kptlist, xlist, knames, Xs, allevals = load_magnon_json(filename)
    for i, (kpts, xs) in enumerate(zip(kptlist, xlist)):
        evals = allevals[i]
        # Plot band structure
        nbands = evals.shape[1]
        emin = np.min(evals[:, 0])
        for i in range(nbands):
            ax.plot(xs, (evals[:, i]) / 1.6e-22, color=color)

    ax.set_ylabel("Energy (meV)")
    ax.set_xlim(xlist[0][0], xlist[-1][-1])
    ax.set_xticks(Xs)
    knames = [x if x != "G" else "$\Gamma$" for x in knames]
    ax.set_xticklabels(knames)
    for x in Xs:
        ax.axvline(x, linewidth=0.6, color="gray")
    if show:
        plt.show()


if __name__ == "__main__":
    plot_magnon_json("magnon_band.json", ax=None, color="k", show=True)
