kofta-plot

#!python3
# -*- coding: utf-8 -*-

import argparse
import csv

from pathlib import Path


import matplotlib.pyplot as plt
import numpy as np


from debug import pok, pwarn, phello


def plot(plot_data: dict[str, dict[str, list[int]]], title: str | None, output: Path | None) -> None:

    for name, data in plot_data.items():
        color = plt.plot(data["time"], data["paths"], label=name)[-1].get_color()
        if "paths_min" in data and "paths_max" in data:
            plt.fill_between(data["time"], data["paths_min"], data["paths_max"], alpha=0.2, color=color)

    plt.xlabel("Time (seconds)")
    plt.ylabel("Paths found")
    plt.legend()

    if title:
        plt.title(title)

    if output:
        plt.savefig(output, bbox_inches='tight')
    else:
        plt.show()


def read_plot_data(dir: Path, time: int | None, sample: int) -> tuple[str | None, dict[str, list[int]]]:

    data = {"time": [], "paths": []}
    with (dir / "plot_data").open("r") as f:

        reader = csv.DictReader(f)
        row = next(reader, None)
        if row is None:
            return None, {}

        data["name"] = dir.name
        data["time"].append(0)
        data["paths"].append(int(row[" paths_total"]))

        start = int(row["# unix_time"])
        now = 0

        while row := next(reader, None):
            t = int(row["# unix_time"]) - start
            p = int(row[" paths_total"])
            if time is not None and t > time:
                break
            while t >= now + sample:
                now += sample
                data["paths"].append(data["paths"][-1])
                data["time"].append(now)
            data["paths"][-1] = p

        return dir.name, data


def read_group_data(group: Path, time: int | None, sample: int) -> tuple[str | None, dict[str, list[int]]]:

    plot_data: dict[str, dict[str, list[int]]] = {}

    for dir in group.iterdir():
        if not dir.is_dir():
            continue
        name, data = read_plot_data(dir, time, sample)
        if name is None:
            continue
        plot_data[name] = data

    if not plot_data:
        return None, {}

    pok(f"Grp: {group.name}:", [data["paths"][-1] for data in plot_data.values()])

    min_size = min(len(data["time"]) for data in plot_data.values())
    for data in plot_data.values():
        data["time"] = data["time"][:min_size]
        data["paths"] = data["paths"][:min_size]

    np_data = np.array([data["paths"] for data in plot_data.values()])
    group_data: dict[str, list[int]] = {
        "time": plot_data[next(iter(plot_data))]["time"],
        "paths": np_data.mean(axis=0).tolist(),
        "paths_min": np_data.min(axis=0).tolist(),
        "paths_max": np_data.max(axis=0).tolist()
    }

    return group.name, group_data


def parse_args() -> argparse.Namespace:

    parser = argparse.ArgumentParser(description="Plot KOFTA paths coverage over time.")

    parser.add_argument("-d", "--dir",    metavar="STATE_DIR",   type=Path, dest="dirs",   action="append", help="a directory containing AFL state data")
    parser.add_argument("-g", "--group",  metavar="STATE_GROUP", type=Path, dest="groups", action="append", help="a group containing multiple AFL states")

    parser.add_argument("-l", "--labels", metavar="LABEL_LIST",                                             help="comma-separated list of labels to replace in the plot (e.g., 'd1,d2,.,g2')")
    parser.add_argument(      "--title",  metavar="TITLE",                                                  help="title of the plot")

    parser.add_argument("-t", "--time",   metavar="SECOND",      type=int,                                  help="time in seconds to plot")
    parser.add_argument("-s", "--sample", metavar="SECOND",      type=int,  default=5,                      help="sample time in seconds")

    parser.add_argument("output",         metavar="OUTPUT",      type=Path, nargs="?",                      help="output file to save the plot")

    return parser.parse_args()


def main():

    phello("kofta-plot")

    args = parse_args()

    groups: list[Path] = args.groups if args.groups else []
    dirs: list[Path] = args.dirs if args.dirs else []

    labels: list[str | None] = args.labels.split(",") if args.labels else []
    title: str | None = args.title

    time: int | None = args.time
    sample: int = args.sample

    output: Path | None = args.output

    nplot = len(dirs) + len(groups)
    if len(labels) < nplot:
        labels += [None] * (nplot - len(labels))
    elif len(labels) > nplot:
        labels = labels[:nplot]
    lidx = 0

    plot_data: dict[str, dict[str, list[int]]] = {}

    for dir in dirs:
        name, data = read_plot_data(dir, time, sample)
        if name is None:
            continue
        pok(f"Dir: {dir.name}: {data['paths'][-1]}")
        if labels[lidx] is not None and labels[lidx] != '.':
            name = labels[lidx]
        lidx += 1
        plot_data[name] = data

    for group in groups:
        name, data = read_group_data(group, time, sample)
        if name is None:
            continue
        if labels[lidx] is not None and labels[lidx] != '.':
            name = labels[lidx]
        lidx += 1
        plot_data[name] = data

    plot(plot_data, title, output)


if __name__ == "__main__":
    main()