task_commons.py

import glob
import os
import re

import joblib
import numpy as np
import pandas as pd
import torch
from sklearn.decomposition import IncrementalPCA
from sklearn.linear_model import LogisticRegression, LinearRegression, Ridge
from sklearn.model_selection import StratifiedKFold, GridSearchCV
from torch.utils.data import DataLoader

import globals
from DataSetManager import customDataLoader
from datasets.adni_datset import collapse_adni_labels, filter_adni_dataset_based_on_task, filter_labels

from trainer import get_shape_and_adapter_based_on_model

FEATURE_CACHE_DIR = os.environ.get("FEATURE_CACHE_DIR", "/scratch/tmp/features_cache")
os.makedirs(FEATURE_CACHE_DIR, exist_ok=True)

PROJECT_DIR = os.environ.get("PROJECT_DIR", "/scratch/progetto")
CHECKPOINT_DIR = os.path.join(PROJECT_DIR, "checkpoints")
def get_dataset_type(task_label):
    if task_label == "age" or task_label == "nage":
        from datasets.AgeRegressionDataset import AgeRegressionDataset
        return AgeRegressionDataset
    elif task_label == "sex":
        from datasets.sex_classification_dataset import sex_classification_Dataset
        return sex_classification_Dataset
    elif task_label in ["gm","csf","cAvgThick"]:
        from datasets.brain_roi_datset import brainRoiDataset
        return  brainRoiDataset
    elif task_label in  ["adni","adniEasier","ADvsCN","adni-hard"]:
        from datasets.adni_datset import adni_classification_Dataset
        return adni_classification_Dataset
    else:
        raise NotImplementedError(f"Dataset for task {task_label} not implemented")

