"""
visualize_results.py
Visualizes the output of run_model.cpp (predictions_latency.csv).

Expected input CSV columns:
    latency_us

Outputs:
    latency_hist_counts.png   – occurrence histogram
    latency_pdf_cdf.png       – density histogram + empirical CDF
"""

import sys
import numpy as np
import pandas as pd

# ─────────────────────────────────────────────────────────────────────────────
# Configuration
# ─────────────────────────────────────────────────────────────────────────────
CSV_INPUT_PATH = "model_t4_inferences_latency.csv"
OUT_PNG_COUNTS = "latency_hist_counts.png"
OUT_PNG_PDF_CDF = "latency_pdf_cdf.png"

# Plot toggles
SHOW_P99 = False
SHOW_CENTRAL_LINE = True

# Select which central-value line to display:
#   "mean" or "trimmed_mean"
CENTRAL_LINE_MODE = "mean"

# Used only when CENTRAL_LINE_MODE == "trimmed_mean"
TRIM_FRACTION = 0.10

# Histogram appearance
COUNT_HIST_FIGSIZE = (8.0, 5.5)
PDF_CDF_FIGSIZE = (13.0, 5.0)
HIST_ALPHA_COUNTS = 0.90
HIST_ALPHA_DENSITY = 0.80
HIST_EDGE_COLOR = "black"
HIST_LINE_WIDTH = 0.8

# Central line appearance
CENTRAL_LINE_COLOR = "red"
CENTRAL_LINE_WIDTH = 0.8

# P99 line appearance
P99_LINE_COLOR = "crimson"
P99_LINE_WIDTH = 1.0

# ─────────────────────────────────────────────────────────────────────────────
# 1. Load results
# ─────────────────────────────────────────────────────────────────────────────
try:
    df = pd.read_csv(CSV_INPUT_PATH)
except FileNotFoundError:
    sys.exit(
        f"ERROR: '{CSV_INPUT_PATH}' not found.\n"
        "Build and run run_model.cpp first to generate this file."
    )

if "latency_us" not in df.columns:
    sys.exit(f"ERROR: Missing required column 'latency_us' in {CSV_INPUT_PATH}")

lat_us = df["latency_us"].values.astype(np.float64)

print(f"Loaded {len(df)} samples from '{CSV_INPUT_PATH}'")

# ─────────────────────────────────────────────────────────────────────────────
# 2. Basic latency statistics
# ─────────────────────────────────────────────────────────────────────────────
if lat_us.size == 0:
    sys.exit("No latency samples found; nothing to analyze.")

print("\nLatency statistics (µs):")
print(f"  Average : {int(np.mean(lat_us))}")
print(f"  Min     : {int(np.min(lat_us))}")
print(f"  Max     : {int(np.max(lat_us))}")

# ─────────────────────────────────────────────────────────────────────────────
# 3. Extended descriptive statistics
# ─────────────────────────────────────────────────────────────────────────────
n = int(lat_us.size)

mean_us = float(np.mean(lat_us))
std_us = float(np.std(lat_us, ddof=1)) if n > 1 else 0.0
median_us = float(np.median(lat_us))

q1 = float(np.quantile(lat_us, 0.25))
q3 = float(np.quantile(lat_us, 0.75))
iqr = q3 - q1

p50 = float(np.quantile(lat_us, 0.50))
p90 = float(np.quantile(lat_us, 0.90))
p95 = float(np.quantile(lat_us, 0.95))
p99 = float(np.quantile(lat_us, 0.99))

mad = float(np.median(np.abs(lat_us - median_us)))

n_trim = int(np.floor(TRIM_FRACTION * n))
lat_sorted_full = np.sort(lat_us)
trimmed_mean = (
    float(np.mean(lat_sorted_full[n_trim:n - n_trim]))
    if n - 2 * n_trim > 0
    else mean_us
)

if CENTRAL_LINE_MODE == "mean":
    central_value = mean_us
    central_label = f"Mean={central_value:.3f} µs"
elif CENTRAL_LINE_MODE == "trimmed_mean":
    central_value = trimmed_mean
    central_label = f"Trimmed Mean ({int(TRIM_FRACTION * 100)}%)={central_value:.3f} µs"
else:
    sys.exit("ERROR: CENTRAL_LINE_MODE must be either 'mean' or 'trimmed_mean'")

print("\nExtended latency statistics (µs):")
print(f"  Samples       : {n}")
print(f"  Mean          : {mean_us:.3f}")
print(f"  Trimmed Mean  : {trimmed_mean:.3f}")
print(f"  Std           : {std_us:.3f}")
print(f"  Median        : {median_us:.3f}")
print(f"  IQR           : {iqr:.3f} (Q1={q1:.3f}, Q3={q3:.3f})")
print(f"  MAD           : {mad:.3f}")
print(f"  P50           : {p50:.3f}")
print(f"  P90           : {p90:.3f}")
print(f"  P95           : {p95:.3f}")
print(f"  P99           : {p99:.3f}")
print(f"  Display Line  : {central_label}")

