transfer_learning.py

"""
transfer_learning.py
====================
Fine-tune a pre-trained binding-change CNN to your own TF/protein dataset.

The pipeline trains two binary classifiers in sequence:
  1. Change vs No-Change  — does binding change at all?
  2. Gain vs Loss         — if it changes, is it a gain or a loss?

At prediction time the two models are combined to produce 3-class labels:
  0 = Gain, 1 = Loss, 2 = No Change

QUICK START
-----------
  python transfer_learning.py \
      --tf MyTF \
      --data Resources/MyTF_Transfer_data.tsv \
      --base_gain_loss Models/BenchPlus_GainvsLoss_General.keras \
      --base_change    Models/BenchPlus_ChangevsNone_General.keras

  # Predict only (skip training, load saved transfer models)
  python transfer_learning.py \
      --tf MyTF \
      --data Resources/MyTF_Transfer_data.tsv \
      --base_gain_loss Models/BenchPlus_GainvsLoss_General.keras \
      --base_change    Models/BenchPlus_ChangevsNone_General.keras \
      --predict_only

EXPERT OVERRIDES (full list below)
------------------------------------
  --target_n         Resampling target per class during training (default 3100)
  --transfer_epochs  Max epochs for transfer learning (default 100)
  --batch_size       Mini-batch size (default 32)
  --patience         Early-stopping patience in epochs (default 10)
  --max_trials       Keras Tuner Bayesian trials (default 50)
  --layers_to_keep   Number of base-model layers to unfreeze at the end (default 2)
  --no_gvat          Skip adding the GVAT normalised data during training
  --no_tune          Skip hyperparameter search; use default HP values instead
"""

# ── Standard library ──────────────────────────────────────────────────────────
import argparse
import os
import pickle
import random
import sys
from collections import Counter
from pathlib import Path

# ── Third-party ───────────────────────────────────────────────────────────────
import numpy as np
import pandas as pd
import tensorflow as tf
import keras_tuner as kt
from sklearn.metrics import (
    classification_report,
    confusion_matrix,
    precision_recall_curve,
)
from sklearn.model_selection import train_test_split
from sklearn.utils import resample, shuffle
from sklearn.utils.class_weight import compute_class_weight
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam, AdamW
import tensorflow.keras.backend as K

# ── Reproducibility ───────────────────────────────────────────────────────────
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
tf.random.set_seed(SEED)
os.environ["PYTHONHASHSEED"] = str(SEED)


# =============================================================================
# SECTION 1 — COMMAND-LINE INTERFACE
# =============================================================================