def select_n_components_batchwise(dataloader, training_set, target_variance=0.95, batch_size=5, max_components=200):

    print("🚀 Starting Incremental PCA (batch-wise)")
    ipca = IncrementalPCA(n_components=max_components, batch_size=batch_size)

    # Fit PCA batch-wise
    for i in range(0, len(training_set), batch_size):
        batch = training_set[i:i + batch_size]
        X_batch = np.array([
            dataloader.load_npz_as_tensor(d["image_path"], normalize=True, downsample=globals.DOWNSAMPLE)
            .squeeze().flatten().numpy().astype(np.float32)
            for d in batch
        ])
        ipca.partial_fit(X_batch)

        if (i // batch_size + 1) % 10 == 0:
            print(f"  → Processed batch {i // batch_size + 1}/{len(training_set) // batch_size}")

    # Determine number of components needed
    cum_var = np.cumsum(ipca.explained_variance_ratio_)
    n_components_needed = np.searchsorted(cum_var, target_variance) + 1

    print(f"✅ Batch-wise PCA: {n_components_needed} components capture ≥ {target_variance * 100:.1f}% variance")
    return ipca, n_components_needed, cum_var

def parse_model_filename(filename):
    match = re.search(r"([\w\-]+)_model_train([\d.]+)_batch(\d+)_epoch(\d+)\.pt", filename)
    if match:
        model_type = match.group(1)
        train_val = float(match.group(2))
        batch = int(match.group(3))
        epoch = int(match.group(4))
        return model_type, train_val, batch, epoch
    return None

def group_checkpoints(target_dir):
    ckpts = glob.glob(os.path.join(CHECKPOINT_DIR, f"{target_dir}","*_model_train*_batch*_epoch*.pt"))
    grouped = {}
    for f in ckpts:
        parsed = parse_model_filename(os.path.basename(f))
        if parsed:
            mt, tr, bs, ep = parsed
            grouped.setdefault((mt, tr, bs), []).append((ep, f))
    for k in grouped:
        grouped[k] = sorted(grouped[k], key=lambda x: x[0])
    return grouped


def PCA_maker(batch_size, max_components, target_variance, task_label="sex"):
    dataloader = customDataLoader(train_ratio=0, test_ratio=0, batch_size=batch_size)
    print("🔍 Preparing datasets...")
    test_set = dataloader.get_test_data(forModelTraining=False)
    training_set = dataloader.get_test_data(forModelTraining=True)

    # PCA cache paths
    ipca_train_path = os.path.join(FEATURE_CACHE_DIR, "train_ipca_features.npz")
    ipca_test_path = os.path.join(FEATURE_CACHE_DIR, "test_ipca_features.npz")
    ipca_model_path = os.path.join(FEATURE_CACHE_DIR, "ipca_model.pkl")
    force_recompute = False


    if not os.path.exists(ipca_model_path) or force_recompute:
        print("🔁 Computing PCA model and features from scratch...")
        ipca, n_components_needed, cum_var = select_n_components_batchwise(
            dataloader, training_set, target_variance=target_variance,
            batch_size=batch_size, max_components=max_components
        )


        X_train = np.array([
            dataloader.load_npz_as_tensor(d["image_path"], normalize=True, downsample=globals.DOWNSAMPLE)
            .flatten().numpy().astype(np.float32) for d in training_set
        ])
        X_test = np.array([
            dataloader.load_npz_as_tensor(d["image_path"], normalize=True, downsample=globals.DOWNSAMPLE)
            .flatten().numpy().astype(np.float32) for d in test_set
        ])
        X_train_pca = ipca.transform(X_train)[:, :n_components_needed]
        X_test_pca = ipca.transform(X_test)[:, :n_components_needed]


        joblib.dump(ipca, ipca_model_path)
        np.savez(ipca_train_path, X=X_train_pca)
        np.savez(ipca_test_path, X=X_test_pca)
        print(f"✅ PCA model and features saved (components={n_components_needed})")

    else:

        print("📦 PCA model found — loading...")
        ipca = joblib.load(ipca_model_path)
        n_components_used = ipca.n_components_
        cum_var = np.cumsum(ipca.explained_variance_ratio_)
        var_captured = cum_var[-1]


        if os.path.exists(ipca_train_path) and os.path.exists(ipca_test_path):
            print("📦 PCA features found — loading...")
            train_data = np.load(ipca_train_path)
            test_data = np.load(ipca_test_path)
            X_train_pca = train_data["X"]
            X_test_pca = test_data["X"]
        else:

            print("📌 PCA features missing — computing features from model...")
            X_train = np.array([
                dataloader.load_npz_as_tensor(d["image_path"], normalize=True, downsample=globals.DOWNSAMPLE)
                .flatten().numpy().astype(np.float32) for d in training_set
            ])
            X_test = np.array([
                dataloader.load_npz_as_tensor(d["image_path"], normalize=True, downsample=globals.DOWNSAMPLE)
                .flatten().numpy().astype(np.float32) for d in test_set
            ])
            X_train_pca = ipca.transform(X_train)[:, :n_components_used]
            X_test_pca = ipca.transform(X_test)[:, :n_components_used]


            np.savez(ipca_train_path, X=X_train_pca)
            np.savez(ipca_test_path, X=X_test_pca)
            print("💾 PCA features computed and saved.")

        print(f"ℹ️ PCA model loaded: {n_components_used} components capture {var_captured * 100:.2f}% of variance")


    y_train_path = os.path.join(FEATURE_CACHE_DIR, f"y_train_{task_label}.npz")
    y_test_path = os.path.join(FEATURE_CACHE_DIR, f"y_test_{task_label}.npz")

    if os.path.exists(y_train_path) and os.path.exists(y_test_path):
        y_train = np.load(y_train_path)["y"]
        y_test = np.load(y_test_path)["y"]
        print(f"📦 Loaded cached targets for task '{task_label}'.")
    else:
        y_train = np.array([d[task_label] for d in training_set])
        y_test = np.array([d[task_label] for d in test_set])
        np.savez(y_train_path, y=y_train)
        np.savez(y_test_path, y=y_test)
        print(f"💾 Created and saved targets for task '{task_label}'.")


    if task_label == "sex":
        clf = LogisticRegression(
            max_iter=1000,
            solver="lbfgs",
            penalty="l2",
            C=1.0,
            n_jobs=8,
        )
    elif task_label == "age":
        clf = LinearRegression(

            fit_intercept=True,
            copy_X=True,

        )
    else:
        raise ValueError(f"<UNK> Task label {task_label} not recognized.")

    clf.fit(X_train_pca, y_train)

    return X_test_pca, X_train_pca, clf, y_test, y_train



def to_representation_space(device, encoder, train_loader):
    all_x, all_y = [], []
    encoder.eval()
    import globals

    must_downsample = globals.DOWNSAMPLE
    model_type = globals.MODEL_TYPE


    target_shape, vitAdapter = get_shape_and_adapter_based_on_model(device, model_type, must_downsample)
    with torch.no_grad():
        for x, y in train_loader:
            x = x.to(device).squeeze(1)
            if vitAdapter is not None:
                print(f"pre adapter shape {x.shape}")
                x = vitAdapter.crop_batch(x)
                print(f"Successfully adapted to {target_shape}, shape batch {x.shape}.")
            # Encode immagini

            x = encoder(x)
            x_np = x.detach().cpu().numpy()


            # Convertiamo in array numpy coerente
            y_np = np.array(y)

            # Aggiungiamo ogni elemento del batch
            for xi, yi in zip(x_np, y_np):
                all_x.append(xi)
                all_y.append(yi)

        # Concatenamento finale
        X_train = np.stack(all_x, axis=0)
        y_train = np.stack(all_y, axis=0)

        print(f"REP : shape of X_train is {X_train.shape}")
        print(f"REP : shape of y_train is {y_train.shape}")

    return X_train, y_train




def estimate_best_alpha_cv_subset(encoder, dataloader, device, batch_size, seed, size, datasetType, n_splits=5,task=None, doing_global=False, downsample=True):
    if  task not in ["adni","adni-hard", "ADvsCN", "adniEasier"]:
        # Otteniamo l'intero dataset
        test_set = dataloader.get_test_data(forModelTraining=False)
        training_set = dataloader.get_test_data(forModelTraining=True)
        dataset = training_set + test_set
        if "nage" == task:
            dataset = expand_dataset_with_adni(dataloader, dataset)
        df = pd.DataFrame(dataset)

        # Stratificazione basata sull'età
        df["age_group"] = pd.cut(df["age"], bins=[0, 10, 20, 30, 40, 50, 60, 70, 80, 100], labels=False)
        df["strata"] = df["age_group"].astype(str)
        n_classes = df["strata"].nunique()
    else:
        test_set = dataloader.get_adni_test_data(False)
        training_set = dataloader.get_adni_test_data(True)

        dataset = training_set + test_set
        dataset, labels_to_filter = filter_adni_dataset_based_on_task(dataset, task)
        print(f"filtered labels: {labels_to_filter} , length of filtered dataset: {len(dataset)}")
        dataset = collapse_adni_labels(dataset,task=task)

        df = pd.DataFrame(dataset)
        df["age"] = pd.to_numeric(df["age"], errors="coerce")
        df = df.dropna(subset=["age"])

        df["age_group"] = pd.cut(df["age"], bins=[0, 10, 20, 30, 40, 50, 60, 70, 80, 100], labels=False)
        df["diagnosis"] = pd.Categorical(df["group"])
        df["strata"] = (
                df["age_group"].astype(str) + "_" +
                df["diagnosis"].astype(str)
        )
        n_classes = df["strata"].nunique()

    # Campionamento bilanciato
    sampled_df = sample_balanced(-1, n_classes, seed, size, df,adni_logic=task in ["adni", "ADvsCN", "adniEasier","adni-hard"])
    sampled_data = sampled_df.to_dict(orient="records")

    # Loader per il task specifico
    if task in ["adni", "ADvsCN", "adniEasier","adni-hard"]:
        sampled_loader = DataLoader(
            datasetType(sampled_data, dataloader.load_nii_as_tensor, downsample=downsample),
            batch_size=batch_size,
            num_workers=8,
            pin_memory=True
        )
    elif task is None or task == "age":
        sampled_loader = DataLoader(
            datasetType(sampled_data, dataloader.load_npz_as_tensor,secondary_load_fn = None,downsample=downsample),
            batch_size=batch_size,
            num_workers=8,
            pin_memory=True
        )
    elif task== "sex":
        sampled_loader = DataLoader(
            datasetType(sampled_data, dataloader.load_npz_as_tensor, downsample=downsample),
            batch_size=batch_size,
            num_workers=8,
            pin_memory=True
        )
    elif task == "nage":
        sampled_loader = DataLoader(
            datasetType(sampled_data, dataloader.load_npz_as_tensor, downsample=downsample , secondary_load_fn = dataloader.load_nii_as_tensor),
            batch_size=batch_size,
            num_workers=8,
            pin_memory=True
        )
    else:

        sampled_loader = DataLoader(
            datasetType(sampled_data, dataloader.load_npz_as_tensor,target=task,downsample=downsample),
            batch_size=batch_size,
            num_workers=8,
            pin_memory=True,
        )


    # Otteniamo X e y del task
    X_train, y_train = to_representation_space(device, encoder, sampled_loader)
    if doing_global:
        if isinstance( y_train, np.ndarray) and  y_train.ndim > 1:
            y_train =  y_train.sum(axis=1)
            print( y_train)
    if task not in ["adni", "ADvsCN", "adniEasier"]:
        # Stratificazione basata sull'età, indipendente dal target del task
        age_bins = sampled_df["age_group"].values

        # Grid search su alpha
        alphas = 10.0 ** np.arange(-2, 4)
        print(f"Testing alphas: {alphas}")

        ridge = Ridge(fit_intercept=True, random_state=seed, positive=True)
        skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)

        grid = GridSearchCV(
            ridge,
            param_grid={"alpha": alphas},
            cv=skf.split(X_train, age_bins),
            scoring="neg_mean_squared_error",
            n_jobs=8
        )
        grid.fit(X_train, y_train)

        best_alpha = grid.best_params_["alpha"]
        print(f"✅ Best alpha for size={size}: {best_alpha:.4f}")
        return best_alpha
    else:
        param_grid = {"C": 10.0 ** np.arange(-2, 3)}
        skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
        lr = LogisticRegression(
            penalty='l2',
            tol=1e-5,
            fit_intercept=True,
            random_state=seed,
            solver='lbfgs',
            max_iter=10000,
            multi_class='auto',
            class_weight='balanced',
            n_jobs=8
        )
        grid = GridSearchCV(lr, param_grid, cv=skf, scoring='balanced_accuracy', n_jobs=8)
        grid.fit(X_train, y_train)
        best_C = grid.best_params_['C']
        return best_C


