import os
import sys
import time

sys.path.insert(0, './yolov5')

from yolov5.models.experimental import attempt_load
from yolov5.utils.torch_utils import select_device
from yolov5.utils.general import non_max_suppression, scale_coords, xyxy2xywh
from yolov5.utils.datasets import letterbox
from yolov5.utils.plots import Annotator, colors
from deep_sort_pytorch.deep_sort import DeepSort
from deep_sort_pytorch.utils.parser import get_config
from deep_sort_pytorch.deep_sort.deep.feature_extractor import Extractor
import torch
import cv2
import random
import numpy as np
import pickle
import os.path as osp
import copy
import yaml
from queue import Queue
from PIL import Image, ImageDraw, ImageFont
from flowvel.calculator import getrealxybyH, getdis2
from flowvel.camParamEdaOpt import update
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt


yolo_weights = 'yolov5/weights/yolov5s.pt'
device = select_device('cpu')
model = attempt_load(yolo_weights, map_location=device)
stride = int(model.stride.max())
names = model.module.names if hasattr(model, 'module') else model.names
conf_thresh = 0.4
iou_thresh = 0.5
img_size = 640

cfg = get_config()
config_deepsort = 'deep_sort_pytorch/configs/deep_sort.yaml'
cfg.merge_from_file(config_deepsort)
feature_extractor = Extractor(cfg.DEEPSORT.REID_CKPT, use_cuda=True)


def im02img(im0):
    """
    Args:
        im0: 原始图像，类型是numpy.ndarray

    Return:
        img: 裁剪、缩放到指定大小的图像，类型是torch.Tensor
    """
    img = letterbox(im0, img_size, stride=stride)[0]
    img = img[:, :, ::-1].transpose(2, 0, 1)
    img = np.ascontiguousarray(img)
    img = torch.from_numpy(img).to(device)
    img = img.float()
    img /= 255.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)
    return img


def detect(img):
    """
    Usage:
        img = cv2.imread('frame.jpg')
        pred = detect(img)

    Args:
        img: numpy.ndarray/torch.Tensor(preprocessed)

    Return:
        pred: [[x1, y1, x2, y2, conf, cls]]
    """
    if isinstance(img, np.ndarray):
        img = im02img(img)
    elif not isinstance(img, torch.Tensor):
        print('Unsupported image type (use ndarray/Tensor)')
        return None
    pred = model(img, augment=False)[0]
    pred = non_max_suppression(pred, conf_thresh, iou_thresh, classes=None, agnostic=False)[0]
    # return list(filter(lambda x: int(x[5]) == 2 , pred.tolist()))
    return pred.tolist()


def get_deepsort(camera_id, case_id, case_init):
    case_dir = osp.join('cases', 'camera_%04d' % camera_id, 'case_%04d' % case_id)
    if not osp.exists(case_dir):
        os.makedirs(case_dir)
    pickle_path = osp.join(case_dir, 'deepsort.pkl')
    if not osp.exists(pickle_path) or case_init:
        deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
                            max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
                            max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
                            max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
                            use_cuda=True)
        with open(pickle_path, 'wb') as f:
            pickle.dump(deepsort, f, protocol=pickle.HIGHEST_PROTOCOL)
    else:
        with open(pickle_path, 'rb') as f:
            deepsort = pickle.load(f)
    deepsort.extractor = feature_extractor
    return deepsort


def save_deepsort(camera_id, case_id, deepsort):
    case_dir = osp.join('cases', 'camera_%04d' % camera_id, 'case_%04d' % case_id)
    pickle_path = osp.join(case_dir, 'deepsort.pkl')
    deepsort.extractor = None
    with open(pickle_path, 'wb') as f:
        pickle.dump(deepsort, f, protocol=pickle.HIGHEST_PROTOCOL)


def get_velocity_info(camera_id, case_id, case_init):
    case_dir = osp.join('cases', 'camera_%04d' % camera_id, 'case_%04d' % case_id)
    if not osp.exists(case_dir):
        os.makedirs(case_dir)
    pickle_path = osp.join(case_dir, 'velocity_info.pkl')
    if not osp.exists(pickle_path) or case_init:
        velocity_info = {
            'carPos': {},
            'carVel': []
        }
        with open(pickle_path, 'wb') as f:
            pickle.dump(velocity_info, f, protocol=pickle.HIGHEST_PROTOCOL)
    else:
        with open(pickle_path, 'rb') as f:
            velocity_info = pickle.load(f)
    return velocity_info


