from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths
import os
import os.path as osp
import cv2
import logging
import argparse
import motmetrics as mm
import numpy as np
import torch
from lib.tracker.tracker import JDETracker
# from lib.tracker.multitracker import JDETracker
# from lib.tracker.mot20tracker import JDETracker
# from lib.tracker.tracker_test import JDETracker
from lib.tracker.GSI import GSInterpolation
from lib.tracker.AFL import AFLink
from lib.tracker.byte_tracker import BYTETracker
from lib.tracking_utils import visualization as vis
from lib.tracking_utils.log import logger
from lib.tracking_utils.timer import Timer
from lib.tracking_utils.evaluation import Evaluator
from lib.datasets.dataset import jde as datasets

from lib.tracking_utils.utils import mkdir_if_missing
import lib.tracker.det_feat_record as det_feat_record
import lib.tracker.refined_track_vis as refined_track_vis
from opts import opts


def write_results(filename, results, data_type):
    if data_type == 'mot':
        save_format = '{frame},{id},{x1},{y1},{w},{h},1,-1,-1,-1\n'
    elif data_type == 'kitti':
        save_format = '{frame} {id} pedestrian 0 0 -10 {x1} {y1} {x2} {y2} -10 -10 -10 -1000 -1000 -1000 -10\n'
    else:
        raise ValueError(data_type)

    with open(filename, 'w') as f:
        for frame_id, tlwhs, track_ids in results:
            if data_type == 'kitti':
                frame_id -= 1
            for tlwh, track_id in zip(tlwhs, track_ids):
                if track_id < 0:
                    continue
                x1, y1, w, h = tlwh
                x2, y2 = x1 + w, y1 + h
                line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h)
                f.write(line)
    logger.info('save results to {}'.format(filename))


def eval_seq(opt, dataloader, data_type, result_filename,seq, save_dir=None, show_image=True, frame_rate=30):
    if save_dir:
        mkdir_if_missing(save_dir)
    tracker = JDETracker(opt, frame_rate=frame_rate)
    use_mat = {'MOT17-02-SDP': False, 'MOT17-04-SDP': False, "MOT17-05-SDP": True, "MOT17-09-SDP": False,
               "MOT17-10-SDP": True, "MOT17-11-SDP": True, "MOT17-13-SDP": True,
               'MOT17-01-SDP': False, 'MOT17-03-SDP': False, "MOT17-06-SDP": True, "MOT17-07-SDP": True,
               "MOT17-08-SDP": False, "MOT17-12-SDP": True, "MOT17-14-SDP": True,
               'MOT17-01-FRCNN': False, 'MOT17-03-FRCNN': False, "MOT17-06-FRCNN": True, "MOT17-07-FRCNN": True,
               "MOT17-08-FRCNN": False, "MOT17-12-FRCNN": True, "MOT17-14-FRCNN": True,
               'MOT17-01-DPM': False, 'MOT17-03-DPM': False, "MOT17-06-DPM": True, "MOT17-07-DPM": True,
               "MOT17-08-DPM": False, "MOT17-12-DPM": True, "MOT17-14-DPM": True,
               "MOT20-01": False, "MOT20-02": False, "MOT20-03": False,"MOT20-04": False,
               "MOT20-05": False, "MOT20-06": False, "MOT20-07": False, "MOT20-08": False,
               'MOT16-02': False, 'MOT16-04': False, "MOT16-05": True, "MOT16-09": False,
               "MOT16-10": True, "MOT16-11": True, "MOT16-13": True,
               'MOT16-01': False, 'MOT16-03': False, "MOT16-06": True, "MOT16-07": True,
               "MOT16-08": False, "MOT16-12": True, "MOT16-14": True,
               }
    tracker.use_mat = use_mat[seq]

    tracker.recorder =  det_feat_record.det_feat_recorder(seq,'/home/hust/yly/Dataset/MOT17/','get')

    timer = Timer()
    results = []
    len_all = len(dataloader)
    start_frame = len_all // 2
    frame_id = start_frame
    tracker.start_frame_id = start_frame
    for i, (path, img, img0) in enumerate(dataloader):
        if i < start_frame:
            continue
        if frame_id % 20 == 0:
            logger.info('Processing frame {} ({:.2f} fps)'.format(frame_id, 1. / max(1e-5, timer.average_time)))

        # run tracking
        timer.tic()
        blob = torch.from_numpy(img).cuda().unsqueeze(0)
        online_targets,insert_frame = tracker.update(blob, img0)
        online_tlwhs = []
        online_ids = []
        #online_scores = []
        for t in online_targets:
            tlwh = t[0]
            tid = t[1]
            vertical = tlwh[2] / tlwh[3] > 1.6
            if tlwh[2] * tlwh[3] > opt.min_box_area and not vertical:
                online_tlwhs.append(tlwh)
                online_ids.append(tid)
                #online_scores.append(t.score)
        timer.toc()
        # save results
        results.append((frame_id + 1, online_tlwhs, online_ids))
        for track_id in insert_frame.keys():
            frames = insert_frame[track_id]
            for frame in frames:
                tlwh,id = frame
                results[id-1][1].append(tlwh)
                results[id-1][2].append(track_id)
        #results.append((frame_id + 1, online_tlwhs, online_ids, online_scores))
        if show_image or save_dir is not None:
            online_im = vis.plot_tracking(img0, online_tlwhs, online_ids, frame_id=frame_id,
                                          fps=1. / timer.average_time)
        if show_image:
            cv2.imshow('online_im', online_im)
        if save_dir is not None:
            cv2.imwrite(os.path.join(save_dir, '{:05d}.jpg'.format(frame_id)), online_im)
        frame_id += 1
    # save results
    for result in results:
        for i,id in enumerate(result[2]):
            if id in tracker.id_map.keys():
                id = tracker.id_map[id]
            result[2][i] = id


    write_results(result_filename, results, data_type)

    afl = AFLink(result_filename,result_filename,0.4,(0,30))
    afl.feat = tracker.feat_record
    afl.link()
    GSInterpolation(path_in=result_filename,
                path_out=result_filename,
                interval=20,
                tau=10)

    #write_results_score(result_filename, results, data_type)
    return frame_id, timer.average_time, timer.calls


