import os
import shutil
import numpy as np
import torch

from Graph import Grapher
from Data_for_gcn import load_predictionData
from torch_geometric.utils.convert import from_networkx


def clearOutputFolder(test_output_fd):
    shutil.rmtree(test_output_fd)
    if not os.path.exists(test_output_fd):
        os.mkdir(test_output_fd)


save_modelPath = os.path.join("C:\\AsianGCN\\saved_model\\saved_model.pt")
model = torch.load(save_modelPath)

def make_info(document):
    connect = Grapher(document)
    G, _, _ = connect.graph_formation()
    df = connect.relative_distance(700,850)
    individual_data = from_networkx(G)
    return G, df, individual_data

def predictionsUnique(saved_model_folder, prediction_dataset_folder, document_to_predict,labels_tab):
    predictions_data = load_predictionData(save_fd=prediction_dataset_folder)

    model = torch.load(saved_model_folder)

    y_preds = model.max(dim=1)[1].cpu().numpy()

    test_batch = predictions_data.batch.cpu().numpy()
    sample_indexes = np.where(test_batch == 0)[0]
    y_pred = y_preds[sample_indexes]

    print("Beginning of the prediction")
    G, df, individual_data, img = make_info(document_to_predict)

    assert len(y_pred) == df.shape[0]

    predictions = []
    for row_index, row in df.iterrows():
        _y_pred = y_pred[row_index]
        _label = labels_tab[_y_pred]
        if _label != 'undefined':
            _text = row['Object']
            predictions.append({'Text': _text, 'Label': _label})

    return predictions

def load_predictionData(save_fd):
    predictions_data = torch.load(os.path.join(save_fd, 'predictions_data.dataset'))
    return predictions_data

if __name__ == "__main__":
    predictionsUnique()