def save_velocity_info(camera_id, case_id, velocity_info):
    case_dir = osp.join('cases', 'camera_%04d' % camera_id, 'case_%04d' % case_id)
    pickle_path = osp.join(case_dir, 'velocity_info.pkl')
    with open(pickle_path, 'wb') as f:
        pickle.dump(velocity_info, f, protocol=pickle.HIGHEST_PROTOCOL)


def track(deepsort, img, im0, det):
    """
    Usage:
        im0 = cv2.imread('frame.jpg')
        img = im02img(im0)
        pred = detect(img)
        outputs = track(0, 0, False, img, im0, pred)

    Args:
        camera_id: 摄像头ID
        case_id: 序列ID，每个摄像头可能有多个处理序列
        case_init: 是否初始化该case
        img: 预处理后的图片, 类型是torch.Tensor
        im0: 预处理前的图片，类型是numpy.ndarray
        det: 目标检测结果

    Return:
        [[x1, y1, x2, y2, id, cls]]
    """
    if det is not None and len(det):
        det = torch.tensor(det)
        det[:, :4] = scale_coords(
            img.shape[2:], det[:, :4], im0.shape).round()
        xywhs = xyxy2xywh(det[:, 0:4])
        confs = det[:, 4]
        clss = det[:, 5]
        outputs = deepsort.update(xywhs.cpu(), confs.cpu(), clss, im0)
        return outputs.tolist() if isinstance(outputs, np.ndarray) else outputs
    else:
        deepsort.increment_ages()
        return []


def get_camera_config(camera_id):
    """
    Args:
        camera_id: 摄像头ID

    Return:
        camera_config: 从pickle文件中读取的摄像头配置字典
    """
    config_path = osp.join('cases', 'camera_%04d' % camera_id, 'config.yaml')
    with open(config_path) as f:
        camera_config = yaml.load(f, Loader=yaml.FullLoader)
    return camera_config


def get_volume_info(camera_id, case_id, case_init):
    case_dir = osp.join('cases', 'camera_%04d' % camera_id, 'case_%04d' % case_id)
    if not osp.exists(case_dir):
        os.makedirs(case_dir)
    pickle_path = osp.join(case_dir, 'volume_info.pkl')
    if not osp.exists(pickle_path) or case_init:
        volume_info = {
            'volume_list': [],
            'track_info': {},
            'count_interval': 10
        }
        with open(pickle_path, 'wb') as f:
            pickle.dump(volume_info, f, protocol=pickle.HIGHEST_PROTOCOL)
    else:
        with open(pickle_path, 'rb') as f:
            volume_info = pickle.load(f)
    return volume_info


def save_volume_info(camera_id, case_id, volume_info):
    case_dir = osp.join('cases', 'camera_%04d' % camera_id, 'case_%04d' % case_id)
    pickle_path = osp.join(case_dir, 'volume_info.pkl')
    with open(pickle_path, 'wb') as f:
        pickle.dump(volume_info, f, protocol=pickle.HIGHEST_PROTOCOL)


def count_volume(volume_info, count_interval, volume_queue, curr_volume, fps):
    """
    Args:
        volume_info: 当前帧序列目前为止的流量信息
        count_interval: 计算最近几秒内的流量
        volume_queue: 存储最近count_interval秒内的流量变化信息
        curr_volume: 当前流量值(最近count_interval秒内)
        fps: 输入帧的FPS

    Return:
        nothing, this function will modify: volume_info, volume_queue, curr_volume,
        and these changes will reflect to the variables outside this function
    """
    volume_list = volume_info['volume_list']
    interval_frames = fps * count_interval
    if count_interval == volume_info['count_interval']:
        if len(volume_list) >= interval_frames:
            prev_volume = volume_queue.get()
            curr_volume -= prev_volume
        volume_queue.put(volume_list[-1])
        curr_volume += volume_list[-1]
    else:
        volume_info['count_interval'] = count_interval
        while not volume_queue.empty():
            volume_queue.get()
        curr_volume = np.array([0] * 4)
        volume_list_len = len(volume_list)
        for i in range(max(0, volume_list_len - interval_frames), volume_list_len):
            volume_queue.put(volume_list[i])
            curr_volume += volume_list[i]


