#!/usr/bin/env python

import sys

from diag import *
from diag_feast import *
from tbpm import *


def run_diag(test: str) -> int:
    matched = 1
    if test == "diag_all":
        test_diag_bands()
        test_diag_dos()
        test_diag_ldos()
        test_diag_states()
        test_spin()
        test_z2()
        test_berry()
        test_diag_dyn_pol()
        test_diag_ac_cond()
    elif test == "feast":
        test_diag_bands_feast()
        test_diag_dos_feast()
        test_diag_ldos_feast()
        test_spin_feast()
    else:
        matched = 0
    return matched


def run_tbpm(test: str) -> int:
    matched = 1
    if test == "haydock":
        test_haydock()
    elif test == "dos":
        test_dos()
    elif test == "ac_cond":
       test_ac_cond()
    elif test == "dyn_pol":
        test_dyn_pol()
    elif test == "dc_cond":
       test_dc_cond()
    elif test == "hall_cond":
        test_hall_cond()
    elif test == "qe":
        test_qe()
    elif test == "wft":
        test_wft()
    elif test == "qe_speed":
        test_qe_speed()
    elif test == "wft_speed":
        test_wft_speed()
    elif test == "bands":
        test_bands()
    else:
        matched = 0
    return matched


def run_model(test: str) -> int:
    raise NotImplementedError("run_model not implemented yet!")


def main():
    test = sys.argv[1]
    matched = run_diag(test)
    matched += run_tbpm(test)
    # matched += run_model(test)
    if matched == 0:
        raise AssertionError(f"Unknown test: {test}")


if __name__ == "__main__":
    main()