plot_signal_group_comparison.py

"""
plot_signal_group_comparison.py
================================
For every model architecture found in the timefreq_033_* result folders, produce
two PDF figures:

  {arch}_loss.pdf        — 3×2 grid of training/validation loss curves, one
                           subplot per signal group (AC, AC-table, DC, DC-table,
                           ACDC, all). Best fold for each group is used.

  {arch}_predictions.pdf — 3×2 grid of prediction vs ground-truth scatter/line
                           plots for the same best folds.

Layout (rows × cols):
  Row 0: AC        | AC-table
  Row 1: DC        | DC-table
  Row 2: ACDC      | all

All subplots share the same Y-axis range.
  Loss: clipped to [0, 0.5].
  Predictions: determined from the data (min/max across all groups).

Subplot indexed captions (a–f) with the signal group name are embedded in the
x-axis label.  Legend appears only in the top-right subplot (row 0, col 1).

Typography: Times New Roman, 12.5 pt.

Usage
-----
  python plot_signal_group_comparison.py [--arch ARCH] [--out DIR]

  --arch   Only produce plots for this architecture (default: all found).
  --out    Output directory (default: thesis/plots/signal_groups).
"""

import argparse
import glob
import json
import os
import re
import sys

import matplotlib
matplotlib.use("Agg")                          # no display required
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
import numpy as np
import pandas as pd

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# (display_label, folder_name, file_tag)
GROUP_DEFS = [
    ("AC",       "timefreq_033_AC_sw500_ss500",       "AC"),
    ("AC-table", "timefreq_033_AC_table_sw500_ss500", "AC_table"),
    ("DC",       "timefreq_033_DC_sw500_ss500",       "DC"),
    ("DC-table", "timefreq_033_DC_table_sw500_ss500", "DC_table"),
    ("ACDC",     "timefreq_033_ACDC_sw500_ss500",     "ACDC"),
    ("all",      "timefreq_033_all_sw500_ss500",      "all"),
]

# Subplot layout: [[AC, AC-table], [DC, DC-table], [ACDC, all]]
# Expressed as indices into GROUP_DEFS
LAYOUT_IDX = [[0, 1], [2, 3], [4, 5]]

# Alphabet labels (a)–(f) in reading order
SUBPLOT_LETTERS = list("abcdef")

LOSS_YMAX = 0.5          # hard clip for loss plots
LOSS_YMIN = 0.0

FONT_SIZE   = 12.5
FONT_FAMILY = "Times New Roman"
FIGSIZE     = (9.0, 10.3)   # width × height in inches

# Line colours — shared by both figure types
LOSS_COLORS = {
    "train": "#2166ac",   # blue  (GT in pred plots)
    "val":   "#d6604d",   # red-orange  (Prediction in pred plots)
}

# NASA dataset constants (mirrors train.py)
_PAPER_TRAIN_POOL = [1, 2, 3, 4, 5, 7, 8, 9, 10, 13, 14]
_PAPER_TEST_CASES = [11, 12, 15, 16]


def _case_sequence(fold_num: int) -> list:
    """Return case IDs in the order they appear in alldata_pred for a given LOCO fold.

    alldata_pred order: train cases (sorted, fold excluded) → val case → test cases.
    """
    train = sorted(c for c in _PAPER_TRAIN_POOL if c != fold_num)
    return train + [fold_num] + _PAPER_TEST_CASES


def _detect_case_bounds(gt: np.ndarray, drop: float = 0.03, min_wear: float = 0.05) -> list:
    """Return 0-based indices of the first window of each case.

    A new case starts when GT drops by more than ``drop`` after having been
    above ``min_wear`` (i.e. at least one non-zero-wear run completed).
    Index 0 is always included.
    """
    bounds = [0]
    for i in range(1, len(gt)):
        if gt[i-1] > min_wear and (gt[i-1] - gt[i]) > drop:
            bounds.append(i)
    return bounds


# ---------------------------------------------------------------------------
# Font setup
# ---------------------------------------------------------------------------

