import torch
from phate import PHATEAE
from funcs.som import ClusterSOM

from funcs.dataloader import BaseDataset2, read_json_files


DEVICE = torch.device("cpu")

reducer10d = PHATEAE(epochs=30, n_components=10, lr=.0001, batch_size=128, t='auto', knn=8, relax=True, metric='euclidean')
reducer10d.load('models/r10d_2.pth')

cluster_som = ClusterSOM()
cluster_som.load("models/cluster_som2.pkl")

train_x, train_y  = read_json_files('output.json')
# Convert tensors to numpy arrays if necessary
if isinstance(train_x, torch.Tensor):
    train_x = train_x.numpy()
if isinstance(train_y, torch.Tensor):
    train_y = train_y.numpy()

# load the time series slices of the data 4*3*2*64 (feeds+axis*sensor*samples) + 5 for time diff
data = BaseDataset2(train_x.reshape(len(train_x), -1) / 32768, train_y)

#compute the 10 dimensional embeding vector
embedding10d = reducer10d.transform(data)

prediction = cluster_som.predict(embedding10d)
cluster_som.plot_activation(embedding10d)