import pandas as pd
import numpy as np
import cv2, os
import scipy
from scipy import signal
import csv


input_source = "/home/stephen/Desktop/me.MP4"
cap = cv2.VideoCapture(input_source)
out = cv2.VideoWriter('/home/stephen/Desktop/test.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 60, (400,700))
frame_number = 0

# there are 15 points in the skeleton
# 0 head
# 1 neck
# 2, 5 shoulders
# 3, 6 elbows
# 4, 7 hands
# 8, 11 hips
# 9, 12 knees
# 10, 13 ankles
# 14 torso

pairs = [[0,1], #head to neck
         [1,2],[1,5], #neck to sholders
         [2,3],[5,6], #sholders to elbows
         [3,4],[6,7], #elbows to hands
         [1,14], #neck to torso
         [14,8],[14,11], #torso (14) to hips (8 and 11)
         [8,9],[11,12],#hips to knees
         [9,10],[12,13]]#knees to feet


#get pose data - data is generated by open pose video
df = pd.read_csv('/home/stephen/Desktop/swapped_body_parts.csv')

#smooth it out
for i in range(30): df[str(i)] = signal.savgol_filter(df[str(i)], 5, 3)

frame_number = 0
while cv2.waitKey(1000) < 0:

    hasFrame, frame = cap.read()
    frameCopy = np.copy(frame)
    #frame = np.zeros_like(frame)
    #frame[:,:,:] = 123,123,123
    if not hasFrame: break

    values = df.values[frame_number]
    values = np.array(values, int)

    points = []
    points.append((values[0], values[1]))
    points.append((values[2], values[3]))
    points.append((values[4], values[5]))
    points.append((values[6], values[7]))
    points.append((values[8], values[9]))
    points.append((values[10], values[11]))
    points.append((values[12], values[13]))
    points.append((values[14], values[15]))
    points.append((values[16], values[17]))
    points.append((values[18], values[19]))
    points.append((values[20], values[21]))
    points.append((values[22], values[23]))
    points.append((values[24], values[25]))
    points.append((values[26], values[27]))
    points.append((values[28], values[29]))

    # Draw Skeleton
    for pair in pairs:
        partA = pair[0]
        partB = pair[1]
        cv2.line(frame, points[partA], points[partB], (32, 255, 0), 2, lineType=cv2.LINE_AA)
        cv2.circle(frame, points[partA], 2, (0, 0, 255), thickness=-1, lineType=cv2.FILLED)
        cv2.circle(frame, points[partB], 2, (0, 0, 255), thickness=-1, lineType=cv2.FILLED)

    cv2.imshow('Output-Skeleton', frame)
    out.write(frame)
    frame_number+=1