def main(opt, data_root='/data/MOT16/train', det_root=None, seqs=('MOT16-05',), exp_name='demo',
         save_images=False, save_videos=False, show_image=True):
    logger.setLevel(logging.INFO)
    result_root = os.path.join(data_root, '..', 'results', exp_name)
    mkdir_if_missing(result_root)
    data_type = 'mot'

    # run tracking
    accs = []
    n_frame = 0
    timer_avgs, timer_calls = [], []
    for seq in seqs:
        output_dir = os.path.join(data_root, '..', 'outputs', exp_name, seq) if save_images or save_videos else None
        logger.info('start seq: {}'.format(seq))
        dataloader = datasets.LoadImages(osp.join(data_root, seq, 'img1'), opt.img_size)
        result_filename = os.path.join(result_root, '{}.txt'.format(seq))
        meta_info = open(os.path.join(data_root, seq, 'seqinfo.ini')).read()
        frame_rate = int(meta_info[meta_info.find('frameRate') + 10:meta_info.find('\nseqLength')])
        nf, ta, tc= eval_seq(opt, dataloader, data_type, result_filename,seq,
                              save_dir=output_dir, show_image=show_image, frame_rate=frame_rate)
        n_frame += nf
        timer_avgs.append(ta)
        timer_calls.append(tc)

        # eval
        logger.info('Evaluate seq: {}'.format(seq))
        evaluator = Evaluator(data_root, seq, data_type)
        accs.append(evaluator.eval_file(result_filename))
        events = evaluator.acc.events
        events_csv = os.path.join(result_root,'{}.csv'.format(seq))
        events.to_csv(events_csv)
        print(events_csv)
        if save_videos:
            output_video_path = osp.join(output_dir, '{}.mp4'.format(seq))
            cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -c:v copy {}'.format(output_dir, output_video_path)
            os.system(cmd_str)
    timer_avgs = np.asarray(timer_avgs)
    timer_calls = np.asarray(timer_calls)
    all_time = np.dot(timer_avgs, timer_calls)
    avg_time = all_time / np.sum(timer_calls)
    logger.info('Time elapsed: {:.2f} seconds, FPS: {:.2f}'.format(all_time, 1.0 / avg_time))

    # get summary
    metrics = mm.metrics.motchallenge_metrics
    mh = mm.metrics.create()
    summary = Evaluator.get_summary(accs, seqs, metrics)
    strsummary = mm.io.render_summary(
        summary,
        formatters=mh.formatters,
        namemap=mm.io.motchallenge_metric_names
    )
    print(strsummary)
    import json
    Evaluator.save_summary(summary, os.path.join(result_root, 'summary_{}.xlsx'.format(exp_name)))