def _setup_font():
    """Configure matplotlib to use Times New Roman (or Nimbus Roman as fallback)."""
    # Candidate font directories to register with matplotlib
    _extra_dirs = [
        "/usr/share/fonts/opentype/urw-base35",   # Nimbus Roman (Times clone)
        "/usr/share/fonts/truetype/msttcorefonts", # Times New Roman if installed
    ]
    for d in _extra_dirs:
        if os.path.isdir(d):
            for font_path in glob.glob(os.path.join(d, "*.otf")) + glob.glob(os.path.join(d, "*.ttf")):
                try:
                    fm.fontManager.addfont(font_path)
                except Exception:
                    pass

    available = {f.name for f in fm.fontManager.ttflist}

    if FONT_FAMILY in available:
        chosen = FONT_FAMILY
    elif "Nimbus Roman" in available:
        chosen = "Nimbus Roman"
        print(f"[info] '{FONT_FAMILY}' not found; using 'Nimbus Roman' (Times New Roman clone).")
    else:
        chosen = "DejaVu Serif"
        print(f"[warn] '{FONT_FAMILY}' not found; falling back to 'DejaVu Serif'.")

    plt.rcParams.update({
        "font.family": "serif",
        "font.serif":  [chosen, "DejaVu Serif"],
        "font.size":   FONT_SIZE,
    })


# ---------------------------------------------------------------------------
# Data loading helpers
# ---------------------------------------------------------------------------

def _find_best_fold_files(arch: str, folder: str, tag: str):
    """
    Return (pred_df, history_dict, fold_num) for the best fold of arch in folder.

    The 'best fold' alldata_pred CSV is the single file produced by the
    training pipeline for that arch; the corresponding JSON history is
    resolved from the same fold tag in the filename.
    """
    pattern = os.path.join(BASE_DIR, folder, f"DL_{arch}_{tag}_*_rep0_alldata_pred.csv")
    matches = glob.glob(pattern)
    if not matches:
        return None, None, None

    pred_path = matches[0]  # exactly one file per arch

    # Extract fold number from filename (e.g. "fold10" → 10)
    m = re.search(r"_(fold(\d+))_rep0_alldata_pred", pred_path)
    if not m:
        return None, None, None
    fold     = m.group(1)   # "fold10"
    fold_num = int(m.group(2))  # 10

    # Load predictions
    pred_df = pd.read_csv(pred_path)

    # Load training history JSON
    json_path = os.path.join(BASE_DIR, folder, f"DL_{arch}_{tag}_{fold}_rep0.json")
    history = None
    if os.path.exists(json_path):
        with open(json_path) as fh:
            data = json.load(fh)
        history = data.get("history", {})

    return pred_df, history, fold_num


def _load_all_data(arch: str):
    """
    Load pred_df, history, and fold_num for every group.

    Returns list of (group_label, pred_df_or_None, history_or_None, fold_num_or_None)
    in the order defined by GROUP_DEFS.
    """
    results = []
    for label, folder, tag in GROUP_DEFS:
        pred_df, history, fold_num = _find_best_fold_files(arch, folder, tag)
        results.append((label, pred_df, history, fold_num))
    return results


# ---------------------------------------------------------------------------
# Y-axis range helpers
# ---------------------------------------------------------------------------

def _pred_yrange(group_data):
    """Compute shared [ymin, ymax] from all groups' prediction data.
    Applies the same [0, 0.45] physical clip used when rendering.
    """
    all_vals = []
    for _, pred_df, _, _fn in group_data:
        if pred_df is not None:
            all_vals.extend(pred_df["gt"].tolist())
            # clip predictions the same way _plot_pred does
            clipped = np.clip(pred_df["pred"].values, 0.0, 0.5)
            all_vals.extend(clipped.tolist())
    if not all_vals:
        return 0.0, 0.5
    margin = 0.05 * (max(all_vals) - min(all_vals))
    return min(all_vals) - margin, max(all_vals) + margin


# ---------------------------------------------------------------------------
# Individual subplot renderers
# ---------------------------------------------------------------------------

def _plot_loss(ax, history, letter, group_label, show_legend: bool):
    """Draw training + validation loss on ax."""
    ax.grid(True, color="#d4d4d4", linewidth=0.5, linestyle="-", zorder=0)
    ax.set_axisbelow(True)

    if history is None:
        ax.text(0.5, 0.5, "No data", transform=ax.transAxes,
                ha="center", va="center", fontsize=FONT_SIZE)
        ax.set_xlabel(f"Epoch\n({letter}) {group_label}", fontsize=FONT_SIZE)
        ax.set_ylabel("")
        return

    loss     = history.get("loss", [])
    val_loss = history.get("val_loss", [])
    epochs   = range(1, len(loss) + 1)

    ax.plot(epochs, val_loss, color=LOSS_COLORS["val"],   linewidth=1.0, label="Validation")
    ax.plot(epochs, loss,     color=LOSS_COLORS["train"], linewidth=1.0, label="Training")

    if show_legend:
        ax.legend(fontsize=FONT_SIZE, frameon=True, loc="upper right")

    ax.set_xlabel(f"Epoch\n({letter}) {group_label}", fontsize=FONT_SIZE)


