trainer.py

import os.path
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from multiprocessing import get_context
import time
import torch.distributed.nn.functional as dist_nnF
import torch.optim.lr_scheduler
from torch import optim
import globals
from DataSetManager import *
# from scheduler.incremental_and_decay_scheduler import IncrementalAndDecayScheduler
from transforms.augmentations import get_random_augmentation, apply_augmentation
from transforms.resize4Vit import ModelImageAdapter
from losses.SupInfoNCE import EpsInfoNCE
from models import representation as rep

import sys
import logging
logging.basicConfig(
    level=logging.DEBUG,
    stream=sys.stderr,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
import traceback


def excepthook(exc_type, exc_value, exc_tb):
    traceback.print_exception(exc_type, exc_value, exc_tb, file=sys.stderr)

sys.excepthook = excepthook

ctx = get_context("spawn")
# === training loop ===
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

from torch.cuda.amp import GradScaler

PROJECT_DIR=os.environ.get("PROJECT_DIR", "/scratch/progetto")

import subprocess

def slurm_time_remaining():
    import os, subprocess
    job_id = os.environ.get("SLURM_JOB_ID")
    if not job_id:
        return None  # non siamo su SLURM

    result = subprocess.run(
        ["squeue", "-j", job_id, "-h", "-o", "%L"], capture_output=True, text=True
    )
    if result.returncode != 0:
        return None

    time_str = result.stdout.strip()
    total_seconds = 0

    # caso D-HH:MM:SS
    if "-" in time_str:
        days, hms = time_str.split("-")
        parts = list(map(int, hms.split(":")))
        while len(parts) < 3:  # aggiunge zeri se manca
            parts.insert(0, 0)
        hours, minutes, seconds = parts
        total_seconds = int(days) * 86400 + hours*3600 + minutes*60 + seconds
    else:
        parts = list(map(int, time_str.split(":")))
        if len(parts) == 3:
            hours, minutes, seconds = parts
        elif len(parts) == 2:
            hours = 0
            minutes, seconds = parts
        elif len(parts) == 1:
            hours = 0
            minutes = 0
            seconds = parts[0]
        else:
            return None
        total_seconds = hours*3600 + minutes*60 + seconds
        print(f"remaining time: {total_seconds}")
    return total_seconds

def slurm_kill_job():
    """Rank 0 kills the whole Slurm job."""
    job_id = os.environ.get("SLURM_JOB_ID")
    if job_id is None:
        print("[Rank 0] SLURM_JOB_ID not found, cannot kill job automatically.")
        return

    print(f"[Rank 0] Killing Slurm job {job_id} to stop all ranks.")
    try:
        subprocess.run(["scancel", job_id], check=True)
        print(f"[Rank 0] scancel command sent successfully.")
    except subprocess.CalledProcessError as e:
        print(f"[Rank 0] Failed to cancel job {job_id}: {e}")
def save_safe(obj, filename):
    """Salva l'oggetto su disco e aspetta che sia fisicamente scritto prima di continuare."""
    tmp_filename = filename + ".tmp"
    torch.save(obj, tmp_filename)
    # Rinomina in modo atomico
    os.replace(tmp_filename, filename)
    # Forza la scrittura su disco
    with open(filename, 'rb') as f:
        os.fsync(f.fileno())

def parse_args():
    parser = argparse.ArgumentParser(description="Training script parameters")

    parser.add_argument('--train_ratio', type=float, default=0.7, help='Train dataset ratio')
    parser.add_argument('--train_second_chance', type=float, default=0.4, help='Probability of subject to be taken again')
    parser.add_argument('--sampling_second_chance', type=float, default=1, help='Probability of subject to be taken again at sampling')
    parser.add_argument('--batch_size', type=int, default=16, help='Batch size')
    parser.add_argument('--old_total_batch_size', type=int, default=0, help='Batch size pre uptrain')
    parser.add_argument('--num_epochs', type=int, default=100, help='Number of epochs')
    parser.add_argument('--uptraining_start_epoch', type=int, default=101, help='epoch to start uptraining')
    parser.add_argument('--max_same_batch_positives', type=int, default=2, help='Max positives in same batch')
    parser.add_argument('--max_real_same_batch_positives', type=int, default=1, help='Max positives in same batch')
    parser.add_argument('--seed', type=int, default=12, help='Random seed')
    parser.add_argument('--force_augmentation', action='store_true', help='Force Augmentation of real images')
    parser.add_argument('--use_muon', action='store_true', help='Use Muon optimizer')
    parser.add_argument('--model_type', type=str, default="resnet-18", help='Model to be trained')
    parser.add_argument('--target_dir', type=str,required=True, help='the directory inside checkpoints and encoder-losses to save the trained model')
    parser.add_argument('--max_test_batches', type=int, default=10, help='max number of batches used for testing')
    parser.add_argument('--required_augmentations', type=int, default=2, help='number of augmentations added')
    parser.add_argument('--lr',type=float,default=1e-4,help='Learning rate (default: 1e-4)')
    parser.add_argument('--lr_scheduler', type=str, default="None", choices=["StepLR", "CosineAnnealing","None","warmup+CosineAnnealing"], help='Type of LR scheduler to use')
    parser.add_argument('--lr_scheduler_step', type=int, default=10,help='Step size for StepLR or T_max for CosineAnnealing (in epochs)')
    parser.add_argument("--uptraining", action="store_true", help="Enable up-training (only resnet)")
    parser.add_argument("--uptraining_batch", type=int, default=24, help="batch size to use with up-training (singular gpu pre augmentation)")
    # passato da torch run parser.add_argument("--local_rank", type=int, help="Local rank for DistributedDataParallel")

    return parser.parse_args()

def resize_tensor(t, target_shape):
    """
    Resize a tensor of shape [C, D, H, W] to target_shape (D, H, W)
    using trilinear interpolation. Adds missing channel/batch dims if needed.
    """
    # Ensure input is 4D: [C, D, H, W]
    if t.dim() == 3:
        t = t.unsqueeze(0)  # add channel dim
    elif t.dim() == 5:
        t = t.squeeze(0)  # remove batch if already present

    if t.dim() != 4:
        raise ValueError(f"Expected tensor with shape [C, D, H, W], got {t.shape}")

    t = t.unsqueeze(0)  # [1, C, D, H, W] for interpolate
    resized = F.interpolate(t, size=target_shape, mode='trilinear', align_corners=False)
    return resized.squeeze(0)  # back to [C, D, H, W]



def main(args):
    train_ratio = args.train_ratio
    batch_size = args.batch_size
    num_epochs = args.num_epochs
    seed = args.seed
    max_same_batch_positives = args.max_same_batch_positives
    max_real_same_batch_positives = args.max_real_same_batch_positives
    force_augmentation = args.force_augmentation
    t_second_chance = args.train_second_chance
    s_second_chance = args.sampling_second_chance
    required_augmentations = args.required_augmentations
    max_test_batches = args.max_test_batches
    target_directory=args.target_dir
    learning_rate=args.lr
    schedulerToUse=args.lr_scheduler
    schedulerStep=args.lr_scheduler_step
    uptraining=args.uptraining
    uptraining_start_epoch=args.uptraining_start_epoch
    uptraining_batch=args.uptraining_batch
    old_total_batch_size=args.old_total_batch_size
    use_muon=args.use_muon
    print("[MUOONING ]using muon ", use_muon)
    globals.USE_MUON=use_muon

    print("learning_rate: ", learning_rate)
    local_rank = int(os.environ["LOCAL_RANK"])

    # ---------------------- Initialize DDP ----------------------
    dist.init_process_group("nccl")
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    if rank == 0:
        print(f"[Rank {rank}] Using device: {device}, world_size: {world_size}")
    if not uptraining:
        size=batch_size
    else:
        size= uptraining_batch

    # ---------------------- DataLoader ----------------------
    with torch.no_grad():
        dataloader = customDataLoader(
            train_ratio=train_ratio,
            test_ratio=0,
            batch_size=size,
            seed=seed,
            max_same_batch_positives=max_same_batch_positives,
            max_real_same_batch_positives=max_real_same_batch_positives,
            second_chance=s_second_chance,
            ddp_enabled=True,
            downsample=not uptraining,
        )
    if old_total_batch_size == 0:
        file_batch_size = batch_size * world_size
    else:
        file_batch_size = old_total_batch_size
    print("file batch size usato",file_batch_size)
    # ---------------------- Checkpoint paths ----------------------
    if local_rank == 0:
        os.makedirs(f"{PROJECT_DIR}/checkpoints/{target_directory}", exist_ok=True)
        os.makedirs(f"{PROJECT_DIR}/encoder-losses/{target_directory}", exist_ok=True)
        if uptraining:
            os.makedirs(f"{PROJECT_DIR}/checkpoints/{target_directory}_uptrain-start{uptraining_start_epoch}", exist_ok=True)
            os.makedirs(f"{PROJECT_DIR}/encoder-losses/{target_directory}_uptrain-start{uptraining_start_epoch}", exist_ok=True)

    model_type = args.model_type
    model_filename = f"{PROJECT_DIR}/checkpoints/{target_directory}/{model_type}_model_train{train_ratio}_batch{file_batch_size}.pt"
    train_loss_filename = f"{PROJECT_DIR}/encoder-losses/{target_directory}/{model_type}_loss_train{train_ratio}_batch{file_batch_size}.pt"
    test_loss_filename = f"{PROJECT_DIR}/encoder-losses/{target_directory}/{model_type}_loss_test{train_ratio}_batch{file_batch_size}.pt"
    print("model_filename", model_filename)
    start_epoch, test_losses, train_losses = load_losses(test_loss_filename, train_loss_filename)
    print("start_epoch", start_epoch)
    print("uptraining", uptraining )
    print("uptraining_start_epoch",uptraining_start_epoch)
    if uptraining and "resnet" in model_type and start_epoch >= uptraining_start_epoch - 1:
        print("uptraining param refresh")
        #usare direttory dedicata ad uptrain
        up_batch_size=uptraining_batch
        up_file_batch_size=up_batch_size * world_size
        up_target_dir=target_directory+f"_uptrain-start{uptraining_start_epoch}"
        up_model_filename = f"{PROJECT_DIR}/checkpoints/{up_target_dir}/{model_type}_model_train{train_ratio}_batch{up_file_batch_size}.pt"
        up_train_loss_filename = f"{PROJECT_DIR}/encoder-losses/{up_target_dir}/{model_type}_loss_train{train_ratio}_batch{up_file_batch_size}.pt"
        up_test_loss_filename = f"{PROJECT_DIR}/encoder-losses/{up_target_dir}/{model_type}_loss_test{train_ratio}_batch{up_file_batch_size}.pt"
        branch_exists = (
                os.path.isfile(up_model_filename) and
                os.path.isfile(up_train_loss_filename) and
                os.path.isfile(up_test_loss_filename)
        )
        if not branch_exists:

            print("inizializzazione uptraining")
            # inizializziamo i parametri necessari a inizializzazione
            next_epoch = uptraining_start_epoch
            if rank == 0:
                if next_epoch != uptraining_start_epoch:
                    raise RuntimeError("probabile errore di logica")
            #     inizializzazione uptrained dir
                pre_epoch_upping = f"{PROJECT_DIR}/checkpoints/{target_directory}/{model_type}_model_train{train_ratio}_batch{file_batch_size}_epoch{uptraining_start_epoch-1}.pt"
                print("reading from ", pre_epoch_upping)
                checkpoint = torch.load(pre_epoch_upping, map_location=device)

                if "optimizer" in checkpoint:
                    checkpoint.pop("optimizer")
                else:
                    print("[INFO] Nessun optimizer nel checkpoint (ok per finetuning)")
                if "scheduler" in checkpoint:
                    checkpoint.pop("scheduler")
                else:
                    print("[INFO] Nessun scheduler nel checkpoint (ok per finetuning)")
                if "scaler" in checkpoint:
                    checkpoint.pop("scaler")
                else:
                    print("[INFO] Nessun scaler nel checkpoint (ok per finetuning)")


                train_losses= train_losses[:uptraining_start_epoch - 1]
                test_losses = test_losses[:uptraining_start_epoch - 1]

                save_safe(checkpoint, up_model_filename)
                save_safe(train_losses, up_train_loss_filename)
                save_safe(test_losses, up_test_loss_filename)
                print("initialized uptraining fine tuning")

        model_filename=up_model_filename
        train_loss_filename=up_train_loss_filename
        test_loss_filename=up_test_loss_filename
        batch_size=up_batch_size
        file_batch_size=up_file_batch_size
        target_directory=up_target_dir

        print("sovrascrizione per upping fine tuning")
        print("model_filename",model_filename)
        print("train_loss_filename",train_loss_filename)
        print("test_loss_filename",test_loss_filename)
        print("batch_size",batch_size)
        print("file_batch_size",file_batch_size)
        print("target_directory",target_directory)


        dist.barrier()
        start_epoch, test_losses, train_losses = load_losses(test_loss_filename, train_loss_filename)
        print(f"Rank {rank}, start_epoch {start_epoch}, ultima loss {train_losses[-1]}")


    # ---------------------- Model & DDP ----------------------

    model = rep.representation(model_type=model_type,downsampling=dataloader.must_downsample).to(device)
    if not use_muon:
        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    else:
        optimizer = optim.Muon(model.parameters(), lr=learning_rate)
    scaler = GradScaler()
    print(f"Using {schedulerToUse} scheduler with {schedulerStep} steps")
    scheduler = init_scheduler(learning_rate, num_epochs, optimizer, schedulerStep, schedulerToUse, start_epoch)

    resume_checkpoint(device, model, model_filename, optimizer, scaler, scheduler)

    if "vit" in model_type:
        needs_find=True
    else:
        needs_find=False
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=needs_find
    )

    test_set = dataloader.get_test_data(forModelTraining=False)
    test_sampler = dataloader.SimpleBatchSampler(test_set)


    loss_fn = EpsInfoNCE(temperature=0.1, epsilon=0).to(device)
    avg_epoch_time = None
    time_remaining = slurm_time_remaining()
    print(f"Tempo rimanente in secondi {time_remaining}")
    for ep in range(start_epoch, num_epochs):

        epoch_start_time = time.time()

        print(f"Starting Epoch {ep+1}")
        print(f"Are we downsampling?{dataloader.must_downsample}")
        dataloader.set_epoch(ep)
        model.train()
        n_batch = 0
        epoch_loss = []

        avg_train_loss = train(batch_size, dataloader, device, ep, epoch_loss,
                               force_augmentation, loss_fn, model, n_batch,
                               optimizer, scaler, t_second_chance, required_augmentations)
        if scheduler is not None:
            # if schedulerToUse == "IaD":
            #     scheduler.step(avg_train_loss)
            # else:
            scheduler.step()


        print(f"Epoch {ep} avg train loss: {avg_train_loss:.4f}")

        if rank == 0:
            checkpoint = {
                "model": model.module.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scaler": scaler.state_dict(),

            }

            if scheduler is not None:
                checkpoint["scheduler"] = scheduler.state_dict()

            save_safe(checkpoint, model_filename)
            train_losses.append(avg_train_loss)
            print(f"Epoch {ep+1}/{num_epochs} | Avg InfoNCE train loss: {avg_train_loss:.6f}")

            avg_test_loss = test(test_sampler, 50//2, device, model, loss_fn,max_batches=max_test_batches, required_augmentations=required_augmentations,must_downsample=dataloader.must_downsample) #fixed batch size for comparison
            test_losses.append(avg_test_loss)
            print(f"Epoch {ep+1}/{num_epochs} | Avg InfoNCE test loss: {avg_test_loss:.6f}")
            print(
                f"[{time.strftime('%X')}] Learning rate after scheduler step: {optimizer.param_groups[0]['lr']}")


            save_safe(train_losses, train_loss_filename)
            save_safe(test_losses, test_loss_filename)
            if (ep + 1) == 1 or ((ep + 1) >= 5 and (ep + 1) % 5 == 0):
                milestone_filename = f"{PROJECT_DIR}/checkpoints/{target_directory}/{model_type}_model_train{train_ratio}_batch{file_batch_size}_epoch{ep+1}.pt"

                save_safe(checkpoint, milestone_filename)
                print(f"Saved milestone checkpoint: {milestone_filename}")
        epoch_time = time.time() - epoch_start_time
        if avg_epoch_time is None:
            avg_epoch_time = epoch_time
        else:
            avg_epoch_time = 0.6 * avg_epoch_time + 0.4 * epoch_time
        # ------------------ TIME AWARE CHECK ------------------
        time_remaining = slurm_time_remaining()
        safety_margin = 1200  # 20  minuti di margine per non rischiare epoca piu lunga del solito fallisca
        # stop_training = torch.tensor(0, device=device,dtype=torch.int32)
        if time_remaining is not None and avg_epoch_time + safety_margin >= time_remaining:
            if rank == 0:
                print(f"[Rank {rank}] Not enough time for next epoch, exiting...")
                print(f"Remaining time: {time_remaining}")
                print(f"average_epoch_time: {avg_epoch_time}",flush=True)
                # stop_training.fill_(1)
                slurm_kill_job()

        # dist.barrier()
        # print(f"Rank {rank} pre broadcast")
        # torch.cuda.synchronize()
        # dist.all_reduce(stop_training, op=dist.ReduceOp.MAX)
        # stop_training=stop_training
        # print(f"Rank {rank} post broadcast")
        # print(f"Rank  {rank} stop training: {stop_training.item()}")

        # if stop_training.item()==1:
        #     print(f"Rank {rank} stopping")
        #     break
    if rank == 0:
        slurm_kill_job()


def load_losses(test_loss_filename, train_loss_filename):
    start_epoch = 0
    train_losses = []
    test_losses = []
    if os.path.isfile(train_loss_filename):
        print(f"Loading train loss history from {train_loss_filename}")
        train_losses = torch.load(train_loss_filename)
        start_epoch = len(train_losses)
    if os.path.isfile(test_loss_filename):
        print(f"Loading test loss history from {test_loss_filename}")
        test_losses = torch.load(test_loss_filename)
    return start_epoch, test_losses, train_losses


def resume_checkpoint(device, model, model_filename, optimizer, scaler, scheduler):
    if os.path.isfile(model_filename):
        print(f"Loading checkpoint from {model_filename}")
        checkpoint = torch.load(model_filename, map_location=device)
        if "model" in checkpoint:
            model.load_state_dict(checkpoint["model"])
            if "optimizer" in checkpoint:
                optimizer.load_state_dict(checkpoint["optimizer"])
            if "scaler" in checkpoint:
                scaler.load_state_dict(checkpoint["scaler"])
            else:
                print("⚠️ scaler not found in checkpoint, reinitializing")
            if "scheduler" in checkpoint and scheduler is not None:
                scheduler.load_state_dict(checkpoint["scheduler"])
            else:
                print("⚠️ scheduler not found in checkpoint, reinitializing")
        else:
            model.load_state_dict(checkpoint)


def init_scheduler(learning_rate, num_epochs, optimizer, schedulerStep, schedulerToUse, start_epoch):
    if schedulerToUse == "StepLR":
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=schedulerStep, gamma=0.75)
    elif schedulerToUse == "CosineAnnealing":
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=num_epochs,
            eta_min=1e-6
        )
    elif schedulerToUse == "warmup+CosineAnnealing":
        warmup_epochs = 10
        min_lr = 1e-6

        warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
            optimizer,
            start_factor=min_lr / learning_rate,
            end_factor=1.0,
            total_iters=warmup_epochs
        )

        cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=num_epochs - warmup_epochs,
            eta_min=min_lr
        )

        scheduler = torch.optim.lr_scheduler.SequentialLR(
            optimizer,
            schedulers=[warmup_scheduler, cosine_scheduler],
            milestones=[warmup_epochs]
        )

        if start_epoch > 0:
            for _ in range(start_epoch):
                scheduler.step()
    # elif schedulerToUse == "IaD":
    #     scheduler = IncrementalAndDecayScheduler(
    #         optimizer,
    #         increase_factor=1.05,
    #         stagnation_patience=0,  # quante epoche senza miglioramento prima di ridurre LR
    #         halving_factor=0.5,  # fattore di dimezzamento
    #         min_lr=1e-6  # LR minimo
    #     )
    elif schedulerToUse == "None":
        print("No scheduler")
        scheduler = None

    else:
        raise ValueError(f"Scheduler {schedulerToUse} not recognized")
    return scheduler


