Features_analysis.py

"""Timestamp contribution heatmap for SAX/ALSAX VSM models.

Given a fitted .pkl model (SAXVSM or ALSAXVSM) and a UCR dataset name,
this script selects one time series and computes how each timestamp
contributes to a class score.

The contribution is derived from the exact decomposition of:
    score(class c) = vec(ts) dot class_vec(c)

where vec(ts) is the normalized bag-of-words vector returned by model.transform.
Each token contribution is evenly distributed over timestamps covered by each
occurrence window.

Example:
  python Empirical_SAX/Features_analysis.py \
    --model-path saved_models/method=paa/win_ratio=0.8/filter_ratio=0.3/alpha=4/Coffee/ALSAXVSM.pkl \
    --dataset Coffee \
    --split test \
    --series-index 0
"""

from __future__ import annotations

import argparse
import os
import re
import sys
from collections import defaultdict
from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import joblib
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import Normalize
from scipy import sparse


@dataclass
class PatternContribution:
    token: str
    count: int
    tf_weight: float
    class_weight: float
    total_contribution: float


def _register_pickle_classes():
    """Register classes in __main__ so joblib can unpickle legacy models."""
    import __main__

    try:
        import Empirical_SAX.parameters_experiment as pc
    except Exception:
        import parameters_experiment as pc  # type: ignore

    __main__.SAXVSM = getattr(pc, "SAXVSM", None)
    __main__.ALSAXVSM = getattr(pc, "ALSAXVSM", None)
    return pc


def _flatten_ts(ts) -> np.ndarray:
    arr = np.asarray(ts, dtype=float)
    if arr.ndim > 1:
        arr = arr.reshape(-1)
    return arr


def _to_dense_1d(vec, expected_len: Optional[int] = None) -> np.ndarray:
    if sparse.issparse(vec):
        arr = vec.toarray()
    else:
        arr = np.asarray(vec)

    if arr.ndim == 2 and 1 in arr.shape:
        arr = arr.reshape(-1)
    elif arr.ndim > 1:
        arr = arr.ravel()

    arr = arr.astype(float, copy=False)

    if expected_len is not None and arr.shape[0] != expected_len:
        raise ValueError(
            f"Unexpected vector size {arr.shape[0]} (expected {expected_len})."
        )

    return arr


def _detect_model_kind(model, override: str) -> str:
    if override in {"sax", "alsax"}:
        return override

    class_name = model.__class__.__name__.lower()
    if "alsax" in class_name:
        return "alsax"
    if "sax" in class_name:
        return "sax"

    bps = getattr(model, "bps", None)
    if isinstance(bps, (list, tuple, np.ndarray)) and len(bps) > 0:
        first = bps[0]
        if isinstance(first, (list, tuple, np.ndarray)):
            return "alsax"

    return "sax"


def _words_from_series(model, pc, ts: np.ndarray, model_kind: str) -> List[str]:
    if len(ts) < int(model.win):
        raise ValueError(
            f"Time series length ({len(ts)}) is smaller than model.win ({model.win})."
        )

    if model_kind == "alsax":
        overflow_mode = getattr(model, "overflow_mode", None)
        if overflow_mode is None:
            infer_mode = getattr(pc, "infer_alsax_overflow_mode", None)
            if callable(infer_mode):
                overflow_mode = infer_mode(getattr(model, "vocab_index", None))
            else:
                overflow_mode = "top"

        return pc.ts_to_boew(
            ts,
            model.win,
            model.paa_size,
            model.alpha,
            model.bps,
            method=model.reduction_method,
            filter_ratio=model.filter_ratio,
            overflow_mode=overflow_mode,
        )

    return pc.ts_to_bow(
        ts,
        model.win,
        model.paa_size,
        model.alpha,
        model.bps,
        method=model.reduction_method,
        filter_ratio=model.filter_ratio,
    )


def _match_label(raw_label: str, available_labels: Iterable) -> object:
    labels = list(available_labels)

    for label in labels:
        if str(label) == raw_label:
            return label

    try:
        numeric_raw = float(raw_label)
    except ValueError:
        numeric_raw = None

    if numeric_raw is not None:
        for label in labels:
            if isinstance(label, (int, float, np.integer, np.floating)):
                if float(label) == numeric_raw:
                    return label

    available = ", ".join(sorted(str(x) for x in labels))
    raise ValueError(
        f"Label '{raw_label}' not found. Available labels are: {available}"
    )


