import os
import sys
import pandas as pd
import geopandas as gpd
from shapely.geometry import Polygon, Point, LineString
from label_zone import draw
import cv2
from tqdm import tqdm
import numpy as np
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from tssa.engine.line import match_line_by_point

line = LineString([[852, 937], [1500, 743]])
input_video = '/mnt/d/videos/sample/0180-cctv.mp4'
output_video = '/mnt/d/videos/sample/labels/0180-cctv_line.mp4'
tracks = pd.read_csv('/mnt/d/videos/sample/tracks/0180-cctv_track.txt', header=None)
tracks.columns = ['frame', 'track_id', 'x', 'y', 'w', 'h', 'cls', 'r2', 'r3', 'r4']

matched_tracks = match_line_by_point(tracks, line, ref_point='br', ref_buffer=10)
matched_tracks.to_csv('/mnt/d/videos/sample/lanes/0180-cctv_tracks_line.csv', index=False)

'''
geo = tracks.apply(lambda track: Polygon([[track['x'], track['y']], [track['x'] + track['w'], track['y']],
                                        [track['x']+track['w'], track['y']+track['h']], 
                                          [track['x'], track['y']+track['h']]]),
                                        axis=1)
geo_tracks = gpd.GeoDataFrame(tracks, geometry=geo)

match_cnt = 0
for index, track in tracks.iterrows():
    p = Polygon([[track['x'], track['y']], [track['x'] + track['w'], track['y']],
                                        [track['x']+track['w'], track['y']+track['h']], 
                                          [track['x'], track['y']+track['h']]])
    if line.intersects(p):
      match_cnt += 1
    print(match_cnt)

print('ok')

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

pbar = tqdm(total=tot_frames, unit=" frames", desc="Generating labels ...")
results = []
for i in range(tot_frames):
    tracks_at_frame_i = tracks.loc[tracks['frame']==i]
    if len(tracks_at_frame_i) > 0:
        for index, track in tracks_at_frame_i.iterrows():
            final_color = [0, 255, 0]
            label_str = '#'+str(int(track['track_id'])) + ':' + str(int(track['frame']))
            results.append([track['frame'],
                        'bbox',
                        [(track['x'], track['y']),
                         (track['x'] + track['w'], 
                          track['y'] + track['h'])],
                        final_color, 
                        0.5,
                        1,
                        1.0,
                        label_str])
    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)
'''