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.engine import ViewTransform
from tssa.engine.zone import match_zone_by_point
from tssa.io.track_file import read_track
from tssa.engine.line import match_lines_by_point, match_lines_by_bbox
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'
output_video = '/mnt/d/videos/sample/labels/0180-cctv_line_speed.mp4'

tracks = read_track(track_file)
tracks = Tracker.infill_frames(tracks) 
matched_tracks = match_lines_by_bbox(tracks=tracks,
                                      lines=speed_lines,
                                      output_ref_xy=True)
matched_tracks = matched_tracks.loc[matched_tracks['line']>-1].copy()
first_indices = matched_tracks.groupby(['track_id', 'line'])['frame'].idxmin()
intersect_frames = matched_tracks.loc[first_indices].sort_values(by=['track_id','line'])   

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()

results = []
track_ids = intersect_frames['track_id'].unique()
for id in track_ids:
    selected = intersect_frames.loc[intersect_frames['track_id']==id].copy()
    selected['frame_diff'] = selected['frame'].diff()
    selected['time_diff'] = selected['frame_diff'] / fps
    selected['line_diff'] = selected['line'].diff()
    selected['distance'] = selected['line_diff']*40
    selected['speed'] = selected['distance']/selected['time_diff']
    selected.dropna(inplace=True)
    selected['mph'] = selected['speed'].apply(lambda x: int(x*0.681818))
    results.append(selected)

speed_tracks = pd.concat(results)
speed_tracks.to_csv('/mnt/d/videos/sample/speed/0180-cctv_speed_lines.csv', index=False)

# merge into tracks
results = []

for id in track_ids:
    selected = speed_tracks.loc[speed_tracks['track_id']==id].sort_values(by='line').copy()
    frame_ids = selected['frame'].unique()
    for j in range(len(frame_ids)-1):
        beg_frame = frame_ids[j]
        end_frame = frame_ids[j+1]-1
        matched_track_frames = tracks.loc[(tracks['track_id']==id) &
                                          (tracks['frame']>=beg_frame) &
                                    (tracks['frame']<=end_frame)].copy()
        matched_track_frames['line'] = int(selected.loc[selected['frame']==beg_frame]['line'].values[0])
        matched_track_frames['mph'] = int(selected.loc[selected['frame']==beg_frame]['mph'].values[0])
        
        results.append(matched_track_frames)

speed_tracks = pd.concat(results)
speed_tracks.to_csv('/mnt/d/videos/sample/speed/0180-cctv_track_lines_speed.csv', index=False)
'''
speed_tracks = matched_tracks
speed_tracks['mph'] = 0
'''
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['mph']))+" 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)


