import os

import monai.networks.nets
import torch
from transformers import AutoConfig, AutoModel, PreTrainedModel
from vista3d_config import VISTA3DConfig


class VISTA3DModel(PreTrainedModel):
    """VISTA3D model for hugging face"""

    config_class = VISTA3DConfig

    def __init__(self, config):
        super().__init__(config)
        if config.model_type == "VISTA3D":
            self.network = monai.networks.nets.vista3d132(
                encoder_embed_dim=config.encoder_embed_dim,
                in_channels=config.input_channels,
            )

    def forward(self, input):
        return self.network(input)


def register_my_model():
    """Utility function to register VISTA3D model so that it can be instantiate by the AutoModel function."""
    AutoConfig.register("VISTA3D", VISTA3DConfig)
    AutoModel.register(VISTA3DConfig, VISTA3DModel)


if __name__ == "__main__":
    FILE_PATH = os.path.dirname(__file__)
    MODEL_WEIGHT_PATH = os.path.join(FILE_PATH, "models/model.pt")
    MODEL_PATH = os.path.join(FILE_PATH, "vista3d_pretrained_model")
    config = VISTA3DConfig()
    hugging_face_model = VISTA3DModel(config)
    hugging_face_model.network.load_state_dict(torch.load(MODEL_WEIGHT_PATH))
    hugging_face_model.save_pretrained(MODEL_PATH)