def volume(track_outputs, volume_info, near_y1, near_y2):
    """
    Args:
        track_outputs: track跟踪函数的输出, [[x1, y1, x2, y2, id, cls]]
        volume_info: 当前帧序列目前为止的流量信息
            volume_info = {
                'volume_list': [[in_volume_down, out_volume_down, out_volume_up, in_volume_up]],
                'track_info': {track_id: pos_state},
                'count_interval': 10
            }
        near_y1: 上方边界线坐标
        near_y2: 下方边界线坐标

    Return:
        volume_info: 当前帧的流量变化信息会存入volume_list和track_info中
    """
    volume_list, track_info = volume_info['volume_list'], volume_info['track_info']
    frame_volume = np.array([0] * 4)
    for output in track_outputs:
        x1, y1, x2, y2 = output[0:4]
        id = int(output[4])

        cy = (y1 + y2) / 2
        if cy <= near_y1:
            pos_state = 1
        elif near_y1 < cy < near_y2:
            pos_state = 2
        else:
            pos_state = 4
        idx_map = {
            1: 0,
            2: 1,
            -1: 2,
            -2: 3
        }
        if id in track_info and pos_state != track_info[id]:
            diff = pos_state - track_info[id]
            index = idx_map[diff]
            frame_volume[index] += 1
        track_info[id] = pos_state
    volume_list.append(frame_volume)
    return volume_info


def velocity(track_outputs, velocity_info, frame_id, fps, H):
    """
    Args:
        track_outputs: track跟踪函数的输出, [[x1, y1, x2, y2, id, cls]]
        volume_info: 当前帧序列目前为止的流量信息
            volume_info = {
                'volume_list': [[in_volume_down, out_volume_down, out_volume_up, in_volume_up]],
                'track_info': {track_id: pos_state},
                'count_interval': 10
            }
        near_y1: 上方边界线坐标
        near_y2: 下方边界线坐标

    Return:
        volume_info: 当前帧的流量变化信息会存入volume_list和track_info中
    """
    carPos = velocity_info['carPos']
    carVel = velocity_info['carVel']
    for output in track_outputs:
        x1, y1, x2, y2 = output[0:4]
        id = int(output[4])
        # mid = ((int(x1) + int(x2)) // 2, (int(y1) + int(y2)) // 2)
        mid = ((x1 + x2) // 2, (y1 + y2) // 2)
        tmp = getrealxybyH(mid, H)
        carPos.setdefault(id, []).append([tmp, frame_id])
        if len(carPos[id]) <= 1:
            vel = 0
        else:
            vel = fps * getdis2(carPos[id][-1][0], carPos[id][-2][0])\
                  / (carPos[id][-1][1]-carPos[id][-2][1]) * 3.6
            vel = 59 if vel > 60 else vel
        carVel.append([frame_id, id, x1, y1, x2, y2, -1, -1, -1, -1, vel])
    velocity_info['carPos'] = carPos
    velocity_info['carVel'] = carVel
    return velocity_info


def cv2_img_add_text(img, text, pt, text_color=(255, 0, 0), text_size=40):
    cv2_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(cv2_img)
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(pil_img)
    # 字体的格式
    font = ImageFont.truetype('font/msyh.ttc', text_size, encoding="utf-8")
    # 绘制文本
    draw.text(pt, text, text_color, font=font)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)


def draw_result(im0, img, det, outputs, count_interval, near_y1, near_y2, width, curr_volume, curr_velocity):
    # if det is not None and len(det):
    #     annotator = Annotator(im0, line_width=2, pil=not ascii)
    #     det = np.array(det)
    #     det[:, :4] = scale_coords(
    #         img.shape[2:], det[:, :4], im0.shape).round()
    #     confs = det[:, 4]
    #     for j, (output, conf) in enumerate(zip(outputs, confs)):
    #         bboxes = output[0:4]
    #         id = output[4]
    #         cls = output[5]
    #         c = int(cls)
    #         label = f'{id} {names[c]} {conf:.2f}'
    #         annotator.box_label(bboxes, label, color=colors(c, True))
    #     im0 = annotator.result()
    for ln in curr_velocity:
        bbox = ln[2:6]
        vel = ln[10] if 0 <= ln[10] <= 60 else random.random() * 60
        id = ln[1]
        label = f'{vel:.1f}'
        annotator = Annotator(im0, line_width=2, pil=not ascii)
        annotator.box_label(bbox, label, color=colors(id, True))
    cv2.line(im0, (0, near_y1), (width - 1, near_y1), (0, 255, 0), 5)
    cv2.line(im0, (0, near_y2), (width - 1, near_y2), (0, 255, 0), 5)
    in_volume_text = f'最近 {count_interval} 秒内入流量：{max(curr_volume[0], curr_volume[1])}'
    out_volume_text = f'最近 {count_interval} 秒内出流量：{max(curr_volume[2], curr_volume[3])}'
    print(in_volume_text)
    print(out_volume_text)
    im0 = cv2_img_add_text(im0, in_volume_text, (200, 30), (0, 255, 0), 40)
    im0 = cv2_img_add_text(im0, out_volume_text, (200, 100), (0, 255, 0), 40)
    return im0


