"""Test that a model can be saved and loaded again."""

from pathlib import Path

import torch
from dotenv import load_dotenv
from torchinfo import summary

from model.src.models import VisualCommunicationModel
from shared.data_store import connect, put_model
from shared.utils import setup_logging

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if __name__ == '__main__':
    # load in env file
    env_path = Path(__file__).parent.parent / 'server.env'
    assert env_path.exists()
    load_dotenv(env_path)
    # setup logging
    setup_logging()
    # connect to minio
    client = connect()
    # instantiate model
    model = VisualCommunicationModel().to(DEVICE)
    # show model weights
    summary(model)
    # put buffer in minio bucket
    hash_str = put_model(
        client=client,
        model=model,
    )
    print(f"hash string: {hash_str}")
    print('saved to Minio')