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

from io import BytesIO

import torch
from torchinfo import summary

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

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

if __name__ == '__main__':
    # check that environment variables are set
    check_env()
    # setup logging
    setup_logging()
    # instantiate model
    model = VisualCommunicationModel().to(DEVICE)
    # show model weights
    summary(model)
    # for param in model.parameters():
    #     logging.info(param.data)
    # save model to buffer
    buffer = BytesIO()
    torch.save(model.state_dict(), buffer)
    print(f"buffer size: {len(buffer.getvalue())}")
    # connect to minio
    client = connect(bucket='models')
    # put buffer in minio bucket
    hash_str = put(
        client=client,
        buffer=buffer,
    )
    print(f"hash string: {hash_str}")
    print('saved to Minio')