def _reorder_segments(gt: np.ndarray, pred: np.ndarray,
                      bounds: list, case_seq: list) -> tuple:
    """Reorder case segments by ascending case ID, then sort windows within each
    segment by ascending ground-truth VB.

    Returns (gt_out, pred_out, new_bounds, sorted_case_ids).
    """
    ends = bounds[1:] + [len(gt)]
    # Build (case_id, gt_slice, pred_slice) tuples in storage order
    segments = []
    for i, (start, end) in enumerate(zip(bounds, ends)):
        cid = case_seq[i] if case_seq and i < len(case_seq) else i + 1
        gt_seg   = gt[start:end].copy()
        pred_seg = pred[start:end].copy()
        # Sort within segment by ascending VB
        order    = np.argsort(gt_seg, kind="stable")
        segments.append((cid, gt_seg[order], pred_seg[order]))

    # Sort segments by case ID
    segments.sort(key=lambda s: s[0])

    # Concatenate into output arrays and compute new boundary positions
    gt_parts, pred_parts, new_bounds, sorted_ids = [], [], [0], []
    pos = 0
    for cid, gt_seg, pred_seg in segments:
        gt_parts.append(gt_seg)
        pred_parts.append(pred_seg)
        sorted_ids.append(cid)
        pos += len(gt_seg)
        new_bounds.append(pos)

    new_bounds = new_bounds[:-1]   # drop final sentinel (= len of array)
    return (np.concatenate(gt_parts),
            np.concatenate(pred_parts),
            new_bounds,
            sorted_ids)


def _plot_pred(ax, pred_df, fold_num, letter, group_label, show_legend: bool,
               x_label: str = "Case"):
    """Draw ground truth + prediction lines on ax with case-labelled x-axis."""
    # Only horizontal grid lines — vertical separators are the case boundaries
    ax.grid(True, axis="y", color="#d4d4d4", linewidth=0.5, linestyle="-", zorder=0)
    ax.set_axisbelow(True)

    if pred_df is None:
        ax.text(0.5, 0.5, "No data", transform=ax.transAxes,
                ha="center", va="center", fontsize=FONT_SIZE)
        ax.set_xlabel(f"Case\n({letter}) {group_label}", fontsize=FONT_SIZE)
        return

    gt   = pred_df["gt"].values
    pred = pred_df["pred"].values

    # Physical post-processing: VB is bounded to [0, training_cap]
    pred = np.clip(pred, 0.0, 0.45)

    # --- detect storage-order boundaries, then reorder by case ID ---
    bounds   = _detect_case_bounds(gt)
    if fold_num is not None:
        case_seq = _case_sequence(fold_num)
    elif "case_id" in pred_df.columns:
        # Pre-sorted data: extract case ID at the start of each segment
        case_seq = [int(pred_df["case_id"].iloc[b]) for b in bounds]
    else:
        case_seq = None

    gt, pred, bounds, sorted_ids = _reorder_segments(gt, pred, bounds, case_seq)
    ends = bounds[1:] + [len(gt)]

    x = np.arange(len(gt))

    # Use same colours as loss plot: GT → train blue, Prediction → val red-orange
    ax.plot(x, pred, color=LOSS_COLORS["val"],   linewidth=0.8, label="Prediction")
    ax.plot(x, gt,   color=LOSS_COLORS["train"], linewidth=0.8, label="Ground truth")

    # Case boundary lines serve as the x-grid (same style as horizontal grid)
    for b in bounds[1:]:
        ax.axvline(b, color="#d4d4d4", linewidth=0.5, linestyle="-", zorder=0)

    # Place x-ticks at case midpoints, labelled with sorted NASA case IDs
    mids   = [(bounds[i] + ends[i]) / 2 for i in range(len(bounds))]
    labels = [str(c) for c in sorted_ids]

    ax.set_xticks(mids)
    ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=FONT_SIZE - 2)
    ax.tick_params(axis="x", which="both", length=0)   # hide tick marks; labels suffice
    ax.set_xlim(0, len(gt))
    ax.set_ylim(0, 0.5)

    if show_legend:
        ax.legend(fontsize=FONT_SIZE, frameon=True, loc="upper left")

    ax.set_xlabel(f"{x_label}\n({letter}) {group_label}", fontsize=FONT_SIZE)


