plot.py

"""
This script contains code to parse the logs generated by the benchmark runner and plot the results.
"""
import argparse
import json
import matplotlib.pyplot as plt

import numpy as np
import os

from dataclasses import dataclass
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.ticker import FuncFormatter, MaxNLocator

from results_checker import InvariantsComparator
from utils import get_project_root

InvResult = InvariantsComparator.InvContainmentResult

##
# Dataclass to store graph styles that we use.
##
@dataclass
class GraphStyle:
    # Colors
    bg_color: str = "#F5F5F5"
    low_time_color = "#238B45"
    high_time_color  = "#B22222"
    strengthened_color = "#238B45"
    remaining_color = "#B22222"
    blue_color = "#1A4E8A"
    maroon_color = "#A63232"
    ig_color = "#4682B4"

    dpi = 600

    # Sizes
    line_width: float = 0.6
    label_font = 6
    ticks_font = 7
    axis_font = 8.5
    title_font = 10
    legend_font = 8.5
    marker_size=3.5
    x_axis_width = 0.4
    grid_alp = 0.15

    # Label offsets
    label_x1 = 0
    label_y1 = -3
    label_x2 = 0
    label_y2 = 3

def get_json_from_file(filename):
    with open(filename, "r") as f:
        data = json.load(f)
    
    return data

def get_time_from_json_file(filename, to_remove = []):
    """
    Get the total time taken from the json file. If to_remove is provided, 
    then remove the time taken for those benchmarks from the total time.
    """
    time_data = get_json_from_file(filename)

    total_time = time_data["total_time"]

    if len(to_remove) == 0:
        return total_time

    time_to_remove = 0
    for bench in to_remove:
        if bench in time_data["benchmarks"]:
            time_to_remove += time_data["benchmarks"][bench]["time"]
    
    return (total_time - time_to_remove)

def collect_exp1_lin_data(logs_folder, domains):
    def get_containment_counts(res_dict):
        return {
            enum_member: int(res_dict.get(str(enum_member.value), 0))
            for enum_member in InvResult
        }
    
    lin_data = {}

    for dom in domains:
        dom_logs_folder = logs_folder + f"/{dom}/"

        bl_improve = []
        gb_improve = []
        res_improve = []
        run_times = []
        df = {}

        baseline_time = get_time_from_json_file(f"{dom_logs_folder}/baseline/run_log.json")
        df["baseline_time"] = baseline_time

        gb_time = get_time_from_json_file(f"{dom_logs_folder}/aff-gb/run_log.json")
        df["gb_time"] = gb_time

        for ep in range(6):
            this_dual_folder = f"{dom_logs_folder}/aff-dual-{ep}"

            baseline_cmp = get_containment_counts(get_json_from_file(f"{this_dual_folder}/baseline_cmp.json")["containment_checks"])
            gurobi_cmp = get_containment_counts(get_json_from_file(f"{this_dual_folder}/gurobi_cmp.json")["containment_checks"])
            prev_epoch_cmp = get_containment_counts(get_json_from_file(f"{this_dual_folder}/prev_epoch_cmp.json")["containment_checks"])

            bl_improve.append(baseline_cmp[InvResult.WEAKER])
            gb_improve.append(gurobi_cmp[InvResult.STRONGER])
            res_improve.append(prev_epoch_cmp[InvResult.WEAKER])
            
            run_time = get_time_from_json_file(f"{this_dual_folder}/run_log.json")
            run_times.append(run_time)
        
        df["bl_imp"] = bl_improve
        df["gb_imp"] = gb_improve
        df["res_imp"] = res_improve
        df["runtimes"] = run_times

        lin_data[dom] = df
    
    return lin_data