def _select_series(
    X,
    y,
    series_index: int,
    series_label: Optional[str],
) -> Tuple[np.ndarray, object, int]:
    X_arr = np.asarray(X)
    y_arr = np.asarray(y)

    if X_arr.shape[0] != y_arr.shape[0]:
        raise ValueError("X and y must have the same number of samples.")

    if series_label is None:
        if series_index < 0 or series_index >= X_arr.shape[0]:
            raise IndexError(
                f"series-index={series_index} out of bounds [0, {X_arr.shape[0] - 1}]."
            )
        global_idx = int(series_index)
    else:
        label = _match_label(series_label, np.unique(y_arr))
        candidate_indices = np.where(y_arr == label)[0]

        if len(candidate_indices) == 0:
            raise ValueError(
                f"No sample found for label '{series_label}'."
            )

        if series_index < 0 or series_index >= len(candidate_indices):
            raise IndexError(
                "series-index is out of bounds for the filtered class subset. "
                f"Got {series_index}, expected in [0, {len(candidate_indices) - 1}]."
            )

        global_idx = int(candidate_indices[series_index])

    ts = _flatten_ts(X_arr[global_idx])
    true_label = y_arr[global_idx]
    return ts, true_label, global_idx


def compute_timestamp_contributions(
    model,
    pc,
    ts: np.ndarray,
    target_label,
    model_kind: str,
) -> Tuple[np.ndarray, np.ndarray, List[PatternContribution], float, float, int]:
    words = _words_from_series(model, pc, ts, model_kind)

    occurrences: Dict[str, List[int]] = defaultdict(list)
    for start, token in enumerate(words):
        if token in model.vocab_index:
            occurrences[token].append(start)

    matched_windows_total = int(sum(len(starts) for starts in occurrences.values()))

    vocab_size = len(model.vocab)
    ts_vec = _to_dense_1d(model.transform(ts), expected_len=vocab_size)

    if target_label not in model.class_vecs:
        known_classes = ", ".join(str(c) for c in model.class_vecs.keys())
        raise ValueError(
            f"Target class '{target_label}' is not in model.class_vecs. "
            f"Available classes: {known_classes}"
        )

    class_vec = _to_dense_1d(model.class_vecs[target_label], expected_len=vocab_size)

    contribution = np.zeros(len(ts), dtype=float)
    overlap_count = np.zeros(len(ts), dtype=float)
    pattern_rows: List[PatternContribution] = []
    win = int(model.win)

    for token, starts in occurrences.items():
        idx = model.vocab_index[token]
        tf_weight = float(ts_vec[idx])
        class_weight = float(class_vec[idx])
        total_contribution = tf_weight * class_weight

        if total_contribution <= 0:
            continue

        per_occurrence = total_contribution / len(starts)

        for start in starts:
            end = min(start + win, len(ts))
            if end <= start:
                continue

            span = end - start
            contribution[start:end] += per_occurrence / span
            overlap_count[start:end] += 1.0

        pattern_rows.append(
            PatternContribution(
                token=token,
                count=len(starts),
                tf_weight=tf_weight,
                class_weight=class_weight,
                total_contribution=total_contribution,
            )
        )

    pattern_rows.sort(key=lambda row: row.total_contribution, reverse=True)

    class_score = float(np.dot(ts_vec, class_vec))
    explained_score = float(sum(row.total_contribution for row in pattern_rows))

    return (
        contribution,
        overlap_count,
        pattern_rows,
        class_score,
        explained_score,
        matched_windows_total,
    )


def _safe_normalizer(values: np.ndarray) -> Normalize:
    vmin = float(np.min(values))
    vmax = float(np.max(values))
    if np.isclose(vmin, vmax):
        vmax = vmin + 1e-12
    return Normalize(vmin=vmin, vmax=vmax)


def _sanitize(value: object) -> str:
    return re.sub(r"[^0-9A-Za-z_.=-]+", "_", str(value))


def _short_token(token: str, max_len: int = 48) -> str:
    if len(token) <= max_len:
        return token
    return token[: max_len - 3] + "..."


def _extract_model_path_params(model_path: str) -> Dict[str, str]:
    params: Dict[str, str] = {}
    for match in re.finditer(r"([A-Za-z0-9_]+)=([^/\\]+)", model_path):
        key, value = match.group(1), match.group(2)
        params[key] = value
    return params


