clboss/contrib/plot-size-balance

#!/usr/bin/env python3
import argparse
import json
import math
import os
import sqlite3


SECONDS_PER_DAY = 24 * 60 * 60


def load_earnings_filter(path, min_abs_net_msat):
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    history = data.get("history", [])
    if not isinstance(history, list):
        raise ValueError("earnings history JSON missing 'history' array")
    active = set()
    for entry in history:
        if not isinstance(entry, dict):
            continue
        bucket_time = entry.get("bucket_time")
        node = entry.get("node")
        if not bucket_time or bucket_time <= 0 or not node:
            continue
        in_earnings = entry.get("in_earnings", 0)
        out_earnings = entry.get("out_earnings", 0)
        in_expenditures = entry.get("in_expenditures", 0)
        out_expenditures = entry.get("out_expenditures", 0)
        net = (in_earnings + out_earnings) - (in_expenditures + out_expenditures)
        if abs(net) >= min_abs_net_msat:
            active.add((int(bucket_time), node))
    return active


def fetch_points(db_path, min_age_days, earnings_filter):
    min_age_seconds = min_age_days * SECONDS_PER_DAY
    query = """
        WITH firsts AS (
            SELECT peer_id, MIN(ts) AS first_ts
            FROM fee_change_events
            GROUP BY peer_id
        )
        SELECT f.ts,
               p.node_id,
               f.size_less_peers,
               f.size_total_peers,
               f.balance_our_msat,
               f.balance_total_msat,
               f.price_level
        FROM fee_change_events f
        JOIN peers p ON p.id = f.peer_id
        JOIN firsts ON firsts.peer_id = f.peer_id
        WHERE f.size_less_peers IS NOT NULL
          AND f.size_total_peers IS NOT NULL
          AND f.balance_our_msat IS NOT NULL
          AND f.balance_total_msat IS NOT NULL
          AND f.price_level IS NOT NULL
          AND f.ts >= firsts.first_ts + :min_age
    """
    with sqlite3.connect(db_path) as conn:
        rows = conn.execute(query, {"min_age": min_age_seconds}).fetchall()
    if not earnings_filter:
        return rows
    filtered = []
    for (
        ts,
        node_id,
        size_less,
        size_total,
        balance_our,
        balance_total,
        price_level,
    ) in rows:
        bucket = int(math.floor(ts / SECONDS_PER_DAY) * SECONDS_PER_DAY)
        if (bucket, node_id) not in earnings_filter:
            continue
        filtered.append(
            (ts, node_id, size_less, size_total, balance_our, balance_total, price_level)
        )
    return filtered


def main():
    parser = argparse.ArgumentParser(
        description="Scatter plot of size_ratio vs balance_ratio, colored by theory_level."
    )
    parser.add_argument(
        "--db",
        default=os.path.expanduser("./clboss-fee-info.sqlite3"),
        help="Path to sqlite database.",
    )
    parser.add_argument(
        "--min-age-days",
        type=int,
        default=60,
        help="Minimum age in days since a peer's first data point.",
    )
    parser.add_argument(
        "--earnings-json",
        help="Path to clboss-earnings-history all JSON output.",
    )
    parser.add_argument(
        "--min-abs-net-msat",
        type=int,
        default=None,
        help=(
            "Minimum absolute daily net earnings (msat) for a point to count. "
            "Requires --earnings-json; omit to skip earnings filtering."
        ),
    )
    parser.add_argument(
        "--png",
        nargs="?",
        const="__DEFAULT__",
        default=None,
        help="Output image path.",
    )
    parser.add_argument(
        "--show",
        action="store_true",
        help="Show the plot interactively.",
    )

    args = parser.parse_args()

    if not args.png and not args.show:
        raise SystemExit("use --png to save or --show to display")

    earnings_filter = None
    if args.min_abs_net_msat is not None:
        if not args.earnings_json:
            raise SystemExit("--min-abs-net-msat requires --earnings-json")
        earnings_filter = load_earnings_filter(
            args.earnings_json,
            args.min_abs_net_msat,
        )
        if not earnings_filter:
            raise SystemExit(
                "no earnings entries meeting "
                f"|net| >= {args.min_abs_net_msat} msat"
            )

    points = fetch_points(args.db, args.min_age_days, earnings_filter)
    if not points:
        raise SystemExit("no size/balance data found for the selected age window")

    xs = []
    ys = []
    levels = []
    for (
        _,
        _,
        size_less,
        size_total,
        balance_our,
        balance_total,
        price_level,
    ) in points:
        if size_total <= 0 or size_less < 0:
            continue
        if size_less > size_total:
            continue
        if balance_total <= 0 or balance_our < 0:
            continue
        if balance_our > balance_total:
            continue
        xs.append(balance_our / balance_total)
        ys.append(size_less / size_total)
        levels.append(price_level)
    if not xs:
        raise SystemExit("no valid size/balance points after filtering")

    try:
        import matplotlib.pyplot as plt
        from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm
    except ImportError:
        raise SystemExit("matplotlib is required to output graphs.")

    max_abs = max(abs(min(levels)), abs(max(levels)))
    if max_abs == 0:
        max_abs = 1
    norm = TwoSlopeNorm(vmin=-max_abs, vcenter=0, vmax=max_abs)
    cmap = LinearSegmentedColormap.from_list(
        "theory_level",
        ["#c0392b", "#bdbdbd", "#1e8449"],
    )

    fig, ax = plt.subplots(figsize=(8, 6))
    sc = ax.scatter(xs, ys, c=levels, s=8, alpha=0.4, linewidths=0, cmap=cmap, norm=norm)
    ax.set_xlabel("balance_ratio")
    ax.set_ylabel("size_ratio")

    title = f"size_ratio vs balance_ratio (>= {args.min_age_days} days)"
    if earnings_filter is not None:
        title += f", |net| >= {args.min_abs_net_msat} msat"
    ax.set_title(title)

    cbar = fig.colorbar(sc, ax=ax, pad=0.02)
    cbar.set_label("theory_level")

    fig.tight_layout()

    if args.png:
        out_path = args.png
        if out_path == "__DEFAULT__":
            out_path = "size-balance.png"
        plt.savefig(out_path, dpi=150)

    if args.show:
        plt.show()


if __name__ == "__main__":
    main()