from utils.data_loader import *
from encoder.modeling_chart2vec import *

# Load Chart2Vec model
chart2vec_model = Chart2Vec().to(device)
path = os.path.dirname(__file__)
model_save_path = os.path.join(
    path, "models/chart2vec_base.pth")
state_dict = torch.load(model_save_path, map_location=torch.device(device))
chart2vec_model.load_state_dict(state_dict['model'])
chart2vec_model.eval()

# Prepare the chart 7 tuple as input
chart_fact = {"fact_type": "trend", "subspace": [], "breakdown": [{"field": "Year of release", "type": "temporal"}], "measure": [
    {"field": "Rating", "aggregate": "avg", "type": "numerical"}], "focus": [], "chart_type": "Area Chart", "meta": "increasing"}
# Alignment of data (multiple charts placed in a list is also supported)
input_chart = pad_w2v_convert_fact_to_input([chart_fact])

# Represent charts as vectors with Chart2Vec
output_vec = chart2vec_model(
    np.array([input_chart["batch_struct_one_hot"][0]]),
    np.array([input_chart["batch_indexed_tokens"][0]]),
    np.array([input_chart["batch_pos"][0]]))
# Output
print(output_vec[0].tolist())

import torch
import json
from utils.data_loader import *
from encoder.modeling_chart2vec import *
from encoder.ImprovedQuadrupletLoss import *

def get_chart2vec_embeddings(chart2vec_model, chart_data):
    return chart2vec_model(chart_data["batch_struct_one_hot"], chart_data["batch_indexed_tokens"], chart_data["batch_pos"])

# The hyperparameter settings of the model can be adjusted according to the actual situation
epochs = 50
batch_size = 128
lr = 1e-2

# Setting up the machine on which the model will run
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load training data and build batch
with open("./dataset/training_data.json") as f:
        train_data = json.load(f)
batch_gen = quadruplet_batch_generator_word2vec(train_data, batch_size)

# multitasking loss function
loss_fc = ImprovedQuadrupletLoss().to(device)
chart_emb = Chart2Vec().to(device)
optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, chart_emb.parameters()), lr=lr)

