"""Plot the process CPU usage using VaspInteractive pause and resume. 
The results are generated by running
```
python ex13_pause_mpi.py
```
on a 8-core machine. CPU usage is recorded using:
```
top -b -d 0.1 >> top_vpi_pause_0.1.txt
```

"""
import numpy as np
import matplotlib.pyplot as plt


def get_blocks(txt):
    """Get output blocks from GNU top. Time interval 0.1 sec"""
    with open(txt, "r") as fd:
        content = fd.readlines()
    block_lines = [i for i, ln in enumerate(content) if ln.startswith("top")]
    blocks = []
    for i in range(1, len(block_lines)):
        start = block_lines[i - 1]
        end = block_lines[i]
        blocks.append(content[start:end])
    return blocks


def parse_line(
    line,
    name="vasp_std",
    matching={"pid": (0, int), "name": (-1, str), "cpu": (-4, float)},
    discard_pids=[],
):
    """Parse each line to find pid / name / cpu usage"""
    if name not in line:
        return None
    fields = line.split()
    results = {}
    for key, value in matching.items():
        field_id, func = value
        results[key] = func(fields[field_id])
    # Discard results from certain pids
    if results["pid"] in discard_pids:
        return None
    else:
        return results


def parse_block(block, names=["vasp_std", "python"], **kwargs):
    matches = []
    for line in block:
        for name in names:
            res = parse_line(line, name, **kwargs)
            if res is not None:
                matches.append(res)
                break
    return matches


def main(txt, t_interval=0.1, discard_pids=[1099, 1130]):
    """Plot CPU usage"""
    blocks = get_blocks(txt)
    times = np.arange(len(blocks)) * t_interval
    cpu_vasp = []
    cpu_python = []
    for block in blocks:
        matches = parse_block(block, discard_pids=discard_pids)
        # vasp
        vasp_cpus = [match["cpu"] for match in matches if match["name"] == "vasp_std"]
        if len(vasp_cpus) > 0:
            cpu_vasp.append([np.mean(vasp_cpus), np.std(vasp_cpus)])
        else:
            cpu_vasp.append([np.nan, np.nan])

        # python
        python_cpus = [match["cpu"] for match in matches if match["name"] == "python"]
        if len(python_cpus) > 0:
            cpu_python.append([np.mean(python_cpus), np.std(python_cpus)])
        else:
            cpu_python.append([np.nan, np.nan])

    cpu_vasp = np.array(cpu_vasp)
    cpu_python = np.array(cpu_python)

    # align the time
    cutter1 = [5, 85]
    cutter2 = [84, 140]
    cond1 = np.where((times > cutter1[0]) & (times < cutter1[1]))
    cond2 = np.where((times > cutter2[0]) & (times < cutter2[1]))
    times1, cpu_vasp1, cpu_python1 = (
        times[cond1],
        cpu_vasp[:, 0][cond1],
        cpu_python[:, 0][cond1],
    )
    times2, cpu_vasp2, cpu_python2 = (
        times[cond2],
        cpu_vasp[:, 0][cond2],
        cpu_python[:, 0][cond2],
    )
    times1 -= times1[0]
    times2 -= times2[0]

    # time_vasp1
    plt.subplot(211)

    plt.plot(times1, cpu_vasp1, label="vasp")
    plt.plot(times1, cpu_python1, label="python")
    plt.axhline(0, ls="--", color="grey")
    plt.axhline(100, ls="--", color="grey")
    plt.ylim(-5, 105)
    plt.xlim(0, 85)
    plt.ylabel("CPU usage (%)")
    # plt.legend()

    plt.subplot(212)
    plt.plot(times2, cpu_vasp2, label="vasp")
    plt.plot(times2, cpu_python2, label="python")
    plt.axhline(0, ls="--", color="grey")
    plt.axhline(100, ls="--", color="grey")
    plt.ylim(-5, 105)
    plt.xlim(0, 85)
    plt.legend()
    plt.xlabel("Wall time (s)")
    plt.ylabel("Average CPU usage per core (%)")

    plt.tight_layout()
    plt.savefig("ex13_time_cpu.svg")
    plt.savefig("ex13_time_cpu.png")


if __name__ == "__main__":
    txt = "top_vpi_pause_0.1.txt"
    main(txt)