def expand_dataset_with_adni(dataloader, dataset):
    datasetadni = get_adni_adapted(dataloader)
    dataset = dataset + datasetadni

    print("total dataset", len(dataset))
    return dataset


def get_class_age_indices(train_df):
    """
    Restituisce un dizionario gerarchico:
    class -> age_group -> lista di indici
    """
    class_to_indices = {}

    # tutte le classi
    classes = train_df["diagnosis"].unique()

    for cls in classes:
        cls_df = train_df[train_df["diagnosis"] == cls]
        age_groups = cls_df["age_group"].unique()

        # per ogni classe, creo sotto-dizionario per fascia età
        class_to_indices[cls] = {}
        for age in age_groups:
            age_idx = cls_df[cls_df["age_group"] == age].index.tolist()
            class_to_indices[cls][age] = age_idx


    return class_to_indices

def sample_balanced(fold_idx, n_classes, seed, size, train_df,adni_logic=False):
    rng = np.random.RandomState(seed + fold_idx)
    size = min(size, len(train_df))
    if adni_logic:
        selected_indices = []
        diagnosi_age_index = get_class_age_indices(train_df)
        classes = list(diagnosi_age_index.keys())
        n_classes = len(classes)

        # --- calcolo numero massimo per classe
        min_class_count = min(sum(len(indices) for indices in age_dict.values()) for age_dict in diagnosi_age_index.values())
        n_per_class = min(min_class_count, size // n_classes)
        print(f"Numero di elementi da prendere per classe: {n_per_class}")

        # --- estrazione
        for cls in classes[:]:  # [:] per iterare su copia
            age_dict = diagnosi_age_index[cls]
            taken = 0

            while taken < n_per_class and len(age_dict) > 0:
                # scegli una fascia d'età casuale tra quelle con elementi
                age_chosen = rng.choice(list(age_dict.keys()))

                # scegli un elemento casuale in quella fascia
                idx_list = age_dict[age_chosen]
                chosen_idx = rng.choice(idx_list)
                chosen_idx = int(chosen_idx)

                # aggiungi alla lista selezionata
                selected_indices.append(chosen_idx)
                taken += 1

                # rimuovi indice e pulisci fascia vuota
                idx_list.remove(chosen_idx)
                if len(idx_list) == 0:
                    del age_dict[age_chosen]

            # se la classe è vuota, eliminiamola
            if len(age_dict) == 0:
                del diagnosi_age_index[cls]

        print("Totale campioni selezionati con metodo adni logic:", len(selected_indices))
        remaining_to_take = size - len(selected_indices)
        print(f"rimasti {remaining_to_take} da sampling random")
        if remaining_to_take > 0:
            # raccogli tutti gli indici rimasti
            remaining_indices = []
            for cls, age_dict in diagnosi_age_index.items():
                for idx_list in age_dict.values():
                    remaining_indices.extend(idx_list)

            # se ci sono elementi rimasti
            if len(remaining_indices) > 0:
                # se ne rimane meno di remaining_to_take, prendiamo tutti
                n_fill = min(remaining_to_take, len(remaining_indices))
                fill_indices = rng.choice(remaining_indices, size=n_fill, replace=False).tolist()
                selected_indices.extend(fill_indices)
        selected_indices_flat = []

        for idx in selected_indices:
            if isinstance(idx, (list, np.ndarray)):
                selected_indices_flat.extend(list(idx))
            else:
                selected_indices_flat.append(idx)
        print(selected_indices_flat)
        return train_df.loc[selected_indices_flat]


    classes = train_df["strata"].unique()
    print(f"sampling {size} from {len(classes)} classes ")
    if len(classes) != n_classes:
        raise ValueError(
            f"n_classes={n_classes} ma trovate {len(classes)} classi in train_df"
        )

    selected_indices = []

    if size >= n_classes:
        # 1 elemento per ogni classe

        for cls in classes:
            cls_idx = train_df[train_df["strata"] == cls].index
            chosen = rng.choice(cls_idx, size=1, replace=False)
            selected_indices.extend(chosen)
        print("forzate :", len(selected_indices))


        # a caso
        remaining_size = size - n_classes
        print(f"remaining size {remaining_size}")
        if remaining_size > 0:
            remaining_idx = train_df.index.difference(selected_indices)
            extra_idx = rng.choice(
                remaining_idx,
                size=remaining_size,
                replace=False
            )
            selected_indices.extend(extra_idx)
        print(f"Tot sampled size {len(selected_indices)}")

    #caso size troppo piccolo per tutte le classi esistenti
    else:
        # sample classi casuali
        chosen_classes = rng.choice(classes, size=size, replace=False)
        print(f"WARN: TOO MANY CLASSES ")
        print(f"chosen classes {chosen_classes}")
        for cls in chosen_classes:
            cls_idx = train_df[train_df["strata"] == cls].index
            chosen = rng.choice(cls_idx, size=1, replace=False)
            selected_indices.extend(chosen)

    return train_df.loc[selected_indices]

#
# def sample_balanced(
#     fold_idx,
#     n_classes,
#     seed,
#     size,
#     train_df,
#     force_minority_representation=False
# ):
#     size = min(size, len(train_df))
#
#     if size >= n_classes:
#         stratified_indices = []
#
#         if force_minority_representation:
#             last_two_classes = train_df["strata"].unique()[-2:]
#             for cls in last_two_classes:
#                 cls_idx = train_df[train_df["strata"] == cls].index
#                 if len(cls_idx) > 0:
#                     chosen = np.random.RandomState(seed + fold_idx).choice(
#                         cls_idx, size=1, replace=False
#                     )
#                     stratified_indices.extend(chosen)
#
#         remaining_size = size - len(stratified_indices)
#         if remaining_size > 0:
#             remaining_idx = train_df.index.difference(stratified_indices)
#             additional_idx = np.random.RandomState(seed + fold_idx).choice(
#                 remaining_idx, size=remaining_size, replace=False
#             )
#             stratified_indices.extend(additional_idx)
#
#         train_df = train_df.loc[stratified_indices]
#
#     else:
#         if size != 10:
#             raise Exception("Random sampling implemented only for size=10")
#         sampled_idx = np.random.RandomState(seed + fold_idx).choice(
#             len(train_df), size=size, replace=False
#         )
#         train_df = train_df.iloc[sampled_idx]
#
#     return train_df
def get_adni_adapted(dataloader):
    ad_test = dataloader.get_adni_test_data(False)
    ad_train = dataloader.get_adni_test_data(True)
    datasetadni = ad_train + ad_test
    datasetadni = filter_labels(datasetadni, labels_to_keep=["CN"], labels_to_filter=[])
    for item in datasetadni:
        item["sex"]= 1 if item["sex"] == 'M' else 0
        item["age"] = float(item["age"])
    return datasetadni