import pandas as pd
import numpy as np
import cv2, os
import csv

input_source = "/home/stephen/Desktop/me.MP4"
cap = cv2.VideoCapture(input_source)
frame_number = 0
font, scale, colorText, thick = cv2.FONT_HERSHEY_SIMPLEX, .5, (234,234,234), 1
size, color, thickness = 5, (255,255,255), 5
#get pose data - data is generated by open pose video
df = pd.read_csv('/home/stephen/Desktop/open_pose_data.csv')

# 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

data = []
while cv2.waitKey(10) < 0 and frame_number<len(df.values)-2:
    ret, img = cap.read()
    if not ret: break
    try: values = df.values[frame_number]
    except: break
    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]))


    #create a blank list to store the non-swapped poitns
    non_swap_points = []
    for i in range(15): non_swap_points.append((0,0))
    #add the head, that point never changes
    non_swap_points[0] = points[0]
    #add the neck, that point never changes
    non_swap_points[1] = points[1]
    #add the torso, that never changes
    non_swap_points[14] = points[14]

    #swap the left and right shoulders (2 and 5)
    if points[2][0]<points[5][0]:
        non_swap_points[2] = points[2]
        non_swap_points[5] = points[5]
    else:
        non_swap_points[2] = points[5]
        non_swap_points[5] = points[2]
    # swap the elbows
    if points[3][0]<points[6][0]:
        non_swap_points[3] = points[3]
        non_swap_points[6] = points[6]
    else:
        non_swap_points[6] = points[3]
        non_swap_points[3] = points[6]

    #swap the hands
    if points[4][0]<points[7][0]:
        non_swap_points[4] = points[4]
        non_swap_points[7] = points[7]
    else:
        non_swap_points[7] = points[4]
        non_swap_points[4] = points[7]
    #swap the hips
    if points[8][0]<points[11][0]:
        non_swap_points[11] = points[11]
        non_swap_points[8] = points[8]
    else:
        non_swap_points[8] = points[11]
        non_swap_points[11] = points[8]

    #swap the knees
    if points[9][0]<points[12][0]:
        non_swap_points[9] = points[9]
        non_swap_points[12] = points[12]
    else:
        non_swap_points[12] = points[9]
        non_swap_points[9] = points[12]

    #swap the feet
    if points[10][0]<points[13][0]:
        non_swap_points[10] = points[10]
        non_swap_points[13] = points[13]
    else:
        non_swap_points[13] = points[10]
        non_swap_points[10] = points[13]


    for point in non_swap_points:
        cv2.circle(img, point, 3, (0, 0, 255), 3)
        cv2.putText(img, str(non_swap_points.index(point)), point, font, scale, colorText, thick, cv2.LINE_AA)

    cv2.imshow('Output-Skeleton', img)
    frame_number+=1
    data.append(non_swap_points)
cv2.destroyAllWindows()


with open('/home/stephen/Desktop/swapped_body_parts.csv', 'w') as csvfile:
    fieldnames = []
    for i in range(30): fieldnames.append(str(i))
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()

    for trick in data:
        writer.writerow({'0': trick[0][0],
                         '1':trick[0][1],
                         '2':trick[1][0],
                         '3':trick[1][1],
                         '4':trick[2][0],
                         '5':trick[2][1],
                         '6':trick[3][0],
                         '7':trick[3][1],
                         '8':trick[4][0],
                         '9':trick[4][1],
                         '10':trick[5][0],
                         '11':trick[5][1],
                         '12':trick[6][0],
                         '13':trick[6][1],
                         '14':trick[7][0],
                         '15':trick[7][1],
                         '16':trick[8][0],
                         '17':trick[8][1],
                         '18':trick[9][0],
                         '19':trick[9][1],
                         '20':trick[10][0],
                         '21':trick[10][1],
                         '22':trick[11][0],
                         '23':trick[11][1],
                         '24':trick[12][0],
                         '25':trick[12][1],
                         '26':trick[13][0],
                         '27':trick[13][1],
                         '28':trick[14][0],
                         '29':trick[14][1]})