mt/tools/plot_scripts/results_inference_timelines_exp_compare.py

#!/usr/bin/env python3
# results_inference_timelines_exp_compare.py

import json
import pickle
import re
import shutil
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
import polars as pl
from load_results import load_inference_results_dataframe
from matplotlib.lines import Line2D

# =====================================
# User-configurable params
# =====================================

# Root directory that contains per-run outputs (your loader will scan this)
INFERENCE_ROOT = Path("/home/fedex/mt/results/inference/copy")

# Cached stats + manual labels (same location as your earlier scripts)
CACHE_PATH = Path("/home/fedex/mt/plots/results_inference_exp_compare")

# .bag directory (used only to rebuild experiment order for mapping stats)
ALL_DATA_PATH = Path("/home/fedex/mt/data/subter")

# Output base directory (timestamped subfolder will be created here, archived, and copied to latest/)
OUTPUT_PATH = Path("/home/fedex/mt/plots/results_inference_exp_compare")

# Two experiments to compare (exact strings as they appear in your DF’s `experiment` column)
EXPERIMENT_CLEAN = "2_static_no_artifacts_illuminated_2023-01-23-001"
EXPERIMENT_DEGRADED = "3_smoke_human_walking_2023-01-23"

# Shared model configuration for BOTH experiments
LATENT_DIM = 32
SEMI_NORMALS = 0
SEMI_ANOMALOUS = 0

# Comparison y-axis mode for methods: "baseline_z" or "baseline_tailprob"
Y_MODE = "baseline_z"

# Progress axis resolution (number of bins from 0% to 100%)
PROGRESS_BINS = 100

# Frames per second for building time axes before progress-binning (informational only)
FPS = 10.0

# ---- EMA smoothing only ----
# Methods (scores) EMA alpha
EMA_ALPHA_METHODS = 0.1  # (0,1], smaller = smoother
# Stats (absolute %) EMA alpha
EMA_ALPHA_STATS = 0.1  # (0,1], smaller = smoother

# LiDAR points per frame (for stats -> percent)
DATA_RESOLUTION = 32 * 2048

# Copy this script into outputs for provenance (best-effort if not running as a file)
COPY_SELF = True

# =====================================
# Setup output folders
# =====================================
datetime_folder_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
latest_folder_path = OUTPUT_PATH / "latest"
archive_folder_path = OUTPUT_PATH / "archive"
output_datetime_path = OUTPUT_PATH / datetime_folder_name

OUTPUT_PATH.mkdir(exist_ok=True, parents=True)
archive_folder_path.mkdir(exist_ok=True, parents=True)
latest_folder_path.mkdir(exist_ok=True, parents=True)
output_datetime_path.mkdir(exist_ok=True, parents=True)

# =====================================
# Load Polars DataFrame via your helper
# =====================================

df: pl.DataFrame = load_inference_results_dataframe(INFERENCE_ROOT)

required_cols = {
    "experiment",
    "network",
    "latent_dim",
    "semi_normals",
    "semi_anomalous",
    "model",
    "scores",
    "folder",
    "config_json",
}
missing = required_cols - set(df.columns)
if missing:
    raise KeyError(f"DataFrame missing required columns: {sorted(missing)}")


# =====================================
# Rebuild experiment → stats mapping (like your original)
# =====================================
def rebuild_experiment_index():
    normals, anomalies = [], []
    if not ALL_DATA_PATH.exists():
        return [], [], {}
    for bag in ALL_DATA_PATH.iterdir():
        if bag.suffix != ".bag":
            continue
        if "smoke" in bag.name:
            anomalies.append(bag)
        else:
            normals.append(bag)
    normals = sorted(normals, key=lambda p: p.stat().st_size)
    anomalies = sorted(anomalies, key=lambda p: p.stat().st_size)
    mapping = {}
    for i, p in enumerate(normals):
        mapping[p.stem] = (False, i, p)
    for i, p in enumerate(anomalies):
        mapping[p.stem] = (True, i, p)
    return normals, anomalies, mapping


normal_paths, anomaly_paths, exp_map = rebuild_experiment_index()

# Load cached statistical data and manual labels
missing_points_cache = CACHE_PATH / "missing_points.pkl"
near_sensor_cache = CACHE_PATH / "particles_near_sensor_counts_500.pkl"
labels_json_path = CACHE_PATH / "manually_labeled_anomaly_frames.json"