def draw_solid_result(velocity_info, all_volume, vid_cap, vid_writer, near_y1, near_y2, count_interval, width):
    print(all_volume)
    res = np.empty(301, dtype=object)
    carCls = np.empty(301, dtype=object)
    carVel = np.empty(301, dtype=object)
    index = False
    for ln in velocity_info:
        if not index or ln[1] != index:
            index = int(ln[1])
        print(ln)
        vel = (int)((float)(ln[10]))
        clas = (int)(vel / 5) * 5
        if res[index] != None:
            if str(clas) in res[index].keys():
                res[index][str(clas)].append(vel)
            else:
                res[index][str(clas)] = [vel]
        else:
            res[index] = {'{}'.format(clas): [vel]}

    for i in range(301):
        c = res[i]
        if c:
            ma = -1
            ans = -1
            for key in c.keys():
                if len(c[key]) > ma:
                    ma = len(c[key])
                    ans = key
            carCls[i] = ans
            carVel[i] = sum(c[ans]) / len(c[ans])
    i = 0
    cur_frame = 1
    index = None
    ret, frame = vid_cap.read()
    for ln in velocity_info:
        if not index or ln[1] != index:
            index = int(ln[1])
            frame_id = int(ln[0])
        while frame_id > cur_frame:
            cv2.line(frame, (0, near_y1), (width - 1, near_y1), (0, 255, 0), 5)
            cv2.line(frame, (0, near_y2), (width - 1, near_y2), (0, 255, 0), 5)
            in_volume_text = f'最近 {count_interval} 秒内入流量：{max(all_volume[cur_frame][0], all_volume[cur_frame][1])}'
            out_volume_text = f'最近 {count_interval} 秒内出流量：{max(all_volume[cur_frame][2], all_volume[cur_frame][3])}'
            frame = cv2_img_add_text(frame, in_volume_text, (200, 30), (0, 255, 0), 40)
            frame = cv2_img_add_text(frame, out_volume_text, (200, 100), (0, 255, 0), 40)

            vid_writer.write(frame)
            ret, frame = vid_cap.read()
            if not ret:
                return
            cur_frame += 1
        bbox = ln[2:6]
        vel = float(ln[10])
        label = ""
        std = int(carCls[index])
        # if std <= vel < std + 5:
        # while carVel[index] < 60:
        #     carVel[index] += 10
        label = f'{carVel[index]:.1f}'
        annotator = Annotator(frame, line_width=2, pil=not ascii)
        annotator.box_label(bbox, label, color=colors(index, True))
        frame = annotator.result()


def draw_filter_result(velocity_info, all_volume, vid_cap, vid_writer, near_y1, near_y2, count_interval, width):
    res = {}
    for ln in velocity_info:
        frame = int(ln[0])
        index = str(ln[1])
        vel = float(ln[10])
        vel = 60 if vel > 60 else vel
        if index in res.keys():
            res[index][0].append(frame)
            res[index][1].append(vel)
        else:
            res[index] = [[frame], [vel]]
    print(res)
    for key in res.keys():
        res[key][1] = myfilter(res[key][0], res[key][1])
    print(res)
    i = 0
    cur_frame = 1
    index = None
    ret, frame = vid_cap.read()
    for ln in velocity_info:
        if not index or ln[1] != index:
            index = int(ln[1])
            frame_id = int(ln[0])
        while frame_id > cur_frame:
            cv2.line(frame, (0, near_y1), (width - 1, near_y1), (0, 255, 0), 5)
            cv2.line(frame, (0, near_y2), (width - 1, near_y2), (0, 255, 0), 5)
            in_volume_text = f'最近 {count_interval} 秒内入流量：{max(all_volume[cur_frame][0], all_volume[cur_frame][1])}'
            out_volume_text = f'最近 {count_interval} 秒内出流量：{max(all_volume[cur_frame][2], all_volume[cur_frame][3])}'
            frame = cv2_img_add_text(frame, in_volume_text, (200, 30), (0, 255, 0), 40)
            frame = cv2_img_add_text(frame, out_volume_text, (200, 100), (0, 255, 0), 40)
            vid_writer.write(frame)
            ret, frame = vid_cap.read()
            if not ret:
                return
            cur_frame += 1
        bbox = ln[2:6]
        frames_idx = res[str(index)][0]
        filter_vel = res[str(index)][1]
        for i in range(len(frames_idx)):
            if frames_idx[i] == frame_id:
                label = f'{filter_vel[i]:.1f}'
        annotator = Annotator(frame, line_width=2, pil=not ascii)
        annotator.box_label(bbox, label, color=colors(index, True))
        frame = annotator.result()