# Iterative training model
best_loss = 1e7
for i in range(epochs):
    total_loss = 0.0
    for j in range(len(train_data)//batch_size):
        # Represent the charts in each training sample as vectors
        batch_chart_x1, batch_chart_x2, batch_chart_x3, batch_chart_y1 = next(batch_gen)
        res_chart_x1 = get_chart2vec_embeddings(chart_emb,batch_chart_x1)
        res_chart_x2 = get_chart2vec_embeddings(chart_emb,batch_chart_x2)
        res_chart_x3 = get_chart2vec_embeddings(chart_emb,batch_chart_x3)
        res_chart_y1= get_chart2vec_embeddings(chart_emb,batch_chart_y1)
        # calculate the loss function
        chart2vec_loss = loss_fc(res_chart_x1,  res_chart_x2, res_chart_x3, res_chart_y1)
        total_loss += chart2vec_loss.cpu().item()
        # Backpropagation for optimization
        chart2vec_loss.backward()
        optimizer.step()
        chart_emb.zero_grad()
        print("mean_loss:", chart2vec_loss.item())

    # Saving of models
    if total_loss < best_loss:
        best_loss = total_loss      
        torch.save({'model': chart_emb.state_dict()}, os.path.join("YOUR_MODEL_SAVE_PATH/chart2vec_base.pth"))

import os
import json
import torch
from tqdm import tqdm
import numpy
from utils.data_loader import *
from utils.constants import *
from encoder.modeling_chart2vec_word_max_pooling import *


def test_and_evaluate_chart2vec(trained_model_path, testing_data_file):
    """
        test and evaluate the chart2vec model
        Inputs:
            `model_save_path`: The directory where model iterations are saved during the training process, the overall directory is located under models.
            `training_data_file`: Filename of the training data.
    """
    # get the path to the current file
    path = os.path.dirname(__file__)
    # get the testing data
    with open(os.path.join(path, testing_data_file)) as f:
        test_data = json.load(f)

    # load the model and the parameters of the trained model
    chart2vec_model = Chart2Vec().to(device)
    model_save_path = os.path.join(
        path, "models", trained_model_path+"chart2vec_base.pth")
    state_dict = torch.load(model_save_path, map_location=torch.device(device))
    chart2vec_model.load_state_dict(state_dict['model'])
    chart2vec_model.eval()

    # construct test data into the input format required by the model
    input = pad_w2v_convert_fact_to_input(test_data)
    batch_struct_one_hot = input["batch_struct_one_hot"]
    batch_indexed_tokens = input["batch_indexed_tokens"]
    batch_pos = input["batch_pos"]
    input_len = len(batch_struct_one_hot)

    # represent each chart in the test data as a vector and store it in the result_dict
    result_dict = {}
    for i in tqdm(range(input_len), desc="data_num"):
        output_vec = chart2vec_model(
            np.array([batch_struct_one_hot[i]]), np.array([batch_indexed_tokens[i]]), np.array([batch_pos[i]]))
        result_dict[test_data[i]["fact_id"]] = output_vec[0].tolist()

    # evaluate the results
    acc2, acc3, acc_co = evaluate_more_nearest_dis_triplets(
        result_dict, search_num=1)
    print(acc2, acc3, acc_co)


def get_models_folder(dir_name):
    models_folder_list = os.listdir(dir_name)
    for folder in models_folder_list:
        if os.path.isfile(os.path.join(dir_name, folder)):
            models_folder_list.remove(folder)
    return models_folder_list


def evaluate_more_nearest_dis_triplets(data, search_num=1):
    """
        Validates Chart2Vec's modeling results against the vectors generated by the model. Specify a chart, calculate the top k nearest charts, determine if it is in the context window, and calculate the context distance.
    """
    same_dataset_facts, accuracy_list_window2, accuracy_list_window3, accuracy_list_story = {}, {}, {}, {}
    last_dataset_key = ""
    for key in list(data.keys()):
        dataset_label = key.split("-", 1)[0]
        if len(same_dataset_facts) == 0 or (len(same_dataset_facts) > 0 and dataset_label == list(same_dataset_facts.keys())[0].split("-", 1)[0]):
            same_dataset_facts[key] = data[key]
            last_dataset_key = key.split("-", 1)[0]
        else:
            # First calculate the accuracy between facts in the previous dataset
            accuracy_list_window2[last_dataset_key],  accuracy_list_window3[last_dataset_key], accuracy_list_story[last_dataset_key] = cal_same_dataset_facts_more_dis_min(
                same_dataset_facts, search_num)
            same_dataset_facts = {}
            same_dataset_facts[key] = data[key]
    accuracy_list_window2[key], accuracy_list_window3[key], accuracy_list_story[key] = cal_same_dataset_facts_more_dis_min(
        same_dataset_facts, search_num)
    return cal_avg_accuracy_value(accuracy_list_window2), cal_avg_accuracy_value(accuracy_list_window3), cal_avg_accuracy_value(accuracy_list_story)


def cal_avg_accuracy_value(accuracy_list):
    """
        Calculate the overall average
    """
    accuracy_list = dict(sorted(accuracy_list.items(),
                         key=lambda d: d[1], reverse=True))
    return numpy.mean(list(accuracy_list.values()))


def cal_same_dataset_facts_more_dis_min(facts_dict, search_num):
    """
        Calculate the closest chart based on the anchor chart
    """
    dis_map = {}
    facts_dict_key = list(facts_dict.keys())
    for i in range(len(facts_dict)):
        temp_dis_map = {}
        key1 = facts_dict_key[i]
        fact1_id = key1.split("-", 1)[1]
        value1 = numpy.array(facts_dict[key1])
        for j in range(len(facts_dict)):
            if i == j:
                continue
            key2 = facts_dict_key[j]
            fact2_id = key2.split("-", 1)[1]
            value2 = numpy.array(facts_dict[key2])
            dis = eucliDist(value1, value2)
            if len(temp_dis_map.keys()) < search_num:
                temp_dis_map[fact2_id] = dis
            else:
                # If the value is already stored, it looks for the maximum value and replaces it if it is smaller than it.
                temp_max_key = max(temp_dis_map, key=temp_dis_map.get)
                if dis < temp_dis_map[temp_max_key] and dis > 0:
                    del temp_dis_map[temp_max_key]
                    temp_dis_map[fact2_id] = dis
        dis_map[fact1_id] = temp_dis_map
    is_context, is_context_window2, is_context_window3, is_context_story = 0, 0, 0, 0
    for key in dis_map.keys():
        fact1_id = key
        for fact2 in dis_map[key].keys():
            fact2_id = fact2
            if fact1_id.split("-")[0] != fact2_id.split("-")[0]:
                continue
            if int(fact1_id.split("-")[1])+1 == int(fact2_id.split("-")[1]) or int(fact1_id.split("-")[1])-1 == int(fact2_id.split("-")[1]):
                is_context += 1
                is_context_window2 += 1
                is_context_window3 += 1
                is_context_story += 1
                break
            if int(fact1_id.split("-")[1])+2 == int(fact2_id.split("-")[1]) or int(fact1_id.split("-")[1])-2 == int(fact2_id.split("-")[1]):
                is_context_window2 += 1
                is_context_window3 += 1
                is_context_story += 1
                break
            if int(fact1_id.split("-")[1])+3 == int(fact2_id.split("-")[1]) or int(fact1_id.split("-")[1])-3 == int(fact2_id.split("-")[1]):
                is_context_window3 += 1
                is_context_story += 1
                break
            if int(fact1_id.split("-")[0]) == int(fact2_id.split("-")[0]):
                is_context_story += 1
                break
    window2_accuracy = float(is_context_window2/len(dis_map.keys()))
    window3_accuracy = float(is_context_window3/len(dis_map.keys()))
    story_accuracy = float(is_context_story/len(dis_map.keys()))
    return window2_accuracy, window3_accuracy, story_accuracy


def eucliDist(A, B):
    return numpy.sqrt(sum(numpy.power((A - B), 2)))


if __name__ == "__main__":
    test_and_evaluate_chart2vec(trained_model_path="models/chart2vec_base.pth", testing_data_file="dataset/testing_data_50.json")