missing_points_normal = missing_points_anomaly = None
near_sensor_normal = near_sensor_anomaly = None
if missing_points_cache.exists():
    with open(missing_points_cache, "rb") as f:
        missing_points_normal, missing_points_anomaly = pickle.load(f)
if near_sensor_cache.exists():
    with open(near_sensor_cache, "rb") as f:
        near_sensor_normal, near_sensor_anomaly = pickle.load(f)

manual_windows = {}
if labels_json_path.exists():
    with open(labels_json_path, "r") as f:
        labeled_json = json.load(f)
        for file in labeled_json.get("files", []):
            manual_windows[file["filename"]] = (
                file.get("semi_target_begin_frame"),
                file.get("semi_target_end_frame"),
            )


# =====================================
# Helpers
# =====================================
def ema(x: np.ndarray, alpha: float) -> np.ndarray:
    if x is None or len(x) == 0:
        return x
    y = np.empty_like(x, dtype=float)
    y[0] = x[0]
    for i in range(1, len(x)):
        y[i] = alpha * x[i] + (1 - alpha) * y[i - 1]
    return y


def to_np_list(list_cell) -> Optional[np.ndarray]:
    if list_cell is None:
        return None
    return np.asarray(list_cell, dtype=float).ravel()


def normalize_exp_name(name: str) -> str:
    # strip trailing run suffix like -001, -002 if present
    return re.sub(r"-\d{3}$", "", name)


def map_experiment_to_stats_stem(exp_name: str) -> Optional[str]:
    """Try exact match, then prefix match with / without -### suffix stripped."""
    if exp_name in exp_map:
        return exp_name
    base = normalize_exp_name(exp_name)
    if base in exp_map:
        return base
    for stem in exp_map.keys():
        if stem.startswith(exp_name) or stem.startswith(base):
            return stem
    return None


def get_stats_for_experiment(
    exp_name: str,
) -> Tuple[
    Optional[np.ndarray], Optional[np.ndarray], Tuple[Optional[int], Optional[int]]
]:
    key = map_experiment_to_stats_stem(exp_name)
    if key is None:
        return None, None, (None, None)
    is_anomaly, idx, path = exp_map[key]
    missing = near = None
    if missing_points_normal is not None and missing_points_anomaly is not None:
        series = (
            missing_points_anomaly[idx] if is_anomaly else missing_points_normal[idx]
        )
        missing = (np.asarray(series, dtype=float) / DATA_RESOLUTION) * 100.0
    if near_sensor_normal is not None and near_sensor_anomaly is not None:
        series = near_sensor_anomaly[idx] if is_anomaly else near_sensor_normal[idx]
        near = (np.asarray(series, dtype=float) / DATA_RESOLUTION) * 100.0
    npy_key = path.with_suffix(".npy").name
    window = manual_windows.get(npy_key, (None, None))
    return missing, near, window


def _bin_to_progress(x: np.ndarray, bins: int = PROGRESS_BINS) -> np.ndarray:
    """Average x into fixed #bins across its length (progress-normalized timeline)."""
    if x is None or len(x) == 0:
        return x
    n = len(x)
    edges = np.linspace(0, n, bins + 1, dtype=int)
    out = np.empty(bins, dtype=float)
    for i in range(bins):
        a, b = edges[i], edges[i + 1]
        if b <= a:
            out[i] = out[i - 1] if i > 0 else x[0]
        else:
            out[i] = float(np.mean(x[a:b]))
    return out


def _ecdf(x: np.ndarray):
    xs = np.sort(np.asarray(x, dtype=float))
    n = len(xs)

    def F(t):
        return float(np.searchsorted(xs, t, side="right")) / n

    return F


def baseline_transform(clean: np.ndarray, other: np.ndarray, mode: str):
    """Transform using stats from clean only."""
    assert mode in ("baseline_z", "baseline_tailprob")
    if clean is None or len(clean) == 0:
        return clean, other, "raw"
    if mode == "baseline_z":
        mu = float(np.mean(clean))
        sd = float(np.std(clean, ddof=0))
        if sd < 1e-12:
            zc = clean - mu
            zo = other - mu if other is not None else None
        else:
            zc = (clean - mu) / sd
            zo = (other - mu) / sd if other is not None else None
        return zc, zo, "Anomaly score (σ above clean)"
    else:
        F = _ecdf(clean)
        tp_clean = np.array([1.0 - F(v) for v in clean], dtype=float)
        tp_other = (
            np.array([1.0 - F(v) for v in other], dtype=float)
            if other is not None
            else None
        )
        return tp_clean, tp_other, "Tail probability vs clean (1 - F_clean)"