# ---------------------------------------------------------------------------
# Figure builders
# ---------------------------------------------------------------------------

def _make_figure(arch: str, group_data: list, mode: str, out_dir: str,
                 x_label: str = "Case", loss_ymax: float = LOSS_YMAX):
    """
    Build and save a single figure (mode='loss' or mode='pred').

    group_data: list of (label, pred_df, history) in GROUP_DEFS order.
    loss_ymax:  upper Y limit for loss plots (default LOSS_YMAX=0.5).
    """
    fig, axes = plt.subplots(3, 2, figsize=FIGSIZE)

    # --- determine shared Y range ---
    if mode == "loss":
        ymin, ymax = LOSS_YMIN, loss_ymax
    else:
        ymin, ymax = _pred_yrange(group_data)

    letter_idx  = 0
    top_right_ax = axes[0, 1]   # for legend placement

    for row_idx, row in enumerate(LAYOUT_IDX):
        for col_idx, group_pos in enumerate(row):
            ax        = axes[row_idx, col_idx]
            label, pred_df, history, fold_num = group_data[group_pos]
            letter    = SUBPLOT_LETTERS[letter_idx]
            show_lgnd = (ax is top_right_ax)

            if mode == "loss":
                _plot_loss(ax, history, letter, label, show_lgnd)
                ax.set_ylabel("Loss", fontsize=FONT_SIZE)
            else:
                _plot_pred(ax, pred_df, fold_num, letter, label, show_lgnd,
                           x_label=x_label)
                ax.set_ylabel("VB (mm)", fontsize=FONT_SIZE)

            ax.set_ylim(ymin, ymax)
            ax.tick_params(labelsize=FONT_SIZE - 1)

            letter_idx += 1

    # Suppress duplicate y-labels on right column
    for row_idx in range(3):
        axes[row_idx, 1].set_ylabel("")

    fig.tight_layout(pad=1.2, h_pad=1.6, w_pad=1.2)

    # Save
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"{arch}_{mode}.pdf")
    fig.savefig(out_path, format="pdf", bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {out_path}")


# ---------------------------------------------------------------------------
# Architecture discovery
# ---------------------------------------------------------------------------

def _discover_archs():
    """Return sorted list of architectures that have data in any group."""
    archs = set()
    for _, folder, tag in GROUP_DEFS:
        pattern = os.path.join(BASE_DIR, folder, f"DL_*_{tag}_*_rep0_alldata_pred.csv")
        for f in glob.glob(pattern):
            name = os.path.basename(f)
            # DL_{arch}_{tag}_{fold}_rep0_alldata_pred.csv
            m = re.match(rf"DL_(.+?)_{re.escape(tag)}_fold\d+_rep0_alldata_pred\.csv", name)
            if m:
                archs.add(m.group(1))
    return sorted(archs)


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

def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--arch", default=None,
                        help="Only process this architecture (default: all found).")
    parser.add_argument("--out", default=os.path.join(BASE_DIR, "thesis", "plots", "signal_groups"),
                        help="Output directory for PDF files.")
    args = parser.parse_args()

    _setup_font()

    all_archs = _discover_archs()
    if not all_archs:
        sys.exit("No alldata_pred CSV files found. Check BASE_DIR and folder structure.")

    target_archs = [args.arch] if args.arch else all_archs
    missing = [a for a in target_archs if a not in all_archs]
    if missing:
        sys.exit(f"Architecture(s) not found in any group: {missing}")

    print(f"Found {len(all_archs)} architecture(s). Processing {len(target_archs)}.")
    print(f"Output directory: {args.out}\n")

    for arch in target_archs:
        print(f"[{arch}]")
        group_data = _load_all_data(arch)

        missing_groups = [g[0] for g in group_data if g[1] is None and g[2] is None]
        if missing_groups:
            print(f"  [warn] No data for groups: {missing_groups}")

        _make_figure(arch, group_data, mode="loss", out_dir=args.out)
        _make_figure(arch, group_data, mode="pred", out_dir=args.out)

    print("\nDone.")


if __name__ == "__main__":
    main()