import torch
from vnet_light_arch import VNetLight

def load_light_model(checkpoint_path, channels):
  print(f"Loading The Saved Model at {checkpoint_path}")
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  model = VNetLight(in_channels=channels, classes=1)
  model.to(device)
  model.restore_checkpoint(checkpoint_path)
  for param in model.parameters():
    param.requires_grad = False
  return model

def load_vChain_model(wt_chkpt, tc_chkpt, et_chkpt):
  wt_model = load_light_model(wt_chkpt, 4)
  tc_model = load_light_model(tc_chkpt, 8)
  et_model = load_light_model(et_chkpt, 8)
  return wt_model, tc_model, et_model

# see this for modifying loss function:
# https://docs.monai.io/en/latest/losses.html