def pick_method_series(gdf: pl.DataFrame, label: str) -> Optional[np.ndarray]:
    if label == "DeepSAD LeNet":
        sel = gdf.filter(
            (pl.col("network") == "subter_LeNet") & (pl.col("model") == "deepsad")
        )
    elif label == "DeepSAD Efficient":
        sel = gdf.filter(
            (pl.col("network") == "subter_efficient") & (pl.col("model") == "deepsad")
        )
    elif label == "OCSVM":
        sel = gdf.filter(
            (pl.col("network") == "subter_LeNet") & (pl.col("model") == "ocsvm")
        )
    elif label == "Isolation Forest":
        sel = gdf.filter(
            (pl.col("network") == "subter_LeNet") & (pl.col("model") == "isoforest")
        )
    else:
        sel = pl.DataFrame()
    if sel.height == 0:
        return None
    row = sel.row(0)
    row_dict = {c: row[i] for i, c in enumerate(sel.columns)}
    return to_np_list(row_dict["scores"])


def group_slice(
    df: pl.DataFrame,
    experiment: str,
    latent_dim: int,
    semi_normals: int,
    semi_anomalous: int,
) -> pl.DataFrame:
    return df.filter(
        (pl.col("experiment") == experiment)
        & (pl.col("latent_dim") == latent_dim)
        & (pl.col("semi_normals") == semi_normals)
        & (pl.col("semi_anomalous") == semi_anomalous)
    )


