import os
import sys
import numpy as np
import cv2
import pandas as pd
from tqdm import tqdm
import distinctipy
from dnt.track import Tracker
from shapely.geometry import Polygon, LineString
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from tssa.core.speed import SpeedMeterbyLine
from tssa.io.track_file import read_track
from label_zone import draw

ref_lines_coords = [
    np.array([[107, 583], [567, 545]]),
    np.array([[180, 614], [656, 566]]),
    np.array([[271, 656], [796, 589]]),
    np.array([[382, 706], [930, 624]]),
    np.array([[1153, 668], [535, 774]]),
    np.array([[766, 895], [1380, 727]]),
    np.array([[1194, 1076], [1755, 812]])
    ]
speed_lines = [LineString(coords) for coords in ref_lines_coords]
red_dists = [
    40.37,
    39.67,
    40.03,
    39.32,
    41.42,
    40.03]

input_video = '/mnt/d/videos/sample/0180-cctv.mp4'
track_file ='/mnt/d/videos/sample/tracks/0180-cctv_track.txt'
speed_file = '/mnt/d/videos/sample/speed/0180-cctv_speed_lines.csv'
output_video = '/mnt/d/videos/sample/labels/0180-cctv_line_speed.mp4'

cap = cv2.VideoCapture(input_video)
if not cap.isOpened():
    raise IOError("Couldn't open webcam or video")
fps = int(cap.get(cv2.CAP_PROP_FPS))
tot_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
cap.release()

tracks = read_track(track_file)
speeder = SpeedMeterbyLine(speed_lines=speed_lines,
                           line_distances_in_ft= red_dists,
                           fps=30,
                           ref_point='bbox',
                           ref_offset=(0,0),
                           ref_buffer=5,
                           speed_unit='mph',
                           speed_field='speed')

speed_tracks = speeder.detect(tracks=tracks, out_file=speed_file)

colors_rbg = distinctipy.get_colors(10)
colors = [color[::-1] for color in colors_rbg]
colors = [(int(x[0]*255), int(x[1]*255), int(x[2]*255)) for x in colors]

pbar = tqdm(total=tot_frames, unit=" frames", desc="Generating labels ...")
results = []
for i in range(tot_frames):
    tracks_at_frame_i = speed_tracks.loc[speed_tracks['frame']==i]
    if len(tracks_at_frame_i) > 0:
        for index, track in tracks_at_frame_i.iterrows():
            final_color = colors[int(track['track_id']%10)]
            label_str = str(int(track['speed']))+" mph"
            results.append([track['frame'],
                        'bbox',
                        [(track['x'], track['y']),
                         (track['x'] + track['w'], 
                          track['y'] + track['h'])],
                        final_color, 
                        1.2,
                        2,
                        1.0,
                        label_str])
    for line in speed_lines:
            results.append([i, 
                    'polygon',
                    np.array(list(line.coords), np.int32).reshape((-1, 1, 2)),
                    [0, 0, 0],
                    1,
                    1,
                    1,
                    ""])
    pbar.update()
pbar.close()

df = pd.DataFrame(results, columns=['frame','type','coords','color','size','thick','alpha','desc'])
df.sort_values(by='frame', inplace=True)
df.to_csv(output_video.replace('.mp4', '.csv'), index=False)

draw(input_video = input_video,
         output_video = output_video,
         draws = df)

