# WISDM activity recognition with imuposer_wisdm.pt (HiT, right-wrist smartwatch, 18 activities).
# Run from a checkout of IMU2Language at commit 156f014 with the files of this repository in HiT-v1/ (see Setup).
from pathlib import Path

import numpy as np
import torch
from imu2language.eval import har_datasets as H
from imu2language.eval.har_zeroshot import batches
from imu2language.models.hit import load_hit

H.HAR_ROOT = Path('HiT-v1/wisdm')                  # WISDM/{X,y}_{train,test}.npy + WISDM.json from this repository


@torch.no_grad()
def classify(model, clf, samples, device='cuda'):
    """samples: har_datasets.load output -> predicted class index per window."""
    C = clf['class_vectors'].to(device)
    pred = np.zeros(len(samples), int)
    for idx, b in batches(samples):
        b = {k: v.to(device) for k, v in b.items()}
        with torch.autocast('cuda', torch.bfloat16):
            taps, g, pad, _ = model.encode(b['imu'], b['frame_ok'], b['slot'], b['sens_pad'], b['gyro_ok'],
                                           b['token_ok'], b['pad'])
            z = model.span_embed(1, taps, g, b['txt1_span'], pad)[:, 0].float()    # whole window = one span
        pred[idx] = (clf['logit_scale'] * z @ C.T).argmax(1).cpu().numpy()
    return pred


if __name__ == '__main__':
    model, ck = load_hit('HiT-v1/imuposer_wisdm.pt', 'cuda')
    clf = ck['wisdm_classifier']
    test, names = H.load('WISDM', split='test')      # 20 Hz watch windows -> 60 Hz, right-wrist slot, at most 10 s
    assert names == clf['classes']
    y = np.array([s['label'] for s in test])
    p = classify(model, clf, test)
    print(f'WISDM test: {len(y)} windows, accuracy {100 * (p == y).mean():.1f}%')