#!/usr/bin/env python
"""
Plot diagonalization results of full speedtest, which are generated by
either cpp or python version of speedtest.
"""
import os

from plot_diag import main as plot
from utils import print_color


top_dir = os.getcwd()

args = [("diag_bands", "graphene_bands.dat"),
        ("diag_dos", "graphene_dos.dat"),
        ("diag_ldos", "graphene_local_dos.dat"),
        ("spin", "kane_mele_spin_sigma_x.dat"),
        ("spin", "bismuth_spin_sigma_x.dat"),
        ("z2", "kane_mele_z2.dat"),
        ("berry", "haldane_kubo_berry_omega_xy.dat"),
        ("berry", "haldane_wilson_berry_omega_xy.dat"),
        ("diag_dyn_pol", "graphene_dyn_pol.dat"),
        ("diag_dyn_pol", "graphene_epsilon.dat"),
        ("diag_ac_cond", "graphene_ac_cond.dat"),
        ("diag_ac_cond", "graphene_epsilon_q0.dat"),
        ("diag_bands_feast", "graphene_bands.dat"),
        ("diag_dos_feast", "graphene_dos.dat"),
        ("diag_ldos_feast", "graphene_local_dos.dat"),
        ("spin_feast", "kane_mele_spin_sigma_x.dat")]

for a in args:
    try:
        os.chdir(f"{top_dir}/{a[0]}")
        plot(a[1])
    except (FileNotFoundError, AttributeError):  # AttributeError for missing HDF5 support
        print_color(f"WARNING: {a[0]}/{a[1]} and found and skipped!", "r")
        continue