# ─────────────────────────────────────────────────────────────────────────────
# 4. Bin-width selection (Freedman–Diaconis rule)
# ─────────────────────────────────────────────────────────────────────────────
if iqr > 0:
    bin_width = 2.0 * iqr / (n ** (1.0 / 3.0))
    bins = int(np.ceil((lat_us.max() - lat_us.min()) / bin_width)) if bin_width > 0 else 30
    bins = max(10, min(bins, 200))
else:
    bins = min(max(int(np.sqrt(n)), 10), 200)

# ─────────────────────────────────────────────────────────────────────────────
# 5. Empirical CDF
# ─────────────────────────────────────────────────────────────────────────────
lat_sorted = np.sort(lat_us)
cdf = np.arange(1, n + 1, dtype=np.float64) / n

# ─────────────────────────────────────────────────────────────────────────────
# 6. Plotting
# ─────────────────────────────────────────────────────────────────────────────
try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    _PLOTTING_AVAILABLE = True
except Exception as _e:
    _PLOTTING_AVAILABLE = False
    print(f"\n[WARN] Matplotlib not available ({_e}); skipping plots.")

if _PLOTTING_AVAILABLE:
    # ── 6a) Occurrence histogram ─────────────────────────────────────────────
    fig1, ax = plt.subplots(1, 1, figsize=COUNT_HIST_FIGSIZE)
    ax.hist(
        lat_us,
        bins=bins,
        density=False,
        alpha=HIST_ALPHA_COUNTS,
        color="tab:blue",
        edgecolor=HIST_EDGE_COLOR,
        linewidth=HIST_LINE_WIDTH
    )

    if SHOW_P99:
        ax.axvline(
            p99,
            color=P99_LINE_COLOR,
            linestyle="--",
            linewidth=P99_LINE_WIDTH,
            label=f"P99={p99:.3f} µs"
        )

    if SHOW_CENTRAL_LINE:
        ax.axvline(
            central_value,
            color=CENTRAL_LINE_COLOR,
            linestyle="--",
            linewidth=CENTRAL_LINE_WIDTH,
            label=central_label
        )

    ax.set_title("Latency Histogram (counts)")
    ax.set_xlabel("Latency (µs)")
    ax.set_ylabel("Occurrences")
    ax.grid(True, alpha=0.3)

    if SHOW_P99 or SHOW_CENTRAL_LINE:
        ax.legend(loc="best")

    fig1.tight_layout()
    fig1.savefig(OUT_PNG_COUNTS, dpi=200)
    plt.close(fig1)
    print(f"Saved count histogram to: {OUT_PNG_COUNTS}")

    # ── 6b) Density histogram (PDF-like) + empirical CDF ────────────────────
    fig2, axes = plt.subplots(1, 2, figsize=PDF_CDF_FIGSIZE)

    ax0 = axes[0]
    ax0.hist(
        lat_us,
        bins=bins,
        density=True,
        alpha=HIST_ALPHA_DENSITY,
        color="tab:blue",
        edgecolor=HIST_EDGE_COLOR,
        linewidth=HIST_LINE_WIDTH
    )

    if SHOW_P99:
        ax0.axvline(
            p99,
            color=P99_LINE_COLOR,
            linestyle="--",
            linewidth=P99_LINE_WIDTH,
            label=f"P99={p99:.3f} µs"
        )

    if SHOW_CENTRAL_LINE:
        ax0.axvline(
            central_value,
            color=CENTRAL_LINE_COLOR,
            linestyle="-",
            linewidth=CENTRAL_LINE_WIDTH,
            label=central_label
        )

    ax0.set_title("Latency PDF (density histogram)")
    ax0.set_xlabel("Latency (µs)")
    ax0.set_ylabel("Density")
    ax0.grid(True, alpha=0.3)

    if SHOW_P99 or SHOW_CENTRAL_LINE:
        ax0.legend(loc="best")

    ax1 = axes[1]
    ax1.plot(lat_sorted, cdf, color="tab:green", linewidth=2.2)

    if SHOW_P99:
        ax1.axhline(0.99, color=P99_LINE_COLOR, linestyle="--", linewidth=1.5)
        ax1.axvline(p99, color=P99_LINE_COLOR, linestyle="--", linewidth=1.5)
        ax1.scatter(
            [p99],
            [0.99],
            color=P99_LINE_COLOR,
            s=50,
            zorder=5,
            label=f"P99={p99:.3f} µs"
        )

    if SHOW_CENTRAL_LINE:
        ax1.axvline(
            central_value,
            color=CENTRAL_LINE_COLOR,
            linestyle="-",
            linewidth=CENTRAL_LINE_WIDTH,
            label=central_label
        )

    ax1.set_title("Empirical CDF of Latency")
    ax1.set_xlabel("Latency (µs)")
    ax1.set_ylabel("CDF")
    ax1.set_ylim(0, 1.0)
    ax1.grid(True, alpha=0.3)

    if SHOW_P99 or SHOW_CENTRAL_LINE:
        ax1.legend(loc="best")

    fig2.tight_layout()
    fig2.savefig(OUT_PNG_PDF_CDF, dpi=200)
    plt.close(fig2)
    print(f"Saved density+CDF figure to: {OUT_PNG_PDF_CDF}")
