import os
import sys
import numpy as np
import cv2
import pandas as pd
from tqdm import tqdm
import distinctipy
from shapely.geometry import Polygon
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 label_zone import draw

source = np.float32([[199, 611],
                   [510, 581],
                   [1336, 786],
                   [842, 941]])
destination = np.float32([[0, 0],
                       [49, 0],
                       [55, 170],
                       [0, 170]])
transformer = ViewTransform(source, destination)

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_zones_speed.mp4'


tracks = pd.read_csv(track_file)
speed_zone = Polygon(source)
matched_tracks = match_zone_by_point(tracks, speed_zone, ref_point='br', output_ref_xy=True)
matched_tracks = matched_tracks.loc[matched_tracks['in_zone']==True].copy()
matched_tracks[['anchor_x', 'anchor_y']] = \
                matched_tracks.apply(lambda x:
                                tuple(transformer.transform(
                                    np.float32([x['ref_x'], x['ref_y']]))[0]
                                    ), 
                                axis=1,
                                result_type='expand') 
track_ids = matched_tracks['track_id'].unique()

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 = []
for id in tqdm(track_ids):
    track_seq = matched_tracks.loc[matched_tracks['track_id']==id].copy().sort_values(by='frame')
    track_seq['frame_diff'] = track_seq['frame'].diff()
    track_seq['x_diff'] = track_seq['anchor_x'].diff()
    track_seq['y_diff'] = track_seq['anchor_y'].diff()
    track_seq.dropna(inplace=True)
    track_seq['distance'] = np.sqrt(track_seq['x_diff']**2 + track_seq['y_diff']**2)
    track_seq['time_diff'] = track_seq['frame_diff']/fps
    track_seq['speed'] = track_seq['distance']/track_seq['time_diff']
    track_seq['rolling_speed'] = track_seq['speed'].rolling(window=30, min_periods=1).mean().interpolate(method='nearest')
    track_seq['mph'] = (track_seq['rolling_speed']*0.681818).round(1)

    results.append(track_seq)

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

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])
            results.append([track['frame'],
                        'circle',
                        [(track['ref_x'], track['ref_y'])],
                        [0, 0, 255], 
                        3,
                        -1,
                        1.0,
                        ""])
    results.append([i, 
                    'polygon',
                    np.array(list(speed_zone.exterior.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)