def plot_og_rg(data, graphfile, style):
    times = sorted(data["runtimes"])
    bl_imp = data["bl_imp"]

    if "gb_imp" in data:
        gb_imp = data["gb_imp"]
    else:
        gb_imp = None

    fig, ax = plt.subplots(figsize=(3, 2))
    # fig, ax = plt.subplots(figsize=(2.1, 1.9))

    fig.patch.set_facecolor("white")
    ax.set_facecolor(style.bg_color)

    if gb_imp is None:
        ax.margins(x=0.05, y=0.1)

    for side in ["top", "right", "left"]:
        ax.spines[side].set_visible(False)
    for side in ["bottom"]:
        ax.spines[side].set_linewidth(style.x_axis_width)
        ax.spines['bottom'].set_color('black')

    ax.tick_params(axis="both", labelsize=style.ticks_font)

    # --- Plot OG and RG ---
    ax.plot(times, bl_imp, marker="s", linewidth=style.line_width, markersize=style.marker_size,
            color=style.blue_color, label="Overall Gain (OG)")
    
    if gb_imp is not None:
        ax.plot(times, gb_imp, marker="o", linewidth=style.line_width, markersize=style.marker_size,
                color=style.remaining_color, label="Remaining Gap (RG)")

    # --- Annotate values ---
    for x, y in zip(times, bl_imp):
        ax.annotate(f"{y}", xy=(x, y),
                    xytext=(style.label_x1, style.label_y1),
                    textcoords="offset points", ha="center",
                    va="top" if style.label_y1 < 0 else "bottom",
                    color=style.blue_color, fontsize=style.label_font)

    if gb_imp is not None:
        for x, y in zip(times, gb_imp):
            ax.annotate(f"{y}", xy=(x, y),
                        xytext=(style.label_x2, style.label_y2),
                        textcoords="offset points", ha="center",
                        va="top" if style.label_y2 < 0 else "bottom",
                        color=style.remaining_color, fontsize=style.label_font)

    # --- Labels and layout ---
    ax.set_xlabel("Runtime (s)", fontsize=style.axis_font)
    if gb_imp is not None:
        ax.set_ylabel("# Invariants", fontsize=style.axis_font)
    else:
        ax.set_ylabel("Overall Gain (OG)", fontsize=style.axis_font)
    ax.grid(True, linestyle="--", alpha=style.grid_alp)

    if gb_imp is not None:
        ax.legend(fontsize=style.legend_font, frameon=False, loc="best")

    if gb_imp is not None:
        ax.set_ylim(bottom=0)

    plt.tight_layout()
    plt.savefig(graphfile, bbox_inches="tight")
    plt.close(fig)

def plot_mg(data, graphfile, style):
    res_imp = data["res_imp"]

    # intervals correspond to transitions: 0→1, 1→2, ..., (n-2)→(n-1)
    R = range(len(res_imp))
    interval_labels = []
    for i in R:
        if i == 0:
            interval_labels.append(fr"$BL \rightarrow {i}$")
        else:
            interval_labels.append(fr"${i - 1} \rightarrow {i}$")

    # Use the configured MG bar color when available.
    bar_color = getattr(style, "ig_color", "#6C7A89")

    with plt.rc_context({'text.usetex': False,
                         'font.family': 'serif'}):
        fig, ax = plt.subplots(figsize=(3.0, 2.0))
        fig.patch.set_facecolor("white")
        ax.set_facecolor(style.bg_color)

        # --- Bars ---
        ax.bar(
            R, res_imp,
            color=bar_color,
            alpha=0.9,
            width=0.55,
            edgecolor="black",
            linewidth=0.4,
        )

        # --- Labels above bars ---
        if res_imp:
            y_offset = max(res_imp) * 0.03
            for x, y in zip(R, res_imp):
                if y > 0:
                    ax.text(
                        x, y + y_offset, f"{y}",
                        ha="center", va="bottom",
                        fontsize=style.label_font + 1,
                        color="black"
                    )

        # --- Axis labels ---
        ax.set_xlabel("Progression of Budget $\\mathcal{R}$", fontsize=style.axis_font)
        ax.set_ylabel("Marginal Gain (MG)", fontsize=style.axis_font)

        # --- X-axis ticks use interval labels ---
        ax.set_xticks(list(R))
        ax.set_xticklabels(interval_labels)

        # --- Grid, ticks, formatting ---
        ax.yaxis.set_major_locator(MaxNLocator(integer=True, nbins=5))
        ax.tick_params(axis="both", labelsize=style.ticks_font, width=0.8)
        ax.grid(axis="y", linestyle="--", alpha=style.grid_alp)
        ax.set_axisbelow(True)

        # Clean spines
        for side in ["top", "right", "left"]:
            ax.spines[side].set_visible(False)
        ax.spines["bottom"].set_linewidth(style.x_axis_width)
        ax.spines["bottom"].set_color("black")

        plt.tight_layout()
        plt.savefig(graphfile, bbox_inches="tight")
        plt.close(fig)

