from io import BytesIO
import threading
import ffmpeg
import sys
from PIL import Image
from numpy import average, dot, linalg
import json
import time
from numpy.lib.utils import info
import os

class pageThread (threading.Thread):
    def __init__(self, threadID, name, bframe, eframe,total_frame,step):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.name = name
        self.bindex = bframe
        self.eindex = eframe
        self.totalframe = total_frame
        self.step = step

    def run(self):
        # print("开始线程：" + self.name)
        frame_compare(self.bindex,self.eindex,self.totalframe,self.step)
        print("退出线程：" + self.name)

def startThread(id, name, bframe, eframe,totalframe,step):
    print(str(bframe) + ' ' + str(eframe))
    thread = pageThread(id, name, bframe, eframe+1,totalframe,step)
    thread.start()

# 对图片进行统一化处理
def get_thum(image, size=(64, 64), greyscale=False):
    # 利用image对图像大小重新设置, Image.ANTIALIAS为高质量的
    image = image.resize(size, Image.ANTIALIAS)
    if greyscale:
        # 将图片转换为L模式，其为灰度图，其每个像素用8个bit表示
        image = image.convert('L')
    return image

# 计算图片的余弦距离
def image_similarity_vectors_via_numpy(image1, image2):
    image1 = get_thum(image1)
    image2 = get_thum(image2)
    images = [image1, image2]
    vectors = []
    norms = []
    for image in images:
        vector = []
        for pixel_tuple in image.getdata():
            vector.append(average(pixel_tuple))
        vectors.append(vector)
        # linalg=linear（线性）+algebra（代数），norm则表示范数
        # 求图片的范数？？
        norms.append(linalg.norm(vector, 2))
    a, b = vectors
    a_norm, b_norm = norms
    # dot返回的是点积，对二维数组（矩阵）进行计算
    res = dot(a / a_norm, b / b_norm)
    return res


def read_frame_by_time(in_file, time):
    """
    指定时间节点读取任意帧
    """
    out, err = (
        ffmpeg.input(in_file, ss=time,loglevel='quiet')
              .output('pipe:', vframes=1, format='image2', vcodec='mjpeg')
              .run(capture_stdout=True)
    )
    return out


def save_time_frame(filepath, savepath, time):
    """
    保存视频的指定帧
    """
    out = read_frame_by_time(filepath, time)
    with open(savepath, 'wb') as img:
        img.write(out)


def read_frame_as_jpeg(in_file, frame_num):
    """
    指定帧数读取任意帧
    """
    out, err = (
        ffmpeg.input(in_file,loglevel='quiet')
              .filter('select', 'gte(n,{})'.format(frame_num))
              .output('pipe:', vframes=1, format='image2', vcodec='mjpeg')
              .run(capture_stdout=True)
    )
    return out


def save_frame(filepath, savepath, frame_num):
    """
    保存视频的指定帧
    """
    out = read_frame_as_jpeg(filepath, frame_num)
    with open(savepath, 'wb') as img:
        img.write(out)


