import os
import torch

from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer


class NTUHLPLTrainer(nnUNetTrainer):
    """nnU-Net trainer for iterative pseudo-labeling on NTUH T1c tumor segmentation.

    Env knobs (read in __init__, applied before initialize() builds optimizer/LR):
      NNU_PL_EPOCHS   total epochs of this run   (default 250)
      NNU_PL_LR       initial lr, PolyLR decay   (default 1e-2)

    Warm start (round-to-round fine-tuning):
      NNU_PL_WARMSTART  path to an nnU-Net checkpoint. Its FULL network weights,
      segmentation head included, are loaded in on_train_start(). The CLI
      -pretrained_weights flag deliberately skips `.seg_layers.` keys, which is
      wrong for iterative pseudo-labeling, hence this hook.
    """

    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict,
                 device: torch.device = torch.device("cuda")):
        super().__init__(plans, configuration, fold, dataset_json, device)
        self.num_epochs = int(float(os.environ.get("NNU_PL_EPOCHS", "250")))
        self.initial_lr = float(os.environ.get("NNU_PL_LR", "1e-2"))
        self._pl_warmstart_done = False

    def on_train_start(self):
        super().on_train_start()
        if self._pl_warmstart_done:
            return
        self._pl_warmstart_done = True
        ws = os.environ.get("NNU_PL_WARMSTART")
        if not ws:
            return
        mod = self.network
        if hasattr(mod, "module"):
            mod = mod.module
        if hasattr(mod, "_orig_mod"):
            mod = mod._orig_mod
        ckpt = torch.load(ws, map_location=self.device, weights_only=False)
        w = ckpt["network_weights"]
        mod.load_state_dict(w, strict=True)
        torch.cuda.empty_cache()
        print(f"[NTUHLPLTrainer] warm-started full network weights from {ws}", flush=True)
        self.print_to_log_file(f"warm-started full network weights from {ws}")