def compare_two_experiments_progress(
    df: pl.DataFrame,
    experiment_clean: str,
    experiment_degraded: str,
    latent_dim: int,
    semi_normals: int,
    semi_anomalous: int,
    y_mode: str = "baseline_z",
    include_stats: bool = True,
):
    methods = [
        "DeepSAD LeNet",
        "DeepSAD Efficient",
        "OCSVM",
        "Isolation Forest",
    ]

    g_clean = group_slice(
        df, experiment_clean, latent_dim, semi_normals, semi_anomalous
    )
    g_deg = group_slice(
        df, experiment_degraded, latent_dim, semi_normals, semi_anomalous
    )
    if g_clean.is_empty() or g_deg.is_empty():
        print(
            f"[WARN] Missing one of the experiment groups: clean({g_clean.height}), degraded({g_deg.height}). Skipping."
        )
        return 0

    # Stats (% absolute, EMA smoothed later)
    mp_clean, ns_clean, _ = get_stats_for_experiment(experiment_clean)
    mp_deg, ns_deg, _ = get_stats_for_experiment(experiment_degraded)

    # Build baseline-anchored, progress-binned curves per method
    curves_clean: Dict[str, np.ndarray] = {}
    curves_deg: Dict[str, np.ndarray] = {}
    y_label = "Anomaly"

    for label in methods:
        s_clean = pick_method_series(g_clean, label)
        s_deg = pick_method_series(g_deg, label)
        if s_clean is None or s_deg is None:
            continue

        # Smooth raw with EMA for stability before fitting baseline
        s_clean_sm = ema(s_clean.astype(float), EMA_ALPHA_METHODS)
        s_deg_sm = ema(s_deg.astype(float), EMA_ALPHA_METHODS)

        t_clean, t_deg, y_label = baseline_transform(s_clean_sm, s_deg_sm, y_mode)

        # Progress-bin both
        curves_clean[label] = _bin_to_progress(t_clean, PROGRESS_BINS)
        curves_deg[label] = _bin_to_progress(t_deg, PROGRESS_BINS)

    if not curves_clean:
        print("[WARN] No method curves available for comparison in this config.")
        return 0

    x = np.linspace(0, 100, PROGRESS_BINS)

    # Prep stats: absolute %, EMA, progress-binned
    def prep_stat_pair(a, b):
        if a is None or len(a) == 0 or b is None or len(b) == 0:
            return None, None
        a_s = ema(a.astype(float), EMA_ALPHA_STATS)
        b_s = ema(b.astype(float), EMA_ALPHA_STATS)
        return _bin_to_progress(a_s, PROGRESS_BINS), _bin_to_progress(
            b_s, PROGRESS_BINS
        )

    mp_c, mp_d = prep_stat_pair(mp_clean, mp_deg)
    ns_c, ns_d = prep_stat_pair(ns_clean, ns_deg)

    # Colors & styles
    COLOR_METHOD = "#d62728"  # vibrant red
    COLOR_MISSING = "#9ecae1"  # pale blue
    COLOR_NEAR = "#a1d99b"  # pale green

    LS_CLEAN = "--"  # dashed for normal/clean
    LS_DEG = "-"  # solid for anomalous/degraded
    LW_METHOD = 1.8
    LW_METHOD_CLEAN = 1.2
    LW_STATS = 1.6
    ALPHA_STATS = 0.95

    # Build the 2x2 subplots
    fig, axes = plt.subplots(
        4, 1, figsize=(12, 16), constrained_layout=True, sharex=False
    )
    axes = axes.ravel()

    method_to_axidx = {
        "DeepSAD LeNet": 0,
        "DeepSAD Efficient": 1,
        "OCSVM": 2,
        "Isolation Forest": 3,
    }

    stats_available = (
        mp_c is not None and mp_d is not None and ns_c is not None and ns_d is not None
    )
    if not stats_available:
        print("[WARN] One or both stats missing. Subplots will include methods only.")

    letters = ["a", "b", "c", "d"]

    for label, axidx in method_to_axidx.items():
        ax = axes[axidx]
        yc = curves_clean.get(label)
        yd = curves_deg.get(label)
        if yc is None or yd is None:
            ax.text(
                0.5, 0.5, "No data", ha="center", va="center", transform=ax.transAxes
            )
            ax.set_title(f"({letters[axidx]}) {label}")
            ax.grid(True, alpha=0.3)
            continue

        # Left axis: method score (z or tailprob)
        ax.plot(
            x,
            yd,
            linestyle=LS_DEG,
            color=COLOR_METHOD,
            linewidth=LW_METHOD,
            label=f"{label} — degraded",
        )
        ax.plot(
            x,
            yc,
            linestyle=LS_CLEAN,
            color=COLOR_METHOD,
            linewidth=LW_METHOD_CLEAN,
            label=f"{label} — clean",
        )
        ax.set_ylabel(y_label)
        ax.set_title(label)
        ax.set_title(f"({letters[axidx]}) {label}")
        ax.grid(True, alpha=0.3)

        # Right axis #1 (closest to plot): Missing points (%)
        axy_miss = ax.twinx()
        if mp_c is not None and mp_d is not None:
            axy_miss.plot(
                x,
                mp_d,
                linestyle=LS_DEG,
                color=COLOR_MISSING,
                alpha=ALPHA_STATS,
                linewidth=LW_STATS,
                label="Missing points — degraded (%)",
            )
            axy_miss.plot(
                x,
                mp_c,
                linestyle=LS_CLEAN,
                color=COLOR_MISSING,
                alpha=ALPHA_STATS,
                linewidth=LW_STATS,
                label="Missing points — clean (%)",
            )
        axy_miss.set_ylabel("Missing points (%)")
        axy_miss.tick_params(axis="y")  # , colors=COLOR_MISSING)
        # axy_miss.spines["right"].set_edgecolor(COLOR_MISSING)

        # Right axis #2 (slightly offset): Near-sensor points (%)
        axy_near = ax.twinx()
        # push this spine outward so it doesn't overlap the first right axis
        axy_near.spines["right"].set_position(("axes", 1.08))
        # make patch invisible so only spine shows
        axy_near.set_frame_on(True)
        axy_near.patch.set_visible(False)

        if ns_c is not None and ns_d is not None:
            axy_near.plot(
                x,
                ns_d,
                linestyle=LS_DEG,
                color=COLOR_NEAR,
                alpha=ALPHA_STATS,
                linewidth=LW_STATS,
                label="Near-sensor — degraded (%)",
            )
            axy_near.plot(
                x,
                ns_c,
                linestyle=LS_CLEAN,
                color=COLOR_NEAR,
                alpha=ALPHA_STATS,
                linewidth=LW_STATS,
                label="Near-sensor — clean (%)",
            )
        axy_near.set_ylabel("Near-sensor points (%)")
        axy_near.tick_params(axis="y")  # , colors=COLOR_NEAR)
        # axy_near.spines["right"].set_edgecolor(COLOR_NEAR)

        # Compose legend: show *method name* explicitly, plus the two stats
        handles = [
            Line2D(
                [0],
                [0],
                color=COLOR_METHOD,
                lw=LW_METHOD,
                ls=LS_DEG,
                label=f"{label} — degraded",
            ),
            Line2D(
                [0],
                [0],
                color=COLOR_METHOD,
                lw=LW_METHOD_CLEAN,
                ls=LS_CLEAN,
                label=f"{label} — clean",
            ),
            Line2D(
                [0],
                [0],
                color=COLOR_MISSING,
                lw=LW_STATS,
                ls=LS_DEG,
                label="Missing points — degraded",
            ),
            Line2D(
                [0],
                [0],
                color=COLOR_MISSING,
                lw=LW_STATS,
                ls=LS_CLEAN,
                label="Missing points — clean",
            ),
            Line2D(
                [0],
                [0],
                color=COLOR_NEAR,
                lw=LW_STATS,
                ls=LS_DEG,
                label="Near-sensor — degraded",
            ),
            Line2D(
                [0],
                [0],
                color=COLOR_NEAR,
                lw=LW_STATS,
                ls=LS_CLEAN,
                label="Near-sensor — clean",
            ),
        ]
        ax.legend(handles=handles, loc="upper left", fontsize=9, framealpha=0.9)

    # Shared labels / super-title
    for ax in axes:
        ax.set_xlabel("Progress through experiment (%)")

    # fig.suptitle(
    #     f"AD Method vs Stats Inference — progress-normalized\n"
    #     f"Transform: z-score normalized to non-degraded experiment | EMA(α={EMA_ALPHA_METHODS})",
    #     fontsize=14,
    # )
    fig.tight_layout(rect=[0, 0, 1, 0.99])

    out_name = (
        f"4up_{EXPERIMENT_CLEAN}_vs_{EXPERIMENT_DEGRADED}"
        f"_ld{latent_dim}_sn{semi_normals}_sa{semi_anomalous}_{y_mode}_methods_vs_stats.png"
    )
    fig.savefig(output_datetime_path / out_name, dpi=150)
    plt.close(fig)

    return 1