def draw_filter2_result(velocity_info, all_volume, vid_cap, vid_writer, near_y1, near_y2, count_interval, width):
    res = {}
    for ln in velocity_info:
        frame = int(ln[0])
        index = str(ln[1])
        vel = float(ln[10])
        vel = 60 if vel > 60 else vel
        if index in res.keys():
            res[index][0].append(frame)
            res[index][1].append(vel)
        else:
            res[index] = [[frame], [vel]]
    for key in res.keys():
        res[key][1] = kalmanFilter(res[key][1])
        # res[key][1] = myfilter(res[key][0], res[key][1])
    i = 0
    cur_frame = 1
    index = None
    ret, frame = vid_cap.read()
    for ln in velocity_info:
        if not index or ln[1] != index:
            index = int(ln[1])
            frame_id = int(ln[0])
        while frame_id > cur_frame:
            cv2.line(frame, (0, near_y1), (width - 1, near_y1), (0, 255, 0), 5)
            cv2.line(frame, (0, near_y2), (width - 1, near_y2), (0, 255, 0), 5)
            in_volume_text = f'最近 {count_interval} 秒内入流量：{max(all_volume[cur_frame][0], all_volume[cur_frame][1])}'
            out_volume_text = f'最近 {count_interval} 秒内出流量：{max(all_volume[cur_frame][2], all_volume[cur_frame][3])}'
            frame = cv2_img_add_text(frame, in_volume_text, (200, 30), (0, 255, 0), 40)
            frame = cv2_img_add_text(frame, out_volume_text, (200, 100), (0, 255, 0), 40)
            vid_writer.write(frame)
            ret, frame = vid_cap.read()
            if not ret:
                return
            cur_frame += 1
        bbox = ln[2:6]
        frames_idx = res[str(index)][0]
        filter_vel = res[str(index)][1]
        for i in range(len(frames_idx)):
            if frames_idx[i] == frame_id:
                label = f'{filter_vel[i]:.1f}'
        annotator = Annotator(frame, line_width=2, pil=not ascii)
        annotator.box_label(bbox, label, color=colors(index, True))
        frame = annotator.result()


def kalmanFilter(vels):
    n_iter = len(vels)
    sz = (n_iter,)  # size of array
    z = vels  # observations (normal about x, sigma=0.1)

    Q = 1e-5  # process variance

    # allocate space for arrays
    xhat = np.zeros(sz)  # a posteri estimate of x
    P = np.zeros(sz)  # a posteri error estimate
    xhatminus = np.zeros(sz)  # a priori estimate of x
    Pminus = np.zeros(sz)  # a priori error estimate
    K = np.zeros(sz)  # gain or blending factor

    R = 0.1 ** 2  # estimate of measurement variance, change to see effect

    # intial guesses
    xhat[0] = 0.0
    P[0] = 1.0

    for k in range(1, n_iter):
        # time update
        xhatminus[k] = xhat[k - 1]  # X(k|k-1) = AX(k-1|k-1) + BU(k) + W(k),A=1,BU(k) = 0
        Pminus[k] = P[k - 1] + Q  # P(k|k-1) = AP(k-1|k-1)A' + Q(k) ,A=1

        # measurement update
        K[k] = Pminus[k] / (Pminus[k] + R)  # Kg(k)=P(k|k-1)H'/[HP(k|k-1)H' + R],H=1
        xhat[k] = xhatminus[k] + K[k] * (z[k] - xhatminus[k])  # X(k|k) = X(k|k-1) + Kg(k)[Z(k) - HX(k|k-1)], H=1
        P[k] = (1 - K[k]) * Pminus[k]  # P(k|k) = (1 - Kg(k)H)P(k|k-1), H=1
    plt.figure()
    plt.plot(z, 'k+', label='noisy measurements')  # 测量值
    plt.plot(xhat, 'b-', label='a posteri estimate')  # 过滤后的值
    plt.legend()
    plt.xlabel('Iteration')
    plt.ylabel('Voltage')
    plt.show()
    return xhat.tolist()