def get_video_info(in_file):
    """
    获取视频基本信息
    """
    try:
        probe = ffmpeg.probe(in_file)
        video_stream = next(
            (stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
        if video_stream is None:
            print('No video stream found', file=sys.stderr)
            sys.exit(1)
        return video_stream
    except ffmpeg.Error as err:
        print(str(err.stderr, encoding='utf8', errors='ignore'))
        sys.exit(1)

def save_loop_info(start_time,end_time,path):
    """
    写入loop信息到本地
    """
    global config_file_path

    infos = {}
    infos['startTime'] = start_time * 1000
    infos['endTime'] = end_time * 1000
    infos['path'] = path

    infoStr = json.dumps(infos)
    with open(config_file_path,'w') as fp:
        fp.write(infoStr)

def frame_compare(part_start_frame, end_frame, total_frame, step):
    """
    图像相似度比较
    """
    global start_frame, isfinded,most_same_most_similar_info,file_path,amout_totalThread

    sim_value = 0

    try:
        for frame in range(part_start_frame, end_frame+1, step):

            # 如果有一个线程已经找到了匹配的结果，设置全局变量isfinished=True,以便其它线程退出
            if isfinded:
                break
            out2 = read_frame_as_jpeg(file_path, frame)
            img2 = Image.open(BytesIO(out2))

            sim_value = image_similarity_vectors_via_numpy(img1, img2)
            print('---------index: ' + str(frame) + ' 相似度： ' + str(sim_value))

            #记录所有关键帧中匹配度最高的帧
            if(sim_value >= most_same_most_similar_info['simvalue']):
                most_same_most_similar_info['simvalue'] = sim_value
                most_same_most_similar_info['frame_index'] = frame

            #相似度超过了阈值，匹配完成
            if sim_value >= threhold_sim_value:
                print('找到了符合的关键帧 end frame: ' + str(frame))
                end_frame = frame + extra_offset
                isfinded = True

                #写入找到的信息
                p_start_time = start_frame / total_frame * total_duration
                p_end_time = end_frame / total_frame * total_duration
                p_len_time = p_end_time - p_start_time
                print('匹配成功，循环起止时间：%s ~ %s 起止帧：%s ~ %s  总时长：%s 相似度：%s' %
                    (p_start_time, p_end_time,start_frame,end_frame, p_len_time, sim_value))
                save_loop_info(p_start_time,p_end_time,file_path)
                print('准备退出...')

                break    
    except:
        print('线程出错')
    finally:
        amout_totalThread-=1
       	


if __name__ == '__main__':
    # infoStr = r'{{MatchedFiles}}'
    # infoObjects = json.loads(infoStr)
    # finfo = infoObjects[0]
    # file_path = finfo['path']
    # file_path = 'D:/SCRIPT_CREATE/PYTHON_SCRIPTS/VideoPerfectLoop/test.mp4'
    print(sys.argv[0])
    config_file_path = sys.argv[0].replace('VideoPerfectLoop.py','config.json')
    file_path = sys.argv[1]

    isfinded = False
    video_info = get_video_info(file_path)
    total_duration = float(video_info['duration'])
    total_frame = int(video_info['nb_frames'])
    start_offset = int(sys.argv[2])
    end_frame = int(sys.argv[3])
    
    threhold_sim_value = 0.9999

    start_frame = 0
    extra_offset = 0
    step = 1
    amountPerThread = 20
    amout_totalThread = 0
    most_same_most_similar_info = {'simvalue':0,'frame_index':0}
    # 保存原始图
    out1 = read_frame_as_jpeg(file_path, start_frame)
    img1 = Image.open(BytesIO(out1))
    img2 = None

    # 开启下载线程
    count = 0
    bindex = start_frame + start_offset
    for i in range(bindex,end_frame + 1):
        #已经到了末尾页
        if (end_frame+1) - bindex <= amountPerThread:
            eindex = end_frame
            amout_totalThread+=1
            startThread(i,str(i),bindex,eindex,total_frame,step)
            break

        count += 1
        if count >= amountPerThread:
            count = 0
            eindex = i
            amout_totalThread+=1
            startThread(i,str(i),bindex,eindex,total_frame,step)
            bindex = eindex + 1

    while True:
        time.sleep(1)
        if amout_totalThread == 0:
            if not isfinded:
                end_frame = most_same_most_similar_info['frame_index']
                p_start_time = start_frame / total_frame * total_duration
                p_end_time = end_frame / total_frame * total_duration
                p_len_time = p_end_time - p_start_time
                print('没有匹配成功，已知最佳起止时间：%s ~ %s 起止帧：%s ~ %s  总时长：%s 相似度：%s' %
                    (p_start_time, p_end_time,start_frame,end_frame, p_len_time, most_same_most_similar_info['simvalue']))
                save_loop_info(p_start_time,p_end_time,file_path)
            break
    
    os.system('pause')