import sys

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

from detection import detect, im02img, get_detection_log
from tracking import get_camera_config, track, getMatrixFromTxt
from volume import get_deepsort, get_volume_info, save_deepsort, volume, count_volume, save_volume_info
from velocity import get_velocity_info, save_velocity_info, velocity, getcurrvelocity, draw_result, draw_solid_result, draw_filter2_result
from congestion import get_congestion_info, congestion, save_congestion_info
from flowvel.camParamEdaOpt import update

import time
import csv
import requests
import socket
import argparse
from threading import Timer
from queue import Queue
import cv2
import json
import numpy as np
import copy
import base64


def print_pretty_dict(my_dict):
    pretty_dict = json.dumps(my_dict, indent=2, ensure_ascii=False)
    print(pretty_dict)


def get_model_list():
    return {
        'source': {
            'id': 0,
            'name': 'UA-DETRAC'
        },
        'detection': {
            'id': 0,
            'name': 'yolov5s',
            'params': [
                {
                    'id': 0,
                    'name': 'yolov5/models/yolov5s.yaml'
                }
            ]
        },
        'tracking': {
            'id': 0,
            'name': 'DeepSort',
            'params': [
                {
                    'id': 0,
                    'name': 'deep_sort_pytorch/configs/deep_sort.yaml'
                }
            ]
        }
    }


def read_loss_mAP(train_result):
    csv_reader = csv.reader(open(train_result, 'r'))
    next(csv_reader)
    box_loss_list, mAP50_list = [], []
    for line in csv_reader:
        box_loss_list.append(float(line[1]))
        mAP50_list.append(float(line[6]))
    return box_loss_list, mAP50_list


train_result = 'docs/yolov5s_detrac_300.csv'
box_loss_list, mAP50_list = read_loss_mAP(train_result)


def get_execute_video_info(camera_id, case_id, camera_config, congestion_info, pred):
    start_time = time.time()
    congestion_box, congestion_level, origin_fps = camera_config['congestion_box'], camera_config['congestion_level'], \
                                                   camera_config['origin_fps']
    sample_interval = camera_config['sample_interval']
    target_fps = origin_fps / sample_interval
    prev_level, duration = congestion_info['prev_level'], congestion_info['duration']

    num_in_area, area_level = congestion(pred, congestion_box, congestion_level)
    if area_level == prev_level:
        duration += 1
    else:
        duration = 1
    prev_level = area_level
    congestion_info['prev_level'] = prev_level
    congestion_info['duration'] = duration
    save_congestion_info(camera_id, case_id, congestion_info)
    end_time = time.time()
    return {
        'algorithm': {
            'name': 'yolov5s+DeepSort',
            'edition': '1.0',
            'author': 'BUAA'
        },
        'warning': {
            'time': time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
            'level': area_level,
            'duration': duration,
            'content': f'当前拥堵等级 {area_level}，已持续时间 {duration / target_fps} 秒'
        },
        'jam': {
            'level': area_level,
            'time': duration
        },
        'model': {
            'accuracy': mAP50_list,
            'loss': box_loss_list
        },
        'log': {
            'start_time': time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(start_time)),
            'duration': end_time - start_time
        }
    }


live_info = {
    9107: 'http://117.141.132.122:10590/live/Y4vq46I7R',
    9355: 'http://117.141.132.122:10736/live/yT8Um~WnR'
}


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--ip', type=str, required=True,
                        help='后端监听IP')
    parser.add_argument('--port', type=int, required=True,
                        help='后端监听端口')
    parser.add_argument('--camera', type=int, required=True,
                        help='摄像头ID')
    parser.add_argument('--case', type=int, default=0,
                        help='Case ID，可选，默认是0')
    parser.add_argument('--init', type=bool, default=True,
                        help='是否初始化Case，可选，默认是True')
    parser.add_argument('--model', type=str, required=True,
                        help='模型ID')
    parser.add_argument('--matrix', type=str, default=False,
                        help='是否重新计算单应性矩阵H')
    return parser.parse_args()


ts_queue = Queue()
request_interval = 30

result_queue = Queue()
send_interval = 1 / 12.5


def update_ts_list(link):
    """每隔30s获取5个TS视频放入队列"""
    res = requests.get(link)
    if res.text:
        for line in res.text.split('\n'):
            if line.endswith('.ts'):
                ts_queue.put(line)
    timer = Timer(request_interval, update_ts_list, args=[link])
    timer.start()


def send_result(sock):
    if not result_queue.empty():
        result_queue.get()
        # 后端监听就绪后取消注释
        # sock.send(result_queue.get())
    timer = Timer(send_interval, send_result, args=[sock])
    timer.start()


