import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from dnt.filter import Filter
from dnt.track import Tracker
from dnt.detect import Detector
from dnt.label.labeler2 import Labeler
from shapely import LineString, Polygon
import os, pandas as pd
import glob
import numpy as np
from tssa.engine.zone import match_zone_by_point
from tssa.core.count import Count
from tssa.core.speed import SpeedMeterbyLine
from tssa.io.track_file import read_track


base_path = '/mnt/e/25-120037/25-120037-001-NB'
input_videos = []
input_videos.extend(glob.glob(os.path.join(base_path, '02.04', '*.mp4'), recursive=True))
input_videos.extend(glob.glob(os.path.join(base_path, '02.05', '*.mp4'), recursive=True))
input_videos.extend(glob.glob(os.path.join(base_path, '02.06', '*.mp4'), recursive=True))

track_path = os.path.join(base_path, 'tracks')
os.makedirs(track_path, exist_ok=True)
veh_path = os.path.join(base_path, 'vehs')
os.makedirs(veh_path, exist_ok=True)
label_path = os.path.join(base_path, 'labels_veh')
os.makedirs(label_path, exist_ok=True)

side_coords = [
    np.array([[1, 236], [233, 174], [481, 132], [655, 111], [980, 100], [1202, 115], [1274, 124], [1277, 630], [497, 630], [8, 630]])
] 
side_zone = Polygon(side_coords[0])

spd_coords = [
    np.array([[473, 126], [468, 716]]),
    np.array([[631, 110], [1184, 716]])
]
spd_lines = [LineString(line) for line in spd_coords]
spd_dist = [11]

cnt_line = spd_lines[0]

counter = Count(line=cnt_line)

cnt = 1
tot = len(input_videos)
for input_video in input_videos:

    print(f'Processing {cnt}/{tot}: {input_video}')
    # Generate paths
    basename = os.path.splitext(os.path.basename(input_video))[0]
    veh_track_file = os.path.join(track_path, basename+'_veh_track.txt')
    fps = Detector.get_fps(input_video)

    if not os.path.exists(veh_track_file):
        continue
    if os.stat(veh_track_file).st_size == 0:
        continue
    tracks = read_track(veh_track_file)
    tracks = match_zone_by_point(tracks, side_zone, ref_point='bc')
    tracks = tracks[tracks['in_zone'] == 1].copy().drop(columns=['in_zone'])
    if tracks is None:
        continue
    
    count_file = os.path.join(veh_path, basename+'_veh_count.csv')
    veh_count = counter.count_tracks_by_line(tracks=tracks)
    veh_count.to_csv(count_file, index=False)

    speeder = SpeedMeterbyLine(speed_lines=spd_lines,
                           line_distances_in_ft= spd_dist,
                           fps=fps,
                           ref_point='bbox',
                           ref_offset=(0,0),
                           ref_buffer=5,
                           speed_unit='mph',
                           speed_field='speed')
    speed_file = os.path.join(veh_path, basename+'_veh_speed.csv')
    speeder.detect(tracks=tracks, out_file=speed_file)
    
    cnt += 1
'''
ids = cross_tracks['track'].unique()
for id in tqdm(ids):
    selected = cross_tracks[cross_tracks['track']==id]
    start_frame = selected['frame'].min()
    end_frame = selected['frame'].max()
    out_video = os.path.join(label_path, str(id)+'.mp4')
    labeler.draw_tracks(input_video=input_video, output_video=out_video, tracks = selected, 
                        start_frame=start_frame, end_frame=end_frame, verbose=False)
   
    plt.plot(*lines[0].xy)
    plt.plot(*lines[1].xy)
    plt.plot(*lines[2].xy)
    for index, s in selected.iterrows():
        plt.plot(*Polygon([(s[2], s[3]), (s[2] + s[4], s[3]), 
                                (s[2] + s[4], s[3] + s[5]), (s[2], s[3] + s[5])]).exterior.xy)
    plt.savefig(os.path.join(pathname, 'frames', str(id)+'.jpg'))
    plt.close()
'''
print('ok')