import json
import matplotlib.pyplot as plt
import argparse
import sys
# Load the data from the provided JSON file
def draw_pdf(json_file_path,pdf_file_path,bk):
    plt.switch_backend('Agg')

    json_file_path = json_file_path

    with open(json_file_path, 'r') as file:
        data = json.load(file)

    # Extracting the relevant data
    losses = []
    accuracy = []
    memory = []

    for entry in data:
        result = entry['result']
        losses.append(result['Losses'])
        accuracy.append(result['Accuracy'])
        memory.append(result['MemoryInfoList'])

    # Plotting the data
    plt.figure(figsize=(10, 6))

    # Plot Losses
    plt.plot(losses, label='Losses', marker='o')

    # Plot Accuracy
    plt.plot(accuracy, label='Accuracy', marker='x')

    # Plot Memory Info List
    plt.plot(memory, label='Memory', marker='s')

    # Adding title and labels
    plt.title(bk)
    plt.xlabel('Iterations')
    plt.ylabel('Values')
    plt.legend()

    # Display the plot
    plt.grid(True)

    pdf_file_path = pdf_file_path
    plt.savefig(pdf_file_path)

    plt.close()

def draw_pdf_all(tensorflow_json_path,mxnet_json_path,losses_pdf_path, accuracy_pdf_path, memory_pdf_path):

    with open(tensorflow_json_path, 'r') as file:
        tensorflow_data = json.load(file)

    with open(mxnet_json_path, 'r') as file:
        mxnet_data = json.load(file)

    # Extracting the relevant data
    tf_losses = [entry['result']['Losses'] for entry in tensorflow_data]
    tf_accuracy = [entry['result']['Accuracy'] for entry in tensorflow_data]
    tf_memory = [entry['result']['MemoryInfoList'] for entry in tensorflow_data]

    mx_losses = [entry['result']['Losses'] for entry in mxnet_data]
    mx_accuracy = [entry['result']['Accuracy'] for entry in mxnet_data]
    mx_memory = [entry['result']['MemoryInfoList'] for entry in mxnet_data]

    # Creating plots for Losses, Accuracy, and Memory
    # Separate the plots into individual figures and save each as a PDF

    # Plot Losses
    plt.figure(figsize=(10, 6))
    plt.plot(tf_losses, label='TensorFlow Losses', marker='o')
    plt.plot(mx_losses, label='MXNet Losses', marker='x')
    plt.title('Losses Over Iterations')
    plt.xlabel('Iterations')
    plt.ylabel('Losses')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()

    plt.savefig(losses_pdf_path)
    plt.close()

    # Plot Accuracy
    plt.figure(figsize=(10, 6))
    plt.plot(tf_accuracy, label='TensorFlow Accuracy', marker='o')
    plt.plot(mx_accuracy, label='MXNet Accuracy', marker='x')
    plt.title('Accuracy Over Iterations')
    plt.xlabel('Iterations')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()

    plt.savefig(accuracy_pdf_path)
    plt.close()

    # Plot Memory
    plt.figure(figsize=(10, 6))
    plt.plot(tf_memory, label='TensorFlow Memory', marker='o')
    plt.plot(mx_memory, label='MXNet Memory', marker='x')
    plt.title('Memory Usage Over Iterations')
    plt.xlabel('Iterations')
    plt.ylabel('Memory Info List')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()

    plt.savefig(memory_pdf_path)
    plt.close()



if __name__ == "__main__":

    """Parser of command args"""
    parse = argparse.ArgumentParser()
    parse.add_argument("--backend", type=str, help="name of backends")
    parse.add_argument("--model_path", type=str, help="redis db port")
    flags, unparsed = parse.parse_known_args(sys.argv[1:])

    bk = ['tensorflow', 'mxnet']
    model_path = flags.model_path
    for i in bk:
        if 'svhn' in model_path or 'fashion2' in model_path:
            txt_path = model_path.split("\\")[-1][:-5].split("mut_model")[0]+i+"_train.json"
        else:
            txt_path = model_path.split("\\")[-1][:-3].split("mut_model")[0]+i+"_train.json"

        pdf_path = txt_path.replace(".json",".jpg")
        draw_pdf(txt_path,pdf_path,i)

    tensorflow_json_path = model_path.split("\\")[-1][:-5].split("mut_model")[0]+'tensorflow_train'+".json"
    mxnet_json_path = model_path.split("\\")[-1][:-5].split("mut_model")[0]+'mxnet_train'+".json"

    losses_pdf_path = model_path.split("\\")[-1][:-5].split("mut_model")[0]+"losses"+".jpg"
    accuracy_pdf_path = model_path.split("\\")[-1][:-5].split("mut_model")[0]+"accuracy"+".jpg"
    memory_pdf_path = model_path.split("\\")[-1][:-5].split("mut_model")[0]+"memory"+".jpg"
    

    draw_pdf_all(tensorflow_json_path,mxnet_json_path,losses_pdf_path, accuracy_pdf_path, memory_pdf_path)