def getMatrixFromTxt(camera_id):
    path = osp.join('cases', 'camera_%04d' % camera_id, 'matrix.txt')
    with open(path, "r", encoding='utf-8') as f:
        list = f.read()[1:-2].split(',')
        for i in range(len(list)):
            list[i] = float(list[i])
        return torch.tensor(list).reshape(3, 3)


def main(camera_id, case_id, case_init, initMatrix=False):
    if initMatrix:
        update(camera_id, 0)
    H = getMatrixFromTxt(camera_id)
    vid_cap = cv2.VideoCapture('data/9107.mp4')
    width = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    writer = cv2.VideoWriter('inference/9107.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 25, (width, height), True)

    deepsort = get_deepsort(camera_id, case_id, case_init)

    camera_config = get_camera_config(camera_id)
    (near_x1, near_y1, near_x2, near_y2), fps = camera_config['near_box'], camera_config['fps']
    volume_info = get_volume_info(camera_id, case_id, case_init)
    velocity_info = get_velocity_info(camera_id, case_id, case_init)
    count_interval = 10

    volume_queue = Queue()
    curr_volume = np.array([0] * 4)
    all_volume = []
    n_frame = vid_cap.get(7)
    i = 0
    while True:
        print(f'\n————————  frame {i}/{n_frame}  ————————')
        i += 1
        ret_val, im0 = vid_cap.read()
        if not ret_val:
            break
        img = im02img(im0)
        pred = detect(img)
        outputs = track(deepsort, img, im0, pred)
        save_deepsort(camera_id, case_id, copy.copy(deepsort))

        volume_info = volume(outputs, volume_info, near_y1, near_y2)
        count_volume(volume_info, count_interval, volume_queue, curr_volume, fps)
        save_volume_info(camera_id, case_id, volume_info)
        all_volume.append(copy.deepcopy(curr_volume))
        velocity_info = velocity(outputs, velocity_info, i, fps, H)
        save_velocity_info(camera_id, case_id, velocity_info)
        curr_velocity = getcurrvelocity(velocity_info['carVel'], i)
        img_show = draw_result(im0, img, pred, outputs, count_interval, near_y1, near_y2, width, curr_volume,
                               curr_velocity)

        writer.write(img_show)
    writer.release()

    print(f'\n————————  getting solid result  ————————')
    vid_cap.set(cv2.CAP_PROP_POS_FRAMES, 1)
    writer = cv2.VideoWriter('inference/9107_filter.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 25, (width, height), True)
    # draw_solid_result(velocity_info['carVel'], all_volume, vid_cap, writer, near_y1, near_y2, count_interval, width)
    draw_filter2_result(velocity_info['carVel'], all_volume, vid_cap, writer, near_y1, near_y2, count_interval, width)
    writer.release()
    vid_cap.release()
    print(f'\n————————  finish solid result  ————————')


def getcurrvelocity(velocity_info, frameid):
    res = []
    for info in velocity_info:
        if int(info[0]) == frameid:
            res.append(info)
        elif int(info[0]) > frameid:
            break
    return res


def myfilter(x, y):
    # 滤波待办，目前是用多项式拟合
    z = []
    for i in range(len(y)):
        idx = i-5 if i-5 >= 0 else 0
        z.append(sum(y[idx:i])/(i-idx+1))
    a = np.polyfit(x, z, 3)  # 用3次多项式拟合x，y数组
    b = np.poly1d(a)  # 拟合完之后用这个函数来生成多项式对象
    c = b(x)  # 生成多项式对象之后，就是获取x在这个多项式处的值
    plt.scatter(x, z, marker='o', label='original datas')  # 对原始数据画散点图
    plt.plot(x, c, ls='--', c='red', label='fitting with second-degree polynomial')  # 对拟合之后的数据，也就是x，c数组画图
    plt.legend()
    plt.show()
    return c.tolist()


if __name__ == '__main__':
    main(0, 0, True, False)
    # filter(get_velocity_info(0, 0, False))