def _collect_model_metadata(model, model_kind: str, model_path: str) -> Dict[str, object]:
    path_params = _extract_model_path_params(model_path)

    metadata: Dict[str, object] = {
        "class": model.__class__.__name__,
        "kind": model_kind,
        "method": getattr(model, "reduction_method", path_params.get("method")),
        "win_ratio": path_params.get("win_ratio"),
        "win": getattr(model, "win", None),
        "paa": getattr(model, "paa_size", None),
        "alpha": getattr(model, "alpha", path_params.get("alpha")),
        "filter_ratio": getattr(model, "filter_ratio", path_params.get("filter_ratio")),
    }

    if model_kind == "alsax":
        metadata["overflow"] = getattr(model, "overflow_mode", None)

    return metadata


def _format_model_metadata(metadata: Dict[str, object]) -> str:
    keys = [
        "class",
        "kind",
        "method",
        "win_ratio",
        "win",
        "paa",
        "alpha",
        "filter_ratio",
        "overflow",
    ]

    parts: List[str] = []
    for key in keys:
        value = metadata.get(key)
        if value is None:
            continue
        parts.append(f"{key}={value}")

    return ", ".join(parts) if parts else "unknown_model"


def _metadata_for_plot_title(metadata: Dict[str, object]) -> Dict[str, object]:
    plot_metadata = dict(metadata)
    method_raw = plot_metadata.get("method")
    method = str(method_raw).lower() if method_raw is not None else ""

    # Keep only metadata relevant to the selected reduction method in the figure title.
    if method == "fra":
        plot_metadata.pop("paa", None)
    elif method == "paa":
        plot_metadata.pop("filter_ratio", None)

    return plot_metadata


def _model_tag_for_filename(metadata: Dict[str, object]) -> str:
    key_map = [
        ("class", "mdl"),
        ("kind", "kind"),
        ("method", "meth"),
        ("win_ratio", "wr"),
        ("win", "w"),
        ("paa", "paa"),
        ("alpha", "a"),
        ("filter_ratio", "fr"),
        ("overflow", "ov"),
    ]

    parts: List[str] = []
    for src_key, dst_key in key_map:
        value = metadata.get(src_key)
        if value is None:
            continue
        parts.append(f"{dst_key}={_sanitize(value)}")

    return "_".join(parts) if parts else "mdl=unknown"


def _build_default_output_path(
    output_dir: str,
    dataset: str,
    split: str,
    global_index: int,
    target_label,
    method,
    model_tag: str,
    output_ext: str,
) -> str:
    method_dir = f"method={_sanitize(method) if method is not None else 'unknown'}"
    extension = output_ext.lstrip(".") or "pdf"
    filename = (
        f"{_sanitize(dataset)}_{_sanitize(split)}_idx={global_index}_"
        f"class={_sanitize(target_label)}_{model_tag}_heatmap.{extension}"
    )
    return os.path.join(output_dir, method_dir, filename)


