import torch

from EfficientNet_2d.EfficientNet_2d import get_pretrained_EfficientNet, get_pretrained_DAR

if __name__ == "__main__":

    # Phase 1:

    # pre-train prd-net, cf-net, and lr-net on CR-set, IC-set, and LR-set, respectively, and save the pre-trained model

    prd_net = get_pretrained_EfficientNet(num_classes=5)

    cf_net = get_pretrained_EfficientNet(num_classes=5)

    lr_net = get_pretrained_EfficientNet(num_classes=5)

    # Phase 2:

    # fine-tune dar on CR-set

    prd_params_path = "../your_checkpoint_path/prd_net_checkpoint.pth"

    cf_params_path = "../your_checkpoint_path/cf_net_checkpoint.pth"

    lr_params_path = "../your_checkpoint_path/lr_net_checkpoint.pth"

    prd_params = torch.load(prd_params_path)

    cf_params = torch.load(cf_params_path)

    lr_params = torch.load(lr_params_path)

    model = get_pretrained_DAR(prd_params, cf_params, lr_params, num_classes=5)

    # prediction

    imgs = torch.rand([4, 3, 224, 224])

    prd_preds, cf_preds, lr_preds = model(imgs)

    _, preds = torch.softmax(prd_preds, dim=1).max(dim=1)