import os
import sys
import pandas as pd
import numpy as np
from tqdm import tqdm
from shapely.geometry import LineString, Polygon
import cv2
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from tssa.core.count import Count
from tssa.io.definition import track_fields
from label_zone import draw

input_video = '/mnt/d/videos/sample/0180-cctv.mp4'
track_file = '/mnt/d/videos/sample/tracks/0180-cctv_track.txt'
lane_track_file ='/mnt/d/videos/sample/lanes/0180-cctv_track_zones.csv'
count_file = '/mnt/d/videos/sample/count/0180-cctv_count.csv'
output_video = '/mnt/d/videos/sample/labels/0180-cctv_zones_count.mp4'

'''
count_line = LineString(np.array([[642, 938], [1146, 653]]))
lanes = [
    Polygon(np.array([[403, 716], [477, 678], [1253, 1016], [1123, 1078]])),
    Polygon(np.array([[481, 677], [535, 652], [1359, 957], [1253, 1016], [1253, 1016]])),
    Polygon(np.array([[538, 651], [577, 634], [1430, 909], [1363, 955]])),
    Polygon(np.array([[575, 634], [627, 610], [1530, 836], [1430, 907]])),
    Polygon(np.array([[660, 593], [1541, 817], [1610, 765], [689, 569]]))]
'''
count_line = LineString(np.array([[535, 852], [1006, 639]]))
lanes = [
    Polygon(np.array([[110, 533], [532, 706], [973, 895], [1294, 1033], [1214, 1077], [1123, 1076], [1123, 1076], [464, 742], [67, 551]])),
    Polygon(np.array([[134, 515], [391, 608], [817, 757], [821, 759], [1392, 965], [1295, 1031], [970, 895], [525, 704], [107, 531]])),
    Polygon(np.array([[172, 508], [917, 744], [1478, 915], [1392, 965], [817, 757], [387, 606], [132, 513]])),
    Polygon(np.array([[193, 497], [1597, 847], [1481, 913], [917, 742], [168, 505]])),
    Polygon(np.array([[233, 477], [1680, 785], [1597, 845], [585, 594], [192, 495]]))
]

counter = Count(line=count_line)
veh_count = counter.count_tracks_by_line(track_file=track_file, with_header=False)

lane_tracks = pd.read_csv(lane_track_file)
count_track_ids = veh_count['track_id'].unique()
lane_track_ids = lane_tracks['track_id'].unique()
final_ids = [int(id) for id in count_track_ids if id in lane_track_ids]

results = []
for id in final_ids:
    cnt_frame = veh_count.loc[veh_count['track_id']==id]['frame'].values[0]
    cnt = veh_count.loc[veh_count['track_id']==id]['count'].values[0]
    direction = veh_count.loc[veh_count['track_id']==id]['direction'].values[0]
    selected_lane_tracks = lane_tracks.loc[lane_tracks['track_id']==id].copy()
    selected_lane_tracks['diff'] = selected_lane_tracks['frame'].apply(lambda x: abs(x-cnt_frame))
    selected_lane_tracks.sort_values(by='diff', inplace=True)
    lane = selected_lane_tracks.iloc[0]['lane']
    lane_frame = selected_lane_tracks.iloc[0]['frame']
    ref_x = selected_lane_tracks.iloc[0]['ref_x'] 
    ref_y = selected_lane_tracks.iloc[0]['ref_y']
    results.append([id, cnt_frame, cnt, direction, lane, lane_frame, ref_x, ref_y])
df = pd.DataFrame(results, columns=['track_id', 'frame', 'count', 'direction', 'lane', 'lane_frame', 'ref_x', 'ref_y'])
df.to_csv(count_file, index=False)
veh_count = pd.read_csv(count_file)

tracks = pd.read_csv('/mnt/d/videos/sample/tracks/0180-cctv_track.txt', header=None)
tracks.columns = track_fields
tracks = tracks.loc[tracks['track_id'].isin(count_track_ids)].copy()

colors = [[255, 0, 0],
          [0, 255, 0],
          [128, 0, 128],
          [0, 127, 255],
          [255, 255, 0]
        ]
points = [(874, 897),
          (962, 842),
          (1027, 803),
          (1096, 756),
          (1182, 707)]

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 = []
cnt = 0
cnts = [0, 0, 0, 0, 0]
for i in range(tot_frames):
    tracks_at_frame_i = lane_tracks.loc[lane_tracks['frame']==i]
    if len(tracks_at_frame_i) > 0:
        for index, track in tracks_at_frame_i.iterrows():
            final_color = colors[int(track['lane'])]
            label_str = '#'+str(int(track['track_id']))
            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([track['frame'],
                        'circle',
                        [(track['ref_x'], track['ref_y'])],
                        [0, 0, 255], 
                        3,
                        -1,
                        1.0,
                        ""])
    for j in range(len(lanes)):
        lane = lanes[j]
        results.append([i, 
                        'text',
                        [(points[j][0], points[j][1])], 
                        colors[j],
                        1,
                        2,
                        1,
                        "L"+str(j+1)])
        results.append([i, 
                        'polygon',
                        np.array(list(lane.exterior.coords), np.int32).reshape((-1, 1, 2)),
                        [0, 0, 0],
                        1,
                        1,
                        1,
                        ""])
        results.append([i, 
                        'filled_polygon',
                        np.array(list(lane.exterior.coords), np.int32).reshape((-1, 1, 2)),
                        [218, 177, 218],
                        1,
                        1,
                        0.2, 
                        ""])
    veh_count_at_frame_i = veh_count.loc[veh_count['frame']==i]
    if len(veh_count_at_frame_i) > 0:
        cnt = veh_count_at_frame_i.iloc[-1]['count']
        for index, t in veh_count_at_frame_i.iterrows():
            if t['lane'] == 0:
                cnts[0] = cnts[0] + 1
            elif t['lane'] == 1:
                cnts[1] = cnts[1] + 1
            elif t['lane'] == 2:
                cnts[2] = cnts[2] + 1
            elif t['lane'] == 3:
                cnts[3] = cnts[3] + 1
            elif t['lane'] == 4:
                cnts[4] = cnts[4] + 1
 
    results.append([i,
                    'text',
                    [(1200, 100)],
                    [0, 0, 0],
                    2,
                    3,
                    1,
                    "NB Count:"+str(int(cnt))])
    results.append([i,
                    'line',
                    np.array(list(count_line.coords), np.int32),
                    [0, 0, 255],
                    1,
                    1,
                    1,
                    ""])
    for j in range(len(lanes)):
        lane = lanes[j]
        results.append([i, 
                        'text',
                        [(points[j][0]+ 100, points[j][1] + 30 - ((j-5)*5))], 
                        [0, 0, 0],
                        1,
                        2,
                        1,
                        str(cnts[j])])      
    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)