def load_image_worker(image_path, dataloader, normalize, downsampled, adapter:ModelImageAdapter):
    # carica immagine
    if downsampled:
        t = dataloader.load_tensor(image_path, normalize=normalize)
    else:
        t = dataloader.load_nii_as_tensor(image_path,downsample=downsampled, normalize=normalize)
    if adapter is not None:

        t=adapter.crop_batch(t)

    return image_path, t

def process_entry_from_already_loaded_tensor(entry, tensor, forbidden,
                              force_augmentation,
                              required_augmentations,
                              safe_max_value,
                              device):

    subject = entry["subject"]

    # da spostare ancora dentro le augmentazioni in un secondo momento
    tensor = tensor.to(device)

    with torch.no_grad():
        if entry.get("require_aug", force_augmentation):
            aug = get_random_augmentation(forbidden[subject])
            tensor = apply_augmentation(tensor, aug)
            forbidden[subject].append(aug)

            tmp = []
            tensor = generate_further_augs(
                aug,
                required_augmentations,
                safe_max_value,
                subject,
                tensor,
                tmp
            )

        if tensor.dim() == 4:
            tensor = tensor.unsqueeze(0)

    return tensor.detach().cpu(), subject

def train(batch_size, dataloader : customDataLoader, device, ep, epoch_loss, force_augmentation,
          loss_fn, model, n_batch, optimizer, scaler, t_second_chance, required_augmentations):
    import time
    torch.cuda.empty_cache()
    model_type = model.module.model_type

    must_downsample = dataloader.must_downsample
    print("must downsample in train ", must_downsample)
    target_shape, vitAdapter = get_shape_and_adapter_based_on_model(device, model_type, must_downsample)

    nan_happenings = 0
    safe_max_value = 1e3
    rank = dist.get_rank() if dist.is_initialized() else 0
    world_size = dist.get_world_size() if dist.is_initialized() else 1


    real_epoch_loss = []
    local_finished = False

    for name, param in model.named_parameters():
        assert param.requires_grad, f"{name} is frozen!"

    dataloader.generate_batch_init(epoch=ep)

    max_workers = batch_size // 2
    num_workers = (min(max_workers, (os.cpu_count()) // max(1, world_size)))*2


    if rank == 0:
        print(f"[{time.strftime('%X')}] [Rank {rank}] Workers starting to work: {num_workers}")

    num_workers_aug=(num_workers*3)//4
    num_workers_load = num_workers - num_workers_aug
    load_pool = ProcessPoolExecutor(max_workers=num_workers_load, mp_context=ctx)
    aug_pool = ProcessPoolExecutor(max_workers=num_workers_aug, mp_context=ctx)
    all_ranks_finished=False

    while not all_ranks_finished:

        forbidden = defaultdict(list)
        n_batch += 1
        print(f"[{time.strftime('%X')}] [Rank {rank}] Starting batch {n_batch}")
        print(
            f"[{time.strftime('%X')}] using Learning rate : {optimizer.param_groups[0]['lr']}")

        batch = dataloader.generate_batch(ep,second_chance=t_second_chance)
        is_real_batch = bool(batch)
        print(f"[{time.strftime('%X')}] [Rank {rank}] Generated batch of size {len(batch)}")

        if len(batch) < batch_size and rank == 0:
            print(f"[{time.strftime('%X')}] [Rank {rank}] Last batch detected: {len(batch)}")

        try:
            with torch.no_grad():
                batch_tensors = []
                labels = []

                if batch:
                    print(f"[{time.strftime('%X')}] [Rank {rank}] Starting preload for batch with {len(batch)} items")

                    # 1) DEDUP dei path
                    unique_paths = list({entry["image_path"] for entry in batch})

                    # 2) PRELOAD PARALLELO (POOL A)
                    preload_futures = {
                        path: load_pool.submit(
                            load_image_worker,
                            path,
                            dataloader,
                            True,  # normalize
                            downsampled=must_downsample,
                            adapter=vitAdapter,
                        )
                        for path in unique_paths
                    }

                    print(f"[{time.strftime('%X')}] [Rank {rank}] Submitted {len(unique_paths)} unique loads")

                    # POOL AUGMENTAZIONE
                    aug_futures = []
                    for entry in batch:
                        path = entry["image_path"]

                        # appena immagine caricata la passo alle augmentation
                        _, tensor = preload_futures[path].result()

                        fut = aug_pool.submit(
                            process_entry_from_already_loaded_tensor,
                            entry,
                            tensor,
                            forbidden,
                            force_augmentation,
                            required_augmentations,
                            safe_max_value,
                            device
                        )
                        aug_futures.append(fut)

                    print(f"[{time.strftime('%X')}] [Rank {rank}] Submitted {len(aug_futures)} augmentation jobs")

                    # 4) RACCOLTA RISULTATI
                    for fut in aug_futures:
                        tensor, subject = fut.result()
                        batch_tensors.append(tensor.to(device, non_blocking=True))
                        labels.append(subject)

                    print(
                        f"[{time.strftime('%X')}] [Rank {rank}] Completed augmentations: {len(batch_tensors)} tensors")



                #  dummy batch participates in forward/backward so DDP collectives remain balanced.
                if not batch_tensors:
                    print(f"[{time.strftime('%X')}] [Rank {rank}] No real data, using dummy batch")
                    D, H, W = target_shape
                    dummy_input = torch.zeros(1, 1, D, H, W, device=device)


                    dummy_label = f"dummy_{rank}"


                    batch_tensors = [dummy_input.clone() for _ in range(batch_size)]
                    labels = [dummy_label] * batch_size

                    is_real_batch = False
                    local_finished = True


                filler_called = False
                if batch and batch_tensors:

                    filler_called = filler(batch, batch_size, batch_tensors, dataloader, device,
                                           forbidden, labels, required_augmentations, safe_max_value)
                    if filler_called:

                        is_real_batch = False
                        print(f"[{time.strftime('%X')}] [Rank {rank}] Filler appended synthetic items")



                # Resize tensors to same shape
                all_views = [t[0] for t in batch_tensors]
                max_D = max(v.shape[1] for v in all_views)
                max_H = max(v.shape[2] for v in all_views)
                max_W = max(v.shape[3] for v in all_views)
                target_shape = (max_D, max_H, max_W)
                all_views_resized = [resize_tensor(v, target_shape).to(device, non_blocking=True)
                                     for v in all_views]


            print(f"[{time.strftime('%X')}] [Rank {rank}] Stacking input tensors")
            inputs = torch.stack(all_views_resized)

            del all_views, all_views_resized
            torch.cuda.empty_cache()
            # if "vit" in model_type: qua se dopo augmentation
            #     inputs=vitAdapter.crop_batch(inputs)

            print(f"[{time.strftime('%X')}] [Rank {rank}] Forward pass")
            print("shape inputs", inputs.shape)
            with torch.cuda.amp.autocast():
                projections = model(inputs).unsqueeze(1)

            print("projections ",projections.shape)

            if dist.is_initialized():


                gathered = dist_nnF.all_gather(projections)  # tuple di [B, D]
                projections_all = torch.cat(list(gathered), dim=0)  # [B*world_size, D]


                all_labels = [None for _ in range(world_size)]
                dist.all_gather_object(all_labels, labels)
                labels_all = sum(all_labels, [])  # concatena tutte le liste

            else:
                projections_all = projections
                labels_all = labels


            unique_subjects = {subj: idx for idx, subj in enumerate(dict.fromkeys(labels_all))}
            numeric_labels_global = torch.tensor(
                [unique_subjects[subj] for subj in labels_all],
                dtype=torch.long, device=device
            )


            loss_value =  loss_fn(projections_all, numeric_labels_global)
            if torch.isnan(loss_value) or torch.isinf(loss_value):
                nan_happenings += 1
                print(f"[{time.strftime('%X')}] [Rank {rank}] ⚠️ NaN/Inf loss; replacing with zeros")
                loss_value = torch.zeros_like(loss_value)


            scaler.scale(loss_value).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)


            # record losses only if real
            epoch_loss.append(loss_value.item())
            if is_real_batch:
                real_epoch_loss.append(loss_value.item())

            print(f"[{time.strftime('%X')}] [Rank {rank}] Loss: {loss_value.item()} | NaN happenings: {nan_happenings}")

            # cleanup per-batch tensors
            del loss_value, inputs, projections, batch_tensors #numeric_labels

        finally:

            try:
                del forbidden
            except Exception:
                pass
            try:
                del batch
            except Exception:
                pass
            torch.cuda.empty_cache()


        if dist.is_initialized():

            finished_tensor = torch.tensor([1 if local_finished else 0], dtype=torch.long, device=device)

            dist.all_reduce(finished_tensor, op=dist.ReduceOp.SUM)
            total_finished = int(finished_tensor.item())
            if rank == 0:
                print(f"[{time.strftime('%X')}] [Rank {rank}] total_finished={total_finished}/{world_size}")

            if total_finished >= world_size:
                print(f"[{time.strftime('%X')}] [Rank {rank}] All ranks finished, breaking training loop")
                all_ranks_finished=True
        else:

            if local_finished:
                all_ranks_finished=True

    # Shutdown executor after loop ends
    print(f"[{time.strftime('%X')}] [Rank {rank}] Shutting down executor")
    load_pool.shutdown(wait=True)
    aug_pool.shutdown(wait=True)
    print(f"[{time.strftime('%X')}] [Rank {rank}] Executor shutdown complete")


    local_sum = sum(real_epoch_loss)
    local_count = len(real_epoch_loss)

    if dist.is_initialized():
        stats = torch.tensor([local_sum, local_count], dtype=torch.float64, device=device)
        dist.all_reduce(stats, op=dist.ReduceOp.SUM)
        global_sum = float(stats[0].item())
        global_count = int(stats[1].item())
        global_avg = (global_sum / global_count) if global_count > 0 else 0.0
    else:
        global_avg = (local_sum / local_count) if local_count > 0 else 0.0

    print(f"[{time.strftime('%X')}] [Rank {rank}] ✅ Epoch complete. Global avg real-batch loss: {global_avg:.6f}")
    return global_avg


def get_shape_and_adapter_based_on_model(device, model_type, must_downsample):
    vitAdapter=None
    D = 91 if must_downsample else 182
    H = 109 if must_downsample else 218
    W = 91 if must_downsample else 182
    dummy_input = torch.zeros(1, 1, D, H, W, device=device)
    if "vit" in model_type:
        vitAdapter = ModelImageAdapter((16, 16, 16))
        _, _, D, H, W = vitAdapter.crop_batch(dummy_input).shape
    target_shape = (D, H, W)

    return target_shape, vitAdapter


def filler(batch, batch_size, batch_tensors, dataloader : customDataLoader, device, forbidden, labels, required_augmentations,
           safe_max_value):
    downsampled = dataloader.must_downsample
    called=False
    while len(labels) < batch_size:
        called=True
        entry_to_duplicate = random.choice(batch)
        subject = entry_to_duplicate["subject"]
        print("Filling with downsample?",downsampled)

        if downsampled:
            tensor = dataloader.load_tensor(entry_to_duplicate["image_path"], normalize=True).to(device)
        else:
            tensor = dataloader.load_nii_as_tensor(entry_to_duplicate["image_path"],downsample=downsampled, normalize=True).to(device)

        aug = get_random_augmentation(forbidden[subject])

        tensor = apply_augmentation(tensor, aug)
        forbidden[subject].append(aug)
        tensor = generate_further_augs(aug, required_augmentations, safe_max_value, subject, tensor, [])

        if tensor.dim() == 4:
            tensor = tensor.unsqueeze(0)

        batch_tensors.append(tensor)
        labels.append(subject)
    return called


def generate_further_augs(aug, required_augmentations, safe_max_value, subject, tensor, tmp):
    for _ in range(required_augmentations - 1):
        if aug == "sobel":
            tmp.append("random_noise")
        if aug == "random_noise":
            tmp.append("sobel")
        aug = get_random_augmentation(tmp)
        # print(f"[{subject}] Applying augmentation: {aug}")

        if torch.isnan(tensor).any() or torch.isinf(tensor).any():
            print(f"⚠️ Augmentation '{aug}' produced NaN or Inf values. Replacing with zeros.")
            tensor = torch.zeros_like(tensor)


        max_val = tensor.abs().max()
        if max_val > safe_max_value:
            print(
                f"⚠️ Augmentation '{aug}' produced too high values ({max_val:.2e}). Clamping to ±{safe_max_value}.")
            tensor = torch.clamp(tensor, min=-safe_max_value, max=safe_max_value)

        tensor = apply_augmentation(tensor, aug)
    return tensor
def process_test_entry(entry, test_sampler, forbidden,device, must_downsample, adapter: ModelImageAdapter):

    subject = entry["subject"]
    image_path = entry["image_path"]


    tensor = test_sampler.load_npz_as_tensor(image_path, normalize=True,downsample=must_downsample).to(device)
    print("load test is downsample", must_downsample)
    # Ensure shape [1, C, D, H, W]
    while tensor.dim() < 5:
        tensor = tensor.unsqueeze(0)
    if adapter is not None:

        tensor=adapter.crop_batch(tensor)

    results = []
    for _ in range(2):  # two augments for one img
        aug = get_random_augmentation(forbidden[subject])
        forbidden[subject].append(aug)


        t = apply_augmentation(tensor.clone(), aug).squeeze(0)
        results.append((t, subject))

    return results

@torch.no_grad()
def test(test_sampler, batch_size, device, model, loss_fn, max_batches=30,must_downsample=True, required_augmentations=2):
    torch.cuda.empty_cache()
    model.eval()
    model_type = model.module.model_type
    vitAdapter = None

    D = 91 if must_downsample else 182
    H = 109 if must_downsample else 218
    W = 91 if must_downsample else 182
    dummy_input = torch.zeros(1, 1, D, H, W, device=device)

    if "vit" in model_type:
        vitAdapter = ModelImageAdapter((16, 16, 16))
        _, _, D, H, W = vitAdapter.crop_batch(dummy_input).shape

    target_shape = (D, H, W)
    print("target_shape", target_shape)

    test_sampler.reset()

    epoch_loss = []
    n_batch = 0

    max_workers = batch_size
    num_workers = min(max_workers, os.cpu_count())
    executor = ProcessPoolExecutor(max_workers=num_workers, mp_context=ctx)

    while test_sampler.dataset_copy:
        n_batch += 1
        if n_batch > max_batches:
            break


        batch = test_sampler.sample_batch(batch_size)
        if not batch:
            break
        if len(batch) < batch_size:
            print(f"--- last batch: {len(batch)} ---")

        forbidden = defaultdict(list)
        cpu_views = []
        labels = []


        futures = [executor.submit(process_test_entry, entry, test_sampler, forbidden, device, must_downsample,vitAdapter)
                   for entry in batch]
        for future in as_completed(futures):
            results = future.result()
            for tensor, subject in results:
                cpu_views.append(tensor)  # keep on CPU
                labels.append(subject)

        del futures
        torch.cuda.empty_cache()

        if not cpu_views:


            dummy_input = torch.zeros(1, 1, D, H, W)
            cpu_views = [dummy_input.clone(), dummy_input.clone()]
            labels = [0, 1]


        max_D = max(v.shape[1] for v in cpu_views)
        max_H = max(v.shape[2] for v in cpu_views)
        max_W = max(v.shape[3] for v in cpu_views)
        target_shape = (max_D, max_H, max_W)

        resized_views = [resize_tensor(v, target_shape) for v in cpu_views]
        del cpu_views
        torch.cuda.empty_cache()


        batch_tensor = torch.stack(resized_views).to(device)
        del resized_views
        torch.cuda.empty_cache()


        unique_subjects = {subj: idx for idx, subj in enumerate(dict.fromkeys(labels))}
        numeric_labels = torch.tensor([unique_subjects[subj] for subj in labels],
                                      dtype=torch.long, device=device)
        del labels
        torch.cuda.empty_cache()


        with torch.cuda.amp.autocast(), torch.no_grad():
            projection = model(batch_tensor)
            projection = projection.unsqueeze(1)
            #print("Projection stats:", projection.shape, projection.min().item(), projection.max().item())
            # print("Any NaN in projection?", torch.isnan(projection).any())

            loss_value = loss_fn(projection, numeric_labels)

        epoch_loss.append(loss_value.item())
        print(f"Batch Loss: {loss_value.item()}")


        del batch_tensor, projection, numeric_labels
        torch.cuda.empty_cache()

    executor.shutdown(wait=True)
    avg_loss = sum(epoch_loss) / len(epoch_loss) if epoch_loss else 0.0
    print(f"✅ TEST complete. Avg Loss: {avg_loss:.4f}")
    return avg_loss




if __name__ == "__main__":
    import argparse
    args = parse_args()
    main(args)

#
# def pad_tensors_to_max_shape(tensor_list):
#     """
#     Pads a list of 4D tensors [C, D, H, W] to the max shape along each spatial dimension.
#     Returns a single tensor of shape [len(tensor_list), C, max_D, max_H, max_W].
#     """
#     # Find max size along each dimension
#     max_C = max(t.size(0) for t in tensor_list)
#     max_D = max(t.size(1) for t in tensor_list)
#     max_H = max(t.size(2) for t in tensor_list)
#     max_W = max(t.size(3) for t in tensor_list)
#
#     padded_tensors = []
#     for t in tensor_list:
#         C, D, H, W = t.shape
#
#         # Check channel dimension matches
#         if C != max_C:
#             raise ValueError(f"Expected channel dimension {max_C}, but got {C}")
#
#         # Calculate padding: (pad_W_left, pad_W_right, pad_H_left, pad_H_right, pad_D_left, pad_D_right)
#         pad = (0, max_W - W,
#                0, max_H - H,
#                0, max_D - D)
#
#         # F.pad expects padding in reverse order of dims, so order is (W,H,D)
#         padded = F.pad(t, pad, mode='constant', value=0)
#         padded_tensors.append(padded)
#
#     return torch.stack(padded_tensors)  # Shape: [n_tensors, C, max_D, max_H, max_W]
# def pad_views_to_max_views(proj_list):
#     """
#     proj_list: list of tensors [V_i, proj_dim]
#     pads each to max V with zeros so we can stack into [B, max_V, proj_dim]
#     """
#     max_views = max(p.shape[0] for p in proj_list)
#     proj_dim = proj_list[0].shape[1]
#     padded = []
#     for p in proj_list:
#         if p.shape[0] < max_views:
#             pad_size = max_views - p.shape[0]
#             pad_tensor = torch.zeros(pad_size, proj_dim, device=p.device, dtype=p.dtype)
#             p = torch.cat([p, pad_tensor], dim=0)
#         padded.append(p)
#     return padded
#


#  old train for backup
# def train(batch_size, dataloader, device, ep, epoch_loss, force_augmentation,
#           loss_fn, model, n_batch, optimizer, scaler, t_second_chance, required_augmentations):
#     import time
#     nan_happenings = 0
#     safe_max_value = 1e3
#     rank = dist.get_rank() if dist.is_initialized() else 0
#     world_size = dist.get_world_size() if dist.is_initialized() else 1
#
#     target_shape = (64, 64, 64)  # default for dummy batches
#     real_epoch_loss = []  # track only real batch losses (per-rank)
#     local_finished = False   # becomes True when this rank has exhausted local data (inserted dummy)
#
#     for name, param in model.named_parameters():
#         assert param.requires_grad, f"{name} is frozen!"
#
#     dataloader.generate_batch_init()
#
#     max_workers = batch_size // 2
#     num_workers = min(max_workers, os.cpu_count() // max(1, world_size))
#     if rank == 0:
#         print(f"[{time.strftime('%X')}] [Rank {rank}] Workers starting to work: {num_workers}")
#
#     executors = ProcessPoolExecutor(max_workers=num_workers, mp_context=ctx)
#
#     # Loop until *all* ranks finished
#     while True:
#         # At loop start, clear per-batch state
#         forbidden = defaultdict(list)
#         n_batch += 1
#         print(f"[{time.strftime('%X')}] [Rank {rank}] Starting batch {n_batch}")
#
#         batch = dataloader.generate_batch(second_chance=t_second_chance)
#         is_real_batch = bool(batch)
#         print(f"[{time.strftime('%X')}] [Rank {rank}] Generated batch of size {len(batch)}")
#
#         if len(batch) < batch_size and rank == 0:
#             print(f"[{time.strftime('%X')}] [Rank {rank}] Last batch detected: {len(batch)}")
#
#         try:
#             with torch.no_grad():
#                 batch_tensors = []
#                 labels = []
#
#                 # Process batch entries if available
#                 if batch:
#                     print(f"[{time.strftime('%X')}] [Rank {rank}] Mapping batch entries in executor")
#                     results = list(executors.map(
#                         process_batch_entry,
#                         batch,
#                         [dataloader] * len(batch),
#                         [forbidden] * len(batch),
#                         [force_augmentation] * len(batch),
#                         [required_augmentations] * len(batch),
#                         [safe_max_value] * len(batch),
#                         [device] * len(batch),
#                     ))
#                     print(f"[{time.strftime('%X')}] [Rank {rank}] Executor mapping done")
#                     for tensor, subject in results:
#                         batch_tensors.append(tensor.to(device, non_blocking=True))
#                         labels.append(subject)
#                     print(f"[{time.strftime('%X')}] [Rank {rank}] Collected {len(batch_tensors)} tensors")
#
#                 # If no real data for THIS rank, insert dummy batch and mark local_finished.
#                 # Note: dummy batch participates in forward/backward so DDP collectives remain balanced.
#                 if not batch_tensors:
#                     print(f"[{time.strftime('%X')}] [Rank {rank}] No real data, using dummy batch")
#                     D, H, W = target_shape
#                     dummy_input = torch.zeros(1, 1, D, H, W, device=device)
#                     # duplicate so loss function has something valid to compute (pairs)
#                     batch_tensors = [dummy_input.clone(), dummy_input.clone()]
#                     labels = [0, 1]
#                     is_real_batch = False
#                     local_finished = True   # THIS RANK has no more local data
#
#                 # Call filler only if we had initial real batch entries.
#                 # filler may append synthetic items to reach batch_size.
#                 filler_called = False
#                 if batch and batch_tensors:
#
#                     filler_called = filler(batch, batch_size, batch_tensors, dataloader, device,
#                                            forbidden, labels, required_augmentations, safe_max_value)
#                     if filler_called:
#                         # filled batches are synthetic (don't count in real_epoch_loss),
#                         # but DO NOT mark local_finished here. We must continue participating
#                         # in DDP until all ranks are done.
#                         is_real_batch = False
#                         print(f"[{time.strftime('%X')}] [Rank {rank}] Filler appended synthetic items")
#
#                 # Prepare numeric labels
#                 unique_subjects = {subject: idx for idx, subject in enumerate(dict.fromkeys(labels))}
#                 numeric_labels = torch.tensor([unique_subjects[subj] for subj in labels],
#                                               dtype=torch.long, device=device)
#
#                 # Resize tensors to same shape
#                 all_views = [t[0] for t in batch_tensors]
#                 max_D = max(v.shape[1] for v in all_views)
#                 max_H = max(v.shape[2] for v in all_views)
#                 max_W = max(v.shape[3] for v in all_views)
#                 target_shape = (max_D, max_H, max_W)
#                 all_views_resized = [resize_tensor(v, target_shape).to(device, non_blocking=True)
#                                      for v in all_views]
#
#             # Stack and forward/backward as usual (participate)
#             print(f"[{time.strftime('%X')}] [Rank {rank}] Stacking input tensors")
#             inputs = torch.stack(all_views_resized)
#             del all_views, all_views_resized
#             torch.cuda.empty_cache()
#
#             print(f"[{time.strftime('%X')}] [Rank {rank}] Forward pass")
#             with torch.cuda.amp.autocast():
#                 projections = model(inputs).unsqueeze(1)
#
#             loss_value = loss_fn(projections, numeric_labels)
#             if torch.isnan(loss_value) or torch.isinf(loss_value):
#                 nan_happenings += 1
#                 print(f"[{time.strftime('%X')}] [Rank {rank}] ⚠️ NaN/Inf loss; replacing with zeros")
#                 loss_value = torch.zeros_like(loss_value)
#
#             # Normal backward / optimizer step
#             scaler.scale(loss_value).backward()
#             scaler.unscale_(optimizer)
#             torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
#             scaler.step(optimizer)
#             scaler.update()
#             optimizer.zero_grad(set_to_none=True)
#
#             # record losses: epoch_loss (maybe used elsewhere) and real_epoch_loss only if real
#             epoch_loss.append(loss_value.item())
#             if is_real_batch:
#                 real_epoch_loss.append(loss_value.item())
#
#             print(f"[{time.strftime('%X')}] [Rank {rank}] Loss: {loss_value.item()} | NaN happenings: {nan_happenings}")
#
#             # cleanup per-batch tensors
#             del loss_value, inputs, projections, numeric_labels, batch_tensors
#
#         finally:
#             # ensure we always clear these so subsequent iterations are clean
#             try:
#                 del forbidden
#             except Exception:
#                 pass
#             try:
#                 del batch
#             except Exception:
#                 pass
#             torch.cuda.empty_cache()
#
#         # At the end of each batch make sure all ranks exchange their 'finished' state:
#         # each rank contributes 1 if it has no local data left (local_finished), else 0.
#         if dist.is_initialized():
#             # must use a CUDA tensor for NCCL backend
#             finished_tensor = torch.tensor([1 if local_finished else 0], dtype=torch.long, device=device)
#             # all_reduce sum of finished flags
#             dist.all_reduce(finished_tensor, op=dist.ReduceOp.SUM)
#             total_finished = int(finished_tensor.item())
#             if rank == 0:
#                 print(f"[{time.strftime('%X')}] [Rank {rank}] total_finished={total_finished}/{world_size}")
#             # If every rank reported finished -> safe to exit loop
#             if total_finished >= world_size:
#                 print(f"[{time.strftime('%X')}] [Rank {rank}] All ranks finished, breaking training loop")
#                 break
#         else:
#             # single-process: break if local_finished
#             if local_finished:
#                 break
#
#     # Shutdown executor after loop ends
#     print(f"[{time.strftime('%X')}] [Rank {rank}] Shutting down executor")
#     executors.shutdown(wait=True)
#     print(f"[{time.strftime('%X')}] [Rank {rank}] Executor shutdown complete")
#
#     # Aggregate real batch loss across ranks (global average over only real batches)
#     local_sum = sum(real_epoch_loss)
#     local_count = len(real_epoch_loss)
#
#     if dist.is_initialized():
#         stats = torch.tensor([local_sum, local_count], dtype=torch.float64, device=device)
#         dist.all_reduce(stats, op=dist.ReduceOp.SUM)
#         global_sum = float(stats[0].item())
#         global_count = int(stats[1].item())
#         global_avg = (global_sum / global_count) if global_count > 0 else 0.0
#     else:
#         global_avg = (local_sum / local_count) if local_count > 0 else 0.0
#
#     print(f"[{time.strftime('%X')}] [Rank {rank}] ✅ Epoch complete. Global avg real-batch loss: {global_avg:.6f}")
#     return global_avg