import tensorflow.compat.v1 as tf
import csv
from copy import deepcopy
import time
import json
import numpy as np
import string
import random
import logging
import os
from datetime import datetime

# 配置日志记录器


tf.disable_v2_behavior()
TF_XLA_FLAGS="--tf_xla_auto_jit=1"

DTYPE = {
    'float32': tf.float32, 
    'int32': tf.int32, 
    'int64': tf.int64,
    'string': tf.string
}


def create_directory(path):
    if not os.path.exists(path):
        os.makedirs(path)
        print(f"Directory '{path}' created.")
    else:
        print(f"Directory '{path}' already exists.")


def read_csv_data(file_path):
    with open(file_path, 'r') as f:
        reader = csv.reader(f)
        next(reader)
        datas = list(reader)
    for data in datas:
        data[2] = data[2][1:-1].split(",")
    
    return datas


def load_graph(model_file):
    with tf.gfile.GFile(model_file, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name="")
    
    return graph
 
    
def generate_graph_datas(graph, tensors, batch_size):
    graph_datas = {}
    graph_names = {}
    for tensor in tensors:
        graph_names[tensor[1].split(':')[0]] = graph.get_tensor_by_name(tensor[1])
        _dtype = DTYPE[tensor[-1]]
        shapes = deepcopy(tensor[2])
        if shapes == ['']:
            graph_datas[graph.get_tensor_by_name(tensor[1])] = tf.constant(3.14, dtype=_dtype)
            continue
        if shapes[0] == "None":
            shapes[0] = batch_size
        for i in range(len(shapes)):
            if shapes[i] == "None":
                shapes[i] = batch_size
            if shapes[i] != "":
                shapes[i] = int(shapes[i])
        
        shapes = tuple([i for i in shapes])
       
        if tensor[-1] == "int32":
            random_tensor = tf.random.uniform(shape=shapes, minval=0, maxval=10, dtype=_dtype)

        elif tensor[-1] == "int64":
            random_tensor = tf.random.uniform(shape=shapes, minval=0, maxval=10, dtype=_dtype)
    
        elif tensor[-1] == "string":
            # 生成字符串张量
            batch_size = shapes[0]
            sequence_length = shapes[1] if len(shapes) > 1 else 1
            random_tensor = tf.constant([["example_string"] * sequence_length] * batch_size, dtype=tf.string)
        else:
            random_tensor = tf.random.normal(shape=shapes, mean=0.0, stddev=1.0, dtype=_dtype)

        
        graph_datas[graph.get_tensor_by_name(tensor[1])] = random_tensor

    return graph_datas, graph_names       


def load_datasets(datasets_path):
    with open(datasets_path, 'r') as f:
        datasets = json.load(f)
    return datasets


def main():
    model = "model_1"
    model_dir = "./models"
    random_flag = False
    
    logs_dir = os.path.join(model_dir, f"{model}/logs")
    create_directory(logs_dir)
    
    logging.basicConfig(filename=os.path.join(logs_dir, f"TF-{model}-{datetime.now().strftime('%Y%m%d%H%M%S')}.log"),
                        filemode='a', 
                        format='%(asctime)s - %(levelname)s - %(message)s',
                        level=logging.INFO)
    
    input_tensors_path = os.path.join(model_dir, f'{model}/input_tensors.csv')
    output_tensors_path = os.path.join(model_dir, f'{model}/output_tensors.csv')
    model_path = os.path.join(model_dir, f'{model}/model.pb')
    
    
    input_tensors = read_csv_data(input_tensors_path)
    output_tensors = read_csv_data(output_tensors_path)
    
    batch_size = [1,2,4,8,16,32,64,128,256,512,1024,2048]   
    
    with tf.device('/gpu:0'):
        graph = load_graph(model_path)
        with graph.as_default():
            with tf.Session(graph=graph) as sess:
                for bs in batch_size:
                    if random_flag:
                        print("random input data")
                        input_datasets, _ = generate_graph_datas(graph, input_tensors, bs)
                        _, output_names = generate_graph_datas(graph, output_tensors, bs)

                        input_values = sess.run(list(input_datasets.values()))
                        feed_dict = dict(zip(input_datasets.keys(), input_values))                             
                    else:
                        input_datasets_path = os.path.join(model_dir, f"{model}/dataset/input_tensor_datas_{bs}.json")
                        output_datasets_path = os.path.join(model_dir, f"{model}/dataset/output_tensor_datas_{bs}.json")
                        
                        print(f"Load input data from json file: {input_datasets_path}")
                        
                        input_datasets = load_datasets(input_datasets_path)
                        output_datasets = load_datasets(output_datasets_path)
                        input_tensors_dict = {graph.get_tensor_by_name(k): tf.convert_to_tensor(v, dtype=graph.get_tensor_by_name(k).dtype) for k, v in input_datasets.items()}
                        output_names = {k: graph.get_tensor_by_name(k) for k in output_datasets.keys()}
                        
                        input_values = sess.run(list(input_tensors_dict.values()))
                        feed_dict = dict(zip(input_tensors_dict.keys(), input_values))
                    for i in range(130):
                        if i < 30:
                            times = time.time()
                        sess.run(output_names, feed_dict=feed_dict)
                    QPS = 100/(time.time() - times) * bs
                    logging.info(f"*******batch_size: {bs}  *******QPS: {QPS}")
                    print(f"*******batch_size: {bs}  *******QPS: {QPS}")


if __name__ == '__main__':
    
    main()