#!/usr/bin/env python
"""
Post-process and plot TBPM results of full speedtest, which are generated by
either cpp or python version of speedtest.

NOTE:
1. This script saves post-processed results to *_cpp.dat
2. The python version of speedtest also calls Python TBPMSolver and Analyzer
   automatically and produces *_py.dat.
"""
import os

from plot_tbpm import main as plot
from utils import print_color


top_dir = os.getcwd()

args = [("haydock", "sample_ldos.dat"),
        ("dos", "sample_corr_dos.dat"),
        ("ac_cond", "sample_corr_ac_cond.dat"),
        ("dyn_pol", "sample_corr_dyn_pol.dat"),
        ("dc_cond", "sample_corr_dc_cond_xx.dat"),
        ("hall_cond", "sample_hall_cond.dat"),
        ("qe", "sample_quasi_eigen.dat"),
        ("wft", "sample_wft.dat"),
        ("bands", "sample_corr_bands.dat")]

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