#!/usr/bin/env python3
"""
plot_fio_latency.py — visualize fio json+ latency CSVs.

Designed for the wide-format CSV produced by fio_jsonplus_clat2csv in
fio 3.x, which has columns like:

    nsec, read_slat_ns_count, read_slat_ns_cumulative, read_slat_ns_percentile,
          read_clat_ns_count, read_clat_ns_cumulative, read_clat_ns_percentile,
          ... (one triple per metric) ...

Only the metric(s) that were actually exercised contain data.

Examples
--------
    python plot_fio_latency.py scan_latency_job0.csv
    python plot_fio_latency.py scan_latency_job0.csv --metric read_clat --unit ms
    python plot_fio_latency.py scan_latency_job0.csv --metric read_clat --cdf --log-y
    python plot_fio_latency.py scan_latency_job0.csv --all-metrics --cdf-only
"""

import argparse
import os
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter


UNIT_FACTORS = {"ns": 1.0, "us": 1e3, "ms": 1e6, "s": 1e9}
UNIT_LABELS  = {"ns": "ns", "us": "µs", "ms": "ms", "s": "s"}



# ---------------------------------------------------------------------------
# Loading
# ---------------------------------------------------------------------------

def load_latency_csv(path: str) -> pd.DataFrame:
    """
    Parse the wide fio_jsonplus CSV into a long-format DataFrame with columns:
        metric, latency_ns, count, cumulative, percentile
    
    Transforms wide-format columns (e.g., read_clat_ns_count) into rows with
    a metric identifier, making it easier to group and analyze by operation type.
    """
    df = pd.read_csv(path, skipinitialspace=True)
    df.columns = [str(c).strip() for c in df.columns]
    # Drop trailing unnamed / empty columns
    df = df.loc[:, ~df.columns.str.match(r"^Unnamed")]
    df = df.loc[:, df.columns != ""]

    nsec_col = df.columns[0]
    if "nsec" not in nsec_col.lower():
        raise ValueError(
            f"expected first column to be 'nsec', got {nsec_col!r}"
        )

    # Coerce the nsec column to numeric (it should always be numeric, but
    # be defensive about stray whitespace).
    df[nsec_col] = pd.to_numeric(df[nsec_col], errors="coerce")
    df = df.dropna(subset=[nsec_col])

    # Discover metric triples: <metric>_count / _cumulative / _percentile
    metrics = {}
    for col in df.columns:
        if col == nsec_col:
            continue
        for suffix in ("_count", "_cumulative", "_percentile"):
            if col.endswith(suffix):
                base = col[: -len(suffix)]           # e.g. read_clat_ns
                metrics.setdefault(base, {})[suffix[1:]] = col

    if not metrics:
        raise ValueError("no *_count columns found in CSV")

    frames = []
    for base, cols in metrics.items():
        count_col = cols["count"]

        # Coerce count to numeric — blanks become NaN, whitespace is ignored.
        counts = pd.to_numeric(df[count_col], errors="coerce")

        # Keep only rows where the count column actually has a number.
        mask = counts.notna()
        if not mask.any():
            continue

        sub = pd.DataFrame({
            "latency_ns": df.loc[mask, nsec_col].values,
            "count": counts[mask].astype(int).values,
        })
        if "cumulative" in cols:
            sub["cumulative"] = pd.to_numeric(
                df.loc[mask, cols["cumulative"]], errors="coerce"
            ).values
        if "percentile" in cols:
            sub["percentile"] = pd.to_numeric(
                df.loc[mask, cols["percentile"]], errors="coerce"
            ).values

        # read_clat_ns -> read_clat
        sub["metric"] = base.replace("_ns", "")
        frames.append(sub)

    if not frames:
        raise ValueError("no metric had any numeric count data")

    out = pd.concat(frames, ignore_index=True)
    out = out[out["count"] > 0].reset_index(drop=True)
    return out


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def to_unit(values_ns, unit: str) -> np.ndarray:
    """Convert nanosecond values to the specified unit (us, ms, or s)."""
    return np.asarray(values_ns, dtype=float) / UNIT_FACTORS[unit]


def human_count(x, _pos=None):
    """Format large numbers with K/M/G suffixes for axis labels."""
    if x >= 1e9: return f"{x/1e9:.1f}G"
    if x >= 1e6: return f"{x/1e6:.1f}M"
    if x >= 1e3: return f"{x/1e3:.0f}k"
    return f"{x:.0f}"


def percentile_table(df: pd.DataFrame, metric: str, unit: str, pcts) -> str:
    """
    Generate a formatted table showing latency values at specified percentiles
    (e.g., p50, p99, p99.9) for a given metric.
    """
    g = df[df["metric"] == metric].sort_values("latency_ns")
    if g.empty:
        return f"(no data for metric {metric})"
    counts = g["count"].values.astype(float)
    total = counts.sum()
    cdf = np.cumsum(counts) / total * 100.0
    lat = to_unit(g["latency_ns"].values, unit)
    lines = [f"\nMetric: {metric}   (total I/Os: {int(total):,})",
             f"  {'percentile':>12}  {'latency':>14}"]
    for p in pcts:
        idx = min(np.searchsorted(cdf, p), len(lat) - 1)
        lines.append(f"  {p:>11.2f}%  {lat[idx]:>12.3f} {UNIT_LABELS[unit]}")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Plotting
# ---------------------------------------------------------------------------

