#!/usr/bin/env python

import argparse
import os
from pprint import pprint

import phonopy
from matplotlib import pyplot as plt
from twinpy.analysis.phonon_analyzer import PhononAnalyzer
from twinpy.common.band_path import get_band_paths_from_labels, get_seekpath
from twinpy.plot.band_structure import BandPlot

from my_codes.graphs.bandv2 import get_band_labels

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "file_name", default="phonopy.yaml", help="phonopy.yamls to read"
    )
    parser.add_argument("--auto", action="store_true", help="same as --band auto tag")
    parser.add_argument("-f", "--figure", default="band.png", help="graph name to save")
    parser.add_argument("-p", "--plot", action="store_true", help="plot mode")
    args = parser.parse_args()

    phonon = phonopy.load(args.file_name)
    phonon_analyzer = PhononAnalyzer(phonon)

    # dos
    # phonon_analyzer.run_mesh(interval=0.2)
    # dos = phonon_analyzer.get_total_dos()
    # tdp = TotalDosPlot(dos, flip_xy=True)

    primitive_cell = phonon_analyzer.primitive_cell
    seek_path = get_seekpath(primitive_cell)
    labels_qpoints = seek_path["point_coords"]
    labels = get_band_labels(seek_path)
    if args.auto is False:
        labels_qpoints["X_1"] = [0.5, 1, 0.5]
        labels = ["GAMMA", "X", "", "X_1", "GAMMA", "L"]
    pprint(labels_qpoints)
    labels, band_paths = get_band_paths_from_labels(labels, labels_qpoints)
    band_structure = phonon_analyzer.get_band_structure(
        band_paths=band_paths, labels=labels
    )
    bp = BandPlot(band_structure)
    bp.set_ylim(-2.0, 10.0)
    fig, axes = bp.plot_band_structure(figsize=(6.0, 4.2), c="#1f77b4")
    axes[0].yaxis.label.set_size(20)
    plt.setp(axes[0].get_yticklabels(), fontsize=18)
    for ax in axes:
        plt.setp(ax.get_xticklabels(), fontsize=18)
    band_data_dir = os.getenv("HOME") + "/mlp-Fe/input/Fe/band_data/table/"
    with open(band_data_dir + "gamma_x.txt") as f:
        s = [line.split() for line in f]
    x = [float(line[0]) for line in s]
    y = [float(line[1]) for line in s]
    axes[0].scatter(x, y, s=20, marker="x", c="black", linewidth=1.0)
    with open(band_data_dir + "x_gamma_l.txt") as f:
        s = [line.split() for line in f]
    x = [float(line[0]) for line in s]
    y = [float(line[1]) for line in s]
    axes[1].scatter(x, y, s=20, marker="x", c="black", linewidth=1.0)
    if args.plot:
        plt.show()
    else:
        plt.savefig(args.figure, transparent=True)
