import os.path
import shutil


from ytracking.track_ import *
from contrast.test_logic import group_image, inference
from tools.Interface import AiInterface, AiClass

from tools.initModel import models
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

models.initModel()
ai_obj = AiClass()
distance_lists = []
all_distance_lists = []


def showImg(newx):
    plt.scatter(newx[:, 0], newx[:, 1])
    plt.show()


def cosine_similarity(vec1, vec2):
    vec1 = np.array(vec1)
    vec2 = np.array(vec2)
    cos_sim = vec1.dot(vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
    return cos_sim


def get_kclusters(data, k):
    Kmeans = KMeans(n_clusters=k, n_init='auto')
    y_pred = Kmeans.fit_predict(data)
    return y_pred


def move_file(pth, dirs, labels, y_pred):
    for label, y_label in zip(labels, y_pred):
        if not os.path.isdir(os.sep.join([pth, dirs, str(y_label)])):
            os.mkdir(os.sep.join([pth, dirs, str(y_label)]))
        try:
            shutil.move(os.sep.join([pth, dirs, label]),
                        os.sep.join([pth, dirs,  str(y_label), label]))
        except Exception as e:
            print(e)


def createDic(pth, dirs):
    all_dics = {}
    imgs = os.sep.join([pth, dirs])
    for img in os.listdir(imgs):
        imgPth = os.sep.join([imgs, img])
        feat_tensor = inference([imgPth], models.similarityModel, actionModel=False)
        all_dics[img] = feat_tensor
    return all_dics


def main(pth, dirs):
    global distance_lists
    allDics = createDic(pth, dirs)
    labels = []
    temp = None
    for key, value in allDics.items():
        value = value.cpu().detach().numpy()
        labels.append(key)
        if temp is None:
            temp = value
        else:
            temp = np.vstack((temp, value))
    y_pred = get_kclusters(temp, 3)
    move_file(pth, dirs, labels, y_pred)


if __name__ == '__main__':
    pth = 'C:/Users/HP/Desktop/40084'
    dirs = '1'
    main(pth, dirs)