def draw_histogram(ax, df, metrics, unit, log_y, log_x, normalize):
    """
    Draw a histogram showing the distribution of latencies for each metric.
    
    Uses a step plot (where="post") to show the discrete nature of the data.
    If normalize=True, Y-axis shows probability density instead of counts.
    """
    for metric in metrics:
        g = df[df["metric"] == metric].sort_values("latency_ns")
        if g.empty:
            continue
        x = to_unit(g["latency_ns"].values, unit)
        y = g["count"].values.astype(float)
        if normalize:
            s = y.sum()
            if s > 0:
                y = y / s
        ax.step(x, y, where="post", lw=1.4, label=metric)
        ax.fill_between(x, y, step="post", alpha=0.25)

    ax.set_xlabel(f"Latency ({UNIT_LABELS[unit]})")
    ax.set_ylabel("Probability" if normalize else "Number of I/Os")
    ax.set_title("Completion latency distribution")
    ax.grid(True, which="both", ls=":", alpha=0.5)
    if not normalize:
        ax.yaxis.set_major_formatter(FuncFormatter(human_count))
    if log_y:
        ax.set_yscale("log")
    if log_x:
        ax.set_xscale("log")
    if len(metrics) > 1:
        ax.legend(title="metric", fontsize=9)


def draw_cdf(ax, df, metrics, unit, log_x):
    """
    Draw a Cumulative Distribution Function (CDF) showing the percentage of
    I/Os completed within each latency threshold.
    
    Each metric gets its own line; the Y-axis ranges from 0% to 100%.
    """
    for metric in metrics:
        g = df[df["metric"] == metric].sort_values("latency_ns")
        if g.empty:
            continue
        x = to_unit(g["latency_ns"].values, unit)
        counts = g["count"].values.astype(float)
        total = counts.sum()
        if total <= 0:
            continue
        cdf = np.cumsum(counts) / total * 100.0
        ax.plot(x, cdf, lw=1.6, label=f"{metric} (n={int(total):,})")

    ax.set_xlabel(f"Latency ({UNIT_LABELS[unit]})")
    ax.set_ylabel("Cumulative % of I/Os")
    ax.set_title("Cumulative latency distribution")
    ax.set_ylim(0, 100)
    ax.grid(True, which="both", ls=":", alpha=0.5)
    ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:.0f}%"))
    if log_x:
        ax.set_xscale("log")
    ax.legend(fontsize=9)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    p = argparse.ArgumentParser(
        description="Plot fio json+ wide-format latency CSV."
    )
    p.add_argument("csv", help="CSV from fio_jsonplus_clat2csv")
    p.add_argument("--metric", default=None,
                   help="Metric to plot (e.g. read_clat). "
                        "If omitted, the metric with the most I/Os is used.")
    p.add_argument("--all-metrics", action="store_true",
                   help="Plot every metric that has data")
    p.add_argument("--unit", choices=list(UNIT_FACTORS), default="us",
                   help="X-axis latency unit (default: us)")
    p.add_argument("--log-y", action="store_true")
    p.add_argument("--log-x", action="store_true")
    p.add_argument("--cdf", action="store_true",
                   help="Add a CDF subplot below the histogram")
    p.add_argument("--cdf-only", action="store_true")
    p.add_argument("--normalize", action="store_true",
                   help="Normalize histogram to a probability density")
    p.add_argument("--percentiles", default="50,90,99,99.9,99.99")
    p.add_argument("--no-percentiles", action="store_true")
    p.add_argument("--output", "-o", default=None)
    p.add_argument("--dpi", type=int, default=150)
    p.add_argument("--title", default="")
    p.add_argument("--show", action="store_true")
    args = p.parse_args()

    if not os.path.isfile(args.csv):
        sys.exit(f"error: file not found: {args.csv}")

    df = load_latency_csv(args.csv)
    if df.empty:
        sys.exit("error: no latency data in CSV")

    available = sorted(df["metric"].unique())
    print("Metrics with data:", ", ".join(available))

    # Select which metrics to plot
    if args.all_metrics:
        metrics = available
    elif args.metric:
        if args.metric not in available:
            sys.exit(f"error: metric {args.metric!r} not in {available}")
        metrics = [args.metric]
    else:
        # Auto-select: pick the metric with the largest total count
        totals = df.groupby("metric")["count"].sum().sort_values(ascending=False)
        metrics = [totals.index[0]]
        print(f"Auto-selected metric: {metrics[0]}")

    # ------------------------------------------------------------------
    # Build figure
    # ------------------------------------------------------------------
    # Create subplot layout based on user's request:
    #   --cdf-only: CDF only
    #   --cdf: Histogram + CDF stacked vertically
    #   default: Histogram only
    if args.cdf_only:
        fig, ax = plt.subplots(figsize=(11, 6))
        draw_cdf(ax, df, metrics, args.unit, args.log_x)
    elif args.cdf:
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(11, 10))
        draw_histogram(ax1, df, metrics, args.unit, args.log_y, args.log_x,
                       normalize=args.normalize)
        draw_cdf(ax2, df, metrics, args.unit, args.log_x)
    else:
        fig, ax = plt.subplots(figsize=(11, 6))
        draw_histogram(ax, df, metrics, args.unit, args.log_y, args.log_x,
                       normalize=args.normalize)

    if args.title:
        fig.suptitle(args.title, fontsize=13, y=0.995)
    fig.tight_layout()

    # Save plot to file (default: <input>_latency.png)
    out = args.output or f"{Path(args.csv).stem}_latency.png"
    fig.savefig(out, dpi=args.dpi, bbox_inches="tight")
    print(f"Saved plot to {out}")

    # Print percentile table (p50, p90, p99, p99.9, p99.99 by default)
    if not args.no_percentiles:
        try:
            pcts = [float(x) for x in args.percentiles.split(",") if x.strip()]
        except ValueError:
            sys.exit("error: --percentiles must be comma-separated numbers")
        for m in metrics:
            print(percentile_table(df, m, args.unit, pcts))

    # Optionally display the plot interactively
    if args.show:
        plt.show()


if __name__ == "__main__":
    main()