def plot_contribution_heatmap(
    ts: np.ndarray,
    contribution: np.ndarray,
    dataset: str,
    split: str,
    global_index: int,
    true_label,
    predicted_label,
    target_label,
    model_title: str,
    output_path: str,
    show_plot: bool,
) -> None:
    if len(ts) == 0:
        raise ValueError("Cannot plot an empty time series.")

    x = np.arange(len(ts), dtype=float)
    fig = plt.figure(figsize=(14, 8), constrained_layout=True)
    grid = fig.add_gridspec(2, 1, height_ratios=[7, 1], hspace=0.15)
    ax_line = fig.add_subplot(grid[0, 0])
    ax_heat = fig.add_subplot(grid[1, 0], sharex=ax_line)

    if len(ts) >= 2:
        points = np.column_stack((x, ts)).reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        segment_contrib = 0.5 * (contribution[:-1] + contribution[1:])
        norm = _safe_normalizer(segment_contrib)

        lc = LineCollection(segments, cmap="YlOrRd", norm=norm)
        lc.set_array(segment_contrib)
        lc.set_linewidth(5.4)
        ax_line.add_collection(lc)
        mappable = lc

        ax_line.set_xlim(0, len(ts) - 1)
    else:
        norm = _safe_normalizer(np.array([contribution[0], contribution[0] + 1e-12]))
        scatter = ax_line.scatter(
            x,
            ts,
            c=contribution,
            cmap="YlOrRd",
            norm=norm,
            s=40,
        )
        mappable = scatter
        ax_line.set_xlim(-0.5, 0.5)

    ax_line.plot(x, ts, color="black", linewidth=5, alpha=0.25)

    y_span = float(np.max(ts) - np.min(ts))
    y_pad = 0.05 * (y_span + 1e-12)
    ax_line.set_ylim(float(np.min(ts) - y_pad), float(np.max(ts) + y_pad))
    ax_line.grid(True, alpha=0.2)
    ax_line.set_ylabel("Value")
    ax_line.set_title(
        f"{dataset} | split={split} | series_idx={global_index} | "
        f"true={true_label} | pred={predicted_label} | explained_class={target_label}\n"
        f"model: {model_title}"
    )

    x_right = len(ts) - 1 if len(ts) > 1 else 1
    ax_heat.imshow(
        contribution[np.newaxis, :],
        aspect="auto",
        cmap="YlOrRd",
        norm=norm,
        extent=(0, x_right, 0, 1),
    )
    ax_heat.set_yticks([])
    ax_heat.set_xlabel("Time index")
    ax_heat.set_title("Timestamp contribution heatmap")

    cbar = fig.colorbar(mappable, ax=[ax_line, ax_heat], fraction=0.03, pad=0.02)
    cbar.set_label("Contribution level")

    out_dir = os.path.dirname(output_path)
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
    fig.savefig(output_path, dpi=180)

    if show_plot:
        plt.show()
    else:
        plt.close(fig)


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Generate a timestamp heatmap for one time series using a fitted "
            "SAXVSM or ALSAXVSM .pkl model."
        )
    )
    parser.add_argument(
        "--model-path",
        required=True,
        help="Path to fitted .pkl model (SAXVSM or ALSAXVSM)",
    )
    parser.add_argument(
        "--dataset",
        required=True,
        help="UCR dataset name (used by aeon load_classification)",
    )
    parser.add_argument(
        "--split",
        choices=["train", "test"],
        default="test",
        help="Dataset split used to pick the time series",
    )
    parser.add_argument(
        "--series-index",
        type=int,
        default=0,
        help=(
            "Index of the selected series. If --series-label is set, index is "
            "relative to the subset with that label."
        ),
    )
    parser.add_argument(
        "--series-label",
        default=None,
        help="Optional label filter to select a series from one class only",
    )
    parser.add_argument(
        "--target-class",
        default=None,
        help="Class for contribution computation. Default: predicted class.",
    )
    parser.add_argument(
        "--model-kind",
        choices=["auto", "sax", "alsax"],
        default="auto",
        help="Force model family if auto-detection is incorrect",
    )
    parser.add_argument(
        "--output-path",
        default=None,
        help="Figure output path. If omitted, a default path is generated.",
    )
    parser.add_argument(
        "--output-dir",
        default="Resultats/pattern_timestamp_heatmaps",
        help="Used only when --output-path is not provided",
    )
    parser.add_argument(
        "--output-format",
        choices=["png", "pdf"],
        default="pdf",
        help="Output format when using --output-dir (default: png)",
    )
    parser.add_argument(
        "--top-k-patterns",
        type=int,
        default=0,
        help=(
            "Number of predominant patterns printed in the textual summary. "
            "Use 0 to show all patterns (default)."
        ),
    )
    parser.add_argument(
        "--show-plot",
        action="store_true",
        help="Display the matplotlib window in addition to saving the figure",
    )
    return parser.parse_args(argv)


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)

    if args.series_index < 0:
        print("Error: --series-index must be >= 0", file=sys.stderr)
        return 2

    if args.top_k_patterns < 0:
        print("Error: --top-k-patterns must be >= 0", file=sys.stderr)
        return 2

    if not os.path.isfile(args.model_path):
        print(f"Error: model file not found: {args.model_path}", file=sys.stderr)
        return 2

    pc = _register_pickle_classes()

    try:
        model = joblib.load(args.model_path)
    except Exception as exc:
        print(f"Error: unable to load model: {exc}", file=sys.stderr)
        return 1

    required_attrs = ["vocab", "vocab_index", "class_vecs", "win", "paa_size", "alpha", "bps"]
    missing = [name for name in required_attrs if not hasattr(model, name)]
    if missing:
        print(
            "Error: model does not expose expected fitted attributes: "
            f"{missing}",
            file=sys.stderr,
        )
        return 1

    X_train, y_train, X_test, y_test = pc.load_dataset(args.dataset)
    if args.split == "train":
        X_split, y_split = X_train, y_train
    else:
        X_split, y_split = X_test, y_test

    try:
        ts, true_label, global_index = _select_series(
            X_split,
            y_split,
            series_index=args.series_index,
            series_label=args.series_label,
        )
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1

    predicted_label = model.predict(ts)

    if args.target_class is None:
        target_label = predicted_label
    else:
        try:
            target_label = _match_label(args.target_class, model.class_vecs.keys())
        except Exception as exc:
            print(f"Error: {exc}", file=sys.stderr)
            return 1

    model_kind = _detect_model_kind(model, args.model_kind)
    model_metadata = _collect_model_metadata(model, model_kind, args.model_path)
    model_title = _format_model_metadata(_metadata_for_plot_title(model_metadata))
    model_tag = _model_tag_for_filename(model_metadata)

    try:
        (
            contribution,
            overlap_count,
            pattern_rows,
            class_score,
            explained_score,
            matched_windows_total,
        ) = compute_timestamp_contributions(
            model=model,
            pc=pc,
            ts=ts,
            target_label=target_label,
            model_kind=model_kind,
        )
    except Exception as exc:
        print(f"Error during contribution computation: {exc}", file=sys.stderr)
        return 1

    output_path = args.output_path
    if output_path is None:
        output_path = _build_default_output_path(
            output_dir=args.output_dir,
            dataset=args.dataset,
            split=args.split,
            global_index=global_index,
            target_label=target_label,
            method=model_metadata.get("method"),
            model_tag=model_tag,
            output_ext=args.output_format,
        )

    try:
        plot_contribution_heatmap(
            ts=ts,
            contribution=contribution,
            dataset=args.dataset,
            split=args.split,
            global_index=global_index,
            true_label=true_label,
            predicted_label=predicted_label,
            target_label=target_label,
            model_title=model_title,
            output_path=output_path,
            show_plot=args.show_plot,
        )
    except Exception as exc:
        print(f"Error during plotting: {exc}", file=sys.stderr)
        return 1

    non_zero_timestamps = int(np.count_nonzero(contribution > 0))

    print("Contribution heatmap generated")
    print(f"  model_path: {args.model_path}")
    print(f"  model: {model_title}")
    print(f"  model_kind: {model_kind}")
    print(f"  dataset: {args.dataset}")
    print(f"  split: {args.split}")
    print(f"  global_series_index: {global_index}")
    print(f"  true_label: {true_label}")
    print(f"  predicted_label: {predicted_label}")
    print(f"  explained_class: {target_label}")
    if args.target_class is None:
        print("  explained_class_source: predicted_label (default)")
    else:
        print("  explained_class_source: user-specified --target-class")
        if str(target_label) != str(predicted_label):
            print(
                "  note: explained_class differs from predicted_label because "
                "--target-class forces the class to explain"
            )
    print(f"  class_score(vec dot class_vec): {class_score:.8f}")
    print(f"  explained_score(sum pattern contributions): {explained_score:.8f}")
    total_windows = max(0, len(ts) - int(model.win) + 1)
    print(f"  matched_windows_with_vocab: {matched_windows_total}/{total_windows}")
    print(f"  contribution_nonzero_timestamps: {non_zero_timestamps}/{len(ts)}")
    print(f"  mean_window_overlap_per_timestamp: {float(np.mean(overlap_count)):.4f}")
    print(f"  output_path: {output_path}")

    if pattern_rows:
        if args.top_k_patterns == 0:
            rows_to_print = pattern_rows
            print(f"  top_predominant_patterns: all ({len(rows_to_print)})")
        else:
            rows_to_print = pattern_rows[: args.top_k_patterns]
            print(
                "  top_predominant_patterns: "
                f"{len(rows_to_print)}/{len(pattern_rows)}"
            )

        for rank, row in enumerate(rows_to_print, start=1):
            print(
                "    "
                f"{rank:>2}. token={_short_token(row.token)} | "
                f"count={row.count} | "
                f"tf_weight={row.tf_weight:.6f} | "
                f"class_weight={row.class_weight:.6f} | "
                f"contribution={row.total_contribution:.8f}"
            )
    else:
        if matched_windows_total == 0:
            print("  top_predominant_patterns: none (no token overlap with model vocabulary)")
            if model_kind == "alsax":
                print("  hint: possible ALSAX tokenizer/model version mismatch (legacy vs current overflow tokens)")
        else:
            print("  top_predominant_patterns: none (token overlap exists but contributions for this class are non-positive)")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())