def main():
    # 获取命令行参数
    args = get_args()
    # 取出命令行参数
    sock_ip, sock_port, camera_id, case_id, case_init, model_id, matrix_init = args.ip, args.port, args.camera, args.case,\
                                                                               args.init, args.model, args.matrix
    if matrix_init:
        update(camera_id, case_id)
    H = getMatrixFromTxt(camera_id)
    # 从本地config.yaml文件读取摄像头配置
    camera_config = get_camera_config(camera_id)
    # 原始视频FPS
    origin_fps = camera_config['origin_fps']
    # 每隔几帧采样一次
    sample_interval = camera_config['sample_interval']
    # 原视频FPS转化为采样后的FPS
    target_fps = origin_fps / sample_interval
    # 计算每帧之间的时间间隔
    global send_interval
    send_interval = 1 / target_fps

    congestion_info = get_congestion_info(camera_id, case_id, case_init)

    deepsort = get_deepsort(camera_id, case_id, case_init)
    near_x1, near_y1, near_x2, near_y2 = camera_config['near_box']
    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 = []

    # socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # 后端监听就绪后取消注释
    # sock.connect((sock_ip, sock_port))

    # 摄像头直播链接前缀
    link_prefix = live_info[camera_id]
    # m3u8文件完整HTTP链接
    m3u8_link = f'{link_prefix}/index.m3u8'
    # 开始TS视频的获取
    update_ts_list(m3u8_link)

    # 开始每隔固定时间发送每帧的处理结果
    timer = Timer(send_interval, send_result, args=[sock])
    timer.start()

    ts_idx = 0
    frame_idx = 0
    while True:
        # 如果目前所有视频都处理完了，自旋一段时间
        while ts_queue.empty():
            pass

        ts_start = time.time()

        print(f'TS {ts_idx} Start')
        ts_filename = ts_queue.get()
        vid_cap = cv2.VideoCapture(f'{link_prefix}/{ts_filename}')
        width = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = vid_cap.get(5)
        n_frame = vid_cap.get(7)

        while True:
            frame_start = time.time()

            t1 = time.time()

            ret_val, im0 = vid_cap.read()

            t2 = time.time()

            if not ret_val:
                break

            if frame_idx % sample_interval != 0:
                frame_idx += 1
                continue

            img = im02img(im0)
            pred = detect(img)

            log_str = '%gx%g ' % img.shape[2:]
            log_str += get_detection_log(pred, img, im0)

            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, target_fps)
            save_volume_info(camera_id, case_id, volume_info)
            all_volume.append(copy.deepcopy(curr_volume))

            velocity_info = velocity(outputs, velocity_info, frame_idx, fps, H)
            save_velocity_info(camera_id, case_id, velocity_info)
            curr_velocity = getcurrvelocity(velocity_info['carVel'], frame_idx)
            img_show = draw_result(im0, img, pred, outputs, count_interval, near_y1, near_y2, width, curr_volume,
                                   curr_velocity)
            # img_show = draw_result(im0, img, pred, outputs, count_interval, near_y1, near_y2, width, curr_volume)
            success, img_str = cv2.imencode('.jpg', img_show)
            img_bytes = img_str.tobytes()
            b64_code = str(base64.b64encode(img_bytes))

            execute_video_info = get_execute_video_info(camera_id, case_id, camera_config, congestion_info, pred)

            t3 = time.time()

            execute_video_info['model_id'] = model_id
            execute_video_info['result_img'] = b64_code

            t4 = time.time()

            time_str = f'read frame: {t2 - t1}'
            time_str += f'video info: {t3 - t2}'
            time_str += f'result postprocess: {t4 - t3}'
            # print(time_str)

            frame_end = time.time()
            curr_in_volume = max(curr_volume[0], curr_volume[1])
            curr_out_volume = max(curr_volume[2], curr_volume[3])
            log_str += ' Volume: %d IN, %d OUT. (%.3fs)' % (curr_in_volume, curr_out_volume, frame_end - frame_start)
            for ln in curr_velocity:
                vel = ln[10] if 0 <= ln[10] <= 60 else 6
                id = ln[1]
                log_str += '\n Velocity: frame%d car%d %.3fs' % (frame_idx, id, vel)
            print(log_str)

            execute_video_info['log_str'] = log_str
            json_str = json.dumps(execute_video_info) + '\n'
            result_data = json_str.encode()
            result_queue.put(result_data)

            frame_idx += 1

        ts_end = time.time()
        print(f'TS {ts_idx} COST TIME: {ts_end - ts_start}')

        print(f'\n————————  getting filter result  ————————')

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

        ts_idx += 1


if __name__ == '__main__':
    main()
