plotter.py

#!/usr/bin/env python3
"""
Reproduce CryptoEntropy figures from the released datasets.

This script reads the CSV files in datasets/entropy/ and generates PNG/PDF
figures in the output folder. It is intentionally focused on reproducing the
figures associated with the DLT2025 CryptoEntropy paper.
"""

from __future__ import annotations

import argparse
import shutil
from pathlib import Path
from typing import Iterable

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import pandas as pd


DEFAULT_ASSETS = ["btc", "eth", "xrp", "usdc", "doge", "ada"]
DEFAULT_SETS = ["set1", "set2", "set3"]


def read_monthly_series(path: Path, value_col: str) -> pd.DataFrame | None:
    """Read a CSV with a time column and return monthly mean values."""
    if not path.exists():
        return None

    df = pd.read_csv(path)
    if "time" not in df.columns or value_col not in df.columns:
        raise ValueError(f"Expected columns 'time' and '{value_col}' in {path}")

    df["time"] = pd.to_datetime(df["time"], errors="coerce", utc=True)
    df = df.dropna(subset=["time"])
    df = df.set_index("time")[[value_col]]
    return df.resample("ME").mean()


def read_probability_distribution(entropy_dir: Path, asset: str, set_name: str) -> pd.DataFrame | None:
    """Read an asset/set probability-distribution file and attach timestamps.

    The released `*-prob.csv` files do not contain their own `time` column.
    The matching `{asset}_{set}.csv` file contains the corresponding timestamps.
    """
    prob_path = entropy_dir / f"{asset}_{set_name}-prob.csv"
    base_path = entropy_dir / f"{asset}_{set_name}.csv"

    if not prob_path.exists() or not base_path.exists():
        return None

    prob = pd.read_csv(prob_path)
    base = pd.read_csv(base_path, usecols=["time"])

    n = min(len(prob), len(base))
    prob = prob.iloc[:n].copy()
    prob.insert(0, "time", base["time"].iloc[:n].values)

    prob["time"] = pd.to_datetime(prob["time"], errors="coerce", utc=True)
    prob = prob.dropna(subset=["time"]).set_index("time")
    return prob.resample("ME").mean()


def save_current_figure(output_dir: Path, stem: str) -> None:
    """Save current Matplotlib figure as PNG and PDF."""
    output_dir.mkdir(parents=True, exist_ok=True)
    plt.savefig(output_dir / f"{stem}.png", bbox_inches="tight")
    plt.savefig(output_dir / f"{stem}.pdf", bbox_inches="tight")


def plot_all_assets(
    entropy_dir: Path,
    output_dir: Path,
    assets: Iterable[str],
    set_name: str,
    file_prefix: str,
    value_col: str,
    ylabel: str,
) -> None:
    """Plot one entropy-like measure for all available assets."""
    series: list[pd.DataFrame] = []

    plt.figure(figsize=(10, 6))

    for asset in assets:
        path = entropy_dir / f"{file_prefix}_{asset}_{set_name}.csv"
        df = read_monthly_series(path, value_col)
        if df is None or df.empty:
            continue

        series.append(df)
        plt.plot(df.index, df[value_col], label=asset)

    if not series:
        plt.close()
        print(f"[skip] No data found for {file_prefix}, {set_name}")
        return

    all_index = series[0].index
    for df in series[1:]:
        all_index = all_index.union(df.index)

    plt.xlim(all_index.min(), all_index.max())
    plt.xticks(fontsize=17)
    plt.ylim(0, 1)
    plt.yticks(fontsize=17)
    plt.grid(axis="y", linestyle="--", alpha=0.7)
    plt.gca().set_axisbelow(True)
    plt.xlabel(set_name, fontsize=20)
    plt.ylabel(ylabel, fontsize=20)
    plt.legend(fontsize=12)
    plt.tight_layout()

    if file_prefix == "entropy":
        stem = f"entropy_all_assets_{set_name}"
    else:
        stem = f"weighted_entropy_all_assets_{set_name}"

    save_current_figure(output_dir, stem)
    plt.close()


def plot_probability_distributions(
    entropy_dir: Path,
    output_dir: Path,
    assets: Iterable[str],
    set_name: str,
) -> None:
    """Plot probability-distribution components for each available asset/set."""
    for asset in assets:
        df = read_probability_distribution(entropy_dir, asset, set_name)
        if df is None or df.empty:
            print(f"[skip] Missing probability data for {asset}, {set_name}")
            continue

        plt.figure(figsize=(10, 6))
        for col in df.columns:
            plt.plot(df.index, df[col], label=col)

        plt.xlim(df.index.min(), df.index.max())
        plt.xticks(fontsize=22)
        plt.ylim(0, 1)
        plt.yticks(fontsize=22)
        plt.grid(axis="y", linestyle="--", alpha=0.7)
        plt.gca().set_axisbelow(True)
        plt.xlabel(f"{asset.upper()} - {set_name}", fontsize=25)
        plt.ylabel("Prob. Distribution", fontsize=25)
        plt.legend(fontsize=16)
        plt.tight_layout()

        save_current_figure(output_dir, f"prob-distr-{asset}-{set_name}")
        plt.close()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Reproduce CryptoEntropy figures from released entropy datasets."
    )
    parser.add_argument(
        "--entropy-dir",
        default="datasets/entropy",
        type=Path,
        help="Path to the entropy dataset folder. Default: datasets/entropy",
    )
    parser.add_argument(
        "--output-dir",
        default="img",
        type=Path,
        help="Folder where figures are written. Default: img",
    )
    parser.add_argument(
        "--assets",
        nargs="+",
        default=DEFAULT_ASSETS,
        help="Assets to plot. Default: btc eth xrp usdc doge ada",
    )
    parser.add_argument(
        "--sets",
        nargs="+",
        default=DEFAULT_SETS,
        help="Attribute sets to plot. Default: set1 set2 set3",
    )
    parser.add_argument(
        "--no-clean",
        action="store_true",
        help="Do not remove the output folder before generating figures.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if not args.entropy_dir.exists():
        raise FileNotFoundError(f"Entropy folder not found: {args.entropy_dir}")

    if args.output_dir.exists() and not args.no_clean:
        shutil.rmtree(args.output_dir)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    for set_name in args.sets:
        plot_all_assets(
            entropy_dir=args.entropy_dir,
            output_dir=args.output_dir,
            assets=args.assets,
            set_name=set_name,
            file_prefix="entropy",
            value_col="entropy",
            ylabel="Entropy",
        )

        plot_all_assets(
            entropy_dir=args.entropy_dir,
            output_dir=args.output_dir,
            assets=args.assets,
            set_name=set_name,
            file_prefix="norm_weighted_entropy",
            value_col="weighted_entropy",
            ylabel="Weighted Entropy",
        )

        plot_probability_distributions(
            entropy_dir=args.entropy_dir,
            output_dir=args.output_dir,
            assets=args.assets,
            set_name=set_name,
        )

    print(f"FINISH PLOT! Figures written to: {args.output_dir}")


if __name__ == "__main__":
    main()