# =====================================
# Run comparison & save
# =====================================
plots_made = compare_two_experiments_progress(
    df=df,
    experiment_clean=EXPERIMENT_CLEAN,
    experiment_degraded=EXPERIMENT_DEGRADED,
    latent_dim=LATENT_DIM,
    semi_normals=SEMI_NORMALS,
    semi_anomalous=SEMI_ANOMALOUS,
    y_mode=Y_MODE,
    include_stats=True,
)

# =====================================
# Preserve latest/, archive/, copy script
# =====================================
# delete current latest folder
shutil.rmtree(latest_folder_path, ignore_errors=True)
# create new latest folder
latest_folder_path.mkdir(exist_ok=True, parents=True)

# copy contents of output folder to the latest folder
for file in output_datetime_path.iterdir():
    shutil.copy2(file, latest_folder_path)

# copy this python script to preserve the code used (best effort)
if COPY_SELF:
    try:
        shutil.copy2(__file__, output_datetime_path)
        shutil.copy2(__file__, latest_folder_path)
    except Exception:
        (output_datetime_path / "run_config.json").write_text(
            json.dumps(
                {
                    "INFERENCE_ROOT": str(INFERENCE_ROOT),
                    "CACHE_PATH": str(CACHE_PATH),
                    "ALL_DATA_PATH": str(ALL_DATA_PATH),
                    "EXPERIMENT_CLEAN": EXPERIMENT_CLEAN,
                    "EXPERIMENT_DEGRADED": EXPERIMENT_DEGRADED,
                    "LATENT_DIM": LATENT_DIM,
                    "SEMI_NORMALS": SEMI_NORMALS,
                    "SEMI_ANOMALOUS": SEMI_ANOMALOUS,
                    "Y_MODE": Y_MODE,
                    "PROGRESS_BINS": PROGRESS_BINS,
                    "FPS": FPS,
                    "EMA_ALPHA_METHODS": EMA_ALPHA_METHODS,
                    "EMA_ALPHA_STATS": EMA_ALPHA_STATS,
                    "DATA_RESOLUTION": DATA_RESOLUTION,
                    "timestamp": datetime_folder_name,
                },
                indent=2,
            )
        )

# move output date folder to archive
shutil.move(output_datetime_path, archive_folder_path)

print(f"Done. Wrote {plots_made} figure(s). Archived under: {archive_folder_path}")