def plot_lin_data(data, graphfile_og_rg, graphfile_ig, style):
    # Reuse the same style object across the paired plots.
    diff = 1
    plot_og_rg(data, graphfile_og_rg, style)
    style.axis_font += diff
    plot_mg(data, graphfile_ig, style)
    style.axis_font -= diff

def create_exp1_plots(logs_folder, plots_folder, is_subset = False, domains = ["elina-zones", "oct"]):
    # Parse the logs and collect the data
    if not is_subset:
        lin_data = collect_exp1_lin_data(logs_folder + "/7.1_linear/nla-digbench", domains)
    else:
        lin_data = collect_exp1_lin_data(logs_folder + "/7.1_linear_subset/nla-digbench", domains)

    # Plot the parsed data
    plots_folder = plots_folder + "/7.1_linear/" if not is_subset else plots_folder + "/7.1_linear_subset/"
    os.makedirs(plots_folder, exist_ok=True)

    # Tune label placement separately for each domain plot.
    style = GraphStyle()
    style.marker_size = 3.8
    style.line_width = 0.4
    style.legend_font = 6.5
    style.axis_font -= 0.5
    style.label_font = 5

    for dom in domains:
        if dom == "elina-zones":
            style.label_y1 = -5
            style.label_y2 = 3
            plot_lin_data(lin_data[dom],f"{plots_folder}/zones_lin_ogrg.pdf", f"{plots_folder}/zones_lin_mg.pdf", style)
        elif dom == "oct":
            style.label_x1 = 1
            style.label_y1 = -5.5
            style.label_x2 = 1
            style.label_y2 = 4.5
            plot_lin_data(lin_data[dom],f"{plots_folder}/oct_lin_ogrg.pdf", f"{plots_folder}/oct_lin_mg.pdf", style)

def collect_exp2_quad_data(logs_folder, domains, times_to_remove = []):
    def get_containment_counts(res_dict):
        return {
            enum_member: int(res_dict.get(str(enum_member.value), 0))
            for enum_member in InvResult
        }

    quad_data = {}

    for dom in domains:
        dom_logs_folder = logs_folder + f"/{dom}/"
        bl_improve = []
        res_improve = []
        run_times = []
        df = {}

        baseline_time = get_time_from_json_file(f"{dom_logs_folder}/baseline/run_log.json", times_to_remove)
        df["baseline_time"] = baseline_time

        for ep in range(6):
            this_dual_folder = f"{dom_logs_folder}/aff-quad-dual-{ep}"

            baseline_cmp = get_containment_counts(get_json_from_file(f"{this_dual_folder}/baseline_cmp.json")["containment_checks"])
            prev_epoch_cmp = get_containment_counts(get_json_from_file(f"{this_dual_folder}/prev_epoch_cmp.json")["containment_checks"])

            bl_improve.append(baseline_cmp[InvResult.WEAKER])
            res_improve.append(prev_epoch_cmp[InvResult.WEAKER])
            
            run_time = get_time_from_json_file(f"{this_dual_folder}/run_log.json", times_to_remove)
            run_times.append(run_time)
        
        df["bl_imp"] = bl_improve
        df["res_imp"] = res_improve
        df["runtimes"] = run_times

        quad_data[dom] = df
    
    return quad_data