if __name__ == '__main__':
    torch.cuda.set_device(0)
    opt = opts().init()
    print(opt)

    if  opt.val_mot16:
        seqs_str = '''MOT16-02
                      MOT16-04
                      MOT16-05
                      MOT16-09
                      MOT16-10
                      MOT16-11
                      MOT16-13'''
        data_root = os.path.join('/home/hust/yly/Dataset', 'MOT16/train')
    if opt.test_mot16:
        seqs_str = '''MOT16-01
                      MOT16-03
                      MOT16-06
                      MOT16-07
                      MOT16-08
                      MOT16-12
                      MOT16-14'''
        data_root = os.path.join('/home/hust/yly/Dataset', 'MOT16/test')
    if opt.test_mot15:
        seqs_str = '''ADL-Rundle-1
                      ADL-Rundle-3
                      AVG-TownCentre
                      ETH-Crossing
                      ETH-Jelmoli
                      ETH-Linthescher
                      KITTI-16
                      KITTI-19
                      PETS09-S2L2
                      TUD-Crossing
                      Venice-1'''
        data_root = os.path.join(opt.data_dir, 'MOT15/images/test')
    if opt.test_mot17:
        seqs_str = '''MOT17-01-DPM
                      MOT17-03-DPM
                      MOT17-06-DPM
                      MOT17-07-DPM
                      MOT17-08-DPM
                      MOT17-12-DPM
                      MOT17-14-DPM
                      MOT17-01-FRCNN
                      MOT17-03-FRCNN
                      MOT17-06-FRCNN
                      MOT17-07-FRCNN
                      MOT17-08-FRCNN
                      MOT17-12-FRCNN
                      MOT17-14-FRCNN
                      MOT17-01-SDP
                      MOT17-03-SDP
                      MOT17-06-SDP
                      MOT17-07-SDP
                      MOT17-08-SDP
                      MOT17-12-SDP
                      MOT17-14-SDP'''
        data_root = os.path.join('/home/hust/yly/Dataset', 'MOT17/test')
    if opt.val_mot17:
        seqs_str = '''MOT17-02-SDP
                      MOT17-04-SDP
                      MOT17-05-SDP
                      MOT17-09-SDP
                      MOT17-10-SDP
                      MOT17-11-SDP
                      MOT17-13-SDP'''
        # seqs_str = '''
        #               MOT17-05-SDP
        #               MOT17-10-SDP
        #               MOT17-11-SDP
        #               MOT17-13-SDP
        #             '''
        data_root = os.path.join('/home/hust/yly/Dataset', 'MOT17/train')
    if opt.val_mot15:
        seqs_str = '''Venice-2
                      KITTI-13
                      KITTI-17
                      ETH-Bahnhof
                      ETH-Sunnyday
                      PETS09-S2L1
                      TUD-Campus
                      TUD-Stadtmitte
                      ADL-Rundle-6
                      ADL-Rundle-8
                      ETH-Pedcross2
                      TUD-Stadtmitte'''
        data_root = os.path.join(opt.data_dir, 'MOT15/images/train')
    if opt.val_mot20:
        seqs_str = '''MOT20-01
                      MOT20-02
                      MOT20-03
                      MOT20-05
                      '''
        data_root = os.path.join('/home/hust/yly/Dataset', 'MOT20/train')
    if opt.test_mot20:
        seqs_str = '''MOT20-04
                      MOT20-06
                      MOT20-07
                      MOT20-08
                      '''
        data_root = os.path.join('/home/hust/yly/Dataset', 'MOT20/test')
    seqs = [seq.strip() for seq in seqs_str.split()]

    main(opt,
         data_root=data_root,
         seqs=seqs,
         exp_name='mot17_half_yolov5s',
         show_image=False,
         save_images=False,
         save_videos=False)
