import torch
import sys
sys.path.append("/home/pranayr_umass_edu/imu2clip")
from lib.imu_models import MW2StackRNNPooling

if __name__ == "__main__":

    # Generate random IMU-like motions as examples
    # imu_motions: array <n_samples x 6 x 1000>
    imu_motions = torch.rand(1, 6, 1000)
    print(imu_motions.dtype)
    print("Generated random IMU-like motions:", imu_motions)

    # Load the IMU encoder
    """
        The following example .pt model is configured as
        - i2c: IMU2CLIP
        - s_i: source modality = IMU
        - t_v: target modality for alignment = Video
        - t_t: target modality for alignment = Text
        - mw2: MW2StackRNNPooling as the encoder
        - w_5.0: window size of 2.5 x 2 seconds
    """
    #path_imu_encoder = "./i2c_s_i_t_v_ie_mw2_w_5.0_master_imu_encoder.pt"
    # path_imu_encoder = "./i2c_s_i_t_t_ie_mw2_w_2.5_master_imu_encoder.pt"
    path_imu_encoder = "/home/pranayr_umass_edu/imu2clip/shane_models/i2c_s_i_t_v_ie_mw2_w_5.0_master_imu_encoder.pt"
    path_imu_encoder = '/home/pranayr_umass_edu/imu2clip/saved/mw2/i2c/i2c_s_i_t_v_se_mw2_w_5.0_master-epoch=12-val_loss=3.14.ckpt'

    loaded_imu_encoder = MW2StackRNNPooling(size_embeddings=512)
    loaded_imu_encoder.load_state_dict(torch.load(path_imu_encoder))
    print("Loaded IMU Encoder:", loaded_imu_encoder)
    loaded_imu_encoder.eval()
    print("Done loading the IMU Encoder")

    # Inference time
    imu2clip_embeddings = loaded_imu_encoder(imu_motions)
    print('Raw IMU Signals (random)', imu_motions.shape)
    print('Encoded IMU2CLIP embeddings', imu2clip_embeddings.shape)

    # # check if raw and encoded are the same
    # print('Are the raw and encoded the same?', torch.allclose(imu_motions, imu2clip_embeddings))