def create_exp2_plots(logs_folder, plots_folder, domains = ["elina-zones", "oct", "pk"]):
    # Parse the logs and collect the data
    quad_data = collect_exp2_quad_data(logs_folder + "/7.2_full/nla-digbench", domains)

    # Plot the parsed data
    plots_folder = plots_folder + "/7.2_full/"
    os.makedirs(plots_folder, exist_ok=True)

    style = GraphStyle()
    style.marker_size = 3.4
    style.line_width = 0.05
    style.legend_font = 6.2
    style.axis_font += 0.5
    style.label_font = 4

    for dom in domains:
        plot_lin_data(quad_data[dom],f"{plots_folder}/{dom}_quad_ogrg.pdf", f"{plots_folder}/{dom}_quad_ig.pdf", style)

def create_app_exp2_no_collate_plots(logs_folder, plots_folder, domains = ["elina-zones", "oct", "pk"]):
    # Parse the logs and collect the data
    quad_data = collect_exp2_quad_data(logs_folder + "/appendix_d.3_full_no_collation/nla-digbench", domains, times_to_remove = ["ps5-ll.c"])

    # Plot the parsed data
    plots_folder = plots_folder + "/appendix_d.3_full_no_collation/"
    os.makedirs(plots_folder, exist_ok=True)

    style = GraphStyle()
    style.marker_size = 3.4
    style.line_width = 0.05
    style.legend_font = 6.2
    style.axis_font += 0.5
    style.label_font = 4

    for dom in domains:
        plot_lin_data(quad_data[dom],f"{plots_folder}/{dom}_quad_ogrg.pdf", f"{plots_folder}/{dom}_quad_ig.pdf", style)

