#! /usr/bin/env python
"""
Plot *_cpp.dat and *_py.dat in one figure for comparison.
"""
import numpy as np
import matplotlib.pyplot as plt


def plot(cpp_file: str, py_file: str, i_x: int = 0, i_y: int = 1):
    """
    :param cpp_file: path to cpp data file
    :param py_file: path to py data file
    :param i_x: column index of x-data, -1 indicates automatically generated
        linspace
    :param i_y: column index of y-data
    """
    cpp_data = np.loadtxt(cpp_file)
    py_data = np.loadtxt(py_file)
    if i_x == -1:
        cx = np.linspace(0, cpp_data.shape[0]-1, cpp_data.shape[0])
        px = cx
    else:
        cx = cpp_data[:, i_x]
        px = py_data[:, i_x]
    plt.plot(cx, cpp_data[:, i_y], "r+", linewidth=0.5, label="cpp")
    plt.plot(px, py_data[:, i_y], "b-", linewidth=0.5, label="python")
    plt.legend()
    plt.show()
    plt.close()


def main():
    cpp_dir = "cpp"
    py_dir = "py"
    tags = ["haydock/ldos", "dos/dos", "ac_cond/ac_cond", "dyn_pol/dyn_pol",
            "dyn_pol/epsilon", "dc_cond/dc_cond_xx", "hall_cond/hall_cond",
            "qe/qe", "wft/wft"]
    for tag in tags:
        if tag == "qe/qe":
            plot(f"{cpp_dir}/{tag}_cpp.dat", f"{py_dir}/{tag}_py.dat", -1, 1)
        elif tag == "wft/wft":
            for i in (0, 2, 4):
                plot(f"{cpp_dir}/wft/wft_cpp.dat", f"{py_dir}/wft/wft_py.dat", -1, i)
        else:
            plot(f"{cpp_dir}/{tag}_cpp.dat", f"{py_dir}/{tag}_cpp.dat")


if __name__ == "__main__":
    main()