def parse_args():
    """Parse command-line arguments with sensible defaults."""
    p = argparse.ArgumentParser(
        description="Fine-tune a binding-change CNN to a new transcription factor.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # ── Required ──────────────────────────────────────────────────────────────
    required = p.add_argument_group("required arguments")
    required.add_argument("--tf",   required=True,
                          help="Transcription factor name / run identifier (e.g. GATA1)")
    required.add_argument("--data", required=True,
                          help="Path to the primary training TSV file")
    required.add_argument("--base_gain_loss", required=True,
                          help="Path to the base Gain-vs-Loss .keras model")
    required.add_argument("--base_change",    required=True,
                          help="Path to the base Change-vs-NoChange .keras model")

    # ── Data options ──────────────────────────────────────────────────────────
    data_grp = p.add_argument_group("data options")
    data_grp.add_argument("--gvat_data",
                          help="Path to an additional GVAT-normalised TSV to merge "
                               "before training.  Skipped silently when absent.")
    data_grp.add_argument("--adastra_test",
                          help="Path to an AdAstra test-set TSV to enrich the "
                               "balanced test split.  Falls back to a placeholder "
                               "when absent.")
    data_grp.add_argument("--no_gvat", action="store_true",
                          help="Ignore the GVAT addition even if the file exists.")
    data_grp.add_argument("--test_per_class", type=int, default=400,
                          help="Max samples per class when building the test split.")

    # ── Resampling ────────────────────────────────────────────────────────────
    resamp_grp = p.add_argument_group("resampling options")
    resamp_grp.add_argument("--target_n", type=int, default=3100,
                            help="Target number of samples per class after up/down-sampling.")

    # ── Tuner options ─────────────────────────────────────────────────────────
    tune_grp = p.add_argument_group("hyperparameter tuning options")
    tune_grp.add_argument("--no_tune", action="store_true",
                          help="Skip the Bayesian HP search; reuse cached results "
                               "if available, otherwise use default hyperparameters.")
    tune_grp.add_argument("--max_trials", type=int, default=50,
                          help="Number of Bayesian Optimisation trials per tuner.")

    # ── Training options ──────────────────────────────────────────────────────
    train_grp = p.add_argument_group("training options")
    train_grp.add_argument("--transfer_epochs", type=int, default=100,
                           help="Maximum training epochs for transfer learning.")
    train_grp.add_argument("--batch_size", type=int, default=32,
                           help="Mini-batch size.")
    train_grp.add_argument("--patience", type=int, default=10,
                           help="Early-stopping patience (epochs without improvement).")
    train_grp.add_argument("--layers_to_keep", type=int, default=2,
                           help="Number of base-model layers (from the end) to keep "
                                "trainable during transfer learning.")

    # ── Prediction / evaluation ───────────────────────────────────────────────
    pred_grp = p.add_argument_group("prediction / evaluation options")
    pred_grp.add_argument("--predict_only", action="store_true",
                          help="Skip training; load existing transfer models and "
                               "evaluate on the test set.")
    pred_grp.add_argument("--output_dir", default="Models/Transfer",
                          help="Directory where trained transfer models are saved.")
    pred_grp.add_argument("--tuning_dir", default="tuning_results",
                          help="Directory used by Keras Tuner to cache trial results.")

    return p.parse_args()


# =============================================================================
# SECTION 2 — FEATURE ENGINEERING
# =============================================================================

def create_difference_features(dataframe: pd.DataFrame) -> pd.DataFrame:
    """
    Compute per-position Reference-minus-Mutant difference features.

    The input dataframe is expected to have 112 feature columns followed by
    a Label column and an ID column.  Features are arranged in groups of 7
    (one per nucleotide position window), alternating Reference / Mutant rows.

    Returns a new dataframe with 56 difference columns + Label + ID.
    """
    # Build index lists for reference and mutant columns
    reference_cols, mutant_cols = [], []
    column_groups = [list(range(start, start + 7)) for start in range(0, 112, 7)]
    for i in range(0, 16, 2):
        reference_cols.extend(column_groups[i])
    for i in range(1, 16, 2):
        mutant_cols.extend(column_groups[i])

    diff_data = {}
    for idx, (ref_col, mut_col) in enumerate(zip(reference_cols, mutant_cols)):
        diff_data[f"Feature_{idx}"] = (
            dataframe[f"Feature_{ref_col}"] - dataframe[f"Feature_{mut_col}"]
        )

    result = pd.DataFrame(diff_data)
    result["Label"] = dataframe.iloc[:, 112].values
    result["ID"]    = dataframe.iloc[:, 113].values
    return result


# =============================================================================
# SECTION 3 — DATA LOADING & SPLITTING
# =============================================================================

def load_and_prepare_data(args):
    """
    Load TSVs, compute difference features, and return a balanced test split
    plus train/val arrays ready for model input.

    Returns
    -------
    X_train, X_val, X_test : np.ndarray  shape (N, 56, 1)
    y_train, y_val, y_test : np.ndarray  shape (N,)
    """
    print(f"\n[Data] Loading primary data from: {args.data}")
    df = pd.read_csv(args.data, sep="\t")

    # Optionally merge GVAT-normalised supplement
    if not args.no_gvat:
        gvat_path = (
            args.gvat_data
            or f"Resources/Processed_Multi_{args.tf}_GVAT_Normalized_NNinput.tsv"
        )
        if Path(gvat_path).exists():
            print(f"[Data] Merging GVAT data from: {gvat_path}")
            gvat_df = pd.read_csv(gvat_path, sep="\t")
            df = pd.concat([df, gvat_df], ignore_index=True)
            print(f"[Data] Combined shape after GVAT merge: {df.shape}")
        else:
            print(f"[Data] GVAT file not found ({gvat_path}); skipping merge.")

    # Compute difference features
    df = create_difference_features(df)

    label_col = df.columns[56]   # "Label"
    id_col    = df.columns[57]   # "ID"

    # ── Build balanced test split ─────────────────────────────────────────────
    per_class_n = min(
        len(df[df[label_col] == c]) for c in [0, 1, 2]
    )
    per_class_n = min(per_class_n, args.test_per_class)

    test_parts = [
        df[df[label_col] == c].sample(per_class_n, random_state=SEED)
        for c in [0, 1, 2]
    ]
    test_df = shuffle(pd.concat(test_parts), random_state=SEED)
    test_ids = set(test_df[id_col])

    # Try to enrich the test set with AdAstra examples
    adastra_path = (
        args.adastra_test
        or f"Resources/AdAstraTest{args.tf}.tsv"
    )
    fallback_path = "Resources/AdAstraTestPlaceholder.tsv"

    for path in [adastra_path, fallback_path]:
        if Path(path).exists():
            print(f"[Data] Enriching test set with AdAstra data from: {path}")
            extra = create_difference_features(pd.read_csv(path, sep="\t"))
            test_df = _merge_adastra_into_test(test_df, extra, label_col, SEED)
            break
    else:
        print("[Data] No AdAstra file found; using plain balanced test set.")

    # ── Feature / label arrays ────────────────────────────────────────────────
    feature_cols = list(range(56))
    X_test = test_df.iloc[:, :56].values[:, feature_cols]
    y_test = test_df.iloc[:, 56].values

    print(f"[Data] Test set: {X_test.shape}  |  class counts: {Counter(y_test)}")

    # Remove test IDs from the rest
    keep_mask = ~df[id_col].isin(test_ids)
    X_rest = df.iloc[:, :56].values[keep_mask][:, feature_cols]
    y_rest = df.iloc[:, 56].values[keep_mask]

    X_train_raw, X_val_raw, y_train_raw, y_val_raw = train_test_split(
        X_rest, y_rest, test_size=0.1, random_state=SEED
    )

    # ── Up/down-sample training set to balance classes ─────────────────────────
    X_train_raw, y_train_raw = _balance_classes(
        X_train_raw, y_train_raw, target_n=args.target_n
    )

    # ── Carve out a balanced validation subset from training data ─────────────
    X_train, y_train, X_val, y_val = _create_balanced_val(
        X_train_raw, y_train_raw
    )

    # Reshape to (N, 56, 1) for Conv1D
    X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
    X_val   = X_val.reshape(  X_val.shape[0],   X_val.shape[1],   1)
    X_test  = X_test.reshape( X_test.shape[0],  X_test.shape[1],  1)

    print(f"[Data] Train: {X_train.shape}  |  Val: {X_val.shape}  |  "
          f"Test: {X_test.shape}")
    print(f"[Data] Train class counts: {Counter(y_train)}")
    print(f"[Data] Val   class counts: {Counter(y_val)}")

    return X_train, X_val, X_test, y_train, y_val, y_test


def _merge_adastra_into_test(
    original_test: pd.DataFrame,
    extra: pd.DataFrame,
    label_col: str,
    seed: int,
) -> pd.DataFrame:
    """Merge AdAstra examples into the test set and re-balance across classes."""
    combined = {}
    for cls in [0, 1]:
        combined[cls] = pd.concat(
            [original_test[original_test[label_col] == cls],
             extra[extra[label_col] == cls]],
            ignore_index=True,
        )
    combined[2] = original_test[original_test[label_col] == 2].copy()

    final_n = min(len(v) for v in combined.values())
    parts = [combined[c].sample(final_n, random_state=seed) for c in [0, 1, 2]]
    return shuffle(pd.concat(parts, ignore_index=True), random_state=seed)


def _balance_classes(
    X: np.ndarray, y: np.ndarray, target_n: int
) -> tuple[np.ndarray, np.ndarray]:
    """
    Upsample minority classes and downsample the majority class so that every
    class has exactly `target_n` samples.
    """
    parts_X, parts_y = [], []
    for cls in np.unique(y):
        idx = np.where(y == cls)[0]
        X_cls, y_cls = resample(
            X[idx], y[idx],
            replace=(len(idx) < target_n),
            n_samples=target_n,
            random_state=SEED,
        )
        parts_X.append(X_cls)
        parts_y.append(y_cls)

    X_bal, y_bal = shuffle(
        np.concatenate(parts_X), np.concatenate(parts_y), random_state=SEED
    )
    print(f"[Data] After resampling: {X_bal.shape}  |  "
          f"class counts: {Counter(y_bal)}")
    return X_bal, y_bal


def _create_balanced_val(
    X: np.ndarray, y: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Split off a small balanced validation set from the training data.
    Each class contributes 10 % / n_classes of the total training size.
    """
    Xy = list(zip(X, y))
    by_class = {c: [(x, lbl) for x, lbl in Xy if lbl == c] for c in np.unique(y)}
    val_per_class = min(
        len(v) for v in by_class.values()
    )
    val_per_class = min(val_per_class, int(0.1 * len(y) / len(by_class)))

    val_data, train_data = [], []
    for cls, items in by_class.items():
        val_data.extend(items[:val_per_class])
        train_data.extend(items[val_per_class:])

    val_data   = shuffle(val_data,   random_state=SEED)
    train_data = shuffle(train_data, random_state=SEED)

    X_tr = np.array([x for x, _ in train_data])
    y_tr = np.array([lbl for _, lbl in train_data])
    X_v  = np.array([x for x, _ in val_data])
    y_v  = np.array([lbl for _, lbl in val_data])
    return X_tr, y_tr, X_v, y_v


# =============================================================================
# SECTION 4 — MODEL DEFINITIONS
# =============================================================================

def focal_loss(gamma: float = 2.0, alpha: float = 0.25):
    """
    Focal loss for binary classification.  Reduces the relative loss for
    well-classified examples so that the model focuses on hard ones.

    Parameters
    ----------
    gamma : focusing parameter (higher = more focus on hard examples)
    alpha : class-weight balancing factor
    """
    def focal_loss_fixed(y_true, y_pred):
        eps = K.epsilon()
        y_pred = K.clip(y_pred, eps, 1.0 - eps)
        pt_pos = tf.where(K.equal(y_true, 1), y_pred, K.ones_like(y_pred))
        pt_neg = tf.where(K.equal(y_true, 0), y_pred, K.zeros_like(y_pred))
        return (
            -K.mean(alpha       * K.pow(1.0 - pt_pos, gamma) * K.log(pt_pos))
            -K.mean((1 - alpha) * K.pow(pt_neg,        gamma) * K.log(1.0 - pt_neg))
        )
    return focal_loss_fixed


def build_base_cnn(hp, input_shape: tuple = (56, 1)):
    """
    Define the shared convolutional base used by both classifiers.

    Hyperparameters (all tunable via Keras Tuner):
      filters_1 / filters_2   : number of Conv1D filters
      kernel_size_1 / _2      : kernel sizes
      dropout_1 / dropout_2   : dropout rates after each pooling block
    """
    inputs = keras.Input(shape=input_shape)

    # Block 1
    x = layers.Conv1D(
        filters=hp.Int("filters_1", 32, 128, step=32),
        kernel_size=hp.Choice("kernel_size_1", [2, 3, 5]),
        activation="relu",
        padding="same",
    )(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPooling1D(pool_size=2)(x)
    x = layers.Dropout(hp.Float("dropout_1", 0.2, 0.5, step=0.1))(x)

    # Block 2
    x = layers.Conv1D(
        filters=hp.Int("filters_2", 32, 128, step=32),
        kernel_size=hp.Choice("kernel_size_2", [2, 3]),
        activation="relu",
        padding="same",
    )(x)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPooling1D(pool_size=2)(x)
    x = layers.Dropout(hp.Float("dropout_2", 0.2, 0.5, step=0.1))(x)

    x = layers.Flatten()(x)
    return keras.Model(inputs, x, name="cnn_base")


def _build_classification_head(base_model, hp, num_classes: int = 1):
    """
    Attach a fully-connected classification head to a CNN base model.

    Hyperparameters:
      dense_units    : number of units in the Dense layer
      dropout_dense  : dropout rate in the Dense layer
      lr             : AdamW learning rate
      use_focal_loss : whether to use focal loss instead of binary cross-entropy
      focal_gamma    : focal loss gamma parameter
      focal_alpha    : focal loss alpha parameter
    """
    x = base_model.output
    x = layers.Dense(
        hp.Int("dense_units", 32, 128, step=32), activation="relu"
    )(x)
    x = layers.Dropout(hp.Float("dropout_dense", 0.2, 0.5, step=0.1))(x)

    if num_classes == 1:
        outputs = layers.Dense(1, activation="sigmoid")(x)
        if hp.Boolean("use_focal_loss"):
            loss_fn = focal_loss(
                gamma=hp.Float("focal_gamma", 1.0, 4.0, step=0.5),
                alpha=hp.Float("focal_alpha", 0.25, 0.75, step=0.25),
            )
        else:
            loss_fn = "binary_crossentropy"
    else:
        outputs = layers.Dense(num_classes, activation="softmax")(x)
        loss_fn = "categorical_crossentropy"

    model = keras.Model(base_model.input, outputs)
    model.compile(
        optimizer=AdamW(learning_rate=hp.Choice("lr", [1e-3, 5e-4, 1e-4])),
        loss=loss_fn,
        metrics=["accuracy"],
    )
    return model


# Tuner entry points — one per classifier
def build_model_gain_loss(hp):
    return _build_classification_head(build_base_cnn(hp), hp, num_classes=1)

def build_model_change(hp):
    return _build_classification_head(build_base_cnn(hp), hp, num_classes=1)


# =============================================================================
# SECTION 5 — TRANSFER LEARNING
# =============================================================================

def build_feature_extractor(trained_model, layers_to_keep: int = 2):
    """
    Strip the final classification head from a trained model and return the
    remaining feature-extraction trunk.

    The last `layers_to_keep` layers are left trainable; all earlier layers are
    frozen so that the backbone is not destroyed by the new data.
    """
    extractor = Sequential(
        trained_model.layers[:-layers_to_keep],
        name="feature_extractor",
    )
    # Freeze all but the last few layers for fine-tuning
    for layer in extractor.layers[:-layers_to_keep]:
        layer.trainable = False
    for layer in extractor.layers[-layers_to_keep:]:
        layer.trainable = True
    return extractor


def build_transfer_model(
    feature_extractor,
    hidden_units: int = 32,
    dropout_rate: float = 0.3,
    learning_rate: float = 1e-4,
    loss_fn: str = "binary_crossentropy",
):
    """
    Add a new classification head on top of the frozen feature extractor.

    The head consists of a single Dense layer followed by Dropout and a
    sigmoid output.  Compiled with the Adam optimiser.
    """
    model = Sequential(
        [
            feature_extractor,
            layers.Dense(hidden_units, activation="relu"),
            layers.Dropout(dropout_rate),
            layers.Dense(1, activation="sigmoid"),
        ],
        name="transfer_model",
    )
    model.compile(
        optimizer=Adam(learning_rate=learning_rate),
        loss=loss_fn,
        metrics=["accuracy"],
    )
    return model


def run_transfer_learning(
    base_model_path: str,
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: np.ndarray,
    y_val: np.ndarray,
    class_weights_dict: dict,
    transfer_epochs: int = 100,
    batch_size: int = 32,
    patience: int = 10,
    layers_to_keep: int = 2,
) -> keras.Model:
    """
    Full transfer-learning pipeline for one classifier.

    1. Load the pre-trained base model.
    2. Build a frozen feature extractor from it.
    3. Attach a new classification head.
    4. Train with early stopping.

    Parameters
    ----------
    base_model_path    : path to the pre-trained .keras model
    X_train / y_train  : training data and labels
    X_val   / y_val    : validation data and labels
    class_weights_dict : per-class weights to counteract imbalance
    transfer_epochs    : maximum training epochs (early stopping may stop sooner)
    batch_size         : mini-batch size
    patience           : early-stopping patience
    layers_to_keep     : how many base-model layers (from the end) to leave trainable

    Returns
    -------
    Trained Keras Model
    """
    print(f"\n[Transfer] Loading base model from: {base_model_path}")
    base_model = keras.models.load_model(base_model_path)

    extractor      = build_feature_extractor(base_model, layers_to_keep)
    transfer_model = build_transfer_model(extractor)

    early_stop = EarlyStopping(
        monitor="val_loss",
        patience=patience,
        restore_best_weights=True,
        mode="min",
    )

    transfer_model.fit(
        X_train, y_train,
        validation_data=(X_val, y_val),
        epochs=transfer_epochs,
        batch_size=batch_size,
        class_weight=class_weights_dict,
        callbacks=[early_stop],
        verbose=1,
    )
    return transfer_model


# =============================================================================
# SECTION 6 — HYPERPARAMETER TUNING
# =============================================================================

def get_or_run_tuner(
    build_fn,
    project_name: str,
    tuning_dir: str,
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: np.ndarray,
    y_val: np.ndarray,
    max_trials: int = 50,
    run_search: bool = True,
) -> kt.BayesianOptimization:
    """
    Return a Keras Tuner with the best hyperparameters for `build_fn`.

    If `run_search` is True (default), a fresh Bayesian search is executed.
    If False, the function attempts to reload a previously completed search
    from `tuning_dir/project_name`.  If no cached results are found, it
    falls back to the default hyperparameter values.

    Parameters
    ----------
    build_fn     : model-builder function accepted by Keras Tuner
    project_name : unique name for this tuning run (also the cache key)
    tuning_dir   : parent directory for tuning artefacts
    run_search   : whether to actually run the search (set False to reuse cache)
    max_trials   : number of Bayesian Optimisation trials

    Returns
    -------
    A Keras Tuner instance whose `get_best_hyperparameters()` can be called.
    """
    tuner = kt.BayesianOptimization(
        build_fn,
        objective="val_accuracy",
        max_trials=max_trials,
        executions_per_trial=1,
        directory=tuning_dir,
        project_name=project_name,
        seed=SEED,
        overwrite=run_search,   # fresh search overwrites old results
    )

    if run_search:
        print(f"\n[Tuner] Running HP search for '{project_name}' "
              f"({max_trials} trials) …")
        tuner.search(
            X_train, y_train,
            validation_data=(X_val, y_val),
            epochs=20,
            callbacks=[EarlyStopping(monitor="val_loss", patience=5)],
            verbose=0,
        )
    else:
        print(f"[Tuner] Reusing cached HP results for '{project_name}'.")

    try:
        best_hps = tuner.get_best_hyperparameters(num_trials=1)[0]
        print(f"[Tuner] Best hyperparameters: {best_hps.values}")
    except Exception:
        print("[Tuner] No cached results found; using default hyperparameters.")

    return tuner


# =============================================================================
# SECTION 7 — PREDICTION & EVALUATION
# =============================================================================

def find_best_threshold(
    y_true: np.ndarray, y_pred_proba: np.ndarray
) -> tuple[float, float]:
    """
    Sweep the decision threshold and return the one that maximises F1.

    Returns
    -------
    (best_threshold, best_f1)
    """
    precision, recall, thresholds = precision_recall_curve(y_true, y_pred_proba)
    f1 = 2 * (precision * recall) / (precision + recall + 1e-9)
    best_idx = np.argmax(f1)
    return float(thresholds[best_idx]), float(f1[best_idx])


def triple_class_predict(
    model_change: keras.Model,
    model_gain_loss: keras.Model,
    X: np.ndarray,
    thresh_change: float = 0.5,
    thresh_gain_loss: float = 0.5,
) -> np.ndarray:
    """
    Combine the two binary classifiers into a single 3-class prediction.

    Decision logic:
      • model_change predicts whether binding changes at all.
        – if change_prob > thresh_change  →  No Change (2)
        – otherwise pass to model_gain_loss
      • model_gain_loss predicts the direction of change.
        – if gain_loss_prob > thresh_gain_loss  →  Loss (1)
        – otherwise                             →  Gain (0)

    Parameters
    ----------
    model_change     : trained Change-vs-NoChange model
    model_gain_loss  : trained Gain-vs-Loss model
    X                : input features, shape (N, 56, 1)
    thresh_change    : decision threshold for the change classifier
    thresh_gain_loss : decision threshold for the gain/loss classifier

    Returns
    -------
    Integer array of shape (N,) with values in {0, 1, 2}.
    """
    change_proba    = model_change.predict(X, verbose=0).flatten()
    gain_loss_proba = model_gain_loss.predict(X, verbose=0).flatten()

    change_class    = (change_proba    > thresh_change).astype(int)
    gain_loss_class = (gain_loss_proba > thresh_gain_loss).astype(int)

    # 0 = Gain, 1 = Loss, 2 = No Change
    return np.where(change_class == 1, 2, gain_loss_class).astype(int)


def evaluate(
    model_change: keras.Model,
    model_gain_loss: keras.Model,
    X_val_change: np.ndarray,
    y_val_change: np.ndarray,
    X_val_gain_loss: np.ndarray,
    y_val_gain_loss: np.ndarray,
    X_test: np.ndarray,
    y_test: np.ndarray,
    label: str = "",
):
    """
    Find optimal per-classifier thresholds on the validation set, then report
    both default (0.5) and optimised results on the test set.

    Prints confusion matrices and classification reports to stdout.
    """
    # ── Optimised thresholds ──────────────────────────────────────────────────
    thresh_change, f1_change = find_best_threshold(
        y_val_change, model_change.predict(X_val_change, verbose=0).flatten()
    )
    thresh_gl, f1_gl = find_best_threshold(
        y_val_gain_loss, model_gain_loss.predict(X_val_gain_loss, verbose=0).flatten()
    )
    print(f"\n[Eval] {label}")
    print(f"  Best threshold — Change vs None : {thresh_change:.3f}  (val F1 {f1_change:.3f})")
    print(f"  Best threshold — Gain  vs Loss  : {thresh_gl:.3f}  (val F1 {f1_gl:.3f})")

    for thresh_c, thresh_g, tag in [
        (0.5,          0.5,     "default thresholds (0.5)"),
        (thresh_change, thresh_gl, "optimised thresholds"),
    ]:
        preds = triple_class_predict(
            model_change, model_gain_loss, X_test,
            thresh_change=thresh_c, thresh_gain_loss=thresh_g,
        )
        print(f"\n  === {tag} ===")
        print("  Confusion matrix:")
        print(confusion_matrix(y_test, preds))
        print(classification_report(
            y_test, preds, target_names=["Gain", "Loss", "No Change"]
        ))


# =============================================================================
# SECTION 8 — MAIN ENTRY POINT
# =============================================================================

def main():
    args = parse_args()

    # ── Output directory ──────────────────────────────────────────────────────
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)

    transfer_change_path    = f"{args.output_dir}/Transfer_Change_None_{args.tf}.keras"
    transfer_gain_loss_path = f"{args.output_dir}/Transfer_Gain_Loss_{args.tf}.keras"

    # ── Load and prepare data ─────────────────────────────────────────────────
    X_train, X_val, X_test, y_train, y_val, y_test = load_and_prepare_data(args)

    # ── Derive task-specific subsets ──────────────────────────────────────────
    # Gain vs Loss: only samples with an actual change (label 0 or 1)
    change_mask = y_train != 2
    X_train_gl  = X_train[change_mask]
    y_train_gl  = y_train[change_mask]

    change_mask_val = y_val != 2
    X_val_gl  = X_val[change_mask_val]
    y_val_gl  = y_val[change_mask_val]

    # Change vs No Change: remap labels to binary (0/1 → 0 = change, 2 → 1 = no change)
    CHANGE_MAP = {0: 0, 1: 0, 2: 1}
    y_train_change = np.array([CHANGE_MAP[int(lbl)] for lbl in y_train])
    y_val_change   = np.array([CHANGE_MAP[int(lbl)] for lbl in y_val])

    # ── PREDICT-ONLY MODE ─────────────────────────────────────────────────────
    if args.predict_only:
        print(f"\n[Mode] Predict-only — loading models from {args.output_dir}")
        if not (Path(transfer_change_path).exists() and
                Path(transfer_gain_loss_path).exists()):
            sys.exit(
                f"[Error] Could not find saved transfer models in {args.output_dir}.\n"
                "Run without --predict_only first to train them."
            )
        model_change    = keras.models.load_model(transfer_change_path)
        model_gain_loss = keras.models.load_model(transfer_gain_loss_path)

        evaluate(
            model_change, model_gain_loss,
            X_val[change_mask_val == False], y_val_change[change_mask_val == False],   # noqa: E712
            X_val_gl, y_val_gl,
            X_test, y_test,
            label=f"Transfer models — {args.tf}",
        )
        return

    # ── TRAINING MODE ─────────────────────────────────────────────────────────
    print(f"\n[Mode] Training transfer models for: {args.tf}")

    # Gain vs Loss tuner
    tuner_gl = get_or_run_tuner(
        build_fn=build_model_gain_loss,
        project_name=f"CNN_GainvsLoss_{args.tf}",
        tuning_dir=args.tuning_dir,
        X_train=X_train_gl, y_train=y_train_gl,
        X_val=X_val_gl,     y_val=y_val_gl,
        max_trials=args.max_trials,
        run_search=not args.no_tune,
    )

    # Change vs None tuner
    tuner_change = get_or_run_tuner(
        build_fn=build_model_change,
        project_name=f"CNN_ChangevsNone_{args.tf}",
        tuning_dir=args.tuning_dir,
        X_train=X_train, y_train=y_train_change,
        X_val=X_val,     y_val=y_val_change,
        max_trials=args.max_trials,
        run_search=not args.no_tune,
    )

    # Class weights for imbalanced tasks
    cw_gl = dict(enumerate(
        compute_class_weight("balanced", classes=np.unique(y_train_gl), y=y_train_gl)
    ))
    cw_change = dict(enumerate(
        compute_class_weight("balanced", classes=np.unique(y_train_change), y=y_train_change)
    ))

    # Transfer learning — Gain vs Loss
    print("\n[Train] Fine-tuning Gain vs Loss classifier …")
    model_gain_loss = run_transfer_learning(
        base_model_path=args.base_gain_loss,
        X_train=X_train_gl, y_train=y_train_gl,
        X_val=X_val_gl,     y_val=y_val_gl,
        class_weights_dict=cw_gl,
        transfer_epochs=args.transfer_epochs,
        batch_size=args.batch_size,
        patience=args.patience,
        layers_to_keep=args.layers_to_keep,
    )

    # Transfer learning — Change vs None
    print("\n[Train] Fine-tuning Change vs No-Change classifier …")
    model_change = run_transfer_learning(
        base_model_path=args.base_change,
        X_train=X_train, y_train=y_train_change,
        X_val=X_val,     y_val=y_val_change,
        class_weights_dict=cw_change,
        transfer_epochs=args.transfer_epochs,
        batch_size=args.batch_size,
        patience=args.patience,
        layers_to_keep=args.layers_to_keep,
    )

    # ── Save models ───────────────────────────────────────────────────────────
    model_change.save(transfer_change_path)
    model_gain_loss.save(transfer_gain_loss_path)
    print(f"\n[Save] Models saved to {args.output_dir}/")

    # ── Evaluate ──────────────────────────────────────────────────────────────
    evaluate(
        model_change, model_gain_loss,
        X_val, y_val_change,
        X_val_gl, y_val_gl,
        X_test, y_test,
        label=f"Transfer models — {args.tf}",
    )

    # Optionally compare against the original base models
    print("\n[Eval] Base model (no transfer) comparison:")
    base_change    = keras.models.load_model(args.base_change)
    base_gain_loss = keras.models.load_model(args.base_gain_loss)
    evaluate(
        base_change, base_gain_loss,
        X_val, y_val_change,
        X_val_gl, y_val_gl,
        X_test, y_test,
        label="Base models (pre-transfer)",
    )


if __name__ == "__main__":
    main()