def plot_transformer_time_data(data, plot_folder, style):
    lp_obj = data["lp_solver_stats"]["per_obj"]
    dual_obj = data["dual_solver_stats"]["per_obj"]
    lp_constr = data["lp_solver_stats"]["per_cons"]
    dual_constr = data["dual_solver_stats"]["per_cons"]

    def compute_avg(d):
        sizes = sorted(map(int, d.keys()))
        avg_times = [d[s][0] / d[s][1] for s in sizes]
        return sizes, avg_times

    def moving_avg(xs, ys, window=6):
        # Smooth local fluctuations for cleaner trend lines.
        smoothed_xs, smoothed_ys = [], []
        for i in range(len(xs)):
            l = max(0, i - window // 2)
            r = min(len(xs), i + window // 2 + 1)
            smoothed_xs.append(xs[i])
            smoothed_ys.append(sum(ys[l:r]) / (r - l))
        return smoothed_xs, smoothed_ys

    def plot_trend(lpx, lpy, dualx, dualy, x_label, graph_name):
        fig, ax = plt.subplots(figsize=(4, 3), dpi=style.dpi)
        fig.patch.set_facecolor("white")
        ax.set_facecolor(style.bg_color)

        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['bottom'].set_linewidth(style.x_axis_width)
        ax.spines['bottom'].set_color('black')
        ax.tick_params(axis='both', labelsize=style.ticks_font)

        ax.plot(lpx, lpy, linewidth=style.line_width, markersize=style.marker_size, 
                marker='o', color=style.high_time_color, label='LP Solver Transformer')
        ax.plot(dualx, dualy, linewidth=style.line_width, markersize=style.marker_size,
                marker='o', color=style.low_time_color, label='AbsEvolve Transformer')
        ax.set_xlabel(x_label, fontsize=style.axis_font)
        ax.set_ylabel("Average Copmute Time (s)", fontsize=style.axis_font)

        ax.grid(True, linestyle='--', alpha=style.grid_alp)
        ax.legend(fontsize=style.legend_font, loc = "upper left", frameon = False)
        plt.tight_layout()
        plt.savefig(f"{plot_folder}/{graph_name}", bbox_inches="tight")
        plt.close(fig)

    # Compute trends
    obj_sizes_lp, obj_avg_lp = compute_avg(lp_obj)
    obj_sizes_dual, obj_avg_dual = compute_avg(dual_obj)
    constr_sizes_lp, constr_avg_lp = compute_avg(lp_constr)
    constr_sizes_dual, constr_avg_dual = compute_avg(dual_constr)

    # Smooth with moving average
    obj_sizes_lp, obj_avg_lp = moving_avg(obj_sizes_lp, obj_avg_lp)
    obj_sizes_dual, obj_avg_dual = moving_avg(obj_sizes_dual, obj_avg_dual)
    constr_sizes_lp, constr_avg_lp = moving_avg(constr_sizes_lp, constr_avg_lp)
    constr_sizes_dual, constr_avg_dual = moving_avg(constr_sizes_dual, constr_avg_dual)

    style.marker_size = 2
    style.line_width = 0.6
    style.axis_font += 0.5
    style.legend_font += 0.5
    plot_trend(obj_sizes_lp, obj_avg_lp, obj_sizes_dual, obj_avg_dual,
               x_label = "Number of Template Directions",
               graph_name = "solver_trend_with_num_directions.png")

    plot_trend(constr_sizes_lp, constr_avg_lp, constr_sizes_dual, constr_avg_dual,
               x_label = "Number of Constraints",
               graph_name = "solver_trend_with_num_constraints.png")

def collect_transformer_time_data(logs_folder, domains):
    def process_raw_states(raw_stats):
        per_num_obj_stats = {}
        per_num_cons_stats = {}
        for stat in raw_stats:
            objectives = stat["problem"]["combinations"]
            objc = len(objectives)

            var_bounds = stat["problem"]["var_bounds"]
            consc = sum(1 for a, b in var_bounds.values() for x in (a, b) if not isinstance(x, str))
            consc += len(stat["problem"]["constraints"])

            solver_time = stat["solver_time"]

            if objc not in per_num_obj_stats:
                per_num_obj_stats[objc] = (solver_time, 1)
            else:
                per_num_obj_stats[objc] = (per_num_obj_stats[objc][0] + solver_time, per_num_obj_stats[objc][1] + 1)
            
            if consc not in per_num_cons_stats:
                per_num_cons_stats[consc] = (solver_time, 1)
            else:
                per_num_cons_stats[consc] = (per_num_cons_stats[consc][0] + solver_time, per_num_cons_stats[consc][1] + 1)

        return {
            "per_obj": per_num_obj_stats,
            "per_cons": per_num_cons_stats
        }

    # Dual method transformer stats.
    dual_raw_stats = []
    for dom in domains:
        dual_raw_stats.extend(get_json_from_file(f"{logs_folder}/{dom}/aff-dual-5/probs.json"))
    dual_stats = process_raw_states(dual_raw_stats)

    # LP solver method transformer stats.
    lp_raw_stats = []
    for dom in domains:
        lp_raw_stats.extend(get_json_from_file(f"{logs_folder}/{dom}/aff-gb/probs.json"))
    lp_stats = process_raw_states(lp_raw_stats)

    return {
        "dual_solver_stats" : dual_stats,
        "lp_solver_stats" : lp_stats
    }

def create_solver_time_comparison_plots(logs_folder, plots_folder):
    domains = ["elina-zones", "oct"]

    # Parse the solver time stats
    solver_time_comp_data = collect_transformer_time_data(logs_folder + "/7.1_solver_comp/nla-digbench", domains)

    # Plot the parsed data
    plots_folder = plots_folder + "/7.1_solver_comp/"
    os.makedirs(plots_folder, exist_ok=True)

    # Slightly shrink legend text for the comparison figures.
    style = GraphStyle()
    style.legend_font -= 1
    plot_transformer_time_data(solver_time_comp_data, plots_folder, style)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Plotting script for experiments.")
    parser.add_argument(
        "--logs_folder",
        default="logs",
        help="Output logs directory relative to project root"
    )
    parser.add_argument(
        "--plots_folder",
        default="plots",
        help="Plots directory relative to project root"
    )
    args = parser.parse_args()
    logs_folder = f"{get_project_root()}/{args.logs_folder}"
    plots_folder = f"{get_project_root()}/{args.plots_folder}"

    args = parser.parse_args()

    create_exp1_plots(logs_folder, plots_folder)
    create_exp2_plots(logs_folder, plots_folder)
    create_app_exp2_no_collate_plots(logs_folder, plots_folder)
    create_solver_time_comparison_plots(logs_folder, plots_folder)