import datetime

from numba import cuda, float32, int32
import numpy as np
import math


def text_read(f):
    file = open(f, 'r')
    lines = file.readlines()
    return lines


def trans_source_data_to_tuple(data):
    ds = data.split()
    return float(ds[0]), float(ds[1]), float(ds[2]), float(ds[3])


def trans_test_data_to_tuple(data):
    ds = data.split()
    return float(ds[1]), float(ds[2]), float(ds[3]), float(ds[4])


def search_file(test, source):
    time_start_load_data = datetime.datetime.now()

    list_source = text_read(source)
    list_test = text_read(test)

    time_end_load_data = datetime.datetime.now()
    print('读文件用时：%s' % (time_end_load_data - time_start_load_data))

    result, result_frame = calculate_source_position_by_gpu(list_source, list_test)
    return result, result_frame


def calculate_source_position_by_gpu(list_source, list_test):
    source_features = np.array(list(map(lambda d: trans_source_data_to_tuple(d), list_source)))
    test_features = np.array(list(map(lambda d: trans_test_data_to_tuple(d), list_test)))
    source_features_len = len(source_features)
    test_features_len = len(test_features)

    source_len = source_features_len - test_features_len
    test_len = test_features_len
    print('source_features_len=%d source_len=%d test_len=%d' % (source_features_len, source_len, test_len))

    threads_per_block = 1024  # TODO: 应从设备信息获取
    blocks_per_grid = math.ceil(source_len / threads_per_block)
    total_thread_numb = threads_per_block * blocks_per_grid
    print('threads_per_block=%d  blocks_per_grid=%d total_thread_numb=%d'
          % (threads_per_block, blocks_per_grid, total_thread_numb))

    # copy memory from Host to Device
    time_copy_mem_h2d_begin = datetime.datetime.now()
    d_source_features = cuda.to_device(source_features)
    d_test_features = cuda.to_device(test_features)
    d_gpu_result = cuda.device_array(source_len)
    cuda.synchronize()
    time_copy_mem_h2d_end = datetime.datetime.now()
    print('内存copy用时：%s' % (time_copy_mem_h2d_end - time_copy_mem_h2d_begin))

    # call kernel thread to calculate
    compare_frame_by_kernel[blocks_per_grid, threads_per_block](source_len, test_len,
                                                                d_source_features, d_test_features,
                                                                d_gpu_result)
    cuda.synchronize()
    time_gpu_calculation_end = datetime.datetime.now()
    print('*** 纯GPU计算用时：%s' % (time_gpu_calculation_end - time_copy_mem_h2d_end))

    results = np.array([0.0] * source_len, dtype=np.float32)
    d_gpu_result.copy_to_host(results)
    cuda.synchronize()

    result, result_frame = find_most_probability(results)
    time_end_cal = datetime.datetime.now()
    print('匹配度概率为: %.2f %%, 匹配位置在抽帧样本的第 %d 帧,即原视频 %d 分处' % (result * 100, result_frame, (result_frame % 300)))
    print('*** 匹配用时：%s' % (time_end_cal - time_copy_mem_h2d_end))
    return result, result_frame


def find_most_probability(results):
    result = 0
    result_frame = 0
    for i in range(len(results)):
        if results[i] > result:
            result_frame = i + 1
            result = results[i]
    return result, result_frame


@cuda.jit
def compare_frame_by_kernel(source_len, test_len,
                            d_src_frames_features, d_test_frames_features,
                            d_results):
    """
     GPU kernel函数。
    :param source_len:
    :param test_len:
    :param d_src_frames_features: 源视频特征数据序列。Array of tuple
    :param d_test_frames_features: 待检测视频的特征数据序列。
    :param d_results:
    :return:
    """
    idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x
    if idx >= source_len:
        return

    src_frame_start_index = idx
    result_of_from_idx = calculate_features_probability_at_src_idx(src_frame_start_index, test_len,
                                                                   d_src_frames_features, d_test_frames_features)
    cuda.atomic.add(d_results, src_frame_start_index, result_of_from_idx)


@cuda.jit(device=True, argtypes=float32(int32, int32, float32[:], float32[:]))
def calculate_features_probability_at_src_idx(source_start_index, test_len,
                                              d_src_frames_features, d_test_frames_features):
    accumulator = 0
    for test_idx in np.arange(test_len):
        src_idx = source_start_index + test_idx
        accumulator = compare_two_frame_feature(d_src_frames_features[src_idx], d_test_frames_features[test_idx], accumulator)
    result_of_from_idx = accumulator / test_len / 8
    return result_of_from_idx


@cuda.jit(device=True)
def compare_two_frame_feature(src_frame_feature, test_frame_feature, init_val):
    """
    计算两帧特征的相似度。
    :param src_frame_feature: 原帧的特征数据。(Top features with possibility. Tuple of float.)
    :param test_frame_feature: 待检测帧的特征数据。
    :param init_val: 初始值。
    :return: 比较结果
    """
    if src_frame_feature[0] == test_frame_feature[0]:
        init_val += 6 - abs((src_frame_feature[1] - test_frame_feature[1])) * 18
    elif src_frame_feature[0] == test_frame_feature[2]:
        init_val += 2 - abs((src_frame_feature[1] - test_frame_feature[3])) * 2
    elif src_frame_feature[2] == test_frame_feature[0]:
        init_val += 2 - abs((src_frame_feature[3] - test_frame_feature[1])) * 2
    elif src_frame_feature[2] == test_frame_feature[2]:
        init_val += 2 - abs((src_frame_feature[3] - test_frame_feature[3])) * 2
    return init_val


def print_gpu_info():
    print(cuda.gpus)


if __name__ == '__main__':
    print_gpu_info()

    start_time = datetime.datetime.now()
    # source_files = ['../data/10_features_of_source.txt']
    # for src_file in source_files:
    #     search_file('../data/3_features_of_test_at5.txt', src_file)
    #     search_file('../data/3_features_of_test_at7.txt', src_file)

    source_files = ['../data/冰上恋人_01.txt']
    for src_file in source_files:
        search_file('../data/test008_AIresult_top5.txt', src_file)

    end_time = datetime.datetime.now()
    print("\n\nTotalTime: %s" % (end_time - start_time))
